Allow rebase of branches during migration

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.")