feat: added proposals for new models and objects to store entity hierarchy information

Author: stevenhsd

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

@@ -156,8 +156,9 @@ def apply_data_contract(
                                 fld, fld_info.annotation
                             ).alias(fld)
                             if fld in record_df.columns
-                            else lit(None).cast(
-                                get_type_from_annotation(fld_info.annotation)).alias(fld)
+                            else lit(None)
+                            .cast(get_type_from_annotation(fld_info.annotation))
+                            .alias(fld)
                         )
                         for fld, fld_info in entity_fields.items()
                     ],

src/dve/core_engine/configuration/v1/__init__.py

@@ -20,12 +20,17 @@
     BusinessFilterSpecConfig,
     BusinessRuleSpecConfig,
 )
-from dve.core_engine.configuration.v1.hierarchy import HierarchyNode, ChildHierarchyNode
 from dve.core_engine.configuration.v1.steps import StepConfigUnion
 from dve.core_engine.message import DataContractErrorDetail
-from dve.core_engine.type_hints import EntityName, ErrorCategory, ErrorType, TemplateVariables
+from dve.core_engine.type_hints import (
+    EntityName,
+    ErrorCategory,
+    ErrorCode,
+    ErrorMessage,
+    ErrorType,
+    TemplateVariables,
+)
 from dve.core_engine.validation import RowValidator
-from dve.metadata_parser.exc import EntityNotFoundError
 from dve.parser.file_handling import joinuri, open_stream, resolve_location
 from dve.parser.type_hints import URI, Extension
 
@@ -40,8 +45,8 @@
 
 FieldName = str
 """The name of a field within a model/schema."""
-JoinFields = Optional[list[str]]
-"""The fields required to join a child entity back to the parent"""
+JoinFields = Optional[dict[str, str]]
+"""The fields required ( parent > child ) to join a child entity back to the parent"""
 TypeOrDef = Union[  # pylint: disable=C0103
     TypeName, "_CallableTypeDefinition", "_ModelTypeDefinition", "_TypeAliasDefinition"
 ]
@@ -87,12 +92,23 @@ class _TypeAliasDefinition(_BaseTypeDefintion):
 
 class _LinkageConfig(BaseModel):
     """Specify how to link entities back to parents if required"""
+
     parent_entity: EntityName
     """The name of the parent entity"""
     join_fields: JoinFields
     """The fields that can be used to link back to the parent entity"""
     mandatory: Optional[bool] = False
     """If the entity is a child, is it a mandatory field of the parent"""
+    no_valid_records_error_code: Optional[ErrorCode] = "NoValidRecords"
+    """The error code to emit if the entity has no valid records and is mandatory in the parent entity"""  # pylint: disable=C0301
+    no_valid_records_error_message: Optional[ErrorMessage] = (
+        "parent record removed as no valid child records"
+    )
+    """The error message to emit if the entity has no valid records and is mandatory in the parent entity"""  # pylint: disable=C0301
+    orphaned_records_error_code: Optional[ErrorCode] = "OrphanedRecords"
+    """The error code to emit if the entity contains records that are orphaned by parent record rejections"""  # pylint: disable=C0301
+    orphaned_records_error_message: Optional[ErrorMessage] = "Orphaned records removed"
+    """The error code to emit if the entity contains records that are orphaned by parent record rejections"""  # pylint: disable=C0301
 
 
 class _SchemaConfig(BaseModel):
@@ -345,7 +361,7 @@ def get_contract_metadata(self) -> DataContractMetadata:
             reader_metadata=reader_metadata,
             validators=validators,
             reporting_fields=reporting_fields,
-            cache_originals=self.contract.cache_originals
+            cache_originals=self.contract.cache_originals,
         )
 
     def load_error_message_info(self, uri):
@@ -367,24 +383,3 @@ def get_rule_metadata(self) -> RuleMetadata:
             global_variables=self.transformations.parameters,  # pylint: disable=E1101
             reference_data_config=self.get_reference_data_config(),
         )
-
-    def get_entity_hierarchy(self) -> dict[str, HierarchyNode]:
-        """Determine the linkage hierarchy using contact config"""
-        top_level_parents = {
-            entity_name: HierarchyNode(entity_name=entity_name) 
-            for entity_name in self.contract.datasets
-            if not entity_name in self.entity_relationships
-            }
-
-        for name, linkage_detail in self.entity_relationships.items():
-            for main_entity, details in top_level_parents.items(): 
-                if (linkage_detail.parent_entity == main_entity 
-                    or linkage_detail.parent_entity in details.get_descendents()): 
-                    top_level_parents[main_entity].add_child_node(linkage_detail.parent_entity,
-                                                                  ChildHierarchyNode(entity_name=name,
-                                                                                     join_fields=linkage_detail.join_fields,
-                                                                                     mandatory=linkage_detail.mandatory))
-                    break
-            else:
-                raise EntityNotFoundError(f"Can't find parent entity {linkage_detail.parent_entity} defined to establish hierarchy for {name} - please ensure it is defined above any child entities in the dischema.")
-        return top_level_parents

src/dve/core_engine/configuration/v1/hierarchy.py

@@ -1,56 +1,131 @@
 """Classes to help determine and store entity hierarchy information."""
 
-from typing import Any, Optional, Union
+import json
+from typing import Any, Iterable, Optional, Union
+
 from pydantic import BaseModel, Field
+
+from dve.core_engine.configuration.v1 import V1EngineConfig, _LinkageConfig
+from dve.core_engine.type_hints import EntityName, ErrorCode, ErrorMessage
 from dve.metadata_parser.exc import EntityNotFoundError
