ADD: Add Python support for live compression

threecgreen · threecgreen · commit b75b7f206f24 · 2026-02-18T00:22:43.000Z
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -5,6 +5,10 @@
 #### Enhancements
 - Added `slow_reader_behavior` field to `AuthenticationRequest` message
 - Added `SlowReaderBehavior` enum
+- Added support for using compression in the live API:
+  - Added `compression` parameter to the `Live` client constructor
+  - Added `compression` property to the `Live` client
+  - Added `compression` field to `AuthenticationRequest`
 - Upgraded `databento-dbn` to 0.49.0:
   - Added support for decompressing Zstd in the Python `DBNDecoder` and new optional `compression` parameter
 
diff --git a/databento/live/client.py b/databento/live/client.py
@@ -14,6 +14,7 @@
 
 import databento_dbn
 import pandas as pd
+from databento_dbn import Compression
 from databento_dbn import DBNRecord
 from databento_dbn import Schema
 from databento_dbn import SType
@@ -68,6 +69,9 @@ class Live:
         The live gateway behavior when the client falls behind real time.
             - "skip": skip records to immediately catch up
             - "warn": send a slow reader warning `SystemMsg` but continue reading every record
+    compression : Compression or str, default "none"
+        The compression format for live data. Set to "zstd" for
+        Zstandard-compressed data from the gateway.
 
     """
 
@@ -88,6 +92,7 @@ def __init__(
         heartbeat_interval_s: int | None = None,
         reconnect_policy: ReconnectPolicy | str = ReconnectPolicy.NONE,
         slow_reader_behavior: SlowReaderBehavior | str | None = None,
+        compression: Compression = Compression.NONE,
     ) -> None:
         if key is None:
             key = os.environ.get("DATABENTO_API_KEY")
@@ -105,6 +110,7 @@ def __init__(
 
         self._dataset: Dataset | str = ""
         self._ts_out = ts_out
+        self._compression = compression
         self._heartbeat_interval_s = heartbeat_interval_s
 
         self._metadata: SessionMetadata = SessionMetadata()
@@ -119,6 +125,7 @@ def __init__(
             user_port=port,
             reconnect_policy=reconnect_policy,
             slow_reader_behavior=slow_reader_behavior,
+            compression=compression,
         )
 
         self._session._user_callbacks.append(ClientRecordCallback(self._map_symbol))
@@ -298,6 +305,18 @@ def ts_out(self) -> bool:
         """
         return self._ts_out
 
