refactor: configured refdata loader to be instantiated when required without need for class vars

Author: stevenhsd

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

@@ -41,14 +41,12 @@ def __init__(  # pylint: disable=unused-argument
         self,
         contract: BaseDataContract[EntityType],
         steps: BaseStepImplementations[EntityType],
-        reference_data_loader_type: Optional[type[BaseRefDataLoader[EntityType]]],
         logger: Optional[logging.Logger] = None,
         **kwargs: Any,
     ) -> None:
         for component_name, component in (
             ("Contract", contract),
             ("Step implementation", steps),
-            ("Reference data loader", reference_data_loader_type),
         ):
             component_entity_type = getattr(component, "__entity_type__", None)
             if component_entity_type != self.__entity_type__:
@@ -61,42 +59,15 @@ def __init__(  # pylint: disable=unused-argument
         """The data contract implementation used by the backend."""
         self.step_implementations = steps
         """The step implementations used by the backend."""
-        self.reference_data_loader_type = reference_data_loader_type
-        """
-        The loader type to use for the reference data. If `None`, do not
-        load any reference data and error if it is provided.
-
-        """
         self.logger = logger or get_logger(type(self).__name__)
         """The `logging.Logger instance for the backend."""
 
     def load_reference_data(
         self,
         reference_entity_config: dict[EntityName, ReferenceConfigUnion],
         submission_info: Optional[SubmissionInfo],
-    ) -> Mapping[EntityName, EntityType]:
-        """Load the reference data as specified in the reference entity config."""
-        sub_info_entity: Optional[EntityType] = None
-        if submission_info:
-            sub_info_entity = self.convert_submission_info(submission_info)
-
-        if self.reference_data_loader_type is None:
-            if reference_entity_config:
-                raise ValueError(
-                    "Reference data has been specified but no reference data loader is "
-                    + "configured for this backend"
-                )
-
-            reference_data_dict = {}
-            if sub_info_entity is not None:
-                reference_data_dict["dve_submission_info"] = sub_info_entity
-            return reference_data_dict
-
-        reference_data_loader = self.reference_data_loader_type(reference_entity_config)
-        if sub_info_entity is not None:
-            reference_data_loader.entity_cache["dve_submission_info"] = sub_info_entity
-
-        return reference_data_loader
+    ) -> BaseRefDataLoader[EntityType]:
+        raise NotImplementedError()
 
     @abstractmethod
     def convert_submission_info(self, submission_info: SubmissionInfo) -> EntityType:

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

@@ -11,6 +11,7 @@
 from dve.core_engine.backends.base.core import get_entity_type
 from dve.core_engine.backends.exceptions import (
     MissingRefDataEntity,
+    NoRefDataConfigSupplied,
     RefdataLacksFileExtensionSupport,
 )
 from dve.core_engine.backends.types import EntityType
@@ -147,11 +148,11 @@ class variable for the subclass.
     # pylint: disable=unused-argument
     def __init__(
         self,
-        reference_entity_config: dict[EntityName, ReferenceConfig],
-        dataset_config_uri: Optional[URI] = None,
+        reference_data_config: dict[EntityName, ReferenceConfig],
+        dataset_config_uri: URI,
         **kwargs,
     ) -> None:
-        self.reference_entity_config = reference_entity_config
+        self.reference_entity_config = reference_data_config
         self.dataset_config_uri = dataset_config_uri
         """
         Configuration options for the reference data. This is likely to vary
@@ -207,6 +208,8 @@ def __getitem__(self, key: EntityName) -> EntityType:
             try:
                 config = self.reference_entity_config[key]
                 return self.load_entity(entity_name=key, config=config)
+            except TypeError:
+                raise NoRefDataConfigSupplied()
             except Exception as err:
                 raise MissingRefDataEntity(entity_name=key) from err
 

src/dve/core_engine/backends/exceptions.py

@@ -118,6 +118,19 @@ def get_message_preamble(self) -> str:
         """
         return f"Missing reference data entity {self.entity_name!r}"
 
+class NoRefDataConfigSupplied(BackendError):
+    """An error raised when trying to load a refdata entity when no refdata
+       config has been supplied.
+
+    """
+
+    def __init__(self, *args: object) -> None:
+        super().__init__(*args)
+
+    def get_message_preamble(self) -> EntityName:
+        """Message for logging purposes"""
+        return f"Refdata loader not supplied with refdata config - unable to load refdata entities"
+
 
 class ConstraintError(ValueError, BackendErrorMixin):
     """Raised when a given constraint is violated."""

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

@@ -7,6 +7,7 @@
 
 from dve.core_engine.backends.base.reference_data import (
     BaseRefDataLoader,
+    ReferenceConfig,
     ReferenceConfigUnion,
     ReferenceTable,
     mark_refdata_file_extension,
@@ -17,19 +18,20 @@
 
 # pylint: disable=too-few-public-methods
 class DuckDBRefDataLoader(BaseRefDataLoader[DuckDBPyRelation]):
-    """A reference data loader using already existing DuckDB tables."""
-
-    connection: DuckDBPyConnection
-    """The DuckDB connection for the backend."""
-    dataset_config_uri: Optional[URI] = None
-    """The location of the dischema file"""
+    """A reference data loader using already existing DuckDB tables.
+       reference_entity_config and dataset_config_uri (if config uses relative paths)
+       should be supplied using setter methods for the dataset being processed before running."""
 
     def __init__(
         self,
-        reference_entity_config: dict[EntityName, ReferenceConfigUnion],
-        **kwargs,
+        connection: DuckDBPyConnection,
+        reference_data_config: dict[EntityName, ReferenceConfig],
+        dataset_config_uri: URI,
+        **kwargs
     ) -> None:
-        super().__init__(reference_entity_config, self.dataset_config_uri, **kwargs)
+        super().__init__(reference_data_config, dataset_config_uri,**kwargs)
+        
+        self.connection = connection
 
         if not self.connection:
             raise AttributeError("DuckDBConnection must be specified")

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

@@ -6,16 +6,18 @@
 from pyspark.sql import DataFrame, SparkSession
 
 from dve.core_engine.backends.base.backend import BaseBackend
+from dve.core_engine.backends.base.reference_data import ReferenceConfigUnion
 from dve.core_engine.backends.implementations.spark.contract import SparkDataContract
 from dve.core_engine.backends.implementations.spark.reference_data import SparkRefDataLoader
 from dve.core_engine.backends.implementations.spark.rules import SparkStepImplementations
 from dve.core_engine.backends.implementations.spark.spark_helpers import get_type_from_annotation
 from dve.core_engine.backends.implementations.spark.types import SparkEntities
+from dve.core_engine.backends.types import EntityType
 from dve.core_engine.constants import RECORD_INDEX_COLUMN_NAME
 from dve.core_engine.loggers import get_child_logger, get_logger
 from dve.core_engine.models import SubmissionInfo
-from dve.core_engine.type_hints import URI, EntityParquetLocations
-from dve.parser.file_handling import get_resource_exists, joinuri
+from dve.core_engine.type_hints import URI, EntityName, EntityParquetLocations
+from dve.parser.file_handling import get_resource_exists, joinuri, get_parent
 
 
 class SparkBackend(BaseBackend[DataFrame]):
@@ -26,7 +28,6 @@ def __init__(
         dataset_config_uri: Optional[URI] = None,
         contract: Optional[SparkDataContract] = None,
         steps: Optional[SparkStepImplementations] = None,
-        reference_data_loader: Optional[type[SparkRefDataLoader]] = None,
         logger: Optional[logging.Logger] = None,
         spark_session: Optional[SparkSession] = None,
         **kwargs: Any,
@@ -36,6 +37,8 @@ def __init__(
 
         self.spark_session = spark_session or SparkSession.builder.getOrCreate()
         """The Spark session for the backend."""
+        self.dataset_config_uri = dataset_config_uri
+        """The uri of the dischema specifying the DVE config"""
 
         if contract is None:
             contract = SparkDataContract(
@@ -46,11 +49,23 @@ def __init__(
             steps = SparkStepImplementations.register_udfs(
                 logger=get_child_logger("SparkStepImplementations", logger)
             )
-        if reference_data_loader is None:
-            reference_data_loader = SparkRefDataLoader
-        reference_data_loader.spark = self.spark_session
-        reference_data_loader.dataset_config_uri = dataset_config_uri
-        super().__init__(contract, steps, reference_data_loader, logger, **kwargs)
+        super().__init__(contract, steps, logger, **kwargs)
+    
+    def load_reference_data(self,
+        reference_entity_config: dict[EntityName, ReferenceConfigUnion],
+        submission_info: Optional[SubmissionInfo],):
+        """Load the reference data as specified in the reference entity config."""
+        sub_info_entity: Optional[EntityType] = None
+        if submission_info:
+            sub_info_entity = self.convert_submission_info(submission_info)
+
+        reference_data_loader = SparkRefDataLoader(spark=self.spark_session,
+                                                   reference_data_config=reference_entity_config,
+                                                   dataset_config_uri=self.dataset_config_uri)
+        if sub_info_entity is not None:
+            reference_data_loader.entity_cache["dve_submission_info"] = sub_info_entity
+
+        return reference_data_loader
 
     def write_entities_to_parquet(
         self, entities: SparkEntities, cache_prefix: URI

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

@@ -17,19 +17,19 @@
 
 # pylint: disable=too-few-public-methods
 class SparkRefDataLoader(BaseRefDataLoader[DataFrame]):
-    """A reference data loader using already existing Apache Spark Tables."""
-
-    spark: SparkSession
-    """The Spark session for the backend."""
-    dataset_config_uri: Optional[URI] = None
-    """The location of the dischema file defining business rules"""
+    """A reference data loader using already existing Apache Spark Tables.
+       reference_entity_config and dataset_config_uri (if config uses relative paths)
+       should be supplied using setter methods for the dataset being processed before running."""
 
     def __init__(
         self,
-        reference_entity_config: dict[EntityName, ReferenceConfig],
+        spark: SparkSession,
+        reference_data_config: dict[EntityName, ReferenceConfig],
+        dataset_config_uri: URI,
         **kwargs,
     ) -> None:
-        super().__init__(reference_entity_config, self.dataset_config_uri, **kwargs)
+        super().__init__(reference_data_config, dataset_config_uri, **kwargs)
+        self.spark = spark
         if not self.spark:
             raise AttributeError("Spark session must be provided")
 

src/dve/pipeline/duckdb_pipeline.py

@@ -5,14 +5,16 @@
 
 from duckdb import DuckDBPyConnection, DuckDBPyRelation
 
-from dve.core_engine.backends.base.reference_data import BaseRefDataLoader
+from dve.core_engine.backends.base.reference_data import BaseRefDataLoader, ReferenceConfig
 from dve.core_engine.backends.implementations.duckdb.auditing import DDBAuditingManager
 from dve.core_engine.backends.implementations.duckdb.contract import DuckDBDataContract
 from dve.core_engine.backends.implementations.duckdb.duckdb_helpers import duckdb_get_entity_count
+from dve.core_engine.backends.implementations.duckdb.reference_data import DuckDBRefDataLoader
 from dve.core_engine.backends.implementations.duckdb.rules import DuckDBStepImplementations
 from dve.core_engine.models import SubmissionInfo
 from dve.core_engine.type_hints import URI
 from dve.pipeline.pipeline import BaseDVEPipeline
+import dve.parser.file_handling as fh
 
 
 # pylint: disable=abstract-method
@@ -30,7 +32,6 @@ def __init__(
         connection: DuckDBPyConnection,
         rules_path: Optional[URI],
         submitted_files_path: Optional[URI],
-        reference_data_loader: Optional[type[BaseRefDataLoader]] = None,
         job_run_id: Optional[int] = None,
         logger: Optional[logging.Logger] = None,
     ):
@@ -42,11 +43,17 @@ def __init__(
             DuckDBStepImplementations.register_udfs(connection=self._connection),
             rules_path,
             submitted_files_path,
-            reference_data_loader,
             job_run_id,
             logger,
         )
 
+    def get_reference_data_loader(self,
+                           reference_data_config: dict[str, ReferenceConfig],
+                           **kwargs) -> BaseRefDataLoader[DuckDBPyRelation]:
+        return DuckDBRefDataLoader(connection=self._connection,
+                                   reference_data_config=reference_data_config,
+                                   dataset_config_uri=fh.get_parent(self._rules_path),
+                                   **kwargs)
     # pylint: disable=arguments-differ
     def write_file_to_parquet(  # type: ignore
         self, submission_file_uri: URI, submission_info: SubmissionInfo, output: URI

src/dve/pipeline/pipeline.py

@@ -26,7 +26,7 @@
 from dve.core_engine.backends.base.auditing import BaseAuditingManager
 from dve.core_engine.backends.base.contract import BaseDataContract
 from dve.core_engine.backends.base.core import EntityManager
-from dve.core_engine.backends.base.reference_data import BaseRefDataLoader
+from dve.core_engine.backends.base.reference_data import BaseRefDataLoader, ReferenceConfig
 from dve.core_engine.backends.base.rules import BaseStepImplementations
 from dve.core_engine.backends.exceptions import MessageBearingError
 from dve.core_engine.backends.readers import BaseFileReader
@@ -36,7 +36,7 @@
 from dve.core_engine.loggers import get_logger
 from dve.core_engine.message import FeedbackMessage
 from dve.core_engine.models import SubmissionInfo, SubmissionStatisticsRecord
-from dve.core_engine.type_hints import URI, DVEStageName, FileURI, InfoURI
+from dve.core_engine.type_hints import URI, DVEStageName, EntityName, FileURI, InfoURI
 from dve.parser import file_handling as fh
 from dve.parser.file_handling.implementations.file import LocalFilesystemImplementation
 from dve.parser.file_handling.service import _get_implementation
@@ -62,14 +62,13 @@ def __init__(
         step_implementations: Optional[BaseStepImplementations[EntityType]],
         rules_path: Optional[URI],
         submitted_files_path: Optional[URI],
-        reference_data_loader: Optional[type[BaseRefDataLoader]] = None,
         job_run_id: Optional[int] = None,
         logger: Optional[logging.Logger] = None,
     ):
         self._submitted_files_path = submitted_files_path
         self._processed_files_path = processed_files_path
         self._rules_path = rules_path
-        self._reference_data_loader = reference_data_loader
+        self._reference_data_loader = None
         self._job_run_id = job_run_id
         self._audit_tables = audit_tables
         self._data_contract = data_contract
@@ -113,6 +112,13 @@ def step_implementations(self) -> Optional[BaseStepImplementations[EntityType]]:
     def get_entity_count(entity: EntityType) -> int:
         """Get a row count of an entity stored as parquet"""
         raise NotImplementedError()
+    
+    def get_reference_data_loader(self,
+                           reference_data_config: dict[EntityName, ReferenceConfig],
+                           **kwargs) -> BaseRefDataLoader[EntityType]:
+        """Get reference data loader if required for business rules"""
+        raise NotImplementedError()
+        
 
     def get_submission_status(
         self, step_name: DVEStageName, submission_id: str
@@ -542,9 +548,6 @@ def apply_business_rules(   # pylint: disable=R0914
         if not self.rules_path:
             raise AttributeError("business rules path not provided.")
 
-        if not self._reference_data_loader:
-            raise AttributeError("reference data loader not provided.")
-
         if not self.processed_files_path:
             raise AttributeError("processed files path has not been provided.")
 
@@ -556,8 +559,8 @@ def apply_business_rules(   # pylint: disable=R0914
             self._processed_files_path, submission_info.submission_id
         )
         ref_data = config.get_reference_data_config()
+        reference_data = self.get_reference_data_loader(reference_data_config=ref_data)
         rules = config.get_rule_metadata()
-        reference_data = self._reference_data_loader(ref_data)  # type: ignore
         entities = {}
         contract = fh.joinuri(
             self.processed_files_path, submission_info.submission_id, "data_contract"