infrahub-server 1.1.5__py3-none-any.whl → 1.1.7__py3-none-any.whl

infrahub/core/diff/query_parser.py

@@ -28,7 +28,6 @@ from .model.path import (
 if TYPE_CHECKING:
     from infrahub.core.branch import Branch
     from infrahub.core.query import QueryResult
-    from infrahub.core.schema import MainSchemaTypes
     from infrahub.core.schema.manager import SchemaManager
     from infrahub.core.schema.relationship_schema import RelationshipSchema
 
@@ -397,8 +396,12 @@ class DiffNodeIntermediate(TrackedStatusUpdates):
     force_action: DiffAction | None
     uuid: str
     kind: str
+    db_id: str
+    from_time: Timestamp
+    status: RelationshipStatus
     attributes_by_name: dict[str, DiffAttributeIntermediate] = field(default_factory=dict)
-    relationships_by_name: dict[str, DiffRelationshipIntermediate] = field(default_factory=dict)
+    # {(name, identifier): DiffRelationshipIntermediate}
+    relationships_by_identifier: dict[tuple[str, str], DiffRelationshipIntermediate] = field(default_factory=dict)
 
     def to_diff_node(self, from_time: Timestamp, include_unchanged: bool) -> DiffNode:
         attributes = []
@@ -408,7 +411,7 @@ class DiffNodeIntermediate(TrackedStatusUpdates):
                 attributes.append(diff_attr)
         action, changed_at = self.get_action_and_timestamp(from_time=from_time)
         relationships = []
-        for rel in self.relationships_by_name.values():
+        for rel in self.relationships_by_identifier.values():
             diff_rel = rel.to_diff_relationship(include_unchanged=include_unchanged)
             if include_unchanged or diff_rel.action is not DiffAction.UNCHANGED:
                 relationships.append(diff_rel)
@@ -431,7 +434,7 @@ class DiffNodeIntermediate(TrackedStatusUpdates):
 
     @property
     def is_empty(self) -> bool:
-        return len(self.attributes_by_name) == 0 and len(self.relationships_by_name) == 0
+        return len(self.attributes_by_name) == 0 and len(self.relationships_by_identifier) == 0
 
 
 @dataclass
@@ -495,7 +498,7 @@ class DiffQueryParser:
         for node in diff_root.nodes_by_id.values():
             for attribute_name in node.attributes_by_name:
                 node_field_specifiers_map[node.uuid].add(attribute_name)
-            for relationship_diff in node.relationships_by_name.values():
+            for relationship_diff in node.relationships_by_identifier.values():
                 node_field_specifiers_map[node.uuid].add(relationship_diff.identifier)
         return node_field_specifiers_map
 
@@ -567,35 +570,53 @@ class DiffQueryParser:
             diff_root.nodes_by_id[node_id] = DiffNodeIntermediate(
                 uuid=node_id,
                 kind=database_path.node_kind,
+                db_id=database_path.node_db_id,
+                from_time=database_path.node_changed_at,
+                status=database_path.node_status,
                 force_action=DiffAction.UPDATED
                 if database_path.node_branch_support is BranchSupportType.AGNOSTIC
                 else None,
             )
         diff_node = diff_root.nodes_by_id[node_id]
+        # 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
+            or (
+                database_path.node_changed_at >= diff_node.from_time
+                and (diff_node.status, database_path.node_status)
+                == (RelationshipStatus.DELETED, RelationshipStatus.ACTIVE)
+            )
+        ):
+            diff_node.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
         diff_node.track_database_path(database_path=database_path)
         return diff_node
 
