infrahub-server 1.5.0b0__py3-none-any.whl → 1.5.0b2__py3-none-any.whl

infrahub/core/migrations/query/__init__.py

@@ -8,7 +8,12 @@ if TYPE_CHECKING:
     from ..shared import AttributeSchemaMigration, SchemaMigration
 
 
-class MigrationQuery(Query):
+class MigrationBaseQuery(Query):
+    def get_nbr_migrations_executed(self) -> int:
+        return self.num_of_results
+
+
+class MigrationQuery(MigrationBaseQuery):
     type: QueryType = QueryType.WRITE
 
     def __init__(
@@ -19,11 +24,8 @@ class MigrationQuery(Query):
         self.migration = migration
         super().__init__(**kwargs)
 
-    def get_nbr_migrations_executed(self) -> int:
-        return self.num_of_results
-
 
-class AttributeMigrationQuery(Query):
+class AttributeMigrationQuery(MigrationBaseQuery):
     type: QueryType = QueryType.WRITE
 
     def __init__(
@@ -33,6 +35,3 @@ class AttributeMigrationQuery(Query):
     ):
         self.migration = migration
         super().__init__(**kwargs)
-
-    def get_nbr_migrations_executed(self) -> int:
-        return self.num_of_results

infrahub/core/migrations/query/attribute_add.py

@@ -17,26 +17,27 @@ class AttributeAddQuery(Query):
 
     def __init__(
         self,
-        node_kind: str,
+        node_kinds: list[str],
         attribute_name: str,
         attribute_kind: str,
         branch_support: str,
         default_value: Any | None = None,
+        uuids: list[str] | None = None,
         **kwargs: Any,
     ) -> None:
-        self.node_kind = node_kind
+        self.node_kinds = node_kinds
         self.attribute_name = attribute_name
         self.attribute_kind = attribute_kind
         self.branch_support = branch_support
         self.default_value = default_value
-
+        self.uuids = uuids
         super().__init__(**kwargs)
 
     async def query_init(self, db: InfrahubDatabase, **kwargs: dict[str, Any]) -> None:  # noqa: ARG002
         branch_filter, branch_params = self.branch.get_query_filter_path(at=self.at.to_string())
         self.params.update(branch_params)
 
-        self.params["node_kind"] = self.node_kind
+        self.params["node_kinds"] = self.node_kinds
         self.params["attr_name"] = self.attribute_name
         self.params["branch_support"] = self.branch_support
         self.params["current_time"] = self.at.to_string()
@@ -79,12 +80,13 @@ class AttributeAddQuery(Query):
             LIMIT 1
             """ % {"attr_value_label": attr_value_label}
 
+        node_kinds_str = "|".join(self.node_kinds)
         query = """
         %(match_query)s
         MERGE (is_protected_value:Boolean { value: $is_protected_default })
         MERGE (is_visible_value:Boolean { value: $is_visible_default })
         WITH av, is_protected_value, is_visible_value
-        MATCH p = (n:%(node_kind)s)
+        MATCH (n:%(node_kinds_str)s)
         CALL (n) {
             MATCH (:Root)<-[r:IS_PART_OF]-(n)
             WHERE %(branch_filter)s
@@ -110,7 +112,7 @@ class AttributeAddQuery(Query):
         """ % {
             "match_query": match_query,
             "branch_filter": branch_filter,
-            "node_kind": self.node_kind,
+            "node_kinds_str": node_kinds_str,
             "uuid_generation": db.render_uuid_generation(node_label="a", node_attr="uuid"),
         }
 

infrahub/core/migrations/query/attribute_remove.py

@@ -0,0 +1,134 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from infrahub.core.constants import RelationshipStatus
+from infrahub.core.graph.schema import GraphAttributeRelationships
+from infrahub.core.query import Query
+from infrahub.core.schema.generic_schema import GenericSchema
+
+if TYPE_CHECKING:
+    from pydantic.fields import FieldInfo
+
+    from infrahub.database import InfrahubDatabase
+
+
+class AttributeRemoveQuery(Query):
+    name = "attribute_remove"
+    insert_return: bool = False
+
+    def __init__(
+        self,
+        attribute_name: str,
+        node_kinds: list[str],
+        **kwargs: Any,
+    ) -> None:
+        self.attribute_name = attribute_name
+        self.node_kinds = node_kinds
+        super().__init__(**kwargs)
+
+    async def query_init(self, db: InfrahubDatabase, **kwargs: dict[str, Any]) -> None:  # noqa: ARG002
+        branch_filter, branch_params = self.branch.get_query_filter_path(at=self.at.to_string())
+        self.params.update(branch_params)
+
+        kinds_to_ignore = []
+        profile_kinds_to_update = []
+
+        for node_kind in self.node_kinds:
+            new_schema = db.schema.get(name=node_kind, branch=self.branch, duplicate=False)
+
+        if isinstance(new_schema, GenericSchema):
+            for inheriting_schema_kind in new_schema.used_by:
+                node_schema = db.schema.get_node_schema(
+                    name=inheriting_schema_kind, branch=self.branch, duplicate=False
+                )
+                attr_schema = node_schema.get_attribute_or_none(name=self.attribute_name)
+                if attr_schema and not attr_schema.inherited:
+                    kinds_to_ignore.append(inheriting_schema_kind)
+                else:
+                    profile_kinds_to_update.append(f"Profile{inheriting_schema_kind}")
+
+        self.params["kinds_to_ignore"] = kinds_to_ignore
+        self.params["attr_name"] = self.attribute_name
+        self.params["current_time"] = self.at.to_string()
+        self.params["branch_name"] = self.branch.name
+
+        self.params["rel_props"] = {
+            "branch": self.branch.name,
+            "branch_level": self.branch.hierarchy_level,
+            "status": RelationshipStatus.DELETED.value,
+            "from": self.at.to_string(),
+        }
+
+        def render_sub_query_per_rel_type(rel_type: str, rel_def: FieldInfo) -> str:
+            subquery = [
+                "WITH peer_node, rb, active_attr",
+                f'WHERE type(rb) = "{rel_type}"',
+            ]
+            if rel_def.default.direction.value == "outbound":
+                subquery.append(f"CREATE (active_attr)-[:{rel_type} $rel_props ]->(peer_node)")
+            elif rel_def.default.direction.value == "inbound":
+                subquery.append(f"CREATE (active_attr)<-[:{rel_type} $rel_props ]-(peer_node)")
+            else:
+                subquery.append(f"CREATE (active_attr)-[:{rel_type} $rel_props ]-(peer_node)")
+
+            subquery.append("RETURN peer_node as p2")
+            return "\n".join(subquery)
+
+        sub_queries = [
+            render_sub_query_per_rel_type(rel_type, rel_def)
+            for rel_type, rel_def in GraphAttributeRelationships.model_fields.items()
+        ]
+        sub_query_all = "\nUNION\n".join(sub_queries)
+
+        node_kinds_str = "|".join(self.node_kinds + profile_kinds_to_update)
+        query = """
+        // Find all the active nodes
+        MATCH (node:%(node_kinds)s)
+        WHERE (size($kinds_to_ignore) = 0 OR NOT any(l IN labels(node) WHERE l IN $kinds_to_ignore))
+        AND exists((node)-[:HAS_ATTRIBUTE]-(:Attribute { name: $attr_name }))
+        CALL (node) {
+            MATCH (root:Root)<-[r:IS_PART_OF]-(node)
+            WHERE %(branch_filter)s
+            RETURN node as n1, r as r1
+            ORDER BY r.branch_level DESC, r.from DESC
+            LIMIT 1
+        }
+        WITH n1 as active_node, r1 as rb
+        WHERE rb.status = "active"
+        // Find all the attributes that need to be updated
+        CALL (active_node) {
+            MATCH (active_node)-[r:HAS_ATTRIBUTE]-(attr:Attribute { name: $attr_name })
+            WHERE %(branch_filter)s
+            RETURN active_node as n1, r as r1, attr as attr1
+            ORDER BY r.branch_level DESC, r.from DESC
+            LIMIT 1
+        }
+        WITH n1 as active_node, r1 as rb, attr1 as active_attr
+        WHERE rb.status = "active"
+        WITH active_attr
+        MATCH (active_attr)-[]-(peer)
+        WITH DISTINCT active_attr, peer
+        CALL (active_attr, peer) {
+            MATCH (active_attr)-[r]-(peer)
+            WHERE %(branch_filter)s
+            RETURN active_attr as a1, r as r1, peer as p1
+            ORDER BY r.branch_level DESC, r.from DESC
+            LIMIT 1
+        }
+        WITH a1 as active_attr, r1 as rb, p1 as peer_node
+        WHERE rb.status = "active"
+        CALL (peer_node, rb, active_attr) {
+            %(sub_query_all)s
+        }
+        WITH p2 as peer_node, rb, active_attr
+        FOREACH (i in CASE WHEN rb.branch = $branch_name THEN [1] ELSE [] END |
+            SET rb.to = $current_time
+        )
+        RETURN DISTINCT active_attr
+        """ % {
+            "branch_filter": branch_filter,
+            "sub_query_all": sub_query_all,
+            "node_kinds": node_kinds_str,
+        }
+        self.add_to_query(query)

infrahub/core/migrations/runner.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Sequence
+
+from infrahub.core import registry
+from infrahub.core.constants import GLOBAL_BRANCH_NAME
+from infrahub.core.migrations.graph import MIGRATIONS
+
+from .exceptions import MigrationFailureError
+from .shared import MigrationRequiringRebase
+
+if TYPE_CHECKING:
+    from infrahub.core.branch import Branch
+    from infrahub.database import InfrahubDatabase
+
+
+class MigrationRunner:
+    def __init__(self, branch: Branch) -> None:
+        if branch.name in (registry.default_branch, GLOBAL_BRANCH_NAME):
+            raise ValueError("MigrationRunner cannot be used to apply migration on default branches")
+
+        self.branch = branch
+        self.applicable_migrations = self._get_applicable_migrations()
+
+    def _get_applicable_migrations(self) -> Sequence[MigrationRequiringRebase]:
+        applicable_migrations = []
+        for migration_class in [m for m in MIGRATIONS if issubclass(m, MigrationRequiringRebase)]:
+            migration = migration_class.init()
+            if self.branch.graph_version and self.branch.graph_version > migration.minimum_version:
+                continue
+            applicable_migrations.append(migration)
+
+        return applicable_migrations
+
+    def has_migrations(self) -> bool:
+        return bool(self.applicable_migrations)
+
+    async def run(self, db: InfrahubDatabase) -> 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)
+            validation_result = None
+
+            if execution_result.success:
+                validation_result = await migration.validate_migration(db=db)
+
+            if not execution_result.success or (validation_result and not validation_result.success):
+                if execution_result.errors:
+                    raise MigrationFailureError(errors=execution_result.errors)
+
+                if validation_result and not validation_result.success:
+                    raise MigrationFailureError(errors=validation_result.errors)

infrahub/core/migrations/schema/attribute_kind_update.py

@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING, Any, Sequence
 
 from infrahub.types import is_large_attribute_type
 
-from ..query import AttributeMigrationQuery
+from ..query import AttributeMigrationQuery, MigrationBaseQuery
 from ..shared import AttributeSchemaMigration, MigrationResult
 
 if TYPE_CHECKING:
@@ -147,10 +147,16 @@ class AttributeKindUpdateMigration(AttributeSchemaMigration):
     name: str = "attribute.kind.update"
     queries: Sequence[type[AttributeMigrationQuery]] = [AttributeKindUpdateMigrationQuery]  # type: ignore[assignment]
 
-    async def execute(self, db: InfrahubDatabase, branch: Branch, at: Timestamp | str | None = None) -> MigrationResult:
+    async def execute(
+        self,
+        db: InfrahubDatabase,
+        branch: Branch,
+        at: Timestamp | str | None = None,
+        queries: Sequence[type[MigrationBaseQuery]] | None = None,
+    ) -> 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)
+        return await super().execute(db=db, branch=branch, at=at, queries=queries)

infrahub/core/migrations/schema/attribute_supports_profile.py

@@ -0,0 +1,90 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Sequence
+
+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
+
+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]:
+    if not isinstance(schema, (NodeSchema, GenericSchema)):
+        return [schema.kind]
+    schema_kinds = [f"Profile{schema.kind}"]
+    if isinstance(schema, GenericSchema) and schema.used_by:
+        schema_kinds += [f"Profile{kind}" for kind in schema.used_by]
+    return schema_kinds
+
+
+class ProfilesAttributeAddMigrationQuery(AttributeMigrationQuery, AttributeAddQuery):
+    name = "migration_profiles_attribute_add"
+
+    def __init__(
+        self,
+        migration: AttributeSchemaMigration,
+        **kwargs: Any,
+    ):
+        node_kinds = _get_node_kinds(migration.new_schema)
+        super().__init__(
+            migration=migration,
+            node_kinds=node_kinds,
+            attribute_name=migration.new_attribute_schema.name,
+            attribute_kind=migration.new_attribute_schema.kind,
+            branch_support=migration.new_attribute_schema.get_branch(),
+            default_value=migration.new_attribute_schema.default_value,
+            **kwargs,
+        )
+
+
+class ProfilesAttributeRemoveMigrationQuery(AttributeMigrationQuery, AttributeRemoveQuery):
+    name = "migration_profiles_attribute_remove"
+
+    def __init__(
+        self,
+        migration: AttributeSchemaMigration,
+        **kwargs: Any,
+    ):
+        node_kinds = _get_node_kinds(migration.new_schema)
+        super().__init__(
+            migration=migration,
+            attribute_name=migration.new_attribute_schema.name,
+            node_kinds=node_kinds,
+            **kwargs,
+        )
+
+
+class AttributeSupportsProfileUpdateMigration(AttributeSchemaMigration):
+    name: str = "attribute.supports_profile.update"
+    queries: Sequence[type[MigrationBaseQuery]] = []
+
+    async def execute(
+        self,
+        db: InfrahubDatabase,
+        branch: Branch,
+        at: Timestamp | str | None = None,
+        queries: Sequence[type[MigrationBaseQuery]] | None = None,  # noqa: ARG002
+    ) -> MigrationResult:
+        if (
+            # no change in whether the attribute should be used on profiles
+            self.previous_attribute_schema.support_profiles == self.new_attribute_schema.support_profiles
+            # the attribute is new, so there cannot be existing profiles to update
+            or self.previous_attribute_schema.id is None
+        ):
+            return MigrationResult()
+        profiles_queries: list[type[AttributeMigrationQuery]] = []
+        if self.new_attribute_schema.support_profiles:
+            profiles_queries.append(ProfilesAttributeAddMigrationQuery)
+        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)

infrahub/core/migrations/schema/node_attribute_add.py

@@ -4,15 +4,19 @@ from typing import TYPE_CHECKING, Any, Sequence
 
 from infrahub.core import registry
 from infrahub.core.node import Node
+from infrahub.core.schema.generic_schema import GenericSchema
+from infrahub.core.schema.node_schema import NodeSchema
 from infrahub.exceptions import PoolExhaustedError
 from infrahub.tasks.registry import update_branch_registry
 
-from ..query import AttributeMigrationQuery
+from ..query import AttributeMigrationQuery, MigrationBaseQuery
 from ..query.attribute_add import AttributeAddQuery
 from ..shared import AttributeSchemaMigration, 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
@@ -22,14 +26,27 @@ if TYPE_CHECKING:
 class NodeAttributeAddMigrationQuery01(AttributeMigrationQuery, AttributeAddQuery):
     name = "migration_node_attribute_add_01"
 
+    def _get_node_kinds(self, schema: MainSchemaTypes, new_attribute_schema: AttributeSchema) -> list[str]:
+        schema_kinds = [schema.kind]
+        if not isinstance(schema, (NodeSchema, GenericSchema)):
+            return schema_kinds
+        if new_attribute_schema.support_profiles:
+            schema_kinds.append(f"Profile{schema.kind}")
+            if isinstance(schema, GenericSchema) and schema.used_by:
+                schema_kinds.extend([f"Profile{kind}" for kind in schema.used_by])
+        return schema_kinds
+
     def __init__(
         self,
         migration: AttributeSchemaMigration,
         **kwargs: Any,
     ):
+        node_kinds = self._get_node_kinds(
+            schema=migration.new_schema, new_attribute_schema=migration.new_attribute_schema
+        )
         super().__init__(
             migration=migration,
-            node_kind=migration.new_schema.kind,
+            node_kinds=node_kinds,
             attribute_name=migration.new_attribute_schema.name,
             attribute_kind=migration.new_attribute_schema.kind,
             branch_support=migration.new_attribute_schema.get_branch().value,
@@ -42,6 +59,17 @@ class NodeAttributeAddMigration(AttributeSchemaMigration):
     name: str = "node.attribute.add"
     queries: Sequence[type[AttributeMigrationQuery]] = [NodeAttributeAddMigrationQuery01]  # type: ignore[assignment]
 
+    async def execute(
+        self,
+        db: InfrahubDatabase,
+        branch: Branch,
+        at: Timestamp | str | None = None,
+        queries: Sequence[type[MigrationBaseQuery]] | None = None,
+    ) -> MigrationResult:
+        if self.new_attribute_schema.inherited is True:
+            return MigrationResult()
+        return await super().execute(db=db, branch=branch, at=at, queries=queries)
+
     async def execute_post_queries(
         self,
         db: InfrahubDatabase,
@@ -52,8 +80,11 @@ class NodeAttributeAddMigration(AttributeSchemaMigration):
         if self.new_attribute_schema.kind != "NumberPool":
             return result
 
-        number_pool: CoreNumberPool = await Node.fetch_or_create_number_pool(  # type: ignore[assignment]
-            db=db, branch=branch, schema_node=self.new_schema, schema_attribute=self.new_attribute_schema
+        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,
         )
 
         await update_branch_registry(db=db, branch=branch)

infrahub/core/migrations/schema/node_attribute_remove.py

@@ -1,123 +1,27 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Sequence
-
-from infrahub.core.constants import RelationshipStatus
-from infrahub.core.graph.schema import GraphAttributeRelationships
-from infrahub.core.schema.generic_schema import GenericSchema
+from typing import Any, Sequence
 
 from ..query import AttributeMigrationQuery
+from ..query.attribute_remove import AttributeRemoveQuery
 from ..shared import AttributeSchemaMigration
 
-if TYPE_CHECKING:
-    from pydantic.fields import FieldInfo
-
-    from infrahub.database import InfrahubDatabase
-
 
-class NodeAttributeRemoveMigrationQuery01(AttributeMigrationQuery):
+class NodeAttributeRemoveMigrationQuery01(AttributeMigrationQuery, AttributeRemoveQuery):
     name = "migration_node_attribute_remove_01"
     insert_return: bool = False
 
-    async def query_init(self, db: InfrahubDatabase, **kwargs: dict[str, Any]) -> None:  # noqa: ARG002
-        branch_filter, branch_params = self.branch.get_query_filter_path(at=self.at.to_string())
-        self.params.update(branch_params)
-
-        attr_name = self.migration.schema_path.field_name
-        kinds_to_ignore = []
-        if isinstance(self.migration.new_node_schema, GenericSchema) and attr_name is not None:
-            for inheriting_schema_kind in self.migration.new_node_schema.used_by:
-                node_schema = db.schema.get_node_schema(
-                    name=inheriting_schema_kind, branch=self.branch, duplicate=False
-                )
-                attr_schema = node_schema.get_attribute_or_none(name=attr_name)
-                if attr_schema and not attr_schema.inherited:
-                    kinds_to_ignore.append(inheriting_schema_kind)
-
-        self.params["node_kind"] = self.migration.new_schema.kind
-        self.params["kinds_to_ignore"] = kinds_to_ignore
-        self.params["attr_name"] = attr_name
-        self.params["current_time"] = self.at.to_string()
-        self.params["branch_name"] = self.branch.name
-        self.params["branch_support"] = self.migration.previous_attribute_schema.get_branch().value
-
-        self.params["rel_props"] = {
-            "branch": self.branch.name,
-            "branch_level": self.branch.hierarchy_level,
-            "status": RelationshipStatus.DELETED.value,
-            "from": self.at.to_string(),
-        }
-
-        def render_sub_query_per_rel_type(rel_type: str, rel_def: FieldInfo) -> str:
-            subquery = [
-                "WITH peer_node, rb, active_attr",
-                f'WHERE type(rb) = "{rel_type}"',
-            ]
-            if rel_def.default.direction.value == "outbound":
-                subquery.append(f"CREATE (active_attr)-[:{rel_type} $rel_props ]->(peer_node)")
-            elif rel_def.default.direction.value == "inbound":
-                subquery.append(f"CREATE (active_attr)<-[:{rel_type} $rel_props ]-(peer_node)")
-            else:
-                subquery.append(f"CREATE (active_attr)-[:{rel_type} $rel_props ]-(peer_node)")
-
-            subquery.append("RETURN peer_node as p2")
-            return "\n".join(subquery)
-
-        sub_queries = [
-            render_sub_query_per_rel_type(rel_type, rel_def)
-            for rel_type, rel_def in GraphAttributeRelationships.model_fields.items()
-        ]
-        sub_query_all = "\nUNION\n".join(sub_queries)
-
-        query = """
-        // Find all the active nodes
-        MATCH (node:%(node_kind)s)
-        WHERE (size($kinds_to_ignore) = 0 OR NOT any(l IN labels(node) WHERE l IN $kinds_to_ignore))
-        AND exists((node)-[:HAS_ATTRIBUTE]-(:Attribute { name: $attr_name }))
-        CALL (node) {
-            MATCH (root:Root)<-[r:IS_PART_OF]-(node)
-            WHERE %(branch_filter)s
-            RETURN node as n1, r as r1
-            ORDER BY r.branch_level DESC, r.from DESC
-            LIMIT 1
-        }
-        WITH n1 as active_node, r1 as rb
-        WHERE rb.status = "active"
-        // Find all the attributes that need to be updated
-        CALL (active_node) {
-            MATCH (active_node)-[r:HAS_ATTRIBUTE]-(attr:Attribute { name: $attr_name })
-            WHERE %(branch_filter)s
-            RETURN active_node as n1, r as r1, attr as attr1
-            ORDER BY r.branch_level DESC, r.from DESC
-            LIMIT 1
-        }
-        WITH n1 as active_node, r1 as rb, attr1 as active_attr
-        WHERE rb.status = "active"
-        WITH active_attr
-        MATCH (active_attr)-[]-(peer)
-        CALL (active_attr, peer) {
-            MATCH (active_attr)-[r]-(peer)
-            WHERE %(branch_filter)s
-            RETURN active_attr as a1, r as r1, peer as p1
-            ORDER BY r.branch_level DESC, r.from DESC
-            LIMIT 1
-        }
-        WITH a1 as active_attr, r1 as rb, p1 as peer_node
-        WHERE rb.status = "active"
-        CALL (peer_node, rb, active_attr) {
-            %(sub_query_all)s
-        }
-        WITH p2 as peer_node, rb, active_attr
-        FOREACH (i in CASE WHEN rb.branch = $branch_name THEN [1] ELSE [] END |
-            SET rb.to = $current_time
+    def __init__(
+        self,
+        migration: AttributeSchemaMigration,
+        **kwargs: Any,
+    ):
+        super().__init__(
+            migration=migration,
+            attribute_name=migration.previous_attribute_schema.name,
+            node_kinds=[migration.new_schema.kind],
+            **kwargs,
         )
-        RETURN DISTINCT active_attr
-        """ % {
-            "branch_filter": branch_filter,
-            "sub_query_all": sub_query_all,
-            "node_kind": self.migration.new_schema.kind,
-        }
-        self.add_to_query(query)
 
 
 class NodeAttributeRemoveMigration(AttributeSchemaMigration):

infrahub/core/migrations/schema/node_kind_update.py

@@ -2,8 +2,9 @@ from __future__ import annotations
 
 from typing import Any, Sequence
 
+from ..query import MigrationQuery
 from ..query.node_duplicate import NodeDuplicateQuery, SchemaNodeInfo
-from ..shared import MigrationQuery, SchemaMigration
+from ..shared import SchemaMigration
 
 
 class NodeKindUpdateMigrationQuery01(MigrationQuery, NodeDuplicateQuery):

infrahub/core/migrations/schema/node_remove.py

@@ -5,7 +5,8 @@ from typing import TYPE_CHECKING, Any, Sequence
 from infrahub.core.constants import RelationshipStatus
 from infrahub.core.graph.schema import GraphNodeRelationships, GraphRelDirection
 
-from ..shared import MigrationQuery, SchemaMigration
+from ..query import MigrationQuery
+from ..shared import SchemaMigration
 
 if TYPE_CHECKING:
     from pydantic.fields import FieldInfo

infrahub/core/migrations/schema/placeholder_dummy.py

@@ -4,9 +4,10 @@ from typing import Sequence
 
 from pydantic import Field
 
-from ..shared import MigrationQuery, SchemaMigration
+from ..query import MigrationBaseQuery  # noqa: TC001
+from ..shared import SchemaMigration
 
 
 class PlaceholderDummyMigration(SchemaMigration):
     name: str = "dummy.placeholder"
-    queries: Sequence[type[MigrationQuery]] = Field(default_factory=list)
+    queries: Sequence[type[MigrationBaseQuery]] = Field(default_factory=list)