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

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

@@ -45,11 +45,7 @@
     TableUnion,
 )
 from dve.core_engine.backends.types import Entities, EntityType, StageSuccessful
-from dve.core_engine.configuration.v1.hierarchy import (
-    ChildHierarchyNode,
-    EntityHierarchy,
-    HierarchyNode,
-)
+from dve.core_engine.configuration.v1.hierarchy import EntityHierarchy, HierarchyNode
 from dve.core_engine.constants import ORPHANED_RECORD_ENTITY_NAME
 from dve.core_engine.exceptions import CriticalProcessingError
 from dve.core_engine.loggers import get_logger
@@ -390,7 +386,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 +396,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 +415,9 @@ 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 +431,34 @@ 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([
+                        # 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.orphaned_records_error_message,
+                                    error_message=node.missing_parent_id_error_message,
                                     failure_type="record",
                                     error_type="record",
-                                    error_code=node.orphaned_records_error_code,
+                                    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)
 
-
         for root_node in entity_hierarchy.entity_trees.values():
             process_node(root_node, parent_entity_name=None)
 

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

@@ -553,6 +553,7 @@ class OrphanIdentification(AbstractConditionalJoin):
 
     """
 
+
 Step = Union[AbstractStep, Literal["sync"]]
 """A step within a rule. This is either a rule config or the literal string 'sync'."""
 

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,23 +94,48 @@ 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 and (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 and 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):
     """Configuration for a component schema within a dataset."""

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

@@ -16,7 +16,15 @@ class HierarchyNode(BaseModel):
     """Stores entity hierarchy information"""
 
     entity_name: str
-    children: Optional[list["ChildHierarchyNode"]] = Field(default_factory=list)
+    children: list["HierarchyNode"] = Field(default_factory=list)
+    mandatory: bool = False
+    join_fields: dict[str, str] = Field(default_factory=dict)
+    no_valid_records_error_code: ErrorCode = "NoValidRecords"
+    no_valid_records_error_message: 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"""
@@ -58,19 +66,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,12 +77,35 @@ 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 entity_name not 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():
                 if (
@@ -96,7 +114,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"},