ADD: Add exception handler to live client

Author: nmacholl

CHANGELOG.md

@@ -4,6 +4,8 @@
 
 ### Enhancements
 - Added `symbology_map` property to `Live` client
+- Changed `Live.add_callback` and `Live.add_stream` to accept an exception callback
+- Changed `Live.add_callback` and `Live.add_stream` `func` parameter to `record_callback`
 
 ## 0.14.1 - 2023-06-16
 

databento/live/__init__.py

@@ -1,11 +1,12 @@
-from typing import Union
+from typing import Callable, Union
 
 import databento_dbn
 
 
 AUTH_TIMEOUT_SECONDS: float = 2
 CONNECT_TIMEOUT_SECONDS: float = 5
 
+
 DBNRecord = Union[
     databento_dbn.MBOMsg,
     databento_dbn.MBP1Msg,
@@ -19,3 +20,6 @@
     databento_dbn.SystemMsg,
     databento_dbn.ErrorMsg,
 ]
+
+RecordCallback = Callable[[DBNRecord], None]
+ExceptionCallback = Callable[[Exception], None]

databento/live/client.py

@@ -8,7 +8,7 @@
 from collections.abc import Iterable
 from concurrent import futures
 from numbers import Number
-from typing import IO, Callable
+from typing import IO
 
 import databento_dbn
 
@@ -23,6 +23,8 @@
 from databento.common.validation import validate_enum
 from databento.common.validation import validate_semantic_string
 from databento.live import DBNRecord
+from databento.live import ExceptionCallback
+from databento.live import RecordCallback
 from databento.live.session import DEFAULT_REMOTE_PORT
 from databento.live.session import DBNQueue
 from databento.live.session import Session
@@ -31,9 +33,6 @@
 
 
 logger = logging.getLogger(__name__)
-
-UserCallback = Callable[[DBNRecord], None]
-
 DEFAULT_QUEUE_SIZE = 2048
 
 
