feat(events): improve connection handling and error management across event backends

Author: cofin

sqlspec/adapters/asyncpg/events/backend.py

@@ -1,3 +1,4 @@
+# pyright: reportPrivateUsage=false
 """Native and hybrid PostgreSQL backends for EventChannel."""
 
 import asyncio
@@ -68,7 +69,8 @@ async def _ensure_listener(self, channel: str) -> Any:
         if self._listen_connection is None:
             self._listen_connection_cm = self._config.provide_connection()
             self._listen_connection = await self._listen_connection_cm.__aenter__()
-            await self._listen_connection.execute(f"LISTEN {channel}")
+            if self._listen_connection is not None:
+                await self._listen_connection.execute(f"LISTEN {channel}")
         return self._listen_connection
 
     async def ack_async(self, event_id: str) -> None:
@@ -240,7 +242,8 @@ async def _ensure_listener(self, channel: str) -> Any:
                 self._notify_mode = "add_listener"
             elif getattr(self._listen_connection, "notifies", None) is not None:
                 self._notify_mode = "notifies"
-                await self._listen_connection.execute(f"LISTEN {channel}")
+                if self._listen_connection is not None:
+                    await self._listen_connection.execute(f"LISTEN {channel}")
             else:
                 msg = "PostgreSQL connection does not support LISTEN/NOTIFY callbacks"
                 raise EventChannelError(msg)

sqlspec/adapters/oracledb/events/backend.py

