sqlspec 0.47.0__cp314-cp314-win_amd64.whl

sqlspec/adapters/psqlpy/events/_hub.py

@@ -0,0 +1,204 @@
+"""Persistent listener hub used by psqlpy event backends.
+
+Owns the psqlpy Listener handle for the backend's lifetime. Each channel
+gets a single driver-level callback registered once via ``add_callback``
+that dispatches into one queue per consumer task. Subscribe / unsubscribe
+are serialized under an asyncio.Lock; ``clear_channel_callbacks`` only
+runs on unsubscribe / shutdown so concurrent peers cannot wipe each
+other's callbacks mid-iteration.
+"""
+
+# pyright: reportPrivateUsage=false
+
+import asyncio
+import contextlib
+from typing import TYPE_CHECKING, Any, cast
+from weakref import WeakKeyDictionary
+
+from sqlspec.extensions.events import normalize_event_channel_name
+from sqlspec.protocols import NotificationProtocol
+from sqlspec.utils.logging import get_logger
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from sqlspec.adapters.psqlpy._typing import PsqlpyListener
+    from sqlspec.adapters.psqlpy.config import PsqlpyConfig
+
+logger = get_logger("sqlspec.adapters.psqlpy.events.hub")
+
+__all__ = ("PsqlpyListenerHub",)
+
+_PSQLPY_CALLBACK_MIN_ARGS = 3
+
+
+class _PsqlpyHubCallback:
+    """Driver-facing callback that normalizes psqlpy notify args and dispatches."""
+
+    __slots__ = ("_channel", "_dispatch")
+
+    def __init__(self, channel: str, dispatch: "Callable[[str, str], None]") -> None:
+        self._channel = channel
+        self._dispatch = dispatch
+
+    async def __call__(self, *args: Any) -> None:
+        if not args:
+            return
+        notified_channel: str | None = None
+        payload: str | None = None
+        if len(args) == 1:
+            message = args[0]
+            if isinstance(message, NotificationProtocol):
+                notified_channel = message.channel
+                payload = message.payload
+        elif len(args) >= _PSQLPY_CALLBACK_MIN_ARGS:
+            value1 = cast("str", args[1])
+            value2 = cast("str", args[2])
+            if value1 == self._channel:
+                notified_channel = value1
+                payload = value2
+            elif value2 == self._channel:
+                notified_channel = value2
+                payload = value1
+        if notified_channel is None or notified_channel != self._channel or payload is None:
+            return
+        self._dispatch(self._channel, payload)
+
+
+class PsqlpyListenerHub:
+    """Per-channel persistent listener for psqlpy event backends."""
+
+    __slots__ = (
+        "_callbacks",
+        "_config",
+        "_listener",
+        "_listener_started",
+        "_lock",
+        "_pool_destroying_registered",
+        "_queues",
+        "_shutting_down",
+    )
+
+    def __init__(self, config: "PsqlpyConfig") -> None:
+        self._config = config
+        self._lock = asyncio.Lock()
+        self._queues: dict[str, WeakKeyDictionary[asyncio.Task[Any], asyncio.Queue[str]]] = {}
+        self._callbacks: dict[str, _PsqlpyHubCallback] = {}
+        self._listener: Any | None = None
+        self._listener_started = False
+        self._shutting_down = False
+        self._pool_destroying_registered = False
+
+    async def subscribe(self, channel: str) -> None:
+        async with self._lock:
+            if self._shutting_down:
+                msg = "PsqlpyListenerHub is shutting down"
+                raise RuntimeError(msg)
+            if channel in self._queues:
+                return
+            normalize_event_channel_name(channel)
+            listener = await self._ensure_listener_locked()
+            callback = _PsqlpyHubCallback(channel, self._dispatch)
+            self._queues[channel] = WeakKeyDictionary()
+            self._callbacks[channel] = callback
+            try:
+                await listener.add_callback(channel=channel, callback=callback.__call__)
+            except Exception:
+                self._queues.pop(channel, None)
+                self._callbacks.pop(channel, None)
+                raise
+            if not self._listener_started:
+                listener.listen()
+                self._listener_started = True
+                # Give psqlpy a brief moment to actually start its consumer
+                await asyncio.sleep(0.05)
+
+    async def unsubscribe(self, channel: str) -> None:
+        async with self._lock:
+            if channel not in self._queues:
+                return
+            listener = self._listener
+            self._queues.pop(channel, None)
+            self._callbacks.pop(channel, None)
+            if listener is None:
+                return
+            with contextlib.suppress(Exception):
+                await listener.clear_channel_callbacks(channel=channel)
+
+    async def dequeue(self, channel: str, poll_interval: float) -> "str | None":
+        if channel not in self._queues:
+            await self.subscribe(channel)
+        queue = self._get_consumer_queue(channel)
+        if queue is None:
+            return None
+        try:
+            return await asyncio.wait_for(queue.get(), timeout=poll_interval)
+        except asyncio.TimeoutError:
+            return None
+
+    async def shutdown(self) -> None:
+        async with self._lock:
+            if self._shutting_down:
+                return
+            self._shutting_down = True
+            listener = self._listener
+            channels = list(self._queues.keys())
+            was_started = self._listener_started
+            self._queues.clear()
+            self._callbacks.clear()
+            self._listener = None
+            self._listener_started = False
+        if listener is not None:
+            for channel in channels:
+                with contextlib.suppress(Exception):
+                    await listener.clear_channel_callbacks(channel=channel)
+            if was_started:
+                with contextlib.suppress(Exception):
+                    listener.abort_listen()
+            with contextlib.suppress(Exception):
+                await listener.shutdown()
+        self._shutting_down = False
+
+    def is_subscribed(self, channel: str) -> bool:
+        return channel in self._queues
+
+    def _get_consumer_queue(self, channel: str) -> "asyncio.Queue[str] | None":
+        task = asyncio.current_task()
+        if task is None:  # pragma: no cover - coroutine dequeue calls always run in a task
+            msg = "PsqlpyListenerHub.dequeue requires an active asyncio task"
+            raise RuntimeError(msg)
+        queues = self._queues.get(channel)
+        if queues is None:
+            return None
+        queue = queues.get(task)
+        if queue is None:
+            queue = asyncio.Queue()
+            queues[task] = queue
+        return queue
+
+    async def _ensure_listener_locked(self) -> "PsqlpyListener":
+        if self._listener is not None:
+            return cast("PsqlpyListener", self._listener)
+        pool = await self._config.provide_pool()
+        listener = pool.listener()
+        await listener.startup()
+        self._listener = listener
+        self._register_pool_destroying()
+        return listener
+
+    def _register_pool_destroying(self) -> None:
+        if self._pool_destroying_registered:
+            return
+        runtime = self._config.get_observability_runtime()
+        runtime.register_lifecycle_hook("on_pool_destroying", self._pool_destroying_hook)
+        self._pool_destroying_registered = True
+
+    def _pool_destroying_hook(self, _context: "dict[str, Any]") -> "Any":
+        return self.shutdown()
+
+    def _dispatch(self, channel: str, payload: str) -> None:
+        queues = self._queues.get(channel)
+        if queues is None:
+            return
+        for queue in list(queues.values()):
+            queue.put_nowait(payload)

