fix(backend): delete branch data in bounded batches

backend/infrahub/core/branch/data_deleter.py

@@ -0,0 +1,151 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Protocol
+
+from infrahub.core.branch.enums import BranchStatus
+from infrahub.core.constants.database import DatabaseEdgeType
+from infrahub.core.query.branch import (
+    DeleteBranchAgnosticAttributesQuery,
+    DeleteBranchAgnosticRelationshipsQuery,
+    DeleteBranchEdgesQuery,
+)
+from infrahub.core.query.standard_node import StandardNodeDeleteQuery
+from infrahub.exceptions import ValidationError
+from infrahub.log import get_logger
+
+if TYPE_CHECKING:
+    from infrahub.core.branch.models import Branch
+    from infrahub.database import InfrahubDatabase
+
+# The agnostic cleanup batches Nodes, and each one can drag an unbounded number of peer vertices
+# into the transaction with it, so its batch is capped low.
+MAX_AGNOSTIC_PEER_BATCH_SIZE = 500
+
+
+@dataclass(frozen=True)
+class BranchDeleteResult:
+    """What a delete attempt actually did.
+
+    `branch_deleted` is false when the branch had already been removed by the time this attempt got
+    to it, which is how a caller knows not to repeat the work that follows a delete.
+    """
+
+    branch_deleted: bool
+    edges_removed: int
+
+
+class BranchDataDeleterInterface(Protocol):
+    """The database side of a branch delete."""
+
+    async def delete(self, branch: Branch) -> BranchDeleteResult: ...
+
+
+class LoggerInterface(Protocol):
+    """Just enough of a logger for progress reporting."""
+
+    def info(self, message: str, /) -> Any: ...
+
+
+class BranchDataDeleter:
+    """Remove a branch, every edge belonging to it, and the vertices that only it kept alive.
+
+    The graph work is split into one bounded query per batch so that no single transaction has to
+    hold the whole branch in memory. Each query is its own auto-commit transaction, which also means
+    an interrupted delete can be resumed by running the whole thing again.
+    """
+
+    def __init__(self, db: InfrahubDatabase, batch_size: int, log: LoggerInterface | None = None) -> None:
+        self.db = db
+        self.batch_size = batch_size
+        self.log = log or get_logger()
+
+    async def delete(self, branch: Branch) -> BranchDeleteResult:
+        """Remove the branch's data and then the branch itself.
+
+        Returns whether the Branch object was actually deleted in case multiple processes try to
+        delete concurrently so the caller can know which delete really succeeded.
+
+        Raises:
+            ValidationError: When the branch is the default branch or an internal one.
+
+        """
+        if branch.is_default:
+            raise ValidationError(f"Unable to delete {branch.name} it is the default branch.")
+        if branch.is_global:
+            raise ValidationError(f"Unable to delete {branch.name} this is an internal branch.")
+
+        if branch.status != BranchStatus.DELETING:
+            branch.status = BranchStatus.DELETING
+            await branch.save(db=self.db)
+
+        edges_removed = await self.delete_branch_data(branch_name=branch.name)
+
+        query = await StandardNodeDeleteQuery.init(db=self.db, node=branch)
+        await query.execute(db=self.db)
+        branch_deleted = query.stats.get_counter("nodes_deleted") > 0
+
+        return BranchDeleteResult(branch_deleted=branch_deleted, edges_removed=edges_removed)
+
+    async def delete_branch_data(self, branch_name: str) -> int:
+        """Remove a branch's data without requiring the branch itself to still exist.
+
+        Returns the number of edges removed, so a caller whose own logging is the only thing the
+        operator can see is able to report progress.
+        """
+        agnostic_edges_count = await self._delete_agnostic_peers(branch_name=branch_name)
+        branch_edges_count = await self._delete_edges(branch_name=branch_name)
+        return agnostic_edges_count + branch_edges_count
+
+    async def _delete_agnostic_peers(self, branch_name: str) -> int:
+        """Drop the agnostic attributes and relationships of Nodes that exist on no other branch.
+
+        Both queries locate those Nodes through the branch's IS_PART_OF edges, so this has to
+        finish before the edge deletion starts removing them. Resuming a delete that failed part
+        way through this stage is safe for the same reason: no IS_PART_OF edge has been touched yet.
+
+        Returns the number of edges removed, which is every edge of the peers detached here, not
+        only the agnostic ones that led to them.
+        """
+        batch_size = min(self.batch_size, MAX_AGNOSTIC_PEER_BATCH_SIZE)
+
+        relationships_query = await DeleteBranchAgnosticRelationshipsQuery.init(
+            db=self.db, branch_name=branch_name, batch_size=batch_size
+        )
+        await relationships_query.execute(db=self.db)
+
+        attributes_query = await DeleteBranchAgnosticAttributesQuery.init(
+            db=self.db, branch_name=branch_name, batch_size=batch_size
+        )
+        await attributes_query.execute(db=self.db)
+
+        edges_removed = relationships_query.stats.get_counter(
+            "relationships_deleted"
+        ) + attributes_query.stats.get_counter("relationships_deleted")
+        if edges_removed:
+            self.log.info(
+                f"Deleted agnostic peers of nodes only on branch '{branch_name}', {edges_removed} edge(s) removed"
+            )
+        return edges_removed
+
+    async def _delete_edges(self, branch_name: str) -> int:
+        edges_removed = 0
+        for edge_type in DatabaseEdgeType:
+            deleted_total = 0
+            while True:
+                # A fresh query per batch: the stats counters accumulate per instance, so a reused
+                # one would never report zero again and the loop would not end.
+                query = await DeleteBranchEdgesQuery.init(
+                    db=self.db, branch_name=branch_name, edge_type=edge_type, batch_size=self.batch_size
+                )
+                await query.execute(db=self.db)
+                deleted = query.deleted_edge_count()
+                if not deleted:
+                    break
+                deleted_total += deleted
+
+            if deleted_total:
+                edges_removed += deleted_total
+                self.log.info(f"Deleted {deleted_total} {edge_type.value} edge(s) on branch '{branch_name}'")
+
+        return edges_removed

