Coverage for src/lexigram/graphql/dataloader/loader.py: 94%

1"""DataLoaderProtocol implementation. 

2 

3This module provides a DataLoaderProtocol implementation for batching 

4and caching data fetches in GraphQL resolvers to solve the N+1 problem. 

5""" 

6 

7from __future__ import annotations 

8 

9import asyncio 

10from collections.abc import Awaitable, Callable, Hashable, Sequence 

11from typing import ( 

12 Generic, 

13 TypeVar, 

14) 

15 

16from lexigram.graphql.config import DataLoaderConfig 

17from lexigram.graphql.dataloader.cache import InMemoryCache, LoaderCache 

18from lexigram.logging import get_logger 

19 

20logger = get_logger(__name__) 

21 

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

23V = TypeVar("V") 

24 

25 

26class DataLoaderProtocol(Generic[K, V]): 

27 """DataLoaderProtocol for batching and caching. 

28 

29 DataLoaderProtocol is a utility for batching and caching data fetches 

30 to efficiently resolve GraphQL queries and avoid the N+1 problem. 

31 

32 Example: 

33 ```python 

34 async def batch_load_users(ids: list[str]) -> list[User | None]: 

35 users = await db.get_users_by_ids(ids) 

36 # Return in same order as input ids 

37 user_map = {u.id: u for u in users} 

38 return list(map(lambda id: user_map.get(id), ids)) 

39 

40 loader = DataLoaderProtocol(batch_load_users) 

41 

42 # These will be batched into a single call 

43 user1 = await loader.load("1") 

44 user2 = await loader.load("2") 

45 ``` 

46 """ 

47 

48 def __init__( 

49 self, 

50 batch_fn: Callable[[list[K]], Awaitable[list[V | None]]], 

51 config: DataLoaderConfig | None = None, 

52 cache: LoaderCache[K, V] | None = None, 

53 ) -> None: 

54 """Initialize the DataLoaderProtocol. 

55 

56 Args: 

57 batch_fn: Function that loads a batch of values by keys. 

58 config: DataLoaderProtocol configuration. 

59 cache: Optional cache implementation. 

60 """ 

61 self._batch_fn = batch_fn 

62 self._config = config or DataLoaderConfig() 

63 self._cache = cache if cache is not None else InMemoryCache[K, V]() 

64 

65 # Thread-safe state 

66 self._lock = asyncio.Lock() 

67 self._pending: dict[K, asyncio.Future[V | None]] = {} 

68 self._dispatch_task: asyncio.Task[None] | None = None 

69 

70 @property 

71 def config(self) -> DataLoaderConfig: 

72 """Get configuration.""" 

73 return self._config 

74 

75 async def load(self, key: K) -> V | None: 

76 """Load a value by key. 

77 

78 The load will be batched with other loads that occur 

79 in the same tick of the event loop. 

80 

81 Args: 

82 key: The key to load. 

83 

84 Returns: 

85 The loaded value or None if not found. 

86 

87 Raises: 

88 DataLoaderError: If loading fails. 

89 """ 

90 # Check cache first 

91 if self._config.cache_enabled and self._cache.has(key): 

92 return self._cache.get(key) 

93 

94 # Acquire lock for pending dict 

95 async with self._lock: 

96 # Check if already pending 

97 if key in self._pending: 

98 future = self._pending[key] 

99 else: 

100 # Create new future 

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

102 self._pending[key] = future 

103 

104 # Schedule dispatch if not already scheduled 

105 if self._dispatch_task is None or self._dispatch_task.done(): 

106 self._dispatch_task = asyncio.create_task(self._dispatch_batch()) 

107 logger.debug("DataLoaderProtocol dispatch task scheduled") 

108 

109 # Wait for result (outside lock) 

110 return await future 

111 

112 async def load_many(self, keys: Sequence[K]) -> list[V | None]: 

113 """Load multiple values by keys. 

114 

115 Args: 

116 keys: List of keys to load. 

117 

118 Returns: 

119 List of values in the same order as keys. 

120 """ 

121 return await asyncio.gather(*[self.load(key) for key in keys]) 

