MOD: Live client to start session on iteration

Author: nmacholl

CHANGELOG.md

@@ -5,6 +5,7 @@
 #### Enhancements
 - Added `symbology_map` property to `Live` client
 - Changed `Live.add_callback` and `Live.add_stream` to accept an exception callback
+- Changed `Live.__iter__()` and `Live.__aiter__()` to send the session start message if the session is connected but not started
 - Upgraded `databento-dbn` to 0.7.1
 - Removed `Encoding`, `Compression`, `Schema`, and `SType` enums as they are now exposed by `databento-dbn`
 

databento/live/client.py

@@ -128,13 +128,15 @@ async def __anext__(self) -> DBNRecord:
     def __iter__(self) -> Live:
         logger.debug("starting iteration")
         self._dbn_queue._enabled.set()
+        if not self._session.is_started() and self.is_connected():
+            self.start()
         return self
 
     def __next__(self) -> DBNRecord:
         if self._dbn_queue is None:
             raise ValueError("iteration has not started")
 
-        while not self._session.is_disconnected() or self._dbn_queue._qsize() > 0:
+        while not self._session.is_disconnected() or self._dbn_queue.qsize() > 0:
             try:
                 record = self._dbn_queue.get(block=False)
             except queue.Empty:

databento/live/session.py

@@ -258,6 +258,7 @@ def __init__(
         port: int = DEFAULT_REMOTE_PORT,
         ts_out: bool = False,
     ) -> None:
+        self._lock = threading.RLock()
         self._loop = loop
         self._ts_out = ts_out
         self._protocol_factory = protocol_factory
@@ -277,13 +278,14 @@ def is_authenticated(self) -> bool:
         bool
 
         """
-        if self._protocol is None:
-            return False
-        try:
-            self._protocol.authenticated.result()
-        except (asyncio.InvalidStateError, asyncio.CancelledError, BentoError):
-            return False
-        return True
+        with self._lock:
+            if self._protocol is None:
+                return False
+            try:
+                self._protocol.authenticated.result()
+            except (asyncio.InvalidStateError, asyncio.CancelledError, BentoError):
+                return False
+            return True
 
     def is_disconnected(self) -> bool:
         """
@@ -294,9 +296,10 @@ def is_disconnected(self) -> bool:
         bool
 
         """
-        if self._protocol is None:
-            return True
-        return self._protocol.disconnected.done()
+        with self._lock:
+            if self._protocol is None:
+                return True
+            return self._protocol.disconnected.done()
 
     def is_reading(self) -> bool:
         """
@@ -307,9 +310,10 @@ def is_reading(self) -> bool:
         bool
 
         """
-        if self._transport is None:
-            return False
-        return self._transport.is_reading()
+        with self._lock:
+            if self._transport is None:
+                return False
+            return self._transport.is_reading()
 
     def is_started(self) -> bool:
         """
@@ -320,9 +324,10 @@ def is_started(self) -> bool:
         bool
 
         """
-        if self._protocol is None:
-            return False
-        return self._protocol.started.is_set()
+        with self._lock:
+            if self._protocol is None:
+                return False
+            return self._protocol.started.is_set()
 
     @property
     def metadata(self) -> databento_dbn.Metadata | None:
@@ -334,9 +339,10 @@ def metadata(self) -> databento_dbn.Metadata | None:
         databento_dbn.Metadata
 
         """
-        if self._protocol is None:
-            return None
-        return self._protocol._metadata.data
+        with self._lock:
+            if self._protocol is None:
+                return None
+            return self._protocol._metadata.data
 
     def abort(self) -> None:
         """
@@ -347,20 +353,22 @@ def abort(self) -> None:
         Session.close
 
         """
-        if self._transport is None:
-            return
-        self._transport.abort()
-        self._protocol = None
+        with self._lock:
+            if self._transport is None:
+                return
+            self._transport.abort()
+            self._protocol = None
 
     def close(self) -> None:
         """
         Close the current connection.
         """
