diff --git a/src/volcano_sdk/realtime.py b/src/volcano_sdk/realtime.py index d2312739..4f1bdbba 100644 --- a/src/volcano_sdk/realtime.py +++ b/src/volcano_sdk/realtime.py @@ -16,6 +16,7 @@ if TYPE_CHECKING: from ._transport import Transport + from .models import Session MessageCallback = Callable[[Any], Any] RealtimeCallback = Callable[[Any], Any] @@ -47,6 +48,8 @@ SUBSCRIPTION_REGISTRY_UNAVAILABLE = ( "centrifuge client subscription registry is unavailable" ) +NO_ACTIVE_SESSION = "No active session" +CONNECTION_SESSION_UNAVAILABLE = "Realtime connection has no session binding" def _empty_presence_data() -> Mapping[str, JSONValue]: @@ -145,6 +148,12 @@ class _PostgresFetchRequest: row_id: JSONValue +@dataclass(frozen=True, slots=True) +class _PostgresDeliveryIdentity: + session_lineage: int + subscription_epoch: int + + @dataclass(slots=True) class _SessionBoundDatabaseContext: _transport: Transport @@ -213,6 +222,8 @@ def _anon_token(self) -> str: ... def _session_token(self) -> str: ... + def _capture_session_binding(self) -> tuple[int, int, Session | None]: ... + class CentrifugeSubscription(Protocol): """Centrifuge subscription operations used by the SDK.""" @@ -386,17 +397,20 @@ async def on_publication(self, ctx: PublicationContext) -> None: async def on_subscribing(self, ctx: Any) -> None: del ctx self._channel._subscribed = False + self._channel._end_postgres_epoch() await self._channel._presence_unsubscribed() async def on_subscribed(self, ctx: Any) -> None: del ctx self._channel._subscribed = True + self._channel._begin_postgres_epoch() if self._channel._type == "presence": self._channel._schedule_presence_sync() async def on_unsubscribed(self, ctx: Any) -> None: del ctx self._channel._subscribed = False + self._channel._end_postgres_epoch() await self._channel._presence_unsubscribed() async def on_join(self, ctx: Any) -> None: @@ -491,6 +505,8 @@ def __init__( self._active_callback_task: asyncio.Task[None] | None = None self._callback_stop: asyncio.Event | None = None self._pending_presence_sync: Any = NO_PENDING_CALLBACK + self._postgres_epoch = 0 + self._postgres_session_lineage = 0 @property def name(self) -> str: @@ -576,6 +592,34 @@ def _ensure_presence(self) -> None: if self._type != "presence": raise ValueError(PRESENCE_ONLY) + def _capture_postgres_delivery_identity(self) -> _PostgresDeliveryIdentity: + return _PostgresDeliveryIdentity( + session_lineage=self._postgres_session_lineage, + subscription_epoch=self._postgres_epoch, + ) + + def _begin_postgres_epoch(self) -> None: + if self._type == "postgres": + self._postgres_epoch += 1 + self._postgres_session_lineage = self._realtime._connection_lineage() + + def _end_postgres_epoch(self) -> None: + if self._type == "postgres": + self._postgres_epoch += 1 + + def _postgres_delivery_is_current( + self, + identity: _PostgresDeliveryIdentity, + ) -> bool: + _generation, lineage, _session = ( + self._realtime._client_context._capture_session_binding() + ) + return ( + self._subscribed + and identity.subscription_epoch == self._postgres_epoch + and identity.session_lineage == lineage + ) + async def _receive_postgres_change(self, data: Any) -> None: change = _postgres_change(data) if change is None: @@ -842,6 +886,7 @@ def __init__( self._api_url = api_url self._client_factory = client_factory self._connection: _VolcanoCentrifugeConnection | None = None + self._connection_session_lineage: int | None = None self._connection_lock = asyncio.Lock() self._channels: dict[str, Channel] = {} self._removing_channels: set[str] = set() @@ -1031,7 +1076,17 @@ async def _remove_channel(self, channel: Channel) -> None: await channel._reset() async def _token(self) -> str: - return self._client_context._session_token() + _generation, lineage, session = self._client_context._capture_session_binding() + if session is None: + raise RuntimeError(NO_ACTIVE_SESSION) + self._connection_session_lineage = lineage + return session.access_token + + def _connection_lineage(self) -> int: + lineage = self._connection_session_lineage + if lineage is None: + raise RuntimeError(CONNECTION_SESSION_UNAVAILABLE) + return lineage def _address(self) -> str: parsed = urlsplit(self._api_url) @@ -1046,15 +1101,23 @@ async def _connect(self) -> _VolcanoCentrifugeConnection: async def _connect_locked(self) -> _VolcanoCentrifugeConnection: if self._connection is not None: return self._connection + _generation, lineage, session = self._client_context._capture_session_binding() + if session is None: + raise RuntimeError(NO_ACTIVE_SESSION) connection = _VolcanoCentrifugeConnection( self._client_factory( self._address(), events=_ClientEvents(self), - token=self._client_context._session_token(), + token=session.access_token, get_token=self._token, ) ) - await connection.connect() + self._connection_session_lineage = lineage + try: + await connection.connect() + except BaseException: + self._connection_session_lineage = None + raise self._connection = connection return connection @@ -1122,6 +1185,7 @@ async def disconnect(self) -> None: async with self._connection_lock: connection = self._connection self._connection = None + self._connection_session_lineage = None try: if connection is not None: await connection.disconnect() diff --git a/tests/unit/test_realtime.py b/tests/unit/test_realtime.py index 9278e85d..e0f3ed2f 100644 --- a/tests/unit/test_realtime.py +++ b/tests/unit/test_realtime.py @@ -412,6 +412,97 @@ def test_realtime_channel_exposes_its_canonical_name() -> None: ) +def test_realtime_postgres_delivery_identity_changes_on_reauthentication() -> None: + official = FakeCentrifugeClient() + client = VolcanoClient( + anon_key="anon-key", + _transport=AuthTransport(), + _realtime_client_factory=FakeCentrifugeFactory(official), + ) + client.auth.sign_in(email="user@example.com", password="secret") + + async def scenario() -> None: + channel = client.realtime.channel( + "public:messages", + channel_type="postgres", + ) + await channel.subscribe() + identity = channel._capture_postgres_delivery_identity() + + assert channel._postgres_delivery_is_current(identity) + + client.auth.sign_in(email="user@example.com", password="secret") + + identity_after_reauthentication = channel._capture_postgres_delivery_identity() + assert not channel._postgres_delivery_is_current(identity) + assert not channel._postgres_delivery_is_current( + identity_after_reauthentication + ) + await client.realtime.disconnect() + + asyncio.run(scenario()) + + +def test_realtime_postgres_delivery_identity_changes_on_resubscription() -> None: + official = FakeCentrifugeClient() + client = VolcanoClient( + anon_key="anon-key", + _transport=AuthTransport(), + _realtime_client_factory=FakeCentrifugeFactory(official), + ) + client.auth.sign_in(email="user@example.com", password="secret") + + async def scenario() -> None: + channel = client.realtime.channel( + "public:messages", + channel_type="postgres", + ) + await channel.subscribe() + identity = channel._capture_postgres_delivery_identity() + subscription = official.subscription + assert subscription is not None + + await subscription.emit_subscribing() + assert not channel._postgres_delivery_is_current(identity) + + await subscription.emit_subscribed() + next_identity = channel._capture_postgres_delivery_identity() + + assert not channel._postgres_delivery_is_current(identity) + assert channel._postgres_delivery_is_current(next_identity) + await client.realtime.disconnect() + + asyncio.run(scenario()) + + +def test_realtime_postgres_delivery_identity_uses_connection_session() -> None: + official = FakeCentrifugeClient() + client = VolcanoClient( + anon_key="anon-key", + _transport=AuthTransport(), + _realtime_client_factory=FakeCentrifugeFactory(official), + ) + client.auth.sign_in(email="user@example.com", password="secret") + + async def scenario() -> None: + broadcast = client.realtime.channel("updates") + await broadcast.subscribe() + + client.auth.sign_in(email="user@example.com", password="secret") + postgres = client.realtime.channel( + "public:messages", + channel_type="postgres", + ) + await postgres.subscribe() + + identity = postgres._capture_postgres_delivery_identity() + + assert not postgres._postgres_delivery_is_current(identity) + await client.realtime.disconnect() + + asyncio.run(scenario()) + + def test_realtime_routes_immutable_rls_scoped_postgres_changes() -> None: official = FakeCentrifugeClient() client = VolcanoClient(