feat(events): update event backends to use psqlpy's Listener API and refactor connection configurations

cofin · cofin · commit 721fb9affc73 · 2025-12-16T00:27:49.000Z
diff --git a/.claude/agents/docs-vision.md b/.claude/agents/docs-vision.md
@@ -2,7 +2,7 @@
 name: docs-vision
 description: Documentation excellence, quality gate validation, and workspace cleanup specialist - ensures code quality, comprehensive docs, and clean workspace before completion
 tools: mcp__context7__resolve-library-id, mcp__context7__get-library-docs, WebSearch, Read, Write, Edit, Glob, Grep, Bash, Find, Task
-model: sonnet
+model: opus
 standards_uri: ../AGENTS.md#mandatory-code-quality-standards
 guides_root: ../docs/guides/
 workspace_root: ../specs/active/
diff --git a/.claude/agents/expert.md b/.claude/agents/expert.md
@@ -2,7 +2,7 @@
 name: expert
 description: SQLSpec domain expert with comprehensive knowledge of database adapters, SQL parsing, type system, storage backends, and Litestar integration
 tools: mcp__context7__resolve-library-id, mcp__context7__get-library-docs, WebSearch, mcp__zen__analyze, mcp__zen__thinkdeep, mcp__zen__debug, Read, Write, Edit, Glob, Grep, Bash, Find, Task
-model: sonnet
+model: opus
 standards_uri: ../AGENTS.md#mandatory-code-quality-standards
 guides_root: ../docs/guides/
 workspace_root: ../specs/active/
diff --git a/.claude/agents/prd.md b/.claude/agents/prd.md
@@ -2,7 +2,7 @@
 name: prd
 description: Product Requirements and Design agent for complex SQLSpec development spanning multiple files, adapters, and features
 tools: mcp__zen__planner, mcp__context7__resolve-library-id, mcp__context7__get-library-docs, WebSearch, mcp__zen__consensus, Read, Write, Edit, Glob, Grep, Bash, Find, Task
-model: sonnet
+model: opus
 standards_uri: ../AGENTS.md#mandatory-code-quality-standards
 guides_root: ../docs/guides/
 workspace_root: ../specs/active/
diff --git a/.claude/agents/testing.md b/.claude/agents/testing.md
@@ -2,7 +2,7 @@
 name: testing
 description: Comprehensive testing specialist for SQLSpec - creates unit tests, integration tests, fixtures, and validates test coverage across all database adapters
 tools: mcp__context7__resolve-library-id, mcp__context7__get-library-docs, Read, Write, Edit, Glob, Grep, Bash, Find, Task
-model: sonnet
+model: opus
 standards_uri: ../AGENTS.md#mandatory-code-quality-standards
 guides_root: ../docs/guides/
 workspace_root: ../specs/active/
diff --git a/sqlspec/adapters/psqlpy/events/backend.py b/sqlspec/adapters/psqlpy/events/backend.py
@@ -14,6 +14,8 @@
 from sqlspec.utils.serializers import from_json, to_json
 
 if TYPE_CHECKING:
+    from psqlpy import Listener
+
     from sqlspec.adapters.psqlpy.config import PsqlpyConfig
 
 
@@ -27,7 +29,11 @@
 
 
 class PsqlpyEventsBackend:
-    """Native LISTEN/NOTIFY backend for psqlpy adapters."""
+    """Native LISTEN/NOTIFY backend for psqlpy adapters.
+
+    Uses psqlpy's Listener API which provides a dedicated connection for
+    receiving PostgreSQL NOTIFY messages via callbacks or async iteration.
+    """
 
     supports_sync = False
     supports_async = True
@@ -39,8 +45,8 @@ def __init__(self, config: "PsqlpyConfig") -> None:
             raise ImproperConfigurationError(msg)
         self._config = config
         self._runtime = config.get_observability_runtime()
-        self._listen_connection_async: Any | None = None
-        self._listen_connection_async_cm: Any | None = None
+        self._listener: Any | None = None
+        self._listener_started: bool = False
 
     async def publish_async(self, channel: str, payload: "dict[str, Any]", metadata: "dict[str, Any] | None" = None) -> str:
         event_id = uuid.uuid4().hex
@@ -56,25 +62,36 @@ def publish(self, *_: Any, **__: Any) -> str:
         raise ImproperConfigurationError(msg)
 
     async def dequeue_async(self, channel: str, poll_interval: float) -> EventMessage | None:
-        connection = await self._ensure_async_listener(channel)
-        await connection.execute(f"LISTEN {channel}")
-        future: asyncio.Future[str] = asyncio.get_running_loop().create_future()
+        listener = await self._ensure_listener(channel)
+        received_payload: str | None = None
+        event = asyncio.Event()
 
-        def _callback(_conn: Any, _pid: int, notified_channel: str, payload: str) -> None:
-            if notified_channel != channel or future.done():
-                return
-            future.set_result(payload)
+        async def _callback(
+            _connection: Any,
+            payload: str,
+            notified_channel: str,
+            _process_id: int,
+        ) -> None:
+            nonlocal received_payload
+            if notified_channel == channel and received_payload is None:
+                received_payload = payload
+                event.set()
 
-        connection.add_listener(channel, _callback)
-        try:
-            try:
-                payload = await asyncio.wait_for(future, timeout=poll_interval)
-            except asyncio.TimeoutError:
-                return None
-            return self._decode_payload(channel, payload)
-        finally:
-            with contextlib.suppress(Exception):
-                connection.remove_listener(channel, _callback)
+        await listener.add_callback(channel=channel, callback=_callback)
+
+        if not self._listener_started:
+            listener.listen()
+            self._listener_started = True
+            await asyncio.sleep(0.05)
+
+        with contextlib.suppress(asyncio.TimeoutError):
+            await asyncio.wait_for(event.wait(), timeout=poll_interval)
+
+        await listener.clear_channel_callbacks(channel=channel)
+
+        if received_payload is not None:
+            return self._decode_payload(channel, received_payload)
+        return None
 
     def dequeue(self, *_: Any, **__: Any) -> EventMessage | None:
         msg = "dequeue is not supported for sync Psqlpy backends"
@@ -87,15 +104,21 @@ def ack(self, _event_id: str) -> None:
         msg = "ack is not supported for sync Psqlpy backends"
         raise ImproperConfigurationError(msg)
 
-    async def _ensure_async_listener(self, channel: str) -> Any:
-        if self._listen_connection_async is None:
-            self._listen_connection_async_cm = self._config.provide_connection()
-            self._listen_connection_async = await self._listen_connection_async_cm.__aenter__()
-            try:
-                await self._listen_connection_async.set_autocommit(True)  # type: ignore[attr-defined]
-            except Exception:
-                pass
-        return self._listen_connection_async
+    async def _ensure_listener(self, channel: str) -> "Listener":
+        if self._listener is None:
+            pool = await self._config.provide_pool()
+            self._listener = pool.listener()
+            await self._listener.startup()
+        return self._listener
+
+    async def shutdown(self) -> None:
+        """Shutdown the listener and release resources."""
+        if self._listener is not None:
+            if self._listener_started:
+                self._listener.abort_listen()
+                self._listener_started = False
+            await self._listener.shutdown()
+            self._listener = None
 
     @staticmethod
     def _encode_payload(event_id: str, payload: "dict[str, Any]", metadata: "dict[str, Any] | None") -> str:
@@ -147,7 +170,11 @@ def _parse_timestamp(value: Any) -> datetime:
 
 
 class PsqlpyHybridEventsBackend:
-    """Durable hybrid backend combining queue storage with LISTEN/NOTIFY wakeups."""
+    """Durable hybrid backend combining queue storage with LISTEN/NOTIFY wakeups.
+
+    Uses psqlpy's Listener API for real-time notifications while persisting
+    events to a durable queue table.
+    """
 
     supports_sync = False
     supports_async = True
@@ -160,6 +187,8 @@ def __init__(self, config: "PsqlpyConfig", queue_backend: QueueEventBackend) ->
         self._config = config
         self._queue = queue_backend
         self._runtime = config.get_observability_runtime()
+        self._listener: Any | None = None
+        self._listener_started: bool = False
 
     async def publish_async(self, channel: str, payload: "dict[str, Any]", metadata: "dict[str, Any] | None" = None) -> str:
         event_id = uuid.uuid4().hex
@@ -172,31 +201,30 @@ def publish(self, *_: Any, **__: Any) -> str:
         raise ImproperConfigurationError(msg)
 
     async def dequeue_async(self, channel: str, poll_interval: float) -> EventMessage | None:
