Setup automation for Python based computed attributes

Author: ogenstad

backend/infrahub/api/schema.py

@@ -35,7 +35,12 @@
 from infrahub.services import services
 from infrahub.types import ATTRIBUTE_PYTHON_TYPES
 from infrahub.worker import WORKER_IDENTITY
-from infrahub.workflows.catalogue import COMPUTED_ATTRIBUTE_SETUP, SCHEMA_APPLY_MIGRATION, SCHEMA_VALIDATE_MIGRATION
+from infrahub.workflows.catalogue import (
+    COMPUTED_ATTRIBUTE_SETUP,
+    COMPUTED_ATTRIBUTE_SETUP_PYTHON,
+    SCHEMA_APPLY_MIGRATION,
+    SCHEMA_VALIDATE_MIGRATION,
+)
 
 if TYPE_CHECKING:
     from typing_extensions import Self
@@ -350,6 +355,7 @@ async def load_schema(
 
     # Temporary workaround, we need to emit a proper event
     await service.workflow.submit_workflow(workflow=COMPUTED_ATTRIBUTE_SETUP)
+    await service.workflow.submit_workflow(workflow=COMPUTED_ATTRIBUTE_SETUP_PYTHON)
 
     return SchemaUpdate(hash=updated_hash, previous_hash=original_hash, diff=result.diff)
 

backend/infrahub/computed_attribute/models.py

@@ -1,11 +1,18 @@
+from __future__ import annotations
+
 from collections import defaultdict
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
 
-from prefect.events.schemas.automations import Automation
+from prefect.events.schemas.automations import Automation  # noqa: TCH002
 from pydantic import BaseModel, Field
 from typing_extensions import Self
 
 from .constants import AUTOMATION_NAME_PREFIX
 
+if TYPE_CHECKING:
+    from infrahub.core.schema.schema_branch_computed import PythonDefinition
+
 
 class ComputedAttributeAutomations(BaseModel):
     data: dict[str, dict[str, Automation]] = Field(default_factory=lambda: defaultdict(dict))
@@ -37,3 +44,14 @@ def has(self, identifier: str, scope: str) -> bool:
         if identifier in self.data and scope in self.data[identifier]:
             return True
         return False
+
+
+@dataclass
+class PythonTransformComputedAttribute:
+    name: str
+    repository_id: str
+    repository_name: str
+    repository_kind: str
+    query_name: str
+    query_models: list[str]
+    computed_attribute: PythonDefinition

backend/infrahub/computed_attribute/tasks.py

@@ -21,7 +21,7 @@
 from infrahub.workflows.utils import add_branch_tag
 
 from .constants import AUTOMATION_NAME, AUTOMATION_NAME_PREFIX
-from .models import ComputedAttributeAutomations
+from .models import ComputedAttributeAutomations, PythonTransformComputedAttribute
 
 if TYPE_CHECKING:
     from infrahub.core.schema.computed_attribute import ComputedAttribute
@@ -251,3 +251,109 @@ async def computed_attribute_setup() -> None:
             else:
                 await client.create_automation(automation=automation)
                 log.info(f"{computed_attribute.key_name} Created")
+
+
+@flow(
+    name="computed-attribute-setup-python",
+    flow_run_name="Setup computed attributes for Python transforms in task-manager",
+)
+async def computed_attribute_setup_python() -> None:
+    service = services.service
+    schema_branch = registry.schema.get_schema_branch(name=registry.default_branch)
+    log = get_run_logger()
+
+    transform_attributes = schema_branch.computed_attributes.python_attributes_by_transform
+
+    transform_names = list(transform_attributes.keys())
+
+    transforms = await service.client.filters(
+        kind="CoreTransformPython",
+        branch=registry.default_branch,
+        prefetch_relationships=True,
+        populate_store=True,
+        name__values=transform_names,
+    )
+
+    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."
+            )
+
+    computed_attributes: list[PythonTransformComputedAttribute] = []
+    for transform in transforms:
+        for attribute in transform_attributes[transform.name.value]:
+            computed_attributes.append(
+                PythonTransformComputedAttribute(
+                    name=transform.name.value,
+                    repository_id=transform.repository.peer.id,
+                    repository_name=transform.repository.peer.name.value,
+                    repository_kind=transform.repository.peer.typename,
+                    query_name=transform.query.peer.name.value,
+                    query_models=transform.query.peer.models.value,
+                    computed_attribute=attribute,
+                )
+            )
+
+    async with get_client(sync_client=False) as client:
+        deployments = {
+            item.name: item
+            for item in await client.read_deployments(
+                deployment_filter=DeploymentFilter(
+                    name=DeploymentFilterName(any_=[UPDATE_COMPUTED_ATTRIBUTE_TRANSFORM.name])
+                )
+            )
+        }
+        if UPDATE_COMPUTED_ATTRIBUTE_TRANSFORM.name not in deployments:
+            raise ValueError("Unable to find the deployment for UPDATE_COMPUTED_ATTRIBUTE_TRANSFORM")
+
+        deployment_id_python = deployments[UPDATE_COMPUTED_ATTRIBUTE_TRANSFORM.name].id
+
+        automations = await client.read_automations()
+        existing_computed_attr_automations = ComputedAttributeAutomations.from_prefect(automations=automations)
+
+        for computed_attribute in computed_attributes:
+            log.info(f"processing {computed_attribute.computed_attribute.key_name}")
+            scope = "default"
+
+            automation = AutomationCore(
+                name=AUTOMATION_NAME.format(
+                    prefix=AUTOMATION_NAME_PREFIX,
+                    identifier=computed_attribute.computed_attribute.key_name,
+                    scope=scope,
+                ),
+                description=f"Process value of the computed attribute for {computed_attribute.computed_attribute.key_name} [{scope}]",
+                enabled=True,
+                trigger=EventTrigger(
+                    posture=Posture.Reactive,
+                    expect={"infrahub.node.*"},
+                    within=timedelta(0),
+                    match=ResourceSpecification({"infrahub.node.kind": [computed_attribute.computed_attribute.kind]}),
+                    threshold=1,
+                ),
+                actions=[
+                    RunDeployment(
+                        source="selected",
+                        deployment_id=deployment_id_python,
+                        parameters={
+                            "branch_name": "{{ event.resource['infrahub.branch.name'] }}",
+                            "node_kind": "{{ event.resource['infrahub.node.kind'] }}",
+                            "object_id": "{{ event.resource['infrahub.node.id'] }}",
+                        },
+                        job_variables={},
+                    )
+                ],
+            )
+
+            if existing_computed_attr_automations.has(
+                identifier=computed_attribute.computed_attribute.key_name, scope=scope
+            ):
+                existing = existing_computed_attr_automations.get(
+                    identifier=computed_attribute.computed_attribute.key_name, scope=scope
+                )
+                await client.update_automation(automation_id=existing.id, automation=automation)
+                log.info(f"{computed_attribute.computed_attribute.key_name} Updated")
+            else:
+                await client.create_automation(automation=automation)
+                log.info(f"{computed_attribute.computed_attribute.key_name} Created")