-    def _get_relationship_schema(
-        self, database_path: DatabasePath, node_schema: MainSchemaTypes
-    ) -> RelationshipSchema | None:
-        relationship_schemas = node_schema.get_relationships_by_identifier(id=database_path.attribute_name)
-        if len(relationship_schemas) == 1:
-            return relationship_schemas[0]
-        possible_path_directions = database_path.possible_relationship_directions
-        for rel_schema in relationship_schemas:
-            if rel_schema.direction in possible_path_directions:
-                return rel_schema
+    def _get_relationship_schema(self, database_path: DatabasePath) -> RelationshipSchema | None:
+        branches_to_check = [database_path.deepest_branch]
+        if database_path.deepest_branch == self.diff_branch_name:
+            branches_to_check.append(self.base_branch_name)
+        for schema_branch_name in branches_to_check:
+            node_schema = self.schema_manager.get(
+                name=database_path.node_kind, branch=schema_branch_name, duplicate=False
+            )
+            relationship_schemas = node_schema.get_relationships_by_identifier(id=database_path.attribute_name)
+            if len(relationship_schemas) == 1:
+                return relationship_schemas[0]
+            possible_path_directions = database_path.possible_relationship_directions
+            for rel_schema in relationship_schemas:
+                if rel_schema.direction in possible_path_directions:
+                    return rel_schema
         return None
 
     def _update_attribute_level(self, database_path: DatabasePath, diff_node: DiffNodeIntermediate) -> None:
-        node_schema = self.schema_manager.get(
-            name=database_path.node_kind, branch=database_path.deepest_branch, duplicate=False
-        )
         if "Attribute" in database_path.attribute_node.labels:
             diff_attribute = self._get_diff_attribute(database_path=database_path, diff_node=diff_node)
             self._update_attribute_property(database_path=database_path, diff_attribute=diff_attribute)
             return
-        relationship_schema = self._get_relationship_schema(database_path=database_path, node_schema=node_schema)
+        relationship_schema = self._get_relationship_schema(database_path=database_path)
         if not relationship_schema:
             return
         diff_relationship = self._get_diff_relationship(
@@ -649,7 +670,9 @@ class DiffQueryParser:
         relationship_schema: RelationshipSchema,
         database_path: DatabasePath,
     ) -> DiffRelationshipIntermediate:
-        diff_relationship = diff_node.relationships_by_name.get(relationship_schema.name)
+        diff_relationship = diff_node.relationships_by_identifier.get(
+            (relationship_schema.name, relationship_schema.get_identifier())
+        )
         if not diff_relationship:
             branch_name = database_path.deepest_branch
             from_time = self.from_time
@@ -663,7 +686,9 @@ class DiffQueryParser:
                 identifier=relationship_schema.get_identifier(),
                 from_time=from_time,
             )
-            diff_node.relationships_by_name[relationship_schema.name] = diff_relationship
+            diff_node.relationships_by_identifier[relationship_schema.name, relationship_schema.get_identifier()] = (
+                diff_relationship
+            )
         return diff_relationship
 
     def _apply_base_branch_previous_values(self) -> None:
@@ -700,8 +725,8 @@ class DiffQueryParser:
     def _apply_relationship_previous_values(
         self, diff_node: DiffNodeIntermediate, base_diff_node: DiffNodeIntermediate
     ) -> None:
-        for relationship_name, diff_relationship in diff_node.relationships_by_name.items():
-            base_diff_relationship = base_diff_node.relationships_by_name.get(relationship_name)
+        for relationship_key, diff_relationship in diff_node.relationships_by_identifier.items():
+            base_diff_relationship = base_diff_node.relationships_by_identifier.get(relationship_key)
             if not base_diff_relationship:
                 continue
             for db_id, property_set in diff_relationship.properties_by_db_id.items():
@@ -754,7 +779,7 @@ class DiffQueryParser:
                         continue
                     if ordered_diff_values[-1].changed_at >= self.diff_branched_from_time:
                         return
-            for relationship_diff in node_diff.relationships_by_name.values():
+            for relationship_diff in node_diff.relationships_by_identifier.values():
                 for diff_relationship_property_list in relationship_diff.properties_by_db_id.values():
                     for diff_relationship_property in diff_relationship_property_list:
                         if diff_relationship_property.changed_at >= self.diff_branched_from_time:

infrahub/core/diff/repository/deserializer.py

