perf(diff): batch node field summary retrieval over paged queries

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

@@ -121,6 +121,20 @@ def add_relationship_node_uuid(self, name: str, node_uuid: str) -> None:
         """Record that `node_uuid` is a node of this kind whose relationship `name` changed."""
         self.relationship_node_uuids.setdefault(name, set()).add(node_uuid)
 
+    def merge(self, other: NodeDiffFieldSummary) -> None:
+        """Fold another summary of the same kind into this one.
+
+        Raises:
+            ValueError: If the other summary is for a different kind.
+
+        """
+        if other.kind != self.kind:
+            raise ValueError(f"Cannot merge summary for kind {other.kind} into summary for kind {self.kind}")
+        for name, node_uuids in other.attribute_node_uuids.items():
+            self.attribute_node_uuids.setdefault(name, set()).update(node_uuids)
+        for name, node_uuids in other.relationship_node_uuids.items():
+            self.relationship_node_uuids.setdefault(name, set()).update(node_uuids)
+
     @property
     def attribute_names(self) -> set[str]:
         return set(self.attribute_node_uuids)

backend/infrahub/core/diff/query/field_summary.py

@@ -1,4 +1,4 @@
-from dataclasses import dataclass
+from collections.abc import Generator
 from typing import Any
 
 from infrahub.core.constants import DiffAction
@@ -8,29 +8,29 @@
 from ..model.path import NodeDiffFieldSummary, TrackingId
 
 
-@dataclass
-class FieldNodeUuidsRow:
-    """One changed field of a kind and the uuids of the nodes that changed it, as projected by the query."""
-
-    name: str
-    node_uuids: list[str]
-
-
 class EnrichedDiffNodeFieldSummaryQuery(Query):
-    """Get node kind and names of all altered attributes and relationships for each kind."""
+    """Get the names of all altered attributes and relationships for one page of changed nodes.
+
+    Pagination is over the changed nodes, strictly ordered, with each row carrying every altered
+    field name of one node; aggregating the fields of every node in one transaction does not scale
+    with large diffs.
+    """
 
     name = "enriched_diff_node_field_summary"
     type = QueryType.READ
+    insert_limit = False
 
     def __init__(
         self,
         diff_branch_name: str,
+        limit: int,
         tracking_id: TrackingId | None = None,
         diff_id: str | None = None,
         **kwargs: Any,
     ) -> None:
         super().__init__(**kwargs)
         self.diff_branch_name = diff_branch_name
+        self.limit = limit
         self.tracking_id = tracking_id
         self.diff_id = diff_id
 
@@ -42,6 +42,8 @@ async def query_init(self, db: InfrahubDatabase, **kwargs: Any) -> None:  # noqa
             "diff_branch": self.diff_branch_name,
             "tracking_id": self.tracking_id.serialize() if self.tracking_id else None,
             "diff_id": self.diff_id,
+            "node_offset": self.offset or 0,
+            "node_limit": self.limit,
         }
         query = """
         MATCH (diff_root:DiffRoot)
@@ -51,48 +53,40 @@ async def query_init(self, db: InfrahubDatabase, **kwargs: Any) -> None:  # noqa
         AND (diff_root.uuid = $diff_id OR $diff_id IS NULL)
         OPTIONAL MATCH (diff_root)-[:DIFF_HAS_NODE]->(n:DiffNode)
         WHERE n.action <> $unchanged_str
-        WITH DISTINCT diff_root, n.kind AS kind
-        CALL (diff_root, kind) {
-            OPTIONAL MATCH (diff_root)-[:DIFF_HAS_NODE]->(n:DiffNode {kind: kind})-[:DIFF_HAS_ATTRIBUTE]->(a:DiffAttribute)
-            WHERE n.action <> $unchanged_str
-            AND a.action <> $unchanged_str
-            WITH a.name AS attr_name, collect(DISTINCT n.uuid) AS attr_node_uuids
-            WHERE attr_name IS NOT NULL
-            RETURN collect({name: attr_name, node_uuids: attr_node_uuids}) AS attr_name_uuids
+        WITH n
+        ORDER BY n.uuid, elementId(n)
+        SKIP $node_offset
+        LIMIT $node_limit
+        CALL (n) {
+            OPTIONAL MATCH (n)-[:DIFF_HAS_ATTRIBUTE]->(a:DiffAttribute)
+            WHERE a.action <> $unchanged_str
+            RETURN collect(DISTINCT a.name) AS attr_names
         }
-        WITH diff_root, kind, attr_name_uuids
-        CALL (diff_root, kind) {
-            OPTIONAL MATCH (diff_root)-[:DIFF_HAS_NODE]->(n:DiffNode {kind: kind})-[:DIFF_HAS_RELATIONSHIP]->(r:DiffRelationship)
-            WHERE n.action <> $unchanged_str
-            AND r.action <> $unchanged_str
-            WITH r.name AS rel_name, collect(DISTINCT n.uuid) AS rel_node_uuids
-            WHERE rel_name IS NOT NULL
-            RETURN collect({name: rel_name, node_uuids: rel_node_uuids}) AS rel_name_uuids
+        CALL (n) {
+            OPTIONAL MATCH (n)-[:DIFF_HAS_RELATIONSHIP]->(r:DiffRelationship)
+            WHERE r.action <> $unchanged_str
+            RETURN collect(DISTINCT r.name) AS rel_names
         }
         """
         self.add_to_query(query=query)
