feat: added reference data loading of arrow ipc files including enhanced test coverage for reference data loaders

Author: stevenhsd

src/dve/core_engine/backends/base/reference_data.py

@@ -7,23 +7,29 @@
 from pydantic import BaseModel, Field
 from typing_extensions import Annotated, Literal
 
+import dve.parser.file_handling as fh
 from dve.core_engine.backends.base.core import get_entity_type
-from dve.core_engine.backends.exceptions import MissingRefDataEntity, RefdataLacksFileExtensionSupport
+from dve.core_engine.backends.exceptions import (
+    MissingRefDataEntity,
+    RefdataLacksFileExtensionSupport,
+)
 from dve.core_engine.backends.types import EntityType
 from dve.core_engine.type_hints import URI, EntityName
-import dve.parser.file_handling as fh
 from dve.parser.file_handling.implementations.file import LocalFilesystemImplementation
 from dve.parser.file_handling.service import _get_implementation
 
 _FILE_EXTENSION_NAME: str = "_REFDATA_FILE_EXTENSION"
 """Name of attribute added to methods where they relate
    to loading a particular reference file type."""
 
+
 def mark_refdata_file_extension(file_extension):
     """Mark a method for loading a particular file extension"""
+
     def wrapper(func: Callable):
         setattr(func, _FILE_EXTENSION_NAME, file_extension)
         return func
+
     return wrapper
 
 
@@ -52,9 +58,11 @@ class ReferenceFile(BaseModel, frozen=True):
     """The object type."""
     filename: str
     """The path to the reference data relative to the contract."""
+
     @property
     def file_extension(self) -> str:
-        return fh.get_file_suffix(self.filename)
+        """The file extension of the reference file"""
+        return fh.get_file_suffix(self.filename)  # type: ignore
 
 
 class ReferenceURI(BaseModel, frozen=True):
@@ -64,9 +72,11 @@ class ReferenceURI(BaseModel, frozen=True):
     """The object type."""
     uri: str
     """The absolute URI of the reference data (as Parquet)."""
+
     @property
     def file_extension(self) -> str:
-        return fh.get_file_suffix(self.uri)
+        """The file extension of the reference uri"""
+        return fh.get_file_suffix(self.uri)  # type: ignore
 
 
 ReferenceConfig = Union[ReferenceFile, ReferenceTable, ReferenceURI]