@@ -1,5 +1,3 @@
-from typing import Iterable
-
 from neo4j.graph import Node as Neo4jNode
 from neo4j.graph import Path as Neo4jPath
 
@@ -30,37 +28,47 @@ class EnrichedDiffDeserializer:
         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]]] = {}
 
-    def _initialize(self) -> None:
+    def initialize(self) -> None:
         self._diff_root_map = {}
         self._diff_node_map = {}
         self._diff_node_attr_map = {}
         self._diff_node_rel_group_map = {}
         self._diff_node_rel_element_map = {}
         self._diff_prop_map = {}
+        self._parents_path_map = {}
 
-    async def deserialize(
-        self, database_results: Iterable[QueryResult], include_parents: bool
-    ) -> list[EnrichedDiffRoot]:
-        self._initialize()
-        results = list(database_results)
-        for result in results:
-            enriched_root = self._deserialize_diff_root(root_node=result.get_node("diff_root"))
-            node_node = result.get(label="diff_node")
-            if not isinstance(node_node, Neo4jNode):
-                continue
-            enriched_node = self._deserialize_diff_node(node_node=node_node, enriched_root=enriched_root)
-            node_conflict_node = result.get(label="diff_node_conflict")
-            if isinstance(node_conflict_node, Neo4jNode) and not enriched_node.conflict:
-                conflict = self.deserialize_conflict(diff_conflict_node=node_conflict_node)
-                enriched_node.conflict = conflict
-            self._deserialize_attributes(result=result, enriched_root=enriched_root, enriched_node=enriched_node)
-            self._deserialize_relationships(result=result, enriched_root=enriched_root, enriched_node=enriched_node)
+    def _track_parents_path(self, enriched_root: EnrichedDiffRoot, node_uuid: str, 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))
+
+    async def read_result(self, result: QueryResult, include_parents: bool) -> None:
+        enriched_root = self._deserialize_diff_root(root_node=result.get_node("diff_root"))
+        node_node = result.get(label="diff_node")
+        if not isinstance(node_node, Neo4jNode):
+            return
+        enriched_node = self._deserialize_diff_node(node_node=node_node, enriched_root=enriched_root)
 
         if include_parents:
-            for result in results:
-                enriched_root = self._deserialize_diff_root(root_node=result.get_node("diff_root"))
-                self._deserialize_parents(result=result, enriched_root=enriched_root)
+            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
+                )
+
+        node_conflict_node = result.get(label="diff_node_conflict")
+        if isinstance(node_conflict_node, Neo4jNode) and not enriched_node.conflict:
+            conflict = self.deserialize_conflict(diff_conflict_node=node_conflict_node)
+            enriched_node.conflict = conflict
+        self._deserialize_attributes(result=result, enriched_root=enriched_root, enriched_node=enriched_node)
+        self._deserialize_relationships(result=result, enriched_root=enriched_root, enriched_node=enriched_node)
+
+    async def deserialize(self, include_parents: bool = True) -> list[EnrichedDiffRoot]:
+        if include_parents:
+            self._deserialize_parents()
 
         return list(self._diff_root_map.values())
 
@@ -117,30 +125,26 @@ class EnrichedDiffDeserializer:
                 conflict = self.deserialize_conflict(diff_conflict_node=property_conflict)
                 element_property.conflict = conflict
 
