infrahub recover merge CLI tool

backend/infrahub/cli/__init__.py

@@ -9,6 +9,7 @@
 from .db import app as db_app
 from .dev import app as dev_app
 from .events import app as events_app
+from .recover import app as recover_app
 from .server import app as server_app
 from .tasks import app as tasks_app
 from .upgrade import upgrade_cmd
@@ -28,6 +29,7 @@ def common(ctx: typer.Context) -> None:
 app.add_typer(tasks_app, name="tasks", hidden=True)
 app.add_typer(dev_app, name="dev", help="Internal development commands.")
 app.command(name="upgrade")(upgrade_cmd)
+app.add_typer(recover_app, name="recover", help="Recover from failed operations.")
 
 
 async def _init_shell(config_file: str) -> None:

backend/infrahub/cli/recover.py

@@ -0,0 +1,147 @@
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+import typer
+from infrahub_sdk.async_typer import AsyncTyper
+from rich.console import Console
+
+from infrahub import config
+from infrahub.components import ComponentType
+from infrahub.core.branch import Branch
+from infrahub.core.initialization import initialize_registry
+from infrahub.core.merge.failure_identifier import MergeFailureIdentifier
+from infrahub.core.merge.failure_recoverer import MergeFailureRecoverer, RecoveryOutcome, RecoveryReport
+from infrahub.core.merge.write_blocker import MergeWriteBlocker
+from infrahub.core.registry import registry
+from infrahub.core.schema import SchemaRoot, core_models, internal_schema
+from infrahub.core.schema.manager import SchemaManager
+from infrahub.dependencies.registry import build_component_registry
+from infrahub.workers.dependencies import get_cache, get_component, set_component_type
+
+if TYPE_CHECKING:
+    from infrahub.cli.context import CliContext
+
+app = AsyncTyper()
+
+
+# The callback registers this Typer app as a command group so that subcommands (e.g. ``merge``) nest
+# under ``recover`` and it carries the group help. A single-subcommand group would otherwise collapse
+# to the subcommand itself, losing the ``recover`` grouping.
+@app.callback()
+def callback() -> None:
+    """Recover from failed operations."""
+
+
+def _print_report(console: Console, report: RecoveryReport) -> None:
+    if report.outcome is RecoveryOutcome.NOTHING_TO_RECOVER:
+        console.print("[green]No failed merge to recover.[/green]")
+    elif report.outcome is RecoveryOutcome.ORPHANED_CLEARED:
+        console.print("[yellow]Cleared a stale merge-protection marker; there was no branch to recover.[/yellow]")
+    elif report.outcome is RecoveryOutcome.RECOVERED:
+        console.print(f"[bold green]Recovered the failed merge on branch '{report.branch}'.[/bold green]")
+        if report.proposed_change is not None:
+            console.print(f"Proposed change '{report.proposed_change}' was reset to OPEN.")
+        console.print("Writes to the default branch are allowed again.")
+    elif report.outcome is RecoveryOutcome.FAILED:
+        console.print(
+            f"[red]Recovery of branch '{report.branch}' failed. The branch stays protected; "
+            f"review the logs and retry.[/red]"
+        )
+
+
+@app.command(name="merge")
+async def recover_cmd(
+    ctx: typer.Context,
+    branch: str | None = typer.Argument(
+        None, help="Name of the branch to recover; if omitted, the failed merge is auto-detected."
+    ),
+    yes: bool = typer.Option(False, "--yes", "-y", help="Skip the confirmation prompt."),
+    force: bool = typer.Option(
+        False,
+        "--force",
+        "-f",
+        help="Recover even when the merge lock is absent/ambiguous, not just when the worker is confirmed dead.",
+    ),
+    config_file: str = typer.Argument("infrahub.toml", envvar="INFRAHUB_CONFIG"),
+) -> None:
+    """Recover a failed branch merge.
+
+    Roll back the partial graph merge and reset the branch (and any associated proposed change) to
+    OPEN, then lift the write protection so the default branch is writable again. Idempotent: a run
+    with nothing to recover reports so and makes no changes. By default only a merge whose worker is
+    confirmed dead is recovered; pass --force to also recover a merge stuck with an absent/ambiguous
+    lock.
+
+    Raises:
+        Exit: When the recovery does not complete successfully (raises typer.Exit to signal a
+            non-zero exit status to the caller).
+
+    """
+    logging.getLogger("infrahub").setLevel(logging.WARNING)
+    logging.getLogger("neo4j").setLevel(logging.ERROR)
+    logging.getLogger("prefect").setLevel(logging.ERROR)
+
+    console = Console()
+
+    config.load_and_exit(config_file_name=config_file)
+
+    context: CliContext = ctx.obj
+    db = await context.init_db(retry=1)
+
+    try:
+        await initialize_registry(db=db)
+        set_component_type(component_type=ComponentType.API_SERVER)
+        build_component_registry()
+
+        cache = await get_cache()
+        component = await get_component()
+        merge_write_blocker = MergeWriteBlocker(cache=cache)
+        default_branch = await Branch.get_by_name(db=db, name=registry.default_branch)
+
+        # The proposed-change lookup during recovery resolves its schema through ``db.schema``, which
+        # prefers a schema attached to the db over the global registry. Build the internal and core
+        # schema in a fresh, registry-independent manager and attach it to the db, so recovery works
+        # without depending on (or mutating) the global registry or any user-defined schema.
+        schema_manager = SchemaManager()
+        schema_branch = schema_manager.register_schema(schema=SchemaRoot(**internal_schema), branch=default_branch.name)
+        schema_branch.load_schema(schema=SchemaRoot(**core_models))
+        schema_branch.process()
+        db.add_schema(schema=schema_branch, name=default_branch.name)
+
+        identifier = MergeFailureIdentifier(
+            db=db,
+            cache=cache,
+            component=component,
+            merge_write_blocker=merge_write_blocker,
+            default_branch=default_branch,
+            grace_period_seconds=config.SETTINGS.main.merge_failure_grace_period_seconds,
+        )
+        recoverer = MergeFailureRecoverer(
+            db=db,
+            merge_write_blocker=merge_write_blocker,
+            identifier=identifier,
+            default_branch=default_branch,
+            cache=cache,
+        )
+
+        preview = await recoverer.preview(force=force, branch_name=branch)
+        if preview.outcome is RecoveryOutcome.RECOVERABLE:
+            console.print(f"A failed merge was found on branch [bold]'{preview.branch}'[/bold].")
+            if preview.merge_started_at is not None:
+                console.print(f"The merge started at {preview.merge_started_at}.")
+            if preview.proposed_change is not None:
+                console.print(f"Associated proposed change: {preview.proposed_change}.")
+
+            if not yes and not typer.confirm("Recover this failed merge?"):
+                console.print("Aborted; no changes were made.")
+                raise typer.Exit(code=1)
+
+        # Pin recovery to the branch that was previewed and confirmed
+        report = await recoverer.recover(force=force, branch_name=branch or preview.branch)
+        _print_report(console=console, report=report)
+        if report.outcome is RecoveryOutcome.FAILED:
+            raise typer.Exit(code=1)
+    finally:
+        await db.close()

backend/infrahub/core/diff/merger/merger.py

@@ -2,7 +2,6 @@
 
 from typing import TYPE_CHECKING
 
-from infrahub.core import registry
 from infrahub.core.diff.model.path import BranchTrackingId
 from infrahub.core.diff.query.bulk_merge import (
     BulkMergeAttributePropertyEdgesQuery,
@@ -123,11 +122,6 @@ async def merge_graph(self, at: Timestamp) -> None:
                 )
                 await metadata_query.execute(db=self.db)
 
-        branched_from = at.subtract(microseconds=1)
-        self.source_branch.branched_from = branched_from.to_string()
-        await self.source_branch.save(db=self.db)
-        registry.branch[self.source_branch.name] = self.source_branch
-
         log.info("Graph merge complete")
 
     @retry_db_transaction(name="bulk_merge_node_existence")

backend/infrahub/core/merge/failure_identifier.py

@@ -1,7 +1,5 @@
 from __future__ import annotations
 
-from dataclasses import dataclass
-from enum import Enum
 from typing import TYPE_CHECKING
 
 from infrahub import config
@@ -24,22 +22,6 @@
 log = get_logger()
 
 
-class RecoveryOutcome(Enum):
-    NOTHING_TO_RECOVER = "nothing_to_recover"
-    DECLINED = "declined"
-    RECOVERED = "recovered"
-    ORPHANED_CLEARED = "orphaned_cleared"
-    FAILED = "failed"
-
-
-@dataclass(frozen=True)
-class RecoveryReport:
-    outcome: RecoveryOutcome
-    branch: str | None
-    proposed_change: str | None
-    merge_started_at: str | None
-
-
 class MergeFailureIdentifier:
     """Act on a merge whose worker died mid-flight.
 
@@ -55,43 +37,42 @@ def __init__(
         cache: InfrahubCache,
         component: InfrahubComponent,
         merge_write_blocker: MergeWriteBlocker,
+        default_branch: Branch,
         grace_period_seconds: int,
     ) -> None:
         self.db = db
         self.cache = cache
         self.component = component
         self.merge_write_blocker = merge_write_blocker
+        self.default_branch = default_branch
         self.grace_period_seconds = grace_period_seconds
 
-    def should_mark_as_failed_merge(
-        self,
-        *,
-        status: BranchStatus,
-        lock_holder_worker_id: str | None,
-        active_worker_ids: set[str],
-        merge_started_at: Timestamp | None,
-        now: Timestamp,
-    ) -> bool:
-        """Decide whether a branch represents a dead merge that must be flagged ``MERGE_FAILED``.
+    async def should_mark_as_failed_merge(self, *, branch: Branch, now: Timestamp) -> bool:
+        """Decide whether a branch stuck in ``MERGING`` is a dead merge whose worker died mid-flight.
 
         A merge has failed when all of the following hold:
           - the branch is still ``MERGING`` (it never reached ``MERGED`` or rolled back to ``OPEN``),
-          - the global merge lock is *held* (``lock_holder_worker_id is not None``) — a dead worker
-            cannot release it, so a genuine failure leaves it held; an absent lock is ambiguous (a
-            cache flush during a live merge would look the same) and is deliberately not auto-flagged,
+          - the global merge lock is *held* — a dead worker cannot release it, so a genuine failure
+            leaves it held; an absent lock is ambiguous (a cache flush during a live merge would look
+            the same) and is deliberately not flagged,
           - the lock holder is not among the live workers (the holder died), and
           - the merge has been running longer than the grace period, which absorbs a transient
             worker-heartbeat write blip so a healthy merge is never mis-flagged.
         """
-        if status != BranchStatus.MERGING:
+        if branch.status != BranchStatus.MERGING:
             return False
+        lock_holder_worker_id = await self.merge_lock_holder()
         if lock_holder_worker_id is None:
             return False
-        if lock_holder_worker_id in active_worker_ids:
+        if lock_holder_worker_id in await self._active_worker_ids():
             return False
-        if merge_started_at is None:
+        if branch.merge_started_at is None:
             return False
-        return merge_started_at.add(seconds=self.grace_period_seconds) < now
+        return Timestamp(branch.merge_started_at).add(seconds=self.grace_period_seconds) < now
+
+    async def merge_lock_holder(self) -> str | None:
+        """Return the worker id currently holding the global merge lock, or ``None`` if it is unheld."""
+        return await get_merge_lock_holder_worker_id(cache=self.cache)
 
     async def scan(self) -> str | None:
         """Detect a dead merge and reconcile the protection key.
@@ -110,21 +91,9 @@ async def _detect_and_mark(self) -> str | None:
         if not merging:
             return None
 
-        # The liveness signal is the worker holding the global merge lock: a dead holder means a
-        # crashed merge.
-        lock_holder_worker_id = await get_merge_lock_holder_worker_id(cache=self.cache)
-        active_worker_ids = await self._active_worker_ids()
         now = Timestamp()
-
         for branch in merging:
-            merge_started_at = Timestamp(branch.merge_started_at) if branch.merge_started_at else None
-            if not self.should_mark_as_failed_merge(
-                status=branch.status,
-                lock_holder_worker_id=lock_holder_worker_id,
-                active_worker_ids=active_worker_ids,
-                merge_started_at=merge_started_at,
-                now=now,
-            ):
+            if not await self.should_mark_as_failed_merge(branch=branch, now=now):
                 continue
 
             branch.status = BranchStatus.MERGE_FAILED
@@ -175,17 +144,19 @@ async def _reconcile_protection_key(self) -> None:
             await self.merge_write_blocker.set(branch=protected.name, state=expected_state)
 
     async def _active_worker_ids(self) -> set[str]:
-        workers = await self.component.list_workers(branch=registry.default_branch, schema_hash=False)
+        workers = await self.component.list_workers(branch=self.default_branch.name, schema_hash=False)
         return {worker.id for worker in workers if worker.active}
 
 
 async def scan_for_failed_merges(db: InfrahubDatabase, service: InfrahubServices) -> str | None:
-    """Build a ``MergeFailureRecovery`` from the running service and run one detection scan."""
-    recovery = MergeFailureIdentifier(
+    """Build the identifier from the running service and run one detection scan."""
+    default_branch = await registry.get_branch(db=db, branch=registry.default_branch)
+    identifier = MergeFailureIdentifier(
         db=db,
         cache=service.cache,
         component=service.component,
         merge_write_blocker=MergeWriteBlocker(cache=service.cache),
+        default_branch=default_branch,
         grace_period_seconds=config.SETTINGS.main.merge_failure_grace_period_seconds,
     )
-    return await recovery.scan()
+    return await identifier.scan()