@@ -89,8 +88,8 @@ def __init__(
         self._dbn_queue: DBNQueue = DBNQueue(maxsize=DEFAULT_QUEUE_SIZE)
         self._metadata: SessionMetadata = SessionMetadata()
         self._symbology_map: dict[int, str | int] = {}
-        self._user_callbacks: list[UserCallback] = [self._map_symbol]
-        self._user_streams: list[IO[bytes]] = []
+        self._user_callbacks: dict[RecordCallback, ExceptionCallback | None] = {}
+        self._user_streams: dict[IO[bytes], ExceptionCallback | None] = {}
 
         def factory() -> _SessionProtocol:
             return _SessionProtocol(
@@ -269,15 +268,19 @@ def ts_out(self) -> bool:
 
     def add_callback(
         self,
-        func: UserCallback,
+        record_callback: RecordCallback,
+        exception_callback: ExceptionCallback | None = None,
     ) -> None:
         """
         Add a callback for handling records.
 
         Parameters
         ----------
-        func : Callable[[DBNRecord], None]
+        record_callback : Callable[[DBNRecord], None]
             A callback to register for handling live records as they arrive.
+        exception_callback : Callable[[Exception], None], optional
+            An error handling callback to process exceptions that are raised
+            in `record_callback`.
 
         Raises
         ------
@@ -289,20 +292,31 @@ def add_callback(
         Live.add_stream
 
         """
-        if not callable(func):
-            raise ValueError(f"{func} is not callable")
-        callback_name = getattr(func, "__name__", str(func))
+        if not callable(record_callback):
+            raise ValueError(f"{record_callback} is not callable")
+
+        if exception_callback is not None and not callable(exception_callback):
+            raise ValueError(f"{exception_callback} is not callable")
+
+        callback_name = getattr(record_callback, "__name__", str(record_callback))
         logger.info("adding user callback %s", callback_name)
-        self._user_callbacks.append(func)
+        self._user_callbacks[record_callback] = exception_callback
 
-    def add_stream(self, stream: IO[bytes]) -> None:
+    def add_stream(
+        self,
+        stream: IO[bytes],
+        exception_callback: ExceptionCallback | None = None,
+    ) -> None:
         """
         Add an IO stream to write records to.
 
         Parameters
         ----------
         stream : IO[bytes]
             The IO stream to write to when handling live records as they arrive.
+        exception_callback : Callable[[Exception], None], optional
+            An error handling callback to process exceptions that are raised
+            when writing to the stream.
 
         Raises
         ------
@@ -320,11 +334,14 @@ def add_stream(self, stream: IO[bytes]) -> None:
         if not hasattr(stream, "writable") or not stream.writable():
             raise ValueError(f"{type(stream).__name__} is not a writable stream")
 
+        if exception_callback is not None and not callable(exception_callback):
+            raise ValueError(f"{exception_callback} is not callable")
+
         stream_name = getattr(stream, "name", str(stream))
         logger.info("adding user stream %s", stream_name)
         if self.metadata is not None:
             stream.write(bytes(self.metadata))
-        self._user_streams.append(stream)
+        self._user_streams[stream] = exception_callback
 
     def start(
         self,

databento/live/protocol.py

@@ -186,7 +186,7 @@ def eof_received(self) -> bool | None:
         asycnio.BufferedProtocol.eof_received
 
         """
-        logger.info("received EOF file from remote")
+        logger.info("received EOF from remote")
         return super().eof_received()
 
     def get_buffer(self, sizehint: int) -> bytearray:

databento/live/session.py

@@ -20,13 +20,14 @@
 from databento.live import AUTH_TIMEOUT_SECONDS
 from databento.live import CONNECT_TIMEOUT_SECONDS
 from databento.live import DBNRecord
+from databento.live import ExceptionCallback
+from databento.live import RecordCallback
 from databento.live.protocol import DatabentoLiveProtocol
 
 
 logger = logging.getLogger(__name__)
 
 
-UserCallback = Callable[[DBNRecord], None]
 DEFAULT_REMOTE_PORT = 13000
 
 
@@ -122,8 +123,8 @@ def __init__(
         api_key: str,
         dataset: Dataset | str,
         dbn_queue: DBNQueue,
-        user_callbacks: list[UserCallback],
-        user_streams: list[IO[bytes]],
+        user_callbacks: dict[RecordCallback, ExceptionCallback | None],
+        user_streams: dict[IO[bytes], ExceptionCallback | None],
         loop: asyncio.AbstractEventLoop,
         metadata: SessionMetadata,
         ts_out: bool = False,
@@ -140,22 +141,28 @@ def __init__(
     def received_metadata(self, metadata: databento_dbn.Metadata) -> None:
         if not self._metadata:
             self._metadata.data = metadata
-            for stream in self._user_streams:
-                task = self._loop.create_task(self._stream_task(stream, metadata))
+            for stream, exc_callback in self._user_streams.items():
+                task = self._loop.create_task(
+                    self._stream_task(stream, metadata, exc_callback),
+                )
                 task.add_done_callback(self._tasks.remove)
                 self._tasks.add(task)
         else:
             self._metadata.check(metadata)
         return super().received_metadata(metadata)
 
     def received_record(self, record: DBNRecord) -> None:
-        for callback in self._user_callbacks:
-            task = self._loop.create_task(self._callback_task(callback, record))
+        for callback, exc_callback in self._user_callbacks.items():
+            task = self._loop.create_task(
+                self._callback_task(callback, record, exc_callback),
+            )
             task.add_done_callback(self._tasks.remove)
             self._tasks.add(task)
 
-        for stream in self._user_streams:
-            task = self._loop.create_task(self._stream_task(stream, record))
+        for stream, exc_callback in self._user_streams.items():
+            task = self._loop.create_task(
+                self._stream_task(stream, record, exc_callback),
+            )
             task.add_done_callback(self._tasks.remove)
             self._tasks.add(task)
 
@@ -180,26 +187,29 @@ def received_record(self, record: DBNRecord) -> None:
 
     async def _callback_task(
         self,
-        func: UserCallback,
+        record_callback: RecordCallback,
         record: DBNRecord,
+        exception_callback: ExceptionCallback | None,
     ) -> None:
         try:
-            func(record)
+            record_callback(record)
         except Exception as exc:
             logger.error(
                 "error dispatching %s to `%s` callback",
                 type(record).__name__,
-                func.__name__,
+                record_callback.__name__,
                 exc_info=exc,
             )
-            raise
+            if exception_callback is not None:
+                self._loop.call_soon_threadsafe(exception_callback, exc)
 
     async def _stream_task(
         self,
         stream: IO[bytes],
         record: databento_dbn.Metadata | DBNRecord,
+        exc_callback: ExceptionCallback | None,
     ) -> None:
-        has_ts_out = self._metadata and self._metadata.data.ts_out
+        has_ts_out = self._metadata.data and self._metadata.data.ts_out
         try:
             stream.write(bytes(record))
             if not isinstance(record, databento_dbn.Metadata) and has_ts_out:
@@ -212,7 +222,8 @@ async def _stream_task(
                 stream_name,
                 exc_info=exc,
             )
-            raise
+            if exc_callback is not None:
+                self._loop.call_soon_threadsafe(exc_callback, exc)
 
     async def wait_for_processing(self) -> None:
         while self._tasks:

tests/test_live_client.py

@@ -497,7 +497,8 @@ def callback(_: object) -> None:
 
     live_client.add_callback(callback)
     assert callback in live_client._user_callbacks
-    assert live_client._user_streams == []
+    assert live_client._user_callbacks[callback] is None
+    assert live_client._user_streams == {}
 
 
 def test_live_add_stream(
@@ -509,7 +510,9 @@ def test_live_add_stream(
     stream = BytesIO()
 
     live_client.add_stream(stream)
-    assert live_client._user_streams == [stream]
+    assert stream in live_client._user_streams
+    assert live_client._user_streams[stream] is None
+    assert live_client._user_callbacks == {}
 
 
 def test_live_add_stream_invalid(
@@ -1064,3 +1067,55 @@ def test_live_connection_reconnect_cram_failure(
         dataset=Dataset.GLBX_MDP3,
         schema=Schema.MBO,
     )
+
+async def test_live_callback_exception_handler(
+    live_client: client.Live,
+) -> None:
+    """
+    Test exceptions that occur during callbacks are dispatched to the assigned
+    exception handler.
+    """
+    live_client.subscribe(
+        dataset=Dataset.GLBX_MDP3,
+        schema=Schema.MBO,
+        stype_in=SType.RAW_SYMBOL,
+        symbols="TEST",
+    )
+
+    exceptions: list[Exception] = []
+
+    def callback(_: DBNRecord) -> None:
+        raise RuntimeError("this is a test")
+
+    live_client.add_callback(callback, exceptions.append)
+
+    live_client.start()
+
+    await live_client.wait_for_close()
+    assert len(exceptions) == 4
+
+
+async def test_live_stream_exception_handler(
+    live_client: client.Live,
+) -> None:
+    """
+    Test exceptions that occur during stream writes are dispatched to the
+    assigned exception handler.
+    """
+    live_client.subscribe(
+        dataset=Dataset.GLBX_MDP3,
+        schema=Schema.MBO,
+        stype_in=SType.RAW_SYMBOL,
+        symbols="TEST",
+    )
+
+    exceptions: list[Exception] = []
+
+    stream = BytesIO()
+    live_client.add_stream(stream, exceptions.append)
+    stream.close()
+
+    live_client.start()
+
+    await live_client.wait_for_close()
+    assert len(exceptions) == 5  # extra write from metadata