Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
125 changes: 104 additions & 21 deletions src/volcano_sdk/realtime.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,11 @@
from typing import TYPE_CHECKING, Any, Literal, Protocol, TypeAlias, cast
from urllib.parse import quote, urlsplit, urlunsplit

from ._realtime_fetch_worker import (
PostgresFetchJob,
PostgresFetchOutcome,
PostgresFetchWorker,
)
from .database import Database
from .models import JSONValue, _freeze_json

Expand All @@ -34,6 +39,7 @@
importlib.import_module("centrifuge").CentrifugeError,
)
CALLBACK_QUEUE_LIMIT = 128
POSTGRES_QUEUE_LIMIT = 128
NO_PENDING_CALLBACK = object()
CALLBACK_QUEUE_FULL_MESSAGE = (
"Volcano realtime callback queue is full; publication dropped"
Expand Down Expand Up @@ -154,6 +160,19 @@ class _PostgresDeliveryIdentity:
subscription_epoch: int


@dataclass(frozen=True, slots=True)
class _PostgresDelivery:
change: PostgresChange
identity: _PostgresDeliveryIdentity


@dataclass(frozen=True, slots=True)
class _CallbackDelivery:
event: str
data: Any
postgres_identity: _PostgresDeliveryIdentity | None = None


@dataclass(slots=True)
class _SessionBoundDatabaseContext:
_transport: Transport
Expand Down Expand Up @@ -397,20 +416,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._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()
await 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._end_postgres_epoch()
await self._channel._presence_unsubscribed()

async def on_join(self, ctx: Any) -> None:
Expand Down Expand Up @@ -498,7 +517,7 @@ def __init__(
self._presence_lock = asyncio.Lock()
self._presence_sync_task: asyncio.Task[None] | None = None
self._presence_sync_pending = False
self._callback_queue: asyncio.Queue[tuple[str, Any]] = asyncio.Queue(
self._callback_queue: asyncio.Queue[_CallbackDelivery] = asyncio.Queue(
maxsize=CALLBACK_QUEUE_LIMIT
)
self._callback_task: asyncio.Task[None] | None = None
Expand All @@ -507,6 +526,8 @@ def __init__(
self._pending_presence_sync: Any = NO_PENDING_CALLBACK
self._postgres_epoch = 0
self._postgres_session_lineage = 0
self._postgres_lock = asyncio.Lock()
self._postgres_worker: PostgresFetchWorker[_PostgresDelivery] | None = None

@property
def name(self) -> str:
Expand Down Expand Up @@ -598,14 +619,25 @@ def _capture_postgres_delivery_identity(self) -> _PostgresDeliveryIdentity:
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()
async def _begin_postgres_epoch(self) -> None:
if self._type != "postgres":
return
await self._stop_postgres_worker()
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
async def _end_postgres_epoch(self) -> None:
if self._type != "postgres":
return
self._postgres_epoch += 1
await self._stop_postgres_worker()

async def _stop_postgres_worker(self) -> None:
async with self._postgres_lock:
worker = self._postgres_worker
self._postgres_worker = None
if worker is not None:
await worker.close()

def _postgres_delivery_is_current(
self,
Expand All @@ -622,14 +654,41 @@ def _postgres_delivery_is_current(

async def _receive_postgres_change(self, data: Any) -> None:
change = _postgres_change(data)
if change is None:
if change is None or not self._callbacks.get("*"):
return
if change.mode == "lightweight" and change.type == "DELETE":
old_record = change.old_record
if old_record is None and change.id is not None:
old_record = {"id": change.id}
change = replace(change, old_record=old_record, id=None, mode=None)
await self._emit("*", change)
identity = self._capture_postgres_delivery_identity()
if not self._postgres_delivery_is_current(identity):
return
delivery = _PostgresDelivery(change=change, identity=identity)
async with self._postgres_lock:
if not self._postgres_delivery_is_current(identity):
return
worker = self._postgres_worker
if worker is None:
worker = PostgresFetchWorker(
self._realtime._fetch_postgres_row,
self._deliver_postgres,
queue_limit=POSTGRES_QUEUE_LIMIT,
)
self._postgres_worker = worker
await worker.enqueue(PostgresFetchJob(request=None, fallback=delivery))

async def _deliver_postgres(
self,
outcome: PostgresFetchOutcome[_PostgresDelivery],
) -> None:
delivery = outcome.job.fallback
if self._postgres_delivery_is_current(delivery.identity):
await self._emit(
"*",
delivery.change,
postgres_identity=delivery.identity,
)

async def subscribe(self) -> None:
"""Subscribe to this channel."""
Expand All @@ -643,7 +702,13 @@ async def unsubscribe(self) -> None:
"""Unsubscribe from this channel."""
await self._realtime._unsubscribe(self)

async def _emit(self, event: str, data: Any) -> None:
async def _emit(
self,
event: str,
data: Any,
*,
postgres_identity: _PostgresDeliveryIdentity | None = None,
) -> None:
if not self._callbacks.get(event):
return
if event == "presence_sync":
Expand All @@ -657,7 +722,9 @@ async def _emit(self, event: str, data: Any) -> None:
)
self._callback_stop = None
try:
self._callback_queue.put_nowait((event, data))
self._callback_queue.put_nowait(
_CallbackDelivery(event, data, postgres_identity)
)
except asyncio.QueueFull:
if event == "presence_sync":
self._pending_presence_sync = data
Expand All @@ -682,11 +749,15 @@ async def _restart_callback_dispatcher(self, previous: asyncio.Task[None]) -> No

async def _dispatch_callbacks(self, stop: asyncio.Event) -> None:
while not stop.is_set():
event, data = await self._callback_queue.get()
delivery = await self._callback_queue.get()
try:
for callback in tuple(self._callbacks.get(event, [])):
if not self._callback_delivery_is_current(delivery):
continue
for callback in tuple(self._callbacks.get(delivery.event, [])):
if not self._callback_delivery_is_current(delivery):
break
active_task = asyncio.create_task(
self._run_callback(callback, data)
self._run_callback(callback, delivery)
)
self._active_callback_task = active_task
try:
Expand All @@ -707,15 +778,25 @@ async def _dispatch_callbacks(self, stop: asyncio.Event) -> None:
self._callback_queue.task_done()
self._enqueue_pending_presence_sync()

def _callback_delivery_is_current(self, delivery: _CallbackDelivery) -> bool:
identity = delivery.postgres_identity
return identity is None or self._postgres_delivery_is_current(identity)

def _enqueue_pending_presence_sync(self) -> None:
pending = self._pending_presence_sync
if pending is NO_PENDING_CALLBACK or self._callback_queue.full():
return
self._pending_presence_sync = NO_PENDING_CALLBACK
self._callback_queue.put_nowait(("presence_sync", pending))
self._callback_queue.put_nowait(_CallbackDelivery("presence_sync", pending))

async def _run_callback(self, callback: MessageCallback, data: Any) -> None:
result = callback(data)
async def _run_callback(
self,
callback: MessageCallback,
delivery: _CallbackDelivery,
) -> None:
if not self._callback_delivery_is_current(delivery):
return
result = callback(delivery.data)
if inspect.isawaitable(result):
await result

Expand Down Expand Up @@ -850,6 +931,8 @@ async def _cancel_presence_sync(self) -> None:

async def _reset(self) -> None:
self._subscription = None
self._subscribed = False
await self._end_postgres_epoch()
await self._cancel_presence_sync()
task = self._callback_task
active_task = self._active_callback_task
Expand Down
76 changes: 74 additions & 2 deletions tests/unit/test_realtime.py
Original file line number Diff line number Diff line change
Expand Up @@ -503,6 +503,73 @@ async def scenario() -> None:
asyncio.run(scenario())


def test_realtime_drops_queued_postgres_callbacks_from_an_old_epoch() -> 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:
def insert(record_id: int) -> dict[str, Any]:
return {
"type": "INSERT",
"schema": "public",
"table": "messages",
"record": {"id": record_id},
"timestamp": "2026-09-03T12:00:00Z",
}

received: list[int] = []
first_started = asyncio.Event()
release_first = asyncio.Event()
third_received = asyncio.Event()
channel = client.realtime.channel(
"public:messages",
channel_type="postgres",
)

async def on_insert(change: Any) -> None:
record_id = change.record["id"]
if record_id == 1:
first_started.set()
await release_first.wait()
received.append(record_id)
if record_id == 3:
third_received.set()

channel.on_postgres_changes(
"INSERT",
schema="public",
table="messages",
callback=on_insert,
)
await channel.subscribe()
subscription = official.subscription
assert subscription is not None

await subscription.emit(insert(1))
await first_started.wait()
first_worker = channel._postgres_worker
assert first_worker is not None
await subscription.emit(insert(2))

await subscription.emit_subscribing()
assert channel._postgres_worker is None
release_first.set()
await subscription.emit_subscribed()
await subscription.emit(insert(3))
await asyncio.wait_for(third_received.wait(), timeout=0.2)

assert received == [1, 3]
assert channel._postgres_worker is not first_worker
await client.realtime.disconnect()

asyncio.run(scenario())


def test_realtime_routes_immutable_rls_scoped_postgres_changes() -> None:
official = FakeCentrifugeClient()
client = VolcanoClient(
Expand Down Expand Up @@ -1082,15 +1149,20 @@ async def wait_forever() -> None:
blocker = asyncio.create_task(wait_forever())
channel._callback_task = blocker
for _ in range(channel._callback_queue.maxsize):
channel._callback_queue.put_nowait(("presence_sync", {"version": 0}))
channel._callback_queue.put_nowait(
realtime_module._CallbackDelivery(
"presence_sync",
{"version": 0},
)
)

await channel._emit("presence_sync", {"version": 1})
channel._callback_queue.get_nowait()
channel._callback_queue.task_done()
await channel._emit("presence_sync", {"version": 2})
queued: list[Any] = []
while not channel._callback_queue.empty():
queued.append(channel._callback_queue.get_nowait()[1])
queued.append(channel._callback_queue.get_nowait().data)
channel._callback_queue.task_done()
channel._enqueue_pending_presence_sync()
blocker.cancel()
Expand Down
Loading