fix: address some review comments

Author: stevenhsd

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

@@ -1,61 +1,13 @@
 """Metadata classes for the data contract."""
 
-from typing import Any, Optional, Union
+from typing import Any
 
-from pydantic import BaseModel, Field, PrivateAttr, model_validator
+from pydantic import BaseModel, PrivateAttr, model_validator
 
 from dve.core_engine.type_hints import EntityName, ReportingFields
 from dve.core_engine.validation import RowValidator
-from dve.metadata_parser.exc import EntityNotFoundError
 from dve.parser.type_hints import Extension
 
-class HierarchyNode(BaseModel):
-    entity_name: str
-    children: Optional[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]:
-        """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
-        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}")
-    
-    def as_dict(self):
-        ret_dict = {}
-        for node in self.children:
-             ret_dict.update(node.as_dict())
-        return {self.entity_name: {"children": ret_dict}}
-
-
-class ChildHierarchyNode(HierarchyNode):
-    join_fields: list[str]
-    
-    def as_dict(self):
-        ret_value = {self.entity_name: {"join_fields": self.join_fields}}
-        for node in self.children:
-            ret_value[self.entity_name] |= {"children": node.as_dict()}
-        return ret_value
-
 
 class ReaderConfig(BaseModel):
     """Configuration options for a given reader."""
@@ -86,7 +38,6 @@ class DataContractMetadata(BaseModel, frozen=True, arbitrary_types_allowed=True)
     """Whether to cache the original entities after loading."""
     _schemas: dict[EntityName, type[BaseModel]] = PrivateAttr(default_factory=dict)
     """The pydantic models of the schmas."""
-    linkage_hierarchy: dict[EntityName, HierarchyNode] = Field(default_factor=dict)
 
     @property
     def schemas(self) -> dict[EntityName, type[BaseModel]]:

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

@@ -7,7 +7,7 @@
 from typing_extensions import Literal
 
 from dve.core_engine.backends.base.reference_data import ReferenceConfig, ReferenceConfigUnion
-from dve.core_engine.backends.metadata.contract import ChildHierarchyNode, DataContractMetadata, HierarchyNode, ReaderConfig
+from dve.core_engine.backends.metadata.contract import DataContractMetadata, ReaderConfig
 from dve.core_engine.backends.metadata.rules import AbstractStep, Rule, RuleMetadata
 from dve.core_engine.configuration.base import BaseEngineConfig
 from dve.core_engine.configuration.v1.filters import (
@@ -20,6 +20,7 @@
     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
@@ -90,6 +91,8 @@ class _LinkageConfig(BaseModel):
     """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"""
 
 
 class _SchemaConfig(BaseModel):
@@ -123,7 +126,6 @@ class _ModelConfig(_SchemaConfig):
     """Reader configuration options for the model."""
     aliases: dict[FieldName, FieldName] = Field(default_factory=dict)
     """An alias field name mapping."""
-    linkage_details: Optional[_LinkageConfig] = None
 
 
 class _RuleStoreConfig(BaseModel):
@@ -189,6 +191,8 @@ class V1EngineConfig(BaseEngineConfig):
         default_factory=dict
     )
     """Rule store rules from the loaded rule stores."""
+    entity_relationships: dict[EntityName, _LinkageConfig] = Field(default_factory=dict)
+    """The parent-child relationships linking the defined entities"""
 
     @validate_call
     def _update_rule_store(self, rule_store: dict[RuleName, BusinessComponentSpecConfigUnion]):
@@ -341,8 +345,7 @@ def get_contract_metadata(self) -> DataContractMetadata:
             reader_metadata=reader_metadata,
             validators=validators,
             reporting_fields=reporting_fields,
-            cache_originals=self.contract.cache_originals,
-            linkage_hierarchy=self.determine_entity_hierarchy()
+            cache_originals=self.contract.cache_originals
         )
 
     def load_error_message_info(self, uri):
@@ -365,20 +368,22 @@ def get_rule_metadata(self) -> RuleMetadata:
             reference_data_config=self.get_reference_data_config(),
         )
 
-    def determine_entity_hierarchy(self) -> list[HierarchyNode]:
+    def get_entity_hierarchy(self) -> dict[str, HierarchyNode]:
         """Determine the linkage hierarchy using contact config"""
-        linkage_hierarchy = {name: model_conf.linkage_details for name, model_conf in self.contract.datasets.items()}
-        top_level_parents = {}
-        for name, linkage_detail in linkage_hierarchy.items():
-            if not linkage_detail:
-                top_level_parents[name] = HierarchyNode(entity_name=name)
-                continue
+        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))
+                                                                                     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.")

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

@@ -0,0 +1,56 @@
+"""Classes to help determine and store entity hierarchy information."""
+
+from typing import Any, Optional, Union
+from pydantic import BaseModel, Field
+from dve.metadata_parser.exc import EntityNotFoundError
+
+class HierarchyNode(BaseModel):
+    """Stores entity hierarchy information"""
+    entity_name: str
+    mandatory: Optional[bool] = False
+    children: Optional[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]:
+        """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
+        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}")
+    
+    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}
+        
+        return {self.entity_name: ret_dict}
+
+
+class ChildHierarchyNode(HierarchyNode):
+    """Stores child entity hierarchy information"""
+    join_fields: list[str]

tests/test_core_engine/test_config_load.py

