Coverage for core / src / sensorkit / backend / base.py: 98%

1# SPDX-License-Identifier: Apache-2.0 

2from __future__ import annotations 

3 

4import asyncio 

5import traceback 

6from abc import ABC, abstractmethod 

7from collections.abc import AsyncIterator, Callable, Coroutine 

8from dataclasses import dataclass 

9from datetime import datetime 

10from enum import StrEnum 

11from typing import Any, ClassVar, Self 

12 

13from loguru import logger 

14from pydantic import BaseModel 

15 

16 

17class Backend: 

18 """Exposes a system backend implementation.""" 

19 

20 def __init__(self, impl: BackendImpl): 

21 self.impl = impl 

22 

23 async def register_service(self, info: ServiceInfo): 

24 """Announce a service session to the backend.""" 

25 return await self.impl.register_service(info) 

26 

27 async def shutdown_request_listeners(self): 

28 """Remove all registered request listeners from the backend.""" 

29 return await self.impl.request_purge_listeners() 

30 

31 def request(self, entity: Entity | None = None) -> RequestResponseContext: 

32 """Return a request-response context scoped to the given entity.""" 

33 return RequestResponseContext(self.impl, entity or Entity()) 

34 

35 def stream(self, entity: Entity | None = None) -> StreamContext: 

36 """Return a stream context scoped to the given entity.""" 

37 return StreamContext(self.impl, entity or Entity()) 

38 

39 def key_value(self, entity: Entity | None = None) -> KeyValueContext: 

40 """Return a key-value context scoped to the given entity.""" 

41 return KeyValueContext(self.impl, entity or Entity()) 

42 

43 

44class SpecialProperty(StrEnum): 

45 """Well-known subject property tokens used to address groups of keys or event streams.""" 

46 

47 ALL_PROPERTIES = "$ALL$" 

48 ALL_DESCENDANTS = "$WILDCARD$" 

49 EVENTS = "$EVENT$" 

50 NONE = "$NONE$" 

51 

52 

53@dataclass(slots=True, frozen=True, eq=True) 

54class Entity: 

55 """An addressable entity identified by a tuple of path segments.""" 

56 

57 path: tuple[str, ...] = () 

58 

59 @classmethod 

60 def at(cls, *path: str) -> Self: 

61 """Construct an Entity from individual path segments.""" 

62 return cls(path=path) 

63 

64 def subject(self, name: str) -> Subject: 

65 """Return a Subject for the named property on this entity.""" 

66 return Subject(self.path, name) 

67 

68 def __str__(self): 

69 return ".".join(self.path) 

70 

71 def __repr__(self): 

72 return f"Entity({repr(self.path)})" 

73 

74 

75@dataclass(slots=True, frozen=True, eq=True) 

76class Subject(Entity): 

77 """An entity path combined with a named property, forming a fully-qualified address.""" 

78 

79 prop: str | SpecialProperty = SpecialProperty.NONE 

80 

81 def full_path(self): 

82 """Return the complete path tuple including the property segment.""" 

83 return self.path + (self.prop,) 

84 

85 def entity(self): 

86 """Return the parent Entity, stripping the property segment.""" 

87 return Entity(self.path) 

88 

89 def __str__(self): 

90 return f"{".".join(self.path)}.{self.prop}" 

91 

92 def __repr__(self): 

93 return f"Subject({repr(self.path)}, {repr(self.prop)})" 

94 

95 

96# FIXME: This *Context class hierarchy should be collapsed so that only one need exist per entity. 

97# Differentiation between stream, kv, etc., operations can be achieved via Protocols. 

98@dataclass(slots=True, frozen=True) 

99class BaseContext(ABC): 

100 impl: BackendImpl 

101 entity: Entity 

102 

103 

104type RequestCallback = Callable[[bytes], Coroutine[Any, Any, bytes]] 

105 

106 

107class BackendError(Exception): 

108 """Indicates a backend communication error.""" 

109 

110 

111class RemoteRequestError(BackendError): 

112 """Indicates an error raised in the remote request handler.""" 

113 MAGIC: ClassVar[bytes] = b"ERR\0" 

114 

115 def __init__(self, *args, details: str | None = None): 

116 super().__init__(*args) 

117 self.details = details 

118 

119 @classmethod 

120 def parse(cls, name: str, data: bytes) -> Self | None: 

121 """Parse a response payload and return a RemoteRequestError if it contains the error magic prefix.""" 

122 if data.startswith(cls.MAGIC): 

123 _, message, *details = data.split(b"\0", maxsplit=2) 

124 decoded = message.decode("utf-8", errors="replace") 

125 return cls( 

126 f"Remote '{name}' request error: {decoded}", 

127 details=details[0].decode("utf-8", errors="replace") if details else None, 

128 ) 

129 

130 return None 

131 

132 @classmethod 

133 def from_exception(cls, err: BaseException) -> Self: 

134 """Construct a RemoteRequestError from an exception, capturing its traceback as details.""" 

135 return cls( 

136 f"{type(err).__name__}: {err}", 

137 details="".join(traceback.format_exception(type(err), err, err.__traceback__)), 

138 ) 

139 

140 def encode(self) -> bytes: 

141 """Encode this error into a wire-format bytes payload prefixed with the error magic.""" 

142 return ( 

143 self.MAGIC 

144 + str(self).replace("\0", "\ufffd").encode("utf-8", errors="replace") 

145 + (b"\0" + self.details.encode("utf-8", errors="replace") if self.details else b"") 

146 ) 

147 

148 

149class UnregisteredResponder(BackendError): 

150 """Raised when a request is made to a subject with no registered handler.""" 

151 

152 def __init__(self, subject: Subject): 

153 super().__init__(f"Remote `{subject.prop}` request error: responder for {subject} not registered") 

154 

155 

156@dataclass(slots=True, frozen=True) 

157class RequestResponseContext(BaseContext): 

158 """Exposes request-response methods of a backend implementation.""" 

159 

160 async def invoke(self, name: str, payload: bytes): 

161 """Send a request to the named subject and return the response bytes. 

162 

163 Raises RemoteRequestError if the remote handler returned an encoded error. 

164 """ 

165 res = await self.impl.request_invoke(self.entity.subject(name), payload) 

166 

167 if err := RemoteRequestError.parse(name, res): 

168 logger.opt(lazy=True).debug( 

169 f"{err}{{details}}", 

170 details=lambda: "\n\t".join(["", *err.details.split("\n")] if err.details else []), 

171 ) 

172 raise err 

173 

174 return res 

175 

176 async def handle_request(self, name: str, callback: RequestCallback): 

177 """Register a callback to handle incoming requests for the named subject. 

178 

179 Errors raised by the callback are caught and returned to the caller as encoded 

180 RemoteRequestError payloads. 

181 """ 

182 async def _request_wrapper(payload: bytes): 

183 try: 

184 res = await callback(payload) 

185 except asyncio.CancelledError: 

186 logger.debug(f"dropping request {self.entity}.{name} due to cancellation") 

187 raise 

188 except Exception as e: 

189 logger.opt(exception=e).debug(f"error in {self.entity}.{name} request callback") 

190 res = RemoteRequestError.from_exception(e).encode() 

191 

192 return res 

193 

194 return await self.impl.request_listen(self.entity.subject(name), _request_wrapper) 

195 

196 

197@dataclass(slots=True, frozen=True) 

198class StreamContext(BaseContext): 

199 """Exposes stream methods of a backend implementation.""" 

200 

201 async def list_keys(self): 

202 """Return a coroutine that resolves to the list of stream subject keys for this entity.""" 

203 return await self.impl.stream_list(self.entity) 

204 

205 async def consume( 

206 self, 

207 key: str | None = None, 

208 *, 

209 durable_name: str | None = None, 

210 from_sequence: int | None = None, 

211 from_time: datetime | None = None, 

212 include_latest: bool = False, 

213 ): 

214 """Return a coroutine that resolves to an async iterator of StreamMessages. 

215 

216 If key is None, all properties for the entity are consumed. 

217 """ 