@@ -91,7 +101,7 @@ class BaseRefDataLoader(Generic[EntityType], Mapping[EntityName, EntityType], AB
     A mapping between refdata config types and functions to call to load these configs
     into reference data entities
     """
-    
+
     __reader_functions__: ClassVar[dict[str, Callable]] = {}
     """
     A mapping between file extensions and functions to load the file uris
@@ -108,6 +118,9 @@ class variable for the subclass.
         if cls is not BaseRefDataLoader:
             cls.__entity_type__ = get_entity_type(cls, "BaseRefDataLoader")
 
+        # ensure that dicts are specific to each subclass - redefine rather
+        # than keep the same reference
+        cls.__reader_functions__ = {}
         cls.__step_functions__ = {}
 
         for method_name in dir(cls):
@@ -117,7 +130,7 @@ class variable for the subclass.
             method = getattr(cls, method_name, None)
             if method is None or not callable(method):
                 continue
-            
+
             if ext := getattr(method, _FILE_EXTENSION_NAME, None):
                 cls.__reader_functions__[ext] = method
                 continue
@@ -136,7 +149,7 @@ def __init__(
         self,
         reference_entity_config: dict[EntityName, ReferenceConfig],
         dataset_config_uri: Optional[URI] = None,
-        **kwargs
+        **kwargs,
     ) -> None:
         self.reference_entity_config = reference_entity_config
         self.dataset_config_uri = dataset_config_uri
@@ -164,8 +177,8 @@ def load_file(self, config: ReferenceFile) -> EntityType:
         try:
             impl = self.__reader_functions__[config.file_extension]
             return impl(self, target_location)
-        except KeyError:
-            raise RefdataLacksFileExtensionSupport(file_extension=config.file_extension)
+        except KeyError as exc:
+            raise RefdataLacksFileExtensionSupport(file_extension=config.file_extension) from exc
 
     def load_uri(self, config: ReferenceURI) -> EntityType:
         "Load reference entity from an absolute URI"
@@ -176,8 +189,8 @@ def load_uri(self, config: ReferenceURI) -> EntityType:
         try:
             impl = self.__reader_functions__[config.file_extension]
             return impl(self, target_location)
-        except KeyError:
-            raise RefdataLacksFileExtensionSupport(file_extension=config.file_extension)
+        except KeyError as exc:
+            raise RefdataLacksFileExtensionSupport(file_extension=config.file_extension) from exc
 
     def load_entity(self, entity_name: EntityName, config: ReferenceConfig) -> EntityType:
         """Load a reference entity given the reference config"""

src/dve/core_engine/backends/base/rules.py

@@ -406,7 +406,11 @@ def apply_sync_filters(
                                 [
                                     CriticalProcessingError(
                                         "Issue occurred while applying filter logic",
-                                        messages=[msg.error_message for msg in temp_messages if msg.error_message],
+                                        messages=[
+                                            msg.error_message
+                                            for msg in temp_messages
+                                            if msg.error_message
+                                        ],
                                     )
                                 ],
                             )
@@ -440,7 +444,7 @@ def apply_sync_filters(
                             msg_writer.write_queue.put(temp_messages)
                         self.logger.info(
                             f"Filter {rule.reporting.code} found {len(temp_messages)} issues"
-                            )
+                        )
 
                     else:
                         temp_messages, success = self.evaluate(
@@ -466,15 +470,15 @@ def apply_sync_filters(
                             return processing_errors_uri, False
                         if temp_messages:
                             msg_writer.write_queue.put(temp_messages)
-                        
+
                         self.logger.info(
                             f"Filter {rule.reporting.code} found {len(temp_messages)} issues"
-                            )
+                        )
 
                 if filter_column_names:
                     self.logger.info(
-                    f"Filtering records from entity {entity_name} for error code {rule.reporting.code}" # pylint: disable=line-too-long
-                )
+                        f"Filtering records from entity {entity_name} for error code {rule.reporting.code}"  # pylint: disable=line-too-long
+                    )
                     success_condition = " AND ".join(
                         [f"({c_name} IS NOT NULL AND {c_name})" for c_name in filter_column_names]
                     )

src/dve/core_engine/backends/exceptions.py

@@ -108,7 +108,7 @@ def get_joiner(self):
         return "required by"
 
 
-class MissingRefDataEntity(MissingEntity):  # pylint: disable=too-many-ancestors
+class MissingRefDataEntity(MissingEntity, BackendErrorMixin):  # pylint: disable=too-many-ancestors
     """An error to be emitted when a required refdata entity is missing."""
 
     def get_message_preamble(self) -> str:
@@ -165,9 +165,10 @@ def __init__(self, *args: object, entity_type: Any) -> None:
     def get_message_preamble(self) -> EntityName:
         return f"Reader does not support reading directly to entity type {self.entity_type!r}"
 
+
 class RefdataLacksFileExtensionSupport(BackendError):
     """An error raised when trying to load a refdata file where the loader
-       lacks support for the given file type
+    lacks support for the given file type
 
     """
 
@@ -178,6 +179,7 @@ def __init__(self, *args: object, file_extension: str) -> None:
            refdata loader"""
 
     def get_message_preamble(self) -> EntityName:
+        """Message for logging purposes"""
         return f"Refdata loader does not support reading refdata from {self.file_extension} files"
 
 

src/dve/core_engine/backends/implementations/duckdb/contract.py

@@ -169,10 +169,8 @@ def apply_data_contract(
                         if msgs:
                             msg_writer.write_queue.put(msgs)
                             msg_count += len(msgs)
-                
-                self.logger.info(
-                f"Data contract found {msg_count} issues in {entity_name}"
-            )
+
+                self.logger.info(f"Data contract found {msg_count} issues in {entity_name}")
 
                 casting_statements = [
                     (

src/dve/core_engine/backends/implementations/duckdb/reference_data.py

@@ -3,23 +3,15 @@
 from typing import Optional
 
 from duckdb import DuckDBPyConnection, DuckDBPyRelation
-import pyarrow.ipc as ipc
+from pyarrow import ipc  # type: ignore
 
-import dve.parser.file_handling as fh
 from dve.core_engine.backends.base.reference_data import (
     BaseRefDataLoader,
     ReferenceConfigUnion,
-    ReferenceFile,
     ReferenceTable,
-    ReferenceURI,
-    mark_refdata_file_extension
+    mark_refdata_file_extension,
 )
 from dve.core_engine.type_hints import EntityName
-from dve.parser.file_handling.implementations.file import (
-    LocalFilesystemImplementation,
-    file_uri_to_local_path,
-)
-from dve.parser.file_handling.service import _get_implementation
 from dve.parser.type_hints import URI
 
 
@@ -44,13 +36,14 @@ def __init__(
 
     def load_table(self, config: ReferenceTable) -> DuckDBPyRelation:
         """Load reference entity from a database table"""
-        return self.connection.sql(f"select * from {config.fq_table_name}")         
-    
+        return self.connection.sql(f"select * from {config.fq_table_name}")
+
     @mark_refdata_file_extension("parquet")
     def load_parquet_file(self, uri: str) -> DuckDBPyRelation:
+        """Load a parquet file into a duckdb relation"""
         return self.connection.read_parquet(uri)
-    
+
     @mark_refdata_file_extension("arrow")
     def load_arrow_file(self, uri: str) -> DuckDBPyRelation:
-        return self.connection.from_arrow(ipc.open_file(uri).read_all())
-            
+        """Load an arrow ipc file into a duckdb relation"""
+        return self.connection.from_arrow(ipc.open_file(uri).read_all())  # type:ignore

src/dve/core_engine/backends/implementations/spark/reference_data.py

@@ -5,13 +5,10 @@
 
 from pyspark.sql import DataFrame, SparkSession
 
-import dve.parser.file_handling as fh
 from dve.core_engine.backends.base.reference_data import (
     BaseRefDataLoader,
     ReferenceConfig,
-    ReferenceFile,
     ReferenceTable,
-    ReferenceURI,
     mark_refdata_file_extension,
 )
 from dve.core_engine.type_hints import EntityName
@@ -40,5 +37,6 @@ def load_table(self, config: ReferenceTable) -> DataFrame:
         return self.spark.table(f"{config.fq_table_name}")
 
     @mark_refdata_file_extension("parquet")
-    def load_parquet_file(self, uri:str) -> DataFrame:
+    def load_parquet_file(self, uri: str) -> DataFrame:
+        """Load a parquet file into a spark dataframe"""
         return self.spark.read.parquet(uri)

tests/fixtures.py

@@ -113,12 +113,11 @@ def spark_test_database(spark: SparkSession) -> Iterator[str]:
 
 
 
-@pytest.fixture()
+@pytest.fixture(scope="function")
 def temp_ddb_conn() -> Iterator[Tuple[Path, DuckDBPyConnection]]:
     """Temp DuckDB directory for the database"""
-    db = uuid4().hex
+    db = f"dve_{uuid4().hex}"
     with tempfile.TemporaryDirectory(prefix="ddb_audit_testing") as tmp:
         db_file = Path(tmp, db + ".duckdb")
         conn = connect(database=db_file, read_only=False)
-        
         yield db_file, conn