infrahub-server 1.7.0b0__py3-none-any.whl → 1.7.1__py3-none-any.whl

infrahub/core/migrations/graph/m051_subtract_branched_from_microsecond.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from infrahub.core.branch import Branch
+from infrahub.core.migrations.shared import ArbitraryMigration, MigrationInput, MigrationResult
+from infrahub.core.timestamp import Timestamp
+
+if TYPE_CHECKING:
+    from infrahub.database import InfrahubDatabase
+
+
+class Migration051(ArbitraryMigration):
+    name: str = "051_subtract_branched_from_microsecond"
+    minimum_version: int = 50
+
+    async def validate_migration(self, db: InfrahubDatabase) -> MigrationResult:  # noqa: ARG002
+        return MigrationResult()
+
+    async def execute(self, migration_input: MigrationInput) -> MigrationResult:
+        db = migration_input.db
+        result = MigrationResult()
+
+        branches = await Branch.get_list(db=db)
+
+        for branch in branches:
+            if branch.is_default or branch.is_global:
+                continue
+
+            if not branch.branched_from:
+                continue
+
+            original_ts = Timestamp(branch.branched_from)
+            new_ts = original_ts.subtract(microseconds=1)
+            branch.branched_from = new_ts.to_string()
+
+            await branch.save(db=db)
+
+        return result

infrahub/core/migrations/query/__init__.py

@@ -20,7 +20,7 @@ class MigrationQuery(MigrationBaseQuery):
         self,
         migration: SchemaMigration,
         **kwargs: Any,
-    ):
+    ) -> None:
         self.migration = migration
         super().__init__(**kwargs)
 
@@ -32,6 +32,6 @@ class AttributeMigrationQuery(MigrationBaseQuery):
         self,
         migration: AttributeSchemaMigration,
         **kwargs: Any,
-    ):
+    ) -> None:
         self.migration = migration
         super().__init__(**kwargs)

infrahub/core/migrations/query/schema_attribute_update.py

@@ -24,7 +24,7 @@ class SchemaAttributeUpdateQuery(Query):
         new_value: Any,
         previous_value: Any | None = None,
         **kwargs: Any,
-    ):
+    ) -> None:
         self.attr_name = attribute_name
         self.node_name = node_name
         self.node_namespace = node_namespace

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_name_update.py

@@ -14,7 +14,7 @@ class AttributeNameUpdateMigrationQuery01(AttributeMigrationQuery, AttributeRena
         self,
         migration: AttributeSchemaMigration,
         **kwargs: Any,
-    ):
+    ) -> None:
         new_attr = AttributeInfo(
             name=migration.new_attribute_schema.name,
             node_kind=migration.new_schema.kind,

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]:
@@ -34,7 +31,7 @@ class ProfilesAttributeAddMigrationQuery(AttributeMigrationQuery, AttributeAddQu
         self,
         migration: AttributeSchemaMigration,
         **kwargs: Any,
-    ):
+    ) -> None:
         node_kinds = _get_node_kinds(migration.new_schema)
         super().__init__(
             migration=migration,
@@ -54,7 +51,7 @@ class ProfilesAttributeRemoveMigrationQuery(AttributeMigrationQuery, AttributeRe
         self,
         migration: AttributeSchemaMigration,
         **kwargs: Any,
-    ):
+    ) -> None:
         node_kinds = _get_node_kinds(migration.new_schema)
         super().__init__(
             migration=migration,
@@ -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,7 +3,6 @@ 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
@@ -12,16 +11,14 @@ 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
     from infrahub.core.schema import MainSchemaTypes
     from infrahub.core.schema.attribute_schema import AttributeSchema
-    from infrahub.database import InfrahubDatabase
 
     from ...branch import Branch
-    from ...timestamp import Timestamp
 
 
 class NodeAttributeAddMigrationQuery01(AttributeMigrationQuery, AttributeAddQuery):
@@ -41,7 +38,7 @@ class NodeAttributeAddMigrationQuery01(AttributeMigrationQuery, AttributeAddQuer
         self,
         migration: AttributeSchemaMigration,
         **kwargs: Any,
-    ):
+    ) -> None:
         node_kinds = self._get_node_kinds(
             schema=migration.new_schema, new_attribute_schema=migration.new_attribute_schema
         )
@@ -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)
@@ -108,11 +105,11 @@ class NodeAttributeAddMigration(AttributeSchemaMigration):
             return result
 
         for node, number in zip(nodes, numbers, strict=True):
-            await number_pool.reserve(db=db, number=number, identifier=node.get_id())
+            await number_pool.reserve(db=db, number=number, identifier=node.get_id(), at=at)
             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])
+            await node.save(db=db, fields=[self.new_attribute_schema.name], at=at)
 
         return result

infrahub/core/migrations/schema/node_attribute_remove.py

@@ -15,7 +15,7 @@ class NodeAttributeRemoveMigrationQuery01(AttributeMigrationQuery, AttributeRemo
         self,
         migration: AttributeSchemaMigration,
         **kwargs: Any,
-    ):
+    ) -> None:
         super().__init__(
             migration=migration,
             attribute_name=migration.previous_attribute_schema.name,

infrahub/core/migrations/schema/node_kind_update.py

@@ -14,7 +14,7 @@ class NodeKindUpdateMigrationQuery01(MigrationQuery, NodeDuplicateQuery):
         self,
         migration: SchemaMigration,
         **kwargs: Any,
-    ):
+    ) -> None:
         new_node = SchemaNodeInfo(
             name=migration.new_schema.name,
             namespace=migration.new_schema.namespace,

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

@@ -82,7 +82,7 @@ class Node(BaseNode, MetadataInterface, metaclass=BaseNodeMeta):
         _meta.default_filter = default_filter
         super().__init_subclass_with_meta__(_meta=_meta, **options)
 
-    def __init__(self, schema: NodeSchema | ProfileSchema | TemplateSchema, branch: Branch, at: Timestamp):
+    def __init__(self, schema: NodeSchema | ProfileSchema | TemplateSchema, branch: Branch, at: Timestamp) -> None:
         super().__init__()
         self._schema: NodeSchema | ProfileSchema | TemplateSchema = schema
         self._branch: Branch = branch
@@ -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.

infrahub/core/node/base.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from typing import Any
+
 from ..utils import SubclassWithMeta, SubclassWithMeta_Meta
 
 
@@ -9,19 +11,19 @@ class BaseOptions:
 
     _frozen: bool = False
 
-    def __init__(self, class_type):
+    def __init__(self, class_type: type) -> None:
         self.class_type = class_type
 
     def freeze(self) -> None:
         self._frozen = True
 
-    def __setattr__(self, name, value):
+    def __setattr__(self, name: str, value: Any) -> None:
         if not self._frozen:
             super().__setattr__(name, value)
         else:
             raise Exception(f"Can't modify frozen Options {self}")
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f"<{self.__class__.__name__} name={repr(self.name)}>"
 
 
@@ -34,7 +36,9 @@ class BaseNode(SubclassWithMeta):
     #     return type(class_name, (cls,), {"Meta": options})
 
     @classmethod
-    def __init_subclass_with_meta__(cls, name=None, description=None, _meta=None, **_kwargs) -> None:
+    def __init_subclass_with_meta__(
+        cls, name: str | None = None, description: str | None = None, _meta: BaseOptions | None = None, **_kwargs: Any
+    ) -> None:
         assert "_meta" not in cls.__dict__, "Can't assign meta directly"
         if not _meta:
             return
@@ -52,7 +56,7 @@ class BaseNodeOptions(BaseOptions):
 
 
 class ObjectNodeMeta(BaseNodeMeta):
-    def __new__(mcs, name_, bases, namespace, **options):
+    def __new__(mcs, name_: str, bases: tuple[type, ...], namespace: dict[str, Any], **options: Any) -> type:
         # Note: it's safe to pass options as keyword arguments as they are still type-checked by NodeOptions.
 
         # We create this type, to then overload it with the dataclass attrs

infrahub/core/node/delete_validator.py

@@ -117,7 +117,7 @@ class NodeDeleteIndex:
 
 
 class NodeDeleteValidator:
-    def __init__(self, db: InfrahubDatabase, branch: Branch):
+    def __init__(self, db: InfrahubDatabase, branch: Branch) -> None:
         self.db = db
         self.branch = branch
         schema_branch = registry.schema.get_schema_branch(name=self.branch.name)