122 

123 async def _dispatch_batch(self) -> None: 

124 """Dispatch the pending batch.""" 

125 # Extract pending items (thread-safe) 

126 async with self._lock: 

127 if not self._pending: 

128 return 

129 

130 # Get all pending keys 

131 keys = list(self._pending.keys()) 

132 futures = dict(self._pending) 

133 

134 # Clear pending (others can add new ones) 

135 self._pending.clear() 

136 

137 logger.debug("DataLoaderProtocol dispatching batch of %d keys", len(keys)) 

138 

139 # Respect max batch size (outside lock) 

140 if self._config.max_batch_size > 0: 

141 # Split into chunks if needed 

142 for i in range(0, len(keys), self._config.max_batch_size): 

143 chunk_keys = keys[i : i + self._config.max_batch_size] 

144 chunk_futures = {k: futures[k] for k in chunk_keys} 

145 await self._execute_batch(chunk_keys, chunk_futures) 

146 else: 

147 await self._execute_batch(keys, futures) 

148 

149 async def _execute_batch( 

150 self, 

151 keys: list[K], 

152 futures: dict[K, asyncio.Future[V | None]], 

153 ) -> None: 

154 """Execute a batch load. 

155 

156 Args: 

157 keys: Keys to load in this batch. 

158 futures: Futures to resolve for these keys. 

159 """ 

160 try: 

161 # Execute batch function 

162 values = await self._batch_fn(keys) 

163 

164 # Validate response 

165 if len(values) != len(keys): 

166 raise ValueError( 

167 f"batch_fn returned {len(values)} values for {len(keys)} keys. " 

168 f"Must return same number of values as keys (use None for missing).", 

169 ) 

170 

171 # Resolve futures and cache results 

172 for key, value in zip(keys, values, strict=True): 

173 if self._config.cache_enabled and value is not None: 

174 self._cache.set(key, value) 

175 

176 future = futures[key] 

177 if not future.done(): 

178 future.set_result(value) 

179 

180 logger.debug("DataLoaderProtocol batch completed successfully") 

181 

182 except Exception as e: # noqa: BLE001 — dataloader must propagate any batch error to all pending futures 

183 # Reject all pending futures 

184 logger.error("DataLoaderProtocol batch failed: %s", e, exc_info=True) 

185 

186 for future in futures.values(): 

187 if not future.done(): 

188 future.set_exception(e) 

189 

190 def prime(self, key: K, value: V) -> None: 

191 """Prime the cache with a value. 

192 

193 Use this to add values to the cache that were loaded 

194 through other means. 

195 

196 Args: 

197 key: The key. 

198 value: The value to cache. 

199 """ 

200 if self._config.cache_enabled: 

201 self._cache.set(key, value) 

202 

203 def clear(self, key: K) -> None: 

204 """Clear a cached value. 

205 

206 Args: 

207 key: The key to clear. 

208 """ 

209 self._cache.delete(key) 

210 

211 def clear_all(self) -> None: 

212 """Clear all cached values.""" 

213 self._cache.clear() 

214 

215 

216def create_loader( 

217 batch_fn: Callable[[list[K]], Awaitable[list[V | None]]], 

218 batch_size: int = 100, 

219 cache_enabled: bool = True, 

220) -> DataLoaderProtocol[K, V]: 

221 """Create a DataLoaderProtocol (convenience function). 

222 

223 Args: 

224 batch_fn: Batch load function. 

225 batch_size: Maximum batch size. 

226 cache_enabled: Whether to enable caching. 

227 

228 Returns: 

229 Configured DataLoaderProtocol. 

230 

231 Example: 

232 ```python 

233 user_loader = create_loader( 

234 batch_load_users, 

235 batch_size=50, 

236 ) 

237 

238 user = await user_loader.load("123") 

239 ``` 

240 """ 

241 config = DataLoaderConfig( 

242 max_batch_size=batch_size, 

243 cache_enabled=cache_enabled, 

244 ) 

245 return DataLoaderProtocol(batch_fn, config=config) 

246 

247 

248__all__ = ["DataLoaderProtocol", "create_loader"]