infrahub-server 1.1.7__py3-none-any.whl → 1.2.0b1__py3-none-any.whl

infrahub/core/diff/repository/repository.py

@@ -1,10 +1,15 @@
 from collections import defaultdict
 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 DiffSummaryCountsEnricherQuery
+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
@@ -14,6 +19,7 @@ from infrahub.log import get_logger
 from ..model.path import (
     ConflictSelection,
     EnrichedDiffConflict,
+    EnrichedDiffNode,
     EnrichedDiffRoot,
     EnrichedDiffRootMetadata,
     EnrichedDiffs,
@@ -42,11 +48,10 @@ 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,
@@ -154,18 +159,23 @@ class DiffRepository:
             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()],
+            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,
@@ -225,43 +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 | EnrichedDiffsMetadata, do_summary_counts: bool = True) -> None:
+    @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)
         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)
+        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)
+
+    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
-        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})...")
+
+        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}...")
-            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.")
+            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:
-            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,
-            )
+            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,
@@ -325,7 +403,7 @@ class DiffRepository:
         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(
@@ -446,8 +524,48 @@ class DiffRepository:
         diff_id: str | None = None,
         node_uuids: list[str] | None = None,
     ) -> None:
-        log.info("Updating summary counts...")
-        query = await DiffSummaryCountsEnricherQuery.init(
+        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,
@@ -455,4 +573,4 @@ class DiffRepository:
             node_uuids=node_uuids,
         )
         await query.execute(db=self.db)
-        log.info("Summary counts updated...")
+        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)

infrahub/core/merge.py

@@ -19,6 +19,7 @@ if TYPE_CHECKING:
     from infrahub.core.branch import Branch
     from infrahub.core.diff.coordinator import DiffCoordinator
     from infrahub.core.diff.merger.merger import DiffMerger
+    from infrahub.core.diff.model.path import EnrichedDiffRoot
     from infrahub.core.diff.repository.repository import DiffRepository
     from infrahub.core.models import SchemaUpdateConstraintInfo, SchemaUpdateMigrationInfo
     from infrahub.core.schema.manager import SchemaDiff
@@ -26,6 +27,7 @@ if TYPE_CHECKING:
     from infrahub.database import InfrahubDatabase
     from infrahub.services import InfrahubServices
 
+
 log = get_logger()
 
 
@@ -172,7 +174,7 @@ class BranchMerger:
     async def merge(
         self,
         at: Optional[Union[str, Timestamp]] = None,
-    ) -> None:
+    ) -> EnrichedDiffRoot:
         """Merge the current branch into main."""
         if self.source_branch.name == registry.default_branch:
             raise ValidationError(f"Unable to merge the branch '{self.source_branch.name}' into itself")
@@ -198,8 +200,9 @@ class BranchMerger:
         # TODO need to find a way to properly communicate back to the user any issue that could come up during the merge
         # From the Graph or From the repositories
         self._merge_at = Timestamp(at)
-        await self.diff_merger.merge_graph(at=self._merge_at)
+        branch_diff = await self.diff_merger.merge_graph(at=self._merge_at)
         await self.merge_repositories()
+        return branch_diff
 
     async def rollback(self) -> None:
         await self.diff_merger.rollback(at=self._merge_at)

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_add_generate_template_attr 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,
 ]
 
 

infrahub/core/migrations/graph/m001_add_version_to_graph.py

@@ -16,7 +16,7 @@ class Migration001Query01(Query):
     name = "migration_001_01"
     type: QueryType = QueryType.WRITE
 
-    async def query_init(self, db: InfrahubDatabase, **kwargs: dict[str, Any]) -> None:
+    async def query_init(self, db: InfrahubDatabase, **kwargs: dict[str, Any]) -> None:  # noqa: ARG002
         query = """
         MATCH (root:Root)
         SET root.graph_version = 1

infrahub/core/migrations/graph/m002_attribute_is_default.py

@@ -15,7 +15,7 @@ class Migration002Query01(Query):
     name = "migration_002_01"
     type: QueryType = QueryType.WRITE
 
-    async def query_init(self, db: InfrahubDatabase, **kwargs: dict[str, Any]) -> None:
+    async def query_init(self, db: InfrahubDatabase, **kwargs: dict[str, Any]) -> None:  # noqa: ARG002
         query = """
         MATCH (a:AttributeValue)
         WHERE a.is_default IS NULL
@@ -30,7 +30,7 @@ class Migration002(GraphMigration):
     queries: Sequence[type[Query]] = [Migration002Query01]
     minimum_version: int = 1
 
-    async def validate_migration(self, db: InfrahubDatabase) -> MigrationResult:
+    async def validate_migration(self, db: InfrahubDatabase) -> MigrationResult:  # noqa: ARG002
         result = MigrationResult()
 
         return result

infrahub/core/migrations/graph/m003_relationship_parent_optional.py

@@ -16,7 +16,7 @@ class Migration003Query01(Query):
     name = "migration_003_01"
     type: QueryType = QueryType.WRITE
 
-    async def query_init(self, db: InfrahubDatabase, **kwargs: dict[str, Any]) -> None:
+    async def query_init(self, db: InfrahubDatabase, **kwargs: dict[str, Any]) -> None:  # noqa: ARG002
         at = Timestamp()
         filters, params = at.get_query_filter_path()
 
@@ -49,7 +49,7 @@ class Migration003(GraphMigration):
     queries: Sequence[type[Query]] = [Migration003Query01]
     minimum_version: int = 2
 
-    async def validate_migration(self, db: InfrahubDatabase) -> MigrationResult:
+    async def validate_migration(self, db: InfrahubDatabase) -> MigrationResult:  # noqa: ARG002
         result = MigrationResult()
 
         return result

infrahub/core/migrations/graph/m004_add_attr_documentation.py

@@ -43,6 +43,6 @@ class Migration004(InternalSchemaMigration):
         ]
         return cls(migrations=migrations, **kwargs)  # type: ignore[arg-type]
 
-    async def validate_migration(self, db: InfrahubDatabase) -> MigrationResult:
+    async def validate_migration(self, db: InfrahubDatabase) -> MigrationResult:  # noqa: ARG002
         result = MigrationResult()
         return result

infrahub/core/migrations/graph/m005_add_rel_read_only.py

@@ -35,6 +35,6 @@ class Migration005(InternalSchemaMigration):
         ]
         return cls(migrations=migrations, **kwargs)  # type: ignore[arg-type]
 
-    async def validate_migration(self, db: InfrahubDatabase) -> MigrationResult:
+    async def validate_migration(self, db: InfrahubDatabase) -> MigrationResult:  # noqa: ARG002
         result = MigrationResult()
         return result

infrahub/core/migrations/graph/m006_add_rel_on_delete.py

@@ -35,6 +35,6 @@ class Migration006(InternalSchemaMigration):
         ]
         return cls(migrations=migrations, **kwargs)  # type: ignore[arg-type]
 
-    async def validate_migration(self, db: InfrahubDatabase) -> MigrationResult:
+    async def validate_migration(self, db: InfrahubDatabase) -> MigrationResult:  # noqa: ARG002
         result = MigrationResult()
         return result

infrahub/core/migrations/graph/m007_add_rel_allow_override.py

@@ -43,6 +43,6 @@ class Migration007(InternalSchemaMigration):
         ]
         return cls(migrations=migrations, **kwargs)  # type: ignore[arg-type]
 
-    async def validate_migration(self, db: InfrahubDatabase) -> MigrationResult:
+    async def validate_migration(self, db: InfrahubDatabase) -> MigrationResult:  # noqa: ARG002
         result = MigrationResult()
         return result