-        if self._transport is None:
-            return
-        if self._transport.can_write_eof():
-            self._loop.call_soon_threadsafe(self._transport.write_eof)
-        self._loop.call_soon_threadsafe(self._transport.close)
+        with self._lock:
+            if self._transport is None:
+                return
+            if self._transport.can_write_eof():
+                self._loop.call_soon_threadsafe(self._transport.write_eof)
+            self._loop.call_soon_threadsafe(self._transport.close)
 
     def subscribe(
         self,
@@ -389,27 +397,29 @@ def subscribe(
             within 24 hours.
 
         """
-        if self._protocol is None:
-            self._connect(
-                dataset=dataset,
-                port=self._port,
-                loop=self._loop,
-            )
+        with self._lock:
+            if self._protocol is None:
+                self._connect(
+                    dataset=dataset,
+                    port=self._port,
+                    loop=self._loop,
+                )
 
-        self._protocol.subscribe(
-            schema=schema,
-            symbols=symbols,
-            stype_in=stype_in,
-            start=start,
-        )
+            self._protocol.subscribe(
+                schema=schema,
+                symbols=symbols,
+                stype_in=stype_in,
+                start=start,
+            )
 
     def resume_reading(self) -> None:
         """
         Resume reading from the connection.
         """
-        if self._transport is None:
-            return
-        self._loop.call_soon_threadsafe(self._transport.resume_reading)
+        with self._lock:
+            if self._transport is None:
+                return
+            self._loop.call_soon_threadsafe(self._transport.resume_reading)
 
     def start(self) -> None:
         """
@@ -421,9 +431,10 @@ def start(self) -> None:
             If there is no connection.
 
         """
-        if self._protocol is None:
-            raise ValueError("session is not connected")
-        self._protocol.start()
+        with self._lock:
+            if self._protocol is None:
+                raise ValueError("session is not connected")
+            self._protocol.start()
 
     async def wait_for_close(self) -> None:
         """
@@ -433,44 +444,52 @@ async def wait_for_close(self) -> None:
         if self._protocol is None:
             return
 
+        await self._protocol.authenticated
         await self._protocol.disconnected
-        disconnect_exc = self._protocol.disconnected.exception()
-
         await self._protocol.wait_for_processing()
-        self._protocol = self._transport = None
 
-        if disconnect_exc is not None:
-            raise BentoError(disconnect_exc)
+        try:
+            self._protocol.authenticated.result()
+        except Exception as exc:
+            raise BentoError(exc)
+
+        try:
+            self._protocol.disconnected.result()
+        except Exception as exc:
+            raise BentoError(exc)
+
+        self._protocol = self._transport = None
 
     def _connect(
         self,
         dataset: Dataset | str,
         port: int,
         loop: asyncio.AbstractEventLoop,
     ) -> None:
-        if self._user_gateway is None:
-            subdomain = dataset.lower().replace(".", "-")
-            gateway = f"{subdomain}.lsg.databento.com"
-            logger.debug("using default gateway for dataset %s", dataset)
-        else:
-            gateway = self._user_gateway
-            logger.debug("using user specified gateway: %s", gateway)
+        with self._lock:
+            if not self.is_disconnected():
+                return
+            if self._user_gateway is None:
+                subdomain = dataset.lower().replace(".", "-")
+                gateway = f"{subdomain}.lsg.databento.com"
+                logger.debug("using default gateway for dataset %s", dataset)
+            else:
+                gateway = self._user_gateway
+                logger.debug("using user specified gateway: %s", gateway)
 
-        asyncio.run_coroutine_threadsafe(
-            coro=self._connect_task(
-                gateway=gateway,
-                port=port,
-            ),
-            loop=loop,
-        ).result()
+            self._transport, self._protocol = asyncio.run_coroutine_threadsafe(
+                coro=self._connect_task(
+                    gateway=gateway,
+                    port=port,
+                ),
+                loop=loop,
+            ).result()
 
     async def _connect_task(
         self,
         gateway: str,
         port: int,
-    ) -> None:
-        if not self.is_disconnected():
-            return
+    ) -> tuple[asyncio.Transport, _SessionProtocol]:
         logger.info("connecting to remote gateway")
         try:
             transport, protocol = await asyncio.wait_for(
@@ -514,5 +533,4 @@ async def _connect_task(
             "authentication with remote gateway completed",
         )
 
-        self._transport = transport
-        self._protocol = protocol
+        return transport, protocol

tests/test_live_client.py

@@ -547,8 +547,6 @@ async def test_live_async_iteration(
         symbols="TEST",
     )
 
-    live_client.start()
-
     records: list[DBNRecord] = []
     async for record in live_client:
         records.append(record)
@@ -591,13 +589,12 @@ async def test_live_async_iteration_backpressure(
         pause_mock := MagicMock(),
     )
 
-    live_client.start()
-    it = live_client.__iter__()
+    live_it = iter(live_client)
     await live_client.wait_for_close()
 
-    assert pause_mock.called
+    pause_mock.assert_called()
 
-    records = list(it)
+    records: list[DBNRecord] = list(live_it)
     assert len(records) == 4
     assert live_client._dbn_queue.empty()
 
@@ -632,13 +629,12 @@ async def test_live_async_iteration_dropped(
         pause_mock := MagicMock(),
     )
 
-    live_client.start()
-    it = live_client.__iter__()
+    live_it = iter(live_client)
     await live_client.wait_for_close()
 
-    assert pause_mock.called
+    pause_mock.assert_called()
 
-    records = list(it)
+    records = list(live_it)
     assert len(records) == 1
     assert live_client._dbn_queue.empty()
 
@@ -658,8 +654,6 @@ async def test_live_async_iteration_stop(
         symbols="TEST",
     )
 
-    live_client.start()
-
     records = []
     async for record in live_client:
         records.append(record)
@@ -683,8 +677,6 @@ def test_live_sync_iteration(
         symbols="TEST",
     )
 
-    live_client.start()
-
     records = []
     for record in live_client:
         records.append(record)
@@ -918,7 +910,6 @@ async def test_live_iteration_with_reconnect(
     assert live_client.is_connected()
     assert live_client.dataset == Dataset.GLBX_MDP3
 
-    live_client.start()
     my_iter = iter(live_client)
 
     await live_client.wait_for_close()