Coverage for src/lexigram/graphql/dataloader/batch.py: 91%

1"""Batch loading utilities. 

2 

3This module provides utilities for defining and executing 

4batch load functions. 

5""" 

6 

7from __future__ import annotations 

8 

9import asyncio 

10from collections.abc import Awaitable, Callable, Hashable 

11from typing import ( 

12 Any, 

13 Generic, 

14 TypeVar, 

15) 

16 

17from lexigram.graphql.exceptions import GraphQLError 

18from lexigram.logging import get_logger 

19 

20logger = get_logger(__name__) 

21 

22K = TypeVar("K", bound=Hashable) 

23V = TypeVar("V") 

24 

25 

26# Type alias for batch functions 

27BatchFunction = Callable[[list[K]], Awaitable[list[V | None]]] 

28 

29 

30def batch_load( 

31 fn: ( 

32 Callable[[list[K]], Awaitable[list[V]]] 

33 | Callable[[list[K]], Awaitable[dict[K, V]]] 

34 ), 

35 key_fn: Callable[[V], K] | None = None, 

36) -> BatchFunction[K, V]: 

37 """Create a batch load function from a loader. 

38 

39 This decorator/wrapper helps create properly ordered batch 

40 functions from common data loading patterns. 

41 

42 Args: 

43 fn: Function that loads data by keys. 

44 key_fn: Optional function to extract key from value. 

45 

46 Returns: 

47 Batch function that returns values in input order. 

48 

49 Example: 

50 ```python 

51 # From a function returning a list 

52 @batch_load(key_fn=lambda u: u.id) 

53 async def load_users(ids: list[str]) -> list[User]: 

54 return await db.query(User).filter(User.id.in_(ids)).all() 

55 

56 # From a function returning a dict 

57 @batch_load 

58 async def load_users(ids: list[str]) -> dict[str, User]: 

59 users = await db.query(User).filter(User.id.in_(ids)).all() 

60 return {u.id: u for u in users} 

61 ``` 

62 """ 

63 

64 async def wrapper(keys: list[K]) -> list[V | None]: 

65 try: 

66 result = await fn(keys) 

67 

68 # Handle dict result 

69 if isinstance(result, dict): 

70 return [result.get(key) for key in keys] 

71 

72 # Handle list result - need key_fn to map 

73 if key_fn is not None: 

74 value_map = {key_fn(v): v for v in result} 

75 return [value_map.get(key) for key in keys] 

76 

77 # If no key_fn and list returned, assume same order 

78 if len(result) == len(keys): 

79 return list(result) 

80 

81 raise GraphQLError("Cannot map batch results to keys without key_fn") 

82 

83 except Exception as _batch_err: # noqa: BLE001 — batch wrapper must log any error from user-supplied batch functions before re-raising 

84 logger.exception("Batch load error") 

85 raise 

86 

87 return wrapper 

88 

89 

90class BatchScheduler(Generic[K, V]): 

91 """Scheduler for batching async operations. 

92 

93 Collects operations and executes them in batches 

94 for efficiency. 

95 

96 Example: 

97 ```python 

98 scheduler = BatchScheduler( 

99 batch_fn=batch_load_users, 

100 batch_size=50, 

101 batch_delay_ms=10, 

102 ) 

103 

104 # Schedule loads 

105 user1 = await scheduler.schedule("1") 

106 user2 = await scheduler.schedule("2") 

107 ``` 

108 """ 

109 

110 def __init__( 

111 self, 

112 batch_fn: BatchFunction[K, V], 

113 batch_size: int = 100, 

114 batch_delay_ms: float = 2.0, 

115 ) -> None: 

116 """Initialize the scheduler. 

117 

118 Args: 

119 batch_fn: Batch load function. 

120 batch_size: Maximum batch size. 

121 batch_delay_ms: Delay in milliseconds before executing the batch. 

122 A small non-zero value (default 2ms) allows more keys to accumulate, 

123 improving efficiency at a slight latency cost. 

124 """ 

125 self._batch_fn = batch_fn 

126 self._batch_size = batch_size 

127 self._batch_delay_ms = batch_delay_ms 

128 

129 self._pending: list[tuple[K, asyncio.Future[V | None]]] = [] 

130 self._scheduled = False 

131 self._background_tasks: set[asyncio.Task[Any]] = set() 

132 

133 async def schedule(self, key: K) -> V | None: 

134 """Schedule a load operation. 

135 

136 Args: 

137 key: Key to load. 

138 

139 Returns: 

140 Loaded value. 

141 """ 

142 future: asyncio.Future[V | None] = asyncio.Future() 

143 self._pending.append((key, future)) 

144 

145 # Schedule batch execution 

146 if not self._scheduled: 

147 self._scheduled = True 

148 

149 if self._batch_delay_ms > 0: 

150 await asyncio.sleep(self._batch_delay_ms / 1000) 

151 

152 task = asyncio.create_task(self._execute()) 

153 self._background_tasks.add(task) 

154 task.add_done_callback(self._background_tasks.discard) 

155 

156 # Execute immediately if batch is full 

157 if len(self._pending) >= self._batch_size: 

158 await self._execute() 

159 

160 return await future 

161 

162 async def _execute(self) -> None: 

163 """Execute pending batch.""" 

164 self._scheduled = False 

165 

166 if not self._pending: 

167 return 

168 

169 # Get pending items 

170 items = list(self._pending) 

171 self._pending = [] 

172 

173 keys = [item[0] for item in items] 

174 futures = [item[1] for item in items] 

175 

176 try: 

177 values = await self._batch_fn(keys) 

178 

179 for future, value in zip(futures, values, strict=True): 

180 if not future.done(): 

181 future.set_result(value) 

182 

183 except (RuntimeError, TypeError, ValueError) as e: 

184 # Set exception on ALL pending futures that haven't been resolved 

185 for future in futures: 

186 if not future.done(): 

187 future.set_exception(e) 

188 

189 

190__all__ = ["BatchFunction", "BatchScheduler", "batch_load"]