feat: group rejections for mandatory primary keys (#152)

* docs: update jsonschema for entity relationships to include group rej code+message

Author: georgeRobertson

docs/advanced_guidance/json_schemas/entity_relationships.schema.json

@@ -25,6 +25,12 @@
                 },
                 "missing_parent_id_error_message": {
                     "type": "string"
+                },
+                "no_valid_records_error_code": {
+                    "type": "string"
+                },
+                "no_valid_records_error_message": {
+                    "type": "string"
                 }
             },
             "required": [

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

@@ -28,6 +28,7 @@
     CopyEntity,
     DeferredFilter,
     EntityRemoval,
+    GroupIdentification,
     HeaderJoin,
     ImmediateFilter,
     InnerJoin,
@@ -340,6 +341,14 @@ def remove_orphans(self, entities: Entities, *, config: OrphanRemoval) -> Iterat
         """
         raise NotImplementedError
 
+    @abstractmethod
+    def check_mandatory_group(self, entities: Entities, *, config: GroupIdentification) -> Iterator:
+        """
+        Check that a mandatory key in an entity has at least one valid entry in the all the child
+        entities.
+        """
+        raise NotImplementedError
+
     @abstractmethod
     def union(self, entities: Entities, *, config: TableUnion) -> Messages:
         """Union two entities together, taking the columns from each by name.
@@ -378,7 +387,7 @@ def identify_and_remove_orphans(
         entities: Entities,
         entity_hierarchy: EntityHierarchy,
         key_fields: Optional[dict[str, list[str]]] = None,
-    ) -> Messages:
+    ) -> tuple[Messages, 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.
@@ -388,6 +397,7 @@ def identify_and_remove_orphans(
         def process_node(
             node: HierarchyNode,
             orph_messages: Messages | None = None,
+            processed: bool = False,
         ):
             """Identify orphans and remove in a given node"""
 
@@ -416,6 +426,7 @@ def process_node(
                     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,
@@ -453,17 +464,95 @@ def process_node(
                             ]
                         )
 
+            return processed
+
+        processed = False
+
         for tree in entity_hierarchy.entity_trees.values():
             for node in tree.iterate_root_down():
-                process_node(node)
+                processed = process_node(node)
 
         _orph_rel = entities.get(ORPHANED_RECORD_ENTITY_NAME)
-        if _orph_rel:
+        if _orph_rel is not None:
             del entities[ORPHANED_RECORD_ENTITY_NAME]
 
         entities.update(entities)
 
-        return []
+        return [], processed
+
+    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]:
+        """
+        Identify that an entity with a mandatory key has at least one valid child record.
+        """
+
+        def process_node(
+            node: HierarchyNode,
+            processed: bool = False,
+        ) -> 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
+
+            self.logger.info(
+                f"Identifying that mandatory entity `{node.parent_entity}` has at least 1 valid child record"  # pylint: disable=C0301
+            )
+
+            join_expr = " AND ".join(
+                f"{node.parent_entity}.{k} = {node.entity_name}.{v}"
+                for k, v in node.join_fields.items()
+            )
+
+            with BackgroundMessageWriter(
+                working_directory=working_directory,
+                dve_stage=self.__stage_name__,
+                key_fields=key_fields,
+                logger=self.logger,
+            ) as msg_writer:
+                location = next(iter(node.join_fields.values()))
+                missing_children_records = self.check_mandatory_group(
+                    entities=entities,
+                    config=GroupIdentification(
+                        entity_name=node.parent_entity,
+                        target_name=node.entity_name,
+                        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
+
+        processed = False
+
+        for tree in entity_hierarchy.entity_trees.values():
+            for node in tree.iterate_lowest_descendent_up():
+                processed = process_node(node, processed)
+
+        entities.update(entities)
+
+        return [], processed
 
     # pylint: disable=R0912,R0914
     def apply_sync_filters(

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

@@ -43,6 +43,7 @@
     Aggregation,
     AntiJoin,
     ConfirmJoinHasMatch,
+    GroupIdentification,
     HeaderJoin,
     ImmediateFilter,
     InnerJoin,
@@ -450,6 +451,42 @@ def remove_orphans(self, entities: DuckDBEntities, *, config: OrphanRemoval) ->
             orphan_rel.filter(f"entity_name = '{config.entity_name}'")
         )
 
+    def check_mandatory_group(
+        self, entities: DuckDBEntities, *, config: GroupIdentification
+    ) -> Iterator:
+        """
+        Check that a mandatory key in an entity has at least one valid entry in the all the
+        child entities.
+        """
+        source_rel: DuckDBPyRelation = entities[config.entity_name]
+        source_rel = source_rel.set_alias(config.entity_name)
+        target_rel: DuckDBPyRelation = entities[config.target_name]
+        target_rel = target_rel.set_alias(config.target_name)
+
+        source_columns = [f"{config.entity_name}.{c.strip()}" for c in source_rel.columns]
+        _pk, fk = config.join_condition.split("=")
+
+        joined_rel = source_rel.join(target_rel, config.join_condition, "left").select(
+            *source_columns,
+            ColumnExpression(fk.strip()).alias("fk"),
+        )
+
+        missing_children_rel = joined_rel.filter("fk IS NULL")
+        filtered_rel = joined_rel.filter("fk IS NOT NULL").select(StarExpression(exclude=["fk"]))
+
+        _no_valid_child_records: tuple[int] = missing_children_rel.count("*").fetchone()  # type: ignore # pylint: disable=C0301
+        if _no_valid_child_records:
+            _no_valid_children = _no_valid_child_records[0]
+        else:
+            _no_valid_children = 0
+        self.logger.info(
+            f"Found {_no_valid_children} records with no valid children in {config.entity_name}."
+        )  # pylint: disable=C0301
+
+        entities[config.entity_name] = filtered_rel
+
+        return duckdb_rel_to_dictionaries(missing_children_rel)
+
     def union(self, entities: DuckDBEntities, *, config: TableUnion) -> Messages:
         """Union two entities together, taking the columns from each by name.
 

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

@@ -34,6 +34,7 @@
     ColumnAddition,
     ColumnRemoval,
     ConfirmJoinHasMatch,
+    GroupIdentification,
     HeaderJoin,
     ImmediateFilter,
     InnerJoin,
@@ -384,6 +385,12 @@ def remove_orphans(
         # TODO - implement for spark
         raise NotImplementedError
 
+    def check_mandatory_group(
+        self, entities: SparkEntities, *, config: GroupIdentification
+    ) -> Iterator:
+        # TODO - implement for spark
+        raise NotImplementedError
+
     def filter(self, entities: SparkEntities, *, config: ImmediateFilter) -> Messages:
         """Filter an entity immediately, and do not emit any messages.
 

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

@@ -33,13 +33,15 @@
     "CopyEntity",
     "DeferredFilter",
     "EntityRemoval",
+    "GroupIdentification",
     "HeaderJoin",
     "ImmediateFilter",
     "InnerJoin",
     "LeftJoin",
     "OneToOneJoin",
     "OneToOneJoin",
     "OrphanIdentification",
+    "OrphanRemoval",
     "ParentMetadata",
     "RenameEntity",
     "Rule",
@@ -565,6 +567,10 @@ class OrphanRemoval(BaseStep):
     """The reporting information for the row removal."""
 
 
+class GroupIdentification(AbstractConditionalJoin):
+    """Identify mandatory records which do not have any valid child records"""
+
+
 class Rule(BaseModel):
     """A rule, made up of multiple steps."""
 

src/dve/core_engine/type_hints.py

@@ -133,7 +133,9 @@
 """A string indicating the field that the error pertains to."""
 FieldValue = Optional[Any]
 """The value that caused the error."""
-ErrorCategory = Literal["Blank", "Wrong format", "Bad value", "Bad file", "Parent Missing"]
+ErrorCategory = Literal[
+    "Blank", "Wrong format", "Bad value", "Bad file", "Parent Missing", "Children missing"
+]
 """A string indicating the category of the error."""
 RecordIndex = Optional[int]
 """The record index that the error relates to (if applicable)"""

src/dve/parser/file_handling/service.py

@@ -273,9 +273,12 @@ def copy_resource(source_uri: URI, target_uri: URI, overwrite: bool = False) ->
     _transfer_resource(source_uri, target_uri, overwrite, "copy")
 
 
-def move_resource(source_uri: URI, target_uri: URI, overwrite: bool = False) -> None:
-    """Move a resource from one location to another."""
+def move_resource(source_uri: URI, target_uri: URI, overwrite: bool = False) -> URI:
+    """
+    Move a resource from one location to another. Returns the target_uri.
+    """
     _transfer_resource(source_uri, target_uri, overwrite, "move")
+    return target_uri
 
 
 def create_directory(target_uri: URI):

src/dve/pipeline/pipeline.py

@@ -1,4 +1,4 @@
-# pylint: disable=protected-access,too-many-instance-attributes,too-many-arguments,line-too-long
+# pylint: disable=protected-access,too-many-instance-attributes,too-many-arguments,line-too-long,too-many-lines
 """Generic Pipeline object to define how DVE should be interacted with."""
 
 import json
@@ -545,7 +545,7 @@ def data_contract_step(
 
         return processed_files, failed_processing
 
-    def apply_business_rules(  # pylint: disable=R0914
+    def apply_business_rules(  # pylint: disable=R0914,R0915
         self, submission_info: SubmissionInfo, submission_status: Optional[SubmissionStatus] = None
     ) -> tuple[SubmissionInfo, SubmissionStatus]:
         """Apply the business rules to a given submission, the submission may have failed at the
@@ -635,21 +635,69 @@ def apply_business_rules(  # pylint: disable=R0914
                 fh.joinuri(
                     self.processed_files_path,
                     submission_info.submission_id,
-                    "business_rules",
+                    "temp_business_rules",
                     entity_name,
                 ),
             )
             entity_manager.entities[entity_name] = self.step_implementations.read_parquet(  # type: ignore
                 projected
             )
 
-        self.step_implementations.identify_and_remove_orphans(  # type: ignore
+        _, orph_or_group = 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
             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
+            working_directory,
+            entity_manager.entities,
+            entity_hierarchy,
+            key_fields,
+        )
+
+        for entity_name, entity in entity_manager.entities.items():
+            if orph_or_group:
+                self._logger.info(f"Writing {entity_name} out to disk.")
+                final_projection = self._step_implementations.write_parquet(  # type: ignore
+                    entity,
+                    fh.joinuri(
+                        self.processed_files_path,
+                        submission_info.submission_id,
+                        "business_rules",
+                        entity_name,
+                    ),
+                )
+            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
+                    )
+                )
+
+            entity_manager.entities[entity_name] = self.step_implementations.read_parquet(  # type: ignore
+                final_projection
+            )
+
         submission_status.number_of_records = self.get_entity_count(
             entity=entity_manager.entities[f"""Original{rules.global_variables.get(
                                               'entity',

src/dve/pipeline/utils.py

@@ -94,7 +94,7 @@ def unpersist_all_rdds(spark: SparkSession):
         rdd.unpersist()
 
 
-def deadletter_file(source_uri: URI) -> None:
+def deadletter_file(source_uri: URI) -> URI | None:
     """Move files that can't be processed to a deadletter location"""
     try:
         source_parent: URI = source_uri.rsplit("/", 1)[0]