infrahub-server 1.2.0b1__py3-none-any.whl → 1.2.1__py3-none-any.whl

infrahub/cli/upgrade.py

@@ -0,0 +1,146 @@
+from __future__ import annotations
+
+import logging
+import os
+from typing import TYPE_CHECKING
+
+import typer
+from deepdiff import DeepDiff
+from infrahub_sdk.async_typer import AsyncTyper
+from prefect.client.orchestration import get_client
+from rich import print as rprint
+
+from infrahub import config
+from infrahub.core.constants import InfrahubKind
+from infrahub.core.initialization import (
+    create_anonymous_role,
+    create_default_roles,
+    create_super_administrator_role,
+    create_super_administrators_group,
+    initialize_registry,
+)
+from infrahub.core.manager import NodeManager
+from infrahub.menu.menu import default_menu
+from infrahub.menu.models import MenuDict
+from infrahub.menu.repository import MenuRepository
+from infrahub.menu.utils import create_default_menu
+from infrahub.services import InfrahubServices
+from infrahub.services.adapters.message_bus.local import BusSimulator
+from infrahub.services.adapters.workflow.local import WorkflowLocalExecution
+from infrahub.trigger.tasks import trigger_configure_all
+from infrahub.workflows.initialization import (
+    setup_blocks,
+    setup_deployments,
+    setup_worker_pools,
+)
+
+from .db import initialize_internal_schema, migrate_database, update_core_schema
+
+if TYPE_CHECKING:
+    from infrahub.cli.context import CliContext
+    from infrahub.database import InfrahubDatabase
+
+app = AsyncTyper()
+
+
+@app.command(name="upgrade")
+async def upgrade_cmd(
+    ctx: typer.Context,
+    config_file: str = typer.Argument("infrahub.toml", envvar="INFRAHUB_CONFIG"),
+    check: bool = typer.Option(False, help="Check the state of the system without upgrading."),
+) -> None:
+    """Upgrade Infrahub to the latest version."""
+
+    logging.getLogger("infrahub").setLevel(logging.WARNING)
+    logging.getLogger("neo4j").setLevel(logging.ERROR)
+    logging.getLogger("prefect").setLevel(logging.ERROR)
+    os.environ["PREFECT_SERVER_ANALYTICS_ENABLED"] = "false"
+
+    config.load_and_exit(config_file_name=config_file)
+
+    context: CliContext = ctx.obj
+    dbdriver = await context.init_db(retry=1)
+
+    service = await InfrahubServices.new(
+        database=dbdriver, message_bus=BusSimulator(), workflow=WorkflowLocalExecution()
+    )
+    await initialize_registry(db=dbdriver)
+
+    # NOTE add step to validate if the database and the task manager are reachable
+
+    # -------------------------------------------
+    # Add pre-upgrade  validation
+    # -------------------------------------------
+
+    # -------------------------------------------
+    # Upgrade Infrahub Database and Schema
+    # -------------------------------------------
+    await migrate_database(db=dbdriver, initialize=False, check=check)
+
+    await initialize_internal_schema()
+    await update_core_schema(db=dbdriver, service=service, initialize=False)
+
+    # -------------------------------------------
+    # Upgrade Internal Objects, generated and managed by Infrahub
+    # -------------------------------------------
+    await upgrade_menu(db=dbdriver)
+    await upgrade_permissions(db=dbdriver)
+
+    # -------------------------------------------
+    # Upgrade External system : Task Manager
+    # -------------------------------------------
+    async with get_client(sync_client=False) as client:
+        await setup_blocks()
+        await setup_worker_pools(client=client)
+        await setup_deployments(client=client)
+        await trigger_configure_all(service=service)
+
+    await dbdriver.close()
+
+
+async def upgrade_menu(db: InfrahubDatabase) -> None:
+    menu_repository = MenuRepository(db=db)
+    menu_nodes = await menu_repository.get_menu_db()
+    menu_items = await menu_repository.get_menu(nodes=menu_nodes)
+    default_menu_dict = MenuDict.from_definition_list(default_menu)
+
+    if not menu_nodes:
+        await create_default_menu(db=db)
+        return
+
+    diff_menu = DeepDiff(menu_items.to_rest(), default_menu_dict.to_rest(), ignore_order=True)
+
+    if not diff_menu:
+        rprint("Menu Up to date, nothing to update")
+        return
+
+    await menu_repository.update_menu(existing_menu=menu_items, new_menu=default_menu_dict, menu_nodes=menu_nodes)
+    rprint("Menu has been updated")
+
+
+async def upgrade_permissions(db: InfrahubDatabase) -> None:
+    existing_permissions = await NodeManager.query(
+        schema=InfrahubKind.OBJECTPERMISSION,
+        db=db,
+        limit=1,
+    )
+    if existing_permissions:
+        rprint("Permissions Up to date, nothing to update")
+        return
+
+    await setup_permissions(db=db)
+    rprint("Permissions have been updated")
+
+
+async def setup_permissions(db: InfrahubDatabase) -> None:
+    existing_accounts = await NodeManager.query(
+        schema=InfrahubKind.ACCOUNT,
+        db=db,
+        limit=1,
+    )
+    administrator_role = await create_super_administrator_role(db=db)
+    await create_super_administrators_group(db=db, role=administrator_role, admin_accounts=existing_accounts)
+    await create_default_roles(db=db)
+
+    if config.SETTINGS.main.allow_anonymous_access:
+        await create_anonymous_role(db=db)

