infrahub-server 1.7.0rc0__py3-none-any.whl → 1.7.2__py3-none-any.whl

infrahub/core/migrations/graph/m054_cleanup_orphaned_nodes.py

@@ -0,0 +1,87 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Sequence
+
+from infrahub.core.migrations.shared import GraphMigration, MigrationInput, MigrationResult
+
+from ...query import Query, QueryType
+
+if TYPE_CHECKING:
+    from infrahub.database import InfrahubDatabase
+
+
+class CleanupOrphanedNodesQuery(Query):
+    """
+    Clean up orphaned Node vertices (no IS_PART_OF edge to Root) and their linked
+    Attributes and Relationships.
+    """
+
+    name = "cleanup_orphaned_nodes"
+    type = QueryType.WRITE
+    insert_return = False
+    raise_error_if_empty = False
+
+    async def query_init(self, db: InfrahubDatabase, **kwargs: dict[str, Any]) -> None:  # noqa: ARG002
+        query = """
+// Delete attributes of orphaned nodes
+MATCH (n:Node)
+WHERE NOT exists((n)-[:IS_PART_OF]->(:Root))
+OPTIONAL MATCH (n)-[:HAS_ATTRIBUTE]->(attr:Attribute)
+WITH DISTINCT attr
+CALL (attr) {
+    DETACH DELETE attr
+} IN TRANSACTIONS
+
+// reduce the results to a single row
+WITH 1 AS one
+LIMIT 1
+
+// Delete relationships that will have < 2 Node peers after orphaned node removal
+OPTIONAL MATCH (orphan:Node)-[:IS_RELATED]-(rel:Relationship)
+WHERE NOT exists((orphan)-[:IS_PART_OF]->(:Root))
+WITH DISTINCT rel
+CALL (rel) {
+    OPTIONAL MATCH (rel)-[:IS_RELATED]-(peer:Node)
+    WHERE exists((peer)-[:IS_PART_OF]->(:Root))
+    WITH rel, count(peer) AS remaining_peers
+    WHERE remaining_peers < 2
+    DETACH DELETE rel
+} IN TRANSACTIONS
+
+// reduce the results to a single row
+WITH 1 AS one
+LIMIT 1
+
+// Delete the orphaned nodes
+MATCH (n:Node)
+WHERE NOT exists((n)-[:IS_PART_OF]->(:Root))
+CALL (n) {
+    DETACH DELETE n
+} IN TRANSACTIONS
+        """
+        self.add_to_query(query)
+
+
+class Migration054(GraphMigration):
+    """
+    Clean up orphaned Node vertices that have no IS_PART_OF edge to Root.
+
+    This can happen when a branch-aware node is deleted during branch deletion,
+    but its branch-agnostic attributes or relationships are not properly cleaned up.
+
+    The migration:
+    1. DETACH DELETEs Attributes linked to orphaned nodes
+    2. DETACH DELETEs Relationships that would have < 2 Node peers after orphaned node removal
+    3. DETACH DELETEs the orphaned nodes themselves
+    """
+
+    name: str = "054_cleanup_orphaned_nodes"
+    minimum_version: int = 53
+    queries: Sequence[type[Query]] = [CleanupOrphanedNodesQuery]
+
+    async def validate_migration(self, db: InfrahubDatabase) -> MigrationResult:  # noqa: ARG002
+        return MigrationResult()
+
+    async def execute(self, migration_input: MigrationInput) -> MigrationResult:
+        # Override parent class to skip transaction in case there are many nodes to delete
+        return await self.do_execute(migration_input=migration_input)

