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

infrahub/core/ipam/utilization.py

@@ -40,25 +40,18 @@ class PrefixUtilizationGetter:
         )
         await query.execute(db=self.db)
 
-        for result in query.get_results():
-            prefix_node = result.get_node("pfx")
-            prefix_id = str(prefix_node.get("uuid"))
-            branch_name = str(result.get("branch"))
-            child_node = result.get_node("child")
-            if InfrahubKind.IPADDRESS in child_node.labels:
+        for item in query.get_data():
+            if InfrahubKind.IPADDRESS in item.child_labels:
                 child_type = PrefixMemberType.ADDRESS
             else:
                 child_type = PrefixMemberType.PREFIX
-            child_value_node = result.get_node("av")
-            child_prefixlen = child_value_node.get("prefixlen")
-            child_ip_value = child_value_node.get("value")
 
-            if prefix_id not in self._results_by_prefix_id:
-                self._results_by_prefix_id[prefix_id] = {}
-            if branch_name not in self._results_by_prefix_id[prefix_id]:
-                self._results_by_prefix_id[prefix_id][branch_name] = []
-            self._results_by_prefix_id[prefix_id][branch_name].append(
-                PrefixChildDetails(child_type=child_type, prefixlen=child_prefixlen, ip_value=child_ip_value)
+            if item.prefix_uuid not in self._results_by_prefix_id:
+                self._results_by_prefix_id[item.prefix_uuid] = {}
+            if item.branch not in self._results_by_prefix_id[item.prefix_uuid]:
+                self._results_by_prefix_id[item.prefix_uuid][item.branch] = []
+            self._results_by_prefix_id[item.prefix_uuid][item.branch].append(
+                PrefixChildDetails(child_type=child_type, prefixlen=item.prefixlen, ip_value=item.ip_value)
             )
 
     async def get_children(

infrahub/core/manager.py

@@ -15,6 +15,7 @@ from infrahub.core.constants import (
 from infrahub.core.metadata.model import MetadataQueryOptions
 from infrahub.core.node import Node
 from infrahub.core.node.delete_validator import NodeDeleteValidator
+from infrahub.core.order import OrderModel
 from infrahub.core.query.node import (
     AttributeFromDB,
     GroupedPeerNodes,
@@ -39,7 +40,6 @@ from infrahub.core.schema import (
 )
 from infrahub.core.timestamp import Timestamp
 from infrahub.exceptions import NodeNotFoundError, ProcessingError, SchemaNotFoundError
-from infrahub.graphql.models import OrderModel
 
 if TYPE_CHECKING:
     from infrahub.core.branch import Branch
@@ -98,7 +98,6 @@ class NodeManager:
         branch: Branch | str | None = ...,
         include_metadata: MetadataQueryOptions | MetadataOptions = ...,
         prefetch_relationships: bool = ...,
-        account=...,
         partial_match: bool = ...,
         branch_agnostic: bool = ...,
         order: OrderModel | None = ...,
@@ -118,7 +117,6 @@ class NodeManager:
         branch: Branch | str | None = ...,
         include_metadata: MetadataQueryOptions | MetadataOptions = ...,
         prefetch_relationships: bool = ...,
-        account=...,
         partial_match: bool = ...,
         branch_agnostic: bool = ...,
         order: OrderModel | None = ...,
@@ -137,7 +135,6 @@ class NodeManager:
         branch: Branch | str | None = None,
         include_metadata: MetadataQueryOptions | MetadataOptions = MetadataOptions.NONE,
         prefetch_relationships: bool = False,
-        account=None,
         partial_match: bool = False,
         branch_agnostic: bool = False,
         order: OrderModel | None = None,
@@ -171,7 +168,6 @@ class NodeManager:
                 branch=branch,
                 include_metadata=include_metadata,
                 prefetch_relationships=prefetch_relationships,
-                account=account,
                 branch_agnostic=branch_agnostic,
             )
             return [node] if node else []
@@ -230,7 +226,6 @@ class NodeManager:
         filters: dict | None = None,
         at: Timestamp | str | None = None,
         branch: Branch | str | None = None,
-        account=None,  # noqa: ARG003
         partial_match: bool = False,
         branch_agnostic: bool = False,
     ) -> int:
@@ -534,7 +529,6 @@ class NodeManager:
         branch: Branch | str | None = ...,
         include_metadata: MetadataQueryOptions | MetadataOptions = ...,
         prefetch_relationships: bool = ...,
-        account=...,
         branch_agnostic: bool = ...,
     ) -> SchemaProtocol | None: ...
 
@@ -551,7 +545,6 @@ class NodeManager:
         branch: Branch | str | None = ...,
         include_metadata: MetadataQueryOptions | MetadataOptions = ...,
         prefetch_relationships: bool = ...,
-        account=...,
         branch_agnostic: bool = ...,
     ) -> SchemaProtocol: ...
 
@@ -568,7 +561,6 @@ class NodeManager:
         branch: Branch | str | None = ...,
         include_metadata: MetadataQueryOptions | MetadataOptions = ...,
         prefetch_relationships: bool = ...,
-        account=...,
         branch_agnostic: bool = ...,
     ) -> SchemaProtocol | None: ...
 
@@ -585,7 +577,6 @@ class NodeManager:
         branch: Branch | str | None = ...,
         include_metadata: MetadataQueryOptions | MetadataOptions = ...,
         prefetch_relationships: bool = ...,
-        account=...,
         branch_agnostic: bool = ...,
     ) -> Node | None: ...
 
@@ -602,7 +593,6 @@ class NodeManager:
         branch: Branch | str | None = ...,
         include_metadata: MetadataQueryOptions | MetadataOptions = ...,
         prefetch_relationships: bool = ...,
-        account=...,
         branch_agnostic: bool = ...,
     ) -> Node: ...
 
@@ -619,7 +609,6 @@ class NodeManager:
         branch: Branch | str | None = ...,
         include_metadata: MetadataQueryOptions | MetadataOptions = ...,
         prefetch_relationships: bool = ...,
-        account=...,
         branch_agnostic: bool = ...,
     ) -> Node | None: ...
 
@@ -635,7 +624,6 @@ class NodeManager:
         branch: Branch | str | None = None,
         include_metadata: MetadataQueryOptions | MetadataOptions = MetadataOptions.NONE,
         prefetch_relationships: bool = False,
-        account=None,
         branch_agnostic: bool = False,
     ) -> Node | SchemaProtocol | None:
         branch = await registry.get_branch(branch=branch, db=db)
@@ -657,7 +645,6 @@ class NodeManager:
             at=at,
             include_metadata=include_metadata,
             prefetch_relationships=prefetch_relationships,
-            account=account,
             branch_agnostic=branch_agnostic,
             order=OrderModel(disable=True),
         )
@@ -694,7 +681,6 @@ class NodeManager:
         branch: Branch | str | None = ...,
         include_metadata: MetadataQueryOptions | MetadataOptions = ...,
         prefetch_relationships: bool = ...,
-        account=...,
     ) -> SchemaProtocol | None: ...
 
     @overload
@@ -710,7 +696,6 @@ class NodeManager:
         branch: Branch | str | None = ...,
         include_metadata: MetadataQueryOptions | MetadataOptions = ...,
         prefetch_relationships: bool = ...,
-        account=...,
         branch_agnostic: bool = ...,
     ) -> SchemaProtocol: ...
 
@@ -727,7 +712,6 @@ class NodeManager:
         branch: Branch | str | None = ...,
         include_metadata: MetadataQueryOptions | MetadataOptions = ...,
         prefetch_relationships: bool = ...,
-        account=...,
         branch_agnostic: bool = ...,
     ) -> SchemaProtocol | None: ...
 
@@ -744,7 +728,6 @@ class NodeManager:
         branch: Branch | str | None = ...,
         include_metadata: MetadataQueryOptions | MetadataOptions = ...,
         prefetch_relationships: bool = ...,
-        account=...,
         branch_agnostic: bool = ...,
     ) -> Node: ...
 
@@ -761,7 +744,6 @@ class NodeManager:
         branch: Branch | str | None = ...,
         include_metadata: MetadataQueryOptions | MetadataOptions = ...,
         prefetch_relationships: bool = ...,
-        account=...,
         branch_agnostic: bool = ...,
     ) -> Node | None: ...
 
@@ -778,7 +760,6 @@ class NodeManager:
         branch: Branch | str | None = ...,
         include_metadata: MetadataQueryOptions | MetadataOptions = ...,
         prefetch_relationships: bool = ...,
-        account=...,
         branch_agnostic: bool = ...,
     ) -> Node | None: ...
 
@@ -794,7 +775,6 @@ class NodeManager:
         branch: Branch | str | None = None,
         include_metadata: MetadataQueryOptions | MetadataOptions = MetadataOptions.NONE,
         prefetch_relationships: bool = False,
-        account=None,
         branch_agnostic: bool = False,
     ) -> Node | SchemaProtocol | None:
         branch = await registry.get_branch(branch=branch, db=db)
@@ -844,7 +824,6 @@ class NodeManager:
             at=at,
             include_metadata=include_metadata,
             prefetch_relationships=prefetch_relationships,
-            account=account,
             branch_agnostic=branch_agnostic,
             order=OrderModel(disable=True),
         )
@@ -876,7 +855,6 @@ class NodeManager:
         branch: Branch | str | None = ...,
         include_metadata: MetadataQueryOptions | MetadataOptions = ...,
         prefetch_relationships: bool = ...,
-        account=...,
         branch_agnostic: bool = ...,
     ) -> SchemaProtocol: ...
 
@@ -892,7 +870,6 @@ class NodeManager:
         branch: Branch | str | None = ...,
         include_metadata: MetadataQueryOptions | MetadataOptions = ...,
         prefetch_relationships: bool = ...,
-        account=...,
         branch_agnostic: bool = ...,
     ) -> Any: ...
 
@@ -907,7 +884,6 @@ class NodeManager:
         branch: Branch | str | None = None,
         include_metadata: MetadataQueryOptions | MetadataOptions = MetadataOptions.NONE,
         prefetch_relationships: bool = False,
-        account=None,
         branch_agnostic: bool = False,
     ) -> Any:
         branch = await registry.get_branch(branch=branch, db=db)
@@ -935,7 +911,6 @@ class NodeManager:
             branch=branch,
             include_metadata=include_metadata,
             prefetch_relationships=prefetch_relationships,
-            account=account,
             branch_agnostic=branch_agnostic,
         )
         if not node:

infrahub/core/merge.py

@@ -43,7 +43,7 @@ class BranchMerger:
         diff_locker: DiffLocker,
         destination_branch: Branch | None = None,
         workflow: InfrahubWorkflow | None = None,
-    ):
+    ) -> None:
         self.source_branch = source_branch
         self.destination_branch: Branch = destination_branch or registry.get_branch_from_registry()
         self.db = db

infrahub/core/migrations/graph/__init__.py

@@ -52,6 +52,7 @@ from .m047_backfill_or_null_display_label import Migration047
 from .m048_undelete_rel_props import Migration048
 from .m049_remove_is_visible_relationship import Migration049
 from .m050_backfill_vertex_metadata import Migration050
+from .m051_subtract_branched_from_microsecond import Migration051
 
 if TYPE_CHECKING:
     from ..shared import MigrationTypes
@@ -108,6 +109,7 @@ MIGRATIONS: list[type[MigrationTypes]] = [
     Migration048,
     Migration049,
     Migration050,
+    Migration051,
 ]
 
 

infrahub/core/migrations/graph/m012_convert_account_generic.py

@@ -43,7 +43,7 @@ class Migration012RenameTypeAttributeData(AttributeRenameQuery):
     name = "migration_012_rename_attr_type"
     type = QueryType.WRITE
 
-    def __init__(self, **kwargs: Any):
+    def __init__(self, **kwargs: Any) -> None:
         new_attr = AttributeInfo(
             name="account_type",
             node_kind=InfrahubKind.ACCOUNT,
@@ -83,7 +83,7 @@ class Migration012AddLabelData(NodeDuplicateQuery):
     name = "migration_012_add_labels"
     type = QueryType.WRITE
 
-    def __init__(self, **kwargs: Any):
+    def __init__(self, **kwargs: Any) -> None:
         new_node = SchemaNodeInfo(
             name="Account",
             namespace="Core",
@@ -132,7 +132,7 @@ class Migration012RenameTypeAttributeSchema(SchemaAttributeUpdateQuery):
     type = QueryType.WRITE
     insert_return = False
 
-    def __init__(self, **kwargs: Any):
+    def __init__(self, **kwargs: Any) -> None:
         super().__init__(
             attribute_name="name",
             node_name="Account",
@@ -150,7 +150,7 @@ class Migration012RenameRelationshipAccountTokenData(RelationshipDuplicateQuery)
     name = "migration_012_rename_rel_account_token_data"
     type = QueryType.WRITE
 
-    def __init__(self, **kwargs: Any):
+    def __init__(self, **kwargs: Any) -> None:
         new_rel = SchemaRelationshipInfo(
             name="account__token",
             branch_support=BranchSupportType.AGNOSTIC.value,
@@ -173,7 +173,7 @@ class Migration012RenameRelationshipRefreshTokenData(RelationshipDuplicateQuery)
     name = "migration_012_rename_rel_refresh_token_data"
     type = QueryType.WRITE
 
-    def __init__(self, **kwargs: Any):
+    def __init__(self, **kwargs: Any) -> None:
         new_rel = SchemaRelationshipInfo(
             name="account__refreshtoken",
             branch_support=BranchSupportType.AGNOSTIC.value,
@@ -196,7 +196,7 @@ class Migration012RenameRelationshipThreadData(RelationshipDuplicateQuery):
     name = "migration_012_rename_rel_thread_data"
     type = QueryType.WRITE
 
-    def __init__(self, **kwargs: Any):
+    def __init__(self, **kwargs: Any) -> None:
         new_rel = SchemaRelationshipInfo(
             name="thread__account",
             branch_support=BranchSupportType.AGNOSTIC.value,
@@ -219,7 +219,7 @@ class Migration012RenameRelationshipCommentData(RelationshipDuplicateQuery):
     name = "migration_012_rename_rel_comment_data"
     type = QueryType.WRITE
 
-    def __init__(self, **kwargs: Any):
+    def __init__(self, **kwargs: Any) -> None:
         new_rel = SchemaRelationshipInfo(
             name="comment__account",
             branch_support=BranchSupportType.AGNOSTIC.value,
@@ -243,7 +243,7 @@ class Migration012DeleteOldElementsSchema(DeleteElementInSchemaQuery):
     type = QueryType.WRITE
     insert_return = False
 
-    def __init__(self, **kwargs: Any):
+    def __init__(self, **kwargs: Any) -> None:
         kwargs.pop("branch", None)
 
         super().__init__(
@@ -260,7 +260,7 @@ class Migration012UpdateDisplayLabels(SchemaAttributeUpdateQuery):
     type = QueryType.WRITE
     insert_return = False
 
-    def __init__(self, **kwargs: Any):
+    def __init__(self, **kwargs: Any) -> None:
         kwargs.pop("branch", None)
 
         super().__init__(
@@ -277,7 +277,7 @@ class Migration012UpdateOrderBy(SchemaAttributeUpdateQuery):
     type = QueryType.WRITE
     insert_return = False
 
-    def __init__(self, **kwargs: Any):
+    def __init__(self, **kwargs: Any) -> None:
         kwargs.pop("branch", None)
 
         super().__init__(
@@ -294,7 +294,7 @@ class Migration012UpdateDefaultFilter(SchemaAttributeUpdateQuery):
     type = QueryType.WRITE
     insert_return = False
 
-    def __init__(self, **kwargs: Any):
+    def __init__(self, **kwargs: Any) -> None:
         kwargs.pop("branch", None)
 
         super().__init__(
@@ -311,7 +311,7 @@ class Migration012UpdateHFID(SchemaAttributeUpdateQuery):
     type = QueryType.WRITE
     insert_return = False
 
-    def __init__(self, **kwargs: Any):
+    def __init__(self, **kwargs: Any) -> None:
         kwargs.pop("branch", None)
 
         super().__init__(

infrahub/core/migrations/graph/m013_convert_git_password_credential.py

@@ -228,7 +228,7 @@ class Migration013DeleteUsernamePasswordGenericSchema(DeleteElementInSchemaQuery
     type = QueryType.WRITE
     insert_return = False
 
-    def __init__(self, **kwargs: Any):
+    def __init__(self, **kwargs: Any) -> None:
         kwargs.pop("branch", None)
 
         super().__init__(
@@ -245,7 +245,7 @@ class Migration013DeleteUsernamePasswordReadWriteSchema(DeleteElementInSchemaQue
     type = QueryType.WRITE
     insert_return = False
 
-    def __init__(self, **kwargs: Any):
+    def __init__(self, **kwargs: Any) -> None:
         kwargs.pop("branch", None)
 
         super().__init__(
@@ -262,7 +262,7 @@ class Migration013DeleteUsernamePasswordReadOnlySchema(DeleteElementInSchemaQuer
     type = QueryType.WRITE
     insert_return = False
 
-    def __init__(self, **kwargs: Any):
+    def __init__(self, **kwargs: Any) -> None:
         kwargs.pop("branch", None)
 
         super().__init__(
@@ -277,7 +277,7 @@ class Migration013DeleteUsernamePasswordReadOnlySchema(DeleteElementInSchemaQuer
 class Migration013AddInternalStatusData(AttributeAddQuery):
     type = QueryType.WRITE
 
-    def __init__(self, **kwargs: Any):
+    def __init__(self, **kwargs: Any) -> None:
         kwargs.pop("branch", None)
 
         super().__init__(

infrahub/core/migrations/graph/m014_remove_index_attr_value.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 from typing import TYPE_CHECKING, Sequence
 
 from infrahub.constants.database import IndexType
-from infrahub.core.migrations.shared import MigrationResult
+from infrahub.core.migrations.shared import MigrationInput, MigrationResult
 from infrahub.core.query import Query  # noqa: TC001
 from infrahub.database import DatabaseType
 from infrahub.database.index import IndexItem
@@ -23,7 +23,8 @@ class Migration014(GraphMigration):
     queries: Sequence[type[Query]] = []
     minimum_version: int = 13
 
-    async def execute(self, db: InfrahubDatabase) -> MigrationResult:
+    async def execute(self, migration_input: MigrationInput) -> MigrationResult:
+        db = migration_input.db
         result = MigrationResult()
 
         # Only execute this migration for Neo4j

infrahub/core/migrations/graph/m015_diff_format_update.py

@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING
 
 from infrahub.core import registry
 from infrahub.core.diff.repository.repository import DiffRepository
-from infrahub.core.migrations.shared import MigrationResult
+from infrahub.core.migrations.shared import MigrationInput, MigrationResult
 from infrahub.dependencies.registry import build_component_registry, get_component_registry
 from infrahub.log import get_logger
 
@@ -25,7 +25,8 @@ class Migration015(ArbitraryMigration):
 
         return result
 
-    async def execute(self, db: InfrahubDatabase) -> MigrationResult:
+    async def execute(self, migration_input: MigrationInput) -> MigrationResult:
+        db = migration_input.db
         default_branch = registry.get_branch_from_registry()
         build_component_registry()
         component_registry = get_component_registry()

infrahub/core/migrations/graph/m016_diff_delete_bug_fix.py

@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING
 
 from infrahub.core import registry
 from infrahub.core.diff.repository.repository import DiffRepository
-from infrahub.core.migrations.shared import MigrationResult
+from infrahub.core.migrations.shared import MigrationInput, MigrationResult
 from infrahub.dependencies.registry import build_component_registry, get_component_registry
 from infrahub.log import get_logger
 
@@ -25,7 +25,8 @@ class Migration016(ArbitraryMigration):
 
         return result
 
-    async def execute(self, db: InfrahubDatabase) -> MigrationResult:
+    async def execute(self, migration_input: MigrationInput) -> MigrationResult:
+        db = migration_input.db
         default_branch = registry.get_branch_from_registry()
         build_component_registry()
         component_registry = get_component_registry()

infrahub/core/migrations/graph/m017_add_core_profile.py

@@ -3,8 +3,7 @@ from __future__ import annotations
 from typing import TYPE_CHECKING, Sequence
 
 from infrahub.core import registry
-from infrahub.core.constants import SYSTEM_USER_ID
-from infrahub.core.migrations.shared import MigrationResult
+from infrahub.core.migrations.shared import MigrationInput, MigrationResult
 from infrahub.core.schema.definitions.core import core_profile_schema_definition
 from infrahub.core.schema.manager import SchemaManager
 from infrahub.log import get_logger
@@ -27,17 +26,20 @@ class Migration017(InternalSchemaMigration):
 
         return result
 
-    async def execute(self, db: InfrahubDatabase, user_id: str = SYSTEM_USER_ID) -> MigrationResult:
+    async def execute(self, migration_input: MigrationInput) -> MigrationResult:
         """
         Load CoreProfile schema node in db.
         """
+        db = migration_input.db
+        at = migration_input.at
+        user_id = migration_input.user_id
         default_branch = registry.get_branch_from_registry()
         manager = SchemaManager()
         manager.set_schema_branch(name=default_branch.name, schema=self.get_internal_schema())
 
         db.add_schema(manager.get_schema_branch(default_branch.name))
         await manager.load_node_to_db(
-            node=core_profile_schema_definition, db=db, branch=default_branch, user_id=user_id
+            node=core_profile_schema_definition, db=db, branch=default_branch, user_id=user_id, at=at
         )
 
         return MigrationResult()

infrahub/core/migrations/graph/m018_uniqueness_nulls.py

@@ -4,9 +4,8 @@ from collections import defaultdict
 from typing import TYPE_CHECKING, Sequence
 
 from infrahub.core import registry
-from infrahub.core.constants import SYSTEM_USER_ID
 from infrahub.core.diff.payload_builder import get_display_labels_per_kind
-from infrahub.core.migrations.shared import MigrationResult
+from infrahub.core.migrations.shared import MigrationInput, MigrationResult
 from infrahub.core.schema import GenericSchema, NodeSchema, SchemaRoot, internal_schema
 from infrahub.core.schema.manager import SchemaManager
 from infrahub.core.validators.uniqueness.checker import UniquenessChecker
@@ -96,5 +95,5 @@ class Migration018(InternalSchemaMigration):
     async def validate_migration(self, db: InfrahubDatabase) -> MigrationResult:  # noqa: ARG002
         return MigrationResult()
 
-    async def execute(self, db: InfrahubDatabase, user_id: str = SYSTEM_USER_ID) -> MigrationResult:  # noqa: ARG002
-        return await validate_nulls_in_uniqueness_constraints(db=db)
+    async def execute(self, migration_input: MigrationInput) -> MigrationResult:
+        return await validate_nulls_in_uniqueness_constraints(db=migration_input.db)

infrahub/core/migrations/graph/m020_duplicate_edges.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 from typing import TYPE_CHECKING, Any, Sequence
 
 from infrahub.core.constants.database import DatabaseEdgeType
-from infrahub.core.migrations.shared import GraphMigration, MigrationResult
+from infrahub.core.migrations.shared import GraphMigration, MigrationInput, MigrationResult
 from infrahub.log import get_logger
 
 from ...query import Query, QueryType
@@ -140,9 +140,9 @@ class Migration020(GraphMigration):
         DeleteDuplicateIsProtectedEdgesQuery,
     ]
 
-    async def execute(self, db: InfrahubDatabase) -> MigrationResult:
+    async def execute(self, migration_input: MigrationInput) -> MigrationResult:
         # skip the transaction b/c it will run out of memory on a large database
-        return await self.do_execute(db=db)
+        return await self.do_execute(migration_input=migration_input)
 
     async def validate_migration(self, db: InfrahubDatabase) -> MigrationResult:  # noqa: ARG002
         result = MigrationResult()

infrahub/core/migrations/graph/m025_uniqueness_nulls.py

@@ -2,8 +2,7 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING, Sequence
 
-from infrahub.core.constants import SYSTEM_USER_ID
-from infrahub.core.migrations.shared import MigrationResult
+from infrahub.core.migrations.shared import MigrationInput, MigrationResult
 from infrahub.log import get_logger
 
 from ..shared import InternalSchemaMigration, SchemaMigration
@@ -23,5 +22,5 @@ class Migration025(InternalSchemaMigration):
     async def validate_migration(self, db: InfrahubDatabase) -> MigrationResult:  # noqa: ARG002
         return MigrationResult()
 
-    async def execute(self, db: InfrahubDatabase, user_id: str = SYSTEM_USER_ID) -> MigrationResult:  # noqa: ARG002
-        return await validate_nulls_in_uniqueness_constraints(db=db)
+    async def execute(self, migration_input: MigrationInput) -> MigrationResult:
+        return await validate_nulls_in_uniqueness_constraints(db=migration_input.db)

infrahub/core/migrations/graph/m026_0000_prefix_fix.py

@@ -4,12 +4,10 @@ import ipaddress
 from typing import TYPE_CHECKING, Sequence
 
 from infrahub.core.branch.models import Branch
-from infrahub.core.constants import SYSTEM_USER_ID
 from infrahub.core.initialization import initialization
 from infrahub.core.ipam.reconciler import IpamReconciler
 from infrahub.core.manager import NodeManager
-from infrahub.core.migrations.shared import MigrationResult
-from infrahub.core.timestamp import Timestamp
+from infrahub.core.migrations.shared import MigrationInput, MigrationResult
 from infrahub.lock import initialize_lock
 from infrahub.log import get_logger
 
@@ -29,12 +27,13 @@ class Migration026(InternalSchemaMigration):
     async def validate_migration(self, db: InfrahubDatabase) -> MigrationResult:  # noqa: ARG002
         return MigrationResult()
 
-    async def execute(self, db: InfrahubDatabase, user_id: str = SYSTEM_USER_ID) -> MigrationResult:  # noqa: ARG002
+    async def execute(self, migration_input: MigrationInput) -> MigrationResult:
+        db = migration_input.db
+        at = migration_input.at
         # load schemas from database into registry
         initialize_lock()
         await initialization(db=db)
 
-        at = Timestamp()
         for branch in await Branch.get_list(db=db):
             prefix_0000s = await NodeManager.query(
                 db=db, schema="BuiltinIPPrefix", branch=branch, filters={"prefix__values": ["0.0.0.0/0", "::/0"]}

infrahub/core/migrations/graph/m028_delete_diffs.py

@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING
 
 from infrahub.core import registry
 from infrahub.core.diff.repository.repository import DiffRepository
-from infrahub.core.migrations.shared import MigrationResult
+from infrahub.core.migrations.shared import MigrationInput, MigrationResult
 from infrahub.dependencies.registry import build_component_registry, get_component_registry
 from infrahub.log import get_logger
 
@@ -27,7 +27,8 @@ class Migration028(ArbitraryMigration):
 
         return result
 
-    async def execute(self, db: InfrahubDatabase) -> MigrationResult:
+    async def execute(self, migration_input: MigrationInput) -> MigrationResult:
+        db = migration_input.db
         default_branch = registry.get_branch_from_registry()
         build_component_registry()
         component_registry = get_component_registry()

infrahub/core/migrations/graph/m029_duplicates_cleanup.py

@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING, Any, Literal
 
 from infrahub.core.constants import RelationshipDirection
 from infrahub.core.constants.database import DatabaseEdgeType
-from infrahub.core.migrations.shared import MigrationResult
+from infrahub.core.migrations.shared import MigrationInput, MigrationResult
 from infrahub.core.query import Query, QueryType
 from infrahub.log import get_logger
 
@@ -574,8 +574,9 @@ class Migration029(ArbitraryMigration):
 
         return result
 
-    async def execute(self, db: InfrahubDatabase) -> MigrationResult:
+    async def execute(self, migration_input: MigrationInput) -> MigrationResult:
         migration_result = MigrationResult()
+        db = migration_input.db
         limit = self.limit
         offset = 0
         more_nodes_to_process = True

infrahub/core/migrations/graph/m031_check_number_attributes.py

@@ -5,9 +5,9 @@ from typing import TYPE_CHECKING, Sequence
 from infrahub import config
 from infrahub.core import registry
 from infrahub.core.branch import Branch
-from infrahub.core.constants import SYSTEM_USER_ID, SchemaPathType
+from infrahub.core.constants import SchemaPathType
 from infrahub.core.initialization import initialization
-from infrahub.core.migrations.shared import InternalSchemaMigration, MigrationResult, SchemaMigration
+from infrahub.core.migrations.shared import InternalSchemaMigration, MigrationInput, MigrationResult, SchemaMigration
 from infrahub.core.path import SchemaPath
 from infrahub.core.schema import GenericSchema, NodeSchema
 from infrahub.core.schema.attribute_parameters import NumberAttributeParameters
@@ -35,7 +35,7 @@ class Migration031(InternalSchemaMigration):
     minimum_version: int = 30
     migrations: Sequence[SchemaMigration] = []
 
-    async def execute(self, db: InfrahubDatabase, user_id: str = SYSTEM_USER_ID) -> MigrationResult:  # noqa: ARG002
+    async def execute(self, migration_input: MigrationInput) -> MigrationResult:
         """Retrieve all number attributes that have a min/max/excluded_values
         For any of these attributes, check if corresponding existing nodes are valid."""
 
@@ -43,6 +43,7 @@ class Migration031(InternalSchemaMigration):
             return MigrationResult()
 
         # load schemas from database into registry
+        db = migration_input.db
         initialize_lock()
         await initialization(db=db)