infrahub/computed_attribute/gather.py

@@ -0,0 +1,185 @@
+from __future__ import annotations
+
+from collections import defaultdict
+from typing import TYPE_CHECKING
+
+from prefect import task
+from prefect.cache_policies import NONE
+from prefect.logging import get_run_logger
+
+from infrahub.core.constants import InfrahubKind
+from infrahub.core.manager import NodeManager
+from infrahub.core.protocols import CoreGenericRepository, CoreGraphQLQuery
+from infrahub.core.registry import registry
+from infrahub.database import InfrahubDatabase  # noqa: TC001  needed for prefect flow
+from infrahub.git.utils import get_repositories_commit_per_branch
+from infrahub.graphql.analyzer import InfrahubGraphQLQueryAnalyzer
+from infrahub.graphql.initialization import prepare_graphql_params
+
+from .models import (
+    ComputedAttrJinja2TriggerDefinition,
+    ComputedAttrPythonQueryTriggerDefinition,
+    ComputedAttrPythonTriggerDefinition,
+    PythonTransformComputedAttribute,
+)
+
+if TYPE_CHECKING:
+    from infrahub.core.protocols import CoreTransformPython as CoreTransformPythonNode
+    from infrahub.git.models import RepositoryData
+
+
+@task(
+    name="gather-python-transform-attributes",
+    task_run_name="Gather Python transform attributes for {branch_name}",
+    cache_policy=NONE,
+)
+async def gather_python_transform_attributes(
+    db: InfrahubDatabase, branch_name: str, repositories: dict[str, RepositoryData] | None = None
+) -> list[PythonTransformComputedAttribute]:
+    log = get_run_logger()
+    schema_branch = registry.schema.get_schema_branch(name=branch_name)
+    branches_with_diff_from_main = registry.get_altered_schema_branches()
+    branch = registry.get_branch_from_registry(branch=branch_name)
+
+    transform_attributes = schema_branch.computed_attributes.python_attributes_by_transform
+
+    transform_names = list(transform_attributes.keys())
+
+    if not transform_names:
+        return []
+
+    transforms: list[CoreTransformPythonNode] = await NodeManager.query(
+        db=db,
+        schema=InfrahubKind.TRANSFORMPYTHON,
+        branch=branch_name,
+        fields={"id": None, "name": None, "repository": None, "query": None},
+        filters={"name__values": transform_names},
+        prefetch_relationships=True,
+    )
+
+    found_transforms_names = [transform.name.value for transform in transforms]
+    for transform_name in transform_names:
+        if transform_name not in found_transforms_names:
+            log.warning(
+                msg=f"The transform {transform_name} is assigned to a computed attribute but the transform could not be found in the database."
+            )
+    repositories = repositories or await get_repositories_commit_per_branch(db=db)
+    graphql_params = await prepare_graphql_params(db=db, branch=branch)
+
+    computed_attributes: list[PythonTransformComputedAttribute] = []
+    for transform in transforms:
+        repository = await transform.repository.get_peer(db=db, peer_type=CoreGenericRepository, raise_on_error=True)
+        query = await transform.query.get_peer(db=db, peer_type=CoreGraphQLQuery, raise_on_error=True)
+        query_analyzer = InfrahubGraphQLQueryAnalyzer(
+            query=query.query.value,
+            branch=branch,
+            schema_branch=schema_branch,
+            schema=graphql_params.schema,
+        )
+        for attribute in transform_attributes[transform.name.value]:
+            python_transform_computed_attribute = PythonTransformComputedAttribute(
+                name=transform.name.value,
+                branch_name=branch_name,
+                repository_id=repository.get_id(),
+                repository_name=repository.name.value,
+                repository_kind=repository.get_kind(),
+                query_analyzer=query_analyzer,
+                query_name=query.name.value,
+                computed_attribute=attribute,
+                default_schema=branch_name not in branches_with_diff_from_main,
+            )
+            python_transform_computed_attribute.populate_branch_commit(
+                repository_data=repositories.get(repository.name.value)
+            )
+            computed_attributes.append(python_transform_computed_attribute)
+
+    return computed_attributes
+
+
+@task(
+    name="gather-trigger-computed-attribute-jinja2",
+    cache_policy=NONE,
+)
+async def gather_trigger_computed_attribute_jinja2() -> list[ComputedAttrJinja2TriggerDefinition]:
+    log = get_run_logger()
+
+    # Build a list of all branches to process based on which branch is different from main
+    branches_with_diff_from_main = registry.get_altered_schema_branches()
+    branches_to_process: list[tuple[str, list[str]]] = [(branch, []) for branch in branches_with_diff_from_main]
+    branches_to_process.append((registry.default_branch, branches_with_diff_from_main))
+
+    triggers: list[ComputedAttrJinja2TriggerDefinition] = []
+
+    for branch_scope, branches_out_of_scope in branches_to_process:
+        schema_branch = registry.schema.get_schema_branch(name=branch_scope)
+        mapping = schema_branch.computed_attributes.get_jinja2_trigger_nodes()
+
+        log.info(f"Generating {len(mapping)} Jinja2 trigger for {branch_scope} (except {branches_out_of_scope})")
+
+        for computed_attribute, trigger_nodes in mapping.items():
+            for trigger_node in trigger_nodes:
+                trigger = ComputedAttrJinja2TriggerDefinition.from_computed_attribute(
+                    branch=branch_scope,
+                    computed_attribute=computed_attribute,
+                    trigger_node=trigger_node,
+                    branches_out_of_scope=branches_out_of_scope,
+                )
+                triggers.append(trigger)
+
+    return triggers
+
+
+@task(
+    name="gather-trigger-computed-attribute-python",
+    cache_policy=NONE,
+)
+async def gather_trigger_computed_attribute_python(
+    db: InfrahubDatabase,
+) -> tuple[list[ComputedAttrPythonTriggerDefinition], list[ComputedAttrPythonQueryTriggerDefinition]]:
+    triggers_python = []
+    triggers_python_query = []
+
+    repositories = await get_repositories_commit_per_branch(db=db)
+
+    all_computed_attributes: dict[str, dict[str, PythonTransformComputedAttribute]] = defaultdict(dict)
+    for branch in list(registry.branch.values()):
+        if branch.is_global:
+            continue
+
+        computed_attributes = await gather_python_transform_attributes(
+            db=db, branch_name=branch.name, repositories=repositories
+        )
+        for computed_attribute in computed_attributes:
+            all_computed_attributes[computed_attribute.name][branch.name] = computed_attribute
+
+    for branches in all_computed_attributes.values():
+        branches_with_diff_from_main = []
+        if registry.default_branch in branches.keys():
+            commit_main = branches[registry.default_branch].repository_commit
+            branches_with_diff_from_main = [
+                branch_name for branch_name, item in branches.items() if item.repository_commit != commit_main
+            ]
+        else:
+            branches_with_diff_from_main = list(branches.keys())
+
+        branches_to_process: list[tuple[str, list[str]]] = [(branch, []) for branch in branches_with_diff_from_main]
+        branches_to_process.append((registry.default_branch, branches_with_diff_from_main))
+
+        for branch_scope, branches_out_of_scope in branches_to_process:
+            trigger_python = ComputedAttrPythonTriggerDefinition.from_object(
+                computed_attribute=branches[branch_scope],
+                branch=branch_scope,
+                branches_out_of_scope=branches_out_of_scope,
+            )
+            triggers_python.append(trigger_python)
+
+            for kind in branches[branch_scope].query_analyzer.query_report.requested_read.keys():
+                trigger_python_query = ComputedAttrPythonQueryTriggerDefinition.from_object(
+                    kind=kind,
+                    computed_attribute=branches[branch_scope],
+                    branch=branch_scope,
+                    branches_out_of_scope=branches_out_of_scope,
+                )
+                triggers_python_query.append(trigger_python_query)
+
+    return triggers_python, triggers_python_query