-        connection_cm = self._config.provide_connection()
-        connection = await connection_cm.__aenter__()
-        try:
-            listener = getattr(connection, "add_listener", None)
-            if listener is None:
-                return await self._queue.dequeue_async(channel)
-            future: asyncio.Future[str] = asyncio.get_running_loop().create_future()
-
-            def _callback(_conn: Any, _pid: int, notified_channel: str, payload: str) -> None:
-                if notified_channel != channel or future.done():
-                    return
-                future.set_result(payload)
-
-            listener(channel, _callback)
-            try:
-                await asyncio.wait_for(future, timeout=poll_interval)
-            except asyncio.TimeoutError:
-                return await self._queue.dequeue_async(channel)
-            return await self._queue.dequeue_async(channel)
-        finally:
-            with contextlib.suppress(Exception):
-                remove = getattr(connection, "remove_listener", None)
-                if remove:
-                    remove(channel, _callback)
-                await connection_cm.__aexit__(None, None, None)
+        listener = await self._ensure_listener(channel)
+        event = asyncio.Event()
+
+        async def _callback(
+            _connection: Any,
+            _payload: str,
+            notified_channel: str,
+            _process_id: int,
+        ) -> None:
+            if notified_channel == channel:
+                event.set()
+
+        await listener.add_callback(channel=channel, callback=_callback)
+
+        if not self._listener_started:
+            listener.listen()
+            self._listener_started = True
+            await asyncio.sleep(0.05)
+
+        with contextlib.suppress(asyncio.TimeoutError):
+            await asyncio.wait_for(event.wait(), timeout=poll_interval)
+
+        await listener.clear_channel_callbacks(channel=channel)
+        return await self._queue.dequeue_async(channel)
 
     async def ack_async(self, event_id: str) -> None:
         await self._queue.ack_async(event_id)
@@ -206,6 +234,22 @@ def ack(self, _event_id: str) -> None:
         msg = "ack is not supported for sync Psqlpy backends"
         raise ImproperConfigurationError(msg)
 
+    async def _ensure_listener(self, channel: str) -> "Listener":
+        if self._listener is None:
+            pool = await self._config.provide_pool()
+            self._listener = pool.listener()
+            await self._listener.startup()
+        return self._listener
+
+    async def shutdown(self) -> None:
+        """Shutdown the listener and release resources."""
+        if self._listener is not None:
+            if self._listener_started:
+                self._listener.abort_listen()
+                self._listener_started = False
+            await self._listener.shutdown()
+            self._listener = None
+
     async def _publish_durable_async(
         self, channel: str, event_id: str, payload: "dict[str, Any]", metadata: "dict[str, Any] | None"
     ) -> None:
diff --git a/sqlspec/adapters/psycopg/events/backend.py b/sqlspec/adapters/psycopg/events/backend.py
@@ -1,6 +1,5 @@
 """Psycopg LISTEN/NOTIFY and hybrid event backends."""
 
-import asyncio
 import contextlib
 import uuid
 from datetime import datetime, timezone
@@ -89,21 +88,15 @@ async def _ensure_async_listener(self, channel: str) -> Any:
         if self._listen_connection_async is None:
             self._listen_connection_async_cm = self._config.provide_connection()
             self._listen_connection_async = await self._listen_connection_async_cm.__aenter__()
-            try:
-                await self._listen_connection_async.set_autocommit(True)  # type: ignore[attr-defined]
-            except Exception:
-                pass
+            await self._listen_connection_async.set_autocommit(True)
             await self._listen_connection_async.execute(f"LISTEN {channel}")
         return self._listen_connection_async
 
     def _ensure_sync_listener(self, channel: str) -> Any:
         if self._listen_connection_sync is None:
             self._listen_connection_sync_cm = self._config.provide_connection()
             self._listen_connection_sync = self._listen_connection_sync_cm.__enter__()
-            try:
-                self._listen_connection_sync.autocommit = True
-            except Exception:
-                pass
+            self._listen_connection_sync.autocommit = True
             self._listen_connection_sync.execute(f"LISTEN {channel}")
         return self._listen_connection_sync
 
