data-validation-engine 0.7.5__tar.gz → 0.7.6__tar.gz

{data_validation_engine-0.7.5 → data_validation_engine-0.7.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: data-validation-engine
-Version: 0.7.5
+Version: 0.7.6
 Summary: `nhs data validation engine` is a framework used to validate data
 License-Expression: MIT
 License-File: LICENSE

{data_validation_engine-0.7.5 → data_validation_engine-0.7.6}/pyproject.toml

@@ -24,7 +24,7 @@ Issues = "https://github.com/NHSDigital/data-validation-engine/issues"
 Changelog = "https://github.com/NHSDigital/data-validation-engine/blob/main/CHANGELOG.md"
 
 [tool.poetry]
-version = "0.7.5"
+version = "0.7.6"
 packages = [
     { include = "dve", from = "src" },
 ]

{data_validation_engine-0.7.5 → data_validation_engine-0.7.6}/src/dve/core_engine/backends/implementations/duckdb/readers/csv.py

@@ -6,13 +6,7 @@ from typing import Any, Optional
 
 import duckdb as ddb
 import polars as pl
-from duckdb import (
-    DuckDBPyConnection,
-    DuckDBPyRelation,
-    StarExpression,
-    default_connection,
-    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
@@ -61,7 +55,7 @@ class DuckDBCSVReader(BaseFileReader):
         self.header = header
         self.delim = delim
         self.quotechar = quotechar
-        self._connection = connection if connection else default_connection
+        self._connection = connection if connection else ddb.connect(":memory:")
         self.field_check = field_check
         self.field_check_error_code = field_check_error_code
         self.field_check_error_message = field_check_error_message
@@ -181,7 +175,7 @@ class PolarsToDuckDBCSVReader(DuckDBCSVReader):
             ] + [pl.col(RECORD_INDEX_COLUMN_NAME)]
             df = df.select(pl_exprs)
 
-        return ddb.sql("SELECT * FROM df")
+        return self._connection.sql("SELECT * FROM df")
 
 
 class DuckDBCSVRepeatingHeaderReader(PolarsToDuckDBCSVReader):

{data_validation_engine-0.7.5 → data_validation_engine-0.7.6}/src/dve/core_engine/backends/implementations/duckdb/readers/json.py

@@ -4,7 +4,8 @@
 from collections.abc import Iterator
 from typing import Any, Optional
 
-from duckdb import DuckDBPyRelation, read_json
+import duckdb
+from duckdb import DuckDBPyConnection, DuckDBPyRelation
 from pydantic import BaseModel
 
 from dve.core_engine.backends.base.reader import BaseFileReader, read_function
@@ -26,9 +27,11 @@ class DuckDBJSONReader(BaseFileReader):
         self,
         *,
         json_format: Optional[str] = "array",
+        connection: Optional[DuckDBPyConnection] = None,
         **_,
     ):
         self._json_format = json_format
+        self._connection = duckdb.connect(":memory:") if not connection else connection
 
         super().__init__()
 
@@ -50,5 +53,7 @@ class DuckDBJSONReader(BaseFileReader):
         }
 
         return self.add_record_index(
-            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
+            )
         )

{data_validation_engine-0.7.5 → data_validation_engine-0.7.6}/src/dve/core_engine/backends/implementations/duckdb/readers/xml.py

@@ -3,8 +3,9 @@
 
 from typing import Optional
 
+import duckdb
 import polars as pl
-from duckdb import DuckDBPyConnection, DuckDBPyRelation, default_connection
+from duckdb import DuckDBPyConnection, DuckDBPyRelation
 from pydantic import BaseModel
 
 from dve.core_engine.backends.base.reader import read_function
@@ -24,8 +25,8 @@ from dve.core_engine.type_hints import URI
 class DuckDBXMLStreamReader(XMLStreamReader):
     """A reader for XML files"""
 
-    def __init__(self, *, ddb_connection: Optional[DuckDBPyConnection] = None, **kwargs):
-        self.ddb_connection = ddb_connection if ddb_connection else default_connection
+    def __init__(self, *, connection: Optional[DuckDBPyConnection] = None, **kwargs):
+        self._connection = connection if connection else duckdb.connect(":memory:")
         super().__init__(**kwargs)
 
     @read_function(DuckDBPyRelation)
@@ -49,4 +50,4 @@ class DuckDBXMLStreamReader(XMLStreamReader):
                 data=self.read_to_py_iterator(resource, entity_name, schema), schema=polars_schema
             )
         )
-        return self.ddb_connection.sql("select * from _lazy_frame")
+        return self._connection.sql("select * from _lazy_frame")

{data_validation_engine-0.7.5 → data_validation_engine-0.7.6}/src/dve/pipeline/duckdb_pipeline.py

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

{data_validation_engine-0.7.5 → data_validation_engine-0.7.6}/src/dve/pipeline/pipeline.py

@@ -9,7 +9,7 @@ from concurrent.futures import Executor, Future, ThreadPoolExecutor
 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 @@ PERMISSIBLE_EXCEPTIONS: tuple[type[Exception]] = (
 )
 
 
+# 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 @@ class BaseDVEPipeline:
         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 @@ class BaseDVEPipeline:
         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 @@ class BaseDVEPipeline:
         """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 @@ class BaseDVEPipeline:
 
         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(

{data_validation_engine-0.7.5 → data_validation_engine-0.7.6}/src/dve/pipeline/utils.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