infrahub/computed_attribute/models.py

@@ -5,15 +5,34 @@ from dataclasses import dataclass, field
 from typing import TYPE_CHECKING
 
 from prefect.events.schemas.automations import Automation  # noqa: TC002
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, ConfigDict, Field, computed_field
 from typing_extensions import Self
 
+from infrahub.core import registry
+from infrahub.core.schema.schema_branch_computed import (  # noqa: TC001
+    ComputedAttributeTarget,
+    ComputedAttributeTriggerNode,
+    PythonDefinition,
+)
+from infrahub.events import NodeCreatedEvent, NodeUpdatedEvent
+from infrahub.graphql.analyzer import InfrahubGraphQLQueryAnalyzer  # noqa: TC001
+from infrahub.trigger.constants import NAME_SEPARATOR
+from infrahub.trigger.models import (
+    EventTrigger,
+    ExecuteWorkflow,
+    TriggerBranchDefinition,
+    TriggerType,
+)
+from infrahub.workflows.catalogue import (
+    COMPUTED_ATTRIBUTE_PROCESS_JINJA2,
+    COMPUTED_ATTRIBUTE_PROCESS_TRANSFORM,
+    QUERY_COMPUTED_ATTRIBUTE_TRANSFORM_TARGETS,
+)
+
 if TYPE_CHECKING:
     from uuid import UUID
 