218 return await self.impl.stream_consume( 

219 self.entity.subject(key if key else SpecialProperty.ALL_PROPERTIES), 

220 durable_name=durable_name, 

221 start_at=from_time or from_sequence, 

222 include_latest=include_latest, 

223 ) 

224 

225 async def publish(self, key: str, payload: bytes): 

226 """Return a coroutine that publishes payload to the named stream subject.""" 

227 return await self.impl.stream_publish(self.entity.subject(key), payload) 

228 

229 async def publish_event(self, payload: bytes): 

230 """Return a coroutine that publishes payload to the entity's event stream subject.""" 

231 return await self.impl.stream_publish(self.entity.subject(SpecialProperty.EVENTS), payload) 

232 

233 

234@dataclass(slots=True, frozen=True) 

235class KeyValueContext(BaseContext): 

236 """Exposes key-value store methods of a backend implementation.""" 

237 

238 async def monitor_all(self, *, deep=False): 

239 """Return an async generator that yields KVEntry objects as keys are created or updated. 

240 

241 With deep=True, monitors all descendants rather than direct properties only. 

242 """ 

243 if deep: 

244 # FIXME: This option is expeditious at the moment, but can lead to a lot of duplicative 

245 # messages. In general, this backend code needs a major iteration to use multi- 

246 # plexing. Use of ALL_DESCENDANTS is generally problematic. Perhaps a client- 

247 # scoped option: set up a single "firehose" subscription and use a global multi- 

248 # plexer, or, disallow use of ALL_DESCENDANTS and use shared subscriptions per 

249 # entity, possibly via AsyncObserver. In the latter case, calling this method 

250 # with deep=True would raise. 

251 logger.debug("warning: monitor_all(deep=True) may cause excess traffic") 

252 

253 prop = SpecialProperty.ALL_DESCENDANTS if deep else SpecialProperty.ALL_PROPERTIES 

254 monitor = await self.impl.kv_monitor(self.entity.subject(prop)) 

255 

256 async def key_value_monitor(): 

257 async for batch in monitor: 

258 for entry in batch: 

259 yield entry 

260 

261 return key_value_monitor() 

262 

263 async def monitor(self, key: str): 

264 """Return an async generator that yields KVEntry objects for changes to the named key.""" 

265 monitor = await self.impl.kv_monitor(self.entity.subject(key)) 

266 

267 async def key_value_monitor(): 

268 async for batch in monitor: 

269 for entry in batch: 

270 yield entry 

271 

272 return key_value_monitor() 

273 

274 async def get_all(self, *, deep=False) -> list[KVEntry]: 

275 """Return all current (non-deleted) KV entries for this entity. 

276 

277 With deep=True, includes entries from descendant entities as well. 

278 """ 

279 prop = SpecialProperty.ALL_DESCENDANTS if deep else SpecialProperty.ALL_PROPERTIES 

280 subject = self.entity.subject(prop) 

281 out = [] 

282 

283 # Monitor for a single batch, which are the current values. 

284 async for batch in await self.impl.kv_monitor(subject): 

285 for entry in batch: 

286 if not entry.deleted(): 

287 out.append(entry) 

288 

289 break 

290 

291 return out 

292 

293 async def get(self, key: str): 

294 """Fetch the current KVEntry for the named key.""" 

295 return await self.impl.kv_get(self.entity.subject(key)) 

296 

297 async def create( 

298 self, 

299 key: str, 

300 value: bytes, 

301 *, 

302 ttl: float | None = None, 

303 ): 

304 """Create a new KV entry. Raises if the key already exists.""" 

305 return await self.impl.kv_create(self.entity.subject(key), value, ttl) 

306 

307 async def update( 

308 self, 

309 key: str, 

310 value: bytes, 

311 *, 

312 revision: int | None = None, 

313 ttl: float | None = None, 

314 ): 

315 """Update an existing KV entry, optionally requiring a specific revision.""" 

316 return await self.impl.kv_update(self.entity.subject(key), value, revision, ttl) 

317 

318 async def delete( 

319 self, 

320 key: str, 

321 *, 

322 revision: int | None = None, 

323 ): 

