infrahub-server 1.2.7__py3-none-any.whl → 1.2.9__py3-none-any.whl

infrahub/core/diff/query/filters.py

@@ -1,6 +1,10 @@
+from collections import defaultdict
+from typing import Any
+
 from pydantic import BaseModel, Field
 
 from infrahub.core.constants import DiffAction
+from infrahub.core.diff.model.path import NodeIdentifier
 from infrahub.core.query.utils import filter_and, filter_or
 
 
@@ -29,6 +33,7 @@ class IncExclActionFilterOptions(BaseModel):
 class EnrichedDiffQueryFilters(BaseModel):
     ids: list[str] = Field(default_factory=list)
     kind: IncExclFilterOptions = IncExclFilterOptions()
+    identifiers: list[NodeIdentifier] = Field(default_factory=list)
     namespace: IncExclFilterOptions = IncExclFilterOptions()
     status: IncExclActionFilterOptions = IncExclActionFilterOptions()
     only_conflicted: bool = Field(default=False)
@@ -37,6 +42,7 @@ class EnrichedDiffQueryFilters(BaseModel):
     def is_empty(self) -> bool:
         if (
             not self.ids
+            and not self.identifiers
             and self.only_conflicted is False
             and self.kind.is_empty
             and self.namespace.is_empty
@@ -48,12 +54,20 @@ class EnrichedDiffQueryFilters(BaseModel):
     def generate(self) -> tuple[str, dict]:
         default_filter = ""
 
-        params = {}
+        params: dict[str, Any] = {}
 
         if self.ids:
             params["ids"] = self.ids
             return "diff_node.uuid in $ids", params
 
+        if self.identifiers:
+            params["ids"] = [n.uuid for n in self.identifiers]
+            id_kind_map: dict[str, list[str]] = defaultdict(list)
+            for node_identifier in self.identifiers:
+                id_kind_map[node_identifier.uuid].append(node_identifier.kind)
+            params["id_kind_map"] = id_kind_map
+            return "diff_node.uuid in $ids AND diff_node.kind IN $id_kind_map[diff_node.uuid]", params
+
         filters_list = []
 
         if self.is_empty:

infrahub/core/diff/query/save.py

@@ -1,3 +1,4 @@
+import json
 from typing import Any, Iterable
 
 from infrahub.core.query import Query, QueryType
@@ -92,6 +93,7 @@ WITH root_uuid, node_map, diff_node, (node_map.conflict_params IS NOT NULL) AS h
 SET
     diff_node.kind = node_map.node_properties.kind,
     diff_node.label = node_map.node_properties.label,
+    diff_node.db_labels = node_map.node_properties.db_labels,
     diff_node.changed_at = node_map.node_properties.changed_at,
     diff_node.action = node_map.node_properties.action,
     diff_node.path_identifier = node_map.node_properties.path_identifier
@@ -401,6 +403,7 @@ FOREACH (i in CASE WHEN has_property_conflict = TRUE THEN [1] ELSE [] END |
             "node_properties": {
                 "uuid": enriched_node.uuid,
                 "kind": enriched_node.kind,
+                "db_labels": json.dumps(list(enriched_node.identifier.labels)),
                 "label": enriched_node.label,
                 "changed_at": enriched_node.changed_at.to_string() if enriched_node.changed_at else None,
                 "action": enriched_node.action.value,

infrahub/core/diff/query_parser.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-from collections import defaultdict
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any
 from uuid import uuid4
@@ -15,6 +14,7 @@ from infrahub.core.constants import (
 from infrahub.core.constants.database import DatabaseEdgeType
 from infrahub.core.timestamp import Timestamp
 
+from .model.field_specifiers_map import NodeFieldSpecifierMap
 from .model.path import (
     DatabasePath,
     DiffAttribute,
@@ -23,6 +23,7 @@ from .model.path import (
     DiffRelationship,
     DiffRoot,
     DiffSingleRelationship,
+    NodeIdentifier,
 )
 
 if TYPE_CHECKING:
@@ -394,8 +395,7 @@ class DiffRelationshipIntermediate:
 @dataclass
 class DiffNodeIntermediate(TrackedStatusUpdates):
     force_action: DiffAction | None
-    uuid: str
-    kind: str
+    identifier: NodeIdentifier
     db_id: str
     from_time: Timestamp
     status: RelationshipStatus
@@ -403,6 +403,14 @@ class DiffNodeIntermediate(TrackedStatusUpdates):
     # {(name, identifier): DiffRelationshipIntermediate}
     relationships_by_identifier: dict[tuple[str, str], DiffRelationshipIntermediate] = field(default_factory=dict)
 
+    @property
+    def uuid(self) -> str:
+        return self.identifier.uuid
+
+    @property
+    def kind(self) -> str:
+        return self.identifier.kind
+
     def to_diff_node(self, from_time: Timestamp, include_unchanged: bool) -> DiffNode:
         attributes = []
         for attr in self.attributes_by_name.values():
@@ -424,8 +432,7 @@ class DiffNodeIntermediate(TrackedStatusUpdates):
         if self.force_action:
             action = self.force_action
         return DiffNode(
-            uuid=self.uuid,
-            kind=self.kind,
+            identifier=self.identifier,
             changed_at=changed_at,
             action=action,
             attributes=attributes,
@@ -441,11 +448,11 @@ class DiffNodeIntermediate(TrackedStatusUpdates):
 class DiffRootIntermediate:
     uuid: str
     branch: str
-    nodes_by_id: dict[str, DiffNodeIntermediate] = field(default_factory=dict)
+    nodes_by_identifier: dict[NodeIdentifier, DiffNodeIntermediate] = field(default_factory=dict)
 
     def to_diff_root(self, from_time: Timestamp, to_time: Timestamp, include_unchanged: bool) -> DiffRoot:
         nodes = []
-        for node in self.nodes_by_id.values():
+        for node in self.nodes_by_identifier.values():
             if node.is_empty:
                 continue
             diff_node = node.to_diff_node(from_time=from_time, include_unchanged=include_unchanged)
@@ -462,7 +469,7 @@ class DiffQueryParser:
         schema_manager: SchemaManager,
         from_time: Timestamp,
         to_time: Timestamp | None = None,
-        previous_node_field_specifiers: dict[str, set[str]] | None = None,
+        previous_node_field_specifiers: NodeFieldSpecifierMap | None = None,
     ) -> None:
         self.base_branch_name = base_branch.name
         self.diff_branch_name = diff_branch.name
@@ -476,9 +483,9 @@ class DiffQueryParser:
             self.diff_branched_from_time = Timestamp(diff_branch.get_branched_from())
         self._diff_root_by_branch: dict[str, DiffRootIntermediate] = {}
         self._final_diff_root_by_branch: dict[str, DiffRoot] = {}
-        self._previous_node_field_specifiers = previous_node_field_specifiers or {}
-        self._new_node_field_specifiers: dict[str, set[str]] | None = None
-        self._current_node_field_specifiers: dict[str, set[str]] | None = None
+        self._previous_node_field_specifiers = previous_node_field_specifiers or NodeFieldSpecifierMap()
+        self._new_node_field_specifiers: NodeFieldSpecifierMap | None = None
+        self._current_node_field_specifiers: NodeFieldSpecifierMap | None = None
 
     def get_branches(self) -> set[str]:
         return set(self._final_diff_root_by_branch.keys())
@@ -490,48 +497,33 @@ class DiffQueryParser:
             return self._final_diff_root_by_branch[branch]
         return DiffRoot(from_time=self.from_time, to_time=self.to_time, uuid=str(uuid4()), branch=branch, nodes=[])
 
-    def get_diff_node_field_specifiers(self) -> dict[str, set[str]]:
+    def get_diff_node_field_specifiers(self) -> NodeFieldSpecifierMap:
+        node_field_specifiers_map = NodeFieldSpecifierMap()
         if self.diff_branch_name not in self._diff_root_by_branch:
-            return {}
-        node_field_specifiers_map: dict[str, set[str]] = defaultdict(set)
+            return node_field_specifiers_map
         diff_root = self._diff_root_by_branch[self.diff_branch_name]
-        for node in diff_root.nodes_by_id.values():
+        for node in diff_root.nodes_by_identifier.values():
             for attribute_name in node.attributes_by_name:
-                node_field_specifiers_map[node.uuid].add(attribute_name)
+                node_field_specifiers_map.add_entry(node_uuid=node.uuid, kind=node.kind, field_name=attribute_name)
             for relationship_diff in node.relationships_by_identifier.values():
-                node_field_specifiers_map[node.uuid].add(relationship_diff.identifier)
+                node_field_specifiers_map.add_entry(
+                    node_uuid=node.uuid, kind=node.kind, field_name=relationship_diff.identifier
+                )
         return node_field_specifiers_map
 
-    def _remove_node_specifiers(
-        self, node_specifiers: dict[str, set[str]], node_specifiers_to_remove: dict[str, set[str]]
-    ) -> dict[str, set[str]]:
-        final_node_specifiers: dict[str, set[str]] = defaultdict(set)
-        for node_uuid, field_names_set in node_specifiers.items():
-            specifiers_to_remove = node_specifiers_to_remove.get(node_uuid, set())
-            final_specifiers = field_names_set - specifiers_to_remove
-            if final_specifiers:
-                final_node_specifiers[node_uuid] = final_specifiers
-        return final_node_specifiers
-
-    def get_new_node_field_specifiers(self) -> dict[str, set[str]]:
+    def get_new_node_field_specifiers(self) -> NodeFieldSpecifierMap:
         if self._new_node_field_specifiers is not None:
             return self._new_node_field_specifiers
         branch_node_specifiers = self.get_diff_node_field_specifiers()
-        new_node_field_specifiers = self._remove_node_specifiers(
-            branch_node_specifiers, self._previous_node_field_specifiers
-        )
-        self._new_node_field_specifiers = new_node_field_specifiers
-        return new_node_field_specifiers
+        self._new_node_field_specifiers = branch_node_specifiers - self._previous_node_field_specifiers
+        return self._new_node_field_specifiers
 
-    def get_current_node_field_specifiers(self) -> dict[str, set[str]]:
+    def get_current_node_field_specifiers(self) -> NodeFieldSpecifierMap:
         if self._current_node_field_specifiers is not None:
             return self._current_node_field_specifiers
         new_node_field_specifiers = self.get_new_node_field_specifiers()
-        current_node_field_specifiers = self._remove_node_specifiers(
-            self._previous_node_field_specifiers, new_node_field_specifiers
-        )
-        self._current_node_field_specifiers = current_node_field_specifiers
-        return current_node_field_specifiers
+        self._current_node_field_specifiers = self._previous_node_field_specifiers - new_node_field_specifiers
+        return self._current_node_field_specifiers
 
     def read_result(self, query_result: QueryResult) -> None:
         path = query_result.get_path(label="diff_path")
@@ -565,11 +557,12 @@ class DiffQueryParser:
         return self._diff_root_by_branch[branch]
 
     def _get_diff_node(self, database_path: DatabasePath, diff_root: DiffRootIntermediate) -> DiffNodeIntermediate:
-        node_id = database_path.node_id
-        if node_id not in diff_root.nodes_by_id:
-            diff_root.nodes_by_id[node_id] = DiffNodeIntermediate(
-                uuid=node_id,
-                kind=database_path.node_kind,
+        identifier = NodeIdentifier(
+            uuid=database_path.node_id, kind=database_path.node_kind, labels=database_path.node_labels
+        )
+        if identifier not in diff_root.nodes_by_identifier:
+            diff_root.nodes_by_identifier[identifier] = DiffNodeIntermediate(
+                identifier=identifier,
                 db_id=database_path.node_db_id,
                 from_time=database_path.node_changed_at,
                 status=database_path.node_status,
@@ -577,7 +570,7 @@ class DiffQueryParser:
                 if database_path.node_branch_support is BranchSupportType.AGNOSTIC
                 else None,
             )
-        diff_node = diff_root.nodes_by_id[node_id]
+        diff_node = diff_root.nodes_by_identifier[identifier]
         # special handling for nodes that have their kind updated, which results in 2 nodes with the same uuid
         if diff_node.db_id != database_path.node_db_id and (
             database_path.node_changed_at > diff_node.from_time
@@ -587,7 +580,7 @@ class DiffQueryParser:
                 == (RelationshipStatus.DELETED, RelationshipStatus.ACTIVE)
             )
         ):
-            diff_node.kind = database_path.node_kind
+            diff_node.identifier.kind = database_path.node_kind
             diff_node.db_id = database_path.node_db_id
             diff_node.from_time = database_path.node_changed_at
             diff_node.status = database_path.node_status
@@ -634,7 +627,9 @@ class DiffQueryParser:
         from_time = self.from_time
         if branch_name == self.base_branch_name:
             new_node_field_specifiers = self.get_new_node_field_specifiers()
-            if attribute_name in new_node_field_specifiers.get(diff_node.uuid, set()):
+            if new_node_field_specifiers.has_entry(
+                node_uuid=diff_node.uuid, kind=diff_node.kind, field_name=attribute_name
+            ):
                 from_time = self.diff_branched_from_time
         if attribute_name not in diff_node.attributes_by_name:
             diff_node.attributes_by_name[attribute_name] = DiffAttributeIntermediate(
@@ -678,7 +673,9 @@ class DiffQueryParser:
             from_time = self.from_time
             if branch_name == self.base_branch_name:
                 new_node_field_specifiers = self.get_new_node_field_specifiers()
-                if relationship_schema.get_identifier() in new_node_field_specifiers.get(diff_node.uuid, set()):
+                if new_node_field_specifiers.has_entry(
+                    node_uuid=diff_node.uuid, kind=diff_node.kind, field_name=relationship_schema.get_identifier()
+                ):
                     from_time = self.diff_branched_from_time
             diff_relationship = DiffRelationshipIntermediate(
                 name=relationship_schema.name,
@@ -700,8 +697,8 @@ class DiffQueryParser:
             branch_diff_root = self._diff_root_by_branch.get(branch)
             if not branch_diff_root:
                 continue
-            for node_id, diff_node in branch_diff_root.nodes_by_id.items():
-                base_diff_node = base_diff_root.nodes_by_id.get(node_id)
+            for identifier, diff_node in branch_diff_root.nodes_by_identifier.items():
+                base_diff_node = base_diff_root.nodes_by_identifier.get(identifier)
                 if not base_diff_node:
                     continue
                 self._apply_attribute_previous_values(diff_node=diff_node, base_diff_node=base_diff_node)
@@ -771,7 +768,7 @@ class DiffQueryParser:
         base_diff_root = self._diff_root_by_branch.get(self.base_branch_name)
         if not base_diff_root:
             return
-        for node_diff in base_diff_root.nodes_by_id.values():
+        for node_diff in base_diff_root.nodes_by_identifier.values():
             for attribute_diff in node_diff.attributes_by_name.values():
                 for property_diff in attribute_diff.properties_by_type.values():
                     ordered_diff_values = property_diff.get_ordered_values_asc()

infrahub/core/diff/repository/deserializer.py

@@ -1,3 +1,5 @@
+import json
+
 from neo4j.graph import Node as Neo4jNode
 from neo4j.graph import Path as Neo4jPath
 
@@ -16,6 +18,7 @@ from ..model.path import (
     EnrichedDiffRoot,
     EnrichedDiffRootMetadata,
     EnrichedDiffSingleRelationship,
+    NodeIdentifier,
     deserialize_tracking_id,
 )
 from ..parent_node_adder import DiffParentNodeAdder, ParentNodeAddRequest
@@ -25,13 +28,15 @@ class EnrichedDiffDeserializer:
     def __init__(self, parent_adder: DiffParentNodeAdder) -> None:
         self.parent_adder = parent_adder
         self._diff_root_map: dict[str, EnrichedDiffRoot] = {}
-        self._diff_node_map: dict[tuple[str, str], EnrichedDiffNode] = {}
-        self._diff_node_attr_map: dict[tuple[str, str, str], EnrichedDiffAttribute] = {}
-        self._diff_node_rel_group_map: dict[tuple[str, str, str], EnrichedDiffRelationship] = {}
-        self._diff_node_rel_element_map: dict[tuple[str, str, str, str], EnrichedDiffSingleRelationship] = {}
-        self._diff_prop_map: dict[tuple[str, str, str, str] | tuple[str, str, str, str, str], EnrichedDiffProperty] = {}
-        # {EnrichedDiffRoot: [(node_uuid, parents_path: Neo4jPath), ...]}
-        self._parents_path_map: dict[EnrichedDiffRoot, list[tuple[str, Neo4jPath]]] = {}
+        self._diff_node_map: dict[tuple[str, NodeIdentifier], EnrichedDiffNode] = {}
+        self._diff_node_attr_map: dict[tuple[str, NodeIdentifier, str], EnrichedDiffAttribute] = {}
+        self._diff_node_rel_group_map: dict[tuple[str, NodeIdentifier, str], EnrichedDiffRelationship] = {}
+        self._diff_node_rel_element_map: dict[tuple[str, NodeIdentifier, str, str], EnrichedDiffSingleRelationship] = {}
+        self._diff_prop_map: dict[
+            tuple[str, NodeIdentifier, str, str] | tuple[str, str, str, str, str], EnrichedDiffProperty
+        ] = {}
+        # {EnrichedDiffRoot: [(NodeIdentifier, parents_path: Neo4jPath), ...]}
+        self._parents_path_map: dict[EnrichedDiffRoot, list[tuple[NodeIdentifier, Neo4jPath]]] = {}
 
     def initialize(self) -> None:
         self._diff_root_map = {}
@@ -42,10 +47,12 @@ class EnrichedDiffDeserializer:
         self._diff_prop_map = {}
         self._parents_path_map = {}
 
-    def _track_parents_path(self, enriched_root: EnrichedDiffRoot, node_uuid: str, parents_path: Neo4jPath) -> None:
+    def _track_parents_path(
+        self, enriched_root: EnrichedDiffRoot, node_identifier: NodeIdentifier, parents_path: Neo4jPath
+    ) -> None:
         if enriched_root not in self._parents_path_map:
             self._parents_path_map[enriched_root] = []
-        self._parents_path_map[enriched_root].append((node_uuid, parents_path))
+        self._parents_path_map[enriched_root].append((node_identifier, parents_path))
 
     async def read_result(self, result: QueryResult, include_parents: bool) -> None:
         enriched_root = self._deserialize_diff_root(root_node=result.get_node("diff_root"))
@@ -58,7 +65,7 @@ class EnrichedDiffDeserializer:
             parents_path = result.get("parents_path")
             if parents_path and isinstance(parents_path, Neo4jPath):
                 self._track_parents_path(
-                    enriched_root=enriched_root, node_uuid=enriched_node.uuid, parents_path=parents_path
+                    enriched_root=enriched_root, node_identifier=enriched_node.identifier, parents_path=parents_path
                 )
 
         node_conflict_node = result.get(label="diff_node_conflict")
@@ -130,17 +137,21 @@ class EnrichedDiffDeserializer:
     def _deserialize_parents(self) -> None:
         for enriched_root, node_path_tuples in self._parents_path_map.items():
             self.parent_adder.initialize(enriched_diff_root=enriched_root)
-            for node_uuid, parents_path in node_path_tuples:
+            for node_identifier, parents_path in node_path_tuples:
                 # Remove the node itself from the path
                 parents_path_slice = parents_path.nodes[1:]
 
                 # TODO Ensure the list is even
-                current_node_uuid = node_uuid
+                current_node_identifier = node_identifier
                 for rel, parent in zip(parents_path_slice[::2], parents_path_slice[1::2], strict=False):
+                    parent_identifier = NodeIdentifier(
+                        uuid=parent.get("uuid"),
+                        kind=parent.get("kind"),
+                        labels=frozenset(json.loads(parent.get("db_labels"))),
+                    )
                     parent_request = ParentNodeAddRequest(
-                        node_id=current_node_uuid,
-                        parent_id=parent.get("uuid"),
-                        parent_kind=parent.get("kind"),
+                        node_identifier=current_node_identifier,
+                        parent_identifier=parent_identifier,
                         parent_label=parent.get("label"),
                         parent_rel_name=rel.get("name"),
                         parent_rel_identifier=rel.get("identifier"),
@@ -148,7 +159,7 @@ class EnrichedDiffDeserializer:
                         parent_rel_label=rel.get("label"),
                     )
                     self.parent_adder.add_parent(parent_request=parent_request)
-                    current_node_uuid = parent.get("uuid")
+                    current_node_identifier = parent_identifier
 
     @classmethod
     def _get_str_or_none_property_value(cls, node: Neo4jNode, property_name: str) -> str | None:
@@ -189,14 +200,16 @@ class EnrichedDiffDeserializer:
 
     def _deserialize_diff_node(self, node_node: Neo4jNode, enriched_root: EnrichedDiffRoot) -> EnrichedDiffNode:
         node_uuid = str(node_node.get("uuid"))
-        node_key = (enriched_root.uuid, node_uuid)
+        node_kind = str(node_node.get("kind"))
+        node_db_labels = frozenset(json.loads(node_node.get("db_labels")))
+        node_identifier = NodeIdentifier(uuid=node_uuid, kind=node_kind, labels=node_db_labels)
+        node_key = (enriched_root.uuid, node_identifier)
         if node_key in self._diff_node_map:
             return self._diff_node_map[node_key]
 
         timestamp_str = self._get_str_or_none_property_value(node=node_node, property_name="changed_at")
         enriched_node = EnrichedDiffNode(
-            uuid=node_uuid,
-            kind=str(node_node.get("kind")),
+            identifier=node_identifier,
             label=str(node_node.get("label")),
             changed_at=Timestamp(timestamp_str) if timestamp_str else None,
             action=DiffAction(str(node_node.get("action"))),
@@ -215,7 +228,7 @@ class EnrichedDiffDeserializer:
         self, diff_attr_node: Neo4jNode, enriched_root: EnrichedDiffRoot, enriched_node: EnrichedDiffNode
     ) -> EnrichedDiffAttribute:
         attr_name = str(diff_attr_node.get("name"))
-        attr_key = (enriched_root.uuid, enriched_node.uuid, attr_name)
+        attr_key = (enriched_root.uuid, enriched_node.identifier, attr_name)
         if attr_key in self._diff_node_attr_map:
             return self._diff_node_attr_map[attr_key]
 
@@ -238,7 +251,7 @@ class EnrichedDiffDeserializer:
         self, relationship_group_node: Neo4jNode, enriched_root: EnrichedDiffRoot, enriched_node: EnrichedDiffNode
     ) -> EnrichedDiffRelationship:
         diff_rel_name = str(relationship_group_node.get("name"))
-        rel_key = (enriched_root.uuid, enriched_node.uuid, diff_rel_name)
+        rel_key = (enriched_root.uuid, enriched_node.identifier, diff_rel_name)
         if rel_key in self._diff_node_rel_group_map:
             return self._diff_node_rel_group_map[rel_key]
 
@@ -272,7 +285,7 @@ class EnrichedDiffDeserializer:
         diff_element_peer_id = str(relationship_element_node.get("peer_id"))
         rel_element_key = (
             enriched_root.uuid,
-            enriched_node.uuid,
+            enriched_node.identifier,
             enriched_relationship_group.name,
             diff_element_peer_id,
         )
@@ -320,7 +333,7 @@ class EnrichedDiffDeserializer:
         enriched_root: EnrichedDiffRoot,
     ) -> EnrichedDiffProperty:
         diff_prop_type = str(diff_attr_property_node.get("property_type"))
-        attr_property_key = (enriched_root.uuid, enriched_node.uuid, enriched_attr.name, diff_prop_type)
+        attr_property_key = (enriched_root.uuid, enriched_node.identifier, enriched_attr.name, diff_prop_type)
         if attr_property_key in self._diff_prop_map:
             return self._diff_prop_map[attr_property_key]
 

infrahub/core/diff/repository/repository.py

@@ -1,10 +1,10 @@
-from collections import defaultdict
 from typing import AsyncGenerator, Generator, Iterable
 
 from neo4j.exceptions import TransientError
 
 from infrahub import config
 from infrahub.core import registry
+from infrahub.core.diff.query.drop_nodes import EnrichedDiffDropNodesQuery
 from infrahub.core.diff.query.field_summary import EnrichedDiffNodeFieldSummaryQuery
 from infrahub.core.diff.query.summary_counts_enricher import (
     DiffFieldsSummaryCountsEnricherQuery,
@@ -16,6 +16,7 @@ from infrahub.database import InfrahubDatabase, retry_db_transaction
 from infrahub.exceptions import ResourceNotFoundError
 from infrahub.log import get_logger
 
+from ..model.field_specifiers_map import NodeFieldSpecifierMap
 from ..model.path import (
     ConflictSelection,
     EnrichedDiffConflict,
@@ -26,6 +27,7 @@ from ..model.path import (
     EnrichedDiffsMetadata,
     EnrichedNodeCreateRequest,
     NodeDiffFieldSummary,
+    NodeIdentifier,
     TimeRange,
     TrackingId,
 )
@@ -108,7 +110,7 @@ class DiffRepository:
         diff_branch_names: list[str],
         from_time: Timestamp | None = None,
         to_time: Timestamp | None = None,
-        filters: dict | None = None,
+        filters: EnrichedDiffQueryFilters | None = None,
         include_parents: bool = True,
         limit: int | None = None,
         offset: int | None = None,
@@ -180,11 +182,11 @@ class DiffRepository:
     async def hydrate_diff_pair(
         self,
         enriched_diffs_metadata: EnrichedDiffsMetadata,
-        node_uuids: Iterable[str] | None = None,
+        node_identifiers: Iterable[NodeIdentifier] | None = None,
     ) -> EnrichedDiffs:
-        filters = None
-        if node_uuids:
-            filters = {"ids": list(node_uuids) if node_uuids is not None else None}
+        filters = EnrichedDiffQueryFilters()
+        if node_identifiers:
+            filters.identifiers = list(node_identifiers)
         hydrated_base_diff = await self.get_one(
             diff_branch_name=enriched_diffs_metadata.base_branch_name,
             diff_id=enriched_diffs_metadata.base_branch_diff.uuid,
@@ -207,7 +209,7 @@ class DiffRepository:
         diff_branch_name: str,
         tracking_id: TrackingId | None = None,
         diff_id: str | None = None,
-        filters: dict | None = None,
+        filters: EnrichedDiffQueryFilters | None = None,
         include_parents: bool = True,
     ) -> EnrichedDiffRoot:
         enriched_diffs = await self.get(
@@ -274,6 +276,12 @@ class DiffRepository:
                 )
                 await single_node_query.execute(db=self.db)
 
+    async def _drop_nodes(self, diff_root: EnrichedDiffRoot, node_identifiers: list[NodeIdentifier]) -> None:
+        drop_node_query = await EnrichedDiffDropNodesQuery.init(
+            db=self.db, enriched_diff_uuid=diff_root.uuid, node_identifiers=node_identifiers
+        )
+        await drop_node_query.execute(db=self.db)
+
     @retry_db_transaction(name="enriched_diff_hierarchy_update")
     async def _run_hierarchy_links_update_query(self, diff_root_uuid: str, diff_nodes: list[EnrichedDiffNode]) -> None:
         log.info(f"Updating diff hierarchy links, num_nodes={len(diff_nodes)}")
@@ -321,7 +329,12 @@ class DiffRepository:
                 node_uuids=node_uuids,
             )
 
-    async def save(self, enriched_diffs: EnrichedDiffs | EnrichedDiffsMetadata, do_summary_counts: bool = True) -> None:
+    async def save(
+        self,
+        enriched_diffs: EnrichedDiffs | EnrichedDiffsMetadata,
+        do_summary_counts: bool = True,
+        node_identifiers_to_drop: list[NodeIdentifier] | None = None,
+    ) -> None:
         # metadata-only update
         if not isinstance(enriched_diffs, EnrichedDiffs):
             await self._save_root_metadata(enriched_diffs=enriched_diffs)
@@ -336,6 +349,8 @@ class DiffRepository:
             await self._save_node_batch(node_create_batch=node_create_batch)
             count_nodes_remaining -= len(node_create_batch)
             log.info(f"Batch saved. {count_nodes_remaining=}")
+        if node_identifiers_to_drop:
+            await self._drop_nodes(diff_root=enriched_diffs.diff_branch_diff, node_identifiers=node_identifiers_to_drop)
         await self._update_hierarchy_links(enriched_diffs=enriched_diffs)
         if do_summary_counts:
             await self._update_summary_counts(diff_root=enriched_diffs.diff_branch_diff)
@@ -501,21 +516,25 @@ class DiffRepository:
         await query.execute(db=self.db)
         return query.get_num_changes_by_branch()
 
-    async def get_node_field_specifiers(self, diff_id: str) -> dict[str, set[str]]:
+    async def get_node_field_specifiers(self, diff_id: str) -> NodeFieldSpecifierMap:
         limit = config.SETTINGS.database.query_size_limit
         offset = 0
-        specifiers: dict[str, set[str]] = defaultdict(set)
+        specifiers_map = NodeFieldSpecifierMap()
         while True:
             query = await EnrichedDiffFieldSpecifiersQuery.init(db=self.db, diff_id=diff_id, offset=offset, limit=limit)
             await query.execute(db=self.db)
             has_data = False
             for field_specifier_tuple in query.get_node_field_specifier_tuples():
-                specifiers[field_specifier_tuple[0]].add(field_specifier_tuple[1])
+                specifiers_map.add_entry(
+                    node_uuid=field_specifier_tuple[0],
+                    kind=field_specifier_tuple[1],
+                    field_name=field_specifier_tuple[2],
+                )
                 has_data = True
             if not has_data:
                 break
             offset += limit
-        return specifiers
+        return specifiers_map
 
     async def add_summary_counts(
         self,

infrahub/core/graph/__init__.py

@@ -1 +1 @@
-GRAPH_VERSION = 26
+GRAPH_VERSION = 27

infrahub/core/graph/index.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from infrahub.database.constants import IndexType
+from infrahub.constants.database import IndexType
 from infrahub.database.index import IndexItem
 
 node_indexes: list[IndexItem] = [
@@ -17,6 +17,8 @@ node_indexes: list[IndexItem] = [
     IndexItem(name="diff_node_uuid", label="DiffNode", properties=["uuid"], type=IndexType.TEXT),
 ]
 
+attr_value_index = IndexItem(name="attr_value", label="AttributeValue", properties=["value"], type=IndexType.RANGE)
+
 rel_indexes: list[IndexItem] = [
     IndexItem(name="attr_from", label="HAS_ATTRIBUTE", properties=["from"], type=IndexType.RANGE),
     IndexItem(name="attr_branch", label="HAS_ATTRIBUTE", properties=["branch"], type=IndexType.RANGE),

infrahub/core/initialization.py

@@ -1,7 +1,9 @@
 import importlib
+from typing import TYPE_CHECKING
 from uuid import uuid4
 
 from infrahub import config, lock
+from infrahub.constants.database import DatabaseType
 from infrahub.core import registry
 from infrahub.core.branch import Branch
 from infrahub.core.constants import (
@@ -13,6 +15,7 @@ from infrahub.core.constants import (
     PermissionDecision,
 )
 from infrahub.core.graph import GRAPH_VERSION
+from infrahub.core.graph.index import attr_value_index, node_indexes, rel_indexes
 from infrahub.core.manager import NodeManager
 from infrahub.core.node import Node
 from infrahub.core.node.ipam import BuiltinIPPrefix
@@ -25,6 +28,8 @@ from infrahub.core.root import Root
 from infrahub.core.schema import SchemaRoot, core_models, internal_schema
 from infrahub.core.schema.manager import SchemaManager
 from infrahub.database import InfrahubDatabase
+from infrahub.database.memgraph import IndexManagerMemgraph
+from infrahub.database.neo4j import IndexManagerNeo4j
 from infrahub.exceptions import DatabaseError
 from infrahub.graphql.manager import GraphQLSchemaManager
 from infrahub.log import get_logger
@@ -32,6 +37,9 @@ from infrahub.menu.utils import create_default_menu
 from infrahub.permissions import PermissionBackend
 from infrahub.storage import InfrahubObjectStorage
 
+if TYPE_CHECKING:
+    from infrahub.database.index import IndexManagerBase
+
 log = get_logger()
 
 
@@ -115,7 +123,19 @@ async def initialize_registry(db: InfrahubDatabase, initialize: bool = False) ->
     registry.permission_backends = initialize_permission_backends()
 
 
-async def initialization(db: InfrahubDatabase) -> None:
+async def add_indexes(db: InfrahubDatabase) -> None:
+    if db.db_type is DatabaseType.MEMGRAPH:
+        index_manager: IndexManagerBase = IndexManagerMemgraph(db=db)
+    index_manager = IndexManagerNeo4j(db=db)
+
+    if config.SETTINGS.experimental_features.value_db_index:
+        node_indexes.append(attr_value_index)
+    index_manager.init(nodes=node_indexes, rels=rel_indexes)
+    log.debug("Loading database indexes ..")
+    await index_manager.add()
+
+
+async def initialization(db: InfrahubDatabase, add_database_indexes: bool = False) -> None:
     if config.SETTINGS.database.db_type == config.DatabaseType.MEMGRAPH:
         session = await db.session()
         await session.run(query="SET DATABASE SETTING 'log.level' TO 'INFO'")
@@ -129,12 +149,8 @@ async def initialization(db: InfrahubDatabase) -> None:
         log.debug("Checking Root Node")
         await initialize_registry(db=db, initialize=True)
 
-        # Add Indexes to the database
-        if db.manager.index.initialized:
-            log.debug("Loading database indexes ..")
-            await db.manager.index.add()
-        else:
-            log.warning("The database index manager hasn't been initialized.")
+        if add_database_indexes:
+            await add_indexes(db=db)
 
     # ---------------------------------------------------
     # Load all schema in the database into the registry