infrahub/core/migrations/graph/m055_remove_webhook_validate_certificates_default.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Sequence
+
+from infrahub.core.constants import NULL_VALUE
+from infrahub.core.migrations.shared import MigrationResult
+from infrahub.core.query import Query, QueryType
+
+from ..shared import GraphMigration
+
+if TYPE_CHECKING:
+    from infrahub.database import InfrahubDatabase
+
+
+class Migration055Query01(Query):
+    """Remove the default_value from CoreWebhook's validate_certificates attribute.
+
+    This migration finds the CoreWebhook SchemaGeneric, locates its validate_certificates
+    SchemaAttribute, and sets the default_value to NULL. This removes the previous default
+    of True.
+    """
+
+    name = "migration_055_01"
+    type: QueryType = QueryType.WRITE
+
+    async def query_init(self, db: InfrahubDatabase, **kwargs: dict[str, Any]) -> None:  # noqa: ARG002
+        self.params["at"] = self.at.to_string()
+        self.params["null_value"] = NULL_VALUE
+
+        query = """
+// get the generic schema
+MATCH p1 = (sg:SchemaGeneric)-[:HAS_ATTRIBUTE]->(:Attribute {name: "name"})-[:HAS_VALUE]->(:AttributeValueIndexed {value: "Webhook"})
+WHERE all(r IN relationships(p1) WHERE r.status = "active" AND r.to IS NULL)
+
+// for safety, also check that the sg is in the namespace "Core"
+MATCH p2 = (sg)-[:HAS_ATTRIBUTE]->(:Attribute {name: "namespace"})-[:HAS_VALUE]->(:AttributeValueIndexed {value: "Core"})
+WHERE all(r IN relationships(p2) WHERE r.status = "active" AND r.to IS NULL)
+
+// there should only be 1 CoreWebhook schema generic
+WITH sg
+LIMIT 1
+
+// find the validate_certificates attribute
+MATCH p3 = (sg)-[:IS_RELATED]-(:Relationship {name: "schema__node__attributes"})
+    -[:IS_RELATED]-(sa:SchemaAttribute)
+    -[:HAS_ATTRIBUTE]->(:Attribute {name: "name"})
+    -[:HAS_VALUE]->(:AttributeValueIndexed {value: "validate_certificates"})
+WHERE all(r IN relationships(p3) WHERE r.status = "active" AND r.to IS NULL)
+// there should only be 1 validate_certificates attribute
+WITH sa
+LIMIT 1
+
+// get the default_value Attribute
+MATCH (sa)-[ha:HAS_ATTRIBUTE]->(default_value_attr:Attribute {name: "default_value"})-[hv:HAS_VALUE]->(default_value)
+WHERE all(r IN [ha, hv] WHERE r.status = "active" AND r.to IS NULL)
+LIMIT 1
+
+// skip if it is already NULL
+WITH sa, default_value_attr, hv, default_value
+WHERE default_value.value <> $null_value
+
+// close the HAS_VALUE edge for the current default value
+SET hv.to = $at
+
+// get the new value
+MERGE (new_value:AttributeValue:AttributeValueIndexed {value: $null_value, is_default: true})
+LIMIT 1
+
+// link the new value
+CREATE (default_value_attr)-[new_hv:HAS_VALUE]->(new_value)
+SET new_hv = properties(hv)
+SET new_hv.from = $at, new_hv.to = NULL
+        """
+        self.add_to_query(query)
+        self.return_labels = ["new_value"]
+
+
+class Migration055(GraphMigration):
+    name: str = "055_remove_webhook_validate_certificates_default"
+    queries: Sequence[type[Query]] = [Migration055Query01]
+    minimum_version: int = 54
+
+    async def validate_migration(self, db: InfrahubDatabase) -> MigrationResult:  # noqa: ARG002
+        result = MigrationResult()
+
+        return result

infrahub/core/migrations/runner.py

@@ -7,10 +7,11 @@ from infrahub.core.constants import GLOBAL_BRANCH_NAME
 from infrahub.core.migrations.graph import MIGRATIONS
 
 from .exceptions import MigrationFailureError
-from .shared import MigrationRequiringRebase
+from .shared import MigrationInput, MigrationRequiringRebase
 
 if TYPE_CHECKING:
     from infrahub.core.branch import Branch
+    from infrahub.core.timestamp import Timestamp
     from infrahub.database import InfrahubDatabase
 
 
@@ -35,12 +36,14 @@ class MigrationRunner:
     def has_migrations(self) -> bool:
         return bool(self.applicable_migrations)
 
