infrahub-server 1.4.10__py3-none-any.whl → 1.5.0b1__py3-none-any.whl

infrahub/core/manager.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 from copy import copy
-from functools import reduce
 from typing import TYPE_CHECKING, Any, Iterable, Literal, TypeVar, overload
 
 from infrahub_sdk.utils import deep_merge_dict, is_valid_uuid
@@ -11,9 +10,7 @@ from infrahub.core.node import Node
 from infrahub.core.node.delete_validator import NodeDeleteValidator
 from infrahub.core.query.node import (
     AttributeFromDB,
-    AttributeNodePropertyFromDB,
     GroupedPeerNodes,
-    NodeAttributesFromDB,
     NodeGetHierarchyQuery,
     NodeGetListQuery,
     NodeListGetAttributeQuery,
@@ -63,72 +60,21 @@ def identify_node_class(node: NodeToProcess) -> type[Node]:
 
 
 def get_schema(
-    db: InfrahubDatabase, branch: Branch, node_schema: type[SchemaProtocol] | MainSchemaTypes | str
+    db: InfrahubDatabase,
+    branch: Branch,
+    node_schema: type[SchemaProtocol] | MainSchemaTypes | str,
+    duplicate: bool = False,
 ) -> MainSchemaTypes:
     if isinstance(node_schema, str):
-        return db.schema.get(name=node_schema, branch=branch.name)
+        return db.schema.get(name=node_schema, branch=branch.name, duplicate=duplicate)
     if hasattr(node_schema, "_is_runtime_protocol") and node_schema._is_runtime_protocol:
-        return db.schema.get(name=node_schema.__name__, branch=branch.name)
+        return db.schema.get(name=node_schema.__name__, branch=branch.name, duplicate=duplicate)
     if not isinstance(node_schema, (MainSchemaTypes)):
         raise ValueError(f"Invalid schema provided {node_schema}")
 
     return node_schema
 
 
-class ProfileAttributeIndex:
-    def __init__(
-        self,
-        profile_attributes_id_map: dict[str, NodeAttributesFromDB],
-        profile_ids_by_node_id: dict[str, list[str]],
-    ) -> None:
-        self._profile_attributes_id_map = profile_attributes_id_map
-        self._profile_ids_by_node_id = profile_ids_by_node_id
-
-    def apply_profiles(self, node_data_dict: dict[str, Any]) -> dict[str, Any]:
-        updated_data: dict[str, Any] = {**node_data_dict}
-        node_id = node_data_dict.get("id")
-        profile_ids = self._profile_ids_by_node_id.get(node_id, [])
-        if not profile_ids:
-            return updated_data
-        profiles = [
-            self._profile_attributes_id_map[p_id] for p_id in profile_ids if p_id in self._profile_attributes_id_map
-        ]
-
-        def get_profile_priority(nafd: NodeAttributesFromDB) -> tuple[int | float, str]:
-            try:
-                return (int(nafd.attrs.get("profile_priority").value), nafd.node.get("uuid"))
-            except (TypeError, AttributeError):
-                return (float("inf"), "")
-
-        profiles.sort(key=get_profile_priority)
-
-        for attr_name, attr_data in updated_data.items():
-            if not isinstance(attr_data, AttributeFromDB):
-                continue
-            if not attr_data.is_default:
-                continue
-            profile_value, profile_uuid = None, None
-            index = 0
-
-            while profile_value is None and index <= (len(profiles) - 1):
-                try:
-                    profile_value = profiles[index].attrs[attr_name].value
-                    if profile_value != "NULL":
-                        profile_uuid = profiles[index].node["uuid"]
-                        break
-                    profile_value = None
-                except (IndexError, KeyError, AttributeError):
-                    ...
-                index += 1
-
-            if profile_value is not None:
-                attr_data.value = profile_value
-                attr_data.is_from_profile = True
-                attr_data.is_default = False
-                attr_data.node_properties["source"] = AttributeNodePropertyFromDB(uuid=profile_uuid, labels=[])
-        return updated_data
-
-
 class NodeManager:
     @overload
     @classmethod
@@ -1129,21 +1075,11 @@ class NodeManager:
         )
         await query.execute(db=db)
         nodes_info_by_id: dict[str, NodeToProcess] = {node.node_uuid: node async for node in query.get_nodes(db=db)}