backend/infrahub/core/branch/delete_coordinator.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Protocol
+
+from infrahub.events.branch_action import BranchDeletedEvent
+from infrahub.events.models import EventMeta
+from infrahub.exceptions import ValidationError
+from infrahub.workflows.catalogue import BRANCH_CANCEL_PROPOSED_CHANGES, GIT_REPOSITORIES_DELETE_BRANCH
+
+if TYPE_CHECKING:
+    from infrahub.context import InfrahubContext
+    from infrahub.core.branch.data_deleter import BranchDataDeleterInterface, BranchDeleteResult, LoggerInterface
+    from infrahub.core.branch.models import Branch
+    from infrahub.services.adapters.event import InfrahubEventService
+    from infrahub.services.adapters.workflow import InfrahubWorkflow
+
+
+class DiffFreezerInterface(Protocol):
+    """Interface for freezing diffs."""
+
+    async def freeze_diffs_for_branch(self, branch_name: str) -> None: ...
+
+
+class BranchDeleteOrchestrator:
+    """Delete a branch and do the work that follows from it.
+
+    Holds no database of its own: the deletion is delegated, which is what keeps the ordering and the
+    post-delete decisions here testable without one.
+    """
+
+    def __init__(
+        self,
+        data_deleter: BranchDataDeleterInterface,
+        diff_freezer: DiffFreezerInterface,
+        event_service: InfrahubEventService,
+        workflow: InfrahubWorkflow,
+        log: LoggerInterface,
+        global_branch: Branch,
+        delete_git_branch_after_merge: bool,
+    ) -> None:
+        self.data_deleter = data_deleter
+        self.diff_freezer = diff_freezer
+        self.event_service = event_service
+        self.workflow = workflow
+        self.log = log
+        self.global_branch = global_branch
+        self.delete_git_branch_after_merge = delete_git_branch_after_merge
+
+    async def delete(
+        self,
+        branch: Branch,
+        context: InfrahubContext,
+        delete_from_git: bool = False,
+        proposed_change_id: str | None = None,
+    ) -> BranchDeleteResult:
+        """Remove the branch, then cancel its proposed changes, announce it, and drop its Git branch.
+
+        Raises:
+            ValidationError: When the branch is the default branch or an internal one.
+
+        """
+        # Before the freeze, not after: a refused delete has to leave the branch's diffs alone.
+        if branch.is_default:
+            raise ValidationError(f"Unable to delete {branch.name} it is the default branch.")
+        if branch.is_global:
+            raise ValidationError(f"Unable to delete {branch.name} this is an internal branch.")
+
+        # Freezing has to precede the deletion, which takes away the branch name they are found by.
+        await self.diff_freezer.freeze_diffs_for_branch(branch_name=branch.name)
+
+        result = await self.data_deleter.delete(branch=branch)
+
+        if result.branch_deleted:
+            await self.workflow.submit_workflow(
+                workflow=BRANCH_CANCEL_PROPOSED_CHANGES, context=context, parameters={"branch_name": branch.name}
+            )
+            await self.event_service.send(
+                event=BranchDeletedEvent(
+                    branch_name=branch.name,
+                    branch_id=str(branch.uuid),
+                    sync_with_git=branch.sync_with_git,
+                    meta=EventMeta.from_context(context=context.to_event_context(), branch=self.global_branch),
+                    proposed_change_id=proposed_change_id,
+                )
+            )
+        else:
+            # Another attempt removed the branch, so the work above is already its responsibility.
+            self.log.info(f"Branch '{branch.name}' was already deleted")
+
+        # Always execute in case concurrent delete process with delete_from_git=False won the delete race.
+        if (self.delete_git_branch_after_merge or delete_from_git) and branch.sync_with_git:
+            await self.workflow.submit_workflow(
+                workflow=GIT_REPOSITORIES_DELETE_BRANCH, context=context, parameters={"branch": branch.name}
+            )
+
+        return result

