Coverage for src / lexigram / admin / data / data_loader.py: 0%

1"""Batch and cache data loading to avoid N+1 query problems.""" 

2 

3from __future__ import annotations 

4 

5import asyncio 

6from collections.abc import Awaitable, Callable 

7from typing import Any, Generic, TypeVar 

8 

9from lexigram.concurrency import Parallel 

10from lexigram.contracts.core import TaskManagerProtocol 

11from lexigram.di.decorators import inject 

12 

13K = TypeVar("K") # Key type 

14V = TypeVar("V") # Value type 

15 

16 

17@inject 

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

19 """ 

20 DataLoaderProtocol implements the batch-loading pattern. 

21 It collects requests for individual IDs and executes them in a single batch. 

22 """ 

23 

24 def __init__( 

25 self, 

26 batch_fn: Callable[[list[K]], Awaitable[dict[K, V]]], 

27 task_manager: TaskManagerProtocol, 

28 ): 

29 """ 

30 Initialize with a function that takes a list of keys and returns a dict of results. 

31 """ 

32 self._batch_fn = batch_fn 

33 self.task_manager = task_manager 

34 self._cache: dict[K, V] = {} 

35 self._pending_keys: list[K] = [] 

36 self._pending_futures: dict[K, asyncio.Future] = {} 

37 self._lock = asyncio.Lock() 

38 

39 async def load(self, key: K) -> V: 

40 """Load a single value by key, batching with others if possible.""" 

41 if key in self._cache: 

42 return self._cache[key] 

43 

44 async with self._lock: 

45 if key not in self._pending_futures: 

46 self._pending_keys.append(key) 

47 self._pending_futures[key] = asyncio.get_event_loop().create_future() 

48 

49 # If this is the first item in the batch, schedule execution 

50 if len(self._pending_keys) == 1: 

51 self.task_manager.create_background_task(self._execute_batch()) 

52 

53 return await self._pending_futures[key] 

54 

55 async def load_many(self, keys: list[K]) -> list[V]: 

56 """Load multiple values by keys in a single batch.""" 

57 return await Parallel.gather(*map(self.load, keys)) 

58 

59 async def _execute_batch(self) -> Any: 

60 """Internal method to execute the pending batch of keys.""" 

61 # Wait a tiny bit to collect more keys 

62 await asyncio.sleep(0) 

63 

64 async with self._lock: 

65 if not self._pending_keys: 

66 return 

67 

68 keys_to_fetch = list(self._pending_keys) 

69 futures_to_resolve = dict(self._pending_futures) 

70 

71 self._pending_keys.clear() 

72 self._pending_futures.clear() 

73 

74 try: 

75 results = await self._batch_fn(keys_to_fetch) 

76 

77 for key in keys_to_fetch: 

78 value = results.get(key) 

79 self._cache[key] = value # type: ignore[assignment] 

80 if not futures_to_resolve[key].done(): 

81 futures_to_resolve[key].set_result(value) 

82 except (RuntimeError, ValueError, TypeError, OSError) as e: 

83 for future in futures_to_resolve.values(): 

84 if not future.done(): 

85 future.set_exception(e) 

86 

87 def clear(self, key: K | None = None) -> Any: 

88 """Clear the cache.""" 

89 if key: 

90 self._cache.pop(key, None) 

91 else: 

92 self._cache.clear()