Coverage for object_streams/sessions.py: 93%
372 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-02 17:07 +0000
1"""Connection-local subscription coordination."""
3from __future__ import annotations
5from collections.abc import Callable
6from collections.abc import Mapping
7from dataclasses import dataclass
8from dataclasses import replace
9from itertools import count
10from typing import Any
12from asgiref.sync import async_to_sync
13from channels.db import database_sync_to_async
14from django.db import models
16from object_streams.events import EventOperation
17from object_streams.events import ListAction
18from object_streams.events import ObjectRef
19from object_streams.events import StreamEvent
20from object_streams.exceptions import FilterValidationError
21from object_streams.exceptions import NotRegistered
22from object_streams.models import ObjectStreamEvent
23from object_streams.outbox import latest_outbox_cursor
24from object_streams.outbox import outbox_events_after
25from object_streams.outbox import replay_is_complete
26from object_streams.registry import ObjectStreamRegistration
27from object_streams.registry import ObjectStreamRegistry
28from object_streams.registry import registry as default_registry
29from object_streams.subscriptions import ResyncRequired
30from object_streams.subscriptions import SubscriptionKind
31from object_streams.subscriptions import SubscriptionRequest
32from object_streams.transports.base import Transport
35__all__ = ("ActiveSubscription", "AsyncSubscriptionSession", "SubscriptionSession")
38@dataclass(slots=True)
39class ActiveSubscription:
40 """Connection-local state for one active subscription."""
42 request: SubscriptionRequest
43 registration: ObjectStreamRegistration
44 member_pks: set[str]
45 through_cursor: int
47 @property
48 def subscription_id(self) -> str:
49 if self.request.subscription_id is None: 49 ↛ 50line 49 didn't jump to line 50 because the condition on line 49 was never true
50 msg = "Active subscriptions require a subscription id."
51 raise ValueError(msg)
52 return self.request.subscription_id
55class _BaseSubscriptionSession:
56 def __init__(
57 self,
58 *,
59 user: Any,
60 transport: Transport,
61 registry: ObjectStreamRegistry = default_registry,
62 request: Any = None,
63 subscription_id_factory: Callable[[], str] | None = None,
64 replay_limit: int = 1000,
65 max_subscriptions: int | None = 100,
66 max_member_pks: int | None = 10000,
67 ):
68 self.user = user
69 self.transport = transport
70 self.registry = registry
71 self.request = request
72 self.replay_limit = replay_limit
73 self.max_subscriptions = max_subscriptions
74 self.max_member_pks = max_member_pks
75 self._subscriptions: dict[str, ActiveSubscription] = {}
76 self._counter = count(1)
77 self._subscription_id_factory = subscription_id_factory
79 @property
80 def subscriptions(self) -> tuple[ActiveSubscription, ...]:
81 return tuple(self._subscriptions.values())
83 def _coerce_subscription_request(self, message: Mapping[str, Any] | SubscriptionRequest) -> SubscriptionRequest:
84 if isinstance(message, SubscriptionRequest): 84 ↛ 85line 84 didn't jump to line 85 because the condition on line 84 was never true
85 return message
86 return SubscriptionRequest.from_message(message)
88 def _next_subscription_id(self) -> str:
89 if self._subscription_id_factory is not None: 89 ↛ 90line 89 didn't jump to line 90 because the condition on line 89 was never true
90 return self._subscription_id_factory()
91 return f"sub_{next(self._counter)}"
93 def _subscription_queryset(
94 self,
95 registration: ObjectStreamRegistration,
96 request: SubscriptionRequest,
97 ) -> models.QuerySet:
98 if request.kind == SubscriptionKind.FILTER:
99 queryset = registration.get_queryset(
100 self.user,
101 request.filters,
102 request=self.request,
103 )
104 else:
105 queryset = registration.visibility.get_queryset(self.user, registration.model, action="read")
106 if request.kind == SubscriptionKind.OBJECT:
107 queryset = queryset.filter(pk=request.pk)
108 return queryset
110 def _current_member_pks(
111 self,
112 registration: ObjectStreamRegistration,
113 request: SubscriptionRequest,
114 ) -> set[str]:
115 return {str(pk) for pk in self._subscription_queryset(registration, request).values_list("pk", flat=True)}
117 def _member_pks_exceed_limit(
118 self,
119 registration: ObjectStreamRegistration,
120 request: SubscriptionRequest,
121 ) -> bool:
122 """Check membership size with a bounded query, before materializing every pk."""
124 if self.max_member_pks is None: 124 ↛ 125line 124 didn't jump to line 125 because the condition on line 124 was never true
125 return False
126 pks = self._subscription_queryset(registration, request).values_list("pk", flat=True)
127 return len(pks[: self.max_member_pks + 1]) > self.max_member_pks
129 def _connection_rejection(self, requested: SubscriptionRequest) -> tuple[str, str] | None:
130 """Return an error code and message when connection state forbids the subscription."""
132 if requested.subscription_id is not None and requested.subscription_id in self._subscriptions:
133 return ("duplicate_subscription", "Subscription id is already active on this connection.")
134 if self.max_subscriptions is not None and len(self._subscriptions) >= self.max_subscriptions:
135 return ("subscription_limit_exceeded", "Connection has too many active subscriptions.")
136 return None
138 def _is_current_member(self, active: ActiveSubscription, event: StreamEvent) -> bool:
139 if event.op == str(EventOperation.DELETED):
140 return False
141 return self._subscription_queryset(active.registration, active.request).filter(pk=event.subject.pk).exists()
143 def _list_action(
144 self,
145 event: StreamEvent,
146 *,
147 before_member: bool,
148 after_member: bool,
149 ) -> ListAction | None:
150 if before_member and after_member:
151 return ListAction.CHANGED
152 if not before_member and after_member:
153 return ListAction.ADDED
154 if before_member and not after_member:
155 if event.op == str(EventOperation.DELETED):
156 return ListAction.DELETED
157 return ListAction.REMOVED
158 return None
160 def _replay_is_complete(self, requested_cursor: int) -> bool:
161 return replay_is_complete(requested_cursor)
163 def _has_collection_replay_events(
164 self,
165 active: ActiveSubscription,
166 requested_cursor: int,
167 through_cursor: int,
168 ) -> bool:
169 return outbox_events_after(
170 requested_cursor,
171 model=active.registration.model,
172 through_cursor=through_cursor,
173 ).exists()
175 def _object_replay_events(
176 self,
177 active: ActiveSubscription,
178 requested_cursor: int,
179 through_cursor: int,
180 ) -> tuple[StreamEvent, ...] | None:
181 subject = ObjectRef(model=active.request.model, pk=active.request.pk)
182 rows = list(
183 outbox_events_after(
184 requested_cursor,
185 subject=subject,
186 through_cursor=through_cursor,
187 limit=self.replay_limit + 1,
188 )
189 )
190 if len(rows) > self.replay_limit:
191 return None
193 events = []
194 for row in rows:
195 event = row.to_stream_event()
196 list_action = ListAction.DELETED if event.op == str(EventOperation.DELETED) else ListAction.CHANGED
197 events.append(
198 replace(
199 event,
200 subscription_id=active.subscription_id,
201 list_action=list_action,
202 )
203 )
204 return tuple(events)
206 def _subscription_catch_up(
207 self,
208 active: ActiveSubscription,
209 snapshot_cursor: int,
210 ) -> tuple[StreamEvent | ResyncRequired, ...]:
211 through_cursor = latest_outbox_cursor()
212 if through_cursor == snapshot_cursor:
213 return ()
215 messages: tuple[StreamEvent | ResyncRequired, ...] = ()
216 if active.request.kind == SubscriptionKind.OBJECT:
217 replay_events = self._object_replay_events(active, snapshot_cursor, through_cursor)
218 if replay_events is None: 218 ↛ 219line 218 didn't jump to line 219 because the condition on line 218 was never true
219 messages = (
220 ResyncRequired(
221 subscription_id=active.subscription_id,
222 cursor=through_cursor,
223 reason="object_replay_limit_exceeded",
224 ),
225 )
226 else:
227 messages = replay_events
228 elif self._has_collection_replay_events(active, snapshot_cursor, through_cursor): 228 ↛ 236line 228 didn't jump to line 236 because the condition on line 228 was always true
229 messages = (
230 ResyncRequired(
231 subscription_id=active.subscription_id,
232 cursor=through_cursor,
233 ),
234 )
236 active.member_pks = self._current_member_pks(active.registration, active.request)
237 active.through_cursor = through_cursor
238 return messages
241class SubscriptionSession(_BaseSubscriptionSession):
242 """Coordinate subscriptions and event delivery for one connection."""
244 def handle_message(self, message: Mapping[str, Any]) -> SubscriptionRequest | bool | None:
245 """Handle a subscribe or unsubscribe protocol message."""
247 op = message.get("op")
248 if op == "subscribe":
249 return self.subscribe(message)
250 if op == "unsubscribe":
251 subscription_id = message.get("subscription_id")
252 if subscription_id is None:
253 self._send_error("invalid_request", "Unsubscribe messages require a subscription_id.")
254 return None
255 return self.unsubscribe(str(subscription_id))
257 self._send_error("invalid_request", "Messages require a supported op.")
258 return None
260 def subscribe(self, message: Mapping[str, Any] | SubscriptionRequest) -> SubscriptionRequest | None:
261 """Register a subscription and send its acknowledgement."""
263 try:
264 requested = self._coerce_subscription_request(message)
265 registration = self.registry.get(requested.model)
266 except (LookupError, NotRegistered, ValueError) as exc:
267 self._send_error("invalid_request", str(exc))
268 return None
270 rejection = self._connection_rejection(requested)
271 if rejection is not None:
272 self._send_error(*rejection)
273 return None
275 if requested.search is not None:
276 self._send_error("unsupported_search", "Search subscriptions are not supported yet.")
277 return None
279 snapshot_cursor = latest_outbox_cursor()
280 if requested.cursor is not None and requested.cursor > snapshot_cursor:
281 self._send_error("invalid_cursor", "Subscription cursor is newer than the outbox.")
282 return None
284 subscription_id = requested.subscription_id or self._next_subscription_id()
285 acknowledged = replace(requested, subscription_id=subscription_id, cursor=snapshot_cursor)
287 try:
288 if self._member_pks_exceed_limit(registration, acknowledged):
289 self._send_error("subscription_too_large", "Subscription matches too many objects.")
290 return None
291 member_pks = self._current_member_pks(registration, acknowledged)
292 except FilterValidationError as exc:
293 self._send_error("invalid_filter", "Subscription filters are invalid.", details=exc.errors)
294 return None
296 if acknowledged.kind == SubscriptionKind.OBJECT and not member_pks:
297 self._send_error("not_found", "Object does not exist or is not visible.")
298 return None
300 active = ActiveSubscription(
301 request=acknowledged,
302 registration=registration,
303 member_pks=member_pks,
304 through_cursor=snapshot_cursor,
305 )
306 self._subscriptions[subscription_id] = active
307 self._prepare_subscription(acknowledged)
308 catch_up_messages = self._subscription_catch_up(active, snapshot_cursor)
309 self._send_subscribed(acknowledged)
310 self._replay_requested_cursor(active, requested.cursor, through_cursor=snapshot_cursor)
311 self._send_catch_up_messages(catch_up_messages)
312 return acknowledged
314 def unsubscribe(self, subscription_id: str) -> bool:
315 """Remove a subscription and send its acknowledgement."""
317 if subscription_id not in self._subscriptions:
318 self._send_error("not_subscribed", "Subscription is not active.")
319 return False
321 del self._subscriptions[subscription_id]
322 self._send_unsubscribed(subscription_id)
323 return True
325 def publish(self, event_or_row: StreamEvent | ObjectStreamEvent) -> list[StreamEvent]:
326 """Evaluate and deliver one outbox event to active subscriptions."""
328 event = event_or_row.to_stream_event() if isinstance(event_or_row, ObjectStreamEvent) else event_or_row
329 delivered = []
330 for active in tuple(self._subscriptions.values()):
331 if event.cursor is not None and event.cursor <= active.through_cursor:
332 continue
333 subscription_event = self.evaluate(active, event)
334 if event.cursor is not None: 334 ↛ 336line 334 didn't jump to line 336 because the condition on line 334 was always true
335 active.through_cursor = event.cursor
336 if subscription_event is None:
337 continue
338 self._send_event(subscription_event)
339 delivered.append(subscription_event)
340 return delivered
342 def evaluate(self, active: ActiveSubscription, event: StreamEvent) -> StreamEvent | None:
343 """Return the subscription-relative event, or None when it has no effect."""
345 if event.subject.model != active.request.model: 345 ↛ 346line 345 didn't jump to line 346 because the condition on line 345 was never true
346 return None
347 if event.facet not in active.registration.facets: 347 ↛ 348line 347 didn't jump to line 348 because the condition on line 347 was never true
348 return None
349 if active.request.kind == SubscriptionKind.OBJECT and event.subject.pk != active.request.pk: 349 ↛ 350line 349 didn't jump to line 350 because the condition on line 349 was never true
350 return None
352 before_member = event.subject.pk in active.member_pks
353 after_member = self._is_current_member(active, event)
354 list_action = self._list_action(event, before_member=before_member, after_member=after_member)
355 if list_action is None:
356 return None
358 if after_member:
359 active.member_pks.add(event.subject.pk)
360 else:
361 active.member_pks.discard(event.subject.pk)
363 return replace(
364 event,
365 subscription_id=active.subscription_id,
366 list_action=list_action,
367 )
369 def _replay_requested_cursor(
370 self,
371 active: ActiveSubscription,
372 requested_cursor: int | None,
373 *,
374 through_cursor: int,
375 ) -> None:
376 if requested_cursor is None or requested_cursor == through_cursor:
377 return
378 if not self._replay_is_complete(requested_cursor):
379 self._send_resync(
380 ResyncRequired(
381 subscription_id=active.subscription_id,
382 cursor=through_cursor,
383 reason="cursor_pruned",
384 )
385 )
386 return
387 if active.request.kind == SubscriptionKind.OBJECT:
388 self._replay_object_subscription(active, requested_cursor, through_cursor=through_cursor)
389 return
391 if self._has_collection_replay_events(active, requested_cursor, through_cursor): 391 ↛ exitline 391 didn't return from function '_replay_requested_cursor' because the condition on line 391 was always true
392 self._send_resync(
393 ResyncRequired(
394 subscription_id=active.subscription_id,
395 cursor=through_cursor,
396 )
397 )
399 def _send_catch_up_messages(self, messages: tuple[StreamEvent | ResyncRequired, ...]) -> None:
400 for message in messages:
401 if isinstance(message, ResyncRequired): 401 ↛ 404line 401 didn't jump to line 404 because the condition on line 401 was always true
402 self._send_resync(message)
403 else:
404 self._send_event(message)
406 def _replay_object_subscription(
407 self,
408 active: ActiveSubscription,
409 requested_cursor: int,
410 *,
411 through_cursor: int,
412 ) -> None:
413 replay_events = self._object_replay_events(active, requested_cursor, through_cursor)
414 if replay_events is None:
415 self._send_resync(
416 ResyncRequired(
417 subscription_id=active.subscription_id,
418 cursor=through_cursor,
419 reason="object_replay_limit_exceeded",
420 )
421 )
422 return
424 for event in replay_events:
425 self._send_event(event)
427 def _send_subscribed(self, subscription: SubscriptionRequest) -> None:
428 async_to_sync(self.transport.send_subscribed)(subscription)
430 def _prepare_subscription(self, subscription: SubscriptionRequest) -> None:
431 prepare = getattr(self.transport, "prepare_subscription", None)
432 if prepare is not None:
433 async_to_sync(prepare)(subscription)
435 def _send_unsubscribed(self, subscription_id: str) -> None:
436 async_to_sync(self.transport.send_unsubscribed)(subscription_id)
438 def _send_event(self, event: StreamEvent) -> None:
439 async_to_sync(self.transport.send_event)(event)
441 def _send_resync(self, resync: ResyncRequired) -> None:
442 async_to_sync(self.transport.send_resync)(resync)
444 def _send_error(self, code: str, message: str, *, details: Any = None) -> None:
445 async_to_sync(self.transport.send_error)(code, message, details=details)
448class AsyncSubscriptionSession(_BaseSubscriptionSession):
449 """Async subscription coordinator for ASGI consumers."""
451 async def handle_message(self, message: Mapping[str, Any]) -> SubscriptionRequest | bool | None:
452 """Handle a subscribe or unsubscribe protocol message."""
454 op = message.get("op")
455 if op == "subscribe":
456 return await self.subscribe(message)
457 if op == "unsubscribe":
458 subscription_id = message.get("subscription_id")
459 if subscription_id is None:
460 await self.transport.send_error("invalid_request", "Unsubscribe messages require a subscription_id.")
461 return None
462 return await self.unsubscribe(str(subscription_id))
464 await self.transport.send_error("invalid_request", "Messages require a supported op.")
465 return None
467 async def subscribe(self, message: Mapping[str, Any] | SubscriptionRequest) -> SubscriptionRequest | None:
468 """Register a subscription and send its acknowledgement."""
470 try:
471 requested = self._coerce_subscription_request(message)
472 registration = self.registry.get(requested.model)
473 except (LookupError, NotRegistered, ValueError) as exc:
474 await self.transport.send_error("invalid_request", str(exc))
475 return None
477 rejection = self._connection_rejection(requested)
478 if rejection is not None:
479 await self.transport.send_error(*rejection)
480 return None
482 if requested.search is not None:
483 await self.transport.send_error("unsupported_search", "Search subscriptions are not supported yet.")
484 return None
486 snapshot_cursor = await database_sync_to_async(latest_outbox_cursor)()
487 if requested.cursor is not None and requested.cursor > snapshot_cursor:
488 await self.transport.send_error("invalid_cursor", "Subscription cursor is newer than the outbox.")
489 return None
491 subscription_id = requested.subscription_id or self._next_subscription_id()
492 acknowledged = replace(requested, subscription_id=subscription_id, cursor=snapshot_cursor)
494 try:
495 if await database_sync_to_async(self._member_pks_exceed_limit)(registration, acknowledged):
496 await self.transport.send_error("subscription_too_large", "Subscription matches too many objects.")
497 return None
498 member_pks = await database_sync_to_async(self._current_member_pks)(registration, acknowledged)
499 except FilterValidationError as exc:
500 await self.transport.send_error("invalid_filter", "Subscription filters are invalid.", details=exc.errors)
501 return None
503 if acknowledged.kind == SubscriptionKind.OBJECT and not member_pks:
504 await self.transport.send_error("not_found", "Object does not exist or is not visible.")
505 return None
507 active = ActiveSubscription(
508 request=acknowledged,
509 registration=registration,
510 member_pks=member_pks,
511 through_cursor=snapshot_cursor,
512 )
513 self._subscriptions[subscription_id] = active
514 await self._prepare_subscription(acknowledged)
515 catch_up_messages = await database_sync_to_async(self._subscription_catch_up)(active, snapshot_cursor)
516 await self.transport.send_subscribed(acknowledged)
517 await self._replay_requested_cursor(active, requested.cursor, through_cursor=snapshot_cursor)
518 await self._send_catch_up_messages(catch_up_messages)
519 return acknowledged
521 async def unsubscribe(self, subscription_id: str) -> bool:
522 """Remove a subscription and send its acknowledgement."""
524 if subscription_id not in self._subscriptions:
525 await self.transport.send_error("not_subscribed", "Subscription is not active.")
526 return False
528 del self._subscriptions[subscription_id]
529 await self.transport.send_unsubscribed(subscription_id)
530 return True
532 async def publish(self, event_or_row: StreamEvent | ObjectStreamEvent) -> list[StreamEvent]:
533 """Evaluate and deliver one outbox event to active subscriptions."""
535 if isinstance(event_or_row, ObjectStreamEvent): 535 ↛ 538line 535 didn't jump to line 538 because the condition on line 535 was always true
536 event = await database_sync_to_async(event_or_row.to_stream_event)()
537 else:
538 event = event_or_row
540 delivered = []
541 for active in tuple(self._subscriptions.values()):
542 if event.cursor is not None and event.cursor <= active.through_cursor:
543 continue
544 subscription_event = await self.evaluate(active, event)
545 if event.cursor is not None: 545 ↛ 547line 545 didn't jump to line 547 because the condition on line 545 was always true
546 active.through_cursor = event.cursor
547 if subscription_event is None:
548 continue
549 await self.transport.send_event(subscription_event)
550 delivered.append(subscription_event)
551 return delivered
553 async def evaluate(self, active: ActiveSubscription, event: StreamEvent) -> StreamEvent | None:
554 """Return the subscription-relative event, or None when it has no effect."""
556 if event.subject.model != active.request.model: 556 ↛ 557line 556 didn't jump to line 557 because the condition on line 556 was never true
557 return None
558 if event.facet not in active.registration.facets: 558 ↛ 559line 558 didn't jump to line 559 because the condition on line 558 was never true
559 return None
560 if active.request.kind == SubscriptionKind.OBJECT and event.subject.pk != active.request.pk: 560 ↛ 561line 560 didn't jump to line 561 because the condition on line 560 was never true
561 return None
563 before_member = event.subject.pk in active.member_pks
564 after_member = await database_sync_to_async(self._is_current_member)(active, event)
565 list_action = self._list_action(event, before_member=before_member, after_member=after_member)
566 if list_action is None:
567 return None
569 if after_member:
570 active.member_pks.add(event.subject.pk)
571 else:
572 active.member_pks.discard(event.subject.pk)
574 return replace(
575 event,
576 subscription_id=active.subscription_id,
577 list_action=list_action,
578 )
580 async def _replay_requested_cursor(
581 self,
582 active: ActiveSubscription,
583 requested_cursor: int | None,
584 *,
585 through_cursor: int,
586 ) -> None:
587 if requested_cursor is None or requested_cursor == through_cursor:
588 return
589 if not await database_sync_to_async(self._replay_is_complete)(requested_cursor):
590 await self.transport.send_resync(
591 ResyncRequired(
592 subscription_id=active.subscription_id,
593 cursor=through_cursor,
594 reason="cursor_pruned",
595 )
596 )
597 return
598 if active.request.kind == SubscriptionKind.OBJECT:
599 await self._replay_object_subscription(active, requested_cursor, through_cursor=through_cursor)
600 return
602 has_events = await database_sync_to_async(self._has_collection_replay_events)(
603 active,
604 requested_cursor,
605 through_cursor,
606 )
607 if has_events: 607 ↛ exitline 607 didn't return from function '_replay_requested_cursor' because the condition on line 607 was always true
608 await self.transport.send_resync(
609 ResyncRequired(
610 subscription_id=active.subscription_id,
611 cursor=through_cursor,
612 )
613 )
615 async def _prepare_subscription(self, subscription: SubscriptionRequest) -> None:
616 prepare = getattr(self.transport, "prepare_subscription", None)
617 if prepare is not None:
618 await prepare(subscription)
620 async def _send_catch_up_messages(self, messages: tuple[StreamEvent | ResyncRequired, ...]) -> None:
621 for message in messages:
622 if isinstance(message, ResyncRequired): 622 ↛ 623line 622 didn't jump to line 623 because the condition on line 622 was never true
623 await self.transport.send_resync(message)
624 else:
625 await self.transport.send_event(message)
627 async def _replay_object_subscription(
628 self,
629 active: ActiveSubscription,
630 requested_cursor: int,
631 *,
632 through_cursor: int,
633 ) -> None:
634 replay_events = await database_sync_to_async(self._object_replay_events)(
635 active,
636 requested_cursor,
637 through_cursor,
638 )
639 if replay_events is None:
640 await self.transport.send_resync(
641 ResyncRequired(
642 subscription_id=active.subscription_id,
643 cursor=through_cursor,
644 reason="object_replay_limit_exceeded",
645 )
646 )
647 return
649 for event in replay_events:
650 await self.transport.send_event(event)