-        self.order_by = ["kind"]
-        self.return_labels = ["kind", "attr_name_uuids", "rel_name_uuids"]
+        self.return_labels = ["n.kind AS kind", "n.uuid AS node_uuid", "attr_names", "rel_names"]
 
-    async def get_field_summaries(self) -> list[NodeDiffFieldSummary]:
-        field_summaries = []
-        for result in self.get_results():
-            kind = result.get_as_type(label="kind", return_type=str)
-            attribute_node_uuids = self._to_field_uuids(
-                result.get_as_list_of_type(label="attr_name_uuids", return_type=FieldNodeUuidsRow)
-            )
-            relationship_node_uuids = self._to_field_uuids(
-                result.get_as_list_of_type(label="rel_name_uuids", return_type=FieldNodeUuidsRow)
-            )
-            if attribute_node_uuids or relationship_node_uuids:
-                field_summaries.append(
-                    NodeDiffFieldSummary(
-                        kind=kind,
-                        attribute_node_uuids=attribute_node_uuids,
-                        relationship_node_uuids=relationship_node_uuids,
-                    )
-                )
-        return field_summaries
+    def get_node_field_rows(self) -> Generator[NodeDiffFieldSummary, None, None]:
+        """Yield a single-node field summary for each changed node in this page.
 
-    def _to_field_uuids(self, rows: list[FieldNodeUuidsRow]) -> dict[str, set[str]]:
-        return {row.name: set(row.node_uuids) for row in rows}
+        A node whose fields are all unchanged still yields a summary (with empty field maps) so the
+        caller can count the nodes consumed from this page. A diff root with no changed nodes at all
+        produces one node-less row instead; it is skipped here, which is safe for the caller's
+        consumed-node count because null nodes sort after every real node.
+        """
+        for result in self.get_results():
+            kind = result.get_as_str("kind")
+            node_uuid = result.get_as_str("node_uuid")
+            if not kind or not node_uuid:
+                continue
+            summary = NodeDiffFieldSummary(kind=kind)
+            for attr_name in result.get_as_list_of_type(label="attr_names", return_type=str):
+                summary.add_attribute_node_uuid(name=attr_name, node_uuid=node_uuid)
+            for rel_name in result.get_as_list_of_type(label="rel_names", return_type=str):
+                summary.add_relationship_node_uuid(name=rel_name, node_uuid=node_uuid)
+            yield summary

backend/infrahub/core/diff/repository/repository.py

@@ -588,11 +588,32 @@ async def get_affected_node_uuids(
     async def get_node_field_summaries(
         self, diff_branch_name: str, tracking_id: TrackingId | None = None, diff_id: str | None = None
     ) -> list[NodeDiffFieldSummary]:
-        query = await EnrichedDiffNodeFieldSummaryQuery.init(
-            db=self.db, diff_branch_name=diff_branch_name, tracking_id=tracking_id, diff_id=diff_id
-        )
-        await query.execute(db=self.db)
-        return await query.get_field_summaries()
+        node_limit = config.SETTINGS.database.query_size_limit
+        node_offset = 0
+        summaries_by_kind: dict[str, NodeDiffFieldSummary] = {}
+        while True:
+            query = await EnrichedDiffNodeFieldSummaryQuery.init(
+                db=self.db,
+                diff_branch_name=diff_branch_name,
+                tracking_id=tracking_id,
+                diff_id=diff_id,
+                offset=node_offset,
+                limit=node_limit,
+            )
+            await query.execute(db=self.db)
+            num_nodes = 0
+            for node_summary in query.get_node_field_rows():
+                num_nodes += 1
+                if not node_summary.attribute_node_uuids and not node_summary.relationship_node_uuids:
+                    continue
+                kind_summary = summaries_by_kind.setdefault(
+                    node_summary.kind, NodeDiffFieldSummary(kind=node_summary.kind)
+                )
+                kind_summary.merge(node_summary)
+            if num_nodes < node_limit:
+                break
+            node_offset += node_limit
+        return list(summaries_by_kind.values())
 
     async def mark_tracking_ids_merged(self, tracking_ids: list[TrackingId]) -> None:
         query = await EnrichedDiffMergedTrackingIdQuery.init(db=self.db, tracking_ids=tracking_ids)

backend/tests/component/core/diff/repository/test_diff_field_summaries.py

@@ -1,3 +1,4 @@
+from dataclasses import dataclass
 from typing import Generator
 
 import pytest
@@ -11,10 +12,12 @@
     NodeDiffFieldSummary,
 )
 from infrahub.core.diff.parent_node_adder import DiffParentNodeAdder
