fix: ensure missing parent and group rejection records are being removed - add test coverage

Author: stevenhsd

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

@@ -387,119 +387,113 @@ def identify_and_remove_orphans(
         entities: Entities,
         entity_hierarchy: EntityHierarchy,
         key_fields: Optional[dict[str, list[str]]] = None,
-    ) -> tuple[Messages, bool]:
+    ) -> tuple[Messages, dict[EntityName, bool]]:
         """
         Identifies and removes orphan records by traversing the EntityHierarchy object.
         An orphan is a child record whose parent FK does not exist in the parent entity.
         Processes recursively: removes orphans at each level, then processes children.
         """
 
         def process_node(
-            node: HierarchyNode,
-            orph_messages: Messages | None = None,
-            processed: bool = False,
+            node: HierarchyNode
         ):
             """Identify orphans and remove in a given node"""
+            issues_found: bool = False
+            if node.parent_entity is None:
+                return issues_found
 
-            if orph_messages is None:
-                orph_messages = []
+            self.logger.info(f"Identifying orphans in {node.entity_name}")
 
-            if node.parent_entity is not None:
-                self.logger.info(f"Identifying orphans in {node.entity_name}")
+            join_expr = " AND ".join(
+                f"{node.parent_entity}.{k} = {node.entity_name}.{v}"
+                for k, v in node.join_fields.items()
+            )
 
-                join_expr = " AND ".join(
-                    f"{node.parent_entity}.{k} = {node.entity_name}.{v}"
-                    for k, v in node.join_fields.items()
-                )
+            _, no_orphs = self.identify_orphans(
+                entities=entities,
+                config=OrphanIdentification(
+                    id=list(node.join_fields.values())[0],
+                    entity_name=node.entity_name,
+                    target_name=node.parent_entity,
+                    join_condition=join_expr,
+                ),
+            )
 
