Refactor computed_attribute_setup

Author: dgarros

backend/infrahub/computed_attribute/constants.py

@@ -0,0 +1,3 @@
+AUTOMATION_NAME_PREFIX = "Computed-attribute-process"
+
+AUTOMATION_NAME = AUTOMATION_NAME_PREFIX + "::{identifier}::{scope}"

backend/infrahub/computed_attribute/models.py

@@ -0,0 +1,39 @@
+from collections import defaultdict
+
+from prefect.events.schemas.automations import Automation
+from pydantic import BaseModel, Field
+from typing_extensions import Self
+
+from .constants import AUTOMATION_NAME_PREFIX
+
+
+class ComputedAttributeAutomations(BaseModel):
+    data: dict[str, dict[str, Automation]] = Field(default_factory=lambda: defaultdict(dict))
+
+    @classmethod
+    def from_prefect(cls, automations: list[Automation]) -> Self:
+        obj = cls()
+        for automation in automations:
+            if not automation.name.startswith(AUTOMATION_NAME_PREFIX):
+                continue
+
+            name_split = automation.name.split("::")
+            if len(name_split) != 3:
+                continue
+
+            identifier = name_split[1]
+            scope = name_split[2]
+
+            obj.data[identifier][scope] = automation
+
+        return obj
+
+    def get(self, identifier: str, scope: str) -> Automation:
+        if identifier in self.data and scope in self.data[identifier]:
+            return self.data[identifier][scope]
+        raise KeyError(f"Unable to find an automation for {identifier} {scope}")
+
+    def has(self, identifier: str, scope: str) -> bool:
+        if identifier in self.data and scope in self.data[identifier]:
+            return True
+        return False

backend/infrahub/computed_attribute/tasks.py

@@ -17,9 +17,12 @@
 from infrahub.git.repository import get_initialized_repo
 from infrahub.services import services
 from infrahub.support.macro import MacroDefinition
-from infrahub.workflows.catalogue import PROCESS_COMPUTED_MACRO
+from infrahub.workflows.catalogue import PROCESS_COMPUTED_MACRO, UPDATE_COMPUTED_ATTRIBUTE_TRANSFORM
 from infrahub.workflows.utils import add_branch_tag
 
+from .constants import AUTOMATION_NAME, AUTOMATION_NAME_PREFIX
+from .models import ComputedAttributeAutomations
+
 if TYPE_CHECKING:
     from infrahub.core.schema.computed_attribute import ComputedAttribute
     from infrahub.core.schema.schema_branch import ComputedAttributeTarget
@@ -175,24 +178,33 @@ async def process_jinja2(
             print()
 
 
-@flow(name="computed-attribute-setup")
+@flow(name="computed-attribute-setup", flow_run_name="Setup computed attributes in task-manager")
 async def computed_attribute_setup() -> None:
     # service = services.service
     schema_branch = registry.schema.get_schema_branch(name=registry.default_branch)
     log = get_run_logger()
+
     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_=[PROCESS_COMPUTED_MACRO.name]))
+                deployment_filter=DeploymentFilter(
+                    name=DeploymentFilterName(
+                        any_=[PROCESS_COMPUTED_MACRO.name, UPDATE_COMPUTED_ATTRIBUTE_TRANSFORM.name]
+                    )
+                )
             )
         }
         if PROCESS_COMPUTED_MACRO.name not in deployments:
             raise ValueError("Unable to find the deployment for PROCESS_COMPUTED_MACRO")
-        deployment_id = deployments[PROCESS_COMPUTED_MACRO.name].id
+        if UPDATE_COMPUTED_ATTRIBUTE_TRANSFORM.name not in deployments:
+            raise ValueError("Unable to find the deployment for UPDATE_COMPUTED_ATTRIBUTE_TRANSFORM")
 
