infrahub-server 1.1.6__py3-none-any.whl → 1.2.0rc0__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

@@ -133,7 +133,7 @@ class EnrichedDiffDeserializer:
 
                 # 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]):
+                for rel, parent in zip(parents_path_slice[::2], parents_path_slice[1::2], strict=False):
                     enriched_root.add_parent(
                         node_id=current_node_uuid,
                         parent_id=parent.get("uuid"),
@@ -164,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")),
@@ -176,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:
@@ -197,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
@@ -220,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
@@ -247,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",
         )
 
@@ -282,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
@@ -51,7 +52,8 @@ class DiffRepository:
         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 +64,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 +79,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 +94,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 +112,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,21 +139,21 @@ 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,
+            batch_size_limit=batch_size_limit,
             diff_ids=[d.partner_uuid for d in diffs_by_uuid.values()],
         )
         diffs_by_uuid.update({bbr.uuid: bbr for bbr in base_branch_roots})
@@ -159,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,
@@ -217,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,
@@ -277,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
@@ -285,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] = []
@@ -308,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,
@@ -315,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 = []
@@ -374,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(
@@ -400,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...")

infrahub/core/diff/tasks.py

@@ -1,12 +1,15 @@
+from __future__ import annotations
+
 from prefect import flow
 
+from infrahub.context import InfrahubContext  # noqa: TC001  needed for prefect flow
 from infrahub.core import registry
 from infrahub.core.diff.coordinator import DiffCoordinator
-from infrahub.core.diff.models import RequestDiffUpdate
+from infrahub.core.diff.models import RequestDiffUpdate  # noqa: TC001  needed for prefect flow
 from infrahub.core.diff.repository.repository import DiffRepository
 from infrahub.dependencies.registry import get_component_registry
 from infrahub.log import get_logger
-from infrahub.services import services
+from infrahub.services import InfrahubServices  # noqa: TC001  needed for prefect flow
 from infrahub.workflows.catalogue import DIFF_REFRESH
 from infrahub.workflows.utils import add_branch_tag
 
@@ -14,8 +17,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
+async def update_diff(model: RequestDiffUpdate, service: InfrahubServices) -> None:
     await add_branch_tag(branch_name=model.branch_name)
 
     async with service.database.start_session() as db:
@@ -35,8 +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
+async def refresh_diff(branch_name: str, diff_id: str, service: InfrahubServices) -> None:
     await add_branch_tag(branch_name=branch_name)
 
     async with service.database.start_session() as db:
@@ -49,8 +50,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
+async def refresh_diff_all(branch_name: str, context: InfrahubContext, service: InfrahubServices) -> None:
     await add_branch_tag(branch_name=branch_name)
 
     async with service.database.start_session() as db:
@@ -63,5 +63,6 @@ async def refresh_diff_all(branch_name: str) -> None:
             if diff_root.base_branch_name != diff_root.diff_branch_name:
                 await service.workflow.submit_workflow(
                     workflow=DIFF_REFRESH,
+                    context=context,
                     parameters={"branch_name": diff_root.diff_branch_name, "diff_id": diff_root.uuid},
                 )

infrahub/core/enums.py

@@ -2,7 +2,7 @@ import enum
 import re
 from typing import Any
 
-ENUM_NAME_REGEX = re.compile("[_a-zA-Z0-9]+")
+ENUM_NAME_REGEX = re.compile(r"[_a-zA-Z0-9]+")
 
 
 def generate_python_enum(name: str, options: list[Any]) -> type[enum.Enum]:

infrahub/core/integrity/object_conflict/conflict_recorder.py

@@ -17,7 +17,7 @@ class ObjectConflictValidatorRecorder:
         self.validator_label = validator_label
         self.check_schema_kind = check_schema_kind
 
-    async def record_conflicts(self, proposed_change_id: str, conflicts: Sequence[ObjectConflict]) -> list[Node]:  # pylint: disable=too-many-branches
+    async def record_conflicts(self, proposed_change_id: str, conflicts: Sequence[ObjectConflict]) -> list[Node]:
         try:
             proposed_change = await NodeManager.get_one_by_id_or_default_filter(
                 id=proposed_change_id, kind=InfrahubKind.PROPOSEDCHANGE, db=self.db

infrahub/core/ipam/reconciler.py

@@ -104,7 +104,7 @@ class IpamReconciler:
         ip_node_uuid = query.get_ip_node_uuid()
         if not ip_node_uuid:
             node_type = InfrahubKind.IPPREFIX
-            if isinstance(ip_value, (ipaddress.IPv6Interface, ipaddress.IPv4Interface)):
+            if isinstance(ip_value, ipaddress.IPv6Interface | ipaddress.IPv4Interface):
                 node_type = InfrahubKind.IPADDRESS
             raise NodeNotFoundError(node_type=node_type, identifier=str(ip_value))
         current_parent_uuid = query.get_current_parent_uuid()

infrahub/core/ipam/tasks.py

@@ -5,7 +5,7 @@ from prefect import flow
 
 from infrahub.core import registry
 from infrahub.core.ipam.reconciler import IpamReconciler
-from infrahub.services import services
+from infrahub.services import InfrahubServices
 from infrahub.workflows.utils import add_branch_tag
 
 from .model import IpamNodeDetails
@@ -20,8 +20,7 @@ if TYPE_CHECKING:
     description="Ensure the IPAM Tree is up to date",
     persist_result=False,
 )
-async def ipam_reconciliation(branch: str, ipam_node_details: list[IpamNodeDetails]) -> None:
-    service = services.service
+async def ipam_reconciliation(branch: str, ipam_node_details: list[IpamNodeDetails], service: InfrahubServices) -> None:
     async with service.database.start_session() as db:
         branch_obj = await registry.get_branch(db=db, branch=branch)