sqlspec/adapters/psqlpy/events/backend.py

@@ -0,0 +1,210 @@
+"""Psqlpy LISTEN/NOTIFY and hybrid event backends."""
+
+# pyright: reportPrivateUsage=false
+
+import logging
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Any, cast
+
+from sqlspec.adapters.psqlpy.events._hub import PsqlpyListenerHub
+from sqlspec.core import SQL
+from sqlspec.exceptions import ImproperConfigurationError
+from sqlspec.extensions.events import (
+    AsyncTableEventQueue,
+    EventMessage,
+    build_queue_backend,
+    decode_notify_payload,
+    encode_notify_payload,
+)
+from sqlspec.utils.logging import get_logger, log_with_context
+from sqlspec.utils.serializers import from_json, to_json
+from sqlspec.utils.uuids import uuid4
+
+if TYPE_CHECKING:
+    from sqlspec.adapters.psqlpy.config import PsqlpyConfig
+
+
+logger = get_logger("sqlspec.events.psqlpy")
+
+__all__ = ("PsqlpyEventsBackend", "PsqlpyHybridEventsBackend", "create_event_backend")
+
+
+def _extract_event_id(payload: "str | None") -> "str | None":
+    if not payload:
+        return None
+    raw = from_json(payload)
+    if isinstance(raw, dict):
+        event_id = raw.get("event_id")
+        return event_id if isinstance(event_id, str) else None
+    return None
+
+
+class PsqlpyEventsBackend:
+    """Native LISTEN/NOTIFY backend for psqlpy adapters."""
+
+    __slots__ = ("_config", "_hub", "_runtime")
+
+    supports_sync = False
+    supports_async = True
+    backend_name = "listen_notify"
+
+    def __init__(self, config: "PsqlpyConfig") -> None:
+        if "psqlpy" not in type(config).__module__:
+            msg = "Psqlpy events backend requires a Psqlpy adapter"
+            raise ImproperConfigurationError(msg)
+        self._config = config
+        self._runtime = config.get_observability_runtime()
+        self._hub: PsqlpyListenerHub | None = None
+        log_with_context(
+            logger,
+            logging.DEBUG,
+            "event.listen",
+            adapter_name="psqlpy",
+            backend_name=self.backend_name,
+            mode="async",
+            status="backend_ready",
+        )
+
+    async def publish(self, channel: str, payload: "dict[str, Any]", metadata: "dict[str, Any] | None" = None) -> str:
+        event_id = uuid4().hex
+        async with self._config.provide_session() as driver:
+            await driver.execute_script(
+                SQL("SELECT pg_notify($1, $2)", channel, encode_notify_payload(event_id, payload, metadata))
+            )
+            await driver.commit()
+        self._runtime.increment_metric("events.publish.native")
+        return event_id
+
+    async def dequeue(self, channel: str, poll_interval: float) -> EventMessage | None:
+        hub = self._ensure_hub()
+        payload = await hub.dequeue(channel, poll_interval)
+        if payload is None:
+            return None
+        return decode_notify_payload(channel, payload)
+
+    async def ack(self, _event_id: str) -> None:
+        self._runtime.increment_metric("events.ack")
+
+    async def nack(self, _event_id: str) -> None:
+        """Return an event to the queue (no-op for native LISTEN/NOTIFY)."""
+
+    async def shutdown(self) -> None:
+        hub = self._hub
+        if hub is not None:
+            self._hub = None
+            await hub.shutdown()
+
+    def _ensure_hub(self) -> PsqlpyListenerHub:
+        if self._hub is None:
+            self._hub = PsqlpyListenerHub(self._config)
+        return self._hub
+
+
+class PsqlpyHybridEventsBackend:
+    """Durable hybrid backend combining queue storage with LISTEN/NOTIFY wakeups."""
+
+    __slots__ = ("_config", "_hub", "_queue", "_runtime")
+
+    supports_sync = False
+    supports_async = True
+    backend_name = "listen_notify_durable"
+
+    def __init__(self, config: "PsqlpyConfig", queue: "AsyncTableEventQueue") -> None:
+        if "psqlpy" not in type(config).__module__:
+            msg = "Psqlpy hybrid backend requires a Psqlpy adapter"
+            raise ImproperConfigurationError(msg)
+        self._config = config
+        self._queue = queue
+        self._runtime = config.get_observability_runtime()
+        self._hub: PsqlpyListenerHub | None = None
+        log_with_context(
+            logger,
+            logging.DEBUG,
+            "event.listen",
+            adapter_name="psqlpy",
+            backend_name=self.backend_name,
+            mode="async",
+            status="backend_ready",
+        )
+
+    async def publish(self, channel: str, payload: "dict[str, Any]", metadata: "dict[str, Any] | None" = None) -> str:
+        event_id = uuid4().hex
+        await self._publish_durable(channel, event_id, payload, metadata)
+        self._runtime.increment_metric("events.publish.native")
+        return event_id
+
+    async def dequeue(self, channel: str, poll_interval: float) -> EventMessage | None:
+        hub = self._ensure_hub()
+        payload = await hub.dequeue(channel, poll_interval)
+        if payload is None:
+            return await self._queue.dequeue(channel)
+        event_id = _extract_event_id(payload)
+        if event_id is not None:
+            event = await self._queue.dequeue_by_event_id(event_id)
+            if event is not None:
+                return event
+        return await self._queue.dequeue(channel)
+
+    async def ack(self, event_id: str) -> None:
+        await self._queue.ack(event_id)
+        self._runtime.increment_metric("events.ack")
+
+    async def nack(self, event_id: str) -> None:
+        await self._queue.nack(event_id)
+        self._runtime.increment_metric("events.nack")
+
+    async def shutdown(self) -> None:
+        hub = self._hub
+        if hub is not None:
+            self._hub = None
+            await hub.shutdown()
+
+    def _ensure_hub(self) -> PsqlpyListenerHub:
+        if self._hub is None:
+            self._hub = PsqlpyListenerHub(self._config)
+        return self._hub
+
+    async def _publish_durable(
+        self, channel: str, event_id: str, payload: "dict[str, Any]", metadata: "dict[str, Any] | None"
+    ) -> None:
+        now = datetime.now(timezone.utc)
+        async with self._config.provide_session() as driver:
+            await driver.execute(
+                SQL(
+                    self._queue._upsert_sql,  # pyright: ignore[reportPrivateUsage]
+                    {
+                        "event_id": event_id,
+                        "channel": channel,
+                        "payload_json": payload,
+                        "metadata_json": metadata,
+                        "status": "pending",
+                        "available_at": now,
+                        "lease_expires_at": None,
+                        "attempts": 0,
+                        "created_at": now,
+                    },
+                    statement_config=self._queue._statement_config,  # pyright: ignore[reportPrivateUsage]
+                )
+            )
+            await driver.execute_script(SQL("SELECT pg_notify($1, $2)", channel, to_json({"event_id": event_id})))
+            await driver.commit()
+
+
+def create_event_backend(
+    config: "PsqlpyConfig", backend_name: str, extension_settings: "dict[str, Any]"
+) -> PsqlpyEventsBackend | PsqlpyHybridEventsBackend | None:
+    """Factory used by EventChannel to create the native psqlpy backend."""
+    match backend_name:
+        case "listen_notify":
+            try:
+                return PsqlpyEventsBackend(config)
+            except ImproperConfigurationError:
+                return None
+        case "listen_notify_durable":
+            queue = cast("AsyncTableEventQueue", build_queue_backend(config, extension_settings, adapter_name="psqlpy"))
+            try:
+                return PsqlpyHybridEventsBackend(config, queue)
+            except ImproperConfigurationError:
+                return None
+        case _:
+            return None

sqlspec/adapters/psqlpy/events/store.py

@@ -0,0 +1,20 @@
+"""Psqlpy event queue store."""
+
+from sqlspec.adapters.psqlpy.config import PsqlpyConfig
+from sqlspec.extensions.events import BaseEventQueueStore
+
+__all__ = ("PsqlpyEventQueueStore",)
+
+
+class PsqlpyEventQueueStore(BaseEventQueueStore[PsqlpyConfig]):
+    """Provide PostgreSQL column mappings for the queue table.
+
+    PostgreSQL uses JSONB for efficient binary JSON storage with indexing support,
+    and TIMESTAMPTZ for timezone-aware timestamps.
+    """
+
+    __slots__ = ()
+
+    def _column_types(self) -> "tuple[str, str, str]":
+        """Return PostgreSQL-optimized column types for the event queue."""
+        return "JSONB", "JSONB", "TIMESTAMPTZ"

sqlspec/adapters/psqlpy/litestar/__init__.py

@@ -0,0 +1,5 @@
+"""Litestar integration for psqlpy adapter."""
+
+from sqlspec.adapters.psqlpy.litestar.store import PsqlpyStore
+
+__all__ = ("PsqlpyStore",)