-        profile_ids_by_node_id = query.get_profile_ids_by_node_id()
-        all_profile_ids = reduce(
-            lambda all_ids, these_ids: all_ids | set(these_ids), profile_ids_by_node_id.values(), set()
-        )
-
-        if fields and all_profile_ids:
-            if "profile_priority" not in fields:
-                fields["profile_priority"] = {}
-            if "value" not in fields["profile_priority"]:
-                fields["profile_priority"]["value"] = None
 
         # Query list of all Attributes
         query = await NodeListGetAttributeQuery.init(
             db=db,
-            ids=list(nodes_info_by_id.keys()) + list(all_profile_ids),
+            ids=list(nodes_info_by_id.keys()),
             fields=fields,
             branch=branch,
             include_source=include_source,
@@ -1153,17 +1089,7 @@ class NodeManager:
             branch_agnostic=branch_agnostic,
         )
         await query.execute(db=db)
-        all_node_attributes = query.get_attributes_group_by_node()
-        profile_attributes: dict[str, dict[str, AttributeFromDB]] = {}
-        node_attributes: dict[str, dict[str, AttributeFromDB]] = {}
-        for node_id, attribute_dict in all_node_attributes.items():
-            if node_id in all_profile_ids:
-                profile_attributes[node_id] = attribute_dict
-            else:
-                node_attributes[node_id] = attribute_dict
-        profile_index = ProfileAttributeIndex(
-            profile_attributes_id_map=profile_attributes, profile_ids_by_node_id=profile_ids_by_node_id
-        )
+        node_attributes = query.get_attributes_group_by_node()
 
         nodes: dict[str, Node] = {}
 
@@ -1192,11 +1118,10 @@ class NodeManager:
                 for attr_name, attr in node_attributes[node_id].attrs.items():
                     new_node_data[attr_name] = attr
 
-            new_node_data_with_profile_overrides = profile_index.apply_profiles(new_node_data)
             node_class = identify_node_class(node=node)
             node_branch = await registry.get_branch(db=db, branch=node.branch)
             item = await node_class.init(schema=node.schema, branch=node_branch, at=at, db=db)
-            await item.load(**new_node_data_with_profile_overrides, db=db)
+            await item.load(**new_node_data, db=db)
 
             nodes[node_id] = item
 

infrahub/core/migrations/graph/__init__.py

@@ -41,6 +41,9 @@ from .m036_drop_attr_value_index import Migration036
 from .m037_index_attr_vals import Migration037
 from .m038_redo_0000_prefix_fix import Migration038
 from .m039_ipam_reconcile import Migration039
+from .m040_profile_attrs_in_db import Migration040
+from .m041_create_hfid_display_label_in_db import Migration041
+from .m042_backfill_hfid_display_label_in_db import Migration042
 
 if TYPE_CHECKING:
     from infrahub.core.root import Root
