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

infrahub/core/ipam/resource_allocator.py

@@ -0,0 +1,229 @@
+from __future__ import annotations
+
+import ipaddress
+from typing import TYPE_CHECKING, Iterable
+
+from infrahub.core.query.ipam import (
+    IPAddressData,
+    IPPrefixData,
+    IPPrefixIPAddressFetch,
+    IPPrefixIPAddressFetchFree,
+    IPPrefixSubnetFetch,
+    IPPrefixSubnetFetchFree,
+    IPv6PrefixIPAddressFetchFree,
+    IPv6PrefixSubnetFetchFree,
+)
+
+if TYPE_CHECKING:
+    from infrahub.core.branch import Branch
+    from infrahub.core.ipam.constants import IPAddressType, IPNetworkType
+    from infrahub.core.node import Node
+    from infrahub.core.timestamp import Timestamp
+    from infrahub.database import InfrahubDatabase
+
+
+class IPAMResourceAllocator:
+    """Allocator for IPAM resources (prefixes and addresses) within pools.
+
+    This class provides optimized methods for finding the next available
+    IP prefix or address within a parent prefix. It uses database-side
+    Cypher computation for performance, which is especially important
+    for IPv6 allocations where the address space is very large.
+    """
+
+    def __init__(
+        self,
+        db: InfrahubDatabase,
+        namespace: Node | str | None = None,
+        branch: Branch | None = None,
+        branch_agnostic: bool = False,
+    ) -> None:
+        """Initialize the IPAM resource allocator.
+
+        Args:
+            db: Database connection to use for queries.
+            namespace: IP namespace node or its ID. If None, uses the default namespace.
+            branch: Branch to query. If None, uses the default branch.
+            branch_agnostic: If True, queries across all branches.
+        """
+        self.db = db
+        self.namespace = namespace
+        self.branch = branch
+        self.branch_agnostic = branch_agnostic
+
+    async def _get_next_ipv4_prefix(
+        self, ip_prefix: IPNetworkType, target_prefix_length: int, at: Timestamp | str | None = None
+    ) -> IPNetworkType | None:
+        """Get the next available free IPv4 prefix.
+
+        Uses integer arithmetic in Cypher for efficient gap detection in the
+        IPv4 address space.
+
+        Args:
+            ip_prefix: Parent prefix to allocate from.
+            target_prefix_length: Desired prefix length for the new prefix (0-32).
+            at: Optional timestamp for point-in-time queries.
+
+        Returns:
+            The next available IPv4 prefix, or None if no space is available
+            or if the target_prefix_length is invalid.
+        """
+        if target_prefix_length < 0 or target_prefix_length > 32:
+            return None
+        if target_prefix_length < ip_prefix.prefixlen:
+            return None
+
+        query = await IPPrefixSubnetFetchFree.init(
+            db=self.db,
+            branch=self.branch,
+            obj=ip_prefix,
+            target_prefixlen=target_prefix_length,
+            namespace=self.namespace,
+            at=at,
+            branch_agnostic=self.branch_agnostic,
+        )
+        await query.execute(db=self.db)
+
+        result_data = query.get_prefix_data()
+        if not result_data:
+            return None
+
+        network_address = ipaddress.IPv4Address(result_data.free_start)
+        return ipaddress.ip_network(f"{network_address}/{target_prefix_length}")
+
+    async def _get_next_ipv6_prefix(
+        self, ip_prefix: IPNetworkType, target_prefix_length: int, at: Timestamp | str | None = None
+    ) -> IPNetworkType | None:
+        """Get the next available free IPv6 prefix.
+
+        Uses binary string operations in Cypher to handle IPv6's 128-bit address
+        space, as integer values would overflow Neo4j's 64-bit integer type.
+
+        Args:
+            ip_prefix: Parent prefix to allocate from.
+            target_prefix_length: Desired prefix length for the new prefix (0-128).
+            at: Optional timestamp for point-in-time queries.
+
+        Returns:
+            The next available IPv6 prefix, or None if no space is available
+            or if the target_prefix_length is invalid.
+        """
+        if target_prefix_length < 0 or target_prefix_length > 128:
+            return None
+        if target_prefix_length < ip_prefix.prefixlen:
+            return None
+
+        query = await IPv6PrefixSubnetFetchFree.init(
+            db=self.db,
+            branch=self.branch,
+            obj=ip_prefix,
+            target_prefixlen=target_prefix_length,
+            namespace=self.namespace,
+            at=at,
+            branch_agnostic=self.branch_agnostic,
+        )
+        await query.execute(db=self.db)
+
+        result_data = query.get_prefix_data()
+        if not result_data:
+            return None
+
+        # Convert binary string to IPv6 address
+        addr_int = int(result_data.free_start_bin, 2)
+        network_address = ipaddress.IPv6Address(addr_int)
+        return ipaddress.ip_network(f"{network_address}/{target_prefix_length}")
+
+    async def get_next_prefix(
+        self, ip_prefix: IPNetworkType, target_prefix_length: int, at: Timestamp | str | None = None
+    ) -> IPNetworkType | None:
+        """Get the next available free prefix of specified length within a parent prefix.
+
+        Automatically selects the appropriate method based on IP version (IPv4 vs IPv6).
+
+        Args:
+            ip_prefix: Parent prefix to allocate from.
+            target_prefix_length: Desired prefix length for the new prefix.
+            at: Optional timestamp for point-in-time queries.
+
+        Returns:
+            The next available prefix, or None if no space is available.
+        """
+        if ip_prefix.version == 4:
+            return await self._get_next_ipv4_prefix(
+                ip_prefix=ip_prefix, target_prefix_length=target_prefix_length, at=at
+            )
+        return await self._get_next_ipv6_prefix(ip_prefix=ip_prefix, target_prefix_length=target_prefix_length, at=at)
+
+    async def get_next_address(
+        self, ip_prefix: IPNetworkType, at: Timestamp | str | None = None, is_pool: bool = False
+    ) -> IPAddressType | None:
+        """Get the next available free IP address within a prefix.
+
+        Automatically selects the appropriate query based on IP version (IPv4 vs IPv6).
+
+        Args:
+            ip_prefix: Prefix to allocate an address from.
+            at: Optional timestamp for point-in-time queries.
+            is_pool: If True, includes network and broadcast addresses as allocatable.
+                If False (default), reserves the first and last addresses.
+
+        Returns:
+            The next available IP address, or None if no addresses are available.
+        """
+        # Use IPv6-specific query for IPv6 to avoid 64-bit integer overflow
+        query_class = IPv6PrefixIPAddressFetchFree if ip_prefix.version == 6 else IPPrefixIPAddressFetchFree
+        query = await query_class.init(
+            db=self.db,
+            branch=self.branch,
+            obj=ip_prefix,
+            namespace=self.namespace,
+            at=at,
+            branch_agnostic=self.branch_agnostic,
+            is_pool=is_pool,
+        )
+        await query.execute(db=self.db)
+        return query.get_address()
+
+    async def get_subnets(self, ip_prefix: IPNetworkType, at: Timestamp | str | None = None) -> Iterable[IPPrefixData]:
+        """Get all subnets within a parent prefix.
+
+        Args:
+            ip_prefix: Parent prefix to query.
+            at: Optional timestamp for point-in-time queries.
+
+        Returns:
+            Iterable of IPPrefixData objects representing child subnets.
+        """
+        query = await IPPrefixSubnetFetch.init(
+            db=self.db,
+            branch=self.branch,
+            obj=ip_prefix,
+            namespace=self.namespace,
+            at=at,
+            branch_agnostic=self.branch_agnostic,
+        )
+        await query.execute(db=self.db)
+        return query.get_subnets()
+
+    async def get_ip_addresses(
+        self, ip_prefix: IPNetworkType, at: Timestamp | str | None = None
+    ) -> Iterable[IPAddressData]:
+        """Get all IP addresses within a prefix.
+
+        Args:
+            ip_prefix: Prefix to query.
+            at: Optional timestamp for point-in-time queries.
+
+        Returns:
+            Iterable of IPAddressData objects representing addresses in the prefix.
+        """
+        query = await IPPrefixIPAddressFetch.init(
+            db=self.db,
+            branch=self.branch,
+            obj=ip_prefix,
+            namespace=self.namespace,
+            at=at,
+            branch_agnostic=self.branch_agnostic,
+        )
+        await query.execute(db=self.db)
+        return query.get_addresses()