@@ -51,7 +51,8 @@ def __init__(self, config: "DatabaseConfigProtocol[Any, Any, Any]", settings: di
     def publish_sync(self, channel: str, payload: dict[str, Any], metadata: dict[str, Any] | None = None) -> str:
         event_id = uuid.uuid4().hex
         envelope = self._build_envelope(channel, event_id, payload, metadata)
-        with self._config.provide_session() as driver:
+        session_cm = self._config.provide_session()
+        with session_cm as driver:  # type: ignore[union-attr]
             connection = getattr(driver, "connection", None)
             if connection is None:
                 msg = "Oracle driver does not expose a raw connection"
@@ -67,7 +68,8 @@ async def publish_async(self, *_: Any, **__: Any) -> str:  # pragma: no cover -
         raise ImproperConfigurationError(msg)
 
     def dequeue_sync(self, channel: str, poll_interval: float) -> EventMessage | None:
-        with self._config.provide_session() as driver:
+        session_cm = self._config.provide_session()
+        with session_cm as driver:  # type: ignore[union-attr]
             connection = getattr(driver, "connection", None)
             if connection is None:
                 msg = "Oracle driver does not expose a raw connection"
@@ -76,10 +78,13 @@ def dequeue_sync(self, channel: str, poll_interval: float) -> EventMessage | Non
             options = oracledb.AQDequeueOptions()  # type: ignore[attr-defined]
             options.wait = max(int(self._wait_seconds), 0)
             if self._visibility:
-                options.visibility = getattr(oracledb, self._visibility, None) or oracledb.AQMSG_VISIBLE
+                default_visibility = getattr(oracledb, "AQMSG_VISIBLE", None)
+                options.visibility = getattr(oracledb, self._visibility, None) or default_visibility
             try:
                 message = queue.deqone(options=options)
-            except oracledb.DatabaseError as error:  # pragma: no cover - driver surfaced runtime
+            except Exception as error:  # pragma: no cover - driver surfaced runtime
+                if oracledb is None or not isinstance(error, oracledb.DatabaseError):
+                    raise
                 logger.warning("Oracle AQ dequeue failed: %s", error)
                 driver.rollback()
                 return None

sqlspec/adapters/psycopg/events/backend.py

@@ -1,3 +1,4 @@
+# pyright: reportPrivateUsage=false
 """Psycopg LISTEN/NOTIFY and hybrid event backends."""
 
 import contextlib
@@ -55,7 +56,8 @@ async def publish_async(
     ) -> str:
         event_id = uuid.uuid4().hex
         envelope = self._encode_payload(event_id, payload, metadata)
-        async with self._config.provide_session() as driver:
+        session_cm = self._config.provide_session()
+        async with session_cm as driver:  # type: ignore[union-attr]
             await driver.execute(SQL("SELECT pg_notify(:channel, :payload)", {"channel": channel, "payload": envelope}))
             await driver.commit()
         self._runtime.increment_metric("events.publish.native")
@@ -64,7 +66,8 @@ async def publish_async(
     def publish_sync(self, channel: str, payload: "dict[str, Any]", metadata: "dict[str, Any] | None" = None) -> str:
         event_id = uuid.uuid4().hex
         envelope = self._encode_payload(event_id, payload, metadata)
-        with self._config.provide_session() as driver:
+        session_cm = self._config.provide_session()
+        with session_cm as driver:  # type: ignore[union-attr]
             driver.execute(SQL("SELECT pg_notify(:channel, :payload)", {"channel": channel, "payload": envelope}))
             driver.commit()
         self._runtime.increment_metric("events.publish.native")
@@ -122,9 +125,10 @@ async def _ensure_async_listener(self, channel: str) -> Any:
         if self._listen_connection_async is None:
             validated_channel = normalize_event_channel_name(channel)
             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 {validated_channel}")
+            self._listen_connection_async = await self._listen_connection_async_cm.__aenter__()  # type: ignore[union-attr]
+            if self._listen_connection_async is not None:
+                await self._listen_connection_async.set_autocommit(True)
+                await self._listen_connection_async.execute(f"LISTEN {validated_channel}")
         return self._listen_connection_async
 
     def _ensure_sync_listener(self, channel: str) -> Any:
@@ -139,9 +143,10 @@ def _ensure_sync_listener(self, channel: str) -> Any:
         if self._listen_connection_sync is None:
             validated_channel = normalize_event_channel_name(channel)
             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 {validated_channel}")
+            self._listen_connection_sync = self._listen_connection_sync_cm.__enter__()  # type: ignore[union-attr]
+            if self._listen_connection_sync is not None:
+                self._listen_connection_sync.autocommit = True
+                self._listen_connection_sync.execute(f"LISTEN {validated_channel}")
         return self._listen_connection_sync
 
     async def shutdown_async(self) -> None:
@@ -280,9 +285,10 @@ async def _ensure_async_listener(self, channel: str) -> Any:
         if self._listen_connection_async is None:
             validated_channel = normalize_event_channel_name(channel)
             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 {validated_channel}")
+            self._listen_connection_async = await self._listen_connection_async_cm.__aenter__()  # type: ignore[union-attr]
+            if self._listen_connection_async is not None:
+                await self._listen_connection_async.set_autocommit(True)
+                await self._listen_connection_async.execute(f"LISTEN {validated_channel}")
         return self._listen_connection_async
 
     def _ensure_sync_listener(self, channel: str) -> Any:
@@ -297,9 +303,10 @@ def _ensure_sync_listener(self, channel: str) -> Any:
         if self._listen_connection_sync is None:
             validated_channel = normalize_event_channel_name(channel)
             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 {validated_channel}")
+            self._listen_connection_sync = self._listen_connection_sync_cm.__enter__()  # type: ignore[union-attr]
+            if self._listen_connection_sync is not None:
+                self._listen_connection_sync.autocommit = True
+                self._listen_connection_sync.execute(f"LISTEN {validated_channel}")
         return self._listen_connection_sync
 
     async def ack_async(self, event_id: str) -> None:
@@ -364,7 +371,8 @@ async def _publish_durable_async(
         """
         now = datetime.now(timezone.utc)
         queue = self._get_table_queue()
-        async with self._config.provide_session() as driver:
+        session_cm = self._config.provide_session()
+        async with session_cm as driver:  # type: ignore[union-attr]
             await driver.execute(
                 SQL(
                     queue._upsert_sql,
@@ -403,7 +411,8 @@ def _publish_durable_sync(
         """
         now = datetime.now(timezone.utc)
         queue = self._get_table_queue()
-        with self._config.provide_session() as driver:
+        session_cm = self._config.provide_session()
+        with session_cm as driver:  # type: ignore[union-attr]
             driver.execute(
                 SQL(
                     queue._upsert_sql,

sqlspec/base.py

@@ -165,7 +165,11 @@ def event_channel(self, config: "type[SyncConfigT | AsyncConfigT] | SyncConfigT
         from sqlspec.extensions.events import EventChannel
 
         if isinstance(config, type):
-            config_obj = self._configs.get(config)
+            config_obj: DatabaseConfigProtocol[Any, Any, Any] | None = None
+            for registered_config in self._configs.values():
+                if isinstance(registered_config, config):
+                    config_obj = registered_config
+                    break
             if config_obj is None:
                 msg = f"Configuration {self._get_config_name(config)} is not registered"
                 raise ImproperConfigurationError(msg)

sqlspec/extensions/events/_queue.py

@@ -270,8 +270,9 @@ def _cleanup_sync(self, reference: "datetime") -> None:
 
     async def _fetch_candidate_async(self, channel: str) -> "dict[str, Any] | None":
         current_time = self._utcnow()
-        async with self._config.provide_session() as driver:
-            return await driver.select_one_or_none(
+        session_cm = self._config.provide_session()
+        async with session_cm as driver:  # type: ignore[union-attr]
+            result: dict[str, Any] | None = await driver.select_one_or_none(
                 SQL(
                     self._select_sql,
                     {
@@ -284,11 +285,13 @@ async def _fetch_candidate_async(self, channel: str) -> "dict[str, Any] | None":
                     statement_config=self._statement_config,
                 )
             )
+            return result
 
     def _fetch_candidate_sync(self, channel: str) -> "dict[str, Any] | None":
         current_time = self._utcnow()
-        with self._config.provide_session() as driver:
-            return driver.select_one_or_none(
+        session_cm = self._config.provide_session()
+        with session_cm as driver:  # type: ignore[union-attr]
+            result: dict[str, Any] | None = driver.select_one_or_none(
                 SQL(
                     self._select_sql,
                     {
@@ -301,20 +304,23 @@ def _fetch_candidate_sync(self, channel: str) -> "dict[str, Any] | None":
                     statement_config=self._statement_config,
                 )
             )
+            return result
 
     async def _execute_async(self, sql: str, parameters: "dict[str, Any]") -> int:
-        async with self._config.provide_session() as driver:
+        session_cm = self._config.provide_session()
+        async with session_cm as driver:  # type: ignore[union-attr]
             result = await driver.execute(SQL(sql, parameters, statement_config=self._statement_config))
             if result.rows_affected:
                 await driver.commit()
-        return result.rows_affected
+            return int(result.rows_affected)
 
     def _execute_sync(self, sql: str, parameters: "dict[str, Any]") -> int:
-        with self._config.provide_session() as driver:
+        session_cm = self._config.provide_session()
+        with session_cm as driver:  # type: ignore[union-attr]
             result = driver.execute(SQL(sql, parameters, statement_config=self._statement_config))
             if result.rows_affected:
                 driver.commit()
-        return result.rows_affected
+            return int(result.rows_affected)
 
     @staticmethod
     def _coerce_datetime(value: Any) -> "datetime":

sqlspec/extensions/events/channel.py

@@ -149,7 +149,8 @@ def publish_sync(self, channel: str, payload: "dict[str, Any]", metadata: "dict[
         channel = self._normalize_channel_name(channel)
         if self._is_async:
             if self._should_bridge_sync_calls():
-                return self._bridge_sync_call(self.publish_async, channel, payload, metadata)
+                result: str = self._bridge_sync_call(self.publish_async, channel, payload, metadata)
+                return result
             msg = "publish_sync requires a sync configuration"
             raise ImproperConfigurationError(msg)
         if not getattr(self._backend, "supports_sync", False):
@@ -456,11 +457,14 @@ def _dequeue_for_sync(self, channel: str, poll_interval: float) -> "EventMessage
             except Exception as error:
                 self._end_event_span(span, error=error)
                 raise
-            result = "empty" if event is None else "delivered"
-            self._end_event_span(span, result=result)
+            span_result = "empty" if event is None else "delivered"
+            self._end_event_span(span, result=span_result)
             return event
         if self._should_bridge_sync_calls():
-            return self._bridge_sync_call(self._dequeue_async_with_span, channel, poll_interval)
+            bridged_result: EventMessage | None = self._bridge_sync_call(
+                self._dequeue_async_with_span, channel, poll_interval
+            )
+            return bridged_result
         return None
 
     def _ack_for_sync(self, event_id: str) -> None:
@@ -489,7 +493,7 @@ def _bridge_sync_call(self, func: Any, *args: Any, **kwargs: Any) -> Any:
 
     def _ensure_portal(self) -> Any:
         if self._portal is None:
-            self._portal = get_global_portal()
+            self._portal = get_global_portal()  # type: ignore[assignment]
         return self._portal
 
     @staticmethod

tests/integration/test_adapters/test_asyncmy/test_extensions/test_events.py

@@ -1,3 +1,4 @@
+# pyright: reportPrivateUsage=false, reportAttributeAccessIssue=false
 """AsyncMy integration tests for the EventChannel queue backend."""
 
 import pytest

tests/integration/test_adapters/test_asyncpg/test_extensions/test_events.py

@@ -1,3 +1,4 @@
+# pyright: reportPrivateUsage=false
 """AsyncPG integration tests for the EventChannel native backend."""
 
 import pytest

tests/integration/test_adapters/test_asyncpg/test_extensions/test_events_listen_notify.py

@@ -1,3 +1,4 @@
+# pyright: reportPrivateUsage=false
 """PostgreSQL LISTEN/NOTIFY event channel tests for asyncpg adapter."""
 
 import asyncio
@@ -58,8 +59,8 @@ async def test_asyncpg_listen_notify_message_delivery(postgres_service: "Any") -
 
     received: list[Any] = []
 
-    async def _handler(msg: Any) -> None:
-        received.append(msg)
+    async def _handler(message: Any) -> None:
+        received.append(message)
 
     listener = channel.listen_async("notifications", _handler, poll_interval=0.2)
     event_id = await channel.publish_async("notifications", {"action": "async_delivery"})
@@ -105,8 +106,8 @@ async def test_asyncpg_hybrid_listen_notify_durable(postgres_service: "Any", tmp
 
     received: list[Any] = []
 
-    async def _handler(msg: Any) -> None:
-        received.append(msg)
+    async def _handler(message: Any) -> None:
+        received.append(message)
 
     listener = channel.listen_async("alerts", _handler, poll_interval=0.2)
     event_id = await channel.publish_async("alerts", {"action": "hybrid_async"})
@@ -143,8 +144,8 @@ async def test_asyncpg_listen_notify_metadata(postgres_service: "Any") -> None:
 
     received: list[Any] = []
 
-    async def _handler(msg: Any) -> None:
-        received.append(msg)
+    async def _handler(message: Any) -> None:
+        received.append(message)
 
     listener = channel.listen_async("meta_channel", _handler, poll_interval=0.2)
     event_id = await channel.publish_async(

tests/integration/test_adapters/test_oracledb/test_extensions/test_events.py

@@ -1,3 +1,4 @@
+# pyright: reportPossiblyUnboundVariable=false, reportAttributeAccessIssue=false
 """OracleDB integration tests for the EventChannel queue backend."""
 
 import asyncio