-    def _deserialize_parents(self, result: QueryResult, enriched_root: EnrichedDiffRoot) -> None:
-        parents_path = result.get("parents_path")
-        if not parents_path or not isinstance(parents_path, Neo4jPath):
-            return
-
-        node_uuid = result.get(label="diff_node").get("uuid")
-
-        # Remove the node itself from the path
-        parents_path = parents_path.nodes[1:]  # type: ignore[union-attr]
-
-        # TODO Ensure the list is even
-        current_node_uuid = node_uuid
-        for rel, parent in zip(parents_path[::2], parents_path[1::2]):
-            enriched_root.add_parent(
-                node_id=current_node_uuid,
-                parent_id=parent.get("uuid"),
-                parent_kind=parent.get("kind"),
-                parent_label=parent.get("label"),
-                parent_rel_name=rel.get("name"),
-                parent_rel_identifier=rel.get("identifier"),
-                parent_rel_cardinality=RelationshipCardinality(rel.get("cardinality")),
-                parent_rel_label=rel.get("label"),
-            )
-            current_node_uuid = parent.get("uuid")
+    def _deserialize_parents(self) -> None:
+        for enriched_root, node_path_tuples in self._parents_path_map.items():
+            for node_uuid, 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
+                for rel, parent in zip(parents_path_slice[::2], parents_path_slice[1::2]):
+                    enriched_root.add_parent(
+                        node_id=current_node_uuid,
+                        parent_id=parent.get("uuid"),
+                        parent_kind=parent.get("kind"),
+                        parent_label=parent.get("label"),
+                        parent_rel_name=rel.get("name"),
+                        parent_rel_identifier=rel.get("identifier"),
+                        parent_rel_cardinality=RelationshipCardinality(rel.get("cardinality")),
+                        parent_rel_label=rel.get("label"),
+                    )
+                    current_node_uuid = parent.get("uuid")
 
     @classmethod
     def _get_str_or_none_property_value(cls, node: Neo4jNode, property_name: str) -> str | None:
@@ -160,10 +164,8 @@ class EnrichedDiffDeserializer:
     def build_diff_root_metadata(cls, root_node: Neo4jNode) -> EnrichedDiffRootMetadata:
         from_time = Timestamp(str(root_node.get("from_time")))
         to_time = Timestamp(str(root_node.get("to_time")))
-        tracking_id_str = cls._get_str_or_none_property_value(node=root_node, property_name="tracking_id")
-        tracking_id = None
-        if tracking_id_str:
-            tracking_id = deserialize_tracking_id(tracking_id_str=tracking_id_str)
+        tracking_id_str = str(root_node.get("tracking_id"))
+        tracking_id = deserialize_tracking_id(tracking_id_str=tracking_id_str)
         return EnrichedDiffRootMetadata(
             base_branch_name=str(root_node.get("base_branch")),
             diff_branch_name=str(root_node.get("diff_branch")),
@@ -172,11 +174,12 @@ class EnrichedDiffDeserializer:
             uuid=str(root_node.get("uuid")),
             partner_uuid=str(root_node.get("partner_uuid")),
             tracking_id=tracking_id,
-            num_added=int(root_node.get("num_added")),
-            num_updated=int(root_node.get("num_updated")),
-            num_removed=int(root_node.get("num_removed")),
-            num_conflicts=int(root_node.get("num_conflicts")),
+            num_added=int(root_node.get("num_added", 0)),
+            num_updated=int(root_node.get("num_updated", 0)),
+            num_removed=int(root_node.get("num_removed", 0)),
+            num_conflicts=int(root_node.get("num_conflicts", 0)),
             contains_conflict=str(root_node.get("contains_conflict")).lower() == "true",
+            exists_on_database=True,
         )
 
     def _deserialize_diff_node(self, node_node: Neo4jNode, enriched_root: EnrichedDiffRoot) -> EnrichedDiffNode:
@@ -193,10 +196,10 @@ class EnrichedDiffDeserializer:
             changed_at=Timestamp(timestamp_str) if timestamp_str else None,
             action=DiffAction(str(node_node.get("action"))),
             path_identifier=str(node_node.get("path_identifier")),
-            num_added=int(node_node.get("num_added")),
-            num_updated=int(node_node.get("num_updated")),
-            num_removed=int(node_node.get("num_removed")),
-            num_conflicts=int(node_node.get("num_conflicts")),
+            num_added=int(node_node.get("num_added", 0)),
+            num_updated=int(node_node.get("num_updated", 0)),
+            num_removed=int(node_node.get("num_removed", 0)),
+            num_conflicts=int(node_node.get("num_conflicts", 0)),
             contains_conflict=str(node_node.get("contains_conflict")).lower() == "true",
         )
         self._diff_node_map[node_key] = enriched_node
@@ -216,10 +219,10 @@ class EnrichedDiffDeserializer:
             changed_at=Timestamp(str(diff_attr_node.get("changed_at"))),
             path_identifier=str(diff_attr_node.get("path_identifier")),
             action=DiffAction(str(diff_attr_node.get("action"))),
-            num_added=int(diff_attr_node.get("num_added")),
-            num_updated=int(diff_attr_node.get("num_updated")),
-            num_removed=int(diff_attr_node.get("num_removed")),
-            num_conflicts=int(diff_attr_node.get("num_conflicts")),
+            num_added=int(diff_attr_node.get("num_added", 0)),
+            num_updated=int(diff_attr_node.get("num_updated", 0)),
+            num_removed=int(diff_attr_node.get("num_removed", 0)),
+            num_conflicts=int(diff_attr_node.get("num_conflicts", 0)),
             contains_conflict=str(diff_attr_node.get("contains_conflict")).lower() == "true",
         )
         self._diff_node_attr_map[attr_key] = enriched_attr
@@ -243,10 +246,10 @@ class EnrichedDiffDeserializer:
             changed_at=Timestamp(timestamp_str) if timestamp_str else None,
             action=DiffAction(str(relationship_group_node.get("action"))),
             path_identifier=str(relationship_group_node.get("path_identifier")),
-            num_added=int(relationship_group_node.get("num_added")),
-            num_conflicts=int(relationship_group_node.get("num_conflicts")),
-            num_removed=int(relationship_group_node.get("num_removed")),
-            num_updated=int(relationship_group_node.get("num_updated")),
+            num_added=int(relationship_group_node.get("num_added", 0)),
+            num_conflicts=int(relationship_group_node.get("num_conflicts", 0)),
+            num_removed=int(relationship_group_node.get("num_removed", 0)),
+            num_updated=int(relationship_group_node.get("num_updated", 0)),
             contains_conflict=str(relationship_group_node.get("contains_conflict")).lower() == "true",
         )
 
@@ -278,10 +281,10 @@ class EnrichedDiffDeserializer:
             peer_id=diff_element_peer_id,
             peer_label=peer_label,
             path_identifier=str(relationship_element_node.get("path_identifier")),
-            num_added=int(relationship_element_node.get("num_added")),
-            num_updated=int(relationship_element_node.get("num_updated")),
-            num_removed=int(relationship_element_node.get("num_removed")),
-            num_conflicts=int(relationship_element_node.get("num_conflicts")),
+            num_added=int(relationship_element_node.get("num_added", 0)),
+            num_updated=int(relationship_element_node.get("num_updated", 0)),
+            num_removed=int(relationship_element_node.get("num_removed", 0)),
+            num_conflicts=int(relationship_element_node.get("num_conflicts", 0)),
             contains_conflict=str(relationship_element_node.get("contains_conflict")).lower() == "true",
         )
         enriched_relationship_group.relationships.add(enriched_rel_element)

infrahub/core/diff/repository/repository.py

@@ -1,9 +1,10 @@
 from collections import defaultdict
-from typing import AsyncGenerator, Generator
+from typing import AsyncGenerator, Generator, Iterable
 
 from infrahub import config
 from infrahub.core import registry
 from infrahub.core.diff.query.field_summary import EnrichedDiffNodeFieldSummaryQuery
+from infrahub.core.diff.query.summary_counts_enricher import DiffSummaryCountsEnricherQuery
 from infrahub.core.query.diff import DiffCountChanges
 from infrahub.core.timestamp import Timestamp
 from infrahub.database import InfrahubDatabase, retry_db_transaction
@@ -26,13 +27,13 @@ from ..query.all_conflicts import EnrichedDiffAllConflictsQuery
 from ..query.delete_query import EnrichedDiffDeleteQuery
 from ..query.diff_get import EnrichedDiffGetQuery
 from ..query.diff_summary import DiffSummaryCounters, DiffSummaryQuery