-    async def run(self, db: InfrahubDatabase) -> None:
+    async def run(self, db: InfrahubDatabase, at: Timestamp) -> None:
         if not self.has_migrations():
             return
 
         for migration in self.applicable_migrations:
-            execution_result = await migration.execute_against_branch(db=db, branch=self.branch)
+            execution_result = await migration.execute_against_branch(
+                migration_input=MigrationInput(db=db, at=at), branch=self.branch
+            )
             validation_result = None
 
             if execution_result.success:

infrahub/core/migrations/schema/attribute_kind_update.py

@@ -2,15 +2,13 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING, Any, Sequence
 
-from infrahub.core.constants import SYSTEM_USER_ID
 from infrahub.types import is_large_attribute_type
 
 from ..query import AttributeMigrationQuery, MigrationBaseQuery
-from ..shared import AttributeSchemaMigration, MigrationResult
+from ..shared import AttributeSchemaMigration, MigrationInput, MigrationResult
 
 if TYPE_CHECKING:
     from infrahub.core.branch.models import Branch
-    from infrahub.core.timestamp import Timestamp
     from infrahub.database import InfrahubDatabase
 
 
@@ -40,7 +38,7 @@ class AttributeKindUpdateMigrationQuery(AttributeMigrationQuery):
 // ------------
 // start with all the Attribute vertices we might care about
 // ------------
-MATCH (n:%(schema_kind)s)-[:HAS_ATTRIBUTE]->(attr:Attribute)
+MATCH (n:%(schema_kinds)s)-[:HAS_ATTRIBUTE]->(attr:Attribute)
 WHERE attr.name = $attr_name
 WITH DISTINCT n, attr
 
