ADD: Add subscription ID to clients

Author: threecgreen

CHANGELOG.md

@@ -1,5 +1,10 @@
 # Changelog
 
+## 0.52.0 - TBD
+
+#### Enhancements
+- Added new optional `id` field to `SubcriptionRequest` class which will be used for improved error messages
+
 ## 0.51.0 - 2025-04-08
 
 #### Enhancements

databento/live/gateway.py

@@ -132,6 +132,7 @@ class SubscriptionRequest(GatewayControl):
     symbols: str
     start: int | None = None
     snapshot: int = 0
+    id: int | None = None
 
 
 @dataclasses.dataclass

databento/live/protocol.py

@@ -284,6 +284,7 @@ def subscribe(
         stype_in: SType | str = SType.RAW_SYMBOL,
         start: str | int | None = None,
         snapshot: bool = False,
+        subscription_id: int | None = None,
     ) -> list[SubscriptionRequest]:
         """
         Send a SubscriptionRequest to the gateway. Returns a list of all
@@ -302,6 +303,8 @@ def subscribe(
             within 24 hours.
         snapshot: bool, default to 'False'
             Request subscription with snapshot. The `start` parameter must be `None`.
+        subscription_id : int, optional
+            A numerical identifier to associate with this subscription.
 
         Returns
         -------
@@ -329,6 +332,7 @@ def subscribe(
                 symbols=batch_str,
                 start=optional_datetime_to_unix_nanoseconds(start),
                 snapshot=int(snapshot),
+                id=subscription_id,
             )
             subscriptions.append(message)
 

databento/live/session.py

@@ -332,6 +332,7 @@ def __init__(
         self._transport: asyncio.Transport | None = None
         self._session_id: str | None = None
 
+        self._subscription_counter = 0
         self._subscriptions: list[SubscriptionRequest] = []
         self._reconnect_policy = ReconnectPolicy(reconnect_policy)
         self._reconnect_task: asyncio.Task[None] | None = None
@@ -499,13 +500,15 @@ def subscribe(
             if self._protocol is 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,
                 ),
             )
 
@@ -672,6 +675,7 @@ async def _reconnect(self) -> None:
                             stype_in=sub.stype_in,
                             snapshot=bool(sub.snapshot),
                             start=None,
+                            subscription_id=sub.id,
                         )
 
                     if should_restart:

tests/test_live_gateway_messages.py

@@ -299,15 +299,15 @@ def test_serialize_session_start(
     "line, expected",
     [
         pytest.param(
-            "schema=trades|" "stype_in=instrument_id|" "symbols=1,2,3\n",
-            ("trades", "instrument_id", "1,2,3", None),
+            "schema=trades|" "stype_in=instrument_id|" "symbols=1,2,3|" "id=23\n",
+            ("trades", "instrument_id", "1,2,3", None, "23"),
         ),
         pytest.param(
             "schema=trades|"
             "stype_in=instrument_id|"
             "symbols=1,2,3|"
             "start=1671717080706865759\n",
-            ("trades", "instrument_id", "1,2,3", "1671717080706865759"),
+            ("trades", "instrument_id", "1,2,3", "1671717080706865759", None),
         ),
         pytest.param(
             "schema=trades|" "stype_in=instrument_id|" "symbols=1,2,3",
@@ -336,6 +336,7 @@ def test_parse_subscription_request(
             msg.stype_in,
             msg.symbols,
             msg.start,
+            msg.id,
         ) == expected
     else:
         with pytest.raises(expected):
@@ -374,8 +375,13 @@ def test_parse_subscription_request(
                 symbols="1234,5678,90",
                 start=None,
                 snapshot=1,
+                id=5,
             ),
-            b"schema=mbo|" b"stype_in=instrument_id|" b"symbols=1234,5678,90|" b"snapshot=1\n",
+            b"schema=mbo|"
+            b"stype_in=instrument_id|"
+            b"symbols=1234,5678,90|"
+            b"snapshot=1|"
+            b"id=5\n",
         ),
     ],
 )