infrahub-server 1.3.1__py3-none-any.whl → 1.3.3__py3-none-any.whl

infrahub/cli/db.py

@@ -8,7 +8,7 @@ from csv import DictReader, DictWriter
 from datetime import datetime, timezone
 from enum import Enum
 from pathlib import Path
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Sequence
 
 import typer
 import ujson
@@ -38,7 +38,7 @@ from infrahub.core.initialization import (
     initialization,
     initialize_registry,
 )
-from infrahub.core.migrations.graph import get_graph_migrations
+from infrahub.core.migrations.graph import get_graph_migrations, get_migration_by_number
 from infrahub.core.migrations.schema.models import SchemaApplyMigrationData
 from infrahub.core.migrations.schema.tasks import schema_apply_migrations
 from infrahub.core.schema import SchemaRoot, core_models, internal_schema
@@ -58,8 +58,15 @@ from infrahub.services.adapters.workflow.local import WorkflowLocalExecution
 from .constants import ERROR_BADGE, FAILED_BADGE, SUCCESS_BADGE
 from .patch import patch_app
 
+
+def get_timestamp_string() -> str:
+    """Generate a timestamp string in the format YYYYMMDD-HHMMSS."""
+    return datetime.now(tz=timezone.utc).strftime("%Y%m%d-%H%M%S")
+
+
 if TYPE_CHECKING:
     from infrahub.cli.context import CliContext
+    from infrahub.core.migrations.shared import ArbitraryMigration, GraphMigration, InternalSchemaMigration
     from infrahub.database import InfrahubDatabase
     from infrahub.database.index import IndexManagerBase
 
@@ -154,6 +161,7 @@ async def migrate_cmd(
     ctx: typer.Context,
     check: bool = typer.Option(False, help="Check the state of the database without applying the migrations."),
     config_file: str = typer.Argument("infrahub.toml", envvar="INFRAHUB_CONFIG"),
+    migration_number: int | None = typer.Option(None, help="Apply a specific migration by number"),
 ) -> None:
     """Check the current format of the internal graph and apply the necessary migrations"""
     logging.getLogger("infrahub").setLevel(logging.WARNING)
@@ -165,7 +173,7 @@ async def migrate_cmd(
     context: CliContext = ctx.obj
     dbdriver = await context.init_db(retry=1)
 
-    await migrate_database(db=dbdriver, initialize=True, check=check)
+    await migrate_database(db=dbdriver, initialize=True, check=check, migration_number=migration_number)
 
     await dbdriver.close()
 
@@ -287,7 +295,9 @@ async def index(
     await dbdriver.close()
 
 
-async def migrate_database(db: InfrahubDatabase, initialize: bool = False, check: bool = False) -> bool:
+async def migrate_database(
+    db: InfrahubDatabase, initialize: bool = False, check: bool = False, migration_number: int | str | None = None
+) -> bool:
     """Apply the latest migrations to the database, this function will print the status directly in the console.
 
     Returns a boolean indicating whether a migration failed or if all migrations succeeded.
@@ -295,6 +305,7 @@ async def migrate_database(db: InfrahubDatabase, initialize: bool = False, check
     Args:
         db: The database object.
         check: If True, the function will only check the status of the database and not apply the migrations. Defaults to False.
+        migration_number: If provided, the function will only apply the migration with the given number. Defaults to None.
     """
     rprint("Checking current state of the Database")
 
@@ -302,15 +313,28 @@ async def migrate_database(db: InfrahubDatabase, initialize: bool = False, check
         await initialize_registry(db=db)
 
     root_node = await get_root_node(db=db)
-    migrations = await get_graph_migrations(root=root_node)
-
-    if not migrations:
-        rprint(f"Database up-to-date (v{root_node.graph_version}), no migration to execute.")
-        return True
+    if migration_number:
+        migration = get_migration_by_number(migration_number)
+        migrations: Sequence[GraphMigration | InternalSchemaMigration | ArbitraryMigration] = [migration]
+        if check:
+            if root_node.graph_version > migration.minimum_version:
+                rprint(
+                    f"Migration {migration_number} already applied. To apply again, run the command without the --check flag."
+                )
+                return True
+            rprint(
+                f"Migration {migration_number} needs to be applied. Run `infrahub db migrate` to apply all outstanding migrations."
+            )
+            return False
+    else:
+        migrations = await get_graph_migrations(root=root_node)
+        if not migrations:
+            rprint(f"Database up-to-date (v{root_node.graph_version}), no migration to execute.")
+            return True
 
-    rprint(
-        f"Database needs to be updated (v{root_node.graph_version} -> v{GRAPH_VERSION}), {len(migrations)} migrations pending"
-    )
+        rprint(
+            f"Database needs to be updated (v{root_node.graph_version} -> v{GRAPH_VERSION}), {len(migrations)} migrations pending"
+        )
 
     if check:
         return True
@@ -529,7 +553,7 @@ WITH n, edges + root_edges AS edges, CASE
 END AS vertices
 RETURN vertices, edges
     """ % {"id_func": db.get_id_function_name()}
-    timestamp_str = datetime.now(tz=timezone.utc).strftime("%Y%m%d-%H%M%S")
+    timestamp_str = get_timestamp_string()
     export_dir /= Path(f"export-{timestamp_str}")
     if not export_dir.exists():
         export_dir.mkdir(parents=True)
@@ -737,3 +761,160 @@ async def load_export(db: InfrahubDatabase, export_dir: Path, query_limit: int =
             await load_edges(db=db, edge_type=edge_type, edge_dicts=edge_dicts)
     rprint("Edges loaded")
     rprint(f"{SUCCESS_BADGE} Export loaded")
+
+
+@app.command(name="check")
+async def check_cmd(
+    ctx: typer.Context,
+    output_dir: Path = typer.Option(  # noqa: B008
+        None, help="Directory to save detailed check results (defaults to infrahub_db_check_YYYYMMDD-HHMMSS)"
+    ),
+    config_file: str = typer.Option(
+        "infrahub.toml", envvar="INFRAHUB_CONFIG", help="Location of the configuration file to use for Infrahub"
+    ),
+) -> None:
+    """Run database sanity checks and output the results to the CSV files."""
+    logging.getLogger("infrahub").setLevel(logging.WARNING)
+    logging.getLogger("neo4j").setLevel(logging.ERROR)
+    logging.getLogger("prefect").setLevel(logging.ERROR)
+
+    config.load_and_exit(config_file_name=config_file)
+
+    # Create output directory if not specified
+    if output_dir is None:
+        timestamp_str = get_timestamp_string()
+        output_dir = Path(f"infrahub_db_check_{timestamp_str}")
+
+    if not output_dir.exists():
+        output_dir.mkdir(parents=True)
+
+    context: CliContext = ctx.obj
+    dbdriver = await context.init_db(retry=1)
+
+    await run_database_checks(db=dbdriver, output_dir=output_dir)
+
+    await dbdriver.close()
+
+
+async def run_database_checks(db: InfrahubDatabase, output_dir: Path) -> None:
+    """Run a series of database health checks and output the results to the terminal.
+
+    Args:
+        db: The database object.
+        output_dir: Directory to save detailed check results.
+    """
+    rprint("Running database health checks...")
+
+    # Check 1: Duplicate active relationships
+    rprint("\n[bold cyan]Check 1: Duplicate Active Relationships[/bold cyan]")
+    duplicate_active_rels_query = """
+    MATCH (a:Node)-[e1:IS_RELATED {status: "active"}]-(r:Relationship)-[e2:IS_RELATED {branch: e1.branch, status: "active"}]-(b:Node)
+    WHERE a.uuid < b.uuid
+    AND e1.to IS NULL
+    AND e2.to IS NULL
+    WITH DISTINCT a.uuid AS a_uuid,
+        b.uuid AS b_uuid,
+        r.name AS r_name,
+        e1.branch AS branch,
+        CASE
+        WHEN startNode(e1) = a AND startNode(e2) = r THEN "out"
+            WHEN startNode(e1) = r AND startNode(e2) = b THEN "in"
+            ELSE "bidir"
+        END AS direction,
+        count(*) AS num_paths,
+        collect(DISTINCT a.kind) AS a_kinds,
+        collect(DISTINCT b.kind) AS b_kinds
+    WHERE num_paths > 1
+    RETURN a_uuid, a_kinds, b_uuid, b_kinds, r_name, branch, direction, num_paths
+    """
+
+    results = await db.execute_query(query=duplicate_active_rels_query)
+    if results:
+        rprint(f"[red]Found {len(results)} duplicate active relationships[/red]")
+        # Write detailed results to file
+        output_file = output_dir / "duplicate_active_relationships.csv"
+        with output_file.open(mode="w", newline="") as f:
+            writer = DictWriter(
+                f, fieldnames=["a_uuid", "a_kinds", "b_uuid", "b_kinds", "r_name", "branch", "direction", "num_paths"]
+            )
+            writer.writeheader()
+            for result in results:
+                writer.writerow(dict(result))
+        rprint(f"  Detailed results written to: {output_file}")
+    else:
+        rprint(f"{SUCCESS_BADGE} No duplicate active relationships found")
+
+    # Check 2: Duplicated relationship nodes
+    rprint("\n[bold cyan]Check 2: Duplicated Relationship Nodes[/bold cyan]")
+    duplicate_rel_nodes_query = """
+    MATCH (r:Relationship)
+    WITH r.uuid AS r_uuid, COUNT(*) AS num_rels
+    WHERE num_rels > 1
+    MATCH (n:Node)-[:IS_RELATED]-(r:Relationship {uuid: r_uuid})
+    WITH DISTINCT r_uuid, n.uuid AS n_uuid, n.kind AS n_kind
+    WITH r_uuid, collect([n_uuid, n_kind]) AS node_details
+    RETURN r_uuid, node_details
+    """
+
+    results = await db.execute_query(query=duplicate_rel_nodes_query)
+    if results:
+        rprint(f"[red]Found {len(results)} duplicated relationship nodes[/red]")
+        # Write detailed results to file
+        output_file = output_dir / "duplicated_relationship_nodes.csv"
+        with output_file.open(mode="w", newline="") as f:
+            writer = DictWriter(f, fieldnames=["r_uuid", "node_details"])
+            writer.writeheader()
+            for result in results:
+                writer.writerow(dict(result))
+        rprint(f"  Detailed results written to: {output_file}")
+    else:
+        rprint(f"{SUCCESS_BADGE} No duplicated relationship nodes found")
+
+    # Check 3: Duplicated edges
+    rprint("\n[bold cyan]Check 3: Duplicated Edges[/bold cyan]")
+    duplicate_edges_query = """
+    MATCH (a)
+    CALL (a) {
+        MATCH (a)-[e]->(b)
+        WHERE elementId(a) < elementId(b)
+        WITH DISTINCT a, b, type(e) AS e_type, count(*) AS total_num_edges
+        WHERE total_num_edges > 1
+        MATCH (a)-[e]->(b)
+        WHERE type(e) = e_type
+        WITH
+            elementId(a) AS a_id,
+            labels(a) AS a_labels,
+            elementId(b) AS b_id,
+            labels(b) AS b_labels,
+            type(e) AS e_type,
+            e.branch AS branch,
+            e.status AS status,
+            e.from AS time,
+            collect(e) AS edges
+        WITH a_id, a_labels, b_id, b_labels, e_type, branch, status, time, size(edges) AS num_edges
+        WHERE num_edges > 1
+        WITH a_id, a_labels, b_id, b_labels, e_type, branch, status, time, num_edges
+        RETURN a_id, a_labels, b_id, b_labels, e_type, branch, status, time, num_edges
+    }
+    RETURN a_id, a_labels, b_id, b_labels, e_type, branch, status, time, num_edges
+    """
+
+    results = await db.execute_query(query=duplicate_edges_query)
+    if results:
+        rprint(f"[red]Found {len(results)} sets of duplicated edges[/red]")
+        # Write detailed results to file
+        output_file = output_dir / "duplicated_edges.csv"
+        with output_file.open(mode="w", newline="") as f:
+            writer = DictWriter(
+                f,
+                fieldnames=["a_id", "a_labels", "b_id", "b_labels", "e_type", "branch", "status", "time", "num_edges"],
+            )
+            writer.writeheader()
+            for result in results:
+                writer.writerow(dict(result))
+        rprint(f"  Detailed results written to: {output_file}")
+    else:
+        rprint(f"{SUCCESS_BADGE} No duplicated edges found")
+
+    rprint(f"\n{SUCCESS_BADGE} Database health checks completed")
+    rprint(f"Detailed results saved to: {output_dir.absolute()}")

infrahub/core/branch/enums.py

@@ -0,0 +1,8 @@
+from infrahub.utils import InfrahubStringEnum
+
+
+class BranchStatus(InfrahubStringEnum):
+    OPEN = "OPEN"
+    NEED_REBASE = "NEED_REBASE"
+    CLOSED = "CLOSED"
+    DELETING = "DELETING"

infrahub/core/branch/models.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import re
-from typing import TYPE_CHECKING, Any, Optional, Union
+from typing import TYPE_CHECKING, Any, Optional, Self, Union
 
 from pydantic import Field, field_validator
 
@@ -21,6 +21,8 @@ from infrahub.core.registry import registry
 from infrahub.core.timestamp import Timestamp
 from infrahub.exceptions import BranchNotFoundError, InitializationError, ValidationError
 
+from .enums import BranchStatus
+
 if TYPE_CHECKING:
     from infrahub.database import InfrahubDatabase
 
@@ -29,7 +31,7 @@ class Branch(StandardNode):
     name: str = Field(
         max_length=250, min_length=3, description="Name of the branch (git ref standard)", validate_default=True
     )
-    status: str = "OPEN"  # OPEN, CLOSED
+    status: BranchStatus = BranchStatus.OPEN
     description: str = ""
     origin_branch: str = "main"
     branched_from: Optional[str] = Field(default=None, validate_default=True)
@@ -131,14 +133,17 @@ class Branch(StandardNode):
         return True
 
     @classmethod
-    async def get_by_name(cls, name: str, db: InfrahubDatabase) -> Branch:
+    async def get_by_name(cls, name: str, db: InfrahubDatabase, ignore_deleting: bool = True) -> Branch:
         query = """
         MATCH (n:Branch)
         WHERE n.name = $name
+        AND NOT n.status IN $ignore_statuses
         RETURN n
         """
 
-        params = {"name": name}
+        params: dict[str, Any] = {"name": name}
+        if ignore_deleting:
+            params["ignore_statuses"] = [BranchStatus.DELETING.value]
 
         results = await db.execute_query(query=query, params=params, name="branch_get_by_name", type=QueryType.READ)
 
@@ -147,6 +152,20 @@ class Branch(StandardNode):
 
         return cls.from_db(results[0].values()[0])
 
+    @classmethod
+    async def get_list(
+        cls,
+        db: InfrahubDatabase,
+        limit: int = 1000,
+        ids: list[str] | None = None,
+        name: str | None = None,
+        **kwargs: dict[str, Any],
+    ) -> list[Self]:
+        branches = await super().get_list(db=db, limit=limit, ids=ids, name=name, **kwargs)
+        branches = [branch for branch in branches if branch.status != BranchStatus.DELETING]
+
+        return branches
+
     @classmethod
     def isinstance(cls, obj: Any) -> bool:
         return isinstance(obj, cls)
@@ -248,9 +267,13 @@ class Branch(StandardNode):
             raise ValidationError(f"Unable to delete {self.name} it is the default branch.")
         if self.is_global:
             raise ValidationError(f"Unable to delete {self.name} this is an internal branch.")
-        await super().delete(db=db)
+
+        self.status = BranchStatus.DELETING
+        await self.save(db=db)
+
         query = await DeleteBranchRelationshipsQuery.init(db=db, branch_name=self.name)
         await query.execute(db=db)
+        await super().delete(db=db)
 
     def get_query_filter_relationships(
         self, rel_labels: list, at: Optional[Union[Timestamp, str]] = None, include_outside_parentheses: bool = False

infrahub/core/branch/tasks.py

@@ -15,6 +15,7 @@ from infrahub.core.branch import Branch
 from infrahub.core.changelog.diff import DiffChangelogCollector, MigrationTracker
 from infrahub.core.constants import MutationAction
 from infrahub.core.diff.coordinator import DiffCoordinator
+from infrahub.core.diff.diff_locker import DiffLocker
 from infrahub.core.diff.ipam_diff_parser import IpamDiffParser
 from infrahub.core.diff.merger.merger import DiffMerger
 from infrahub.core.diff.model.path import BranchTrackingId, EnrichedDiffRoot, EnrichedDiffRootMetadata
@@ -31,7 +32,7 @@ from infrahub.dependencies.registry import get_component_registry
 from infrahub.events.branch_action import BranchCreatedEvent, BranchDeletedEvent, BranchMergedEvent, BranchRebasedEvent
 from infrahub.events.models import EventMeta, InfrahubEvent
 from infrahub.events.node_action import get_node_event
-from infrahub.exceptions import BranchNotFoundError, MergeFailedError, ValidationError
+from infrahub.exceptions import BranchNotFoundError, ValidationError
 from infrahub.graphql.mutations.models import BranchCreateModel  # noqa: TC001
 from infrahub.services import InfrahubServices  # noqa: TC001  needed for prefect flow
 from infrahub.workflows.catalogue import (
@@ -65,6 +66,7 @@ async def rebase_branch(branch: str, context: InfrahubContext, service: Infrahub
             diff_merger=diff_merger,
             diff_repository=diff_repository,
             source_branch=obj,
+            diff_locker=DiffLocker(),
             service=service,
         )
 
@@ -221,14 +223,10 @@ async def merge_branch(
                 diff_merger=diff_merger,
                 diff_repository=diff_repository,
                 source_branch=obj,
+                diff_locker=DiffLocker(),
                 service=service,
             )
-            try:
-                branch_diff = await merger.merge()
-            except Exception as exc:
-                log.exception("Merge failed, beginning rollback")
-                await merger.rollback()
-                raise MergeFailedError(branch_name=branch) from exc
+            branch_diff = await merger.merge()
             await merger.update_schema()
 
         changelog_collector = DiffChangelogCollector(diff=branch_diff, branch=obj, db=db)

infrahub/core/diff/calculator.py

@@ -181,8 +181,11 @@ class DiffCalculator:
         log.info("Diff property-level calculation queries for branch complete")
 
         if base_branch.name != diff_branch.name:
-            current_node_field_specifiers = diff_parser.get_current_node_field_specifiers()
             new_node_field_specifiers = diff_parser.get_new_node_field_specifiers()
+            current_node_field_specifiers = None
+            if previous_node_specifiers is not None:
+                current_node_field_specifiers = previous_node_specifiers - new_node_field_specifiers
+
             base_calculation_request = DiffCalculationRequest(
                 base_branch=base_branch,
                 diff_branch=base_branch,

infrahub/core/diff/coordinator.py

@@ -6,8 +6,7 @@ from uuid import uuid4
 
 from prefect import flow
 
-from infrahub import lock
-from infrahub.core.branch import Branch  # noqa: TC001
+from infrahub.core.branch import Branch
 from infrahub.core.timestamp import Timestamp
 from infrahub.exceptions import ValidationError
 from infrahub.log import get_logger
@@ -26,12 +25,14 @@ from .model.path import (
 
 if TYPE_CHECKING:
     from infrahub.core.node import Node
+    from infrahub.database import InfrahubDatabase
 
     from .calculator import DiffCalculator
     from .combiner import DiffCombiner
     from .conflict_transferer import DiffConflictTransferer
     from .conflicts_enricher import ConflictsEnricher
     from .data_check_synchronizer import DiffDataCheckSynchronizer
+    from .diff_locker import DiffLocker
     from .enricher.aggregated import AggregatedDiffEnricher
     from .enricher.labels import DiffLabelsEnricher
     from .repository.repository import DiffRepository
@@ -59,10 +60,9 @@ class EnrichedDiffRequest:
 
 
 class DiffCoordinator:
-    lock_namespace = "diff-update"
-
     def __init__(
         self,
+        db: InfrahubDatabase,
         diff_repo: DiffRepository,
         diff_calculator: DiffCalculator,
         diff_enricher: AggregatedDiffEnricher,
@@ -71,7 +71,9 @@ class DiffCoordinator:
         labels_enricher: DiffLabelsEnricher,
         data_check_synchronizer: DiffDataCheckSynchronizer,
         conflict_transferer: DiffConflictTransferer,
+        diff_locker: DiffLocker,
     ) -> None:
+        self.db = db
         self.diff_repo = diff_repo
         self.diff_calculator = diff_calculator
         self.diff_enricher = diff_enricher
@@ -80,7 +82,7 @@ class DiffCoordinator:
         self.labels_enricher = labels_enricher
         self.data_check_synchronizer = data_check_synchronizer
         self.conflict_transferer = conflict_transferer
-        self.lock_registry = lock.registry
+        self.diff_locker = diff_locker
 
     async def run_update(
         self,
@@ -113,37 +115,35 @@ class DiffCoordinator:
             name=name,
         )
 
-    def _get_lock_name(self, base_branch_name: str, diff_branch_name: str, is_incremental: bool) -> str:
-        lock_name = f"{base_branch_name}__{diff_branch_name}"
-        if is_incremental:
-            lock_name += "__incremental"
-        return lock_name
-
     async def update_branch_diff(self, base_branch: Branch, diff_branch: Branch) -> EnrichedDiffRootMetadata:
+        tracking_id = BranchTrackingId(name=diff_branch.name)
         log.info(f"Received request to update branch diff for {base_branch.name} - {diff_branch.name}")
-        incremental_lock_name = self._get_lock_name(
-            base_branch_name=base_branch.name, diff_branch_name=diff_branch.name, is_incremental=True
-        )
-        existing_incremental_lock = self.lock_registry.get_existing(
-            name=incremental_lock_name, namespace=self.lock_namespace
+        existing_incremental_lock = self.diff_locker.get_existing_lock(
+            target_branch_name=base_branch.name, source_branch_name=diff_branch.name, is_incremental=True
         )
         if existing_incremental_lock and await existing_incremental_lock.locked():
             log.info(f"Branch diff update for {base_branch.name} - {diff_branch.name} already in progress")
-            async with self.lock_registry.get(name=incremental_lock_name, namespace=self.lock_namespace):
+            async with self.diff_locker.acquire_lock(
+                target_branch_name=base_branch.name, source_branch_name=diff_branch.name, is_incremental=True
+            ):
                 log.info(f"Existing branch diff update for {base_branch.name} - {diff_branch.name} complete")
-                return await self.diff_repo.get_one(
-                    tracking_id=BranchTrackingId(name=diff_branch.name), diff_branch_name=diff_branch.name
-                )
-        general_lock_name = self._get_lock_name(
-            base_branch_name=base_branch.name, diff_branch_name=diff_branch.name, is_incremental=False
-        )
+                return await self.diff_repo.get_one(tracking_id=tracking_id, diff_branch_name=diff_branch.name)
         from_time = Timestamp(diff_branch.get_branched_from())
         to_time = Timestamp()
-        tracking_id = BranchTrackingId(name=diff_branch.name)
         async with (
-            self.lock_registry.get(name=general_lock_name, namespace=self.lock_namespace),
-            self.lock_registry.get(name=incremental_lock_name, namespace=self.lock_namespace),
+            self.diff_locker.acquire_lock(
+                target_branch_name=base_branch.name, source_branch_name=diff_branch.name, is_incremental=True
+            ),
+            self.diff_locker.acquire_lock(
+                target_branch_name=base_branch.name, source_branch_name=diff_branch.name, is_incremental=False
+            ),
         ):
+            refreshed_branch = await Branch.get_by_name(db=self.db, name=diff_branch.name)
+            if refreshed_branch.get_branched_from() != diff_branch.get_branched_from():
+                log.info(
+                    f"Branch {diff_branch.name} was merged or rebased while waiting for lock, returning latest diff"
+                )
+                return await self.diff_repo.get_one(tracking_id=tracking_id, diff_branch_name=diff_branch.name)
             log.info(f"Acquired lock to run branch diff update for {base_branch.name} - {diff_branch.name}")
             enriched_diffs, node_identifiers_to_drop = await self._update_diffs(
                 base_branch=base_branch,
@@ -169,10 +169,9 @@ class DiffCoordinator:
         name: str,
     ) -> EnrichedDiffRootMetadata:
         tracking_id = NameTrackingId(name=name)
-        general_lock_name = self._get_lock_name(
-            base_branch_name=base_branch.name, diff_branch_name=diff_branch.name, is_incremental=False
-        )
-        async with self.lock_registry.get(name=general_lock_name, namespace=self.lock_namespace):
+        async with self.diff_locker.acquire_lock(
+            target_branch_name=base_branch.name, source_branch_name=diff_branch.name, is_incremental=False
+        ):
             log.info(f"Acquired lock to run arbitrary diff update for {base_branch.name} - {diff_branch.name}")
             enriched_diffs, node_identifiers_to_drop = await self._update_diffs(
                 base_branch=base_branch,
@@ -196,10 +195,9 @@ class DiffCoordinator:
         diff_branch: Branch,
         diff_id: str,
     ) -> EnrichedDiffRoot:
-        general_lock_name = self._get_lock_name(
-            base_branch_name=base_branch.name, diff_branch_name=diff_branch.name, is_incremental=False
-        )
-        async with self.lock_registry.get(name=general_lock_name, namespace=self.lock_namespace):
+        async with self.diff_locker.acquire_lock(
+            target_branch_name=base_branch.name, source_branch_name=diff_branch.name, is_incremental=False
+        ):
             log.info(f"Acquired lock to recalculate diff for {base_branch.name} - {diff_branch.name}")
             current_branch_diff = await self.diff_repo.get_one(diff_branch_name=diff_branch.name, diff_id=diff_id)
             current_base_diff = await self.diff_repo.get_one(

infrahub/core/diff/diff_locker.py

@@ -0,0 +1,26 @@
+from infrahub import lock
+
+
+class DiffLocker:
+    lock_namespace = "diff-update"
+
+    def __init__(self) -> None:
+        self.lock_registry = lock.registry
+
+    def get_lock_name(self, base_branch_name: str, diff_branch_name: str, is_incremental: bool) -> str:
+        lock_name = f"{base_branch_name}__{diff_branch_name}"
+        if is_incremental:
+            lock_name += "__incremental"
+        return lock_name
+
+    def get_existing_lock(
+        self, target_branch_name: str, source_branch_name: str, is_incremental: bool = False
+    ) -> lock.InfrahubLock | None:
+        name = self.get_lock_name(target_branch_name, source_branch_name, is_incremental)
+        return self.lock_registry.get_existing(name=name, namespace=self.lock_namespace)
+
+    def acquire_lock(
+        self, target_branch_name: str, source_branch_name: str, is_incremental: bool = False
+    ) -> lock.InfrahubLock:
+        name = self.get_lock_name(target_branch_name, source_branch_name, is_incremental)
+        return self.lock_registry.get(name=name, namespace=self.lock_namespace)