Coverage for core / src / sensorkit / webapi / forwarder.py: 95%

1# SPDX-License-Identifier: Apache-2.0 

2import asyncio 

3import collections 

4import contextlib 

5from abc import ABC, abstractmethod 

6from datetime import UTC, datetime 

7from typing import Any, AsyncIterator, Literal, NamedTuple, override 

8 

9from loguru import logger 

10from pydantic import TypeAdapter 

11from pydantic_core import from_json 

12 

13from sensorkit import api as sk 

14from sensorkit.webapi.serve import ServeHandler 

15 

16 

17class SKRecord(NamedTuple): 

18 kind: Literal["stream", "event", "state", "product"] 

19 time: datetime 

20 subject: sk.Subject 

21 payload: dict[str, Any] | None 

22 

23 def serialize(self) -> str: 

24 """Serialize to a JSON string.""" 

25 return _record_adapter.dump_json(self).decode() 

26 

27 

28_record_adapter = TypeAdapter(SKRecord) 

29type RecordQueueSet = set[asyncio.Queue[SKRecord | None]] 

30 

31 

32class Forwarder(ABC): 

33 """Base class that caches backend updates and broadcasts them to target queues. 

34 

35 Subclasses implement `_monitor` to yield `SKRecord` updates from a particular 

36 backend source; the base loop keeps a per-entity cache (treating a `None` payload 

37 as a deletion) and fans each update out to every registered target queue. 

38 """ 

39 

40 def __init__(self, *, targets: RecordQueueSet): 

41 self.targets = targets 

42 self.cache: dict[str, dict[str, SKRecord]] = collections.defaultdict(dict) 

43 self.task: asyncio.Task | None = None 

44 

45 async def start(self, *, task_group: asyncio.TaskGroup): 

46 """Start background monitoring task for this forwarder.""" 

47 self.task = task_group.create_task(self._run()) 

48 return self.task 

49 

50 async def stop(self): 

51 """Stop the background monitoring task, if one is running.""" 

52 if self.task is None: 

53 return 

54 

55 self.task.cancel() 

56 

57 with contextlib.suppress(asyncio.CancelledError): 

58 await self.task 

59 

60 async def _run(self): 

61 async for update in self._monitor(): 

62 entity = str(update.subject.entity()) 

63 attr = update.subject.prop 

64 

65 # Update the cache. Note the update payload being None is used as a deletion marker. 

66 # This assumes that all incoming data are JSON objects (non-primitives) and therefore 

67 # can never legitimately be `null`. 

68 attr_cache = self.cache[entity] 

69 

70 if update.payload is None: 

71 attr_cache.pop(attr, None) 

72 else: 

73 attr_cache[attr] = update 

74 

75 # Broadcast to all targets. 

76 for queue in tuple(self.targets): 

77 try: 

78 queue.put_nowait(update) 

79 except asyncio.QueueFull: 

80 # A subscriber that has stopped draining is cut loose. 

81 logger.warning("dropping a subscriber that is not keeping up") 

82 self.targets.discard(queue) 

83 

84 while not queue.empty(): 

85 queue.get_nowait() 

86 

87 queue.put_nowait(None) 

88 

89 @abstractmethod 

90 def _monitor(self) -> AsyncIterator[SKRecord]: 

91 """Yield updates as they are received from the backend.""" 

92 

93 def snapshot(self, entity_id: str | None = None) -> list[SKRecord]: 

94 """Return the cached records, optionally limited to a single entity.""" 

95 caches = self.cache.values() if entity_id is None else (self.cache.get(entity_id, {}),) 

96 return [record for attrs in caches for record in attrs.values()] 

97 

98 

99class KeyValueForwarder(Forwarder): 

100 """Forwarder specialization that monitors Key-Value changes.""" 

101 

102 def __init__(self, kit: sk.SensorKit, *, targets: RecordQueueSet): 

103 super().__init__(targets=targets) 

104 self.kit = kit 

105 

106 @override 

107 async def _monitor(self) -> AsyncIterator[SKRecord]: 

108 monitor = await self.kit.entity()._kv.monitor_all(deep=True) 

109 

110 async for entry in monitor: 

111 ts = datetime.now(UTC) # FIXME: Pending backend exposure of KV update timestamp. 

112 payload = from_json(entry.value) if not entry.deleted() else None 

113 

114 if not isinstance(payload, dict) and payload is not None: 

115 logger.warning(f"Invalid KV payload for entity {entry.key.entity()}: {payload}") 

116 continue 

117 

118 yield SKRecord( 

119 kind="state", 

120 subject=entry.key, 

121 time=ts, 

122 payload=payload, 

123 ) 

124 

125 

126class StreamForwarder(Forwarder): 

127 """Forwarder specialization that monitors stream updates (state/events).""" 

128 

129 def __init__(self, kit: sk.SensorKit, *, targets: RecordQueueSet): 

130 super().__init__(targets=targets) 

131 self.kit = kit 

132 

133 @override 

134 async def _monitor(self) -> AsyncIterator[SKRecord]: 

135 consumer = await self.kit.entity()._stream.consume( 

136 key=sk.SpecialProperty.ALL_DESCENDANTS, 

137 include_latest=True, 

138 ) 

139 

140 async for msg in consumer: 

141 payload = from_json(msg.data) 

142 

143 if not isinstance(payload, dict): 

144 logger.warning(f"Invalid stream payload for entity {msg.subject}: {payload}") 

145 continue 

146 

147 yield SKRecord( 

148 kind="event" if msg.subject.prop == sk.SpecialProperty.EVENTS else "stream", 

149 subject=msg.subject, 

150 time=msg.timestamp, 

151 payload=payload, 

152 ) 

153 

154 

155class ProductForwarder(Forwarder): 

156 """Forwarder specialization that monitors data products from a ServeHandler.""" 

157 

158 def __init__(self, serve_handler: ServeHandler, *, targets: RecordQueueSet): 

159 super().__init__(targets=targets) 

160 self.serve_handler = serve_handler 

161 

162 @override 

163 async def _monitor(self) -> AsyncIterator[SKRecord]: 

164 async for info in self.serve_handler.watch_listing(): 

165 yield SKRecord( 

166 kind="product", 

167 subject=sk.Subject(path=(info.controller_id,), prop=info.product_id), 

168 time=datetime.now(UTC), 

169 payload=info.model_dump(), 

170 )