@@ -87,6 +90,9 @@ MIGRATIONS: list[type[GraphMigration | InternalSchemaMigration | ArbitraryMigrat
     Migration037,
     Migration038,
     Migration039,
+    Migration040,
+    Migration041,
+    Migration042,
 ]
 
 

infrahub/core/migrations/graph/m040_profile_attrs_in_db.py

@@ -0,0 +1,166 @@
+from __future__ import annotations
+
+from collections import defaultdict
+from typing import TYPE_CHECKING, Any
+
+from rich.console import Console
+from rich.progress import Progress
+
+from infrahub.core.branch.models import Branch
+from infrahub.core.initialization import initialization
+from infrahub.core.manager import NodeManager
+from infrahub.core.migrations.shared import MigrationResult
+from infrahub.core.query import Query, QueryType
+from infrahub.core.timestamp import Timestamp
+from infrahub.lock import initialize_lock
+from infrahub.log import get_logger
+from infrahub.profiles.node_applier import NodeProfilesApplier
+
+from ..shared import ArbitraryMigration
+
+if TYPE_CHECKING:
+    from infrahub.core.node import Node
+    from infrahub.database import InfrahubDatabase
+
+log = get_logger()
+
+
+class GetProfilesByBranchQuery(Query):
+    """
+    Get CoreProfile UUIDs by which branches they have attribute updates on
+    """
+
+    name = "get_profiles_by_branch"
+    type = QueryType.READ
+    insert_return = False
+
+    async def query_init(self, db: InfrahubDatabase, **kwargs: dict[str, Any]) -> None:  # noqa: ARG002
+        query = """
+MATCH (profile:CoreProfile)-[:HAS_ATTRIBUTE]->(attr:Attribute)-[e:HAS_VALUE]->(:AttributeValue)
+WITH DISTINCT profile.uuid AS profile_uuid, e.branch AS branch
+RETURN profile_uuid, collect(branch) AS branches
+        """
+        self.add_to_query(query)
+        self.return_labels = ["profile_uuid", "branches"]
+
+    def get_profile_ids_by_branch(self) -> dict[str, set[str]]:
+        """Get dictionary of branch names to set of updated profile UUIDs"""
+        profiles_by_branch = defaultdict(set)
+        for result in self.get_results():
+            profile_uuid = result.get_as_type("profile_uuid", str)
+            branches = result.get_as_type("branches", list[str])
+            for branch in branches:
+                profiles_by_branch[branch].add(profile_uuid)
+        return profiles_by_branch
+
+
+class GetNodesWithProfileUpdatesByBranchQuery(Query):
+    """
+    Get Node UUIDs by which branches they have updated profiles on
+    """
+
+    name = "get_nodes_with_profile_updates_by_branch"
+    type = QueryType.READ
+    insert_return = False
+
+    async def query_init(self, db: InfrahubDatabase, **kwargs: dict[str, Any]) -> None:  # noqa: ARG002
+        query = """
+MATCH (node:Node)-[e1:IS_RELATED]->(:Relationship {name: "node__profile"})
+WHERE NOT node:CoreProfile
+WITH DISTINCT node.uuid AS node_uuid, e1.branch AS branch
+RETURN node_uuid, collect(branch) AS branches
+        """
+        self.add_to_query(query)
+        self.return_labels = ["node_uuid", "branches"]
+
+    def get_node_ids_by_branch(self) -> dict[str, set[str]]:
+        """Get dictionary of branch names to set of updated node UUIDs"""
+        nodes_by_branch = defaultdict(set)
+        for result in self.get_results():
+            node_uuid = result.get_as_type("node_uuid", str)
+            branches = result.get_as_type("branches", list[str])
+            for branch in branches:
+                nodes_by_branch[branch].add(node_uuid)
+        return nodes_by_branch
+
+
+class Migration040(ArbitraryMigration):
+    """
+    Save profile attribute values on each node using the profile in the database
+    For any profile that has updates on a given branch (including default branch)
+    - run NodeProfilesApplier.apply_profiles on each node related to the profile on that branch
+    For any node that has an updated relationship to a profile on a given branch
+    - run NodeProfilesApplier.apply_profiles on the node on that branch
+    """
+
+    name: str = "040_profile_attrs_in_db"
+    minimum_version: int = 39
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self._appliers_by_branch: dict[str, NodeProfilesApplier] = {}
+
+    async def _get_profile_applier(self, db: InfrahubDatabase, branch_name: str) -> NodeProfilesApplier:
+        if branch_name not in self._appliers_by_branch:
+            branch = await Branch.get_by_name(db=db, name=branch_name)
+            self._appliers_by_branch[branch_name] = NodeProfilesApplier(db=db, branch=branch)
+        return self._appliers_by_branch[branch_name]
+
+    async def validate_migration(self, db: InfrahubDatabase) -> MigrationResult:  # noqa: ARG002
+        return MigrationResult()
+
+    async def execute(self, db: InfrahubDatabase) -> MigrationResult:
+        console = Console()
+        result = MigrationResult()
+        # load schemas from database into registry
+        initialize_lock()
+        await initialization(db=db)
+
+        console.print("Gathering profiles for each branch...", end="")
+        get_profiles_by_branch_query = await GetProfilesByBranchQuery.init(db=db)
+        await get_profiles_by_branch_query.execute(db=db)
+        profiles_ids_by_branch = get_profiles_by_branch_query.get_profile_ids_by_branch()
+
+        profiles_by_branch: dict[str, list[Node]] = {}
+        for branch_name, profile_ids in profiles_ids_by_branch.items():
+            profiles_map = await NodeManager.get_many(db=db, branch=branch_name, ids=list(profile_ids))
+            profiles_by_branch[branch_name] = list(profiles_map.values())
+        console.print("done")
+
+        node_ids_to_update_by_branch: dict[str, set[str]] = defaultdict(set)
+        total_size = sum(len(profiles) for profiles in profiles_by_branch.values())
+        with Progress() as progress:
+            gather_nodes_task = progress.add_task(
+                "Gathering affected objects for each profile on each branch...", total=total_size
+            )
+
+            for branch_name, profiles in profiles_by_branch.items():
+                for profile in profiles:
+                    node_relationship_manager = profile.get_relationship("related_nodes")
+                    node_peers = await node_relationship_manager.get_db_peers(db=db)
+                    node_ids_to_update_by_branch[branch_name].update({str(peer.peer_id) for peer in node_peers})
+                    progress.update(gather_nodes_task, advance=1)
+
+        console.print("Identifying nodes with profile updates by branch...", end="")
+        get_nodes_with_profile_updates_by_branch_query = await GetNodesWithProfileUpdatesByBranchQuery.init(db=db)
+        await get_nodes_with_profile_updates_by_branch_query.execute(db=db)
+        nodes_ids_by_branch = get_nodes_with_profile_updates_by_branch_query.get_node_ids_by_branch()
+        for branch_name, node_ids in nodes_ids_by_branch.items():
+            node_ids_to_update_by_branch[branch_name].update(node_ids)
+        console.print("done")
+
+        right_now = Timestamp()
+        total_size = sum(len(node_ids) for node_ids in node_ids_to_update_by_branch.values())
+        with Progress() as progress:
+            apply_task = progress.add_task("Applying profiles to nodes...", total=total_size)
+            for branch_name, node_ids in node_ids_to_update_by_branch.items():
+                applier = await self._get_profile_applier(db=db, branch_name=branch_name)
+                for node_id in node_ids:
+                    node = await NodeManager.get_one(db=db, branch=branch_name, id=node_id, at=right_now)
+                    if node:
+                        updated_field_names = await applier.apply_profiles(node=node)
+                        if updated_field_names:
+                            await node.save(db=db, fields=updated_field_names, at=right_now)
+                    progress.update(apply_task, advance=1)
+
+        return result

infrahub/core/migrations/graph/m041_create_hfid_display_label_in_db.py

@@ -0,0 +1,97 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from rich.progress import Progress
+from typing_extensions import Self
+
+from infrahub.core import registry
+from infrahub.core.constants import SchemaPathType
+from infrahub.core.initialization import initialization
+from infrahub.core.migrations.schema.node_attribute_add import NodeAttributeAddMigration
+from infrahub.core.migrations.shared import InternalSchemaMigration, MigrationResult
+from infrahub.core.path import SchemaPath
+from infrahub.lock import initialize_lock
+
+if TYPE_CHECKING:
+    from infrahub.database import InfrahubDatabase
+
+
+class Migration041(InternalSchemaMigration):
+    name: str = "041_create_hfid_display_label_in_db"
+    minimum_version: int = 40
+
+    @classmethod
+    def init(cls, **kwargs: Any) -> Self:
+        internal_schema = cls.get_internal_schema()
+        schema_node = internal_schema.get_node(name="SchemaNode")
+        schema_generic = internal_schema.get_node(name="SchemaGeneric")
+
+        cls.migrations = [
+            # HFID is not needed, it was introduced at graph v8
+            NodeAttributeAddMigration(
+                new_node_schema=schema_node,
+                previous_node_schema=schema_node,
+                schema_path=SchemaPath(
+                    schema_kind="SchemaNode", path_type=SchemaPathType.ATTRIBUTE, field_name="display_label"
+                ),
+            ),
+            NodeAttributeAddMigration(
+                new_node_schema=schema_generic,
+                previous_node_schema=schema_generic,
+                schema_path=SchemaPath(
+                    schema_kind="SchemaGeneric", path_type=SchemaPathType.ATTRIBUTE, field_name="display_label"
+                ),
+            ),
+        ]
+        return cls(migrations=cls.migrations, **kwargs)  # type: ignore[arg-type]
+
+    async def execute(self, db: InfrahubDatabase) -> MigrationResult:
+        result = MigrationResult()
+
+        # load schemas from database into registry
+        initialize_lock()
+        await initialization(db=db)
+
+        default_branch = registry.get_branch_from_registry()
+        schema_branch = await registry.schema.load_schema_from_db(db=db, branch=default_branch)
+
+        migrations = list(self.migrations)
+
+        for node_schema_kind in schema_branch.node_names:
+            schema = schema_branch.get(name=node_schema_kind, duplicate=False)
+            migrations.extend(
+                [
+                    NodeAttributeAddMigration(
+                        new_node_schema=schema,
+                        previous_node_schema=schema,
+                        schema_path=SchemaPath(
+                            schema_kind=schema.kind, path_type=SchemaPathType.ATTRIBUTE, field_name="human_friendly_id"
+                        ),
+                    ),
+                    NodeAttributeAddMigration(
+                        new_node_schema=schema,
+                        previous_node_schema=schema,
+                        schema_path=SchemaPath(
+                            schema_kind=schema.kind, path_type=SchemaPathType.ATTRIBUTE, field_name="display_label"
+                        ),
+                    ),
+                ]
+            )
+
+        with Progress() as progress:
+            update_task = progress.add_task("Adding HFID and display label to nodes", total=len(migrations))
+
+            for migration in migrations:
+                try:
+                    execution_result = await migration.execute(db=db, branch=default_branch)
+                    result.errors.extend(execution_result.errors)
+                    progress.update(update_task, advance=1)
+                except Exception as exc:
+                    result.errors.append(str(exc))
+                    return result
+
+        return result
+
+    async def validate_migration(self, db: InfrahubDatabase) -> MigrationResult:  # noqa: ARG002
+        return MigrationResult()

infrahub/core/migrations/graph/m042_backfill_hfid_display_label_in_db.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Sequence
+
+from rich.progress import Progress, TaskID
+
+from infrahub.core import registry
+from infrahub.core.initialization import initialization
+from infrahub.core.manager import NodeManager
+from infrahub.core.migrations.shared import MigrationResult
+from infrahub.lock import initialize_lock
+
+from ..shared import ArbitraryMigration
+
+if TYPE_CHECKING:
+    from infrahub.core.node import Node
+    from infrahub.core.schema import MainSchemaTypes
+    from infrahub.database import InfrahubDatabase
+
+
+class Migration042(ArbitraryMigration):
+    """
+    Backfill `human_friendly_id` and `display_label` attributes for nodes with schemas that define them.
+    """
+
+    name: str = "042_backfill_hfid_display_label_in_db"
+    minimum_version: int = 41
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+
+    async def validate_migration(self, db: InfrahubDatabase) -> MigrationResult:  # noqa: ARG002
+        return MigrationResult()
+
+    async def _update_batch(
+        self,
+        db: InfrahubDatabase,
+        node_schema: MainSchemaTypes,
+        nodes: Sequence[Node],
+        progress: Progress,
+        update_task: TaskID,
+    ) -> None:
+        for node in nodes:
+            fields = []
+            if node_schema.human_friendly_id:
+                await node.add_human_friendly_id(db=db)
+                fields.append("human_friendly_id")
+            if node_schema.display_label:
+                await node.add_display_label(db=db)
+                fields.append("display_label")
+
+            if fields:
+                await node.save(db=db, fields=fields)
+
+            progress.update(task_id=update_task, advance=1)
+
+    async def execute(self, db: InfrahubDatabase) -> MigrationResult:
+        result = MigrationResult()
+        # load schemas from database into registry
+        initialize_lock()
+        await initialization(db=db)
+
+        schemas_to_update: dict[MainSchemaTypes, int] = {}
+        for node_schema in registry.get_full_schema(duplicate=False).values():
+            if node_schema.is_generic_schema or (not node_schema.human_friendly_id and not node_schema.display_label):
+                continue
+
+            node_count = await NodeManager.count(db=db, schema=node_schema)
+            if node_count:
+                schemas_to_update[node_schema] = node_count
+
+        with Progress() as progress:
+            batch_size = 1000
+            update_task = progress.add_task(
+                "Backfill HFID and display_label for nodes in database", total=sum(schemas_to_update.values())
+            )
+
+            for schema, count in schemas_to_update.items():
+                for offset in range(0, count, batch_size):
+                    limit = min(batch_size, count - offset)
+                    nodes: list[Node] = await NodeManager.query(db=db, schema=schema, offset=offset, limit=limit)
+                    await self._update_batch(
+                        db=db, node_schema=schema, nodes=nodes, progress=progress, update_task=update_task
+                    )
+
+        return result

infrahub/core/migrations/schema/node_attribute_add.py

@@ -52,8 +52,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/shared.py

@@ -10,8 +10,7 @@ from infrahub.core.path import SchemaPath  # noqa: TC001
 from infrahub.core.query import Query  # noqa: TC001
 from infrahub.core.schema import (
     AttributeSchema,
-    GenericSchema,
-    NodeSchema,
+    MainSchemaTypes,
     RelationshipSchema,
     SchemaRoot,
     internal_schema,
@@ -43,8 +42,8 @@ class SchemaMigration(BaseModel):
     name: str = Field(..., description="Name of the migration")
     queries: Sequence[type[MigrationQuery]] = Field(..., description="List of queries to execute for this migration")
 
-    new_node_schema: NodeSchema | GenericSchema | None = None
-    previous_node_schema: NodeSchema | GenericSchema | None = None
+    new_node_schema: MainSchemaTypes | None = None
+    previous_node_schema: MainSchemaTypes | None = None
     schema_path: SchemaPath
 
     async def execute_pre_queries(
@@ -91,13 +90,13 @@ class SchemaMigration(BaseModel):
         return result
 
     @property
-    def new_schema(self) -> NodeSchema | GenericSchema:
+    def new_schema(self) -> MainSchemaTypes:
         if self.new_node_schema:
             return self.new_node_schema
         raise ValueError("new_node_schema hasn't been initialized")
 
     @property
-    def previous_schema(self) -> NodeSchema | GenericSchema:
+    def previous_schema(self) -> MainSchemaTypes:
         if self.previous_node_schema:
             return self.previous_node_schema
         raise ValueError("previous_node_schema hasn't been initialized")