|
| 1 | +"""Classes to help determine and store entity hierarchy information.""" |
| 2 | + |
| 3 | +import json |
| 4 | +from typing import Any, Iterable, Optional, Union |
| 5 | + |
| 6 | +from pydantic import BaseModel, Field |
| 7 | + |
| 8 | +from dve.core_engine.configuration.v1 import V1EngineConfig, _LinkageConfig |
| 9 | +from dve.core_engine.type_hints import EntityName, ErrorCode, ErrorMessage |
| 10 | +from dve.metadata_parser.exc import EntityNotFoundError |
| 11 | +from dve.parser.file_handling.service import open_stream |
| 12 | +from dve.parser.type_hints import URI |
| 13 | + |
| 14 | + |
| 15 | +class HierarchyNode(BaseModel): |
| 16 | + """Stores entity hierarchy information""" |
| 17 | + |
| 18 | + entity_name: str |
| 19 | + children: list["HierarchyNode"] = Field(default_factory=list) |
| 20 | + |
| 21 | + def get_descendents(self) -> list[str]: |
| 22 | + """Recursively list all descendents of the node""" |
| 23 | + descendents = [] |
| 24 | + for node in self.children: |
| 25 | + descendents.append(node.entity_name) |
| 26 | + descendents.extend(node.get_descendents()) |
| 27 | + return descendents |
| 28 | + |
| 29 | + def get_node(self, entity_name: str) -> Union["HierarchyNode", None]: |
| 30 | + """Recursively search for node and return if found""" |
| 31 | + node = None |
| 32 | + if self.entity_name == entity_name: |
| 33 | + return self |
| 34 | + for child in self.children: |
| 35 | + node = child.get_node(entity_name) |
| 36 | + if node: |
| 37 | + break |
| 38 | + return node |
| 39 | + |
| 40 | + def add_child_node(self, parent_entity: str, child_info: "HierarchyNode") -> None: |
| 41 | + """Add a child node if the parent exists in the hierarchy""" |
| 42 | + try: |
| 43 | + self.get_node(parent_entity).children.append(child_info) # type: ignore |
| 44 | + except AttributeError as exc: |
| 45 | + raise EntityNotFoundError( |
| 46 | + f"Can't find parent node {parent_entity} in {self.entity_name}" |
| 47 | + ) from exc |
| 48 | + |
| 49 | + def as_dict(self) -> dict[str, dict[str, Any]]: |
| 50 | + """Get dictionary representation of entity hierarchy""" |
| 51 | + child_dict = {} |
| 52 | + for node in self.children: |
| 53 | + child_dict.update(node.as_dict()) |
| 54 | + |
| 55 | + ret_dict = self.model_dump(exclude={"entity_name", "children"}) |
| 56 | + ret_dict.update({"children": child_dict}) |
| 57 | + |
| 58 | + return {self.entity_name: ret_dict} |
| 59 | + |
| 60 | + |
| 61 | +class ChildHierarchyNode(HierarchyNode): |
| 62 | + """Stores child entity hierarchy information""" |
| 63 | + |
| 64 | + join_fields: dict[str, str] |
| 65 | + mandatory: Optional[bool] = False |
| 66 | + no_valid_records_error_code: Optional[ErrorCode] = "NoValidRecords" |
| 67 | + no_valid_records_error_message: Optional[ErrorMessage] = ( |
| 68 | + "parent record removed as no valid child records" |
| 69 | + ) |
| 70 | + orphaned_records_error_code: Optional[ErrorCode] = "OrphanedRecords" |
| 71 | + orphaned_records_error_message: Optional[ErrorMessage] = "Orphaned records removed" |
| 72 | + |
| 73 | + |
| 74 | +class EntityHierarchy: |
| 75 | + """Determines and stores entity hierarchy information from config""" |
| 76 | + |
| 77 | + def __init__(self, entity_trees: dict[EntityName, HierarchyNode]): |
| 78 | + self.entity_trees = entity_trees |
| 79 | + |
| 80 | + @staticmethod |
| 81 | + def determine_trees( |
| 82 | + all_datasets: Iterable[str], entity_relationships: dict[str, _LinkageConfig] |
| 83 | + ) -> dict[EntityName, HierarchyNode]: |
| 84 | + """Determine the entity hierarchy trees and store as HierarchyNodes""" |
| 85 | + top_level_parents: dict[EntityName, HierarchyNode] = { |
| 86 | + entity_name: HierarchyNode(entity_name=entity_name) |
| 87 | + for entity_name in all_datasets |
| 88 | + if entity_name not in entity_relationships |
| 89 | + } |
| 90 | + |
| 91 | + for name, linkage_detail in entity_relationships.items(): |
| 92 | + for main_entity, parent_node in top_level_parents.items(): |
| 93 | + if ( |
| 94 | + linkage_detail.parent_entity == main_entity |
| 95 | + or linkage_detail.parent_entity in parent_node.get_descendents() |
| 96 | + ): |
| 97 | + parent_node.add_child_node( |
| 98 | + linkage_detail.parent_entity, |
| 99 | + ChildHierarchyNode( |
| 100 | + entity_name=name, **linkage_detail.model_dump(exclude={"parent_entity"}) |
| 101 | + ), |
| 102 | + ) |
| 103 | + break |
| 104 | + else: |
| 105 | + raise EntityNotFoundError( |
| 106 | + f"Can't find parent entity {linkage_detail.parent_entity} defined to " |
| 107 | + + f"establish hierarchy for {name} - please ensure it is defined above " |
| 108 | + + "any child entities in the dischema." |
| 109 | + ) |
| 110 | + return top_level_parents |
| 111 | + |
| 112 | + @classmethod |
| 113 | + def from_dischema(cls, dischema_uri: URI): |
| 114 | + """Create entity hierarchy direct from dischema""" |
| 115 | + with open_stream(dischema_uri) as dischema: |
| 116 | + config_dict = json.load(dischema) |
| 117 | + all_datasets = config_dict.get("contract", {}).get("datasets", {}).keys() |
| 118 | + entity_relationships = { |
| 119 | + k: _LinkageConfig(**v) for k, v in config_dict.get("entity_relationships", {}).items() |
| 120 | + } |
| 121 | + return cls(entity_trees=cls.determine_trees(all_datasets, entity_relationships)) |
| 122 | + |
| 123 | + @classmethod |
| 124 | + def from_engine_config(cls, engine_config: V1EngineConfig): |
| 125 | + """Create entity hierarchy direct from engine config""" |
| 126 | + return cls( |
| 127 | + entity_trees=cls.determine_trees( |
| 128 | + all_datasets=engine_config.contract.datasets.keys(), |
| 129 | + entity_relationships=engine_config.entity_relationships, |
| 130 | + ) |
| 131 | + ) |
0 commit comments