fix: initial draft for storing entity hierarchy

Author: stevenhsd

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

@@ -1,13 +1,61 @@
 """Metadata classes for the data contract."""
 
-from typing import Any
+from typing import Any, Optional, Union
 
-from pydantic import BaseModel, PrivateAttr, model_validator
+from pydantic import BaseModel, Field, 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."""
@@ -38,6 +86,7 @@ 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

@@ -1,13 +1,13 @@
 """The loader for the first JSON-based dataset configuration."""
 
 import json
-from typing import Any, Optional, Union
+from typing import Any, Optional, Type, Union
 
 from pydantic import BaseModel, Field, PrivateAttr, validate_call
 from typing_extensions import Literal
 
 from dve.core_engine.backends.base.reference_data import ReferenceConfig, ReferenceConfigUnion
-from dve.core_engine.backends.metadata.contract import DataContractMetadata, ReaderConfig
+from dve.core_engine.backends.metadata.contract import ChildHierarchyNode, DataContractMetadata, HierarchyNode, 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 (
@@ -24,6 +24,7 @@
 from dve.core_engine.message import DataContractErrorDetail
 from dve.core_engine.type_hints import EntityName, ErrorCategory, 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
 
@@ -38,6 +39,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"""
 TypeOrDef = Union[  # pylint: disable=C0103
     TypeName, "_CallableTypeDefinition", "_ModelTypeDefinition", "_TypeAliasDefinition"
 ]
@@ -81,6 +84,14 @@ class _TypeAliasDefinition(_BaseTypeDefintion):
     """The name of the Python type."""
 
 
+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"""
+
+
 class _SchemaConfig(BaseModel):
     """Configuration for a component schema within a dataset."""
 
@@ -112,6 +123,7 @@ 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):
@@ -330,6 +342,7 @@ def get_contract_metadata(self) -> DataContractMetadata:
             validators=validators,
             reporting_fields=reporting_fields,
             cache_originals=self.contract.cache_originals,
+            linkage_hierarchy=self.determine_entity_hierarchy()
         )
 
     def load_error_message_info(self, uri):
@@ -351,3 +364,22 @@ def get_rule_metadata(self) -> RuleMetadata:
             global_variables=self.transformations.parameters,  # pylint: disable=E1101
             reference_data_config=self.get_reference_data_config(),
         )
+
+    def determine_entity_hierarchy(self) -> list[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
+            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))
+                    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/metadata_parser/exc.py

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