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

@@ -1,22 +1,24 @@
+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
+from infrahub.workflows.utils import add_tags
 
 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)
+async def update_diff(model: RequestDiffUpdate, service: InfrahubServices) -> None:
+    await add_tags(branches=[model.branch_name])
 
     async with service.database.start_session() as db:
         component_registry = get_component_registry()
@@ -35,9 +37,8 @@ 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)
+async def refresh_diff(branch_name: str, diff_id: str, service: InfrahubServices) -> None:
+    await add_tags(branches=[branch_name])
 
     async with service.database.start_session() as db:
         component_registry = get_component_registry()
@@ -49,9 +50,8 @@ 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)
+async def refresh_diff_all(branch_name: str, context: InfrahubContext, service: InfrahubServices) -> None:
+    await add_tags(branches=[branch_name])
 
     async with service.database.start_session() as db:
         component_registry = get_component_registry()
@@ -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/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/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)
 

infrahub/core/manager.py

@@ -23,7 +23,14 @@ from infrahub.core.query.node import (
 from infrahub.core.query.relationship import RelationshipGetPeerQuery
 from infrahub.core.registry import registry
 from infrahub.core.relationship import Relationship, RelationshipManager
-from infrahub.core.schema import GenericSchema, MainSchemaTypes, NodeSchema, ProfileSchema, RelationshipSchema
+from infrahub.core.schema import (
+    GenericSchema,
+    MainSchemaTypes,
+    NodeSchema,
+    ProfileSchema,
+    RelationshipSchema,
+    TemplateSchema,
+)
 from infrahub.core.timestamp import Timestamp
 from infrahub.exceptions import NodeNotFoundError, ProcessingError, SchemaNotFoundError
 from infrahub.graphql.models import OrderModel
@@ -35,8 +42,6 @@ if TYPE_CHECKING:
 
 SchemaProtocol = TypeVar("SchemaProtocol")
 
-# pylint: disable=redefined-builtin,too-many-lines
-
 
 def identify_node_class(node: NodeToProcess) -> type[Node]:
     """Identify the proper class to use to create the NodeObject.
@@ -61,7 +66,7 @@ def get_schema(
 ) -> MainSchemaTypes:
     if isinstance(node_schema, str):
         return db.schema.get(name=node_schema, branch=branch.name)
-    if hasattr(node_schema, "_is_runtime_protocol") and getattr(node_schema, "_is_runtime_protocol"):
+    if hasattr(node_schema, "_is_runtime_protocol") and node_schema._is_runtime_protocol:
         return db.schema.get(name=node_schema.__name__, branch=branch.name)
     if not isinstance(node_schema, (MainSchemaTypes)):
         raise ValueError(f"Invalid schema provided {node_schema}")
@@ -129,7 +134,7 @@ class NodeManager:
     async def query(
         cls,
         db: InfrahubDatabase,
-        schema: Union[NodeSchema, GenericSchema, ProfileSchema, str],
+        schema: Union[NodeSchema, GenericSchema, ProfileSchema, TemplateSchema, str],
         filters: dict | None = ...,
         fields: dict | None = ...,
         offset: int | None = ...,
@@ -267,11 +272,11 @@ class NodeManager:
     async def count(
         cls,
         db: InfrahubDatabase,
-        schema: Union[type[SchemaProtocol], NodeSchema, GenericSchema, ProfileSchema, str],
+        schema: Union[type[SchemaProtocol], NodeSchema, GenericSchema, ProfileSchema, TemplateSchema, str],
         filters: Optional[dict] = None,
         at: Optional[Union[Timestamp, str]] = None,
         branch: Optional[Union[Branch, str]] = None,
-        account=None,  # pylint: disable=unused-argument
+        account=None,  # noqa: ARG003
         partial_match: bool = False,
         branch_agnostic: bool = False,
     ) -> int:
@@ -446,7 +451,7 @@ class NodeManager:
         limit: Optional[int] = None,
         at: Optional[Union[Timestamp, str]] = None,
         branch: Optional[Union[Branch, str]] = None,
-    ) -> dict[str, Any]:
+    ) -> dict[str, Node]:
         branch = await registry.get_branch(branch=branch, db=db)
         at = Timestamp(at)
 
@@ -489,7 +494,7 @@ class NodeManager:
         kind: type[SchemaProtocol],
         at: Timestamp | str | None = ...,
         branch: Branch | str | None = ...,
-        id: str | None = ...,  # pylint: disable=redefined-builtin
+        id: str | None = ...,
         hfid: list[str] | None = ...,
     ) -> SchemaProtocol: ...
 
@@ -501,7 +506,7 @@ class NodeManager:
         kind: str,
         at: Timestamp | str | None = ...,
         branch: Branch | str | None = ...,
-        id: str | None = ...,  # pylint: disable=redefined-builtin
+        id: str | None = ...,
         hfid: list[str] | None = ...,
     ) -> Any: ...
 
@@ -512,7 +517,7 @@ class NodeManager:
         kind: type[SchemaProtocol] | str,
         at: Timestamp | str | None = None,
         branch: Branch | str | None = None,
-        id: str | None = None,  # pylint: disable=redefined-builtin
+        id: str | None = None,
         hfid: list[str] | None = None,
     ) -> Any:
         if not id and not hfid:
@@ -802,7 +807,7 @@ class NodeManager:
             raise NodeNotFoundError(branch_name=branch.name, node_type=kind_str, identifier=hfid_str)
 
         filters = {}
-        for key, item in zip(node_schema.human_friendly_id, hfid):
+        for key, item in zip(node_schema.human_friendly_id, hfid, strict=False):
             path = node_schema.parse_schema_path(path=key, schema=registry.schema.get_schema_branch(name=branch.name))
 
             if path.is_type_relationship:
@@ -1007,7 +1012,7 @@ class NodeManager:
         cls,
         id: str,
         db: InfrahubDatabase,
-        kind: Literal[None] = ...,
+        kind: None = ...,
         raise_on_error: bool = ...,
         fields: dict | None = ...,
         at: Timestamp | str | None = ...,
@@ -1148,7 +1153,7 @@ class NodeManager:
 
         nodes: dict[str, Node] = {}
 
-        for node_id in ids:  # pylint: disable=too-many-nested-blocks
+        for node_id in ids:
             if node_id not in nodes_info_by_id:
                 continue
 
@@ -1310,7 +1315,7 @@ class NodeManager:
         nodes: list[Node],
         branch: Optional[Union[Branch, str]] = None,
         at: Optional[Union[Timestamp, str]] = None,
-    ) -> list[Any]:
+    ) -> list[Node]:
         """Returns list of deleted nodes because of cascading deletes"""
         branch = await registry.get_branch(branch=branch, db=db)
         node_delete_validator = NodeDeleteValidator(db=db, branch=branch)