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

infrahub/api/schema.py

@@ -36,6 +36,7 @@ from infrahub.core.schema import (
     TemplateSchema,
 )
 from infrahub.core.schema.constants import SchemaNamespace  # noqa: TC001
+from infrahub.core.timestamp import Timestamp
 from infrahub.core.validators.models.validate_migration import (
     SchemaValidateMigrationData,
     SchemaValidatorPathResponseData,
@@ -353,6 +354,8 @@ async def load_schema(
         if error_messages:
             raise SchemaNotValidError(",\n".join(error_messages))
 
+        schema_load_at = Timestamp()
+
         # ----------------------------------------------------------
         # Update the schema
         # ----------------------------------------------------------
@@ -367,6 +370,7 @@ async def load_schema(
                 limit=result.diff.all,
                 update_db=True,
                 user_id=account_session.account_id,
+                at=schema_load_at,
             )
             branch.update_schema_hash()
             log.info("Schema has been updated", branch=branch.name, hash=branch.active_schema_hash.main)
@@ -389,6 +393,7 @@ async def load_schema(
             previous_schema=origin_schema,
             migrations=result.migrations,
             user_id=account_session.account_id,
+            at=schema_load_at,
         )
         migration_error_msgs = await service.workflow.execute_workflow(
             workflow=SCHEMA_APPLY_MIGRATION,

infrahub/cli/db.py

@@ -40,10 +40,11 @@ from infrahub.core.migrations.exceptions import MigrationFailureError
 from infrahub.core.migrations.graph import get_graph_migrations, get_migration_by_number
 from infrahub.core.migrations.schema.models import SchemaApplyMigrationData
 from infrahub.core.migrations.schema.tasks import schema_apply_migrations
-from infrahub.core.migrations.shared import get_migration_console
+from infrahub.core.migrations.shared import MigrationInput, get_migration_console
 from infrahub.core.schema import SchemaRoot, core_models, internal_schema
 from infrahub.core.schema.definitions.deprecated import deprecated_models
 from infrahub.core.schema.manager import SchemaManager
+from infrahub.core.timestamp import Timestamp
 from infrahub.core.validators.models.validate_migration import SchemaValidateMigrationData
 from infrahub.core.validators.tasks import schema_validate_migrations
 from infrahub.database import DatabaseType
@@ -357,7 +358,7 @@ async def migrate_database(
     root_node = await get_root_node(db=db)
 
     for migration in migrations:
-        execution_result = await migration.execute(db=db)
+        execution_result = await migration.execute(migration_input=MigrationInput(db=db))
         validation_result = None
 
         if execution_result.success:
@@ -502,6 +503,7 @@ async def update_core_schema(db: InfrahubDatabase, initialize: bool = True, debu
     schema_default_branch.process()
     registry.schema.set_schema_branch(name=default_branch.name, schema=schema_default_branch)
 
+    update_at = Timestamp()
     async with db.start_transaction() as dbt:
         await registry.schema.update_schema_branch(
             schema=candidate_schema,
@@ -510,6 +512,7 @@ async def update_core_schema(db: InfrahubDatabase, initialize: bool = True, debu
             diff=result.diff,
             limit=result.diff.all,
             update_db=True,
+            at=update_at,
         )
         default_branch.update_schema_hash()
         get_migration_console().log(
@@ -527,6 +530,7 @@ async def update_core_schema(db: InfrahubDatabase, initialize: bool = True, debu
         new_schema=candidate_schema,
         previous_schema=origin_schema,
         migrations=result.migrations,
+        at=update_at,
     )
     migration_error_msgs = await schema_apply_migrations(message=apply_migration_data)
 

infrahub/core/branch/models.py

@@ -333,7 +333,7 @@ class Branch(StandardNode):
                     f"({rel}.branch IN $branch{idx} AND {rel}.from <= $time{idx} AND {rel}.to IS NULL)"
                 )
                 filters_per_rel.append(
-                    f"({rel}.branch IN $branch{idx} AND {rel}.from <= $time{idx} AND {rel}.to >= $time{idx})"
+                    f"({rel}.branch IN $branch{idx} AND {rel}.from <= $time{idx} AND {rel}.to > $time{idx})"
                 )
 
             if not include_outside_parentheses:
@@ -366,7 +366,9 @@ class Branch(StandardNode):
         at = Timestamp(at)
         at_str = at.to_string()
         if branch_agnostic:
-            filter_str = f"{variable_name}.from < ${pp}time1 AND ({variable_name}.to IS NULL or {variable_name}.to >= ${pp}time1)"
+            filter_str = (
+                f"{variable_name}.from < ${pp}time1 AND ({variable_name}.to IS NULL or {variable_name}.to > ${pp}time1)"
+            )
             params[f"{pp}time1"] = at_str
             return filter_str, params
 
@@ -380,132 +382,24 @@ class Branch(StandardNode):
         for idx in range(len(branches_times)):
             filters.append(
                 f"({variable_name}.branch IN ${pp}branch{idx} "
-                f"AND {variable_name}.from < ${pp}time{idx} AND {variable_name}.to IS NULL)"
+                f"AND {variable_name}.from <= ${pp}time{idx} AND {variable_name}.to IS NULL)"
             )
             filters.append(
                 f"({variable_name}.branch IN ${pp}branch{idx} "
-                f"AND {variable_name}.from < ${pp}time{idx} "
-                f"AND {variable_name}.to >= ${pp}time{idx})"
+                f"AND {variable_name}.from <= ${pp}time{idx} "
+                f"AND {variable_name}.to > ${pp}time{idx})"
             )
 
         filter_str = "(" + "\n OR ".join(filters) + ")"
 
         return filter_str, params
 
-    def get_query_filter_relationships_range(
-        self,
-        rel_labels: list,
-        start_time: Timestamp,
-        end_time: Timestamp,
-        include_outside_parentheses: bool = False,
-        include_global: bool = False,
-    ) -> tuple[list, dict]:
-        """Generate a CYPHER Query filter based on a list of relationships to query a range of values in the graph.
-        The goal is to return all the values that are valid during this timerange.
-        """
-
-        filters = []
-        params = {}
-
-        if not isinstance(rel_labels, list):
-            raise TypeError(f"rel_labels must be a list, not a {type(rel_labels)}")
-
-        start_time = Timestamp(start_time)
-        end_time = Timestamp(end_time)
-
-        if include_global:
-            branches_times = self.get_branches_and_times_to_query_global(at=start_time)
-        else:
-            branches_times = self.get_branches_and_times_to_query(at=start_time)
-
-        params["branches"] = list({branch for branches in branches_times for branch in branches})
-        params["start_time"] = start_time.to_string()
-        params["end_time"] = end_time.to_string()
-
-        for rel in rel_labels:
-            filters_per_rel = [
-                f"({rel}.branch IN $branches AND {rel}.from <= $end_time AND {rel}.to IS NULL)",
-                f"({rel}.branch IN $branches AND ({rel}.from <= $end_time OR ({rel}.to >= $start_time AND {rel}.to <= $end_time)))",
-            ]
-
-            if not include_outside_parentheses:
-                filters.append("\n OR ".join(filters_per_rel))
-
-            filters.append("(" + "\n OR ".join(filters_per_rel) + ")")
-
-        return filters, params
-
-    def get_query_filter_relationships_diff(
-        self, rel_labels: list, diff_from: Timestamp, diff_to: Timestamp
-    ) -> tuple[list, dict]:
-        """
-        Generate a CYPHER Query filter to query all events that are applicable to a given branch based
-        - The time when the branch as created
-        - The branched_from time of the branch
-        - The diff_to and diff_from time as provided
-        """
-
-        if not isinstance(rel_labels, list):
-            raise TypeError(f"rel_labels must be a list, not a {type(rel_labels)}")
-
-        start_times, end_times = self.get_branches_and_times_for_range(start_time=diff_from, end_time=diff_to)
-
-        filters = []
-        params = {}
-
-        for idx, branch_name in enumerate(start_times.keys()):
-            params[f"branch{idx}"] = branch_name
-            params[f"start_time{idx}"] = start_times[branch_name]
-            params[f"end_time{idx}"] = end_times[branch_name]
-
-        for rel in rel_labels:
-            filters_per_rel = []
-            for idx in range(len(start_times)):
-                filters_per_rel.extend(
-                    [
-                        f"""({rel}.branch = $branch{idx}
-                             AND {rel}.from >= $start_time{idx}
-                             AND {rel}.from <= $end_time{idx}
-                             AND ( r2.to is NULL or r2.to >= $end_time{idx}))""",
-                        f"""({rel}.branch = $branch{idx} AND {rel}.from >= $start_time{idx}
-                            AND {rel}.to <= $start_time{idx})""",
-                    ]
-                )
-
-            filters.append("(" + "\n OR ".join(filters_per_rel) + ")")
-
-        return filters, params
-
-    def get_query_filter_range(self, rel_label: list, start_time: Timestamp, end_time: Timestamp) -> tuple[list, dict]:
-        """
-        Generate a CYPHER Query filter to query a range of values in the graph between start_time and end_time."""
-
-        filters = []
-        params = {}
-
-        start_time = Timestamp(start_time)
-        end_time = Timestamp(end_time)
-
-        params["branches"] = self.get_branches_in_scope()
-        params["start_time"] = start_time.to_string()
-        params["end_time"] = end_time.to_string()
-
-        filters_per_rel = [
-            f"""({rel_label}.branch IN $branches AND {rel_label}.from >= $start_time
-                 AND {rel_label}.from <= $end_time AND {rel_label}.to IS NULL)""",
-            f"""({rel_label}.branch IN $branches AND (({rel_label}.from >= $start_time
-                 AND {rel_label}.from <= $end_time) OR ({rel_label}.to >= $start_time
-                 AND {rel_label}.to <= $end_time)))""",
-        ]
-
-        filters.append("(" + "\n OR ".join(filters_per_rel) + ")")
-
-        return filters, params
-
-    async def rebase(self, db: InfrahubDatabase, user_id: str = SYSTEM_USER_ID) -> None:
+    async def rebase(
+        self, db: InfrahubDatabase, at: Optional[Union[str, Timestamp]] = None, user_id: str = SYSTEM_USER_ID
+    ) -> None:
         """Rebase the current Branch with its origin branch"""
 
-        at = Timestamp()
+        at = Timestamp(at)
 
         await self.rebase_graph(db=db, at=at)
 

infrahub/core/branch/tasks.py

@@ -88,7 +88,7 @@ async def migrate_branch(branch: str, context: InfrahubContext, send_events: boo
 
         try:
             log.info(f"Running migrations for branch '{obj.name}'")
-            await migration_runner.run(db=db)
+            await migration_runner.run(db=db, at=Timestamp())
         except MigrationFailureError as exc:
             log.error(f"Failed to run migrations for branch '{obj.name}': {exc.errors}")
             raise
@@ -170,7 +170,8 @@ async def rebase_branch(branch: str, context: InfrahubContext, send_events: bool
         migrations = []
         async with lock.registry.global_graph_lock():
             async with db.start_transaction() as dbt:
-                await obj.rebase(db=dbt, user_id=context.account.account_id)
+                rebase_at = Timestamp()
+                await obj.rebase(db=dbt, user_id=context.account.account_id, at=rebase_at)
                 log.info("Branch successfully rebased")
 
             if obj.has_schema_changes:
@@ -199,6 +200,7 @@ async def rebase_branch(branch: str, context: InfrahubContext, send_events: bool
                         previous_schema=schema_in_main_before,
                         migrations=migrations,
                         user_id=context.account.account_id,
+                        at=rebase_at,
                     )
                 )
                 for error in errors:
@@ -291,7 +293,8 @@ async def merge_branch(branch: str, context: InfrahubContext, proposed_change_id
                 diff_locker=DiffLocker(),
                 workflow=get_workflow(),
             )
-            branch_diff = await merger.merge()
+            merge_at = Timestamp()
+            branch_diff = await merger.merge(at=merge_at)
             await merger.update_schema()
 
         changelog_collector = DiffChangelogCollector(diff=branch_diff, branch=obj, db=db)
@@ -304,6 +307,7 @@ async def merge_branch(branch: str, context: InfrahubContext, proposed_change_id
                     previous_schema=merger.initial_source_schema,
                     migrations=merger.migrations,
                     user_id=context.account.account_id,
+                    at=merge_at,
                 )
             )
             for error in errors:

infrahub/core/diff/merger/merger.py

@@ -125,7 +125,11 @@ class DiffMerger:
                 )
                 await metadata_query.execute(db=self.db)
 
-        self.source_branch.branched_from = at.to_string()
+        # set the branched_from time to the previous microsecond to prevent duplicated
+        # relationships on the branch after the merge
+        branched_from = at.subtract(microseconds=1)
+
+        self.source_branch.branched_from = branched_from.to_string()
         await self.source_branch.save(db=self.db)
         registry.branch[self.source_branch.name] = self.source_branch
         return enriched_diff

infrahub/core/graph/__init__.py

@@ -1 +1 @@
-GRAPH_VERSION = 50
+GRAPH_VERSION = 51

infrahub/core/initialization.py

@@ -30,6 +30,7 @@ from infrahub.core.protocols import CoreAccount, CoreAccountGroup, CoreAccountRo
 from infrahub.core.root import Root
 from infrahub.core.schema import SchemaRoot, core_models, internal_schema
 from infrahub.core.schema.manager import SchemaManager
+from infrahub.core.timestamp import Timestamp
 from infrahub.database import InfrahubDatabase
 from infrahub.database.memgraph import IndexManagerMemgraph
 from infrahub.database.neo4j import IndexManagerNeo4j
@@ -527,7 +528,7 @@ async def first_time_initialization(db: InfrahubDatabase) -> None:
     schema_branch = registry.schema.register_schema(schema=schema, branch=default_branch.name)
     schema_branch.load_schema(schema=SchemaRoot(**core_models))
     schema_branch.process()
-    await registry.schema.load_schema_to_db(schema=schema_branch, branch=default_branch, db=db)
+    await registry.schema.load_schema_to_db(schema=schema_branch, branch=default_branch, db=db, at=Timestamp())
     registry.schema.set_schema_branch(name=default_branch.name, schema=schema_branch)
     default_branch.update_schema_hash()
     await default_branch.save(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/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)