backend/infrahub/core/branch/models.py

@@ -14,7 +14,6 @@
 from infrahub.core.query import Query, QueryType
 from infrahub.core.query.branch import (
     BranchNodeGetListQuery,
-    DeleteBranchRelationshipsQuery,
     RebaseBranchQuery,
 )
 from infrahub.core.registry import registry
@@ -315,17 +314,16 @@ async def create(self, db: InfrahubDatabase, user_id: str = SYSTEM_USER_ID) -> b
         return await super().create(db=db, user_id=user_id)
 
     async def delete(self, db: InfrahubDatabase) -> None:
-        if self.is_default:
-            raise ValidationError(f"Unable to delete {self.name} it is the default branch.")
-        if self.is_global:
-            raise ValidationError(f"Unable to delete {self.name} this is an internal branch.")
+        """Not supported on a Branch.
 
-        self.status = BranchStatus.DELETING
-        await self.save(db=db)
+        The inherited implementation would drop the Branch vertex and silently leave every edge and
+        vertex belonging to the branch behind, so it is refused rather than overridden.
 
-        query = await DeleteBranchRelationshipsQuery.init(db=db, branch_name=self.name)
-        await query.execute(db=db)
-        await super().delete(db=db)
+        Raises:
+            NotImplementedError: Always.
+
+        """
+        raise NotImplementedError("Unable to delete a Branch directly, use BranchDataDeleter instead.")
 
     def get_query_filter_relationships(
         self, rel_labels: list, at: Optional[Timestamp] = None, include_outside_parentheses: bool = False

backend/infrahub/core/branch/tasks.py

@@ -13,6 +13,8 @@
 from infrahub.core import registry
 from infrahub.core.branch import Branch
 from infrahub.core.branch.creator import BranchCreator
+from infrahub.core.branch.data_deleter import BranchDataDeleter
+from infrahub.core.branch.delete_coordinator import BranchDeleteOrchestrator
 from infrahub.core.branch.enums import BranchStatus
 from infrahub.core.changelog.diff import DiffChangelogCollector, MigrationTracker
 from infrahub.core.constants import DiffAction, MutationAction
@@ -36,7 +38,6 @@
 from infrahub.core.validators.tasks import schema_validate_migrations
 from infrahub.dependencies.registry import get_component_registry
 from infrahub.events.branch_action import (
-    BranchDeletedEvent,
     BranchMergedEvent,
     BranchMigratedEvent,
     BranchRebasedEvent,
@@ -54,7 +55,6 @@
     BRANCH_MERGE_POST_PROCESS,
     DIFF_REFRESH_ALL,
     DIFF_UPDATE,
-    GIT_REPOSITORIES_DELETE_BRANCH,
     IPAM_RECONCILIATION,
     TRIGGER_ARTIFACT_DEFINITION_GENERATE,
     TRIGGER_GENERATOR_DEFINITION_RUN,
@@ -538,37 +538,30 @@ async def delete_branch(
 ) -> None:
     await add_tags(branches=[branch], nodes=[proposed_change_id] if proposed_change_id else None)
     database = await get_database()
+    workflow = get_workflow()
+    event_service = await get_event_service()
     async with database.start_session() as db:
-        obj = await Branch.get_by_name(db=db, name=str(branch))
+        # ignore_deleting=False so that a delete which failed part way through can be run again:
+        obj = await Branch.get_by_name(db=db, name=str(branch), ignore_deleting=False)
 
         component_registry = get_component_registry()
         diff_repository = await component_registry.get_component(DiffRepository, db=db, branch=obj)
-        await diff_repository.freeze_diffs_for_branch(branch_name=branch)
-
-        await obj.delete(db=db)
 
-        event_context = context.to_event_context()
-        event = BranchDeletedEvent(
-            branch_name=branch,
-            branch_id=str(obj.uuid),
-            sync_with_git=obj.sync_with_git,
-            meta=EventMeta.from_context(context=event_context, branch=registry.get_global_branch()),
-            proposed_change_id=proposed_change_id,
-        )
-
-        await get_workflow().submit_workflow(
-            workflow=BRANCH_CANCEL_PROPOSED_CHANGES, context=context, parameters={"branch_name": branch}
+        log = get_run_logger()
+        orchestrator = BranchDeleteOrchestrator(
+            data_deleter=BranchDataDeleter(db=db, batch_size=config.SETTINGS.database.query_size_limit, log=log),
+            diff_freezer=diff_repository,
+            event_service=event_service,
+            workflow=workflow,
+            log=log,
+            global_branch=registry.get_global_branch(),
+            delete_git_branch_after_merge=config.SETTINGS.git.delete_git_branch_after_merge,
         )
-
-        event_service = await get_event_service()
-        await event_service.send(event=event)
-
-    should_delete_git = (config.SETTINGS.git.delete_git_branch_after_merge or delete_from_git) and obj.sync_with_git
-    if should_delete_git:
-        await get_workflow().submit_workflow(
-            workflow=GIT_REPOSITORIES_DELETE_BRANCH,
+        await orchestrator.delete(
+            branch=obj,
             context=context,
-            parameters={"branch": branch},
+            delete_from_git=delete_from_git,
+            proposed_change_id=proposed_change_id,
         )
 
 

backend/infrahub/core/graph/__init__.py

@@ -1 +1 @@
-GRAPH_VERSION = 74
+GRAPH_VERSION = 75

backend/infrahub/core/migrations/graph/m032_cleanup_orphaned_branch_relationships.py

@@ -2,9 +2,10 @@
 
 from typing import TYPE_CHECKING, Any
 
+from infrahub import config
+from infrahub.core.branch.data_deleter import BranchDataDeleter
 from infrahub.core.migrations.shared import MigrationInput, MigrationResult
 from infrahub.core.query import Query, QueryType
-from infrahub.core.query.branch import DeleteBranchRelationshipsQuery
 from infrahub.log import get_logger
 
 from ..shared import ArbitraryMigration
@@ -83,10 +84,10 @@ async def execute(self, migration_input: MigrationInput) -> MigrationResult:
 
             log.info(f"Found {len(orphaned_branch_names)} orphaned branch names: {orphaned_branch_names}")
 
+            deleter = BranchDataDeleter(db=db, batch_size=config.SETTINGS.database.query_size_limit)
             for branch_name in orphaned_branch_names:
                 log.info(f"Cleaning up branch '{branch_name}'...")
-                delete_query = await DeleteBranchRelationshipsQuery.init(db=db, branch_name=branch_name)
-                await delete_query.execute(db=db)
+                await deleter.delete_branch_data(branch_name=branch_name)
                 log.info(f"Branch '{branch_name}' cleaned up.")
 
             log.info("Deleting orphaned relationships...")