-        # TODO need to pull the existing automation to see if we need to create or update each object
-        # automations = await client.read_automations()
+        deployment_id_jinja = deployments[PROCESS_COMPUTED_MACRO.name].id
+        # deployment_id_python = deployments[UPDATE_COMPUTED_ATTRIBUTE_TRANSFORM.name].id
+
+        automations = await client.read_automations()
+        existing_computed_attr_automations = ComputedAttributeAutomations.from_prefect(automations=automations)
 
         computed_attributes: dict[str, ComputedAttributeTarget] = {}
         for item in schema_branch._computed_jinja2_attribute_map.values():
@@ -205,9 +217,11 @@ async def computed_attribute_setup() -> None:
 
         for identifier, computed_attribute in computed_attributes.items():
             log.info(f"processing {computed_attribute.key_name}")
+            scope = "default"
+
             automation = AutomationCore(
-                name=f"computed-attribute-process-{identifier}",
-                description=f"Process value of the computed attribute for {identifier}",
+                name=AUTOMATION_NAME.format(prefix=AUTOMATION_NAME_PREFIX, identifier=identifier, scope=scope),
+                description=f"Process value of the computed attribute for {identifier} [{scope}]",
                 enabled=True,
                 trigger=EventTrigger(
                     posture=Posture.Reactive,
@@ -219,7 +233,7 @@ async def computed_attribute_setup() -> None:
                 actions=[
                     RunDeployment(
                         source="selected",
-                        deployment_id=deployment_id,
+                        deployment_id=deployment_id_jinja,
                         parameters={
                             "branch_name": "{{ event.resource['infrahub.branch.name'] }}",
                             "node_kind": "{{ event.resource['infrahub.node.kind'] }}",
@@ -230,5 +244,10 @@ async def computed_attribute_setup() -> None:
                 ],
             )
 
-            response = await client.create_automation(automation=automation)
-            log.info(f"Processed: {computed_attribute.key_name} : {response}")
+            if existing_computed_attr_automations.has(identifier=identifier, scope=scope):
+                existing = existing_computed_attr_automations.get(identifier=identifier, scope=scope)
+                await client.update_automation(automation_id=existing.id, automation=automation)
+                log.info(f"{computed_attribute.key_name} Updated")
+            else:
+                await client.create_automation(automation=automation)
+                log.info(f"{computed_attribute.key_name} Created")

backend/tests/unit/computed_attribute/__init__.py

@@ -0,0 +1,42 @@
+import uuid
+from datetime import timedelta
+
+from prefect.events.schemas.automations import Automation, EventTrigger, Posture
+
+from infrahub.computed_attribute.constants import AUTOMATION_NAME
+from infrahub.computed_attribute.models import ComputedAttributeAutomations
+
+
+def generate_automation(
+    name: str, description: str = "", trigger: EventTrigger | None = None, actions: list | None = None
+) -> Automation:
+    default_trigger = EventTrigger(
+        posture=Posture.Reactive,
+        expect={"infrahub.node.*"},
+        within=timedelta(0),
+        threshold=1,
+    )
+
+    return Automation(
+        id=uuid.uuid4(),
+        name=name,
+        description=description,
+        enabled=True,
+        trigger=trigger or default_trigger,
+        actions=actions or [],
+    )
+
+
+async def test_load_from_prefect():
+    automations: list[Automation] = [
+        generate_automation(name=AUTOMATION_NAME.format(identifier="AAAAA", scope="default")),
+        generate_automation(name=AUTOMATION_NAME.format(identifier="AAAAA", scope="yyyy")),
+        generate_automation(name=AUTOMATION_NAME.format(identifier="BBBBB", scope="default")),
+        generate_automation(name="anothername"),
+    ]
+
+    obj = ComputedAttributeAutomations.from_prefect(automations=automations)
+
+    assert obj.has(identifier="AAAAA", scope="default")
+    assert obj.has(identifier="AAAAA", scope="yyyy")
+    assert obj.has(identifier="BBBBB", scope="default")