Add new flow to run migrations on a branch

This is done by instantiating a `MigrationRunner` which will take care
of determining what migrations need to run to then run them if any.

Author: gmazoyer

backend/infrahub/cli/db.py

@@ -20,9 +20,11 @@
 from rich.table import Table
 
 from infrahub import config
+from infrahub.auth import AccountSession, AuthType
+from infrahub.context import InfrahubContext
 from infrahub.core import registry
 from infrahub.core.branch import Branch
-from infrahub.core.branch.enums import BranchStatus
+from infrahub.core.branch.tasks import rebase_branch
 from infrahub.core.constants import GLOBAL_BRANCH_NAME
 from infrahub.core.graph import GRAPH_VERSION
 from infrahub.core.graph.constraints import ConstraintManagerBase, ConstraintManagerMemgraph, ConstraintManagerNeo4j
@@ -44,7 +46,6 @@
 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.migrations.shared import MigrationWithRebase
 from infrahub.core.schema import SchemaRoot, core_models, internal_schema
 from infrahub.core.schema.definitions.deprecated import deprecated_models
 from infrahub.core.schema.manager import SchemaManager
@@ -54,6 +55,7 @@
 from infrahub.database import DatabaseType
 from infrahub.database.memgraph import IndexManagerMemgraph
 from infrahub.database.neo4j import IndexManagerNeo4j
+from infrahub.exceptions import ValidationError
 from infrahub.log import get_logger
 
 from .constants import ERROR_BADGE, FAILED_BADGE, SUCCESS_BADGE
@@ -69,7 +71,12 @@ def get_timestamp_string() -> str:
 
 if TYPE_CHECKING:
     from infrahub.cli.context import CliContext
-    from infrahub.core.migrations.shared import ArbitraryMigration, GraphMigration, InternalSchemaMigration
+    from infrahub.core.migrations.shared import (
+        ArbitraryMigration,
+        GraphMigration,
+        InternalSchemaMigration,
+        MigrationWithRebase,
+    )
     from infrahub.database import InfrahubDatabase
     from infrahub.database.index import IndexManagerBase
 
@@ -421,51 +428,30 @@ async def migrate_database(
     return True
 
 
-async def rebase_and_migrate_branches(db: InfrahubDatabase, current_graph_version: int) -> bool:
-    """Only applies migrations that aim at rebasing branches."""
-    branches = [
-        b
-        for b in await Branch.get_list(db=db)
-        if b.name not in [registry.default_branch, GLOBAL_BRANCH_NAME]
-        and (not b.graph_version or b.graph_version < current_graph_version)
-    ]
-
+async def trigger_rebase_branches(db: InfrahubDatabase) -> None:
+    """Trigger rebase of non-default branches, also triggering migrations in the process."""
+    branches = [b for b in await Branch.get_list(db=db) if b.name not in [registry.default_branch, GLOBAL_BRANCH_NAME]]
     if not branches:
-        return True
+        return
 
     rprint(f"Planning rebase and migrations for {len(branches)} branches: {', '.join([b.name for b in branches])}")
 
     for branch in branches:
-        migrations = [
-            m
-            for m in await detect_migration_to_run(current_graph_version=branch.graph_version or 0)
-            if isinstance(m, MigrationWithRebase)
-        ]
-        rprint(
-            f"Detected {len(migrations)} migrations to run against '{branch.name}' (ID: {branch.uuid}): {', '.join([m.name for m in migrations])}"
-        )
-
-        for migration in migrations:
-            execution_result = await migration.execute_against_branch(db=db, branch=branch)
-            validation_result = None
-
-            if execution_result.success:
-                validation_result = await migration.validate_migration(db=db)
-                if validation_result.success and branch.status != BranchStatus.NEED_UPGRADE_REBASE:
-                    branch.graph_version = migration.minimum_version + 1
-                    await branch.save(db=db)
-                    rprint(f"Migration: {migration.name} {SUCCESS_BADGE}")
-
-            if not execution_result.success or (validation_result and not validation_result.success):
-                rprint(f"Migration: {migration.name} {FAILED_BADGE}")
-                for error in execution_result.errors:
-                    rprint(f"  {error}")
-                if validation_result and not validation_result.success:
-                    for error in validation_result.errors:
-                        rprint(f"  {error}")
-                return False
-
-    return True
+        if branch.graph_version == GRAPH_VERSION:
+            continue
+
+        rprint(f"Rebasing branch '{branch.name}' (ID: {branch.uuid})...", end="")
+        try:
+            await rebase_branch(
+                branch=branch.name,
+                context=InfrahubContext.init(
+                    branch=branch, account=AccountSession(auth_type=AuthType.NONE, authenticated=False, account_id="")
+                ),
+                send_events=False,
+            )
+            rprint("done")
+        except ValidationError:
+            rprint("failed")
 
 
 async def initialize_internal_schema() -> None:

backend/infrahub/cli/upgrade.py

@@ -35,7 +35,7 @@
     detect_migration_to_run,
     initialize_internal_schema,
     migrate_database,
-    rebase_and_migrate_branches,
+    trigger_rebase_branches,
     update_core_schema,
 )
 
@@ -112,13 +112,8 @@ async def upgrade_cmd(
     # -------------------------------------------
     # Perform branch rebase and apply migrations to them
     # -------------------------------------------
-    if rebase_branches and not await rebase_and_migrate_branches(
-        db=dbdriver, current_graph_version=root_node.graph_version
-    ):
-        # A migration failed, stop the upgrade process
-        rprint("Upgrade cancelled due to branch rebase and migration failure.")
-        await dbdriver.close()
-        return
+    if rebase_branches:
+        await trigger_rebase_branches(db=dbdriver)
 
     await dbdriver.close()
 

backend/infrahub/core/branch/tasks.py

@@ -12,6 +12,7 @@
 from infrahub.context import InfrahubContext  # noqa: TC001  needed for prefect flow
 from infrahub.core import registry
 from infrahub.core.branch import Branch
+from infrahub.core.branch.enums import BranchStatus
 from infrahub.core.changelog.diff import DiffChangelogCollector, MigrationTracker
 from infrahub.core.constants import MutationAction
 from infrahub.core.diff.coordinator import DiffCoordinator
@@ -21,7 +22,10 @@
 from infrahub.core.diff.model.path import BranchTrackingId, EnrichedDiffRoot, EnrichedDiffRootMetadata
 from infrahub.core.diff.models import RequestDiffUpdate
 from infrahub.core.diff.repository.repository import DiffRepository
+from infrahub.core.graph import GRAPH_VERSION
 from infrahub.core.merge import BranchMerger
+from infrahub.core.migrations.exceptions import MigrationFailureError
+from infrahub.core.migrations.runner import MigrationRunner
 from infrahub.core.migrations.schema.models import SchemaApplyMigrationData
 from infrahub.core.migrations.schema.tasks import schema_apply_migrations
 from infrahub.core.timestamp import Timestamp
@@ -39,6 +43,7 @@
 from infrahub.workflows.catalogue import (
     BRANCH_CANCEL_PROPOSED_CHANGES,
     BRANCH_MERGE_POST_PROCESS,
+    BRANCH_MIGRATE,
     DIFF_REFRESH_ALL,
     DIFF_UPDATE,
     GIT_REPOSITORIES_CREATE_BRANCH,
@@ -51,6 +56,7 @@
 
 @flow(name="branch-rebase", flow_run_name="Rebase branch {branch}")
 async def rebase_branch(branch: str, context: InfrahubContext, send_events: bool = True) -> None:  # noqa: PLR0915
+    workflow = get_workflow()
     database = await get_database()
     async with database.start_session() as db:
         log = get_run_logger()
@@ -69,7 +75,7 @@ async def rebase_branch(branch: str, context: InfrahubContext, send_events: bool
             diff_repository=diff_repository,
             source_branch=obj,
             diff_locker=DiffLocker(),
-            workflow=get_workflow(),
+            workflow=workflow,
         )
 
         enriched_diff_metadata = await diff_coordinator.update_branch_diff(base_branch=base_branch, diff_branch=obj)
@@ -156,41 +162,75 @@ async def rebase_branch(branch: str, context: InfrahubContext, send_events: bool
             target_branch_name=registry.default_branch,
         )
         if ipam_node_details:
-            await get_workflow().submit_workflow(
+            await workflow.submit_workflow(
                 workflow=IPAM_RECONCILIATION,
                 context=context,
                 parameters={"branch": obj.name, "ipam_node_details": ipam_node_details},
             )
 
-    await get_workflow().submit_workflow(
-        workflow=DIFF_REFRESH_ALL, context=context, parameters={"branch_name": obj.name}
-    )
+    await workflow.submit_workflow(workflow=DIFF_REFRESH_ALL, context=context, parameters={"branch_name": obj.name})
+    await workflow.submit_workflow(workflow=BRANCH_MIGRATE, context=context, parameters={"branch_name": obj.name})
 
-    if send_events:
-        # -------------------------------------------------------------
-        # Generate an event to indicate that a branch has been rebased
-        # -------------------------------------------------------------
-        rebase_event = BranchRebasedEvent(
-            branch_name=obj.name, branch_id=str(obj.uuid), meta=EventMeta(branch=obj, context=context)
-        )
-        events: list[InfrahubEvent] = [rebase_event]
-        changelog_collector = DiffChangelogCollector(
-            diff=default_branch_diff, branch=obj, db=db, migration_tracker=MigrationTracker(migrations=migrations)
+    if not send_events:
+        return
+
+    # -------------------------------------------------------------
+    # Generate an event to indicate that a branch has been rebased
+    # -------------------------------------------------------------
+    rebase_event = BranchRebasedEvent(
+        branch_name=obj.name, branch_id=str(obj.uuid), meta=EventMeta(branch=obj, context=context)
+    )
+    events: list[InfrahubEvent] = [rebase_event]
+    changelog_collector = DiffChangelogCollector(
+        diff=default_branch_diff, branch=obj, db=db, migration_tracker=MigrationTracker(migrations=migrations)
+    )
+    for action, node_changelog in changelog_collector.collect_changelogs():
+        node_event_class = get_node_event(MutationAction.from_diff_action(diff_action=action))
+        mutate_event = node_event_class(
+            kind=node_changelog.node_kind,
+            node_id=node_changelog.node_id,
+            changelog=node_changelog,
+            fields=node_changelog.updated_fields,
+            meta=EventMeta.from_parent(parent=rebase_event, branch=obj),
         )
-        for action, node_changelog in changelog_collector.collect_changelogs():
-            node_event_class = get_node_event(MutationAction.from_diff_action(diff_action=action))
-            mutate_event = node_event_class(
-                kind=node_changelog.node_kind,
-                node_id=node_changelog.node_id,
-                changelog=node_changelog,
-                fields=node_changelog.updated_fields,
-                meta=EventMeta.from_parent(parent=rebase_event, branch=obj),
-            )
-            events.append(mutate_event)
+        events.append(mutate_event)
 
-        event_service = await get_event_service()
-        for event in events:
-            await event_service.send(event)
+    event_service = await get_event_service()
+    for event in events:
+        await event_service.send(event)
+
+
+@flow(name="migrate_branch", flow_run_name="Apply migrations to branch {branch}")
+async def migrate_branch(branch_name: str, context: InfrahubContext) -> None:  # noqa: ARG001
+    db = await get_database()
+    log = get_run_logger()
+
+    branch = await registry.get_branch(db=db, branch=branch_name)
+
+    if branch.graph_version == GRAPH_VERSION:
+        log.info(f"Branch '{branch.name}' is up-to-date")
+        return
+
+    migration_runner = MigrationRunner(branch=branch)
+    if not migration_runner.has_migrations():
+        log.info(f"No migrations detected for branch '{branch.name}'")
+        return
+
+    # Branch status will remain as so if the migration process fails
+    # This will help user to know that a branch is in an invalid state to be used properly and that actions need to be taken
+    branch.status = BranchStatus.NEED_UPGRADE_REBASE
+    await branch.save(db=db)
+
+    try:
+        await migration_runner.run(db=db)
+    except MigrationFailureError as exc:
+        log.error(f"Failed to migrate branch '{branch.name}': {exc.errors}")
+        return
+
+    if branch.status == BranchStatus.NEED_UPGRADE_REBASE:
+        branch.status = BranchStatus.OPEN
+    branch.graph_version = GRAPH_VERSION
+    await branch.save(db=db)
 
 
 @flow(name="branch-merge", flow_run_name="Merge branch {branch} into main")

backend/infrahub/core/migrations/exceptions.py

@@ -0,0 +1,4 @@
+class MigrationFailureError(Exception):
+    def __init__(self, errors: list[str]) -> None:
+        super().__init__()
+        self.errors = errors

backend/infrahub/core/migrations/graph/m043_backfill_hfid_display_label_in_db.py

@@ -447,9 +447,7 @@ async def _do_one_schema_batch(
 
         print("done")
 
-    async def process_branch(self, db: InfrahubDatabase, branch: Branch) -> MigrationResult:
-        root_node = await get_root_node(db=db, initialize=False)
-        default_branch = root_node.default_branch
+    async def execute_against_branch(self, db: InfrahubDatabase, branch: Branch) -> MigrationResult:
         schema_manager = SchemaManager()
         internal_schema_root = SchemaRoot(**internal_schema)
         schema_manager.register_schema(schema=internal_schema_root)

backend/infrahub/core/migrations/runner.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Sequence
+
+from infrahub.core.migrations.graph import MIGRATIONS
+
+from .exceptions import MigrationFailureError
+from .shared import MigrationWithRebase
+
+if TYPE_CHECKING:
+    from infrahub.core.branch import Branch
+    from infrahub.database import InfrahubDatabase
+
+
+class MigrationRunner:
+    def __init__(self, branch: Branch) -> None:
+        self.branch = branch
+        self.applicable_migrations = self._get_applicable_migrations()
+
+    def _get_applicable_migrations(self) -> Sequence[MigrationWithRebase]:
+        applicable_migrations = []
+        for migration_class in [m for m in MIGRATIONS if issubclass(m, MigrationWithRebase)]:
+            migration = migration_class.init()
+            if self.branch.graph_version and self.branch.graph_version > migration.minimum_version:
+                continue
+            applicable_migrations.append(migration)
+
+        return applicable_migrations
+
+    def has_migrations(self) -> bool:
+        return bool(self.applicable_migrations)
+
+    async def run(self, db: InfrahubDatabase) -> None:
+        if not self.has_migrations():
+            return
+
+        for migration in self.applicable_migrations:
+            execution_result = await migration.execute_against_branch(db=db, branch=self.branch)
+            validation_result = None
+
+            if execution_result.success:
+                validation_result = await migration.validate_migration(db=db)
+                # FIXME: graph_version should be changed elsewhere?
+                # if validation_result.success and self.branch.status != BranchStatus.NEED_UPGRADE_REBASE:
+                #     self.branch.graph_version = migration.minimum_version + 1
+                #     await self.branch.save(db=db)
+
+            if not execution_result.success or (validation_result and not validation_result.success):
+                if execution_result.errors:
+                    raise MigrationFailureError(errors=execution_result.errors)
+
+                if validation_result and not validation_result.success:
+                    raise MigrationFailureError(errors=validation_result.errors)