+from infrahub.core.diff.query.field_summary import EnrichedDiffNodeFieldSummaryQuery
 from infrahub.core.diff.repository.deserializer import EnrichedDiffDeserializer
 from infrahub.core.diff.repository.repository import DiffRepository
 from infrahub.core.timestamp import Timestamp
 from infrahub.database import InfrahubDatabase
+from tests.helpers.db_query_counter import CountingInfrahubDatabase
 from tests.helpers.diff_factories import (
     EnrichedAttributeFactory,
     EnrichedNodeFactory,
@@ -25,25 +28,58 @@
 from .base import DiffRepositoryTestBase
 
 
+@dataclass
+class PageSizeCase:
+    name: str
+    query_size_limit: int
+    expected_query_count: int
+
+
+NUM_CHANGED_NODES = 10
+
+# One query per page, where the last page is the first to come back shorter than the page size. A
+# node count that is an exact multiple of the page size therefore needs one extra, empty page.
+PAGE_SIZE_CASES = [
+    PageSizeCase(name="partial_last_page", query_size_limit=4, expected_query_count=3),
+    PageSizeCase(name="node_count_exact_multiple_of_page", query_size_limit=5, expected_query_count=3),
+    PageSizeCase(name="single_page", query_size_limit=50, expected_query_count=1),
+]
+
+
 class TestDiffNodeFieldSummaries(DiffRepositoryTestBase):
     base_branch_name: str = "main"
     diff_branch_name: str = "diff"
     diff_from_time = Timestamp("2024-06-15T18:35:20Z")
     diff_to_time = Timestamp("2024-06-15T18:49:40Z")
 
     @pytest.fixture
-    def diff_repository(self, db: InfrahubDatabase) -> Generator[DiffRepository, None, None]:
+    def database_settings(self) -> Generator[None, None, None]:
         original_depth = config.SETTINGS.database.max_depth_search_hierarchy
         original_size = config.SETTINGS.database.query_size_limit
         config.SETTINGS.database.max_depth_search_hierarchy = 10
         config.SETTINGS.database.query_size_limit = 50
-        diff_repository = DiffRepository(
-            db=db, deserializer=EnrichedDiffDeserializer(DiffParentNodeAdder()), max_save_batch_size=30
-        )
-        yield diff_repository
+        yield
         config.SETTINGS.database.max_depth_search_hierarchy = original_depth
         config.SETTINGS.database.query_size_limit = original_size
 
+    @pytest.fixture
+    def diff_repository(self, db: InfrahubDatabase, database_settings: None) -> DiffRepository:
+        return DiffRepository(
+            db=db, deserializer=EnrichedDiffDeserializer(DiffParentNodeAdder()), max_save_batch_size=30
+        )
+
+    @pytest.fixture
+    def counting_db(self, db: InfrahubDatabase) -> CountingInfrahubDatabase:
+        return CountingInfrahubDatabase.from_db(db=db)
+
+    @pytest.fixture
+    def counting_diff_repository(
+        self, counting_db: CountingInfrahubDatabase, database_settings: None
+    ) -> DiffRepository:
+        return DiffRepository(
+            db=counting_db, deserializer=EnrichedDiffDeserializer(DiffParentNodeAdder()), max_save_batch_size=30
+        )
+
     def _build_named_field_node(
         self,
         kind: str,
@@ -176,3 +212,104 @@ async def test_get_node_field_summaries_excludes_other_diffs(
             diff_branch_name=requested_branch_name, diff_id=requested_diff.uuid
         )
         assert retrieved_by_diff_id == expected_summaries
+
+    async def test_get_node_field_summaries_empty_diff(
+        self,
+        counting_diff_repository: DiffRepository,
+        counting_db: CountingInfrahubDatabase,
+        reset_database: None,
+    ) -> None:
+        """A diff root with no diff nodes at all yields no summaries, and costs a single query.
+
+        Such a root still produces one node-less row, which occupies a page slot without being a node.
+        Counting that row as a consumed node would keep pagination going past the only page, so the
+        page size here is one: it is the only size at which that miscount changes the query count.
+        """
+        tracking_id = BranchTrackingId(name=self.diff_branch_name)
+        enriched_diff = EnrichedRootFactory.build(
+            base_branch_name=self.base_branch_name,
+            diff_branch_name=self.diff_branch_name,
+            from_time=self.diff_from_time,
+            to_time=self.diff_to_time,
+            nodes=set(),
+            tracking_id=tracking_id,
+        )
+        await self._save_single_diff(
+            diff_repository=counting_diff_repository, enriched_diff=enriched_diff, do_summary_counts=False
+        )
+
+        config.SETTINGS.database.query_size_limit = 1
+        counting_db.reset_counts()
+        retrieved = await counting_diff_repository.get_node_field_summaries(
+            diff_branch_name=self.diff_branch_name, tracking_id=tracking_id
+        )
+        assert retrieved == []
+        assert counting_db.count_for(EnrichedDiffNodeFieldSummaryQuery.name) == 1
+
+    @pytest.mark.parametrize("case", PAGE_SIZE_CASES, ids=lambda c: c.name)
+    async def test_get_node_field_summaries_batched(
+        self,
+        counting_diff_repository: DiffRepository,
+        counting_db: CountingInfrahubDatabase,
+        reset_database: None,
+        case: PageSizeCase,
+    ) -> None:
+        """Per-kind summaries are complete, and cost one query per page, when nodes span pages.
+
+        Ten changed nodes of two kinds are saved, so each page size below ten forces every kind
+        across a page boundary; one node has only unchanged fields, so it fills a page slot without
+        contributing a summary.
+
+        The query count pins the retrieval cost. Improperly configured Query subclasses can cause
+        duplicate queries to run.
+        """
+        tracking_id = BranchTrackingId(name=self.diff_branch_name)
+        kinds = ["TestingKindAlpha", "TestingKindBravo"]
+        nodes: set[EnrichedDiffNode] = set()
+        expected_by_kind: dict[str, NodeDiffFieldSummary] = {}
+        for index in range(NUM_CHANGED_NODES - 1):
+            kind = kinds[index % len(kinds)]
+            node = self._build_named_field_node(
+                kind=kind,
+                node_action=DiffAction.UPDATED,
+                attribute_actions={
+                    "shared_attr": DiffAction.UPDATED,
+                    f"attr_{index}": DiffAction.ADDED,
+                    "quiet_attr": DiffAction.UNCHANGED,
+                },
+                relationship_actions={"shared_rel": DiffAction.UPDATED},
+            )
+            nodes.add(node)
+            expected = expected_by_kind.setdefault(kind, NodeDiffFieldSummary(kind=kind))
+            expected.add_attribute_node_uuid(name="shared_attr", node_uuid=node.uuid)
+            expected.add_attribute_node_uuid(name=f"attr_{index}", node_uuid=node.uuid)
+            expected.add_relationship_node_uuid(name="shared_rel", node_uuid=node.uuid)
+        nodes.add(
+            self._build_named_field_node(
+                kind=kinds[0],
+                node_action=DiffAction.UPDATED,
+                attribute_actions={"quiet_attr": DiffAction.UNCHANGED},
+                relationship_actions={"quiet_rel": DiffAction.UNCHANGED},
+            )
+        )
+        enriched_diff = EnrichedRootFactory.build(
+            base_branch_name=self.base_branch_name,
+            diff_branch_name=self.diff_branch_name,
+            from_time=self.diff_from_time,
+            to_time=self.diff_to_time,
+            nodes=nodes,
+            tracking_id=tracking_id,
+        )
+        await self._save_single_diff(
+            diff_repository=counting_diff_repository, enriched_diff=enriched_diff, do_summary_counts=False
+        )
+
+        config.SETTINGS.database.query_size_limit = case.query_size_limit
+        counting_db.reset_counts()
+        retrieved = await counting_diff_repository.get_node_field_summaries(
+            diff_branch_name=self.diff_branch_name, tracking_id=tracking_id
+        )
+
+        assert len(retrieved) == len(expected_by_kind)
+        assert {summary.kind: summary for summary in retrieved} == expected_by_kind
+        assert counting_db.count_for(EnrichedDiffNodeFieldSummaryQuery.name) == case.expected_query_count