infrahub/core/migrations/graph/__init__.py

@@ -52,6 +52,11 @@ 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
+from .m052_fix_global_branch_level import Migration052
+from .m053_fix_branch_level_zero import Migration053
+from .m054_cleanup_orphaned_nodes import Migration054
+from .m055_remove_webhook_validate_certificates_default import Migration055
 
 if TYPE_CHECKING:
     from ..shared import MigrationTypes
@@ -108,6 +113,11 @@ MIGRATIONS: list[type[MigrationTypes]] = [
     Migration048,
     Migration049,
     Migration050,
+    Migration051,
+    Migration052,
+    Migration053,
+    Migration054,
+    Migration055,
 ]
 
 

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)
 

infrahub/core/migrations/graph/m032_cleanup_orphaned_branch_relationships.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING, Any
 
-from infrahub.core.migrations.shared import MigrationResult
+from infrahub.core.migrations.shared import MigrationInput, MigrationResult
 from infrahub.core.query import Query, QueryType
 from infrahub.core.query.branch import DeleteBranchRelationshipsQuery
 from infrahub.log import get_logger
@@ -68,7 +68,8 @@ class Migration032(ArbitraryMigration):
     async def validate_migration(self, db: InfrahubDatabase) -> MigrationResult:  # noqa: ARG002
         return MigrationResult()
 
-    async def execute(self, db: InfrahubDatabase) -> MigrationResult:
+    async def execute(self, migration_input: MigrationInput) -> MigrationResult:
+        db = migration_input.db
         migration_result = MigrationResult()
 
         try:

infrahub/core/migrations/graph/m034_find_orphaned_schema_fields.py

@@ -5,7 +5,7 @@ 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.migrations.shared import ArbitraryMigration, MigrationInput, MigrationResult
 from infrahub.core.timestamp import Timestamp
 from infrahub.lock import initialize_lock
 from infrahub.log import get_logger
@@ -59,7 +59,8 @@ class Migration034(ArbitraryMigration):
     async def validate_migration(self, db: InfrahubDatabase) -> MigrationResult:  # noqa: ARG002
         return MigrationResult()
 
-    async def execute(self, db: InfrahubDatabase) -> MigrationResult:
+    async def execute(self, migration_input: MigrationInput) -> MigrationResult:
+        db = migration_input.db
         try:
             initialize_lock()
             await initialization(db=db)

infrahub/core/migrations/graph/m035_orphan_relationships.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING, Any, Sequence
 
-from infrahub.core.migrations.shared import GraphMigration, MigrationResult
+from infrahub.core.migrations.shared import GraphMigration, MigrationInput, MigrationResult
 
 from ...query import Query, QueryType
 
@@ -38,6 +38,6 @@ class Migration035(GraphMigration):
     async def validate_migration(self, db: InfrahubDatabase) -> MigrationResult:  # noqa: ARG002
         return MigrationResult()
 
-    async def execute(self, db: InfrahubDatabase) -> MigrationResult:
+    async def execute(self, migration_input: MigrationInput) -> MigrationResult:
         # overrides parent class to skip transaction in case there are a lot of relationships to delete
-        return await self.do_execute(db=db)
+        return await self.do_execute(migration_input=migration_input)

infrahub/core/migrations/graph/m036_drop_attr_value_index.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 Migration036(GraphMigration):
     queries: Sequence[type[Query]] = []
     minimum_version: int = 35
 
-    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/m037_index_attr_vals.py

@@ -5,7 +5,7 @@ from typing import TYPE_CHECKING, Any
 
 from infrahub.constants.database import IndexType
 from infrahub.core.attribute import MAX_STRING_LENGTH
-from infrahub.core.migrations.shared import MigrationResult, get_migration_console
+from infrahub.core.migrations.shared import MigrationInput, MigrationResult, get_migration_console
 from infrahub.core.query import Query, QueryType
 from infrahub.database.index import IndexItem
 from infrahub.database.neo4j import IndexManagerNeo4j
@@ -463,8 +463,9 @@ class Migration037(ArbitraryMigration):
 
         return result
 
-    async def execute(self, db: InfrahubDatabase) -> MigrationResult:  # noqa: PLR0915
+    async def execute(self, migration_input: MigrationInput) -> MigrationResult:  # noqa: PLR0915
         console = get_migration_console()
+        db = migration_input.db
         result = MigrationResult()
 
         # find the active schema attributes that have a LARGE_ATTRIBUTE_TYPE kind on all branches

infrahub/core/migrations/graph/m038_redo_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
 
@@ -38,12 +36,13 @@ class Migration038(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/m039_ipam_reconcile.py

@@ -10,7 +10,7 @@ from infrahub.core.branch.models import Branch
 from infrahub.core.constants import InfrahubKind
 from infrahub.core.initialization import initialization
 from infrahub.core.ipam.reconciler import IpamReconciler
-from infrahub.core.migrations.shared import MigrationResult, get_migration_console
+from infrahub.core.migrations.shared import MigrationInput, MigrationResult, get_migration_console
 from infrahub.core.query import Query, QueryType
 from infrahub.lock import initialize_lock
 from infrahub.log import get_logger
@@ -233,7 +233,8 @@ class Migration039(ArbitraryMigration):
     async def validate_migration(self, db: InfrahubDatabase) -> MigrationResult:  # noqa: ARG002
         return MigrationResult()
 
-    async def execute(self, db: InfrahubDatabase) -> MigrationResult:
+    async def execute(self, migration_input: MigrationInput) -> MigrationResult:
+        db = migration_input.db
         console = get_migration_console()
         result = MigrationResult()
         # load schemas from database into registry