@@ -123,11 +123,7 @@
                 "mandatory_fields": [
                     "ds_003_id",
                     "ds_001_id"
-                ],
-                "linkage_details": {
-                    "parent_entity": "ds_001",
-                    "join_fields": ["ds_001_id"]
-                }
+                ]
             },
              "ds_101": {
                 "fields": {
@@ -147,11 +143,7 @@
                 "mandatory_fields": [
                     "referral_id",
                     "ds_001_id"
-                ],
-                "linkage_details": {
-                    "parent_entity": "ds_001",
-                    "join_fields": ["ds_001_id"]
-                }
+                ]
             },
              "ds_201": {
                 "fields": {
@@ -172,11 +164,7 @@
                     "referral_id",
                     "ds_201_id",
                     "contact_date"
-                ],
-                "linkage_details": {
-                    "parent_entity": "ds_101",
-                    "join_fields": ["referral_id"]
-                }
+                ]
             },
              "ds_202": {
                 "fields": {
@@ -196,11 +184,7 @@
                 "mandatory_fields": [
                     "ds_202_id",
                     "ds_201_id"
-                ],
-                "linkage_details": {
-                    "parent_entity": "ds_201",
-                    "join_fields": ["ds_201_id"]
-                }
+                ]
             }
         }
     },
@@ -214,38 +198,60 @@
                 "failure_message": "Record rejected - `{{ name }}` is not valid."
             }
         ]
-    }
+    },
+    "entity_relationships": {
+            "ds_003": {
+                    "parent_entity": "ds_001",
+                    "join_fields": ["ds_001_id"],
+                    "mandatory": false
+                },
+            "ds_101": {
+                    "parent_entity": "ds_001",
+                    "join_fields": ["ds_001_id"],
+                    "mandatory_entity": true
+                },
+            "ds_201": {
+                    "parent_entity": "ds_101",
+                    "join_fields": ["referral_id"],
+                    "mandatory": false
+                },
+            "ds_202": {
+                    "parent_entity": "ds_201",
+                    "join_fields": ["ds_201_id"],
+                    "mandatory": true
+                }
+        }
 }"""
 
 def test_no_linkage_config_load():
     config = V1EngineConfig(location="", 
                             **json.loads(CONFIG_WITHOUT_LINKAGE))
     assert len(config.contract.datasets) == 1
-    dc_metadata = config.get_contract_metadata()
-    assert len(dc_metadata.linkage_hierarchy) == 1
-    assert not dc_metadata.linkage_hierarchy.get("animals").children
+    hierarchy = config.get_entity_hierarchy()
+    assert len(hierarchy) == 1
+    assert not hierarchy.get("animals").children
 
 
 def test_linkage_config_load():
     config = V1EngineConfig(location="", 
                             **json.loads(CONFIG_WITH_LINKAGE))
     assert len(config.contract.datasets) == 6
-    dc_metadata = config.get_contract_metadata()
-    assert len(dc_metadata.linkage_hierarchy) == 2
-    assert not dc_metadata.linkage_hierarchy.get("ds_002").children
-    assert len(dc_metadata.linkage_hierarchy.get("ds_001").get_descendents()) == 4
-    children_001 = sorted(dc_metadata.linkage_hierarchy.get("ds_001").children, key=lambda x: x.entity_name)
-    dict_rep_001 = dc_metadata.linkage_hierarchy.get("ds_001").as_dict()
+    hierarchy = config.get_entity_hierarchy()
+    assert len(hierarchy) == 2
+    assert not hierarchy.get("ds_002").children
+    assert len(hierarchy.get("ds_001").get_descendents()) == 4
+    children_001 = sorted(hierarchy.get("ds_001").children, key=lambda x: x.entity_name)
+    dict_rep_001 = hierarchy.get("ds_001").as_dict()
     assert len(children_001) == 2
     assert children_001[0].entity_name == "ds_003"
     assert not children_001[0].children
     assert children_001[1].entity_name == "ds_101"
-    assert dict_rep_001 == {'ds_001': {'children': {'ds_003': {'join_fields': ['ds_001_id']}, 'ds_101': {'join_fields': ['ds_001_id'], 'children': {'ds_201': {'join_fields': ['referral_id'], 'children': {'ds_202': {'join_fields': ['ds_201_id']}}}}}}}}
+    assert dict_rep_001 == {'ds_001': {'children': {'ds_003': {'children': {}, 'mandatory': False, 'join_fields': ['ds_001_id']}, 'ds_101': {'children': {'ds_201': {'children': {'ds_202': {'children': {}, 'mandatory': True, 'join_fields': ['ds_201_id']}}, 'mandatory': False, 'join_fields': ['referral_id']}}, 'mandatory': False, 'join_fields': ['ds_001_id']}}, 'mandatory': False}}
     children_101 = children_001[1].children
     dict_rep_101 = dict_rep_001["ds_001"]["children"]["ds_101"]
     assert len(children_101) == 1
     assert children_101[0].entity_name == "ds_201"
     assert children_101[0].children[0].entity_name == "ds_202"
     assert not children_101[0].children[0].children
-    assert dict_rep_101 == {'join_fields': ['ds_001_id'], 'children': {'ds_201': {'join_fields': ['referral_id'], 'children': {'ds_202': {'join_fields': ['ds_201_id']}}}}}
+    assert dict_rep_101 == {'children': {'ds_201': {'children': {'ds_202': {'children': {}, 'mandatory': True, 'join_fields': ['ds_201_id']}}, 'mandatory': False, 'join_fields': ['referral_id']}}, 'mandatory': False, 'join_fields': ['ds_001_id']}