+    @property
+    def compression(self) -> Compression:
+        """
+        Returns the compression mode for this live client.
+
+        Returns
+        -------
+        Compression
+
+        """
+        return self._compression
+
     def add_callback(
         self,
         record_callback: RecordCallback,
diff --git a/databento/live/gateway.py b/databento/live/gateway.py
@@ -7,6 +7,7 @@
 from typing import SupportsBytes
 from typing import TypeVar
 
+from databento_dbn import Compression
 from databento_dbn import Encoding
 from databento_dbn import Schema
 from databento_dbn import SType
@@ -118,6 +119,7 @@ class AuthenticationRequest(GatewayControl):
     encoding: Encoding = Encoding.DBN
     details: str | None = None
     ts_out: str = "0"
+    compression: Compression | str = Compression.NONE
     heartbeat_interval_s: int | None = None
     slow_reader_behavior: SlowReaderBehavior | str | None = None
     client: str = USER_AGENT
diff --git a/databento/live/protocol.py b/databento/live/protocol.py
@@ -7,6 +7,7 @@
 from typing import Final
 
 import databento_dbn
+from databento_dbn import Compression
 from databento_dbn import DBNRecord
 from databento_dbn import Metadata
 from databento_dbn import Schema
@@ -61,6 +62,8 @@ class DatabentoLiveProtocol(asyncio.BufferedProtocol):
     heartbeat_interval_s: int, optional
         The interval in seconds at which the gateway will send heartbeat records if no
         other data records are sent.
+    compression : Compression, default Compression.NONE
+        The compression format for the session.
 
     See Also
     --------
@@ -75,6 +78,7 @@ def __init__(
         ts_out: bool = False,
         heartbeat_interval_s: int | None = None,
         slow_reader_behavior: SlowReaderBehavior | str | None = None,
+        compression: Compression = Compression.NONE,
     ) -> None:
         self.__api_key = api_key
         self.__transport: asyncio.Transport | None = None
@@ -84,9 +88,11 @@ def __init__(
         self._ts_out = ts_out
         self._heartbeat_interval_s = heartbeat_interval_s
         self._slow_reader_behavior: SlowReaderBehavior | str | None = slow_reader_behavior
+        self._compression = compression
 
         self._dbn_decoder = databento_dbn.DBNDecoder(
             upgrade_policy=VersionUpgradePolicy.UPGRADE_TO_V3,
+            compression=compression,
         )
         self._gateway_decoder = GatewayDecoder()
 
@@ -443,15 +449,17 @@ def _(self, message: ChallengeRequest) -> None:
             auth=response,
             dataset=self._dataset,
             ts_out=str(int(self._ts_out)),
+            compression=str(self._compression).lower(),
             heartbeat_interval_s=self._heartbeat_interval_s,
             slow_reader_behavior=self._slow_reader_behavior,
         )
         logger.debug(
-            "sending CRAM challenge response auth='%s' dataset=%s encoding=%s ts_out=%s heartbeat_interval_s=%s client='%s'",
+            "sending CRAM challenge response auth='%s' dataset=%s encoding=%s ts_out=%s compression=%s heartbeat_interval_s=%s client='%s'",
             auth_request.auth,
             auth_request.dataset,
             auth_request.encoding,
             auth_request.ts_out,
+            auth_request.compression,
             auth_request.heartbeat_interval_s,
             auth_request.client,
         )
diff --git a/databento/live/session.py b/databento/live/session.py
@@ -14,6 +14,7 @@
 
 import databento_dbn
 import pandas as pd
+from databento_dbn import Compression
 from databento_dbn import DBNRecord
 from databento_dbn import Schema
 from databento_dbn import SType
@@ -207,8 +208,16 @@ def __init__(
         ts_out: bool = False,
         heartbeat_interval_s: int | None = None,
         slow_reader_behavior: SlowReaderBehavior | str | None = None,
+        compression: Compression = Compression.NONE,
     ):
-        super().__init__(api_key, dataset, ts_out, heartbeat_interval_s, slow_reader_behavior)
+        super().__init__(
+            api_key,
+            dataset,
+            ts_out,
+            heartbeat_interval_s,
+            slow_reader_behavior,
+            compression,
+        )
 
         self._dbn_queue = dbn_queue
         self._loop = loop
@@ -304,6 +313,8 @@ class LiveSession:
         The reconnect policy for the live session.
             - "none": the client will not reconnect (default)
             - "reconnect": the client will reconnect automatically
+    compression : Compression, optional
+        The compression format for the session. Defaults to no compression.
     """
 
     def __init__(
@@ -316,6 +327,7 @@ def __init__(
         user_port: int = DEFAULT_REMOTE_PORT,
         reconnect_policy: ReconnectPolicy | str = ReconnectPolicy.NONE,
         slow_reader_behavior: SlowReaderBehavior | str | None = None,
+        compression: Compression = Compression.NONE,
     ) -> None:
         self._dbn_queue = DBNQueue()
         self._lock = threading.RLock()
@@ -333,6 +345,7 @@ def __init__(
         self._ts_out = ts_out
         self._heartbeat_interval_s = heartbeat_interval_s or 30
         self._slow_reader_behavior = slow_reader_behavior
+        self._compression = compression
 
         self._protocol: _SessionProtocol | None = None
         self._transport: asyncio.Transport | None = None
@@ -584,6 +597,7 @@ def _create_protocol(self, dataset: Dataset | str) -> _SessionProtocol:
             ts_out=self.ts_out,
             heartbeat_interval_s=self.heartbeat_interval_s,
             slow_reader_behavior=self._slow_reader_behavior,
+            compression=self._compression,
         )
 
     def _connect(
diff --git a/tests/mockliveserver/server.py b/tests/mockliveserver/server.py
@@ -13,6 +13,8 @@
 from typing import Any
 from typing import Final
 
+import zstandard
+from databento_dbn import Compression
 from databento_dbn import Schema
 
 from databento.common import cram
@@ -79,6 +81,7 @@ def __init__(
         self._data = BytesIO()
         self._state = SessionState.NEW
         self._dataset: Dataset | None = None
+        self._compression: Compression = Compression.NONE
         self._subscriptions: list[SubscriptionRequest] = []
         self._replay_tasks: set[asyncio.Task[None]] = set()
 
@@ -194,6 +197,13 @@ def _(self, message: AuthenticationRequest) -> None:
 
         self.state = SessionState.AUTHENTICATED
         self._dataset = Dataset(message.dataset)
+        # Parse compression from the auth request
+        compression_str = getattr(message, "compression", "none")
+        if compression_str == "zstd":
+            self._compression = Compression.ZSTD
+        else:
+            self._compression = Compression.NONE
+        logger.debug("client requested compression=%s", compression_str)
         self.send_gateway_message(
             self.get_authentication_response(
                 success=True,
@@ -297,15 +307,37 @@ def _replay_done_callback(self, task: asyncio.Task[Any]) -> None:
             self.hangup(reason="all replay tasks completed")
 
     async def _file_replay_task(self) -> None:
+        compressor = None
+        if self._compression == Compression.ZSTD:
+            compressor = zstandard.ZstdCompressor()
+            cctx = compressor.compressobj()
+
         for subscription in self._subscriptions:
             schema = (
                 Schema.from_str(subscription.schema)
                 if isinstance(subscription.schema, str)
                 else subscription.schema
             )
             replay = self._file_replay_table[(self.dataset, schema)]
-            logger.info("starting replay %s for %s", replay.name, self.peer)
+            logger.info(
+                "starting replay %s for %s (compression=%s)",
+                replay.name,
+                self.peer,
+                self._compression,
+            )
             for chunk in replay:
-                self.transport.write(chunk)
+                if compressor is not None:
+                    compressed = cctx.compress(chunk)
+                    if compressed:
+                        self.transport.write(compressed)
+                else:
+                    self.transport.write(chunk)
                 await asyncio.sleep(0)
+
+            # Flush remaining compressed data
+            if compressor is not None:
+                remaining = cctx.flush()
+                if remaining:
+                    self.transport.write(remaining)
+
             logger.info("replay of %s completed for %s", replay.name, self.peer)
diff --git a/tests/test_live_compression.py b/tests/test_live_compression.py
@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+import asyncio
+from unittest.mock import MagicMock
+
+import pytest
+from databento_dbn import Compression
+from databento_dbn import Schema
+
+from databento.common.publishers import Dataset
+from databento.live.protocol import DatabentoLiveProtocol
+from tests.mockliveserver.fixture import MockLiveServerInterface
+
+
+@pytest.mark.parametrize("compression", [Compression.NONE, Compression.ZSTD])
+async def test_protocol_connection_with_compression(
+    mock_live_server: MockLiveServerInterface,
+    test_live_api_key: str,
+    compression: Compression,
+) -> None:
+    """
+    Test protocol connection with different compression settings.
+    """
+    # Arrange
+    transport, protocol = await asyncio.get_event_loop().create_connection(
+        protocol_factory=lambda: DatabentoLiveProtocol(
+            api_key=test_live_api_key,
+            dataset=Dataset.GLBX_MDP3,
+            compression=compression,
+        ),
+        host=mock_live_server.host,
+        port=mock_live_server.port,
+    )
+
+    # Act, Assert
+    await asyncio.wait_for(protocol.authenticated, timeout=1)
+    transport.close()
+    await asyncio.wait_for(protocol.disconnected, timeout=1)
+
+
+@pytest.mark.parametrize("compression", [Compression.NONE, Compression.ZSTD])
+async def test_protocol_streaming_with_compression(
+    monkeypatch: pytest.MonkeyPatch,
+    mock_live_server: MockLiveServerInterface,
+    test_live_api_key: str,
+    compression: Compression,
+) -> None:
+    """
+    Test streaming records with different compression settings.
+    """
+    # Arrange
+    monkeypatch.setattr(
+        DatabentoLiveProtocol,
+        "received_metadata",
+        metadata_mock := MagicMock(),
+    )
+    monkeypatch.setattr(
+        DatabentoLiveProtocol,
+        "received_record",
+        record_mock := MagicMock(),
+    )
+
+    _, protocol = await asyncio.get_event_loop().create_connection(
+        protocol_factory=lambda: DatabentoLiveProtocol(
+            api_key=test_live_api_key,
+            dataset=Dataset.GLBX_MDP3,
+            compression=compression,
+        ),
+        host=mock_live_server.host,
+        port=mock_live_server.port,
+    )
+
+    await asyncio.wait_for(protocol.authenticated, timeout=1)
+
+    # Act
+    protocol.subscribe(
+        schema=Schema.MBO,
+        symbols="ESM4",
+    )
+    protocol.start()
+
+    # Assert
+    await asyncio.wait_for(protocol.disconnected, timeout=5)
+    metadata_mock.assert_called()
+    record_mock.assert_called()
diff --git a/tests/test_live_gateway_messages.py b/tests/test_live_gateway_messages.py
@@ -93,7 +93,7 @@ def test_parse_authentication_request(
                 dataset=Dataset.GLBX_MDP3,
                 client="unittest",
             ),
-            b"auth=abcd1234|dataset=GLBX.MDP3|encoding=dbn|ts_out=0|client=unittest\n",
+            b"auth=abcd1234|dataset=GLBX.MDP3|encoding=dbn|ts_out=0|compression=none|client=unittest\n",
         ),
         pytest.param(
             AuthenticationRequest(
@@ -103,7 +103,16 @@ def test_parse_authentication_request(
                 client="unittest",
                 heartbeat_interval_s=35,
             ),
-            b"auth=abcd1234|dataset=XNAS.ITCH|encoding=dbn|ts_out=1|heartbeat_interval_s=35|client=unittest\n",
+            b"auth=abcd1234|dataset=XNAS.ITCH|encoding=dbn|ts_out=1|compression=none|heartbeat_interval_s=35|client=unittest\n",
+        ),
+        pytest.param(
+            AuthenticationRequest(
+                auth="abc",
+                dataset=Dataset.OPRA_PILLAR,
+                compression="zstd",
+                client="Databento",
+            ),
+            b"auth=abc|dataset=OPRA.PILLAR|encoding=dbn|ts_out=0|compression=zstd|client=Databento\n",
         ),
     ],
 )
