diff OOM errors and performance (#5887)

* try getting the hierarchy for multiple nodes at once

* batch getting display labels if over maximum node limit

* a few more logging statements in the enrichers

* split diff hierarchy linking query into smaller batches

* split diff summary query in two

* more logging, smaller pieces when saving a diff

* Revert "try getting the hierarchy for multiple nodes at once"

This reverts commit 2d624a9.

* query update to improve performance of diff parent linking

* add diff indices

* clean up diff save method

* improve diff save query performance

* improve parent hierarchy enrichment performance

* remove a few unused diff indices

* fix typo

* save query performance improvement

* error message update

Author: ajtmccarty

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

@@ -3,6 +3,7 @@
 from infrahub.core.constants import NULL_VALUE, DiffAction, RelationshipCardinality
 from infrahub.core.constants.database import DatabaseEdgeType
 from infrahub.database import InfrahubDatabase
+from infrahub.log import get_logger
 
 from ..model.path import (
     CalculatedDiffs,
@@ -16,6 +17,8 @@
 if TYPE_CHECKING:
     from infrahub.core.schema import MainSchemaTypes
 
+log = get_logger()
+
 
 class DiffCardinalityOneEnricher(DiffEnricherInterface):
     """Clean up diffs for cardinality=one relationships to make them cleaner and more intuitive
@@ -34,13 +37,15 @@ def __init__(self, db: InfrahubDatabase):
 
     async def enrich(self, enriched_diff_root: EnrichedDiffRoot, calculated_diffs: CalculatedDiffs) -> None:
         self._node_schema_map = {}
+        log.info("Beginning cardinality-one diff enrichment...")
         for diff_node in enriched_diff_root.nodes:
             for relationship_group in diff_node.relationships:
                 if (
                     relationship_group.cardinality is RelationshipCardinality.ONE
                     and len(relationship_group.relationships) > 0
                 ):
                     self.consolidate_cardinality_one_diff_elements(diff_relationship=relationship_group)
+        log.info("Cardinality-one diff enrichment complete.")
 
     def _determine_action(self, previous_value: Any, new_value: Any) -> DiffAction:
         if previous_value == new_value:

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

@@ -7,19 +7,24 @@
 from infrahub.core.query.relationship import RelationshipGetPeerQuery, RelationshipPeerData
 from infrahub.core.schema import ProfileSchema
 from infrahub.database import InfrahubDatabase
+from infrahub.log import get_logger
 
 from ..model.path import (
     CalculatedDiffs,
     EnrichedDiffRoot,
 )
+from ..parent_node_adder import DiffParentNodeAdder, ParentNodeAddRequest
 from .interface import DiffEnricherInterface
 
+log = get_logger()
+
 
 class DiffHierarchyEnricher(DiffEnricherInterface):
     """Add hierarchy and parent/component nodes to diff even if the higher-level nodes are unchanged"""
 
-    def __init__(self, db: InfrahubDatabase):
+    def __init__(self, db: InfrahubDatabase, parent_adder: DiffParentNodeAdder):
         self.db = db
+        self.parent_adder = parent_adder
 
     async def enrich(
         self, enriched_diff_root: EnrichedDiffRoot, calculated_diffs: CalculatedDiffs | None = None
@@ -28,6 +33,8 @@ async def enrich(
         # - A node has a relationship of kind parent
         # - A node is part of a hierarchy
 
+        log.info("Beginning hierarchical diff enrichment...")
+        self.parent_adder.initialize(enriched_diff_root=enriched_diff_root)
         node_rel_parent_map: dict[str, list[str]] = defaultdict(list)
         node_hierarchy_map: dict[str, list[str]] = defaultdict(list)
 
@@ -53,6 +60,7 @@ async def enrich(
 
         await self._enrich_nodes_with_parent(enriched_diff_root=enriched_diff_root, node_map=node_rel_parent_map)
         await self._enrich_hierarchical_nodes(enriched_diff_root=enriched_diff_root, node_map=node_hierarchy_map)
+        log.info("Hierarchical diff enrichment complete.")
 
     async def _enrich_hierarchical_nodes(
         self,
@@ -63,6 +71,7 @@ async def _enrich_hierarchical_nodes(
 
         # Retrieve the ID of all ancestors
         for kind, node_ids in node_map.items():
+            log.info(f"Beginning hierarchy enrichment for {kind} node, num_nodes={len(node_ids)}...")
             hierarchy_schema = self.db.schema.get(
                 name=kind, branch=enriched_diff_root.diff_branch_name, duplicate=False
             )
@@ -87,7 +96,7 @@ async def _enrich_hierarchical_nodes(
 
                 current_node = node
                 for ancestor in ancestors:
-                    parent = enriched_diff_root.add_parent(
+                    parent_request = ParentNodeAddRequest(
                         node_id=current_node.uuid,
                         parent_id=str(ancestor.uuid),
                         parent_kind=ancestor.kind,
@@ -97,6 +106,7 @@ async def _enrich_hierarchical_nodes(
                         parent_rel_cardinality=parent_rel.cardinality,
                         parent_rel_label=parent_rel.label or "",
                     )
+                    parent = self.parent_adder.add_parent(parent_request=parent_request)
 
                     current_node = parent
 
@@ -114,6 +124,7 @@ async def _enrich_nodes_with_parent(
 
         # Query the UUID of the parent
         for kind, ids in node_map.items():
+            log.info(f"Beginning parent enrichment for {kind} node, num_nodes={len(ids)}...")
             schema_node = self.db.schema.get(name=kind, branch=enriched_diff_root.diff_branch_name, duplicate=False)
 
             parent_rel = [rel for rel in schema_node.relationships if rel.kind == RelationshipKind.PARENT][0]
@@ -138,15 +149,16 @@ async def _enrich_nodes_with_parent(
         # Check if the parent are already present
         # If parent is already in the list of node we need to add a relationship
         # If parent is not in the list of node, we need to add it
+        diff_node_map = enriched_diff_root.get_node_map(node_uuids=set(parent_peers.keys()))
         for node_id, peer_parent in parent_peers.items():
             # TODO check if we can optimize this part to avoid querying this multiple times
-            node = enriched_diff_root.get_node(node_uuid=node_id)
+            node = diff_node_map[node_id]
             schema_node = self.db.schema.get(
                 name=node.kind, branch=enriched_diff_root.diff_branch_name, duplicate=False
             )
             parent_rel = [rel for rel in schema_node.relationships if rel.kind == RelationshipKind.PARENT][0]
 
-            enriched_diff_root.add_parent(
+            parent_request = ParentNodeAddRequest(
                 node_id=node.uuid,
                 parent_id=str(peer_parent.peer_id),
                 parent_kind=peer_parent.peer_kind,
@@ -156,6 +168,7 @@ async def _enrich_nodes_with_parent(
                 parent_rel_cardinality=parent_rel.cardinality,
                 parent_rel_label=parent_rel.label or "",
             )
+            self.parent_adder.add_parent(parent_request=parent_request)
 
         if node_parent_with_parent_map:
             await self._enrich_nodes_with_parent(

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

@@ -6,6 +6,7 @@
 from infrahub.core.constants.database import DatabaseEdgeType
 from infrahub.core.query.node import NodeGetKindQuery
 from infrahub.database import InfrahubDatabase
+from infrahub.log import get_logger
 
 from ..model.path import (
     CalculatedDiffs,
@@ -17,6 +18,8 @@
 from ..payload_builder import get_display_labels
 from .interface import DiffEnricherInterface
 
+log = get_logger()
+
 PROPERTY_TYPES_WITH_LABELS = {DatabaseEdgeType.IS_RELATED, DatabaseEdgeType.HAS_OWNER, DatabaseEdgeType.HAS_SOURCE}
 
 
@@ -194,6 +197,7 @@ async def enrich(
         calculated_diffs: CalculatedDiffs | None = None,
         conflicts_only: bool = False,
     ) -> None:
+        log.info("Beginning display labels diff enrichment...")
         self._base_branch_name = enriched_diff_root.base_branch_name
         self._diff_branch_name = enriched_diff_root.diff_branch_name
         self._conflicts_only = conflicts_only
@@ -214,3 +218,4 @@ async def enrich(
             ...
 
         self._update_relationship_labels(enriched_diff=enriched_diff_root)
+        log.info("Display labels diff enrichment complete.")

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

@@ -1,10 +1,13 @@
 from infrahub.core.constants import PathType
 from infrahub.core.path import DataPath
 from infrahub.database import InfrahubDatabase
+from infrahub.log import get_logger
 
 from ..model.path import CalculatedDiffs, EnrichedDiffRoot
 from .interface import DiffEnricherInterface
 
+log = get_logger()
+
 
 class DiffPathIdentifierEnricher(DiffEnricherInterface):
     """Add path identifiers to every element in the diff"""
@@ -20,6 +23,7 @@ def diff_branch_name(self) -> str:
         return self._diff_branch_name
 
     async def enrich(self, enriched_diff_root: EnrichedDiffRoot, calculated_diffs: CalculatedDiffs) -> None:
+        log.info("Beginning path identifier diff enrichment...")
         self._diff_branch_name = enriched_diff_root.diff_branch_name
         for node in enriched_diff_root.nodes:
             node_path = DataPath(
@@ -62,3 +66,4 @@ async def enrich(self, enriched_diff_root: EnrichedDiffRoot, calculated_diffs: C
                         relationship_property_path = relationship_element_path.model_copy()
                         relationship_property_path.property_name = relationship_property.property_type.value
                         relationship_property.path_identifier = relationship_property_path.get_path(with_peer=False)
+        log.info("Path identifier diff enrichment complete.")

backend/infrahub/core/diff/model/path.py

@@ -486,6 +486,13 @@ def has_node(self, node_uuid: str) -> bool:
         except ValueError:
             return False
 
+    def get_node_map(self, node_uuids: set[str] | None = None) -> dict[str, EnrichedDiffNode]:
+        node_map = {}
+        for node in self.nodes:
+            if node_uuids is None or node.uuid in node_uuids:
+                node_map[node.uuid] = node
+        return node_map
+
     def get_all_conflicts(self) -> dict[str, EnrichedDiffConflict]:
         all_conflicts: dict[str, EnrichedDiffConflict] = {}
         for node in self.nodes:

backend/infrahub/core/diff/parent_node_adder.py

@@ -0,0 +1,78 @@
+from dataclasses import dataclass, field
+
+from infrahub.core.constants import DiffAction, RelationshipCardinality
+
+from .model.path import EnrichedDiffNode, EnrichedDiffRelationship, EnrichedDiffRoot
+
+
+@dataclass
+class ParentNodeAddRequest:
+    node_id: str
+    parent_id: str
+    parent_kind: str
+    parent_label: str
+    parent_rel_name: str
+    parent_rel_identifier: str
+    parent_rel_cardinality: RelationshipCardinality
+    parent_rel_label: str = field(default="")
+
+
+class DiffParentNodeAdder:
+    def __init__(self) -> None:
+        self._diff_root: EnrichedDiffRoot | None = None
+        self._node_map: dict[str, EnrichedDiffNode] = {}
+
+    def initialize(self, enriched_diff_root: EnrichedDiffRoot) -> None:
+        self._diff_root = enriched_diff_root
+        self._node_map = enriched_diff_root.get_node_map()
+
+    def get_root(self) -> EnrichedDiffRoot:
+        if not self._diff_root:
+            raise RuntimeError("Must call initialize before using")
+        return self._diff_root
+
+    def get_node(self, node_uuid: str) -> EnrichedDiffNode:
+        return self._node_map[node_uuid]
+
+    def has_node(self, node_uuid: str) -> bool:
+        return node_uuid in self._node_map
+
+    def add_node(self, node: EnrichedDiffNode) -> None:
+        if node.uuid in self._node_map:
+            return
+        self._node_map[node.uuid] = node
+        self.get_root().nodes.add(node)
+
+    def add_parent(self, parent_request: ParentNodeAddRequest) -> EnrichedDiffNode:
+        if not self._diff_root:
+            raise RuntimeError("Must call initialize before using")
+        node = self.get_node(node_uuid=parent_request.node_id)
+        if not self.has_node(node_uuid=parent_request.parent_id):
+            parent = EnrichedDiffNode(
+                uuid=parent_request.parent_id,
+                kind=parent_request.parent_kind,
+                label=parent_request.parent_label,
+                action=DiffAction.UNCHANGED,
+                changed_at=None,
+            )
+            self.add_node(parent)
+        else:
+            parent = self.get_node(node_uuid=parent_request.parent_id)
+
+        try:
+            rel = node.get_relationship(name=parent_request.parent_rel_name)
+            rel.nodes.add(parent)
+        except ValueError:
+            node.relationships.add(
+                EnrichedDiffRelationship(
+                    name=parent_request.parent_rel_name,
+                    identifier=parent_request.parent_rel_identifier,
+                    label=parent_request.parent_rel_label,
+                    cardinality=parent_request.parent_rel_cardinality,
+                    changed_at=None,
+                    action=DiffAction.UNCHANGED,
+                    nodes={parent},
+                )
+            )
+
+        return parent

backend/infrahub/core/diff/payload_builder.py

@@ -2,6 +2,7 @@
 
 from typing import TYPE_CHECKING
 
+from infrahub import config
 from infrahub.core.manager import NodeManager
 from infrahub.core.registry import registry
 from infrahub.exceptions import SchemaNotFoundError
@@ -26,8 +27,18 @@ async def get_display_labels_per_kind(
         if skip_missing_schema:
             return {}
         raise
-    nodes = await NodeManager.get_many(ids=ids, fields=fields, db=db, branch=branch)
-    return {node_id: await node.render_display_label(db=db) for node_id, node in nodes.items()}
+    display_label_map: dict[str, str] = {}
+    offset = 0
+    limit = config.SETTINGS.database.query_size_limit
+    while True:
+        limited_ids = ids[offset : offset + limit]
+        if not limited_ids:
+            break
+        node_map = await NodeManager.get_many(ids=limited_ids, fields=fields, db=db, branch=branch)
+        for node_id, node in node_map.items():
+            display_label_map[node_id] = await node.render_display_label(db=db)
+        offset += limit
+    return display_label_map
 
 
 async def get_display_labels(nodes: dict[str, dict[str, list[str]]], db: InfrahubDatabase) -> dict[str, dict[str, str]]: