handle migrated kind nodes in diff and merge logic (#6401)

* fix bug in rel create for migrated kind nodes

* handle agnostic rels correctly

* WIP on including migrated kinds in diff

* small changes to diff calculation queries for properties of nodes with
migrated kind/inheritance

* expand unit test for migrated kind diff... big time

* finish up giant unit test

* use db_id instead of labels, will need it for merge updates

* move is_node_kind_migration outside of existing diff queries

* format

* fix broken tests/query

* update DiffCombiner with identifiers

* undo some changes in DiffQueryParser that are no longer needed

* comment placement fix

* test and support for merging branch with migrated kind node

* update unit test for is_node_kind_migration

* add changelog

* remove commented out code

* handle migrated node that is later deleted

* add cypher comment

* remove unused property

* update comment, remove unnecessary int() casting

* linting

* one more test

* add migration to delete diffs b/c they must be recalculated

* actually include the migration

Author: ajtmccarty

backend/infrahub/core/diff/calculator.py

@@ -7,6 +7,7 @@
 from infrahub.core.query.diff import (
     DiffCalculationQuery,
     DiffFieldPathsQuery,
+    DiffMigratedKindNodesQuery,
     DiffNodePathsQuery,
     DiffPropertyPathsQuery,
 )
@@ -15,7 +16,7 @@
 from infrahub.log import get_logger
 
 from .model.field_specifiers_map import NodeFieldSpecifierMap
-from .model.path import CalculatedDiffs
+from .model.path import CalculatedDiffs, DiffNode, DiffRoot, NodeIdentifier
 
 log = get_logger()
 
@@ -59,7 +60,7 @@ async def _run_diff_calculation_query(
             )
             log.info(f"Beginning one diff calculation query {limit=}, {offset=}")
             await diff_query.execute(db=self.db)
-            log.info("Diff calculation query complete")
+            log.info(f"Diff calculation query complete {limit=}, {offset=}")
             last_result = None
             for query_result in diff_query.get_results():
                 diff_parser.read_result(query_result=query_result)
@@ -69,6 +70,56 @@ async def _run_diff_calculation_query(
                 has_more_data = last_result.get_as_type("has_more_data", bool)
             offset += limit
 
+    async def _apply_kind_migrated_nodes(
+        self, branch_diff: DiffRoot, calculation_request: DiffCalculationRequest
+    ) -> None:
+        has_more_data = True
+        offset = 0
+        limit = config.SETTINGS.database.query_size_limit
+        diff_nodes_by_identifier = {n.identifier: n for n in branch_diff.nodes}
+        diff_nodes_to_add: list[DiffNode] = []
+        while has_more_data:
+            diff_query = await DiffMigratedKindNodesQuery.init(
+                db=self.db,
+                branch=calculation_request.diff_branch,
+                base_branch=calculation_request.base_branch,
+                diff_branch_from_time=calculation_request.branch_from_time,
+                diff_from=calculation_request.from_time,
+                diff_to=calculation_request.to_time,
+                limit=limit,
+                offset=offset,
+            )
+            log.info(f"Getting one batch of migrated kind nodes {limit=}, {offset=}")
+            await diff_query.execute(db=self.db)
+            log.info(f"Migrated kind nodes query complete {limit=}, {offset=}")
+            last_result = None
+            for migrated_kind_node in diff_query.get_migrated_kind_nodes():
+                migrated_kind_identifier = NodeIdentifier(
+                    uuid=migrated_kind_node.uuid,
+                    kind=migrated_kind_node.kind,
+                    db_id=migrated_kind_node.db_id,
+                )
+                if migrated_kind_identifier in diff_nodes_by_identifier:
+                    diff_node = diff_nodes_by_identifier[migrated_kind_identifier]
+                    diff_node.is_node_kind_migration = True
+                    continue
+                new_diff_node = DiffNode(
+                    identifier=migrated_kind_identifier,
+                    changed_at=migrated_kind_node.from_time,
+                    action=migrated_kind_node.action,
+                    is_node_kind_migration=True,
+                    attributes=[],
+                    relationships=[],
+                )
+                diff_nodes_by_identifier[migrated_kind_identifier] = new_diff_node
+                diff_nodes_to_add.append(new_diff_node)
+                last_result = migrated_kind_node
+            has_more_data = False
+            if last_result:
+                has_more_data = last_result.has_more_data
+            offset += limit
+        branch_diff.nodes.extend(diff_nodes_to_add)
+
     async def calculate_diff(
         self,
         base_branch: Branch,
@@ -92,7 +143,7 @@ async def calculate_diff(
         )
         node_limit = int(config.SETTINGS.database.query_size_limit / 10)
         fields_limit = int(config.SETTINGS.database.query_size_limit / 3)
-        properties_limit = int(config.SETTINGS.database.query_size_limit)
+        properties_limit = config.SETTINGS.database.query_size_limit
 
         calculation_request = DiffCalculationRequest(
             base_branch=base_branch,
@@ -132,7 +183,7 @@ async def calculate_diff(
         if base_branch.name != diff_branch.name:
             current_node_field_specifiers = diff_parser.get_current_node_field_specifiers()
             new_node_field_specifiers = diff_parser.get_new_node_field_specifiers()
-            calculation_request = DiffCalculationRequest(
+            base_calculation_request = DiffCalculationRequest(
                 base_branch=base_branch,
                 diff_branch=base_branch,
                 branch_from_time=diff_branch_from_time,
@@ -146,7 +197,7 @@ async def calculate_diff(
             await self._run_diff_calculation_query(
                 diff_parser=diff_parser,
                 query_class=DiffNodePathsQuery,
-                calculation_request=calculation_request,
+                calculation_request=base_calculation_request,
                 limit=node_limit,
             )
             log.info("Diff node-level calculation queries for base complete")
@@ -155,7 +206,7 @@ async def calculate_diff(
             await self._run_diff_calculation_query(
                 diff_parser=diff_parser,
                 query_class=DiffFieldPathsQuery,
-                calculation_request=calculation_request,
+                calculation_request=base_calculation_request,
                 limit=fields_limit,
             )
             log.info("Diff field-level calculation queries for base complete")
@@ -164,17 +215,19 @@ async def calculate_diff(
             await self._run_diff_calculation_query(
                 diff_parser=diff_parser,
                 query_class=DiffPropertyPathsQuery,
-                calculation_request=calculation_request,
+                calculation_request=base_calculation_request,
                 limit=properties_limit,
             )
             log.info("Diff property-level calculation queries for base complete")
 
         log.info("Parsing calculated diff")
         diff_parser.parse(include_unchanged=include_unchanged)
         log.info("Calculated diff parsed")
+        branch_diff = diff_parser.get_diff_root_for_branch(branch=diff_branch.name)
+        await self._apply_kind_migrated_nodes(branch_diff=branch_diff, calculation_request=calculation_request)
         return CalculatedDiffs(
             base_branch_name=base_branch.name,
             diff_branch_name=diff_branch.name,
             base_branch_diff=diff_parser.get_diff_root_for_branch(branch=base_branch.name),
-            diff_branch_diff=diff_parser.get_diff_root_for_branch(branch=diff_branch.name),
+            diff_branch_diff=branch_diff,
         )

backend/infrahub/core/diff/combiner.py

@@ -14,6 +14,7 @@
     EnrichedDiffRoot,
     EnrichedDiffs,
     EnrichedDiffSingleRelationship,
+    NodeIdentifier,
 )
 
 
@@ -26,30 +27,35 @@ class NodePair:
 class DiffCombiner:
     def __init__(self) -> None:
         # {child_uuid: (parent_uuid, parent_rel_name)}
-        self._child_parent_uuid_map: dict[str, tuple[str, str]] = {}
-        self._parent_node_uuids: set[str] = set()
-        self._earlier_nodes_by_uuid: dict[str, EnrichedDiffNode] = {}
-        self._later_nodes_by_uuid: dict[str, EnrichedDiffNode] = {}
-        self._common_node_uuids: set[str] = set()
+        self._child_parent_identifier_map: dict[NodeIdentifier, tuple[NodeIdentifier, str]] = {}
+        self._parent_node_identifiers: set[NodeIdentifier] = set()
+        self._earlier_nodes_by_identifier: dict[NodeIdentifier, EnrichedDiffNode] = {}
+        self._later_nodes_by_identifier: dict[NodeIdentifier, EnrichedDiffNode] = {}
+        self._common_node_identifiers: set[NodeIdentifier] = set()
         self._diff_branch_name: str | None = None
 
     def _initialize(self, earlier_diff: EnrichedDiffRoot, later_diff: EnrichedDiffRoot) -> None:
         self._diff_branch_name = earlier_diff.diff_branch_name
-        self._child_parent_uuid_map = {}
-        self._earlier_nodes_by_uuid = {}
-        self._later_nodes_by_uuid = {}
-        self._common_node_uuids = set()
+        self._child_parent_identifier_map = {}
+        self._earlier_nodes_by_identifier = {}
+        self._later_nodes_by_identifier = {}
+        self._common_node_identifiers = set()
         # map the parent of each node (if it exists), preference to the later diff
         for diff_root in (earlier_diff, later_diff):
             for child_node in diff_root.nodes:
                 for parent_rel in child_node.relationships:
                     for parent_node in parent_rel.nodes:
-                        self._child_parent_uuid_map[child_node.uuid] = (parent_node.uuid, parent_rel.name)
+                        self._child_parent_identifier_map[child_node.identifier] = (
+                            parent_node.identifier,
+                            parent_rel.name,
+                        )
         # UUIDs of all the parents, removing the stale parents from the earlier diff
-        self._parent_node_uuids = {parent_tuple[0] for parent_tuple in self._child_parent_uuid_map.values()}
-        self._earlier_nodes_by_uuid = {n.uuid: n for n in earlier_diff.nodes}
-        self._later_nodes_by_uuid = {n.uuid: n for n in later_diff.nodes}
-        self._common_node_uuids = set(self._earlier_nodes_by_uuid.keys()) & set(self._later_nodes_by_uuid.keys())
+        self._parent_node_identifiers = {parent_tuple[0] for parent_tuple in self._child_parent_identifier_map.values()}
+        self._earlier_nodes_by_identifier = {n.identifier: n for n in earlier_diff.nodes}
+        self._later_nodes_by_identifier = {n.identifier: n for n in later_diff.nodes}
+        self._common_node_identifiers = set(self._earlier_nodes_by_identifier.keys()) & set(
+            self._later_nodes_by_identifier.keys()
+        )
 
     @property
     def diff_branch_name(self) -> str:
@@ -61,13 +67,13 @@ def _filter_nodes_to_keep(self, earlier_diff: EnrichedDiffRoot, later_diff: Enri
         filtered_node_pairs: list[NodePair] = []
         for earlier_node in earlier_diff.nodes:
             later_node: EnrichedDiffNode | None = None
-            if earlier_node.uuid in self._common_node_uuids:
-                later_node = self._later_nodes_by_uuid[earlier_node.uuid]
+            if earlier_node.identifier in self._common_node_identifiers:
+                later_node = self._later_nodes_by_identifier[earlier_node.identifier]
             # this is an out-of-date parent
             if (
                 earlier_node.action is DiffAction.UNCHANGED
                 and (later_node is None or later_node.action is DiffAction.UNCHANGED)
-                and earlier_node.uuid not in self._parent_node_uuids
+                and earlier_node.identifier not in self._parent_node_identifiers
             ):
                 continue
             if later_node is None:
@@ -79,15 +85,15 @@ def _filter_nodes_to_keep(self, earlier_diff: EnrichedDiffRoot, later_diff: Enri
             filtered_node_pairs.append(NodePair(earlier=earlier_node, later=later_node))
         for later_node in later_diff.nodes:
             # these have already been handled
-            if later_node.uuid in self._common_node_uuids:
+            if later_node.identifier in self._common_node_identifiers:
                 continue
             filtered_node_pairs.append(NodePair(later=later_node))
         return filtered_node_pairs
 
-    def _get_parent_relationship_name(self, node_id: str) -> str | None:
-        if node_id not in self._child_parent_uuid_map:
+    def _get_parent_relationship_name(self, node_id: NodeIdentifier) -> str | None:
+        if node_id not in self._child_parent_identifier_map:
             return None
-        return self._child_parent_uuid_map[node_id][1]
+        return self._child_parent_identifier_map[node_id][1]
 
     def _should_include(self, earlier: DiffAction, later: DiffAction) -> bool:
         actions = {earlier, later}
@@ -284,7 +290,7 @@ def _combine_relationships(
         self,
         earlier_relationships: set[EnrichedDiffRelationship],
         later_relationships: set[EnrichedDiffRelationship],
-        node_id: str,
+        node_id: NodeIdentifier,
     ) -> set[EnrichedDiffRelationship]:
         earlier_rels_by_name = {rel.name: rel for rel in earlier_relationships}
         later_rels_by_name = {rel.name: rel for rel in later_relationships}
@@ -365,7 +371,7 @@ def _combine_nodes(self, node_pairs: list[NodePair]) -> set[EnrichedDiffNode]:
             combined_relationships = self._combine_relationships(
                 earlier_relationships=node_pair.earlier.relationships,
                 later_relationships=node_pair.later.relationships,
-                node_id=node_pair.later.uuid,
+                node_id=node_pair.later.identifier,
             )
             if all(ca.action is DiffAction.UNCHANGED for ca in combined_attributes) and all(
                 cr.action is DiffAction.UNCHANGED for cr in combined_relationships
@@ -380,14 +386,16 @@ def _combine_nodes(self, node_pairs: list[NodePair]) -> set[EnrichedDiffNode]:
                 combined_attributes
                 or combined_relationships
                 or combined_conflict
-                or node_pair.later.uuid in self._parent_node_uuids
+                or node_pair.later.identifier in self._parent_node_identifiers
             ):
                 combined_nodes.add(
                     EnrichedDiffNode(
                         identifier=node_pair.later.identifier,
                         label=node_pair.later.label,
                         changed_at=node_pair.later.changed_at or node_pair.earlier.changed_at,
                         action=combined_action,
+                        is_node_kind_migration=node_pair.earlier.is_node_kind_migration
+                        or node_pair.later.is_node_kind_migration,
                         path_identifier=node_pair.later.path_identifier,
                         attributes=combined_attributes,
                         relationships=combined_relationships,
@@ -397,12 +405,12 @@ def _combine_nodes(self, node_pairs: list[NodePair]) -> set[EnrichedDiffNode]:
         return combined_nodes
 
     def _link_child_nodes(self, nodes: Iterable[EnrichedDiffNode]) -> None:
-        nodes_by_uuid: dict[str, EnrichedDiffNode] = {n.uuid: n for n in nodes}
-        for child_node in nodes_by_uuid.values():
-            if child_node.uuid not in self._child_parent_uuid_map:
+        nodes_by_identifier: dict[NodeIdentifier, EnrichedDiffNode] = {n.identifier: n for n in nodes}
+        for child_node in nodes_by_identifier.values():
+            if child_node.identifier not in self._child_parent_identifier_map:
                 continue
-            parent_uuid, parent_rel_name = self._child_parent_uuid_map[child_node.uuid]
-            parent_node = nodes_by_uuid[parent_uuid]
+            parent_identifier, parent_rel_name = self._child_parent_identifier_map[child_node.identifier]
+            parent_node = nodes_by_identifier[parent_identifier]
             parent_rel = child_node.get_relationship(name=parent_rel_name)
             parent_rel.nodes.add(parent_node)
 

backend/infrahub/core/diff/enricher/hierarchy.py

@@ -99,7 +99,7 @@ async def _enrich_hierarchical_nodes(
 
                 current_node = node
                 for ancestor in ancestors:
-                    ancestor_identifier = NodeIdentifier(uuid=ancestor.uuid, kind=ancestor.kind, labels=ancestor.labels)
+                    ancestor_identifier = NodeIdentifier(uuid=ancestor.uuid, kind=ancestor.kind, db_id=ancestor.db_id)
                     parent_request = ParentNodeAddRequest(
                         node_identifier=current_node.identifier,
                         parent_identifier=ancestor_identifier,
@@ -146,13 +146,11 @@ async def _enrich_nodes_with_parent(
 
             for peer in query.get_peers():
                 source_identifier = NodeIdentifier(
-                    uuid=str(peer.source_id), kind=peer.source_kind, labels=peer.source_labels
+                    uuid=str(peer.source_id), kind=peer.source_kind, db_id=peer.source_db_id
                 )
                 parent_peers[source_identifier] = peer
                 if parent_schema.has_parent_relationship:
-                    peer_identifier = NodeIdentifier(
-                        uuid=str(peer.peer_id), kind=peer.peer_kind, labels=peer.peer_labels
-                    )
+                    peer_identifier = NodeIdentifier(uuid=str(peer.peer_id), kind=peer.peer_kind, db_id=peer.peer_db_id)
                     node_parent_with_parent_map[parent_schema.kind].append(peer_identifier)
 
         # Check if the parent are already present
@@ -170,7 +168,7 @@ async def _enrich_nodes_with_parent(
             parent_rel = [rel for rel in schema_node.relationships if rel.kind == RelationshipKind.PARENT][0]
 
             peer_identifier = NodeIdentifier(
-                uuid=str(peer_parent.peer_id), kind=peer_parent.peer_kind, labels=peer_parent.peer_labels
+                uuid=str(peer_parent.peer_id), kind=peer_parent.peer_kind, db_id=peer_parent.peer_db_id
             )
             parent_request = ParentNodeAddRequest(
                 node_identifier=node.identifier,

backend/infrahub/core/diff/merger/merger.py

@@ -3,8 +3,14 @@
 from typing import TYPE_CHECKING
 
 from infrahub.core import registry
+from infrahub.core.constants import DiffAction
 from infrahub.core.diff.model.path import BranchTrackingId
-from infrahub.core.diff.query.merge import DiffMergePropertiesQuery, DiffMergeQuery, DiffMergeRollbackQuery
+from infrahub.core.diff.query.merge import (
+    DiffMergeMigratedKindsQuery,
+    DiffMergePropertiesQuery,
+    DiffMergeQuery,
+    DiffMergeRollbackQuery,
+)
 from infrahub.log import get_logger
 
 if TYPE_CHECKING:
@@ -53,6 +59,16 @@ async def merge_graph(self, at: Timestamp) -> EnrichedDiffRoot:
         )
         log.info(f"Diff {latest_diff.uuid} retrieved")
         batch_num = 0
+        migrated_kinds_id_map = {}
+        for n in enriched_diff.nodes:
+            if not n.is_node_kind_migration:
+                continue
+            if n.uuid not in migrated_kinds_id_map or (
+                n.uuid in migrated_kinds_id_map and n.action is DiffAction.ADDED
+            ):
+                # make sure that we use the ADDED db_id if it exists
+                # it will not if a node was migrated and then deleted
+                migrated_kinds_id_map[n.uuid] = n.identifier.db_id
         async for node_diff_dicts, property_diff_dicts in self.serializer.serialize_diff(diff=enriched_diff):
             if node_diff_dicts:
                 log.info(f"Merging batch of nodes #{batch_num}")
@@ -62,6 +78,7 @@ async def merge_graph(self, at: Timestamp) -> EnrichedDiffRoot:
                     at=at,
                     target_branch=self.destination_branch,
                     node_diff_dicts=node_diff_dicts,
+                    migrated_kinds_id_map=migrated_kinds_id_map,
                 )
                 await merge_query.execute(db=self.db)
             if property_diff_dicts:
@@ -72,10 +89,21 @@ async def merge_graph(self, at: Timestamp) -> EnrichedDiffRoot:
                     at=at,
                     target_branch=self.destination_branch,
                     property_diff_dicts=property_diff_dicts,
+                    migrated_kinds_id_map=migrated_kinds_id_map,
                 )
                 await merge_properties_query.execute(db=self.db)
             log.info(f"Batch #{batch_num} merged")
             batch_num += 1
+        migrated_kind_uuids = {n.identifier.uuid for n in enriched_diff.nodes if n.is_node_kind_migration}
+        if migrated_kind_uuids:
+            migrated_merge_query = await DiffMergeMigratedKindsQuery.init(
+                db=self.db,
+                branch=self.source_branch,
+                at=at,
+                target_branch=self.destination_branch,
+                migrated_uuids=list(migrated_kind_uuids),
+            )
+            await migrated_merge_query.execute(db=self.db)
 
         self.source_branch.branched_from = at.to_string()
         await self.source_branch.save(db=self.db)

backend/infrahub/core/diff/merger/serializer.py

@@ -39,6 +39,7 @@ def __init__(self, db: InfrahubDatabase, max_batch_size: int) -> None:
 
     def _reset_caches(self) -> None:
         self._attribute_type_cache = {}
+        self._conflicted_cardinality_one_relationships = set()
 
     @property
     def source_branch_name(self) -> str: