infrahub-server 1.1.6__py3-none-any.whl → 1.1.8__py3-none-any.whl

infrahub/core/diff/repository/repository.py

@@ -1,9 +1,15 @@
 from collections import defaultdict
-from typing import AsyncGenerator, Generator
+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.field_summary import EnrichedDiffNodeFieldSummaryQuery
+from infrahub.core.diff.query.summary_counts_enricher import (
+    DiffFieldsSummaryCountsEnricherQuery,
+    DiffNodesSummaryCountsEnricherQuery,
+)
 from infrahub.core.query.diff import DiffCountChanges
 from infrahub.core.timestamp import Timestamp
 from infrahub.database import InfrahubDatabase, retry_db_transaction
@@ -13,6 +19,7 @@ from infrahub.log import get_logger
 from ..model.path import (
     ConflictSelection,
     EnrichedDiffConflict,
+    EnrichedDiffNode,
     EnrichedDiffRoot,
     EnrichedDiffRootMetadata,
     EnrichedDiffs,
@@ -26,13 +33,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
@@ -41,17 +48,17 @@ log = get_logger()
 
 
 class DiffRepository:
-    MAX_SAVE_BATCH_SIZE: int = 100
-
-    def __init__(self, db: InfrahubDatabase, deserializer: EnrichedDiffDeserializer):
+    def __init__(self, db: InfrahubDatabase, deserializer: EnrichedDiffDeserializer, max_save_batch_size: int = 1000):
         self.db = db
         self.deserializer = deserializer
+        self.max_save_batch_size = max_save_batch_size
 
     async def _run_get_diff_query(
         self,
         base_branch_name: str,
         diff_branch_names: list[str],
-        limit: int,
+        batch_size_limit: int,
+        limit: int | None = None,
         from_time: Timestamp | None = None,
         to_time: Timestamp | None = None,
         filters: EnrichedDiffQueryFilters | None = None,
@@ -62,8 +69,13 @@ class DiffRepository:
         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:
+        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,
@@ -72,12 +84,12 @@ class DiffRepository:
                 to_time=to_time,
                 filters=filters,
                 max_depth=max_depth,
-                limit=limit,
+                limit=batch_size_limit,
                 offset=offset,
                 tracking_id=tracking_id,
                 diff_ids=diff_ids,
             )
-            log.info(f"Beginning enriched diff get query {limit=}, {offset=}")
+            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
@@ -87,7 +99,7 @@ class DiffRepository:
             has_more_data = False
             if last_result:
                 has_more_data = last_result.get_as_type("has_more_data", bool)
-            offset += limit
+            offset += batch_size_limit
         return await self.deserializer.deserialize()
 
     async def get(
@@ -105,16 +117,17 @@ class DiffRepository:
         include_empty: bool = False,
     ) -> list[EnrichedDiffRoot]:
         final_max_depth = config.SETTINGS.database.max_depth_search_hierarchy
-        limit = limit or int(config.SETTINGS.database.query_size_limit / 10)
+        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 or 0,
             tracking_id=tracking_id,
             diff_ids=diff_ids,
@@ -131,42 +144,56 @@ class DiffRepository:
         to_time: Timestamp,
     ) -> list[EnrichedDiffs]:
         max_depth = config.SETTINGS.database.max_depth_search_hierarchy
-        limit = int(config.SETTINGS.database.query_size_limit / 10)
+        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,
-            limit=limit,
+            batch_size_limit=batch_size_limit,
         )
         diffs_by_uuid = {dbr.uuid: dbr for dbr in diff_branch_roots}
         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,
-            limit=limit,
-            diff_ids=[d.partner_uuid for d in diffs_by_uuid.values()],
+            batch_size_limit=batch_size_limit,
+            diff_ids=[d.partner_uuid for d in diffs_by_uuid.values() if d.partner_uuid],
         )
         diffs_by_uuid.update({bbr.uuid: bbr for bbr in base_branch_roots})
-        return [
-            EnrichedDiffs(
-                base_branch_name=base_branch_name,
-                diff_branch_name=diff_branch_name,
-                base_branch_diff=diffs_by_uuid[dbr.partner_uuid],
-                diff_branch_diff=dbr,
+        diff_pairs = []
+        for dbr in diff_branch_roots:
+            if dbr.partner_uuid is None:
+                continue
+            base_branch_diff = diffs_by_uuid[dbr.partner_uuid]
+            diff_pairs.append(
+                EnrichedDiffs(
+                    base_branch_name=base_branch_name,
+                    diff_branch_name=diff_branch_name,
+                    base_branch_diff=base_branch_diff,
+                    diff_branch_diff=dbr,
+                )
             )
-            for dbr in diff_branch_roots
-        ]
+        return diff_pairs
 
-    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,
@@ -208,26 +235,111 @@ class DiffRepository:
     ) -> Generator[list[EnrichedNodeCreateRequest], None, None]:
         node_requests = []
         for diff_root in (enriched_diffs.base_branch_diff, enriched_diffs.diff_branch_diff):
+            size_count = 0
             for node in diff_root.nodes:
-                node_requests.append(EnrichedNodeCreateRequest(node=node, root_uuid=diff_root.uuid))
-                if len(node_requests) == self.MAX_SAVE_BATCH_SIZE:
+                node_size_count = node.num_properties
+                if size_count + node_size_count < self.max_save_batch_size:
+                    node_requests.append(EnrichedNodeCreateRequest(node=node, root_uuid=diff_root.uuid))
+                    size_count += node_size_count
+                else:
+                    log.info(f"Num nodes in batch: {len(node_requests)}, num properties in batch: {size_count}")
                     yield node_requests
-                    node_requests = []
+                    size_count = node_size_count
+                    node_requests = [EnrichedNodeCreateRequest(node=node, root_uuid=diff_root.uuid)]
         if node_requests:
+            log.info(f"Num nodes in batch: {len(node_requests)}, num properties in batch: {size_count}")
             yield node_requests
 
-    @retry_db_transaction(name="enriched_diff_save")
-    async def save(self, enriched_diffs: EnrichedDiffs) -> None:
-        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)
+    @retry_db_transaction(name="enriched_diff_metadata_save")
+    async def _save_root_metadata(self, enriched_diffs: EnrichedDiffsMetadata) -> 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)
-        for node_create_batch in self._get_node_create_request_batch(enriched_diffs=enriched_diffs):
-            node_query = await EnrichedNodeBatchCreateQuery.init(db=self.db, node_create_batch=node_create_batch)
+        log.info("Diff metadata updated.")
+
+    async def _save_node_batch(self, node_create_batch: list[EnrichedNodeCreateRequest]) -> None:
+        node_query = await EnrichedNodeBatchCreateQuery.init(db=self.db, node_create_batch=node_create_batch)
+        try:
             await node_query.execute(db=self.db)
-        link_query = await EnrichedNodesLinkQuery.init(db=self.db, enriched_diffs=enriched_diffs)
+        except TransientError as exc:
+            if not exc.code or "OutOfMemoryError".lower() not in str(exc.code).lower():
+                raise
+            log.exception("Database memory error during save. Trying smaller transactions")
+            for node_request in node_create_batch:
+                log.info(
+                    f"Updating node {node_request.node.uuid}, num_properties={node_request.node.num_properties}..."
+                )
+                single_node_query = await EnrichedNodeBatchCreateQuery.init(
+                    db=self.db, node_create_batch=[node_request]
+                )
+                await single_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)}")
+        link_query = await EnrichedNodesLinkQuery.init(db=self.db, diff_root_uuid=diff_root_uuid, diff_nodes=diff_nodes)
         await link_query.execute(db=self.db)
-        log.info("Diff saved.")
+
+    async def _update_hierarchy_links(self, enriched_diffs: EnrichedDiffs) -> None:
+        for diff_root in (enriched_diffs.base_branch_diff, enriched_diffs.diff_branch_diff):
+            nodes_to_update = []
+            for node in diff_root.nodes:
+                if any(r.nodes for r in node.relationships):
+                    nodes_to_update.append(node)
+                if len(nodes_to_update) >= config.SETTINGS.database.query_size_limit:
+                    await self._run_hierarchy_links_update_query(
+                        diff_root_uuid=diff_root.uuid, diff_nodes=nodes_to_update
+                    )
+                    nodes_to_update = []
+            if nodes_to_update:
+                await self._run_hierarchy_links_update_query(diff_root_uuid=diff_root.uuid, diff_nodes=nodes_to_update)
+
+    async def _update_summary_counts(self, diff_root: EnrichedDiffRoot) -> None:
+        max_nodes_limit = config.SETTINGS.database.query_size_limit
+        num_nodes = len(diff_root.nodes)
+        if diff_root.exists_on_database and num_nodes < max_nodes_limit:
+            await self.add_summary_counts(
+                diff_branch_name=diff_root.diff_branch_name,
+                diff_id=diff_root.uuid,
+                node_uuids=None,
+            )
+            return
+        node_uuids: list[str] = []
+        for diff_node in diff_root.nodes:
+            node_uuids.append(diff_node.uuid)
+            if len(node_uuids) >= max_nodes_limit:
+                await self.add_summary_counts(
+                    diff_branch_name=diff_root.diff_branch_name,
+                    diff_id=diff_root.uuid,
+                    node_uuids=node_uuids,
+                )
+                node_uuids = []
+        if node_uuids:
+            await self.add_summary_counts(
+                diff_branch_name=diff_root.diff_branch_name,
+                diff_id=diff_root.uuid,
+                node_uuids=node_uuids,
+            )
+
+    async def save(self, enriched_diffs: EnrichedDiffs | EnrichedDiffsMetadata, do_summary_counts: bool = True) -> None:
+        # metadata-only update
+        if not isinstance(enriched_diffs, EnrichedDiffs):
+            await self._save_root_metadata(enriched_diffs=enriched_diffs)
+            return
+
+        count_nodes_remaining = len(enriched_diffs.base_branch_diff.nodes) + len(enriched_diffs.diff_branch_diff.nodes)
+        log.info(f"Saving diff (num_nodes={count_nodes_remaining})...")
+        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}...")
+            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=}")
+        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)
+        await self._save_root_metadata(enriched_diffs=enriched_diffs)
 
     async def summary(
         self,
@@ -277,6 +389,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
@@ -285,11 +398,12 @@ 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] = []
         for branch_root in empty_roots:
-            if branch_root.base_branch_name == branch_root.diff_branch_name:
+            if branch_root.base_branch_name == branch_root.diff_branch_name or branch_root.partner_uuid is None:
                 continue
             base_root = roots_by_id[branch_root.partner_uuid]
             pairs.append(
@@ -308,6 +422,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,
@@ -315,6 +430,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 = []
@@ -374,8 +490,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(
@@ -400,3 +516,61 @@ 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:
+        await self._add_field_summary_counts(
+            diff_branch_name=diff_branch_name,
+            tracking_id=tracking_id,
+            diff_id=diff_id,
+            node_uuids=node_uuids,
+        )
+        await self._add_node_summary_counts(
+            diff_branch_name=diff_branch_name,
+            tracking_id=tracking_id,
+            diff_id=diff_id,
+            node_uuids=node_uuids,
+        )
+
+    @retry_db_transaction(name="enriched_diff_field_summary_counts")
+    async def _add_field_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 field summary counts...")
+        query = await DiffFieldsSummaryCountsEnricherQuery.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("Field summary counts updated.")
+
+    @retry_db_transaction(name="enriched_diff_node_summary_counts")
+    async def _add_node_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 node summary counts...")
+        query = await DiffNodesSummaryCountsEnricherQuery.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("node summary counts updated.")

infrahub/core/diff/tasks.py

@@ -8,7 +8,7 @@ from infrahub.dependencies.registry import get_component_registry
 from infrahub.log import get_logger
 from infrahub.services import services
 from infrahub.workflows.catalogue import DIFF_REFRESH
-from infrahub.workflows.utils import add_branch_tag
+from infrahub.workflows.utils import add_tags
 
 log = get_logger()
 
@@ -16,7 +16,7 @@ log = get_logger()
 @flow(name="diff-update", flow_run_name="Update diff for branch {model.branch_name}")
 async def update_diff(model: RequestDiffUpdate) -> None:
     service = services.service
-    await add_branch_tag(branch_name=model.branch_name)
+    await add_tags(branches=[model.branch_name])
 
     async with service.database.start_session() as db:
         component_registry = get_component_registry()
@@ -37,7 +37,7 @@ async def update_diff(model: RequestDiffUpdate) -> None:
 @flow(name="diff-refresh", flow_run_name="Recreate diff for branch {branch_name}")
 async def refresh_diff(branch_name: str, diff_id: str) -> None:
     service = services.service
-    await add_branch_tag(branch_name=branch_name)
+    await add_tags(branches=[branch_name])
 
     async with service.database.start_session() as db:
         component_registry = get_component_registry()
@@ -51,7 +51,7 @@ async def refresh_diff(branch_name: str, diff_id: str) -> None:
 @flow(name="diff-refresh-all", flow_run_name="Recreate all diffs for branch {branch_name}")
 async def refresh_diff_all(branch_name: str) -> None:
     service = services.service
-    await add_branch_tag(branch_name=branch_name)
+    await add_tags(branches=[branch_name])
 
     async with service.database.start_session() as db:
         component_registry = get_component_registry()

infrahub/core/graph/__init__.py

@@ -1 +1 @@
-GRAPH_VERSION = 18
+GRAPH_VERSION = 20

infrahub/core/graph/index.py

@@ -12,6 +12,9 @@ node_indexes: list[IndexItem] = [
     IndexItem(name="attr_iphost_bin", label="AttributeIPHost", properties=["binary_address"], type=IndexType.RANGE),
     IndexItem(name="rel_uuid", label="Relationship", properties=["uuid"], type=IndexType.RANGE),
     IndexItem(name="rel_identifier", label="Relationship", properties=["name"], type=IndexType.RANGE),
+    # diff indices
+    IndexItem(name="diff_uuid", label="DiffRoot", properties=["uuid"], type=IndexType.TEXT),
+    IndexItem(name="diff_node_uuid", label="DiffNode", properties=["uuid"], type=IndexType.TEXT),
 ]
 
 rel_indexes: list[IndexItem] = [

infrahub/core/migrations/graph/__init__.py

@@ -20,6 +20,8 @@ from .m015_diff_format_update import Migration015
 from .m016_diff_delete_bug_fix import Migration016
 from .m017_add_core_profile import Migration017
 from .m018_uniqueness_nulls import Migration018
+from .m019_restore_rels_to_time import Migration019
+from .m020_duplicate_edges import Migration020
 
 if TYPE_CHECKING:
     from infrahub.core.root import Root
@@ -45,6 +47,8 @@ MIGRATIONS: list[type[Union[GraphMigration, InternalSchemaMigration, ArbitraryMi
     Migration016,
     Migration017,
     Migration018,
+    Migration019,
+    Migration020,
 ]