-from ..query.drop_tracking_id import EnrichedDiffDropTrackingIdQuery
 from ..query.field_specifiers import EnrichedDiffFieldSpecifiersQuery
 from ..query.filters import EnrichedDiffQueryFilters
 from ..query.get_conflict_query import EnrichedDiffConflictQuery
 from ..query.has_conflicts_query import EnrichedDiffHasConflictQuery
+from ..query.merge_tracking_id import EnrichedDiffMergedTrackingIdQuery
 from ..query.roots_metadata import EnrichedDiffRootsMetadataQuery
-from ..query.save import EnrichedDiffRootsCreateQuery, EnrichedNodeBatchCreateQuery, EnrichedNodesLinkQuery
+from ..query.save import EnrichedDiffRootsUpsertQuery, EnrichedNodeBatchCreateQuery, EnrichedNodesLinkQuery
 from ..query.time_range_query import EnrichedDiffTimeRangeQuery
 from ..query.update_conflict_query import EnrichedDiffConflictUpdateQuery
 from .deserializer import EnrichedDiffDeserializer
@@ -47,6 +48,55 @@ class DiffRepository:
         self.db = db
         self.deserializer = deserializer
 
+    async def _run_get_diff_query(
+        self,
+        base_branch_name: str,
+        diff_branch_names: list[str],
+        batch_size_limit: int,
+        limit: int | None = None,
+        from_time: Timestamp | None = None,
+        to_time: Timestamp | None = None,
+        filters: EnrichedDiffQueryFilters | None = None,
+        offset: int = 0,
+        include_parents: bool = True,
+        max_depth: int | None = None,
+        tracking_id: TrackingId | None = None,
+        diff_ids: list[str] | None = None,
+    ) -> list[EnrichedDiffRoot]:
+        self.deserializer.initialize()
+        final_row_number = None
+        if limit:
+            final_row_number = offset + limit
+        has_more_data = True
+        while has_more_data and (final_row_number is None or offset < final_row_number):
+            if final_row_number is not None and offset + batch_size_limit > final_row_number:
+                batch_size_limit = final_row_number - offset
+            get_query = await EnrichedDiffGetQuery.init(
+                db=self.db,
+                base_branch_name=base_branch_name,
+                diff_branch_names=diff_branch_names,
+                from_time=from_time,
+                to_time=to_time,
+                filters=filters,
+                max_depth=max_depth,
+                limit=batch_size_limit,
+                offset=offset,
+                tracking_id=tracking_id,
+                diff_ids=diff_ids,
+            )
+            log.info(f"Beginning enriched diff get query {batch_size_limit=}, {offset=}")
+            await get_query.execute(db=self.db)
+            log.info("Enriched diff get query complete")
+            last_result = None
+            for query_result in get_query.get_results():
+                await self.deserializer.read_result(result=query_result, include_parents=include_parents)
+                last_result = query_result
+            has_more_data = False
+            if last_result:
+                has_more_data = last_result.get_as_type("has_more_data", bool)
+            offset += batch_size_limit
+        return await self.deserializer.deserialize()
+
     async def get(
         self,
         base_branch_name: str,
@@ -62,23 +112,21 @@ class DiffRepository:
         include_empty: bool = False,
     ) -> list[EnrichedDiffRoot]:
         final_max_depth = config.SETTINGS.database.max_depth_search_hierarchy
-        query = await EnrichedDiffGetQuery.init(
-            db=self.db,
+        batch_size_limit = int(config.SETTINGS.database.query_size_limit / 10)
+        diff_roots = await self._run_get_diff_query(
             base_branch_name=base_branch_name,
             diff_branch_names=diff_branch_names,
+            batch_size_limit=batch_size_limit,
+            limit=limit,
             from_time=from_time,
             to_time=to_time,
             filters=EnrichedDiffQueryFilters(**dict(filters or {})),
+            include_parents=include_parents,
             max_depth=final_max_depth,
-            limit=limit,
-            offset=offset,
+            offset=offset or 0,
             tracking_id=tracking_id,
             diff_ids=diff_ids,
         )
-        await query.execute(db=self.db)
-        diff_roots = await self.deserializer.deserialize(
-            database_results=query.get_results(), include_parents=include_parents
-        )
         if not include_empty:
             diff_roots = [dr for dr in diff_roots if len(dr.nodes) > 0]
         return diff_roots
@@ -91,30 +139,23 @@ class DiffRepository:
         to_time: Timestamp,
     ) -> list[EnrichedDiffs]:
         max_depth = config.SETTINGS.database.max_depth_search_hierarchy