@@ -76,7 +74,7 @@ CALL (av_is_default, av_value) {
 // ------------
 WITH 1 AS one
 LIMIT 1
-MATCH (n:%(schema_kind)s)-[:HAS_ATTRIBUTE]->(attr:Attribute)
+MATCH (n:%(schema_kinds)s)-[:HAS_ATTRIBUTE]->(attr:Attribute)
 WHERE attr.name = $attr_name
 WITH DISTINCT n, attr
 
@@ -94,7 +92,6 @@ CALL (n, attr) {
     RETURN has_value_e, av
 }
 
-
 // ------------
 // create and update the HAS_VALUE edges
 // ------------
@@ -154,7 +151,9 @@ CALL (attr, n) {
     SET n.updated_at = $at, n.updated_by = $user_id
 }
         """ % {
-            "schema_kind": self.migration.previous_schema.kind,
+            "schema_kinds": (
+                f"{self.migration.previous_schema.kind}|Profile{self.migration.previous_schema.kind}|Template{self.migration.previous_schema.kind}"
+            ),
             "branch_filter": branch_filter,
             "new_attr_value_labels": new_attr_value_labels,
         }
@@ -167,15 +166,13 @@ class AttributeKindUpdateMigration(AttributeSchemaMigration):
 
     async def execute(
         self,
-        db: InfrahubDatabase,
+        migration_input: MigrationInput,
         branch: Branch,
-        at: Timestamp | str | None = None,
         queries: Sequence[type[MigrationBaseQuery]] | None = None,
-        user_id: str = SYSTEM_USER_ID,
     ) -> MigrationResult:
         is_indexed_previous = is_large_attribute_type(self.previous_attribute_schema.kind)
         is_indexed_new = is_large_attribute_type(self.new_attribute_schema.kind)
         if is_indexed_previous is is_indexed_new:
             return MigrationResult()
 
-        return await super().execute(db=db, branch=branch, at=at, queries=queries, user_id=user_id)
+        return await super().execute(migration_input=migration_input, branch=branch, queries=queries)

infrahub/core/migrations/schema/attribute_supports_profile.py

@@ -2,20 +2,17 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING, Any, Sequence
 
-from infrahub.core.constants import SYSTEM_USER_ID
 from infrahub.core.migrations.query.attribute_remove import AttributeRemoveQuery
 from infrahub.core.schema.generic_schema import GenericSchema
 from infrahub.core.schema.node_schema import NodeSchema
 
 from ..query import AttributeMigrationQuery, MigrationBaseQuery
 from ..query.attribute_add import AttributeAddQuery
-from ..shared import AttributeSchemaMigration, MigrationResult
+from ..shared import AttributeSchemaMigration, MigrationInput, MigrationResult
 
 if TYPE_CHECKING:
     from infrahub.core.branch.models import Branch
     from infrahub.core.schema import MainSchemaTypes
-    from infrahub.core.timestamp import Timestamp
-    from infrahub.database import InfrahubDatabase
 
 
 def _get_node_kinds(schema: MainSchemaTypes) -> list[str]:
@@ -70,11 +67,9 @@ class AttributeSupportsProfileUpdateMigration(AttributeSchemaMigration):
 
     async def execute(
         self,
-        db: InfrahubDatabase,
+        migration_input: MigrationInput,
         branch: Branch,
-        at: Timestamp | str | None = None,
         queries: Sequence[type[MigrationBaseQuery]] | None = None,  # noqa: ARG002
-        user_id: str = SYSTEM_USER_ID,
     ) -> MigrationResult:
         if (
             # no change in whether the attribute should be used on profiles
@@ -89,4 +84,4 @@ class AttributeSupportsProfileUpdateMigration(AttributeSchemaMigration):
         if not self.new_attribute_schema.support_profiles:
             profiles_queries.append(ProfilesAttributeRemoveMigrationQuery)
 
-        return await super().execute(db=db, branch=branch, at=at, queries=profiles_queries, user_id=user_id)
+        return await super().execute(migration_input=migration_input, branch=branch, queries=profiles_queries)

infrahub/core/migrations/schema/models.py

@@ -7,6 +7,7 @@ from infrahub.core.constants import SYSTEM_USER_ID
 from infrahub.core.models import SchemaUpdateMigrationInfo
 from infrahub.core.path import SchemaPath
 from infrahub.core.schema.schema_branch import SchemaBranch
+from infrahub.core.timestamp import Timestamp
 
 
 class SchemaApplyMigrationData(BaseModel):
@@ -16,6 +17,7 @@ class SchemaApplyMigrationData(BaseModel):
     new_schema: SchemaBranch
     previous_schema: SchemaBranch
     migrations: list[SchemaUpdateMigrationInfo]
+    at: Timestamp
     user_id: str = SYSTEM_USER_ID
 
     @model_serializer()
@@ -26,6 +28,7 @@ class SchemaApplyMigrationData(BaseModel):
             "new_schema": self.new_schema.to_dict_schema_object(),
             "migrations": [migration.model_dump() for migration in self.migrations],
             "user_id": self.user_id,
+            "at": self.at.to_string(),
         }
 
     @field_validator("new_schema", "previous_schema", mode="before")
@@ -33,6 +36,11 @@ class SchemaApplyMigrationData(BaseModel):
     def validate_schema_branch(cls, value: Any) -> SchemaBranch:
         return SchemaBranch.validate(data=value)
 
+    @field_validator("at", mode="before")
+    @classmethod
+    def validate_at(cls, value: Any) -> Timestamp:
+        return Timestamp(value)
+
 
 class SchemaMigrationPathResponseData(BaseModel):
     errors: list[str] = Field(default_factory=list)

infrahub/core/migrations/schema/node_attribute_add.py

@@ -3,16 +3,14 @@ from __future__ import annotations
 from typing import TYPE_CHECKING, Any, Sequence
 
 from infrahub.core import registry
-from infrahub.core.constants import SYSTEM_USER_ID
 from infrahub.core.node import Node
 from infrahub.core.schema.generic_schema import GenericSchema
 from infrahub.core.schema.node_schema import NodeSchema
-from infrahub.exceptions import PoolExhaustedError
 from infrahub.tasks.registry import update_branch_registry
 
 from ..query import AttributeMigrationQuery, MigrationBaseQuery
 from ..query.attribute_add import AttributeAddQuery
-from ..shared import AttributeSchemaMigration, MigrationResult
+from ..shared import AttributeSchemaMigration, MigrationInput, MigrationResult
 
 if TYPE_CHECKING:
     from infrahub.core.node.resource_manager.number_pool import CoreNumberPool
@@ -21,7 +19,6 @@ if TYPE_CHECKING:
     from infrahub.database import InfrahubDatabase
 
     from ...branch import Branch
-    from ...timestamp import Timestamp
 
 
 class NodeAttributeAddMigrationQuery01(AttributeMigrationQuery, AttributeAddQuery):
@@ -62,32 +59,32 @@ class NodeAttributeAddMigration(AttributeSchemaMigration):
 
     async def execute(
         self,
-        db: InfrahubDatabase,
+        migration_input: MigrationInput,
         branch: Branch,
-        at: Timestamp | str | None = None,
         queries: Sequence[type[MigrationBaseQuery]] | None = None,
-        user_id: str = SYSTEM_USER_ID,
     ) -> MigrationResult:
         if self.new_attribute_schema.inherited is True:
             return MigrationResult()
-        return await super().execute(db=db, branch=branch, at=at, queries=queries, user_id=user_id)
+        return await super().execute(migration_input=migration_input, branch=branch, queries=queries)
 
     async def execute_post_queries(
         self,
-        db: InfrahubDatabase,
+        migration_input: MigrationInput,
         result: MigrationResult,
         branch: Branch,
-        at: Timestamp,  # noqa: ARG002
-        user_id: str,  # noqa: ARG002
     ) -> MigrationResult:
         if self.new_attribute_schema.kind != "NumberPool":
             return result
 
+        db = migration_input.db
+        at = migration_input.at
+
         number_pool: CoreNumberPool = await Node.fetch_or_create_number_pool(
             db=db,
             branch=branch,
             schema_node=self.new_schema,  # type: ignore
             schema_attribute=self.new_attribute_schema,
+            at=at,
         )
 
         await update_branch_registry(db=db, branch=branch)
@@ -96,23 +93,21 @@ class NodeAttributeAddMigration(AttributeSchemaMigration):
             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])
+        async def allocate_numbers(db: InfrahubDatabase) -> None:
+            for node in nodes:
+                number = await number_pool.get_resource(
+                    db=db, branch=branch, node=node, attribute=self.new_attribute_schema, at=at
+                )
+                attr = node.get_attribute(name=self.new_attribute_schema.name)
+                attr.value = number
+                attr.set_source(number_pool)
+
+                await node.save(db=db, fields=[self.new_attribute_schema.name], at=at)
+
+        if db.is_transaction:
+            await allocate_numbers(db=db)
+        else:
+            async with db.start_transaction() as dbt:
+                await allocate_numbers(db=dbt)
 
         return result

infrahub/core/migrations/schema/tasks.py

@@ -10,6 +10,7 @@ from prefect.logging import get_run_logger
 from infrahub.core.branch import Branch  # noqa: TC001
 from infrahub.core.constants import SYSTEM_USER_ID
 from infrahub.core.migrations import MIGRATION_MAP
+from infrahub.core.migrations.shared import MigrationInput
 from infrahub.core.path import SchemaPath  # noqa: TC001
 from infrahub.workers.dependencies import get_database
 from infrahub.workflows.utils import add_branch_tag
@@ -18,6 +19,7 @@ from .models import SchemaApplyMigrationData, SchemaMigrationPathResponseData
 
 if TYPE_CHECKING:
     from infrahub.core.schema import MainSchemaTypes
+    from infrahub.core.timestamp import Timestamp
     from infrahub.database import InfrahubDatabase
 
 
@@ -59,6 +61,7 @@ async def schema_apply_migrations(message: SchemaApplyMigrationData) -> list[str
             schema_path=migration.path,
             database=await get_database(),
             user_id=message.user_id,
+            at=message.at,
         )
 
     async for _, result in batch.execute():
@@ -79,6 +82,7 @@ async def schema_path_migrate(
     migration_name: str,
     schema_path: SchemaPath,
     database: InfrahubDatabase,
+    at: Timestamp,
     new_node_schema: MainSchemaTypes | None = None,
     previous_node_schema: MainSchemaTypes | None = None,
     user_id: str = SYSTEM_USER_ID,
@@ -104,7 +108,9 @@ async def schema_path_migrate(
             previous_node_schema=previous_node_schema,  # type: ignore[arg-type]
             schema_path=schema_path,
         )
-        execution_result = await migration.execute(db=db, branch=branch, user_id=user_id)
+        execution_result = await migration.execute(
+            migration_input=MigrationInput(db=db, at=at, user_id=user_id), branch=branch
+        )
 
         log.info(f"Migration completed for {migration_name}")
         log.debug(f"execution_result {execution_result}")

infrahub/core/migrations/shared.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, Sequence
 
 from pydantic import BaseModel, ConfigDict, Field
@@ -46,6 +47,13 @@ class MigrationResult(BaseModel):
         return False
 
 
+@dataclass
+class MigrationInput:
+    db: InfrahubDatabase
+    at: Timestamp = field(default_factory=Timestamp)
+    user_id: str = SYSTEM_USER_ID
+
+
 class SchemaMigration(BaseModel):
     model_config = ConfigDict(arbitrary_types_allowed=True)
     name: str = Field(..., description="Name of the migration")
@@ -59,37 +67,37 @@ class SchemaMigration(BaseModel):
 
     async def execute_pre_queries(
         self,
-        db: InfrahubDatabase,  # noqa: ARG002
+        migration_input: MigrationInput,  # noqa: ARG002
         result: MigrationResult,
         branch: Branch,  # noqa: ARG002
-        at: Timestamp,  # noqa: ARG002
-        user_id: str,  # noqa: ARG002
     ) -> MigrationResult:
         return result
 
     async def execute_post_queries(
         self,
-        db: InfrahubDatabase,  # noqa: ARG002
+        migration_input: MigrationInput,  # noqa: ARG002
         result: MigrationResult,
         branch: Branch,  # noqa: ARG002
-        at: Timestamp,  # noqa: ARG002
-        user_id: str,  # noqa: ARG002
     ) -> MigrationResult:
         return result
 
     async def execute_queries(
         self,
-        db: InfrahubDatabase,
+        migration_input: MigrationInput,
         result: MigrationResult,
         branch: Branch,
-        at: Timestamp,
         queries: Sequence[type[MigrationBaseQuery]],
-        user_id: str,
     ) -> MigrationResult:
         for migration_query in queries:
             try:
-                query = await migration_query.init(db=db, branch=branch, at=at, migration=self, user_id=user_id)
-                await query.execute(db=db)
+                query = await migration_query.init(
+                    db=migration_input.db,
+                    branch=branch,
+                    at=migration_input.at,
+                    migration=self,
+                    user_id=migration_input.user_id,
+                )
+                await query.execute(db=migration_input.db)
                 result.nbr_migrations_executed += query.get_nbr_migrations_executed()
             except Exception as exc:
                 result.errors.append(str(exc))
@@ -99,22 +107,20 @@ class SchemaMigration(BaseModel):
 
     async def execute(
         self,
-        db: InfrahubDatabase,
+        migration_input: MigrationInput,
         branch: Branch,
-        at: Timestamp | str | None = None,
         queries: Sequence[type[MigrationBaseQuery]] | None = None,
-        user_id: str = SYSTEM_USER_ID,
     ) -> MigrationResult:
-        async with db.start_transaction() as ts:
+        async with migration_input.db.start_transaction() as ts:
             result = MigrationResult()
-            at = Timestamp(at)
+            txn_migration_input = MigrationInput(db=ts, at=migration_input.at, user_id=migration_input.user_id)
 
-            await self.execute_pre_queries(db=ts, result=result, branch=branch, at=at, user_id=user_id)
+            await self.execute_pre_queries(migration_input=txn_migration_input, result=result, branch=branch)
             queries_to_execute = queries or self.queries
             await self.execute_queries(
-                db=ts, result=result, branch=branch, at=at, queries=queries_to_execute, user_id=user_id
+                migration_input=txn_migration_input, result=result, branch=branch, queries=queries_to_execute
             )
-            await self.execute_post_queries(db=ts, result=result, branch=branch, at=at, user_id=user_id)
+            await self.execute_post_queries(migration_input=txn_migration_input, result=result, branch=branch)
 
         return result
 
@@ -174,16 +180,17 @@ class GraphMigration(BaseModel):
     async def validate_migration(self, db: InfrahubDatabase) -> MigrationResult:
         raise NotImplementedError
 
-    async def execute(self, db: InfrahubDatabase) -> MigrationResult:
-        async with db.start_transaction() as ts:
-            return await self.do_execute(db=ts)
+    async def execute(self, migration_input: MigrationInput) -> MigrationResult:
+        async with migration_input.db.start_transaction() as ts:
+            txn_migration_input = MigrationInput(db=ts, at=migration_input.at)
+            return await self.do_execute(migration_input=txn_migration_input)
 
-    async def do_execute(self, db: InfrahubDatabase) -> MigrationResult:
+    async def do_execute(self, migration_input: MigrationInput) -> MigrationResult:
         result = MigrationResult()
         for migration_query in self.queries:
             try:
-                query = await migration_query.init(db=db)
-                await query.execute(db=db)
+                query = await migration_query.init(db=migration_input.db, at=migration_input.at)
+                await query.execute(db=migration_input.db)
             except Exception as exc:
                 result.errors.append(str(exc))
                 return result
@@ -216,14 +223,14 @@ class InternalSchemaMigration(BaseModel):
     async def validate_migration(self, db: InfrahubDatabase) -> MigrationResult:
         raise NotImplementedError
 
-    async def execute(self, db: InfrahubDatabase, user_id: str = SYSTEM_USER_ID) -> MigrationResult:
+    async def execute(self, migration_input: MigrationInput) -> MigrationResult:
         result = MigrationResult()
 
         default_branch = registry.get_branch_from_registry()
 
         for migration in self.migrations:
             try:
-                execution_result = await migration.execute(db=db, branch=default_branch, user_id=user_id)
+                execution_result = await migration.execute(migration_input=migration_input, branch=default_branch)
                 result.errors.extend(execution_result.errors)
             except Exception as exc:
                 result.errors.append(str(exc))
@@ -243,7 +250,7 @@ class ArbitraryMigration(BaseModel):
     async def validate_migration(self, db: InfrahubDatabase) -> MigrationResult:
         raise NotImplementedError()
 
-    async def execute(self, db: InfrahubDatabase) -> MigrationResult:
+    async def execute(self, migration_input: MigrationInput) -> MigrationResult:
         raise NotImplementedError()
 
 
@@ -259,11 +266,11 @@ class MigrationRequiringRebase(BaseModel):
     async def validate_migration(self, db: InfrahubDatabase) -> MigrationResult:
         raise NotImplementedError()
 
-    async def execute_against_branch(self, db: InfrahubDatabase, branch: Branch) -> MigrationResult:
+    async def execute_against_branch(self, migration_input: MigrationInput, branch: Branch) -> MigrationResult:
         """Method that will be run against non-default branches, it assumes that the branches have been rebased."""
         raise NotImplementedError()
 
-    async def execute(self, db: InfrahubDatabase) -> MigrationResult:
+    async def execute(self, migration_input: MigrationInput) -> MigrationResult:
         """Method that will be run against the default branch."""
         raise NotImplementedError()
 

infrahub/core/node/__init__.py

@@ -408,6 +408,7 @@ class Node(BaseNode, MetadataInterface, metaclass=BaseNodeMeta):
         schema_node: NodeSchema | GenericSchema,
         schema_attribute: AttributeSchema,
         branch: Branch | None = None,
+        at: Timestamp | None = None,
     ) -> CoreNumberPool:
         """Fetch or create a number pool based on the schema attribute parameters.
 
@@ -456,7 +457,7 @@ class Node(BaseNode, MetadataInterface, metaclass=BaseNodeMeta):
                     end_range=number_pool_parameters.end_range,
                     pool_type=NumberPoolType.SCHEMA.value,
                 )
-                await number_pool.save(db=db)
+                await number_pool.save(db=db, at=at)
 
                 # Do a lookup of the number pool to get the correct mapped type from the registry
                 # without this we don't get access to the .get_resource() method.