@@ -170,6 +163,10 @@ def __init__(self, config: "DatabaseConfigProtocol[Any, Any, Any]", queue_backen
         self._config = config
         self._queue = queue_backend
         self._runtime = config.get_observability_runtime()
+        self._listen_connection_async: Any | None = None
+        self._listen_connection_async_cm: Any | None = None
+        self._listen_connection_sync: Any | None = None
+        self._listen_connection_sync_cm: Any | None = None
 
     async def publish_async(self, channel: str, payload: "dict[str, Any]", metadata: "dict[str, Any] | None" = None) -> str:
         event_id = uuid.uuid4().hex
@@ -184,31 +181,33 @@ def publish(self, channel: str, payload: "dict[str, Any]", metadata: "dict[str,
         return event_id
 
     async def dequeue_async(self, channel: str, poll_interval: float) -> EventMessage | None:
-        connection_cm = self._config.provide_connection()
-        connection = await connection_cm.__aenter__()
-        try:
-            await connection.execute(f"LISTEN {channel}")
-            async for _notify in connection.notifies(timeout=poll_interval, stop_after=1):
-                break
-            return await self._queue.dequeue_async(channel)
-        finally:
-            with contextlib.suppress(Exception):
-                await connection.execute(f"UNLISTEN {channel}")
-                await connection_cm.__aexit__(None, None, None)
+        connection = await self._ensure_async_listener(channel)
+        async for _notify in connection.notifies(timeout=poll_interval, stop_after=1):
+            break
+        return await self._queue.dequeue_async(channel)
 
     def dequeue(self, channel: str, poll_interval: float) -> EventMessage | None:
-        connection_cm = self._config.provide_connection()
-        connection = connection_cm.__enter__()
-        try:
-            connection.execute(f"LISTEN {channel}")
-            notify_iter = connection.notifies(timeout=poll_interval, stop_after=1)
-            with contextlib.suppress(StopIteration):
-                next(notify_iter)
-            return self._queue.dequeue(channel)
-        finally:
-            with contextlib.suppress(Exception):
-                connection.execute(f"UNLISTEN {channel}")
-                connection_cm.__exit__(None, None, None)
+        connection = self._ensure_sync_listener(channel)
+        notify_iter = connection.notifies(timeout=poll_interval, stop_after=1)
+        with contextlib.suppress(StopIteration):
+            next(notify_iter)
+        return self._queue.dequeue(channel)
+
+    async def _ensure_async_listener(self, channel: str) -> Any:
+        if self._listen_connection_async is None:
+            self._listen_connection_async_cm = self._config.provide_connection()
+            self._listen_connection_async = await self._listen_connection_async_cm.__aenter__()
+            await self._listen_connection_async.set_autocommit(True)
+            await self._listen_connection_async.execute(f"LISTEN {channel}")
+        return self._listen_connection_async
+
+    def _ensure_sync_listener(self, channel: str) -> Any:
+        if self._listen_connection_sync is None:
+            self._listen_connection_sync_cm = self._config.provide_connection()
+            self._listen_connection_sync = self._listen_connection_sync_cm.__enter__()
+            self._listen_connection_sync.autocommit = True
+            self._listen_connection_sync.execute(f"LISTEN {channel}")
+        return self._listen_connection_sync
 
     async def ack_async(self, event_id: str) -> None:
         await self._queue.ack_async(event_id)
diff --git a/tests/integration/test_adapters/test_asyncpg/test_extensions/test_events.py b/tests/integration/test_adapters/test_asyncpg/test_extensions/test_events.py
@@ -15,7 +15,7 @@ async def test_asyncpg_native_event_channel(postgres_service: PostgresService) -
     """AsyncPG configs surface native LISTEN/NOTIFY events."""
 
     config = AsyncpgConfig(
-        pool_config={
+        connection_config={
             "host": postgres_service.host,
             "port": postgres_service.port,
             "user": postgres_service.user,
@@ -33,5 +33,5 @@ async def test_asyncpg_native_event_channel(postgres_service: PostgresService) -
     event_id = await channel.publish_async("notifications", {"action": "native"})
     await channel.ack_async(event_id)
 
-    if config.pool_instance:
+    if config.connection_instance:
         await config.close_pool()
diff --git a/tests/integration/test_adapters/test_psqlpy/test_extensions/test_events_listen_notify.py b/tests/integration/test_adapters/test_psqlpy/test_extensions/test_events_listen_notify.py
diff --git a/tests/integration/test_adapters/test_psycopg/test_extensions/test_events_listen_notify.py b/tests/integration/test_adapters/test_psycopg/test_extensions/test_events_listen_notify.py
