infrahub-server 1.5.0b2__py3-none-any.whl → 1.5.1__py3-none-any.whl

infrahub/api/dependencies.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, AsyncIterator
+from typing import TYPE_CHECKING, AsyncGenerator
 
 from fastapi import Depends, Query, Request
 from fastapi.security import APIKeyHeader, HTTPAuthorizationCredentials, HTTPBearer
@@ -17,8 +17,6 @@ from infrahub.exceptions import AuthorizationError
 from infrahub.permissions import PermissionManager
 
 if TYPE_CHECKING:
-    from neo4j import AsyncSession
-
     from infrahub.models import RefreshTokenData
 
 jwt_scheme = HTTPBearer(auto_error=False)
@@ -36,16 +34,9 @@ class BranchParams(BaseModel):
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
 
-async def get_session(request: Request) -> AsyncIterator[AsyncSession]:
-    session = request.app.state.db.session(database=config.SETTINGS.database.database_name)
-    try:
-        yield session
-    finally:
-        await session.close()
-
-
-async def get_db(request: Request) -> InfrahubDatabase:
-    return request.app.state.db.start_session()
+async def get_db(request: Request) -> AsyncGenerator[InfrahubDatabase, None]:
+    async with request.app.state.db.start_session() as db:
+        yield db
 
 
 async def get_access_token(

infrahub/api/transformation.py

@@ -64,17 +64,18 @@ async def transform_python(
             message="Repository doesn't have a commit",
         )
 
-    gql_params = await prepare_graphql_params(
-        db=request.app.state.db, branch=branch_params.branch, at=branch_params.at, service=request.app.state.service
-    )
+    async with db.start_session(read_only=True) as dbs:
+        gql_params = await prepare_graphql_params(
+            db=dbs, branch=branch_params.branch, at=branch_params.at, service=request.app.state.service
+        )
 
-    result = await graphql(
-        schema=gql_params.schema,
-        source=query.query.value,
-        context_value=gql_params.context,
-        root_value=None,
-        variable_values=params,
-    )
+        result = await graphql(
+            schema=gql_params.schema,
+            source=query.query.value,
+            context_value=gql_params.context,
+            root_value=None,
+            variable_values=params,
+        )
 
     data = extract_data(query_name=query.name.value, result=result)
 
@@ -128,17 +129,18 @@ async def transform_jinja2(
             message="Repository doesn't have a commit",
         )
 
-    gql_params = await prepare_graphql_params(
-        db=request.app.state.db, branch=branch_params.branch, at=branch_params.at, service=request.app.state.service
-    )
+    async with db.start_session(read_only=True) as dbs:
+        gql_params = await prepare_graphql_params(
+            db=dbs, branch=branch_params.branch, at=branch_params.at, service=request.app.state.service
+        )
 
-    result = await graphql(
-        schema=gql_params.schema,
-        source=query.query.value,
-        context_value=gql_params.context,
-        root_value=None,
-        variable_values=params,
-    )
+        result = await graphql(
+            schema=gql_params.schema,
+            source=query.query.value,
+            context_value=gql_params.context,
+            root_value=None,
+            variable_values=params,
+        )
 
     data = extract_data(query_name=query.name.value, result=result)
 

infrahub/cli/db.py

@@ -13,7 +13,6 @@ import typer
 import ujson
 from infrahub_sdk.async_typer import AsyncTyper
 from prefect.testing.utilities import prefect_test_harness
-from rich import print as rprint
 from rich.console import Console
 from rich.table import Table
 
@@ -22,6 +21,7 @@ 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
@@ -40,6 +40,7 @@ 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
+from infrahub.core.migrations.shared import get_migration_console
 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
@@ -284,27 +285,27 @@ 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")
+    get_migration_console().log("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(
+            get_migration_console().log(
                 f"Migration {migration_number} already applied. To apply again, run the command without the --check flag."
             )
             return []
-        rprint(
+        get_migration_console().log(
             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.")
+            get_migration_console().log(f"Database up-to-date (v{current_graph_version}), no migration to execute.")
             return []
 
-    rprint(
+    get_migration_console().log(
         f"Database needs to be updated (v{current_graph_version} -> v{GRAPH_VERSION}), {len(migrations)} migrations pending"
     )
     return migrations
@@ -337,38 +338,57 @@ async def migrate_database(
         if execution_result.success:
             validation_result = await migration.validate_migration(db=db)
             if validation_result.success:
-                rprint(f"Migration: {migration.name} {SUCCESS_BADGE}")
+                get_migration_console().log(f"Migration: {migration.name} {SUCCESS_BADGE}")
                 root_node.graph_version = migration.minimum_version + 1
                 await root_node.save(db=db)
 
         if not execution_result.success or (validation_result and not validation_result.success):
-            rprint(f"Migration: {migration.name} {FAILED_BADGE}")
+            get_migration_console().log(f"Migration: {migration.name} {FAILED_BADGE}")
             for error in execution_result.errors:
-                rprint(f"  {error}")
+                get_migration_console().log(f"  {error}")
             if validation_result and not validation_result.success:
                 for error in validation_result.errors:
-                    rprint(f"  {error}")
+                    get_migration_console().log(f"  {error}")
             return False
 
     return True
 
 
-async def trigger_rebase_branches(db: InfrahubDatabase) -> None:
-    """Trigger rebase of non-default branches, also triggering migrations in the process."""
+async def mark_branches_needing_rebase(db: InfrahubDatabase) -> list[Branch]:
     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 []
+
+    branches_needing_rebase: list[Branch] = []
+    for branch in branches:
+        if branch.graph_version == GRAPH_VERSION:
+            continue
+
+        branch.status = BranchStatus.NEED_UPGRADE_REBASE
+        await branch.save(db=db)
+        branches_needing_rebase.append(branch)
+
+    return branches_needing_rebase
+
+
+async def trigger_rebase_branches(db: InfrahubDatabase, branches: Sequence[Branch]) -> None:
+    """Trigger rebase of non-default branches, also triggering migrations in the process."""
     if not branches:
         return
 
-    rprint(f"Planning rebase and migrations for {len(branches)} branches: {', '.join([b.name for b in branches])}")
+    get_migration_console().log(
+        f"Planning rebase and migrations for {len(branches)} {'branches' if len(branches) != 1 else 'branch'}: "
+        f"{', '.join([b.name for b in branches])}"
+    )
 
     for branch in branches:
         if branch.graph_version == GRAPH_VERSION:
-            rprint(
+            get_migration_console().log(
                 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="")
+        get_migration_console().print(f"Rebasing branch '{branch.name}' (ID: {branch.uuid})...", end="")
         try:
             await registry.schema.load_schema(db=db, branch=branch)
             await rebase_branch(
@@ -378,9 +398,9 @@ async def trigger_rebase_branches(db: InfrahubDatabase) -> None:
                 ),
                 send_events=False,
             )
-            rprint(SUCCESS_BADGE)
+            get_migration_console().log(SUCCESS_BADGE)
         except (ValidationError, MigrationFailureError):
-            rprint(FAILED_BADGE)
+            get_migration_console().log(FAILED_BADGE)
 
 
 async def initialize_internal_schema() -> None:
@@ -417,16 +437,16 @@ async def update_core_schema(db: InfrahubDatabase, initialize: bool = True, debu
     branch_schema.validate_node_deletions(diff=schema_diff)
     result = branch_schema.validate_update(other=candidate_schema, diff=schema_diff, enforce_update_support=False)
     if result.errors:
-        rprint(f"{ERROR_BADGE} | Unable to update the schema, due to failed validations")
+        get_migration_console().log(f"{ERROR_BADGE} | Unable to update the schema, due to failed validations")
         for error in result.errors:
-            rprint(error.to_string())
+            get_migration_console().log(error.to_string())
         raise typer.Exit(1)
 
     if not result.diff.all:
-        rprint("Core Schema Up to date, nothing to update")
+        get_migration_console().log("Core Schema Up to date, nothing to update")
         return
 
-    rprint("Core Schema has diff, will need to be updated")
+    get_migration_console().log("Core Schema has diff, will need to be updated")
     if debug:
         result.diff.print()
 
@@ -441,9 +461,9 @@ async def update_core_schema(db: InfrahubDatabase, initialize: bool = True, debu
     responses = await schema_validate_migrations(message=validate_migration_data)
     error_messages = [violation.message for response in responses for violation in response.violations]
     if error_messages:
-        rprint(f"{ERROR_BADGE} | Unable to update the schema, due to failed validations")
+        get_migration_console().log(f"{ERROR_BADGE} | Unable to update the schema, due to failed validations")
         for message in error_messages:
-            rprint(message)
+            get_migration_console().log(message)
         raise typer.Exit(1)
 
     # ----------------------------------------------------------
@@ -466,9 +486,11 @@ async def update_core_schema(db: InfrahubDatabase, initialize: bool = True, debu
             update_db=True,
         )
         default_branch.update_schema_hash()
-        rprint("The Core Schema has been updated, make sure to rebase any open branches after the upgrade")
+        get_migration_console().log(
+            "The Core Schema has been updated, make sure to rebase any open branches after the upgrade"
+        )
         if debug:
-            rprint(f"New schema hash: {default_branch.active_schema_hash.main}")
+            get_migration_console().log(f"New schema hash: {default_branch.active_schema_hash.main}")
         await default_branch.save(db=dbt)
 
     # ----------------------------------------------------------
@@ -483,9 +505,9 @@ async def update_core_schema(db: InfrahubDatabase, initialize: bool = True, debu
     migration_error_msgs = await schema_apply_migrations(message=apply_migration_data)
 
     if migration_error_msgs:
-        rprint(f"{ERROR_BADGE} | Some error(s) happened while running the schema migrations")
+        get_migration_console().log(f"{ERROR_BADGE} | Some error(s) happened while running the schema migrations")
         for message in migration_error_msgs:
-            rprint(message)
+            get_migration_console().log(message)
         raise typer.Exit(1)
 
 
@@ -609,16 +631,16 @@ RETURN vertices, edges
         edge_csv_writer.writeheader()
 
         while has_more_data:
-            rprint("Retrieving batch of vertices and edges...", end="")
+            get_migration_console().print("Retrieving batch of vertices and edges...", end="")
             results = await db.execute_query(
                 query=query,
                 params={"kinds": kinds, "uuids": uuids, "limit": limit, "offset": offset},
             )
-            rprint("done. ", end="")
+            get_migration_console().print("done. ", end="")
             has_more_data = len(results) >= limit
             offset += limit
 
-            rprint("Writing batch to export files...", end="")
+            get_migration_console().print("Writing batch to export files...", end="")
             for result in results:
                 vertices = result.get("vertices")
                 for vertex in vertices:
@@ -649,10 +671,10 @@ RETURN vertices, edges
                             serial_edge[property_name] = value
                     edge_csv_writer.writerow(serial_edge)
                     all_db_ids.add(edge.element_id)
-            rprint("done.")
+            get_migration_console().log("done.")
 
-    rprint(f"{SUCCESS_BADGE} Export complete")
-    rprint(f"Export directory is here: {export_dir.absolute()}")
+    get_migration_console().log(f"{SUCCESS_BADGE} Export complete")
+    get_migration_console().log(f"Export directory is here: {export_dir.absolute()}")
     return export_dir
 
 
@@ -690,9 +712,9 @@ UNWIND $vertices AS vertex
 CREATE (v:ImportNode:%(node_labels)s {db_id: vertex.db_id})
 SET v = vertex
     """ % {"node_labels": ":".join(vertex_labels)}
-    rprint(f"Loading {len(vertex_dicts)} {vertex_labels} nodes...", end="")
+    get_migration_console().print(f"Loading {len(vertex_dicts)} {vertex_labels} nodes...", end="")
     await db.execute_query(query=vertex_import_query, params={"vertices": vertex_dicts})
-    rprint("done")
+    get_migration_console().log("done")
 
 
 async def load_edges(
@@ -705,17 +727,17 @@ MATCH (b:ImportNode) WHERE b.db_id = toString(edge.end_node_id)
 CREATE (a)-[e:%(edge_type)s]->(b)
 SET e = edge.properties
     """ % {"edge_type": edge_type}
-    rprint(f"Loading {len(edge_dicts)} {edge_type} edges...", end="")
+    get_migration_console().print(f"Loading {len(edge_dicts)} {edge_type} edges...", end="")
     await db.execute_query(query=edges_import_query, params={"edges": edge_dicts})
-    rprint("done")
+    get_migration_console().log("done")
 
 
 async def load_export(db: InfrahubDatabase, export_dir: Path, query_limit: int = 1000) -> None:
     if not export_dir.exists():
-        rprint(f"{ERROR_BADGE} {export_dir} does not exist")
+        get_migration_console().log(f"{ERROR_BADGE} {export_dir} does not exist")
         raise typer.Exit(1)
     if not export_dir.is_dir():
-        rprint(f"{ERROR_BADGE} {export_dir} is not a directory")
+        get_migration_console().log(f"{ERROR_BADGE} {export_dir} is not a directory")
         raise typer.Exit(1)
     vertex_file: Path | None = None
     edge_file: Path | None = None
@@ -726,17 +748,17 @@ async def load_export(db: InfrahubDatabase, export_dir: Path, query_limit: int =
         elif export_file.name == "edges.csv":
             edge_file = export_file
     if not vertex_file or not vertex_file.exists() or not vertex_file.is_file():
-        rprint(f"{ERROR_BADGE} File 'vertices.csv' does not exist in the export directory")
+        get_migration_console().log(f"{ERROR_BADGE} File 'vertices.csv' does not exist in the export directory")
         raise typer.Exit(1)
     if not edge_file or not edge_file.exists() or not edge_file.is_file():
-        rprint(f"{ERROR_BADGE} File 'edges.csv' does not exist in the export directory")
+        get_migration_console().log(f"{ERROR_BADGE} File 'edges.csv' does not exist in the export directory")
         raise typer.Exit(1)
 
     # index massively improves time required to load a large export
     create_index_query = "CREATE RANGE INDEX import_node_db_id IF NOT EXISTS FOR (v:ImportNode) ON (v.db_id)"
     await db.execute_query(query=create_index_query)
 
-    rprint("Loading vertices...")
+    get_migration_console().log("Loading vertices...")
     vertices_by_labels_map: dict[frozenset[str], list[dict[str, Any]]] = defaultdict(list)
     with vertex_file.open() as file:
         csv_reader = DictReader(file)
@@ -751,9 +773,9 @@ async def load_export(db: InfrahubDatabase, export_dir: Path, query_limit: int =
 
         for labels, vertex_rows in vertices_by_labels_map.items():
             await load_vertices(db=db, vertex_labels=list(labels), vertex_dicts=vertex_rows)
-    rprint("Vertices loaded")
+    get_migration_console().log("Vertices loaded")
 
-    rprint("Loading edges...")
+    get_migration_console().log("Loading edges...")
     edges_by_type_map: dict[str, list[dict[str, Any]]] = defaultdict(list)
     with edge_file.open() as file:
         csv_reader = DictReader(file)
@@ -778,8 +800,8 @@ async def load_export(db: InfrahubDatabase, export_dir: Path, query_limit: int =
 
         for edge_type, edge_dicts in edges_by_type_map.items():
             await load_edges(db=db, edge_type=edge_type, edge_dicts=edge_dicts)
-    rprint("Edges loaded")
-    rprint(f"{SUCCESS_BADGE} Export loaded")
+    get_migration_console().log("Edges loaded")
+    get_migration_console().log(f"{SUCCESS_BADGE} Export loaded")
 
 
 @app.command(name="check")
@@ -822,10 +844,10 @@ async def run_database_checks(db: InfrahubDatabase, output_dir: Path) -> None:
         db: The database object.
         output_dir: Directory to save detailed check results.
     """
-    rprint("Running database health checks...")
+    get_migration_console().log("Running database health checks...")
 
     # Check 1: Duplicate active relationships
-    rprint("\n[bold cyan]Check 1: Duplicate Active Relationships[/bold cyan]")
+    get_migration_console().log("\n[bold cyan]Check 1: Duplicate Active Relationships[/bold cyan]")
     duplicate_active_rels_query = """
     MATCH (a:Node)-[e1:IS_RELATED {status: "active"}]-(r:Relationship)-[e2:IS_RELATED {branch: e1.branch, status: "active"}]-(b:Node)
     WHERE a.uuid < b.uuid
@@ -849,7 +871,7 @@ async def run_database_checks(db: InfrahubDatabase, output_dir: Path) -> None:
 
     results = await db.execute_query(query=duplicate_active_rels_query)
     if results:
-        rprint(f"[red]Found {len(results)} duplicate active relationships[/red]")
+        get_migration_console().log(f"[red]Found {len(results)} duplicate active relationships[/red]")
         # Write detailed results to file
         output_file = output_dir / "duplicate_active_relationships.csv"
         with output_file.open(mode="w", newline="") as f:
@@ -859,12 +881,12 @@ async def run_database_checks(db: InfrahubDatabase, output_dir: Path) -> None:
             writer.writeheader()
             for result in results:
                 writer.writerow(dict(result))
-        rprint(f"  Detailed results written to: {output_file}")
+        get_migration_console().log(f"  Detailed results written to: {output_file}")
     else:
-        rprint(f"{SUCCESS_BADGE} No duplicate active relationships found")
+        get_migration_console().log(f"{SUCCESS_BADGE} No duplicate active relationships found")
 
     # Check 2: Duplicated relationship nodes
-    rprint("\n[bold cyan]Check 2: Duplicated Relationship Nodes[/bold cyan]")
+    get_migration_console().log("\n[bold cyan]Check 2: Duplicated Relationship Nodes[/bold cyan]")
     duplicate_rel_nodes_query = """
     MATCH (r:Relationship)
     WITH r.uuid AS r_uuid, COUNT(*) AS num_rels
@@ -877,7 +899,7 @@ async def run_database_checks(db: InfrahubDatabase, output_dir: Path) -> None:
 
     results = await db.execute_query(query=duplicate_rel_nodes_query)
     if results:
-        rprint(f"[red]Found {len(results)} duplicated relationship nodes[/red]")
+        get_migration_console().log(f"[red]Found {len(results)} duplicated relationship nodes[/red]")
         # Write detailed results to file
         output_file = output_dir / "duplicated_relationship_nodes.csv"
         with output_file.open(mode="w", newline="") as f:
@@ -885,12 +907,12 @@ async def run_database_checks(db: InfrahubDatabase, output_dir: Path) -> None:
             writer.writeheader()
             for result in results:
                 writer.writerow(dict(result))
-        rprint(f"  Detailed results written to: {output_file}")
+        get_migration_console().log(f"  Detailed results written to: {output_file}")
     else:
-        rprint(f"{SUCCESS_BADGE} No duplicated relationship nodes found")
+        get_migration_console().log(f"{SUCCESS_BADGE} No duplicated relationship nodes found")
 
     # Check 3: Duplicated edges
-    rprint("\n[bold cyan]Check 3: Duplicated Edges[/bold cyan]")
+    get_migration_console().log("\n[bold cyan]Check 3: Duplicated Edges[/bold cyan]")
     duplicate_edges_query = """
     MATCH (a)
     CALL (a) {
@@ -920,7 +942,7 @@ async def run_database_checks(db: InfrahubDatabase, output_dir: Path) -> None:
 
     results = await db.execute_query(query=duplicate_edges_query)
     if results:
-        rprint(f"[red]Found {len(results)} sets of duplicated edges[/red]")
+        get_migration_console().log(f"[red]Found {len(results)} sets of duplicated edges[/red]")
         # Write detailed results to file
         output_file = output_dir / "duplicated_edges.csv"
         with output_file.open(mode="w", newline="") as f:
@@ -931,12 +953,12 @@ async def run_database_checks(db: InfrahubDatabase, output_dir: Path) -> None:
             writer.writeheader()
             for result in results:
                 writer.writerow(dict(result))
-        rprint(f"  Detailed results written to: {output_file}")
+        get_migration_console().log(f"  Detailed results written to: {output_file}")
     else:
-        rprint(f"{SUCCESS_BADGE} No duplicated edges found")
+        get_migration_console().log(f"{SUCCESS_BADGE} No duplicated edges found")
 
     # Check 4: Orphaned Relationships
-    rprint("\n[bold cyan]Check 4: Orphaned Relationships[/bold cyan]")
+    get_migration_console().log("\n[bold cyan]Check 4: Orphaned Relationships[/bold cyan]")
     orphaned_rels_query = """
     MATCH (r:Relationship)-[:IS_RELATED]-(peer:Node)
     WITH DISTINCT r, peer
@@ -954,7 +976,7 @@ async def run_database_checks(db: InfrahubDatabase, output_dir: Path) -> None:
     """
     results = await db.execute_query(query=orphaned_rels_query)
     if results:
-        rprint(f"[red]Found {len(results)} orphaned Relationships[/red]")
+        get_migration_console().log(f"[red]Found {len(results)} orphaned Relationships[/red]")
         # Write detailed results to file
         output_file = output_dir / "orphaned_relationships.csv"
         with output_file.open(mode="w", newline="") as f:
@@ -965,9 +987,9 @@ async def run_database_checks(db: InfrahubDatabase, output_dir: Path) -> None:
             writer.writeheader()
             for result in results:
                 writer.writerow(dict(result))
-        rprint(f"  Detailed results written to: {output_file}")
+        get_migration_console().log(f"  Detailed results written to: {output_file}")
     else:
-        rprint(f"{SUCCESS_BADGE} No orphaned relationships found")
+        get_migration_console().log(f"{SUCCESS_BADGE} No orphaned relationships found")
 
-    rprint(f"\n{SUCCESS_BADGE} Database health checks completed")
-    rprint(f"Detailed results saved to: {output_dir.absolute()}")
+    get_migration_console().log(f"\n{SUCCESS_BADGE} Database health checks completed")
+    get_migration_console().log(f"Detailed results saved to: {output_dir.absolute()}")

infrahub/cli/upgrade.py

@@ -8,7 +8,6 @@ import typer
 from deepdiff import DeepDiff
 from infrahub_sdk.async_typer import AsyncTyper
 from prefect.client.orchestration import get_client
-from rich import print as rprint
 
 from infrahub import config
 from infrahub.core.initialization import (
@@ -18,6 +17,7 @@ from infrahub.core.initialization import (
     initialize_registry,
 )
 from infrahub.core.manager import NodeManager
+from infrahub.core.migrations.shared import get_migration_console
 from infrahub.core.protocols import CoreAccount, CoreObjectPermission
 from infrahub.dependencies.registry import build_component_registry
 from infrahub.lock import initialize_lock
@@ -35,6 +35,7 @@ from infrahub.workflows.initialization import (
 from .db import (
     detect_migration_to_run,
     initialize_internal_schema,
+    mark_branches_needing_rebase,
     migrate_database,
     trigger_rebase_branches,
     update_core_schema,
@@ -42,9 +43,11 @@ from .db import (
 
 if TYPE_CHECKING:
     from infrahub.cli.context import CliContext
+    from infrahub.core.branch.models import Branch
     from infrahub.database import InfrahubDatabase
 
 app = AsyncTyper()
+console = get_migration_console()
 
 
 @app.command(name="upgrade")
@@ -53,6 +56,9 @@ async def upgrade_cmd(
     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."),
+    interactive: bool = typer.Option(
+        False, help="Use interactive prompt to accept or deny rebase of individual branches."
+    ),
 ) -> None:
     """Upgrade Infrahub to the latest version."""
 
@@ -90,7 +96,7 @@ async def upgrade_cmd(
 
     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.")
+        console.log("Upgrade cancelled due to migration failure.")
         await dbdriver.close()
         return
 
@@ -115,8 +121,22 @@ async def upgrade_cmd(
     # -------------------------------------------
     # Perform branch rebase and apply migrations to them
     # -------------------------------------------
+    branches = await mark_branches_needing_rebase(db=dbdriver)
+    plural = len(branches) != 1
+    get_migration_console().log(
+        f"Found {len(branches)} {'branches' if plural else 'branch'} that {'need' if plural else 'needs'} to be rebased"
+    )
+
     if rebase_branches:
-        await trigger_rebase_branches(db=dbdriver)
+        branches_to_rebase: list[Branch] = []
+        if not interactive:
+            branches_to_rebase = branches
+        else:
+            for branch in branches:
+                if typer.confirm(f"Rebase branch {branch.name}?"):
+                    branches_to_rebase.append(branch)
+
+        await trigger_rebase_branches(db=dbdriver, branches=branches_to_rebase)
 
     await dbdriver.close()
 
@@ -134,21 +154,21 @@ async def upgrade_menu(db: InfrahubDatabase) -> None:
     diff_menu = DeepDiff(menu_items.to_rest(), default_menu_dict.to_rest(), ignore_order=True)
 
     if not diff_menu:
-        rprint("Menu Up to date, nothing to update")
+        console.log("Menu Up to date, nothing to update")
         return
 
     await menu_repository.update_menu(existing_menu=menu_items, new_menu=default_menu_dict, menu_nodes=menu_nodes)
-    rprint("Menu has been updated")
+    console.log("Menu has been updated")
 
 
 async def upgrade_permissions(db: InfrahubDatabase) -> None:
     existing_permissions = await NodeManager.query(schema=CoreObjectPermission, db=db, limit=1)
     if existing_permissions:
-        rprint("Permissions Up to date, nothing to update")
+        console.log("Permissions Up to date, nothing to update")
         return
 
     await setup_permissions(db=db)
-    rprint("Permissions have been updated")
+    console.log("Permissions have been updated")
 
 
 async def setup_permissions(db: InfrahubDatabase) -> None:

infrahub/core/diff/calculator.py

@@ -141,8 +141,8 @@ class DiffCalculator:
             to_time=to_time,
             previous_node_field_specifiers=previous_node_specifiers,
         )
-        node_limit = int(config.SETTINGS.database.query_size_limit / 10)
-        fields_limit = int(config.SETTINGS.database.query_size_limit / 3)
+        node_limit = max(int(config.SETTINGS.database.query_size_limit / 10), 1)
+        fields_limit = max(int(config.SETTINGS.database.query_size_limit / 3), 1)
         properties_limit = config.SETTINGS.database.query_size_limit
 
         calculation_request = DiffCalculationRequest(

infrahub/core/diff/query/delete_query.py

@@ -20,10 +20,14 @@ class EnrichedDiffDeleteQuery(Query):
             diff_filter = "WHERE d_root.uuid IN $diff_root_uuids"
 
         query = """
-        MATCH (d_root:DiffRoot)
-        %(diff_filter)s
-        OPTIONAL MATCH (d_root)-[*]->(diff_thing)
-        DETACH DELETE diff_thing
-        DETACH DELETE d_root
+MATCH (d_root:DiffRoot)
+%(diff_filter)s
+OPTIONAL MATCH (d_root)-[*]->(diff_thing)
+WITH DISTINCT d_root, diff_thing
+ORDER BY elementId(diff_thing)
+CALL (diff_thing) {
+    DETACH DELETE diff_thing
+} IN TRANSACTIONS
+DETACH DELETE d_root
         """ % {"diff_filter": diff_filter}
         self.add_to_query(query=query)