ADD: Add getter for Live subscription requests

Author: nmacholl

CHANGELOG.md

@@ -1,5 +1,14 @@
 # Changelog
 
+## 0.67.0 - TBD
+
+#### Enhancements
+- Added a property `Live.subscription_requests` which returns a list of tuples containing every `SubscriptionRequest` for the live session
+- Changed the return value of `Live.subscribe()` to `int`, the value of the subscription ID, which can be used to index into the `Live.subscription_requests` property
+
+#### Breaking changes
+- Several log messages have been reformatted to improve clarity and reduce redundancy, especially at debug levels
+
 ## 0.66.0 - 2025-11-18
 
 #### Enhancements

databento/live/client.py

@@ -31,6 +31,7 @@
 from databento.common.types import RecordCallback
 from databento.common.validation import validate_enum
 from databento.common.validation import validate_semantic_string
+from databento.live.gateway import SubscriptionRequest
 from databento.live.session import DEFAULT_REMOTE_PORT
 from databento.live.session import LiveSession
 from databento.live.session import SessionMetadata
@@ -228,6 +229,38 @@ def session_id(self) -> str | None:
         """
         return self._session.session_id
 
+    @property
+    def subscription_requests(
+        self,
+    ) -> list[tuple[SubscriptionRequest, ...]]:
+        """
+        Return a list of tuples containing every `SubscriptionRequest` message
+        sent for the session. The list is in order of the subscriptions made
+        and can be indexed using the value returned by each call to
+        `Live.subscribe()`.
+
+        Subscriptions which contain a large
+        list of symbols are batched. Because of this, a single `subscription_id` may have
+        more than one associated `SubscriptionRequest`.
+
+        Returns
+        -------
+        list[tuple[SubscriptionRequest, ...]]
+            A list of tuples containing every subscription request.
+            Each entry in the list corresponds to a single subscription.
+
+        Raises
+        ------
+        IndexError
+            If the subscription ID is invalid.
+
+        See Also
+        --------
+        Live.subscribe()
+
+        """
+        return self._session._subscriptions
+
     @property
     def symbology_map(self) -> dict[int, str | int]:
         """
