|
19 | 19 | from databento.common.enums import ReconnectPolicy |
20 | 20 | from databento.common.error import BentoError |
21 | 21 | from databento.common.publishers import Dataset |
22 | | -from databento.common.types import ClientStream, DBNRecord |
| 22 | +from databento.common.types import ClientRecordCallback |
| 23 | +from databento.common.types import ClientStream |
| 24 | +from databento.common.types import DBNRecord |
23 | 25 | from databento.common.types import ExceptionCallback |
24 | 26 | from databento.common.types import ReconnectCallback |
25 | | -from databento.common.types import RecordCallback |
26 | 27 | from databento.live.gateway import SubscriptionRequest |
27 | 28 | from databento.live.protocol import DatabentoLiveProtocol |
28 | 29 |
|
@@ -195,8 +196,8 @@ def __init__( |
195 | 196 | api_key: str, |
196 | 197 | dataset: Dataset | str, |
197 | 198 | dbn_queue: DBNQueue, |
198 | | - user_callbacks: list[tuple[RecordCallback, ExceptionCallback | None]], |
199 | 199 | user_streams: list[ClientStream], |
| 200 | + user_callbacks: list[ClientRecordCallback], |
200 | 201 | loop: asyncio.AbstractEventLoop, |
201 | 202 | metadata: SessionMetadata, |
202 | 203 | ts_out: bool = False, |
@@ -237,18 +238,16 @@ def received_record(self, record: DBNRecord) -> None: |
237 | 238 | return super().received_record(record) |
238 | 239 |
|
239 | 240 | def _dispatch_callbacks(self, record: DBNRecord) -> None: |
240 | | - for callback, exc_callback in self._user_callbacks: |
| 241 | + for callback in self._user_callbacks: |
241 | 242 | try: |
242 | | - callback(record) |
| 243 | + callback.call(record) |
243 | 244 | except Exception as exc: |
244 | 245 | logger.error( |
245 | 246 | "error dispatching %s to `%s` callback", |
246 | 247 | type(record).__name__, |
247 | | - getattr(callback, "__name__", str(callback)), |
| 248 | + callback.callback_name, |
248 | 249 | exc_info=exc, |
249 | 250 | ) |
250 | | - if exc_callback is not None: |
251 | | - exc_callback(exc) |
252 | 251 |
|
253 | 252 | def _dispatch_writes(self, record: DBNRecord) -> None: |
254 | 253 | record_bytes = bytes(record) |
@@ -315,8 +314,8 @@ def __init__( |
315 | 314 | self._loop = loop |
316 | 315 | self._metadata = SessionMetadata() |
317 | 316 | self._user_gateway: str | None = user_gateway |
318 | | - self._user_callbacks: list[tuple[RecordCallback, ExceptionCallback | None]] = [] |
319 | 317 | self._user_streams: list[ClientStream] = [] |
| 318 | + self._user_callbacks: list[ClientRecordCallback] = [] |
320 | 319 | self._user_reconnect_callbacks: list[tuple[ReconnectCallback, ExceptionCallback | None]] = ( |
321 | 320 | [] |
322 | 321 | ) |
@@ -527,19 +526,20 @@ async def wait_for_close(self) -> None: |
527 | 526 | return |
528 | 527 |
|
529 | 528 | try: |
530 | | - await self._protocol.authenticated |
531 | | - except Exception as exc: |
532 | | - raise BentoError(exc) from None |
533 | | - |
534 | | - try: |
535 | | - if self._reconnect_task is not None: |
536 | | - await self._reconnect_task |
537 | | - else: |
538 | | - await self._protocol.disconnected |
539 | | - except Exception as exc: |
540 | | - raise BentoError(exc) from None |
| 529 | + try: |
| 530 | + await self._protocol.authenticated |
| 531 | + except Exception as exc: |
| 532 | + raise BentoError(exc) from None |
541 | 533 |
|
542 | | - self._cleanup() |
| 534 | + try: |
| 535 | + if self._reconnect_task is not None: |
| 536 | + await self._reconnect_task |
| 537 | + else: |
| 538 | + await self._protocol.disconnected |
| 539 | + except Exception as exc: |
| 540 | + raise BentoError(exc) from None |
| 541 | + finally: |
| 542 | + self._cleanup() |
543 | 543 |
|
544 | 544 | def _cleanup(self) -> None: |
545 | 545 | logger.debug("cleaning up session_id=%s", self.session_id) |
|
0 commit comments