-    from infrahub_sdk.data import RepositoryData
-
-    from infrahub.core.schema.schema_branch_computed import PythonDefinition
+    from infrahub.git.models import RepositoryData
 
 
 class ComputedAttributeAutomations(BaseModel):
@@ -26,7 +45,7 @@ class ComputedAttributeAutomations(BaseModel):
             if not automation.name.startswith(prefix):
                 continue
 
-            name_split = automation.name.split("::")
+            name_split = automation.name.split(NAME_SEPARATOR)
             if len(name_split) != 3:
                 continue
 
@@ -47,9 +66,6 @@ class ComputedAttributeAutomations(BaseModel):
             return True
         return False
 
-    def return_obsolete(self, keep: list[UUID]) -> list[UUID]:
-        return [automation_id for automation_id in self.all_automation_ids if automation_id not in keep]
-
     @property
     def all_automation_ids(self) -> list[UUID]:
         automation_ids: list[UUID] = []
@@ -59,25 +75,237 @@ class ComputedAttributeAutomations(BaseModel):
         return automation_ids
 
 
-@dataclass
-class PythonTransformComputedAttribute:
+class PythonTransformComputedAttribute(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
     name: str
     repository_id: str
     repository_name: str
     repository_kind: str
     query_name: str
-    query_models: list[str]
+    query_analyzer: InfrahubGraphQLQueryAnalyzer
     computed_attribute: PythonDefinition
     default_schema: bool
+    branch_name: str
     branch_commit: dict[str, str] = field(default_factory=dict)
 
+    @computed_field
+    def repository_commit(self) -> str:
+        return self.branch_commit[self.branch_name]
+
     def populate_branch_commit(self, repository_data: RepositoryData | None = None) -> None:
         if repository_data:
             for branch, commit in repository_data.branches.items():
                 self.branch_commit[branch] = commit
 
+    def get_altered_branches(self) -> list[str]:
+        if registry.default_branch in self.branch_commit:
+            default_branch_commit = self.branch_commit[registry.default_branch]
+            return [
+                branch_name for branch_name, commit in self.branch_commit.items() if commit != default_branch_commit
+            ]
+        return list(self.branch_commit.keys())
+
 
 @dataclass
 class PythonTransformTarget:
     kind: str
     object_id: str
+
+
+class ComputedAttrJinja2TriggerDefinition(TriggerBranchDefinition):
+    type: TriggerType = TriggerType.COMPUTED_ATTR_JINJA2
+    computed_attribute: ComputedAttributeTarget
+
+    @classmethod
+    def from_computed_attribute(
+        cls,
+        branch: str,
+        computed_attribute: ComputedAttributeTarget,
+        trigger_node: ComputedAttributeTriggerNode,
+        branches_out_of_scope: list[str] | None = None,
+    ) -> Self:
+        """
+        This function is used to create a trigger definition for a computed attribute of type Jinja2.
+        """
+        event_trigger = EventTrigger()
+        event_trigger.events.add(NodeUpdatedEvent.event_name)
+        if computed_attribute.attribute.optional:
+            # If the computed attribute not optional it means that Infrahub will have assigned the value during
+            # the creation of the node so we don't need to match on node creation events as it would only add
+            # extra work that doesn't accomplish anything. For this reason the filter should only match on the
+            # node creation events if the attribute is optional.
+            event_trigger.events.add(NodeCreatedEvent.event_name)
+
+        event_trigger.match = {"infrahub.node.kind": trigger_node.kind}
+        if branches_out_of_scope:
+            event_trigger.match["infrahub.branch.name"] = [f"!{branch}" for branch in branches_out_of_scope]
+        elif not branches_out_of_scope and branch != registry.default_branch:
+            event_trigger.match["infrahub.branch.name"] = branch
+
+        event_trigger.match_related = {
+            "prefect.resource.role": ["infrahub.node.attribute_update", "infrahub.node.relationship_update"],
+            "infrahub.field.name": trigger_node.fields,
+        }
+
+        workflow = ExecuteWorkflow(
+            workflow=COMPUTED_ATTRIBUTE_PROCESS_JINJA2,
+            parameters={
+                "branch_name": "{{ event.resource['infrahub.branch.name'] }}",
+                "node_kind": "{{ event.resource['infrahub.node.kind'] }}",
+                "object_id": "{{ event.resource['infrahub.node.id'] }}",
+                "computed_attribute_name": computed_attribute.attribute.name,
+                "computed_attribute_kind": computed_attribute.kind,
+                "updated_fields": {
+                    "__prefect_kind": "json",
+                    "value": {
+                        "__prefect_kind": "jinja",
+                        "template": "{{ event.payload['data']['fields'] | tojson }}",
+                    },
+                },
+                "context": {
+                    "__prefect_kind": "json",
+                    "value": {
+                        "__prefect_kind": "jinja",
+                        "template": "{{ event.payload['context'] | tojson }}",
+                    },
+                },
+            },
+        )
+
+        definition = cls(
+            name=f"{computed_attribute.key_name}{NAME_SEPARATOR}kind{NAME_SEPARATOR}{trigger_node.kind}",
+            branch=branch,
+            computed_attribute=computed_attribute,
+            trigger=event_trigger,
+            actions=[workflow],
+        )
+
+        return definition
+
+
+class ComputedAttrPythonTriggerDefinition(TriggerBranchDefinition):
+    type: TriggerType = TriggerType.COMPUTED_ATTR_PYTHON
+    computed_attribute: PythonTransformComputedAttribute
+
+    @classmethod
+    def from_object(
+        cls,
+        branch: str,
+        computed_attribute: PythonTransformComputedAttribute,
+        branches_out_of_scope: list[str] | None = None,
+    ) -> Self:
+        # scope = registry.default_branch
+
+        event_trigger = EventTrigger()
+        event_trigger.events.update({NodeCreatedEvent.event_name, NodeUpdatedEvent.event_name})
+        event_trigger.match = {
+            "infrahub.node.kind": [computed_attribute.computed_attribute.kind],
+        }
+
+        if branches_out_of_scope:
+            event_trigger.match["infrahub.branch.name"] = [f"!{branch}" for branch in branches_out_of_scope]
+        elif not branches_out_of_scope and branch != registry.default_branch:
+            event_trigger.match["infrahub.branch.name"] = branch
+
+        update_fields = computed_attribute.query_analyzer.query_report.fields_by_kind(
+            kind=computed_attribute.computed_attribute.kind
+        )
+        event_trigger.match_related = {
+            "prefect.resource.role": ["infrahub.node.attribute_update", "infrahub.node.relationship_update"],
+        }
+
+        if update_fields and "display_label" not in update_fields:
+            # The GraphQLQuery analyzer doesn't yet support figuring out which updates would match the "display label"
+            # of a query. Because of this we temporarily match any field if the display_label is part of the computed
+            # attribute query
+            event_trigger.match_related["infrahub.field.name"] = update_fields
+
+        definition = cls(
+            name=computed_attribute.computed_attribute.key_name,
+            branch=branch,
+            computed_attribute=computed_attribute,
+            trigger=event_trigger,
+            actions=[
+                ExecuteWorkflow(
+                    workflow=COMPUTED_ATTRIBUTE_PROCESS_TRANSFORM,
+                    parameters={
+                        "branch_name": "{{ event.resource['infrahub.branch.name'] }}",
+                        "node_kind": "{{ event.resource['infrahub.node.kind'] }}",
+                        "object_id": "{{ event.resource['infrahub.node.id'] }}",
+                        "computed_attribute_name": computed_attribute.computed_attribute.attribute.name,
+                        "computed_attribute_kind": computed_attribute.computed_attribute.kind,
+                        "context": {
+                            "__prefect_kind": "json",
+                            "value": {
+                                "__prefect_kind": "jinja",
+                                "template": "{{ event.payload['context'] | tojson }}",
+                            },
+                        },
+                    },
+                ),
+            ],
+        )
+
+        return definition
+
+
+class ComputedAttrPythonQueryTriggerDefinition(TriggerBranchDefinition):
+    type: TriggerType = TriggerType.COMPUTED_ATTR_PYTHON_QUERY
+
+    @classmethod
+    def from_object(
+        cls,
+        branch: str,
+        kind: str,
+        computed_attribute: PythonTransformComputedAttribute,
+        branches_out_of_scope: list[str] | None = None,
+    ) -> Self:
+        # Only matching on node updated events, before nodes are created they won't be a member of the GraphQL query
+        # group regardless so it doesn't make sense to trigger the query on node creation. For the initial object
+        # where the computed attribute belongs that to will need to be created first which will trigger its own initial
+        # process to populated the original members of the GraphQL query group
+        event_trigger = EventTrigger(events={NodeUpdatedEvent.event_name})
+        event_trigger.match = {
+            "infrahub.node.kind": kind,
+        }
+
+        update_fields = computed_attribute.query_analyzer.query_report.fields_by_kind(kind=kind)
+        event_trigger.match_related = {
+            "prefect.resource.role": ["infrahub.node.attribute_update", "infrahub.node.relationship_update"],
+        }
+
+        if update_fields and "display_label" not in update_fields:
+            # The GraphQLQuery analyzer doesn't yet support figuring out which updates would match the "display label"
+            # of a query. Because of this we temporarily match any field if the display_label is part of the computed
+            # attribute query
+            event_trigger.match_related["infrahub.field.name"] = update_fields
+
+        if branches_out_of_scope:
+            event_trigger.match["infrahub.branch.name"] = [f"!{branch}" for branch in branches_out_of_scope]
+        elif not branches_out_of_scope and branch != registry.default_branch:
+            event_trigger.match["infrahub.branch.name"] = branch
+
+        definition = cls(
+            name=f"{computed_attribute.computed_attribute.key_name}{NAME_SEPARATOR}kind{NAME_SEPARATOR}{kind}",
+            branch=branch,
+            trigger=event_trigger,
+            actions=[
+                ExecuteWorkflow(
+                    workflow=QUERY_COMPUTED_ATTRIBUTE_TRANSFORM_TARGETS,
+                    parameters={
+                        "branch_name": "{{ event.resource['infrahub.branch.name'] }}",
+                        "node_kind": "{{ event.resource['infrahub.node.kind'] }}",
+                        "object_id": "{{ event.resource['infrahub.node.id'] }}",
+                        "context": {
+                            "__prefect_kind": "json",
+                            "value": {
+                                "__prefect_kind": "jinja",
+                                "template": "{{ event.payload['context'] | tojson }}",
+                            },
+                        },
+                    },
+                ),
+            ],
+        )
+
+        return definition