324 """Delete the KV entry for the named key, optionally requiring a specific revision.""" 

325 await self.impl.kv_delete(self.entity.subject(key), revision) 

326 

327 

328@dataclass(slots=True, frozen=True, eq=True) 

329class KVEntry: 

330 """A single key-value store entry.""" 

331 key: Subject 

332 value: bytes 

333 revision: int 

334 

335 DELETE_MARKER: ClassVar[bytes] = b"DELETED" 

336 """Sentinel value denoting a deleted KV entry.""" 

337 

338 def deleted(self): 

339 """Return True if this entry has been marked as deleted.""" 

340 return self.value is KVEntry.DELETE_MARKER 

341 

342 

343class KVError(BackendError): 

344 """Base class for key-value store errors.""" 

345 

346 

347class RevisionError(KVError): 

348 """Raised when a KV update or delete is rejected due to a revision mismatch.""" 

349 

350 

351class KeyNotFound(KVError): 

352 """Raised when a requested KV key does not exist.""" 

353 

354 def __init__(self, subject: Subject, *, deleted: bool): 

355 super().__init__(f"key for `{subject}` not found") 

356 self.deleted = deleted 

357 

358 

359class ServiceInfo(BaseModel): 

360 """Information about a service.""" 

361 name: str 

362 version: str 

363 

364 

365@dataclass(slots=True, frozen=True, eq=True) 

366class StreamMessage: 

367 """A message read from a stream, including its subject, sequence number, timestamp, and payload.""" 

368 

369 subject: Subject 

370 sequence: int 

371 timestamp: datetime 

372 data: bytes 

373 

374 

375class BackendImpl(ABC): 

376 """Interface providing access to a backend.""" 

377 

378 @classmethod 

379 @abstractmethod 

380 async def create(cls, *args, **kwargs) -> BackendImpl: 

381 """Factory to create a BackendImpl instance.""" 

382 

383 async def register_service(self, info: ServiceInfo): 

384 """Notify the backend of a new service session. Optional method.""" 

385 

386 @abstractmethod 

387 async def request_invoke(self, target: Subject, payload: bytes) -> bytes: 

388 """Invoke a request.""" 

389 

390 @abstractmethod 

391 async def request_listen(self, target: Subject, coro: RequestCallback): 

392 """Listen for requests.""" 

393 

394 @abstractmethod 

395 async def request_purge_listeners(self): 

396 """Remove all request listeners.""" 

397 

398 @abstractmethod 

399 async def stream_list(self, entity: Entity) -> list[Subject]: 

400 """List defined stream subjects for an entity.""" 

401 

402 @abstractmethod 

403 async def stream_publish(self, target: Subject, payload: bytes): 

404 """Publish to a stream subject.""" 

405 

406 @abstractmethod 

407 async def stream_consume( 

408 self, 

409 target: Subject, 

410 *, 

411 durable_name: str | None = None, 

412 start_at: int | datetime | None = None, 

413 include_latest: bool = False, 

414 ) -> AsyncIterator[StreamMessage]: 

415 """Consume a stream subject.""" 

416 

417 @abstractmethod 

418 async def kv_monitor(self, target: Subject) -> AsyncIterator[list[KVEntry]]: 

419 """Subscribe to KV modifications for the given entity.""" 

420 

421 @abstractmethod 

422 async def kv_get(self, target: Subject) -> KVEntry: 

423 """Return the value associated with a key.""" 

424 

425 @abstractmethod 

426 async def kv_create( 

427 self, 

428 target: Subject, 

429 payload: bytes, 

430 ttl: float | None = None, 

431 ) -> KVEntry: 

432 """Create a KV record.""" 

433 

434 @abstractmethod 

435 async def kv_update( 

436 self, 

437 target: Subject, 

438 payload: bytes, 

439 revision: int | None = None, 

440 ttl: float | None = None, 

441 ) -> KVEntry: 

442 """Update a KV record.""" 

443 

444 @abstractmethod 

445 async def kv_delete(self, target: Subject, revision: int | None = None): 

446 """Delete a KV record."""