refactor: unify HierarchyNode and ChildHierarchyNode also tweaked background writer queue type and use in orphaned records implementation

Author: stevenhsd

docs/advanced_guidance/json_schemas/entity_relationships.schema.json

@@ -20,10 +20,10 @@
                 "mandatory": {
                     "type": "boolean"
                 },
-                "orphaned_records_error_code": {
+                "missing_parent_id_error_code": {
                     "type": "string"
                 },
-                "orphaned_records_error_message": {
+                "missing_parent_id_error_message": {
                     "type": "string"
                 }
             },

src/dve/common/error_utils.py

@@ -5,7 +5,7 @@
 import logging
 from collections.abc import Iterable
 from itertools import chain
-from multiprocessing import Queue
+from queue import Queue
 from threading import Thread
 from typing import Optional, Union
 

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

@@ -46,7 +46,6 @@
 )
 from dve.core_engine.backends.types import Entities, EntityType, StageSuccessful
 from dve.core_engine.configuration.v1.hierarchy import (
-    ChildHierarchyNode,
     EntityHierarchy,
     HierarchyNode,
 )
@@ -390,7 +389,7 @@ def identify_and_remove_orphans(
         """
 
         def process_node(
-            node: HierarchyNode | ChildHierarchyNode,
+            node: HierarchyNode,
             parent_entity_name: Optional[EntityName],
             orph_messages: Messages | None = None,
         ):
@@ -400,7 +399,7 @@ def process_node(
             if orph_messages is None:
                 orph_messages = []
 
-            if isinstance(node, ChildHierarchyNode) and parent_entity_name is not None:
+            if parent_entity_name is not None:
                 self.logger.info(f"Identifying orphans in {current_entity_name}")
 
                 join_expr = " AND ".join(
@@ -419,7 +418,7 @@ def process_node(
                 )
 
                 if no_orphs > 0:
-                    self.logger.info(f"Removing orphan records from {current_entity_name}")
+                    self.logger.info(f"Removing records with missing parent from {current_entity_name}")
                     location = list(node.join_fields.values())[0]
                     with BackgroundMessageWriter(
                         working_directory=working_directory,
@@ -433,32 +432,32 @@ def process_node(
                                 entity_name=current_entity_name,
                                 reporting=ReportingConfig(
                                     emit="record_failure",
-                                    code=node.orphaned_records_error_code,
-                                    message=node.orphaned_records_error_message,
+                                    code=node.missing_parent_id_error_code,
+                                    message=node.missing_parent_id_error_message,
                                     location=location,
                                 )
                             )
                         )
-                        for record in _orph_records:
-                            msg_writer.write_queue.put([
-                                FeedbackMessage(
-                                    entity=current_entity_name,
-                                    record=record,  # type: ignore
-                                    error_location=location,
-                                    error_message=node.orphaned_records_error_message,
-                                    failure_type="record",
-                                    error_type="record",
-                                    error_code=node.orphaned_records_error_code,
-                                    reporting_field=location,
-                                    category="Parent Missing",
-                                )
-                            ])
+                    # moved to batch the write - risky if large number of
+                        msg_writer.write_queue.put([
+                            FeedbackMessage(
+                                entity=current_entity_name,
+                                record=record,  # type: ignore
+                                error_location=location,
+                                error_message=node.missing_parent_id_error_message,
+                                failure_type="record",
+                                error_type="record",
+                                error_code=node.missing_parent_id_error_code,
+                                reporting_field=location,
+                                category="Parent Missing",
+                            )
+                        for record in _orph_records ])
 
             if node.children:
                 for child_node in node.children:
                     process_node(child_node, current_entity_name, orph_messages)
 
-
+        # would a root ever be orphaned?
         for root_node in entity_hierarchy.entity_trees.values():
             process_node(root_node, parent_entity_name=None)
 

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

@@ -3,7 +3,8 @@
 import json
 from typing import Any, Optional, Type, Union
 
-from pydantic import BaseModel, Field, PrivateAttr, validate_call
+from pydantic import BaseModel, Field, PrivateAttr, field_validator, model_validator, validate_call
+from pydantic_core.core_schema import FieldValidationInfo
 from typing_extensions import Literal
 
 from dve.core_engine.backends.base.reference_data import ReferenceConfig, ReferenceConfigUnion
@@ -93,22 +94,47 @@ class _TypeAliasDefinition(_BaseTypeDefintion):
 class _LinkageConfig(BaseModel):
     """Specify how to link entities back to parents if required"""
 
-    parent_entity: EntityName
+    parent_entity: Optional[EntityName] = None
     """The name of the parent entity"""
-    join_fields: JoinFields
+    join_fields: JoinFields = Field(default_factory=dict)
     """The fields that can be used to link back to the parent entity"""
-    mandatory: Optional[bool] = False
+    is_root_entity: bool = False
+    """Whether the entity is the highest level parent in a tree"""
+    mandatory: 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"
+    missing_parent_id_error_code: Optional[ErrorCode] = "MissingParentRecord"
     """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"
+    missing_parent_id_error_message: Optional[ErrorMessage] = (
+        "Records removed due to no valid parent record"
+    )
     """The error code to emit if the entity contains records that are orphaned by parent record rejections"""  # pylint: disable=C0301
+       
+    @model_validator(mode="after")
+    def _check_root_no_parent_or_join_keys(self):
+        if self.is_root_entity:
+            if self.parent_entity or self.join_fields:
+                raise ValueError("If entity is root, neither parent_entity nor join keys should be specified")
+        return self
+    
+    @model_validator(mode="after")
+    def _check_root_mandatory(self):
+        if self.is_root_entity:
+            if not self.mandatory:
+                raise ValueError("If entity is root, it must be labelled mandatory")
+        return self
+    
+    @model_validator(mode="after")
+    def _check_parent_entity_with_join_keys(self):
+        if self.parent_entity or self.join_fields:
+            if not (self.parent_entity and self.join_fields):
+                raise ValueError("Both parent_entity and join_fields must be supplied if one is")
+        return self
 
 
 class _SchemaConfig(BaseModel):

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

@@ -3,7 +3,7 @@
 import json
 from typing import Any, Iterable, Optional, Union
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, field_validator
 
 from dve.core_engine.configuration.v1 import V1EngineConfig, _LinkageConfig
 from dve.core_engine.type_hints import EntityName, ErrorCode, ErrorMessage
@@ -16,8 +16,18 @@ class HierarchyNode(BaseModel):
     """Stores entity hierarchy information"""
 
     entity_name: str
-    children: Optional[list["ChildHierarchyNode"]] = Field(default_factory=list)
-
+    children: Optional[list["HierarchyNode"]] = Field(default_factory=list)
+    mandatory: Optional[bool] = False
+    join_fields: Optional[dict[str, str]] = Field(default_factory=dict)
+    no_valid_records_error_code: Optional[ErrorCode] = "NoValidRecords"
+    no_valid_records_error_message: Optional[ErrorMessage] = (
+        "parent record removed as no valid child records"
+    )
+    missing_parent_id_error_code: Optional[ErrorCode] = "MissingParentRecord"
+    missing_parent_id_error_message: Optional[ErrorMessage] = (
+        "Records removed due to no valid parent record"
+    )
+    
     def get_descendents(self) -> list[str]:
         """Recursively list all descendents of the node"""
         descendents = []
@@ -58,19 +68,6 @@ def as_dict(self) -> dict[str, dict[str, Any]]:
         return {self.entity_name: ret_dict}
 
 
-class ChildHierarchyNode(HierarchyNode):
-    """Stores child entity hierarchy information"""
-
-    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"""
 
@@ -82,11 +79,24 @@ def determine_trees(
         all_datasets: Iterable[str], entity_relationships: dict[str, _LinkageConfig]
     ) -> dict[EntityName, HierarchyNode]:
         """Determine the entity hierarchy trees and store as HierarchyNodes"""
+        root_entities: dict[str, _LinkageConfig] = dict(filter(lambda x: x[1].is_root_entity,
+                                                        entity_relationships.items()))
         top_level_parents: dict[EntityName, HierarchyNode] = {
-            entity_name: HierarchyNode(entity_name=entity_name)
-            for entity_name in all_datasets
-            if entity_name not in entity_relationships
+            entity_name: HierarchyNode(entity_name=entity_name,
+                                       **config.model_dump(exclude={"parent_entity",
+                                                                    "missing_parent_id_error_code",
+                                                                    "missing_parent_id_error_message"}),
+                                       missing_parent_id_error_code=None,
+                                       missing_parent_id_error_message=None)
+            for entity_name, config in root_entities.items()
         }
+        
+        if default_roots := [ entity_name for entity_name in all_datasets 
+                             if not entity_name in entity_relationships]:
+            for entity_name in default_roots:
+                top_level_parents[entity_name] = HierarchyNode(entity_name=entity_name,
+                                                               missing_parent_id_error_code=None,
+                                                               missing_parent_id_error_message=None)
 
         for name, linkage_detail in entity_relationships.items():
             for main_entity, parent_node in top_level_parents.items():
@@ -96,7 +106,7 @@ def determine_trees(
                 ):
                     parent_node.add_child_node(
                         linkage_detail.parent_entity,
-                        ChildHierarchyNode(
+                        HierarchyNode(
                             entity_name=name, **linkage_detail.model_dump(exclude={"parent_entity"})
                         ),
                     )

src/dve/core_engine/constants.py

@@ -7,5 +7,5 @@
 """The name of the field that can be used to extract the field value that caused
    a pydantic validation error"""
 
-ORPHANED_RECORD_ENTITY_NAME: str = "orphaned_records_tracker"
-"""Name to keep track of identified orphaned records"""
+ORPHANED_RECORD_ENTITY_NAME: str = "orphaned_record_tracker"
+"""Name of entity to keep track of records where there is a missing parent record"""

tests/test_core_engine/test_backends/test_implementations/test_duckdb/test_rules.py

@@ -40,8 +40,9 @@
     SemiJoin,
     TableUnion,
 )
+from dve.core_engine.constants import ORPHANED_RECORD_ENTITY_NAME
 from dve.core_engine.configuration.v1.hierarchy import (
-    ChildHierarchyNode, EntityHierarchy, HierarchyNode
+    EntityHierarchy, HierarchyNode
 )
 from dve.core_engine.type_hints import MultipleExpressions
 from tests.test_core_engine.test_backends.fixtures import (
@@ -630,7 +631,7 @@ def test_identify_orphan_record_single_entity(self):
                     join_condition="passengers.flight_id = flights.flight_id"
                 )
             )
-            result = mod_entities["orphaned_records_tracker"]
+            result = mod_entities[ORPHANED_RECORD_ENTITY_NAME]
             assert result.count("*").fetchone()[0] == 1  # type: ignore
             assert result.select("entity_name").unique("*").count("*").fetchone()[0] == 1  # type: ignore
 
@@ -653,10 +654,10 @@ def test_identify_and_remove_orphans(self):
                 "flights": HierarchyNode(
                     entity_name="flights",
                     children=[
-                        ChildHierarchyNode(
+                        HierarchyNode(
                             entity_name="passengers",
                             children=[
-                                ChildHierarchyNode(
+                                HierarchyNode(
                                     entity_name="food",
                                     children=[],
                                     join_fields={"passenger_id": "passenger_id"},