Refine databento venue dataset mapping

Author: faysou

nautilus_trader/adapters/databento/config.py

@@ -13,10 +13,8 @@
 #  limitations under the License.
 # -------------------------------------------------------------------------------------------------
 
-from nautilus_trader.adapters.databento.types import Dataset
 from nautilus_trader.config import LiveDataClientConfig
 from nautilus_trader.model.identifiers import InstrumentId
-from nautilus_trader.model.identifiers import Venue
 
 
 class DatabentoDataClientConfig(LiveDataClientConfig, frozen=True):
@@ -46,7 +44,7 @@ class DatabentoDataClientConfig(LiveDataClientConfig, frozen=True):
         e.g. {'GLBX.MDP3', ['ES.FUT', 'ES.OPT']} (for all E-mini S&P 500 futures and options products).
     instrument_ids : list[InstrumentId], optional
         The instrument IDs to request instrument definitions for on start.
-    venue_dataset_map: dict[Venue, Dataset], optional
+    venue_dataset_map: dict[str, str], optional
         A dictionary to override the default dataset used for a venue.
 
     """
@@ -59,4 +57,4 @@ class DatabentoDataClientConfig(LiveDataClientConfig, frozen=True):
     mbo_subscriptions_delay: float | None = 3.0  # Need to have received all definitions
     instrument_ids: list[InstrumentId] | None = None
     parent_symbols: dict[str, set[str]] | None = None
-    venue_dataset_map: dict[Venue, Dataset] | None = None
+    venue_dataset_map: dict[str, str] | None = None

nautilus_trader/adapters/databento/data.py

@@ -162,7 +162,7 @@ def __init__(
 
         # Cache instrument index
         for instrument_id in config.instrument_ids or []:
-            dataset = self._get_dataset_for_venue(instrument_id.venue)
+            dataset = self._loader.get_dataset_for_venue(instrument_id.venue)
             self._instrument_ids[dataset].add(instrument_id)
 
         # MBO/L3 subscription buffering
@@ -175,11 +175,6 @@ def __init__(
         self._update_dataset_ranges_interval_secs: int = 60 * 60  # Once per hour (hard-coded)
         self._update_dataset_ranges_task: asyncio.Task | None = None
 
-    def _get_dataset_for_venue(self, venue: Venue) -> Dataset:
-        return (
-            self._venue_dataset_map and self._venue_dataset_map.get(venue)
-        ) or self._loader.get_dataset_for_venue(venue)
-
     async def _connect(self) -> None:
         if not self._instrument_ids:
             return  # Nothing else to do yet
@@ -331,7 +326,7 @@ def _send_all_instruments_to_data_engine(self) -> None:
 
     async def _ensure_subscribed_for_instrument(self, instrument_id: InstrumentId) -> None:
         try:
-            dataset: Dataset = self._get_dataset_for_venue(instrument_id.venue)
+            dataset: Dataset = self._loader.get_dataset_for_venue(instrument_id.venue)
             subscribed_instruments = self._instrument_ids[dataset]
 
             if instrument_id in subscribed_instruments:
@@ -420,7 +415,7 @@ async def _subscribe_imbalance(self, data_type: DataType) -> None:
         try:
             # TODO: Create `DatabentoTimeSeriesParams`
             instrument_id: InstrumentId = data_type.metadata["instrument_id"]
-            dataset: Dataset = self._get_dataset_for_venue(instrument_id.venue)
+            dataset: Dataset = self._loader.get_dataset_for_venue(instrument_id.venue)
             live_client = self._get_live_client(dataset)
             live_client.subscribe(
                 schema=DatabentoSchema.IMBALANCE.value,
@@ -434,7 +429,7 @@ async def _subscribe_statistics(self, data_type: DataType) -> None:
         try:
             # TODO: Create `DatabentoTimeSeriesParams`
             instrument_id: InstrumentId = data_type.metadata["instrument_id"]
-            dataset: Dataset = self._get_dataset_for_venue(instrument_id.venue)
+            dataset: Dataset = self._loader.get_dataset_for_venue(instrument_id.venue)
             live_client = self._get_live_client(dataset)
             live_client.subscribe(
                 schema=DatabentoSchema.STATISTICS.value,
@@ -450,7 +445,7 @@ async def _subscribe_instruments(self, command: SubscribeInstruments) -> None:
 
     async def _subscribe_instrument(self, command: SubscribeInstrument) -> None:
         try:
-            dataset: Dataset = self._get_dataset_for_venue(command.instrument_id.venue)
+            dataset: Dataset = self._loader.get_dataset_for_venue(command.instrument_id.venue)
             start: int | None = command.params.get("start")
 
             live_client = self._get_live_client(dataset)
@@ -510,7 +505,7 @@ async def _subscribe_order_book_deltas(self, command: SubscribeOrderBook) -> Non
                 )
                 return
 
-            dataset: Dataset = self._get_dataset_for_venue(command.instrument_id.venue)
+            dataset: Dataset = self._loader.get_dataset_for_venue(command.instrument_id.venue)
 
             if self._is_buffering_mbo_subscriptions:
                 self._log.debug(
@@ -545,7 +540,7 @@ async def _subscribe_order_book_deltas_batch(
             if not instrument_ids:
                 return  # No subscribing instrument IDs were loaded in the cache
 
-            dataset: Dataset = self._get_dataset_for_venue(instrument_ids[0].venue)
+            dataset: Dataset = self._loader.get_dataset_for_venue(instrument_ids[0].venue)
             live_client = self._get_live_client_mbo(dataset)
 
             if dataset == "GLBX.MDP3":
@@ -600,7 +595,7 @@ async def _subscribe_order_book_snapshots(self, command: SubscribeOrderBook) ->
                     )
                     return
 
-            dataset: Dataset = self._get_dataset_for_venue(command.instrument_id.venue)
+            dataset: Dataset = self._loader.get_dataset_for_venue(command.instrument_id.venue)
             live_client = self._get_live_client(dataset)
             live_client.subscribe(
                 schema=schema,
@@ -625,7 +620,7 @@ async def _subscribe_quote_ticks(self, command: SubscribeQuoteTicks) -> None:
 
             start: int | None = command.params.get("start")
 
-            dataset: Dataset = self._get_dataset_for_venue(command.instrument_id.venue)
+            dataset: Dataset = self._loader.get_dataset_for_venue(command.instrument_id.venue)
             live_client = self._get_live_client(dataset)
             live_client.subscribe(
                 schema=schema,
@@ -649,7 +644,7 @@ async def _subscribe_trade_ticks(self, command: SubscribeTradeTicks) -> None:
 
             start: int | None = command.params.get("start")
 
-            dataset: Dataset = self._get_dataset_for_venue(command.instrument_id.venue)
+            dataset: Dataset = self._loader.get_dataset_for_venue(command.instrument_id.venue)
             live_client = self._get_live_client(dataset)
             live_client.subscribe(
                 schema=DatabentoSchema.TRADES.value,
@@ -662,7 +657,7 @@ async def _subscribe_trade_ticks(self, command: SubscribeTradeTicks) -> None:
 
     async def _subscribe_bars(self, command: SubscribeBars) -> None:
         try:
-            dataset: Dataset = self._get_dataset_for_venue(
+            dataset: Dataset = self._loader.get_dataset_for_venue(
                 command.bar_type.instrument_id.venue,
             )
 
@@ -686,7 +681,7 @@ async def _subscribe_bars(self, command: SubscribeBars) -> None:
 
     async def _subscribe_instrument_status(self, command: SubscribeInstrumentStatus) -> None:
         try:
-            dataset: Dataset = self._get_dataset_for_venue(command.instrument_id.venue)
+            dataset: Dataset = self._loader.get_dataset_for_venue(command.instrument_id.venue)
 
             live_client = self._get_live_client(dataset)
             live_client.subscribe(
@@ -770,7 +765,7 @@ async def _request_instrument_status(
         start = data_type.metadata.get("start")
         end = data_type.metadata.get("end")
 
-        dataset: Dataset = self._get_dataset_for_venue(instrument_id.venue)
+        dataset: Dataset = self._loader.get_dataset_for_venue(instrument_id.venue)
         _, available_end = await self._get_dataset_range(dataset)
 
         start = start or available_end - pd.Timedelta(days=2)
@@ -803,7 +798,7 @@ async def _request_imbalance(self, data_type: DataType, correlation_id: UUID4) -
         start = data_type.metadata.get("start")
         end = data_type.metadata.get("end")
 
-        dataset: Dataset = self._get_dataset_for_venue(instrument_id.venue)
+        dataset: Dataset = self._loader.get_dataset_for_venue(instrument_id.venue)
         _, available_end = await self._get_dataset_range(dataset)
 
         start = start or available_end - pd.Timedelta(days=2)
@@ -834,7 +829,7 @@ async def _request_statistics(self, data_type: DataType, correlation_id: UUID4)
         start = data_type.metadata.get("start")
         end = data_type.metadata.get("end")
 
-        dataset: Dataset = self._get_dataset_for_venue(instrument_id.venue)
+        dataset: Dataset = self._loader.get_dataset_for_venue(instrument_id.venue)
         _, available_end = await self._get_dataset_range(dataset)
 
         start = start or available_end - pd.Timedelta(days=2)
@@ -861,7 +856,7 @@ async def _request_statistics(self, data_type: DataType, correlation_id: UUID4)
         )
 
     async def _request_instrument(self, request: RequestInstrument) -> None:
-        dataset: Dataset = self._get_dataset_for_venue(request.instrument_id.venue)
+        dataset: Dataset = self._loader.get_dataset_for_venue(request.instrument_id.venue)
         _, available_end = await self._get_dataset_range(dataset)
 
         start = request.start or available_end - pd.Timedelta(days=2)
@@ -892,7 +887,7 @@ async def _request_instrument(self, request: RequestInstrument) -> None:
         self._handle_instrument(instruments[0], request.id, request.params)
 
     async def _request_instruments(self, request: RequestInstruments) -> None:
-        dataset: Dataset = self._get_dataset_for_venue(request.venue)
+        dataset: Dataset = self._loader.get_dataset_for_venue(request.venue)
         _, available_end = await self._get_dataset_range(dataset)
 
         start = request.start or available_end - pd.Timedelta(days=2)
@@ -925,7 +920,7 @@ async def _request_instruments(self, request: RequestInstruments) -> None:
         self._handle_instruments(instruments, request.venue, request.id, request.params)
 
     async def _request_quote_ticks(self, request: RequestQuoteTicks) -> None:
-        dataset: Dataset = self._get_dataset_for_venue(request.instrument_id.venue)
+        dataset: Dataset = self._loader.get_dataset_for_venue(request.instrument_id.venue)
         _, available_end = await self._get_dataset_range(dataset)
 
         start = request.start or available_end - pd.Timedelta(days=1)
@@ -963,7 +958,7 @@ async def _request_quote_ticks(self, request: RequestQuoteTicks) -> None:
         self._handle_quote_ticks(request.instrument_id, quotes, request.id, request.params)
 
     async def _request_trade_ticks(self, request: RequestTradeTicks) -> None:
-        dataset: Dataset = self._get_dataset_for_venue(request.instrument_id.venue)
+        dataset: Dataset = self._loader.get_dataset_for_venue(request.instrument_id.venue)
         _, available_end = await self._get_dataset_range(dataset)
 
         start = request.start or available_end - pd.Timedelta(days=1)
@@ -991,7 +986,7 @@ async def _request_trade_ticks(self, request: RequestTradeTicks) -> None:
         self._handle_trade_ticks(request.instrument_id, trades, request.id, request.params)
 
     async def _request_bars(self, request: RequestBars) -> None:
-        dataset: Dataset = self._get_dataset_for_venue(request.bar_type.instrument_id.venue)
+        dataset: Dataset = self._loader.get_dataset_for_venue(request.bar_type.instrument_id.venue)
         _, available_end = await self._get_dataset_range(dataset)
 
         start = request.start or available_end - pd.Timedelta(days=1)

nautilus_trader/adapters/databento/factories.py

@@ -146,7 +146,7 @@ def create(  # type: ignore
             gateway=config.http_gateway,
         )
 
-        loader = DatabentoDataLoader()
+        loader = DatabentoDataLoader(config.venue_dataset_map)
         provider = get_cached_databento_instrument_provider(
             http_client=http_client,
             clock=clock,

nautilus_trader/adapters/databento/loaders.py

@@ -55,10 +55,14 @@ class DatabentoDataLoader:
 
     """
 
-    def __init__(self) -> None:
+    def __init__(
+        self,
+        venue_dataset_map: dict[str, str] | None = None,
+    ) -> None:
         self._pyo3_loader: nautilus_pyo3.DatabentoDataLoader = nautilus_pyo3.DatabentoDataLoader(
             str(PUBLISHERS_FILEPATH),
         )
+        self._venue_dataset_map = venue_dataset_map
 
     def load_publishers(self, path: PathLike[str] | str) -> None:
         """
@@ -102,6 +106,9 @@ def get_dataset_for_venue(self, venue: Venue) -> str:
             If `venue` is not in the map of publishers.
 
         """
+        if self._venue_dataset_map and (dataset := self._venue_dataset_map.get(venue.value)):
+            return dataset
+
         dataset = self._pyo3_loader.get_dataset_for_venue(nautilus_pyo3.Venue(venue.value))
 
         if dataset is None: