Allow rebase of branches during migration (#7447)

These changes add a new flow to apply migrations (schema/data ones) inside branches. The upgrade process (via the provided command) is also updated with a new `--rebase-branches` flag to trigger branch rebases as part of the upgrade process.

A new `MigrationRequiringRebase` class is introduced with a couple of methods that need to be defined on classes inheriting it. One method defines what to do on the default branch while the other one defines what to do for all other branches.

The branch nodes can have a new `NEED_UPGRADE_REBASE` status and they also now have a graph_version property to know if a branch is up-to-date regarding the graph structure.

A new prefect event is also introduced and triggered when the `branch-migrate` flow completes for a branch.

Author: gmazoyer

backend/infrahub/cli/db.py

@@ -18,7 +18,12 @@
 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.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
 from infrahub.core.graph.index import node_indexes, rel_indexes
@@ -30,10 +35,8 @@
     GraphRelationshipIsPartOf,
     GraphRelationshipProperties,
 )
-from infrahub.core.initialization import (
-    get_root_node,
-    initialize_registry,
-)
+from infrahub.core.initialization import get_root_node, initialize_registry
+from infrahub.core.migrations.exceptions import MigrationFailureError
 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
@@ -45,6 +48,7 @@
 from infrahub.database import DatabaseType
 from infrahub.database.memgraph import IndexManagerMemgraph
 from infrahub.database.neo4j import IndexManagerNeo4j
+from infrahub.exceptions import ValidationError
 
 from .constants import ERROR_BADGE, FAILED_BADGE, SUCCESS_BADGE
 from .db_commands.check_inheritance import check_inheritance
@@ -59,7 +63,7 @@ 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 MigrationTypes
     from infrahub.database import InfrahubDatabase
     from infrahub.database.index import IndexManagerBase
 
@@ -105,7 +109,15 @@ async def migrate_cmd(
     context: CliContext = ctx.obj
     dbdriver = await context.init_db(retry=1)
 
-    await migrate_database(db=dbdriver, initialize=True, check=check, migration_number=migration_number)
+    root_node = await get_root_node(db=dbdriver)
+    migrations = await detect_migration_to_run(
+        current_graph_version=root_node.graph_version, migration_number=migration_number
+    )
+
+    if check or not migrations:
+        return
+
+    await migrate_database(db=dbdriver, migrations=migrations, initialize=True)
 
     await dbdriver.close()
 
@@ -268,49 +280,55 @@ async def index(
     await dbdriver.close()
 
 
+async def detect_migration_to_run(
+    current_graph_version: int, migration_number: int | str | None = None
+) -> Sequence[MigrationTypes]:
+    """Return a sequence of migrations to apply to upgrade the database."""
+    rprint("Checking current state of the database")
+    migrations: list[MigrationTypes] = []
+
+    if migration_number:
+        migration = get_migration_by_number(migration_number)
+        migrations.append(migration)
+        if current_graph_version > migration.minimum_version:
+            rprint(
+                f"Migration {migration_number} already applied. To apply again, run the command without the --check flag."
+            )
+            return []
+        rprint(
+            f"Migration {migration_number} needs to be applied. Run `infrahub db migrate` to apply all outstanding migrations."
+        )
+    else:
+        migrations.extend(await get_graph_migrations(current_graph_version=current_graph_version))
+        if not migrations:
+            rprint(f"Database up-to-date (v{current_graph_version}), no migration to execute.")
+            return []
+
+    rprint(
+        f"Database needs to be updated (v{current_graph_version} -> v{GRAPH_VERSION}), {len(migrations)} migrations pending"
+    )
+    return migrations
+
+
 async def migrate_database(
-    db: InfrahubDatabase, initialize: bool = False, check: bool = False, migration_number: int | str | None = None
+    db: InfrahubDatabase, migrations: Sequence[MigrationTypes], initialize: bool = False
 ) -> 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.
 
     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.
+        migrations: Sequence of migrations to apply.
+        initialize: Whether to initialize the registry before running migrations.
     """
-    rprint("Checking current state of the Database")
+    if not migrations:
+        return True
 
     if initialize:
         await initialize_registry(db=db)
 
     root_node = await get_root_node(db=db)
-    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"
-        )
-
-    if check:
-        return True
 
     for migration in migrations:
         execution_result = await migration.execute(db=db)
@@ -335,6 +353,36 @@ async def migrate_database(
     return True
 
 
+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
+
+    rprint(f"Planning rebase and migrations for {len(branches)} branches: {', '.join([b.name for b in branches])}")
+
+    for branch in branches:
+        if branch.graph_version == GRAPH_VERSION:
+            rprint(
+                f"Ignoring branch rebase and migrations for '{branch.name}' (ID: {branch.uuid}), it is already up-to-date"
+            )
+            continue
+
+        rprint(f"Rebasing branch '{branch.name}' (ID: {branch.uuid})...", end="")
+        try:
+            await registry.schema.load_schema(db=db, branch=branch)
+            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(SUCCESS_BADGE)
+        except (ValidationError, MigrationFailureError):
+            rprint(FAILED_BADGE)
+
+
 async def initialize_internal_schema() -> None:
     registry.schema = SchemaManager()
     schema = SchemaRoot(**internal_schema)

backend/infrahub/cli/upgrade.py

@@ -11,7 +11,12 @@
 from rich import print as rprint
 
 from infrahub import config
-from infrahub.core.initialization import create_anonymous_role, create_default_account_groups, initialize_registry
+from infrahub.core.initialization import (
+    create_anonymous_role,
+    create_default_account_groups,
+    get_root_node,
+    initialize_registry,
+)
 from infrahub.core.manager import NodeManager
 from infrahub.core.protocols import CoreAccount, CoreObjectPermission
 from infrahub.dependencies.registry import build_component_registry
@@ -26,7 +31,13 @@
     setup_worker_pools,
 )
 
-from .db import initialize_internal_schema, migrate_database, update_core_schema
+from .db import (
+    detect_migration_to_run,
+    initialize_internal_schema,
+    migrate_database,
+    trigger_rebase_branches,
+    update_core_schema,
+)
 
 if TYPE_CHECKING:
     from infrahub.cli.context import CliContext
@@ -40,6 +51,7 @@ async def upgrade_cmd(
     ctx: typer.Context,
     config_file: str = typer.Argument("infrahub.toml", envvar="INFRAHUB_CONFIG"),
     check: bool = typer.Option(False, help="Check the state of the system without upgrading."),
+    rebase_branches: bool = typer.Option(False, help="Rebase and apply migrations to branches if required."),
 ) -> None:
     """Upgrade Infrahub to the latest version."""
 
@@ -57,6 +69,8 @@ async def upgrade_cmd(
 
     build_component_registry()
 
+    root_node = await get_root_node(db=dbdriver)
+
     # NOTE add step to validate if the database and the task manager are reachable
 
     # -------------------------------------------
@@ -67,7 +81,12 @@ async def upgrade_cmd(
     # Upgrade Infrahub Database and Schema
     # -------------------------------------------
 
-    if not await migrate_database(db=dbdriver, initialize=False, check=check):
+    migrations = await detect_migration_to_run(current_graph_version=root_node.graph_version)
+    if check:
+        await dbdriver.close()
+        return
+
+    if not await migrate_database(db=dbdriver, initialize=False, migrations=migrations):
         # A migration failed, stop the upgrade process
         rprint("Upgrade cancelled due to migration failure.")
         await dbdriver.close()
@@ -91,6 +110,12 @@ async def upgrade_cmd(
         await setup_deployments(client=client)
         await trigger_configure_all()
 
+    # -------------------------------------------
+    # Perform branch rebase and apply migrations to them
+    # -------------------------------------------
+    if rebase_branches:
+        await trigger_rebase_branches(db=dbdriver)
+
     await dbdriver.close()
 
 

backend/infrahub/core/branch/enums.py

@@ -4,4 +4,5 @@
 class BranchStatus(InfrahubStringEnum):
     OPEN = "OPEN"
     NEED_REBASE = "NEED_REBASE"
+    NEED_UPGRADE_REBASE = "NEED_UPGRADE_REBASE"
     DELETING = "DELETING"

backend/infrahub/core/branch/models.py

@@ -6,9 +6,8 @@
 from pydantic import Field, field_validator
 
 from infrahub.core.branch.enums import BranchStatus
-from infrahub.core.constants import (
-    GLOBAL_BRANCH_NAME,
-)
+from infrahub.core.constants import GLOBAL_BRANCH_NAME
+from infrahub.core.graph import GRAPH_VERSION
 from infrahub.core.models import SchemaBranchHash  # noqa: TC001
 from infrahub.core.node.standard import StandardNode
 from infrahub.core.query import QueryType
@@ -46,6 +45,7 @@ class Branch(StandardNode):
     is_isolated: bool = True
     schema_changed_at: Optional[str] = None
     schema_hash: Optional[SchemaBranchHash] = None
+    graph_version: int | None = None
 
     _exclude_attrs: list[str] = ["id", "uuid", "owner"]
 
@@ -261,6 +261,10 @@ def get_branches_and_times_for_range(
 
         return start, end
 
+    async def create(self, db: InfrahubDatabase) -> bool:
+        self.graph_version = GRAPH_VERSION
+        return await super().create(db=db)
+
     async def delete(self, db: InfrahubDatabase) -> None:
         if self.is_default:
             raise ValidationError(f"Unable to delete {self.name} it is the default branch.")