backend/infrahub/core/schema/schema_branch.py

@@ -54,6 +54,7 @@
 from infrahub.visuals import select_color
 
 from .constants import INTERNAL_SCHEMA_NODE_KINDS, SchemaNamespace
+from .schema_branch_computed import ComputedAttributes
 
 log = get_logger()
 
@@ -114,6 +115,7 @@ def __init__(self, cache: dict, name: str | None = None, data: dict[str, dict[st
         self.nodes: dict[str, str] = {}
         self.generics: dict[str, str] = {}
         self.profiles: dict[str, str] = {}
+        self.computed_attributes = ComputedAttributes()
         self._computed_jinja2_attribute_map: dict[str, RegisteredNodeComputedAttribute] = {}
 
         if data:
@@ -949,6 +951,7 @@ def validate_kinds(self) -> None:
                     ) from None
 
     def validate_computed_attributes(self) -> None:
+        self.computed_attributes = ComputedAttributes()
         self._computed_jinja2_attribute_map = {}
         for name in self.nodes.keys():
             node_schema = self.get_node(name=name, duplicate=False)
@@ -1016,11 +1019,14 @@ def _validate_computed_attribute(self, node: NodeSchema, attribute: AttributeSch
 
                 self._register_computed_attribute_target(node=node, attribute=attribute, schema_path=schema_path)
 
-        if attribute.computed_attribute.kind == ComputedAttributeKind.TRANSFORM_PYTHON and not attribute.optional:
+        elif attribute.computed_attribute.kind == ComputedAttributeKind.TRANSFORM_PYTHON and not attribute.optional:
             raise ValueError(
                 f"{node.kind}: Attribute {attribute.name!r} is a computed transform, it can't be mandatory"
             )
 
+        elif attribute.computed_attribute.kind == ComputedAttributeKind.TRANSFORM_PYTHON:
+            self.computed_attributes.add_python_attribute(node=node, attribute=attribute)
+
     def get_impacted_macros(self, kind: str, updates: list[str] | None = None) -> list[ComputedAttributeTarget]:
         if mapping := self._computed_jinja2_attribute_map.get(kind):
             return mapping.get_targets(updates=updates)

backend/infrahub/core/schema/schema_branch_computed.py

@@ -0,0 +1,45 @@
+from dataclasses import dataclass
+
+from infrahub.core.schema import (
+    AttributeSchema,
+    NodeSchema,
+)
+
+
+@dataclass
+class PythonDefinition:
+    kind: str
+    attribute: AttributeSchema
+
+    @property
+    def key_name(self) -> str:
+        return f"{self.kind}_{self.attribute.name}"
+
+
+class ComputedAttributes:
+    def __init__(self) -> None:
+        self._computed_python_transform_attribute_map: dict[str, list[AttributeSchema]] = {}
+
+    def add_python_attribute(self, node: NodeSchema, attribute: AttributeSchema) -> None:
+        if node.kind not in self._computed_python_transform_attribute_map:
+            self._computed_python_transform_attribute_map[node.kind] = []
+        self._computed_python_transform_attribute_map[node.kind].append(attribute)
+
+    def get_kinds_python_attributes(self) -> list[str]:
+        """Return kinds that have Python attributes defined"""
+        return list(self._computed_python_transform_attribute_map.keys())
+
+    @property
+    def python_attributes_by_transform(self) -> dict[str, list[PythonDefinition]]:
+        computed_attributes: dict[str, list[PythonDefinition]] = {}
+        for kind, attributes in self._computed_python_transform_attribute_map.items():
+            for attribute in attributes:
+                if attribute.computed_attribute and attribute.computed_attribute.transform:
+                    if attribute.computed_attribute.transform not in computed_attributes:
+                        computed_attributes[attribute.computed_attribute.transform] = []
+
+                    computed_attributes[attribute.computed_attribute.transform].append(
+                        PythonDefinition(kind=kind, attribute=attribute)
+                    )
+
+        return computed_attributes

backend/infrahub/graphql/mutations/main.py

@@ -27,7 +27,6 @@
 from infrahub.exceptions import ValidationError
 from infrahub.log import get_log_data, get_logger
 from infrahub.worker import WORKER_IDENTITY
-from infrahub.workflows.catalogue import UPDATE_COMPUTED_ATTRIBUTE_TRANSFORM
 
 from .node_getter.by_default_filter import MutationNodeGetterByDefaultFilter
 from .node_getter.by_hfid import MutationNodeGetterByHfid
@@ -101,8 +100,6 @@ async def mutate(cls, root: dict, info: GraphQLResolveInfo, data: InputObjectTyp
         # Reset the time of the query to guarantee that all resolvers executed after this point will account for the changes
         context.at = Timestamp()
 
-        # Get relevant macros based on the current change
-        # macros = schema_branch.get_impacted_macros(kind=obj.get_kind(), updates=updated_fields)
         if config.SETTINGS.broker.enable and context.background:
             log_data = get_log_data()
             request_id = log_data.get("request_id", "")
@@ -120,38 +117,6 @@ async def mutate(cls, root: dict, info: GraphQLResolveInfo, data: InputObjectTyp
 
             context.background.add_task(context.service.event.send, event)
 
-            # Temporary workaround until there's a proper deployment for Transforms during schema load / repo sync
-            schema_branch = registry.schema.get_schema_branch(name=context.branch.name)
-            node_schema = schema_branch.get(obj._schema.kind)
-            if isinstance(node_schema, NodeSchema):
-                has_computed_attributes = [
-                    attribute for attribute in node_schema.attributes if attribute.computed_attribute
-                ]
-                if has_computed_attributes:
-                    updated_fields = list(data.keys())
-
-                    await context.service.workflow.submit_workflow(
-                        workflow=UPDATE_COMPUTED_ATTRIBUTE_TRANSFORM,
-                        parameters={
-                            "branch_name": context.branch.name,
-                            "node_kind": obj.get_kind(),
-                            "object_id": obj.get_id(),
-                            "updated_fields": updated_fields,
-                        },
-                    )
-
-            # # Add event
-            # if macros:
-            #     await context.service.workflow.submit_workflow(
-            #         workflow=PROCESS_COMPUTED_MACRO,
-            #         parameters={
-            #             "branch_name": context.branch.name,
-            #             "node_kind": obj.get_kind(),
-            #             "object_id": obj.get_id(),
-            #             "updated_fields": updated_fields,
-            #         },
-            #     )
-
         return mutation
 
     @classmethod

backend/infrahub/workflows/catalogue.py

@@ -218,6 +218,14 @@
     function="computed_attribute_setup",
 )
 
+COMPUTED_ATTRIBUTE_SETUP_PYTHON = WorkflowDefinition(
+    name="computed-attribute-setup-python",
+    type=WorkflowType.INTERNAL,
+    module="infrahub.computed_attribute.tasks",
+    function="computed_attribute_setup_python",
+)
+
+
 UPDATE_COMPUTED_ATTRIBUTE_TRANSFORM = WorkflowDefinition(
     name="process_computed_attribute_transform",
     type=WorkflowType.INTERNAL,
@@ -265,6 +273,7 @@
     GIT_REPOSITORY_ADD_READ_ONLY,
     PROCESS_COMPUTED_MACRO,
     COMPUTED_ATTRIBUTE_SETUP,
+    COMPUTED_ATTRIBUTE_SETUP_PYTHON,
     UPDATE_COMPUTED_ATTRIBUTE_TRANSFORM,
     REQUEST_PROPOSED_CHANGE_DATA_INTEGRITY,
 ]