+from dve.parser.file_handling.service import open_stream
+from dve.parser.type_hints import URI
+
 
 class HierarchyNode(BaseModel):
     """Stores entity hierarchy information"""
+
     entity_name: str
-    mandatory: Optional[bool] = False
-    children: Optional[list["HierarchyNode"]] = Field(default_factory=list)
-    
+    children: list["HierarchyNode"] = Field(default_factory=list)
+
     def get_descendents(self) -> list[str]:
         """Recursively list all descendents of the node"""
         descendents = []
         for node in self.children:
             descendents.append(node.entity_name)
             descendents.extend(node.get_descendents())
         return descendents
-    
-    def get_node(self, entity_name:str) -> Union["HierarchyNode", None]:
+
+    def get_node(self, entity_name: str) -> Union["HierarchyNode", None]:
         """Recursively search for node and return if found"""
         node = None
         if self.entity_name == entity_name:
             return self
-        else:
-            for child in self.children:
-                node = child.get_node(entity_name)
-                if node:
-                    break
+        for child in self.children:
+            node = child.get_node(entity_name)
+            if node:
+                break
         return node
-    
+
     def add_child_node(self, parent_entity: str, child_info: "HierarchyNode") -> None:
         """Add a child node if the parent exists in the hierarchy"""
         try:
-            self.get_node(parent_entity).children.append(child_info)
-        except AttributeError:
-            raise EntityNotFoundError(f"Can't find parent node {parent_entity} in {self.entity_name}")
-    
+            self.get_node(parent_entity).children.append(child_info)  # type: ignore
+        except AttributeError as exc:
+            raise EntityNotFoundError(
+                f"Can't find parent node {parent_entity} in {self.entity_name}"
+            ) from exc
+
     def as_dict(self) -> dict[str, dict[str, Any]]:
         """Get dictionary representation of entity hierarchy"""
         child_dict = {}
         for node in self.children:
-             child_dict.update(node.as_dict())
-        
-        ret_dict = {"children": child_dict,
-                    "mandatory": self.mandatory}
-        if hasattr(self, "join_fields"):
-            ret_dict |= {"join_fields": self.join_fields}
-        
+            child_dict.update(node.as_dict())
+
+        ret_dict = self.model_dump(exclude={"entity_name", "children"})
+        ret_dict.update({"children": child_dict})
+
         return {self.entity_name: ret_dict}
 
 
 class ChildHierarchyNode(HierarchyNode):
     """Stores child entity hierarchy information"""
-    join_fields: list[str]
+
+    join_fields: dict[str, str]
+    mandatory: Optional[bool] = False
+    no_valid_records_error_code: Optional[ErrorCode] = "NoValidRecords"
+    no_valid_records_error_message: Optional[ErrorMessage] = (
+        "parent record removed as no valid child records"
+    )
+    orphaned_records_error_code: Optional[ErrorCode] = "OrphanedRecords"
+    orphaned_records_error_message: Optional[ErrorMessage] = "Orphaned records removed"
+
+
+class EntityHierarchy:
+    """Determines and stores entity hierarchy information from config"""
+
+    def __init__(self, entity_trees: dict[EntityName, HierarchyNode]):
+        self.entity_trees = entity_trees
+
+    @staticmethod
+    def determine_trees(
+        all_datasets: Iterable[str], entity_relationships: dict[str, _LinkageConfig]
+    ) -> dict[EntityName, HierarchyNode]:
+        """Determine the entity hierarchy trees and store as HierarchyNodes"""
+        top_level_parents: dict[EntityName, HierarchyNode] = {
+            entity_name: HierarchyNode(entity_name=entity_name)
+            for entity_name in all_datasets
+            if not entity_name in entity_relationships
+        }
+
+        for name, linkage_detail in entity_relationships.items():
+            for main_entity, parent_node in top_level_parents.items():
+                if (
+                    linkage_detail.parent_entity == main_entity
+                    or linkage_detail.parent_entity in parent_node.get_descendents()
+                ):
+                    parent_node.add_child_node(
+                        linkage_detail.parent_entity,
+                        ChildHierarchyNode(
+                            entity_name=name, **linkage_detail.model_dump(exclude={"parent_entity"})
+                        ),
+                    )
+                    break
+            else:
+                raise EntityNotFoundError(
+                    f"Can't find parent entity {linkage_detail.parent_entity} defined to "
+                    + f"establish hierarchy for {name} - please ensure it is defined above "
+                    + "any child entities in the dischema."
+                )
+        return top_level_parents
+
+    @classmethod
+    def from_dischema(cls, dischema_uri: URI):
+        """Create entity hierarchy direct from dischema"""
+        with open_stream(dischema_uri) as dischema:
+            config_dict = json.load(dischema)
+        all_datasets = config_dict.get("contract", {}).get("datasets", {}).keys()
+        entity_relationships = {
+            k: _LinkageConfig(**v) for k, v in config_dict.get("entity_relationships", {}).items()
+        }
+        return cls(entity_trees=cls.determine_trees(all_datasets, entity_relationships))
+
+    @classmethod
+    def from_engine_config(cls, engine_config: V1EngineConfig):
+        """Create entity hierarchy direct from engine config"""
+        return cls(
+            entity_trees=cls.determine_trees(
+                all_datasets=engine_config.contract.datasets.keys(),
+                entity_relationships=engine_config.entity_relationships,
+            )
+        )

src/dve/metadata_parser/exc.py

@@ -3,7 +3,8 @@
 
 class EntityNotFoundError(KeyError):
     """Error for missing entities"""
-    
+
+
 class LocWarning(UserWarning):
     """Warning class with optional location parameter"""