@@ -446,7 +479,7 @@ def subscribe(
         stype_in: SType | str = SType.RAW_SYMBOL,
         start: pd.Timestamp | datetime | date | str | int | None = None,
         snapshot: bool = False,
-    ) -> None:
+    ) -> int:
         """
         Add a new subscription to the session.
 
@@ -476,6 +509,11 @@ def subscribe(
             Request subscription with snapshot. The `start` parameter must be `None`.
             Only supported with `mbo` schema.
 
+        Returns
+        -------
+        int
+            The numeric identifier for this subscription request.
+
         Raises
         ------
         ValueError
@@ -494,7 +532,7 @@ def subscribe(
 
         """
         logger.info(
-            "subscribing to %s:%s %s start=%s snapshot=%s",
+            "subscribing to schema=%s stype_in=%s symbols='%s' start=%s snapshot=%s",
             schema,
             stype_in,
             symbols,
@@ -509,7 +547,7 @@ def subscribe(
         if snapshot and start is not None:
             raise ValueError("Subscription with snapshot expects start=None")
 
-        self._session.subscribe(
+        return self._session.subscribe(
             dataset=dataset,
             schema=schema,
             stype_in=stype_in,

databento/live/protocol.py

@@ -7,11 +7,9 @@
 from typing import Final
 
 import databento_dbn
-from databento_dbn import DBNError
 from databento_dbn import Metadata
 from databento_dbn import Schema
 from databento_dbn import SType
-from databento_dbn import SystemCode
 from databento_dbn import VersionUpgradePolicy
 
 from databento.common import cram
@@ -313,13 +311,14 @@ def subscribe(
         list[SubscriptionRequest]
 
         """
-        logger.info(
-            "sending subscription to %s:%s %s start=%s snapshot=%s",
+        logger.debug(
+            "sending subscription request schema=%s stype_in=%s symbols='%s' start='%s' snapshot=%s id=%s",
             schema,
             stype_in,
             symbols,
             start if start is not None else "now",
             snapshot,
+            subscription_id,
         )
 
         stype_in_valid = validate_enum(stype_in, SType, "stype_in")
@@ -341,6 +340,12 @@ def subscribe(
             )
             subscriptions.append(message)
 
+        if len(subscriptions) > 1:
+            logger.debug(
+                "batched subscription into %d requests id=%s",
+                len(subscriptions),
+                subscription_id,
+            )
         self.transport.writelines(map(bytes, subscriptions))
         return subscriptions
 
@@ -374,27 +379,20 @@ def _process_dbn(self, data: bytes) -> None:
                     continue
                 if isinstance(record, databento_dbn.ErrorMsg):
                     logger.error(
-                        "gateway error: %s",
+                        "gateway error code=%s err='%s'",
+                        record.code,
                         record.err,
                     )
                     self._error_msgs.append(record.err)
                 elif isinstance(record, databento_dbn.SystemMsg):
                     if record.is_heartbeat():
                         logger.debug("gateway heartbeat")
                     else:
-                        try:
-                            msg_code = record.code
-                        except DBNError:
-                            msg_code = None
-                        if msg_code == SystemCode.SLOW_READER_WARNING:
-                            logger.warning(
-                                record.msg,
-                            )
-                        else:
-                            logger.debug(
-                                "gateway message: %s",
-                                record.msg,
-                            )
+                        logger.info(
+                            "system message code=%s msg='%s'",
+                            record.code,
+                            record.msg,
+                        )
                 self.received_record(record)
 
     def _process_gateway(self, data: bytes) -> None:
@@ -423,34 +421,44 @@ def _handle_gateway_message(self, message: GatewayControl) -> None:
 
     @_handle_gateway_message.register(Greeting)
     def _(self, message: Greeting) -> None:
-        logger.debug("greeting received by remote gateway v%s", message.lsg_version)
+        logger.debug(
+            "greeting received by remote gateway version='%s'",
+            message.lsg_version,
+        )
 
     @_handle_gateway_message.register(ChallengeRequest)
     def _(self, message: ChallengeRequest) -> None:
-        logger.debug("received CRAM challenge: %s", message.cram)
+        logger.debug("received CRAM challenge cram='%s'", message.cram)
         response = cram.get_challenge_response(message.cram, self.__api_key)
         auth_request = AuthenticationRequest(
             auth=response,
             dataset=self._dataset,
             ts_out=str(int(self._ts_out)),
             heartbeat_interval_s=self._heartbeat_interval_s,
         )
-        logger.debug("sending CRAM challenge response: %s", str(auth_request).strip())
+        logger.debug(
+            "sending CRAM challenge response auth='%s' dataset=%s encoding=%s ts_out=%s heartbeat_interval_s=%s client='%s'",
+            auth_request.auth,
+            auth_request.dataset,
+            auth_request.encoding,
+            auth_request.ts_out,
+            auth_request.heartbeat_interval_s,
+            auth_request.client,
+        )
         self.transport.write(bytes(auth_request))
 
     @_handle_gateway_message.register(AuthenticationResponse)
     def _(self, message: AuthenticationResponse) -> None:
         if message.success == "0":
-            logger.error("CRAM authentication failed: %s", message.error)
+            logger.error("CRAM authentication error: %s", message.error)
             self.authenticated.set_exception(
-                BentoError(f"User authentication failed: {message.error}"),
+                BentoError(message.error),
             )
             self.transport.close()
         else:
             session_id = message.session_id
 
             logger.debug(
-                "CRAM authenticated session id assigned `%s`",
-                session_id,
+                "CRAM authentication successful",
             )
             self.authenticated.set_result(session_id)

databento/live/session.py

@@ -2,6 +2,7 @@
 
 import asyncio
 import dataclasses
+import itertools
 import logging
 import queue
 import struct
@@ -329,8 +330,7 @@ def __init__(
         self._transport: asyncio.Transport | None = None
         self._session_id: str | None = None
 
-        self._subscription_counter = 0
-        self._subscriptions: list[SubscriptionRequest] = []
+        self._subscriptions: list[tuple[SubscriptionRequest, ...]] = []
         self._reconnect_policy = ReconnectPolicy(reconnect_policy)
         self._reconnect_task: asyncio.Task[None] | None = None
 
@@ -463,7 +463,7 @@ def subscribe(
         stype_in: SType | str = SType.RAW_SYMBOL,
         start: str | int | None = None,
         snapshot: bool = False,
-    ) -> None:
+    ) -> int:
         """
         Send a subscription request on the current connection. This will create
         a new connection if there is no active connection to the gateway.
@@ -498,17 +498,20 @@ def subscribe(
                 self._session_id = None
                 self._connect(dataset=dataset)
 
-            self._subscription_counter += 1
-            self._subscriptions.extend(
-                self._protocol.subscribe(
-                    schema=schema,
-                    symbols=symbols,
-                    stype_in=stype_in,
-                    start=start,
-                    snapshot=snapshot,
-                    subscription_id=self._subscription_counter,
+            subscription_id = len(self._subscriptions)
+            self._subscriptions.append(
+                tuple(
+                    self._protocol.subscribe(
+                        schema=schema,
+                        symbols=symbols,
+                        stype_in=stype_in,
+                        start=start,
+                        snapshot=snapshot,
+                        subscription_id=subscription_id,
+                    ),
                 ),
             )
+        return subscription_id
 
     def terminate(self) -> None:
         with self._lock:
@@ -542,7 +545,7 @@ async def wait_for_close(self) -> None:
             self._cleanup()
 
     def _cleanup(self) -> None:
-        logger.debug("cleaning up session_id=%s", self.session_id)
+        logger.debug("cleaning up session_id='%s'", self.session_id)
         self._user_callbacks.clear()
         for stream in self._user_streams:
             if not stream.is_closed:
@@ -596,7 +599,7 @@ async def _connect_task(
             logger.debug("using default gateway for dataset %s", dataset)
         else:
             gateway = self._user_gateway
-            logger.debug("using user specified gateway: %s", gateway)
+            logger.debug("user gateway override gateway='%s'", gateway)
 
         logger.info("connecting to remote gateway")
         try:
@@ -638,7 +641,7 @@ async def _connect_task(
 
         self._session_id = session_id
         logger.info(
-            "authenticated session %s",
+            "authenticated session_id='%s'",
             self.session_id,
         )
 
@@ -669,7 +672,7 @@ async def _reconnect(self) -> None:
                         dataset=self._protocol._dataset,
                     )
 
-                    for sub in self._subscriptions:
+                    for sub in itertools.chain(*self._subscriptions):
                         self._protocol.subscribe(
                             schema=sub.schema,
                             symbols=sub.symbols,

tests/test_live_client.py

@@ -163,7 +163,7 @@ def test_live_connection_cram_failure(
         )
 
     # Ensure this was an authentication error
-    exc.match(r"User authentication failed:")
+    exc.match(r"Authentication failed.")
 
 
 @pytest.mark.parametrize(
@@ -554,6 +554,8 @@ async def test_live_subscribe(
     assert message.symbols == symbols
     assert message.start == start
     assert message.snapshot == "0"
+    assert len(live_client.subscription_requests[0]) == 1
+    assert live_client.subscription_requests[0][0].id == int(message.id)
 
 
 @pytest.mark.parametrize(
@@ -645,16 +647,19 @@ async def test_live_subscribe_large_symbol_list(
         symbols=large_symbol_list,
     )
 
+    batched = []
     reconstructed: list[str] = []
     for i in range(8):
         message = await mock_live_server.wait_for_message_of_type(
             message_type=gateway.SubscriptionRequest,
         )
         assert int(message.is_last) == int(i == 7)
         reconstructed.extend(message.symbols.split(","))
+        batched.append(message)
 
     # Assert
     assert reconstructed == large_symbol_list
+    assert len(live_client.subscription_requests[0]) == len(batched)
 
 
 async def test_live_subscribe_from_callback(
@@ -1663,7 +1668,7 @@ async def test_live_connection_reuse_cram_failure(
         )
 
     # Ensure this was an authentication error
-    exc.match(r"User authentication failed:")
+    exc.match(r"Authentication failed.")
 
     async with mock_live_server.api_key_context(test_api_key):
         live_client.subscribe(