infrahub-server 1.3.2__py3-none-any.whl → 1.3.4__py3-none-any.whl

infrahub/core/initialization.py

@@ -6,6 +6,7 @@ from infrahub import config, lock
 from infrahub.constants.database import DatabaseType
 from infrahub.core import registry
 from infrahub.core.branch import Branch
+from infrahub.core.branch.enums import BranchStatus
 from infrahub.core.constants import (
     DEFAULT_IP_NAMESPACE,
     GLOBAL_BRANCH_NAME,
@@ -224,7 +225,7 @@ async def create_root_node(db: InfrahubDatabase) -> Root:
 async def create_default_branch(db: InfrahubDatabase) -> Branch:
     branch = Branch(
         name=registry.default_branch,
-        status="OPEN",
+        status=BranchStatus.OPEN,
         description="Default Branch",
         hierarchy_level=1,
         is_default=True,
@@ -241,7 +242,7 @@ async def create_default_branch(db: InfrahubDatabase) -> Branch:
 async def create_global_branch(db: InfrahubDatabase) -> Branch:
     branch = Branch(
         name=GLOBAL_BRANCH_NAME,
-        status="OPEN",
+        status=BranchStatus.OPEN,
         description="Global Branch",
         hierarchy_level=1,
         is_global=True,
@@ -264,7 +265,7 @@ async def create_branch(
     description = description or f"Branch {branch_name}"
     branch = Branch(
         name=branch_name,
-        status="OPEN",
+        status=BranchStatus.OPEN,
         hierarchy_level=2,
         description=description,
         is_default=False,

infrahub/core/merge.py

@@ -9,7 +9,7 @@ from infrahub.core.models import SchemaUpdateValidationResult
 from infrahub.core.protocols import CoreRepository
 from infrahub.core.registry import registry
 from infrahub.core.timestamp import Timestamp
-from infrahub.exceptions import ValidationError
+from infrahub.exceptions import MergeFailedError, ValidationError
 from infrahub.log import get_logger
 
 from ..git.models import GitRepositoryMerge
@@ -18,6 +18,7 @@ from ..workflows.catalogue import GIT_REPOSITORIES_MERGE
 if TYPE_CHECKING:
     from infrahub.core.branch import Branch
     from infrahub.core.diff.coordinator import DiffCoordinator
+    from infrahub.core.diff.diff_locker import DiffLocker
     from infrahub.core.diff.merger.merger import DiffMerger
     from infrahub.core.diff.model.path import EnrichedDiffRoot
     from infrahub.core.diff.repository.repository import DiffRepository
@@ -39,6 +40,7 @@ class BranchMerger:
         diff_coordinator: DiffCoordinator,
         diff_merger: DiffMerger,
         diff_repository: DiffRepository,
+        diff_locker: DiffLocker,
         destination_branch: Branch | None = None,
         service: InfrahubServices | None = None,
     ):
@@ -48,6 +50,7 @@ class BranchMerger:
         self.diff_coordinator = diff_coordinator
         self.diff_merger = diff_merger
         self.diff_repository = diff_repository
+        self.diff_locker = diff_locker
         self.migrations: list[SchemaUpdateMigrationInfo] = []
         self._merge_at = Timestamp()
 
@@ -185,22 +188,34 @@ class BranchMerger:
         )
         log.info("Diff updated for merge")
 
-        errors: list[str] = []
-        async for conflict_path, conflict in self.diff_repository.get_all_conflicts_for_diff(
-            diff_branch_name=self.source_branch.name, tracking_id=BranchTrackingId(name=self.source_branch.name)
+        log.info("Acquiring lock for merge")
+        async with self.diff_locker.acquire_lock(
+            target_branch_name=self.destination_branch.name,
+            source_branch_name=self.source_branch.name,
+            is_incremental=False,
         ):
-            if conflict.selected_branch is None or conflict.resolvable is False:
-                errors.append(conflict_path)
-
-        if errors:
-            raise ValidationError(
-                f"Unable to merge the branch '{self.source_branch.name}', conflict resolution missing: {', '.join(errors)}"
-            )
-
-        # TODO need to find a way to properly communicate back to the user any issue that could come up during the merge
-        # From the Graph or From the repositories
-        self._merge_at = Timestamp(at)
-        branch_diff = await self.diff_merger.merge_graph(at=self._merge_at)
+            log.info("Lock acquired for merge")
+            try:
+                errors: list[str] = []
+                async for conflict_path, conflict in self.diff_repository.get_all_conflicts_for_diff(
+                    diff_branch_name=self.source_branch.name, tracking_id=BranchTrackingId(name=self.source_branch.name)
+                ):
+                    if conflict.selected_branch is None or conflict.resolvable is False:
+                        errors.append(conflict_path)
+
+                if errors:
+                    raise ValidationError(
+                        f"Unable to merge the branch '{self.source_branch.name}', conflict resolution missing: {', '.join(errors)}"
+                    )
+
+                # TODO need to find a way to properly communicate back to the user any issue that could come up during the merge
+                # From the Graph or From the repositories
+                self._merge_at = Timestamp(at)
+                branch_diff = await self.diff_merger.merge_graph(at=self._merge_at)
+            except Exception as exc:
+                log.exception("Merge failed, beginning rollback")
+                await self.rollback()
+                raise MergeFailedError(branch_name=self.source_branch.name) from exc
         await self.merge_repositories()
         return branch_diff
 

infrahub/core/migrations/graph/__init__.py

@@ -33,6 +33,9 @@ from .m028_delete_diffs import Migration028
 from .m029_duplicates_cleanup import Migration029
 from .m030_illegal_edges import Migration030
 from .m031_check_number_attributes import Migration031
+from .m032_cleanup_orphaned_branch_relationships import Migration032
+from .m033_deduplicate_relationship_vertices import Migration033
+from .m034_find_orphaned_schema_fields import Migration034
 
 if TYPE_CHECKING:
     from infrahub.core.root import Root
@@ -71,6 +74,9 @@ MIGRATIONS: list[type[GraphMigration | InternalSchemaMigration | ArbitraryMigrat
     Migration029,
     Migration030,
     Migration031,
+    Migration032,
+    Migration033,
+    Migration034,
 ]
 
 
@@ -85,3 +91,23 @@ async def get_graph_migrations(
         applicable_migrations.append(migration)
 
     return applicable_migrations
+
+
+def get_migration_by_number(
+    migration_number: int | str,
+) -> GraphMigration | InternalSchemaMigration | ArbitraryMigration:
+    # Convert to string and pad with zeros if needed
+    try:
+        num = int(migration_number)
+        migration_str = f"{num:03d}"
+    except (ValueError, TypeError) as exc:
+        raise ValueError(f"Invalid migration number: {migration_number}") from exc
+
+    migration_name = f"Migration{migration_str}"
+
+    # Find the migration in the MIGRATIONS list
+    for migration_class in MIGRATIONS:
+        if migration_class.__name__ == migration_name:
+            return migration_class.init()
+
+    raise ValueError(f"Migration {migration_number} not found")

infrahub/core/migrations/graph/m012_convert_account_generic.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 from typing import TYPE_CHECKING, Any, Sequence
 
 from infrahub.core.branch import Branch
+from infrahub.core.branch.enums import BranchStatus
 from infrahub.core.constants import GLOBAL_BRANCH_NAME, BranchSupportType, InfrahubKind
 from infrahub.core.migrations.shared import MigrationResult
 from infrahub.core.query import Query, QueryType
@@ -20,7 +21,7 @@ if TYPE_CHECKING:
 
 global_branch = Branch(
     name=GLOBAL_BRANCH_NAME,
-    status="OPEN",
+    status=BranchStatus.OPEN,
     description="Global Branch",
     hierarchy_level=1,
     is_global=True,
@@ -29,7 +30,7 @@ global_branch = Branch(
 
 default_branch = Branch(
     name="main",
-    status="OPEN",
+    status=BranchStatus.OPEN,
     description="Default Branch",
     hierarchy_level=1,
     is_global=False,
@@ -105,7 +106,7 @@ class Migration012AddLabelData(NodeDuplicateQuery):
 
         branch = Branch(
             name=GLOBAL_BRANCH_NAME,
-            status="OPEN",
+            status=BranchStatus.OPEN,
             description="Global Branch",
             hierarchy_level=1,
             is_global=True,

infrahub/core/migrations/graph/m013_convert_git_password_credential.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 from typing import TYPE_CHECKING, Any, Sequence
 
 from infrahub.core.branch import Branch
+from infrahub.core.branch.enums import BranchStatus
 from infrahub.core.constants import (
     GLOBAL_BRANCH_NAME,
     BranchSupportType,
@@ -23,7 +24,7 @@ if TYPE_CHECKING:
 
 default_branch = Branch(
     name="main",
-    status="OPEN",
+    status=BranchStatus.OPEN,
     description="Default Branch",
     hierarchy_level=1,
     is_global=False,
@@ -42,7 +43,7 @@ class Migration013ConvertCoreRepositoryWithCred(Query):
 
         global_branch = Branch(
             name=GLOBAL_BRANCH_NAME,
-            status="OPEN",
+            status=BranchStatus.OPEN,
             description="Global Branch",
             hierarchy_level=1,
             is_global=True,
@@ -176,7 +177,7 @@ class Migration013ConvertCoreRepositoryWithoutCred(Query):
 
         global_branch = Branch(
             name=GLOBAL_BRANCH_NAME,
-            status="OPEN",
+            status=BranchStatus.OPEN,
             description="Global Branch",
             hierarchy_level=1,
             is_global=True,

infrahub/core/migrations/graph/m032_cleanup_orphaned_branch_relationships.py

@@ -0,0 +1,105 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from infrahub.core.migrations.shared import 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
+
+if TYPE_CHECKING:
+    from infrahub.database import InfrahubDatabase
+
+log = get_logger()
+
+
+class DeletedBranchCleanupQuery(Query):
+    """
+    Find all unique edge branch names for which there is no Branch object
+    """
+
+    name = "deleted_branch_cleanup"
+    type = QueryType.WRITE
+    insert_return = False
+
+    async def query_init(self, db: InfrahubDatabase, **kwargs: dict[str, Any]) -> None:  # noqa: ARG002
+        query = """
+MATCH (b:Branch)
+WITH collect(DISTINCT b.name) AS branch_names
+MATCH ()-[e]->()
+WHERE e.branch IS NOT NULL
+AND NOT e.branch IN branch_names
+RETURN DISTINCT (e.branch) AS branch_name
+        """
+        self.add_to_query(query)
+        self.return_labels = ["branch_name"]
+
+
+class DeleteOrphanRelationshipsQuery(Query):
+    """
+    Find all Relationship vertices that link to fewer than 2 Node vertices and delete them
+    """
+
+    name = "delete_orphan_relationships"
+    type = QueryType.WRITE
+    insert_return = False
+
+    async def query_init(self, db: InfrahubDatabase, **kwargs: dict[str, Any]) -> None:  # noqa: ARG002
+        query = """
+MATCH (r:Relationship)-[:IS_RELATED]-(n:Node)
+WITH DISTINCT r, n
+WITH r, count(*) AS node_count
+WHERE node_count < 2
+DETACH DELETE r
+        """
+        self.add_to_query(query)
+
+
+class Migration032(ArbitraryMigration):
+    """
+    Delete edges for branches that were not completely deleted
+    """
+
+    name: str = "032_cleanup_deleted_branches"
+    minimum_version: int = 31
+
+    async def validate_migration(self, db: InfrahubDatabase) -> MigrationResult:  # noqa: ARG002
+        return MigrationResult()
+
+    async def execute(self, db: InfrahubDatabase) -> MigrationResult:
+        migration_result = MigrationResult()
+
+        try:
+            log.info("Get partially deleted branch names...")
+            orphaned_branches_query = await DeletedBranchCleanupQuery.init(db=db)
+            await orphaned_branches_query.execute(db=db)
+
+            orphaned_branch_names = []
+            for result in orphaned_branches_query.get_results():
+                branch_name = result.get_as_type("branch_name", str)
+                orphaned_branch_names.append(branch_name)
+
+            if not orphaned_branch_names:
+                log.info("No partially deleted branches found. All done.")
+                return migration_result
+
+            log.info(f"Found {len(orphaned_branch_names)} orphaned branch names: {orphaned_branch_names}")
+
+            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)
+                log.info(f"Branch '{branch_name}' cleaned up.")
+
+            log.info("Deleting orphaned relationships...")
+            delete_relationships_query = await DeleteOrphanRelationshipsQuery.init(db=db)
+            await delete_relationships_query.execute(db=db)
+            log.info("Orphaned relationships deleted.")
+
+        except Exception as exc:
+            migration_result.errors.append(str(exc))
+            log.exception("Error during branch cleanup")
+
+        return migration_result

infrahub/core/migrations/graph/m033_deduplicate_relationship_vertices.py

@@ -0,0 +1,97 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Sequence
+
+from infrahub.core.migrations.shared import GraphMigration, MigrationResult
+from infrahub.log import get_logger
+
+from ...query import Query, QueryType
+
+if TYPE_CHECKING:
+    from infrahub.database import InfrahubDatabase
+
+log = get_logger()
+
+
+class DeduplicateRelationshipVerticesQuery(Query):
+    """
+    For each group of duplicate Relationships with the same UUID, delete any Relationship that meets the following criteria:
+    - is linked to a deleted node (only if the delete time is before the Relationship's from time)
+    - is linked to a node on an incorrect branch (ie Relationship added on main, but Node is on a branch)
+    """
+
+    name = "deduplicate_relationship_vertices"
+    type = QueryType.WRITE
+    insert_return = False
+
+    async def query_init(self, db: InfrahubDatabase, **kwargs: dict[str, Any]) -> None:  # noqa: ARG002
+        query = """
+MATCH (root:Root)
+WITH root.default_branch AS default_branch_name
+// ------------
+// Find all Relationship vertices with duplicate UUIDs
+// ------------
+MATCH (r:Relationship)
+WITH r.uuid AS r_uuid, default_branch_name, count(*) AS num_dups
+WHERE num_dups > 1
+WITH DISTINCT r_uuid, default_branch_name
+// ------------
+// get the branched_from time for each relationship edge and node
+// ------------
+MATCH (rel:Relationship {uuid: r_uuid})
+CALL (rel) {
+    MATCH (rel)-[is_rel_e:IS_RELATED {status: "active"}]-(n:Node)
+    MATCH (rel_branch:Branch {name: is_rel_e.branch})
+    RETURN is_rel_e, rel_branch.branched_from AS rel_branched_from, n
+}
+// ------------
+// for each IS_RELATED edge of the relationship, check the IS_PART_OF edges of the Node vertex
+// to determine if this side of the relationship is legal
+// ------------
+CALL (n, is_rel_e, rel_branched_from, default_branch_name) {
+    OPTIONAL MATCH (n)-[is_part_of_e:IS_PART_OF {status: "active"}]->(:Root)
+    WHERE (
+        // the Node's create time must precede the Relationship's create time
+        is_part_of_e.from <= is_rel_e.from AND (is_part_of_e.to >= is_rel_e.from OR is_part_of_e.to IS NULL)
+        // the Node must have been created on a branch of equal or lesser depth than the Relationship
+        AND is_part_of_e.branch_level <= is_rel_e.branch_level
+        // if the Node and Relationships were created on branch_level = 2, then they must be on the same branch
+        AND (
+            is_part_of_e.branch_level = 1
+            OR is_part_of_e.branch = is_rel_e.branch
+        )
+        // if the Node was created on the default branch, and the Relationship was created on a branch,
+        // then the Node must have been created after the branched_from time of the Relationship's branch
+        AND (
+            is_part_of_e.branch <> default_branch_name
+            OR is_rel_e.branch_level = 1
+            OR is_part_of_e.from <= rel_branched_from
+        )
+    )
+    WITH is_part_of_e IS NOT NULL AS is_legal
+    ORDER BY is_legal DESC
+    RETURN is_legal
+    LIMIT 1
+}
+WITH rel, is_legal
+ORDER BY rel, is_legal ASC
+WITH rel, head(collect(is_legal)) AS is_legal
+WHERE is_legal = false
+DETACH DELETE rel
+        """
+        self.add_to_query(query)
+
+
+class Migration033(GraphMigration):
+    """
+    Identifies duplicate Relationship vertices that have the same UUID property. Deletes any duplicates that
+    are linked to deleted nodes or nodes on in incorrect branch.
+    """
+
+    name: str = "033_deduplicate_relationship_vertices"
+    minimum_version: int = 31
+    queries: Sequence[type[Query]] = [DeduplicateRelationshipVerticesQuery]
+
+    async def validate_migration(self, db: InfrahubDatabase) -> MigrationResult:  # noqa: ARG002
+        result = MigrationResult()
+        return result

infrahub/core/migrations/graph/m034_find_orphaned_schema_fields.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+from collections import defaultdict
+from typing import TYPE_CHECKING, Any, Sequence
+
+from infrahub.core.initialization import initialization
+from infrahub.core.manager import NodeManager
+from infrahub.core.migrations.shared import ArbitraryMigration, MigrationResult
+from infrahub.core.timestamp import Timestamp
+from infrahub.lock import initialize_lock
+from infrahub.log import get_logger
+
+from ...query import Query, QueryType
+
+if TYPE_CHECKING:
+    from infrahub.database import InfrahubDatabase
+
+log = get_logger()
+
+
+class FindOrphanedSchemaFieldsQuery(Query):
+    name = "find_orphaned_schema_fields"
+    type = QueryType.WRITE
+
+    async def query_init(self, db: InfrahubDatabase, **kwargs: dict[str, Any]) -> None:  # noqa: ARG002
+        query = """
+// ------------
+// Find orphaned SchemaRelationship and SchemaAttribute vertices
+// ------------
+MATCH (schema_field:SchemaRelationship|SchemaAttribute)-[e:IS_RELATED]-(rel:Relationship)
+WHERE rel.name IN ["schema__node__relationships", "schema__node__attributes"]
+AND e.status = "deleted" OR e.to IS NOT NULL
+WITH schema_field, e.branch AS branch, CASE
+    WHEN e.status = "deleted" THEN e.from
+    ELSE e.to
+END AS delete_time
+CALL (schema_field, branch) {
+    OPTIONAL MATCH (schema_field)-[is_part_of:IS_PART_OF {branch: branch}]->(:Root)
+    WHERE is_part_of.status = "deleted" OR is_part_of.to IS NOT NULL
+    RETURN is_part_of IS NOT NULL AS is_deleted
+}
+WITH schema_field, branch, delete_time
+WHERE is_deleted = FALSE
+        """
+        self.add_to_query(query)
+        self.return_labels = ["schema_field.uuid AS schema_field_uuid", "branch", "delete_time"]
+
+
+class Migration034(ArbitraryMigration):
+    """
+    Finds active SchemaRelationship and SchemaAttribute vertices with deleted relationships to SchemaNodes or
+    SchemaGenerics and deletes them on the same branch at the same time
+    """
+
+    name: str = "034_find_orphaned_schema_fields"
+    minimum_version: int = 33
+    queries: Sequence[type[Query]] = [FindOrphanedSchemaFieldsQuery]
+
+    async def validate_migration(self, db: InfrahubDatabase) -> MigrationResult:  # noqa: ARG002
+        return MigrationResult()
+
+    async def execute(self, db: InfrahubDatabase) -> MigrationResult:
+        try:
+            initialize_lock()
+            await initialization(db=db)
+            query = await FindOrphanedSchemaFieldsQuery.init(db=db)
+            await query.execute(db=db)
+            schema_field_uuids_by_branch: dict[str, dict[str, str]] = defaultdict(dict)
+            for result in query.get_results():
+                schema_field_uuid = result.get_as_type("schema_field_uuid", return_type=str)
+                branch = result.get_as_type("branch", return_type=str)
+                delete_time = result.get_as_type("delete_time", return_type=str)
+                schema_field_uuids_by_branch[branch][schema_field_uuid] = delete_time
+
+            for branch, schema_rel_details in schema_field_uuids_by_branch.items():
+                node_map = await NodeManager.get_many(db=db, branch=branch, ids=list(schema_rel_details.keys()))
+                for schema_field_uuid, orphan_schema_rel_node in node_map.items():
+                    delete_time = Timestamp(schema_rel_details[schema_field_uuid])
+                    await orphan_schema_rel_node.delete(db=db, at=delete_time)
+        except Exception as exc:
+            log.exception("Error during orphaned schema field cleanup")
+            return MigrationResult(errors=[str(exc)])
+
+        return MigrationResult()

infrahub/core/migrations/schema/node_attribute_add.py

@@ -1,10 +1,22 @@
 from __future__ import annotations
 
-from typing import Any, Sequence
+from typing import TYPE_CHECKING, Any, Sequence
+
+from infrahub.core import registry
+from infrahub.core.node import Node
+from infrahub.exceptions import PoolExhaustedError
+from infrahub.tasks.registry import update_branch_registry
 
 from ..query import AttributeMigrationQuery
 from ..query.attribute_add import AttributeAddQuery
-from ..shared import AttributeSchemaMigration
+from ..shared import AttributeSchemaMigration, MigrationResult
+
+if TYPE_CHECKING:
+    from infrahub.core.node.resource_manager.number_pool import CoreNumberPool
+    from infrahub.database import InfrahubDatabase
+
+    from ...branch import Branch
+    from ...timestamp import Timestamp
 
 
 class NodeAttributeAddMigrationQuery01(AttributeMigrationQuery, AttributeAddQuery):
@@ -29,3 +41,44 @@ class NodeAttributeAddMigrationQuery01(AttributeMigrationQuery, AttributeAddQuer
 class NodeAttributeAddMigration(AttributeSchemaMigration):
     name: str = "node.attribute.add"
     queries: Sequence[type[AttributeMigrationQuery]] = [NodeAttributeAddMigrationQuery01]  # type: ignore[assignment]
+
+    async def execute_post_queries(
+        self,
+        db: InfrahubDatabase,
+        result: MigrationResult,
+        branch: Branch,
+        at: Timestamp,  # noqa: ARG002
+    ) -> MigrationResult:
+        if self.new_attribute_schema.kind != "NumberPool":
+            return result
+
+        number_pool: CoreNumberPool = await Node.fetch_or_create_number_pool(  # type: ignore[assignment]
+            db=db, branch=branch, schema_node=self.new_schema, schema_attribute=self.new_attribute_schema
+        )
+
+        await update_branch_registry(db=db, branch=branch)
+
+        nodes: list[Node] = await registry.manager.query(
+            db=db, branch=branch, schema=self.new_schema, fields={"id": True, self.new_attribute_schema.name: True}
+        )
+
+        try:
+            numbers = await number_pool.get_next_many(
+                db=db,
+                branch=branch,
+                quantity=len(nodes),
+                attribute=self.new_attribute_schema,
+            )
+        except PoolExhaustedError as exc:
+            result.errors.append(str(exc))
+            return result
+
+        for node, number in zip(nodes, numbers, strict=True):
+            await number_pool.reserve(db=db, number=number, identifier=node.get_id())
+            attr = getattr(node, self.new_attribute_schema.name)
+            attr.value = number
+            attr.source = number_pool.id
+
+            await node.save(db=db, fields=[self.new_attribute_schema.name])
+
+        return result

infrahub/core/migrations/shared.py

@@ -16,13 +16,13 @@ from infrahub.core.schema import (
     SchemaRoot,
     internal_schema,
 )
+from infrahub.core.timestamp import Timestamp
 
 from .query import MigrationQuery  # noqa: TC001
 
 if TYPE_CHECKING:
     from infrahub.core.branch import Branch
     from infrahub.core.schema.schema_branch import SchemaBranch
-    from infrahub.core.timestamp import Timestamp
     from infrahub.database import InfrahubDatabase
 
 
@@ -47,18 +47,46 @@ class SchemaMigration(BaseModel):
     previous_node_schema: NodeSchema | GenericSchema | None = None
     schema_path: SchemaPath
 
+    async def execute_pre_queries(
+        self,
+        db: InfrahubDatabase,  # noqa: ARG002
+        result: MigrationResult,
+        branch: Branch,  # noqa: ARG002
+        at: Timestamp,  # noqa: ARG002
+    ) -> MigrationResult:
+        return result
+
+    async def execute_post_queries(
+        self,
+        db: InfrahubDatabase,  # noqa: ARG002
+        result: MigrationResult,
+        branch: Branch,  # noqa: ARG002
+        at: Timestamp,  # noqa: ARG002
+    ) -> MigrationResult:
+        return result
+
+    async def execute_queries(
+        self, db: InfrahubDatabase, result: MigrationResult, branch: Branch, at: Timestamp
+    ) -> MigrationResult:
+        for migration_query in self.queries:
+            try:
+                query = await migration_query.init(db=db, branch=branch, at=at, migration=self)
+                await query.execute(db=db)
+                result.nbr_migrations_executed += query.get_nbr_migrations_executed()
+            except Exception as exc:
+                result.errors.append(str(exc))
+                return result
+
+        return result
+
     async def execute(self, db: InfrahubDatabase, branch: Branch, at: Timestamp | str | None = None) -> MigrationResult:
         async with db.start_transaction() as ts:
             result = MigrationResult()
+            at = Timestamp(at)
 
-            for migration_query in self.queries:
-                try:
-                    query = await migration_query.init(db=ts, branch=branch, at=at, migration=self)
-                    await query.execute(db=ts)
-                    result.nbr_migrations_executed += query.get_nbr_migrations_executed()
-                except Exception as exc:
-                    result.errors.append(str(exc))
-                    return result
+            await self.execute_pre_queries(db=ts, result=result, branch=branch, at=at)
+            await self.execute_queries(db=ts, result=result, branch=branch, at=at)
+            await self.execute_post_queries(db=ts, result=result, branch=branch, at=at)
 
         return result