Add --rebase-branches flag to upgrade cmd

Author: gmazoyer

backend/infrahub/cli/db.py

@@ -371,7 +371,7 @@ async def rebase_and_migrate_branches(db: InfrahubDatabase, current_graph_versio
     for branch in branches:
         migrations = [
             m
-            for m in await detect_migration_to_run(current_graph_version=branch.graph_version or current_graph_version)
+            for m in await detect_migration_to_run(current_graph_version=branch.graph_version or 0)
             if isinstance(m, MigrationWithRebase)
         ]
         rprint(

backend/infrahub/cli/upgrade.py

@@ -51,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."""
 
@@ -93,12 +94,6 @@ async def upgrade_cmd(
     await initialize_internal_schema()
     await update_core_schema(db=dbdriver, initialize=False)
 
-    if 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
-
     # -------------------------------------------
     # Upgrade Internal Objects, generated and managed by Infrahub
     # -------------------------------------------
@@ -114,6 +109,17 @@ async def upgrade_cmd(
         await setup_deployments(client=client)
         await trigger_configure_all()
 
+    # -------------------------------------------
+    # 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
+
     await dbdriver.close()
 
 

backend/infrahub/core/migrations/shared.py

@@ -277,6 +277,7 @@ async def process_branch(self, db: InfrahubDatabase, branch: Branch) -> Migratio
     async def execute_against_branch(
         self, db: InfrahubDatabase, branch: Branch, skip_rebase: bool = False
     ) -> MigrationResult:
+        initialize_lock()
         await registry.schema.load_schema(db=db, branch=branch)
 
         if not skip_rebase:
@@ -290,7 +291,6 @@ async def execute_against_branch(
     async def execute(self, db: InfrahubDatabase) -> MigrationResult:
         from infrahub.core.initialization import initialization
 
-        initialize_lock()
         await initialization(db=db)
 
         return await self.execute_against_branch(db=db, branch=registry.get_branch_from_registry(), skip_rebase=True)