Use prefect flow for rebasing branches

Author: gmazoyer

backend/infrahub/core/branch/tasks.py

@@ -50,7 +50,7 @@
 
 
 @flow(name="branch-rebase", flow_run_name="Rebase branch {branch}")
-async def rebase_branch(branch: str, context: InfrahubContext) -> None:  # noqa: PLR0915
+async def rebase_branch(branch: str, context: InfrahubContext, send_events: bool = True) -> None:  # noqa: PLR0915
     database = await get_database()
     async with database.start_session() as db:
         log = get_run_logger()
@@ -166,30 +166,31 @@ async def rebase_branch(branch: str, context: InfrahubContext) -> None:  # noqa:
         workflow=DIFF_REFRESH_ALL, context=context, parameters={"branch_name": obj.name}
     )
 
-    # -------------------------------------------------------------
-    # 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),
+    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.append(mutate_event)
+        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),
+            )
+            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="branch-merge", flow_run_name="Merge branch {branch} into main")

backend/infrahub/core/migrations/shared.py

@@ -6,6 +6,8 @@
 from rich.console import Console
 from typing_extensions import Self
 
+from infrahub.auth import AccountSession, AuthType
+from infrahub.context import InfrahubContext
 from infrahub.core import registry
 from infrahub.core.branch.enums import BranchStatus
 from infrahub.core.path import SchemaPath  # noqa: TC001
@@ -18,6 +20,7 @@
     internal_schema,
 )
 from infrahub.core.timestamp import Timestamp
+from infrahub.exceptions import ValidationError
 
 from .query import MigrationBaseQuery  # noqa: TC001
 
@@ -242,24 +245,26 @@ def init(cls, **kwargs: dict[str, Any]) -> Self:
         return cls(**kwargs)  # type: ignore[arg-type]
 
     async def rebase_branch(self, db: InfrahubDatabase, branch: Branch) -> bool:
-        # Circular deps if import inside import block
-        # from infrahub.core.branch.tasks import rebase_branch
-        # await rebase_branch(
-        #     branch=branch.name,
-        #     context=InfrahubContext.init(
-        #         branch=branch,
-        #         # FIXME: or superuser account?
-        #         account=AccountSession(auth_type=AuthType.NONE, authenticated=False, account_id=""),
-        #     ),
-        # )
+        # Circular deps if imported inside import block
+        from infrahub.core.branch.tasks import rebase_branch
+
         # NOTE: need service/bus component up and running to use prefect flow
         console = Console()
         console.print(f"Rebasing branch '{branch.name}' (ID: {branch.uuid})...", end="")
         try:
-            await branch.rebase(db=db)
+            # await branch.rebase(db=db)
+            await rebase_branch(
+                branch=branch.name,
+                context=InfrahubContext.init(
+                    branch=branch,
+                    # FIXME: or superuser account?
+                    account=AccountSession(auth_type=AuthType.NONE, authenticated=False, account_id=""),
+                ),
+                send_events=False,
+            )
             console.print("done")
             branch.graph_version = self.minimum_version + 1
-            # NOTE: Narrow to more accurate exception
+        except ValidationError:
             console.print("failed")
             branch.status = BranchStatus.NEED_UPGRADE_REBASE
             return False
@@ -280,7 +285,7 @@ async def execute_against_branches(self, db: InfrahubDatabase, branches: Sequenc
         for branch in branches:
             if not await self.rebase_branch(db=db, branch=branch):
                 result.errors.append(f"Failed to rebase branch '{branch.name}' ({branch.uuid})")
-                return result
+                continue
 
             r = await self.execute_against_branch(db=db, branch=branch)
             result.nbr_migrations_executed += 1