refactor: add in backend kwargs for readers to allow reader args not determinable at config write time to be passed

Author: stevenhsd

poetry.lock

@@ -6,12 +6,7 @@
 
 import duckdb as ddb
 import polars as pl
-from duckdb import (
-    DuckDBPyConnection,
-    DuckDBPyRelation,
-    StarExpression,
-    read_csv,
-)
+from duckdb import DuckDBPyConnection, DuckDBPyRelation, StarExpression, read_csv
 from pydantic import BaseModel
 
 from dve.core_engine.backends.base.reader import BaseFileReader, read_function

src/dve/core_engine/backends/implementations/duckdb/readers/csv.py

@@ -53,5 +53,7 @@ def read_to_relation(  # pylint: disable=unused-argument
         }
 
         return self.add_record_index(
-            self._connection.read_json(resource, columns=ddb_schema, format=self._json_format)  # type: ignore
+            self._connection.read_json(
+                resource, columns=ddb_schema, format=self._json_format  # type: ignore
+            )
         )

src/dve/core_engine/backends/implementations/duckdb/readers/json.py

@@ -45,6 +45,7 @@ def __init__(
             submitted_files_path,
             job_run_id,
             logger,
+            {"connection": self._connection},
         )
 
     def init_reference_data_loader(

src/dve/pipeline/duckdb_pipeline.py

@@ -9,7 +9,7 @@
 from functools import lru_cache
 from itertools import starmap
 from threading import Lock
-from typing import Optional, Union
+from typing import Any, Optional, Union
 from uuid import uuid4
 
 import polars as pl
@@ -49,6 +49,7 @@
 )
 
 
+# pylint: disable=R0904
 class BaseDVEPipeline:
     """
     Base class for running a DVE Pipeline either by a given step or a full e2e process.
@@ -64,6 +65,7 @@ def __init__(
         submitted_files_path: Optional[URI],
         job_run_id: Optional[int] = None,
         logger: Optional[logging.Logger] = None,
+        backend_reader_kwargs: Optional[dict[str, Any]] = None,
     ):
         self._submitted_files_path = submitted_files_path
         self._processed_files_path = processed_files_path
@@ -76,6 +78,7 @@ def __init__(
         self._summary_lock = Lock()
         self._rec_tracking_lock = Lock()
         self._aggregates_lock = Lock()
+        self._backend_reader_kwargs = backend_reader_kwargs
 
         if self._data_contract:
             self._data_contract.logger = self._logger
@@ -107,6 +110,12 @@ def step_implementations(self) -> Optional[BaseStepImplementations[EntityType]]:
         """The step implementations to apply the business rules to a given dataset"""
         return self._step_implementations
 
+    @property
+    def backend_reader_kwargs(self) -> dict[str, Any] | None:
+        """Important required arguments for all readers related to the specific backend
+        that can't be specified at time of writing config eg. duckdb connection"""
+        return self._backend_reader_kwargs
+
     @staticmethod
     def get_entity_count(entity: EntityType) -> int:
         """Get a row count of an entity stored as parquet"""
@@ -203,7 +212,9 @@ def write_file_to_parquet(
 
         for model_name, model in models.items():
             self._logger.info(f"Transforming {model_name} to stringified parquet")
-            reader: BaseFileReader = load_reader(dataset, model_name, ext)
+            reader: BaseFileReader = load_reader(
+                dataset, model_name, ext, self.backend_reader_kwargs
+            )
             try:
                 if not entity_type:
                     reader.write_parquet(

src/dve/pipeline/pipeline.py

@@ -3,7 +3,7 @@
 
 import json
 from threading import Lock
-from typing import Optional
+from typing import Any, Optional
 
 from pydantic.main import ModelMetaclass
 from pyspark.sql import SparkSession
@@ -45,10 +45,17 @@ def load_config(
     return models, config, dataset
 
 
-def load_reader(dataset: Dataset, model_name: str, file_extension: str):
+def load_reader(
+    dataset: Dataset,
+    model_name: str,
+    file_extension: str,
+    backend_reader_kwargs: Optional[dict[str, Any]] = None,
+):
     """Loads the readers for the diven feed, model name and file extension"""
     reader_config = dataset[model_name].reader_config[f".{file_extension.lower()}"]
-    reader = _READER_REGISTRY[reader_config.reader](**reader_config.kwargs_)
+    reader = _READER_REGISTRY[reader_config.reader](
+        **reader_config.kwargs_, **backend_reader_kwargs if backend_reader_kwargs else {}
+    )
     return reader
 
 

src/dve/pipeline/utils.py

@@ -31,7 +31,7 @@
 @pytest.fixture(scope="function")
 def prep_multithreading_test():
     sub_details: dict[str, tuple[DuckDBPyConnection, str, DDBAuditingManager]] = {}
-    for idx in range(1, 10):
+    for idx in range(1, 4):
         db = f"dve_{uuid4().hex}"
         tmp_dir = tempfile.mkdtemp(prefix="ddb_foundry_testing")
         db_file = Path(tmp_dir, db + ".duckdb")