refactor: Refactor SubmissionStatus for record rejection

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

@@ -32,7 +32,7 @@
     QueueType,
     SubmissionResult,
 )
-from dve.pipeline.utils import SubmissionStatus
+from dve.pipeline.utils import EntityStatistics, SubmissionStatus
 
 AuditReturnType = TypeVar("AuditReturnType")  # pylint: disable=invalid-name
 
@@ -168,12 +168,14 @@ def __init__(
         submission_statistics: AuditorType,
         transfers: AuditorType,
         pool: Optional[ExecutorType] = None,
+        dataset_id: Optional[str] = None,
     ):
         """Audit manager to handle writing of audit information to auditors."""
         self._processing_status = processing_status
         self._submission_info = submission_info
         self._submission_statistics = submission_statistics
         self._transfers = transfers
+        self._dataset_id = dataset_id
         self.pool = pool
         if self.pool is not None:
             thread = isinstance(self.pool, ThreadPoolExecutor)
@@ -521,8 +523,16 @@ def get_submission_status(self, submission_id: str) -> Optional[SubmissionStatus
             sub_status.processing_failed = True
         if processing_rec.submission_result == "validation_failed":
             sub_status.validation_failed = True
-        if sub_stats_rec:
-            sub_status.number_of_records = sub_stats_rec.record_count
+        if sub_stats_rec and sub_stats_rec.record_count:
+            if not self._dataset_id:
+                raise AttributeError(
+                    f"Unable to find dataset id in {type(self).__name__}. Please ensure that " \
+                    +f"dataset id is defined in the setup of the {type(self).__name__} " \
+                    +"before using get_submission_status."
+                )
+            sub_status.entity_stats[self._dataset_id] = EntityStatistics(
+                no_records=sub_stats_rec.record_count
+            )
 
         return sub_status
 

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

@@ -169,6 +169,7 @@ def __init__(
         database_uri: URI,
         pool: Optional[ExecutorType] = None,
         connection: Optional[DuckDBPyRelation] = None,
+        dataset_id: Optional[str] = None,
     ):
         self._database_uri = database_uri
         self._connection = (
@@ -209,6 +210,7 @@ def __init__(
                 name="transfers",
                 connection=self._connection,  # type: ignore
             ),
+            dataset_id=dataset_id,
             pool=self._pool,
         )
 

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

@@ -172,6 +172,7 @@ def __init__(
         pool: Optional[ExecutorType] = None,
         spark: Optional[SparkSession] = None,
         table_format: Optional[SparkTableFormat] = "delta",
+        dataset_id: Optional[str] = None,
     ):
         self._database = database
         self._spark = spark if spark else SparkSession.builder.getOrCreate()
@@ -209,6 +210,7 @@ def __init__(
                 spark=self._spark,
             ),
             pool=self._pool,
+            dataset_id=dataset_id,
         )
 
     def combine_auditor_information(

src/dve/core_engine/models.py

@@ -116,10 +116,12 @@
 
     record_count: Optional[int]
     """Count of records in the submitted file"""
+    total_number_of_records_rejected: Optional[int]
+    """Total number of records rejected in a submitted file"""
     number_submission_rejections: Optional[int]
-    """Number of submission rejections raised following validation"""
+    """Number of submission rejection errors raised following validation"""
     number_record_rejections: Optional[int]
-    """Number of record rejections raised following validation"""
+    """Number of record rejection errors raised following validation"""
     number_warnings: Optional[int]
     """Number of warnings raised following validation"""
 

src/dve/pipeline/pipeline.py

@@ -42,7 +42,7 @@
 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
-from dve.pipeline.utils import SubmissionStatus, deadletter_file, load_config, load_reader
+from dve.pipeline.utils import EntityStatistics, SubmissionStatus, deadletter_file, load_config, load_reader
 from dve.reporting.constants import ErrorReportCategories
 from dve.reporting.error_report import ERROR_SCHEMA, calculate_aggregates
 
@@ -450,10 +450,13 @@ def apply_data_contract(
         entity_locations = {}
 
         for path, _ in fh.iter_prefix(read_from):
-            entity_locations[fh.get_file_name(path)] = path
-            entities[fh.get_file_name(path)] = self.data_contract.add_record_index(
+            entity_name = fh.get_file_name(path)
+            entity_locations[entity_name] = path
+            entity = self.data_contract.add_record_index(
                 self.data_contract.read_parquet(path)
             )
+            entities[entity_name] = entity
+            submission_status.create_new_entity_stat(entity_name, self.get_entity_count(entity))
 
         key_fields = {model: conf.reporting_fields for model, conf in model_config.items()}
 
@@ -620,8 +623,20 @@ def apply_business_rules(  # pylint: disable=R0914
                     entity,
                     entity_name,
                 )
+                entity_ct = self.get_entity_count(filtered_entity)
+                try:
+                    submission_status.entity_stats[entity_name].number_of_record_rejections = (
+                        submission_status.entity_stats[entity_name].number_of_records
+                        - entity_ct
+                    )
+                except KeyError:
+                    # Handling derived entities
+                    submission_status.create_new_entity_stat(entity_name, entity_ct)
             else:
                 self._logger.info(f"Skipping {entity_name}. Marked original.")
+                submission_status.entity_stats[entity_name] = EntityStatistics(
+                    no_records=self.get_entity_count(entity)
+                )
                 filtered_entity = entity
             projected = self._step_implementations.write_parquet(  # type: ignore
                 filtered_entity,
@@ -636,20 +651,6 @@ def apply_business_rules(  # pylint: disable=R0914
                 projected
             )
 
-        submission_status.number_of_records = self.get_entity_count(
-            entity=entity_manager.entities[f"""Original{rules.global_variables.get(
-                                              'entity',
-                                              submission_info.dataset_id)}"""]
-        )
-        submission_status.number_of_records_rejected = (
-            submission_status.number_of_records
-            - self.get_entity_count(
-                entity_manager.entities[
-                    rules.global_variables.get("entity", submission_info.dataset_id)
-                ]
-            )
-        )
-
         return submission_info, submission_status
 
     def business_rule_step(
@@ -823,7 +824,8 @@ def error_report(
         self._logger.info("Reading error dataframes")
         errors_df, aggregates = self._get_error_dataframes(submission_info.submission_id)
 
-        if not submission_status.number_of_records:
+        no_records = submission_status.number_of_records(submission_info.dataset_id)
+        if not no_records:
             sub_stats = None
         else:
             err_types = {
@@ -834,7 +836,10 @@ def error_report(
             }
             sub_stats = SubmissionStatisticsRecord(
                 submission_id=submission_info.submission_id,
-                record_count=submission_status.number_of_records,
+                record_count=no_records,
+                total_number_of_records_rejected=submission_status.number_of_record_rejections(
+                    submission_info.dataset_id
+                ),
                 number_submission_rejections=err_types.get(
                     ErrorReportCategories.FILE_REJECTION.reporting_name, 0
                 ),
@@ -904,7 +909,6 @@ def error_report_step(
             futures.append((info, status, pool.submit(self.error_report, info, status)))
 
         for info_dict, status in failed_file_transformation:
-            status.number_of_records = 0
             futures.append((info_dict, status, pool.submit(self.error_report, info_dict, status)))
 
         for sub_info, status, future in futures:

src/dve/pipeline/utils.py

@@ -14,7 +14,7 @@
 from dve.core_engine.backends.readers import _READER_REGISTRY
 from dve.core_engine.configuration.v1 import SchemaName, V1EngineConfig, _ModelConfig
 from dve.core_engine.loggers import get_logger
-from dve.core_engine.type_hints import URI, SubmissionResult
+from dve.core_engine.type_hints import URI, EntityName, SubmissionResult
 from dve.metadata_parser.model_generator import JSONtoPyd
 
 Dataset = dict[SchemaName, _ModelConfig]
@@ -79,19 +79,45 @@ def deadletter_file(source_uri: URI) -> None:
         return None
 
 
+class EntityStatistics:
+    """Statistics for a given entity"""
+
+    def __init__(self, no_records: int, no_record_rej: Optional[int] = None):
+        self._number_of_records = no_records
+        self._number_of_record_rejections = no_record_rej
+
+    @property
+    def number_of_records(self) -> int:
+        """Get the number of record for the entity"""
+        return self._number_of_records
+
+    @number_of_records.setter
+    def number_of_records(self, no_records: int):
+        """Set the number of records for the entity"""
+        self._number_of_records = no_records
+
+    @property
+    def number_of_record_rejections(self) -> int:
+        """Get the number of record rejections for the entity"""
+        return self._number_of_record_rejections if self._number_of_record_rejections else 0
+
+    @number_of_record_rejections.setter
+    def number_of_record_rejections(self, no_record_rej: int):
+        """Set the number of record rejections for the entity"""
+        self._number_of_record_rejections = no_record_rej
+
+
 class SubmissionStatus:
     """Submission status for a given submission."""
 
     def __init__(
         self,
         validation_failed: bool = False,
-        number_of_records: Optional[int] = None,
-        number_of_records_rejected: Optional[int] = None,
+        entity_stats: Optional[dict[EntityName, EntityStatistics]] = None,
         processing_failed: bool = False,
     ):
         self.validation_failed = validation_failed
-        self.number_of_records = number_of_records
-        self.number_of_records_rejected = number_of_records_rejected
+        self.entity_stats = entity_stats if entity_stats else {}
         self.processing_failed = processing_failed
 
     @property
@@ -103,3 +129,36 @@ def submission_result(self) -> SubmissionResult:
         if self.validation_failed:
             return "validation_failed"
         return "success"
+
+    def number_of_records(self, record_entity_name: EntityName) -> int:
+        """The total number of records across entities for a given submission."""
+        if not self.entity_stats:
+            return 0
+
+        return self.entity_stats[record_entity_name].number_of_records
+
+    def number_of_record_rejections(self, record_entity_name: EntityName) -> int:
+        """The total number of record rejections across entities for a given submission."""
+        if not self.entity_stats:
+            return 0
+
+        return self.entity_stats[record_entity_name].number_of_record_rejections
+
+    def create_new_entity_stat(self, entity_name: str, record_count: int):
+        """Create a new EntityStatistics object for a given entity."""
+        if self.entity_stats.get(entity_name):
+            raise LookupError("Record count is already set for {entity_name}." \
+                              +"Use update_number_of_records method instead")
+        self.entity_stats[entity_name] = EntityStatistics(no_records=record_count)
+
+    def update_number_of_records(self, entity_name: str, record_count: int):
+        """Update the number of records for a given entity"""
+        _new_record = self.entity_stats[entity_name]
+        _new_record.number_of_records = record_count
+        self.entity_stats[entity_name] = _new_record
+
+    def update_number_of_record_rejections(self, entity_name: str, record_rej_count: int):
+        """Update the number of record rejections for a given entity"""
+        _new_record = self.entity_stats[entity_name]
+        _new_record.number_of_record_rejections = record_rej_count
+        self.entity_stats[entity_name] = _new_record

src/dve/reporting/excel_report.py

@@ -147,8 +147,8 @@
                 "",
                 "Total Number of Records Processed",
                 (
-                    self.submission_status.number_of_records
-                    if self.submission_status.number_of_records
+                    self.submission_status.number_of_records(self.summary_dict["Dataset Id"])
+                    if self.submission_status.number_of_records(self.summary_dict["Dataset Id"])
                     else 0
                 ),  # pylint: disable=C0301
             ]
@@ -161,7 +161,9 @@
                 [
                     "",
                     "Total Number of Records Rejected",
-                    self.submission_status.number_of_records_rejected,
+                    self.submission_status.number_of_record_rejections(
+                        self.summary_dict["Dataset Id"]
+                    ),
                 ]
             )
         summary.append(["", ""])

tests/features/books.feature

@@ -11,17 +11,17 @@ Feature: Pipeline tests using the books dataset
         Then the latest audit record for the submission is marked with processing status file_transformation
         When I run the file transformation phase
         Then the header entity is stored as a parquet after the file_transformation phase
-        And the nested_books entity is stored as a parquet after the file_transformation phase
+        And the books entity is stored as a parquet after the file_transformation phase
         And the latest audit record for the submission is marked with processing status data_contract
         When I run the data contract phase
         Then there is 1 record rejection from the data_contract phase
         And the header entity is stored as a parquet after the data_contract phase
-        And the nested_books entity is stored as a parquet after the data_contract phase
+        And the books entity is stored as a parquet after the data_contract phase
         And the latest audit record for the submission is marked with processing status business_rules
         When I run the business rules phase
-        Then The rules restrict "nested_books" to 3 qualifying records
-        And The entity "nested_books" contains an entry for "17.85" in column "total_value_of_books"
-        And the nested_books entity is stored as a parquet after the business_rules phase
+        Then The rules restrict "books" to 3 qualifying records
+        And The entity "books" contains an entry for "17.85" in column "total_value_of_books"
+        And the books entity is stored as a parquet after the business_rules phase
         And the latest audit record for the submission is marked with processing status error_report
         When I run the error report phase
         Then An error report is produced
@@ -57,17 +57,17 @@ Feature: Pipeline tests using the books dataset
         Then the latest audit record for the submission is marked with processing status file_transformation
         When I run the file transformation phase
         Then the header entity is stored as a parquet after the file_transformation phase
-        And the nested_books entity is stored as a parquet after the file_transformation phase
+        And the books entity is stored as a parquet after the file_transformation phase
         And the latest audit record for the submission is marked with processing status data_contract
         When I run the data contract phase
         Then there is 1 record rejection from the data_contract phase
         And the header entity is stored as a parquet after the data_contract phase
-        And the nested_books entity is stored as a parquet after the data_contract phase
+        And the books entity is stored as a parquet after the data_contract phase
         And the latest audit record for the submission is marked with processing status business_rules
         When I run the business rules phase
-        Then The rules restrict "nested_books" to 3 qualifying records
-        And The entity "nested_books" contains an entry for "17.85" in column "total_value_of_books"
-        And the nested_books entity is stored as a parquet after the business_rules phase
+        Then The rules restrict "books" to 3 qualifying records
+        And The entity "books" contains an entry for "17.85" in column "total_value_of_books"
+        And the books entity is stored as a parquet after the business_rules phase
         And the latest audit record for the submission is marked with processing status error_report
         When I run the error report phase
         Then An error report is produced