ADD: Add configurable heartbeat support to clients

Author: threecgreen

CHANGELOG.md

@@ -3,6 +3,8 @@
 ## 0.35.0 - TBD
 
 #### Enhancements
+- Added optional `heartbeat_interval_s` parameter to `Live` client for configuring the
+  interval at which the gateway will send heartbeat records
 - Upgraded `databento-dbn` to 0.18.0
 
 #### Breaking changes

databento/live/client.py

@@ -50,6 +50,9 @@ class Live:
     ts_out: bool, default False
         If set, DBN records will be timestamped when they are sent by the
         gateway.
+    heartbeat_interval_s: int, optional
+        The interval in seconds at which the gateway will send heartbeat records if no
+        other data records are sent.
 
     """
 
@@ -66,6 +69,7 @@ def __init__(
         gateway: str | None = None,
         port: int = DEFAULT_REMOTE_PORT,
         ts_out: bool = False,
+        heartbeat_interval_s: int | None = None,
     ) -> None:
         if key is None:
             key = os.environ.get("DATABENTO_API_KEY")
@@ -83,6 +87,7 @@ def __init__(
 
         self._dataset: Dataset | str = ""
         self._ts_out = ts_out
+        self._heartbeat_interval_s = heartbeat_interval_s
 
         self._dbn_queue: DBNQueue = DBNQueue()
         self._metadata: SessionMetadata = SessionMetadata()
@@ -102,6 +107,7 @@ def factory() -> _SessionProtocol:
                 loop=self._loop,
                 metadata=self._metadata,
                 ts_out=self._ts_out,
+                heartbeat_interval_s=self._heartbeat_interval_s,
             )
 
         self._session: Session = Session(

databento/live/gateway.py

@@ -108,6 +108,7 @@ class AuthenticationRequest(GatewayControl):
     encoding: Encoding = Encoding.DBN
     details: str | None = None
     ts_out: str = "0"
+    heartbeat_interval_s: int | None = None
     client: str = USER_AGENT
 
 

databento/live/protocol.py

@@ -55,6 +55,9 @@ class DatabentoLiveProtocol(asyncio.BufferedProtocol):
         The dataset for authentication.
     ts_out : bool, default False
         Flag for requesting `ts_out` to be appending to all records in the session.
+    heartbeat_interval_s: int, optional
+        The interval in seconds at which the gateway will send heartbeat records if no
+        other data records are sent.
 
     See Also
     --------
@@ -67,13 +70,15 @@ def __init__(
         api_key: str,
         dataset: Dataset | str,
         ts_out: bool = False,
+        heartbeat_interval_s: int | None = None,
     ) -> None:
         self.__api_key = api_key
         self.__transport: asyncio.Transport | None = None
         self.__buffer: bytearray = bytearray(RECV_BUFFER_SIZE)
 
         self._dataset = validate_semantic_string(dataset, "dataset")
         self._ts_out = ts_out
+        self._heartbeat_interval_s = heartbeat_interval_s
 
         self._dbn_decoder = databento_dbn.DBNDecoder(
             upgrade_policy=VersionUpgradePolicy.UPGRADE,
@@ -392,6 +397,7 @@ def _(self, message: ChallengeRequest) -> None:
             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())
         self.transport.write(bytes(auth_request))

databento/live/session.py

@@ -186,6 +186,7 @@ def __init__(
         loop: asyncio.AbstractEventLoop,
         metadata: SessionMetadata,
         ts_out: bool = False,
+        heartbeat_interval_s: int | None = None,
     ):
         super().__init__(api_key, dataset, ts_out)
 

tests/test_live_gateway_messages.py

@@ -30,16 +30,17 @@
     [
         pytest.param(
             "auth=abcd1234|dataset=GLBX.MDP3|encoding=json\n",
-            ("abcd1234", "GLBX.MDP3", "json", None, "0"),
+            ("abcd1234", "GLBX.MDP3", "json", None, "0", None),
         ),
         pytest.param(
-            "auth=abcd1234|dataset=GLBX.MDP3|ts_out=1\n",
+            "auth=abcd1234|dataset=GLBX.MDP3|heartbeat_interval_s=10|ts_out=1\n",
             (
                 "abcd1234",
                 "GLBX.MDP3",
                 str(Encoding.DBN),
                 None,
                 "1",
+                "10",
             ),
         ),
         pytest.param(
@@ -50,10 +51,11 @@
                 str(Encoding.DBN),
                 None,
                 "0",
+                None,
             ),
         ),
         pytest.param(
-            "auth=abcd1234|dataset=GLBX.MDP3|ts_out=1|encoding=csv|extra=key\n",
+            "auth=abcd1234|dataset=GLBX.MDP3|ts_out=1|encoding=csv|heartbeat_interval_s=5|extra=key\n",
             ValueError,
         ),
     ],
@@ -74,6 +76,7 @@ def test_parse_authentication_request(
             msg.encoding,
             msg.details,
             msg.ts_out,
+            msg.heartbeat_interval_s,
         ) == expected
     else:
         with pytest.raises(expected):
@@ -97,8 +100,9 @@ def test_parse_authentication_request(
                 dataset=Dataset.XNAS_ITCH,
                 ts_out="1",
                 client="unittest",
+                heartbeat_interval_s=35,
             ),
-            b"auth=abcd1234|dataset=XNAS.ITCH|encoding=dbn|ts_out=1|client=unittest\n",
+            b"auth=abcd1234|dataset=XNAS.ITCH|encoding=dbn|ts_out=1|heartbeat_interval_s=35|client=unittest\n",
         ),
     ],
 )