-        query = await EnrichedDiffGetQuery.init(
-            db=self.db,
+        batch_size_limit = int(config.SETTINGS.database.query_size_limit / 10)
+        diff_branch_roots = await self._run_get_diff_query(
             base_branch_name=base_branch_name,
             diff_branch_names=[diff_branch_name],
             from_time=from_time,
             to_time=to_time,
             max_depth=max_depth,
-        )
-        await query.execute(db=self.db)
-        diff_branch_roots = await self.deserializer.deserialize(
-            database_results=query.get_results(), include_parents=True
+            batch_size_limit=batch_size_limit,
         )
         diffs_by_uuid = {dbr.uuid: dbr for dbr in diff_branch_roots}
-        base_partner_query = await EnrichedDiffGetQuery.init(
-            db=self.db,
+        base_branch_roots = await self._run_get_diff_query(
             base_branch_name=base_branch_name,
             diff_branch_names=[base_branch_name],
             max_depth=max_depth,
+            batch_size_limit=batch_size_limit,
             diff_ids=[d.partner_uuid for d in diffs_by_uuid.values()],
         )
-        await base_partner_query.execute(db=self.db)
-        base_branch_roots = await self.deserializer.deserialize(
-            database_results=base_partner_query.get_results(), include_parents=True
-        )
         diffs_by_uuid.update({bbr.uuid: bbr for bbr in base_branch_roots})
         return [
             EnrichedDiffs(
@@ -126,14 +167,23 @@ class DiffRepository:
             for dbr in diff_branch_roots
         ]
 
-    async def hydrate_diff_pair(self, enriched_diffs_metadata: EnrichedDiffsMetadata) -> EnrichedDiffs:
+    async def hydrate_diff_pair(
+        self,
+        enriched_diffs_metadata: EnrichedDiffsMetadata,
+        node_uuids: Iterable[str] | None = None,
+    ) -> EnrichedDiffs:
+        filters = None
+        if node_uuids:
+            filters = {"ids": list(node_uuids) if node_uuids is not None else None}
         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,
+            filters=filters,
         )
         hydrated_branch_diff = await self.get_one(
             diff_branch_name=enriched_diffs_metadata.diff_branch_name,
             diff_id=enriched_diffs_metadata.diff_branch_diff.uuid,
+            filters=filters,
         )
         return EnrichedDiffs(
             base_branch_name=enriched_diffs_metadata.base_branch_name,
@@ -184,17 +234,34 @@ class DiffRepository:
             yield node_requests
 
     @retry_db_transaction(name="enriched_diff_save")
-    async def save(self, enriched_diffs: EnrichedDiffs) -> None:
+    async def save(self, enriched_diffs: EnrichedDiffs | EnrichedDiffsMetadata, do_summary_counts: bool = True) -> None:
+        log.info("Updating diff metadata...")
+        root_query = await EnrichedDiffRootsUpsertQuery.init(db=self.db, enriched_diffs=enriched_diffs)
+        await root_query.execute(db=self.db)
+        log.info("Diff metadata updated.")
+        if not isinstance(enriched_diffs, EnrichedDiffs):
+            return
         num_nodes = len(enriched_diffs.base_branch_diff.nodes) + len(enriched_diffs.diff_branch_diff.nodes)
         log.info(f"Saving diff (num_nodes={num_nodes})...")
-        root_query = await EnrichedDiffRootsCreateQuery.init(db=self.db, enriched_diffs=enriched_diffs)
-        await root_query.execute(db=self.db)
-        for node_create_batch in self._get_node_create_request_batch(enriched_diffs=enriched_diffs):
+        for batch_num, node_create_batch in enumerate(
+            self._get_node_create_request_batch(enriched_diffs=enriched_diffs)
+        ):
+            log.info(f"Saving node batch #{batch_num}...")
             node_query = await EnrichedNodeBatchCreateQuery.init(db=self.db, node_create_batch=node_create_batch)
             await node_query.execute(db=self.db)
+            log.info(f"Batch #{batch_num} saved")
         link_query = await EnrichedNodesLinkQuery.init(db=self.db, enriched_diffs=enriched_diffs)
         await link_query.execute(db=self.db)
         log.info("Diff saved.")
+        if do_summary_counts:
+            node_uuids: list[str] | None = None
+            if enriched_diffs.diff_branch_diff.exists_on_database:
+                node_uuids = list(enriched_diffs.branch_node_uuids)
+            await self.add_summary_counts(
+                diff_branch_name=enriched_diffs.diff_branch_name,
+                diff_id=enriched_diffs.diff_branch_diff.uuid,
+                node_uuids=node_uuids,
+            )
 
     async def summary(
         self,
@@ -244,6 +311,7 @@ class DiffRepository:
         base_branch_names: list[str] | None = None,
         from_time: Timestamp | None = None,
         to_time: Timestamp | None = None,
+        tracking_id: TrackingId | None = None,
     ) -> list[EnrichedDiffsMetadata]:
         if diff_branch_names and base_branch_names:
             diff_branch_names += base_branch_names
@@ -252,6 +320,7 @@ class DiffRepository:
             base_branch_names=base_branch_names,
             from_time=from_time,
             to_time=to_time,
+            tracking_id=tracking_id,
         )
         roots_by_id = {root.uuid: root for root in empty_roots}
         pairs: list[EnrichedDiffsMetadata] = []
@@ -275,6 +344,7 @@ class DiffRepository:
         base_branch_names: list[str] | None = None,
         from_time: Timestamp | None = None,
         to_time: Timestamp | None = None,
+        tracking_id: TrackingId | None = None,
     ) -> list[EnrichedDiffRootMetadata]:
         query = await EnrichedDiffRootsMetadataQuery.init(
             db=self.db,
@@ -282,6 +352,7 @@ class DiffRepository:
             base_branch_names=base_branch_names,
             from_time=from_time,
             to_time=to_time,
+            tracking_id=tracking_id,
         )
         await query.execute(db=self.db)
         diff_roots = []
@@ -341,8 +412,8 @@ class DiffRepository:
         await query.execute(db=self.db)
         return await query.get_field_summaries()
 
-    async def drop_tracking_ids(self, tracking_ids: list[TrackingId]) -> None:
-        query = await EnrichedDiffDropTrackingIdQuery.init(db=self.db, tracking_ids=tracking_ids)
+    async def mark_tracking_ids_merged(self, tracking_ids: list[TrackingId]) -> None:
+        query = await EnrichedDiffMergedTrackingIdQuery.init(db=self.db, tracking_ids=tracking_ids)
         await query.execute(db=self.db)
 
     async def get_num_changes_in_time_range_by_branch(
@@ -367,3 +438,21 @@ class DiffRepository:
                 break
             offset += limit
         return specifiers
+
+    async def add_summary_counts(
+        self,
+        diff_branch_name: str,
+        tracking_id: TrackingId | None = None,
+        diff_id: str | None = None,
+        node_uuids: list[str] | None = None,
+    ) -> None:
+        log.info("Updating summary counts...")
+        query = await DiffSummaryCountsEnricherQuery.init(
+            db=self.db,
+            diff_branch_name=diff_branch_name,
+            tracking_id=tracking_id,
+            diff_id=diff_id,
+            node_uuids=node_uuids,
+        )
+        await query.execute(db=self.db)
+        log.info("Summary counts updated...")