-                _, no_orphs = self.identify_orphans(
-                    entities=entities,
-                    config=OrphanIdentification(
-                        id=list(node.join_fields.values())[0],
-                        entity_name=node.entity_name,
-                        target_name=node.parent_entity,
-                        join_condition=join_expr,
-                    ),
+            if no_orphs > 0:
+                self.logger.info(
+                    f"Removing records with missing parent from {node.entity_name}"
                 )
-
-                if no_orphs > 0:
-                    self.logger.info(
-                        f"Removing records with missing parent from {node.entity_name}"
-                    )
-                    processed = True
-                    location = list(node.join_fields.values())[0]
-                    with BackgroundMessageWriter(
-                        working_directory=working_directory,
-                        dve_stage=self.__stage_name__,
-                        key_fields=key_fields,
-                        logger=self.logger,
-                    ) as msg_writer:
-                        _orph_records = self.remove_orphans(
-                            entities=entities,
-                            config=OrphanRemoval(
-                                entity_name=node.entity_name,
-                                reporting=ReportingConfig(
-                                    emit="record_failure",
-                                    code=node.missing_parent_id_error_code,
-                                    message=node.missing_parent_id_error_message,
-                                    location=location,
-                                ),
+                issues_found = True
+                location = list(node.join_fields.values())[0]
+                with BackgroundMessageWriter(
+                    working_directory=working_directory,
+                    dve_stage=self.__stage_name__,
+                    key_fields=key_fields,
+                    logger=self.logger,
+                ) as msg_writer:
+                    _orph_records = self.remove_orphans(
+                        entities=entities,
+                        config=OrphanRemoval(
+                            entity_name=node.entity_name,
+                            reporting=ReportingConfig(
+                                emit="record_failure",
+                                code=node.missing_parent_id_error_code,
+                                message=node.missing_parent_id_error_message,
+                                location=location,
                             ),
-                        )
-                        # moved to batch the write - risky if large number of
-                        msg_writer.write_queue.put(
-                            [
-                                FeedbackMessage(
-                                    entity=node.entity_name,
-                                    record=record,  # type: ignore
-                                    error_location=location,
-                                    error_message=node.missing_parent_id_error_message,
-                                    failure_type="record",
-                                    error_type="record",
-                                    error_code=node.missing_parent_id_error_code,
-                                    reporting_field=location,
-                                    category="Parent Missing",
-                                )
-                                for record in _orph_records
-                            ]
-                        )
+                        ),
+                    )
+                    # moved to batch the write - risky if large number of
+                    msg_writer.write_queue.put(
+                        [
+                            FeedbackMessage(
+                                entity=node.entity_name,
+                                record=record,  # type: ignore
+                                error_location=location,
+                                error_message=node.missing_parent_id_error_message,
+                                failure_type="record",
+                                error_type="record",
+                                error_code=node.missing_parent_id_error_code,
+                                reporting_field=location,
+                                category="Parent Missing",
+                            )
+                            for record in _orph_records
+                        ]
+                    )
 
-            return processed
+            return issues_found
 
-        processed = False
+        entity_issues_found: dict[EntityName, bool] = {}
 
         for tree in entity_hierarchy.entity_trees.values():
             for node in tree.iterate_root_down():
-                processed = process_node(node)
+                entity_issues_found[node.entity_name] = process_node(node)
 
         _orph_rel = entities.get(ORPHANED_RECORD_ENTITY_NAME)
         if _orph_rel is not None:
             del entities[ORPHANED_RECORD_ENTITY_NAME]
 
         entities.update(entities)
 
-        return [], processed
+        return [], entity_issues_found
 
     def identify_and_remove_missing_mandatory_groups(
         self,
         working_directory: URI,
         entities: Entities,
         entity_hierarchy: EntityHierarchy,
         key_fields: Optional[dict[str, list[str]]] = None,
-    ) -> tuple[Messages, bool]:
+    ) -> tuple[Messages, dict[EntityName, bool]]:
         """
         Identify that an entity with a mandatory key has at least one valid child record.
         """
 
         def process_node(
-            node: HierarchyNode,
-            processed: bool = False,
+            node: HierarchyNode
         ) -> bool:
             """Identify at least one valid child for a mandatory entity at a given node."""
             if node.parent_entity is None or not node.mandatory:
-                return processed
-
-            processed = True
+                return False
 
             self.logger.info(
                 f"Identifying that mandatory entity `{node.parent_entity}` has at least 1 valid child record"  # pylint: disable=C0301
@@ -525,34 +519,33 @@ def process_node(
                         join_condition=join_expr,
                     ),
                 )
-                for record in missing_children_records:
-                    msg_writer.write_queue.put(
-                        [
-                            FeedbackMessage(
-                                entity=node.parent_entity,
-                                record=record,  # type: ignore
-                                error_location=location,
-                                error_message=node.no_valid_records_error_message,
-                                failure_type="record",
-                                error_type="record",
-                                error_code=node.no_valid_records_error_code,
-                                reporting_field=location,
-                                category="Children missing",
-                            )
-                        ]
-                    )
-
-            return processed
+                _messages = [
+                        FeedbackMessage(
+                            entity=node.parent_entity,
+                            record=record,  # type: ignore
+                            error_location=location,
+                            error_message=node.no_valid_records_error_message,
+                            failure_type="record",
+                            error_type="record",
+                            error_code=node.no_valid_records_error_code,
+                            reporting_field=location,
+                            category="Children missing",
+                        )
+                    for record in missing_children_records
+                ]
+                msg_writer.write_queue.put(_messages)
+            return len(_messages) > 0
 
-        processed = False
+        entity_issues_found: dict[EntityName, bool] = {}
 
         for tree in entity_hierarchy.entity_trees.values():
             for node in tree.iterate_lowest_descendent_up():
-                processed = process_node(node, processed)
+                if node.parent_entity and node.mandatory:
+                    entity_issues_found[node.parent_entity] = process_node(node)
 
-        entities.update(entities)
+        #entities.update(entities)
 
-        return [], processed
+        return [], entity_issues_found
 
     # pylint: disable=R0912,R0914
     def apply_sync_filters(

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

@@ -3,7 +3,7 @@
 import warnings
 from collections import deque
 from collections.abc import Sequence
-from typing import Optional
+from typing import Iterator, Optional
 
 import pyarrow  # type: ignore
 import pyarrow.parquet as pq  # type: ignore
@@ -144,3 +144,5 @@ def check_if_parquet_file(file_location: URI) -> bool:
         return True
     except (pyarrow.ArrowInvalid, pyarrow.ArrowIOError):
         return False
+
+    

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

@@ -434,7 +434,7 @@ def identify_orphans(
 
     def remove_orphans(self, entities: DuckDBEntities, *, config: OrphanRemoval) -> Iterator:
         """Method to remove identified orphans in the orphan tracker entity."""
-        orphan_rel = entities[ORPHANED_RECORD_ENTITY_NAME].set_alias("orphan")
+        orphan_rel = entities[ORPHANED_RECORD_ENTITY_NAME].filter(f"entity_name = '{config.entity_name}'").set_alias("orphan")
         filtered_rel = (
             entities[config.entity_name]
             .set_alias(config.entity_name)
@@ -448,7 +448,7 @@ def remove_orphans(self, entities: DuckDBEntities, *, config: OrphanRemoval) ->
         entities[config.entity_name] = filtered_rel
 
         return duckdb_rel_to_dictionaries(
-            orphan_rel.filter(f"entity_name = '{config.entity_name}'")
+            orphan_rel
         )
 
     def check_mandatory_group(

src/dve/pipeline/pipeline.py

@@ -643,30 +643,39 @@ def apply_business_rules(  # pylint: disable=R0914,R0915
                 projected
             )
 
-        _, orph_or_group = self.step_implementations.identify_and_remove_orphans(  # type: ignore
+        _, orph_issues_1 = self.step_implementations.identify_and_remove_orphans(  # type: ignore
             working_directory,
             entity_manager.entities,
             entity_hierarchy,
             key_fields,
         )
+        
 
-        _, orph_or_group = self.step_implementations.identify_and_remove_missing_mandatory_groups(  # type: ignore
+        _, grp_issues_1 = self.step_implementations.identify_and_remove_missing_mandatory_groups(  # type: ignore
             working_directory,
             entity_manager.entities,
             entity_hierarchy,
             key_fields,
         )
+        
 
         # Perform a second time incase the mandatory groups result in new orphans
-        _, orph_or_group = self.step_implementations.identify_and_remove_orphans(  # type: ignore
+        _, orph_issues_2 = self.step_implementations.identify_and_remove_orphans(  # type: ignore
             working_directory,
             entity_manager.entities,
             entity_hierarchy,
             key_fields,
         )
+        
+        entity_issues: dict[EntityName, bool] = {
+            entity: any(
+                val for val in (orph_issues_1.get(entity, False), grp_issues_1.get(entity, False), orph_issues_2.get(entity, False)))
+            for entity in orph_issues_1.keys()
+            }
 
+        unchanged_entities: list[EntityName] = []
         for entity_name, entity in entity_manager.entities.items():
-            if orph_or_group:
+            if entity_issues.get(entity_name, False):
                 self._logger.info(f"Writing {entity_name} out to disk.")
                 final_projection = self._step_implementations.write_parquet(  # type: ignore
                     entity,
@@ -677,27 +686,41 @@ def apply_business_rules(  # pylint: disable=R0914,R0915
                         entity_name,
                     ),
                 )
+                
+                entity_manager.entities[entity_name] = self.step_implementations.read_parquet(  # type: ignore
+                final_projection
+            )
             else:
-                self._logger.info(f"Moving {entity_name} from temp_business_rules to business_rules")
-                final_projection = fh.move_resource(
-                    source_uri=fh.joinuri(
-                        self.processed_files_path,
-                        submission_info.submission_id,
-                        "temp_business_rules",
-                        entity_name
-                    ),
-                    target_uri=fh.joinuri(
-                        self.processed_files_path,
-                        submission_info.submission_id,
-                        "business_rules",
-                        entity_name
-                    )
-                )
+                unchanged_entities.append(entity_name)
+        
+        for entity_name in unchanged_entities:
+            self._logger.info(f"Moving {entity_name} from temp_business_rules to business_rules")
+            final_projection = fh.move_resource(
+                source_uri=fh.joinuri(
+                    self.processed_files_path,
+                    submission_info.submission_id,
+                    "temp_business_rules",
+                    entity_name
+                ),
+                target_uri=fh.joinuri(
+                    self.processed_files_path,
+                    submission_info.submission_id,
+                    "business_rules",
+                    entity_name
+                ),
+                overwrite=True
+            )
 
             entity_manager.entities[entity_name] = self.step_implementations.read_parquet(  # type: ignore
                 final_projection
             )
 
+        
+        fh.remove_prefix(fh.joinuri(
+                    self.processed_files_path,
+                    submission_info.submission_id,
+                    "temp_business_rules"))
+        
         submission_status.number_of_records = self.get_entity_count(
             entity=entity_manager.entities[f"""Original{rules.global_variables.get(
                                               'entity',