Merge pull request #6348 from opsmill/pog-refactor-computed-attribute-filter-query

Refactor computed attribute update_computed_attribute_value_jinja2

Author: ogenstad

backend/infrahub/computed_attribute/models.py

@@ -2,13 +2,16 @@
 
 from collections import defaultdict
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
+from infrahub_sdk.graphql import Query
 from prefect.events.schemas.automations import Automation  # noqa: TC002
 from pydantic import BaseModel, ConfigDict, Field, computed_field
 from typing_extensions import Self
 
 from infrahub.core import registry
+from infrahub.core.constants import RelationshipCardinality
+from infrahub.core.schema import AttributeSchema, NodeSchema  # noqa: TC001
 from infrahub.core.schema.schema_branch_computed import (  # noqa: TC001
     ComputedAttributeTarget,
     ComputedAttributeTriggerNode,
@@ -309,3 +312,80 @@ def from_object(
         )
 
         return definition
+
+
+class ComputedAttrJinja2GraphQLResponse(BaseModel):
+    node_id: str
+    computed_attribute_value: str | None
+    variables: dict[str, Any] = Field(default_factory=dict)
+
+
+class ComputedAttrJinja2GraphQL(BaseModel):
+    node_schema: NodeSchema = Field(..., description="The node kind where the computed attribute is defined")
+    attribute_schema: AttributeSchema = Field(..., description="The computed attribute")
+    variables: list[str] = Field(..., description="The list of variable names used within the computed attribute")
+
+    def render_graphql_query(self, query_filter: str, filter_id: str) -> str:
+        query_fields = self.query_fields
+        query_fields["id"] = None
+        query_fields[self.attribute_schema.name] = {"value": None}
+        query = Query(
+            name="ComputedAttributeFilter",
+            query={
+                self.node_schema.kind: {
+                    "@filters": {query_filter: filter_id},
+                    "edges": {"node": query_fields},
+                }
+            },
+        )
+
+        return query.render()
+
+    @property
+    def query_fields(self) -> dict[str, Any]:
+        output: dict[str, Any] = {}
+        for variable in self.variables:
+            field_name, remainder = variable.split("__", maxsplit=1)
+            if field_name in self.node_schema.attribute_names:
+                output[field_name] = {remainder: None}
+            elif field_name in self.node_schema.relationship_names:
+                related_attribute, related_value = remainder.split("__", maxsplit=1)
+                relationship = self.node_schema.get_relationship(name=field_name)
+                if relationship.cardinality == RelationshipCardinality.ONE:
+                    if field_name not in output:
+                        output[field_name] = {"node": {}}
+                    output[field_name]["node"][related_attribute] = {related_value: None}
+        return output
+
+    def parse_response(self, response: dict[str, Any]) -> list[ComputedAttrJinja2GraphQLResponse]:
+        rendered_response: list[ComputedAttrJinja2GraphQLResponse] = []
+        if kind_payload := response.get(self.node_schema.kind):
+            edges = kind_payload.get("edges", [])
+            for node in edges:
+                if node_response := self.to_node_response(node_dict=node):
+                    rendered_response.append(node_response)
+        return rendered_response
+
+    def to_node_response(self, node_dict: dict[str, Any]) -> ComputedAttrJinja2GraphQLResponse | None:
+        if node := node_dict.get("node"):
+            node_id = node.get("id")
+        else:
+            return None
+
+        computed_attribute = node.get(self.attribute_schema.name, {}).get("value")
+        response = ComputedAttrJinja2GraphQLResponse(node_id=node_id, computed_attribute_value=computed_attribute)
+        for variable in self.variables:
+            field_name, remainder = variable.split("__", maxsplit=1)
+            response.variables[variable] = None
+            if field_content := node.get(field_name):
+                if field_name in self.node_schema.attribute_names:
+                    response.variables[variable] = field_content.get(remainder)
+                elif field_name in self.node_schema.relationship_names:
+                    relationship = self.node_schema.get_relationship(name=field_name)
+                    if relationship.cardinality == RelationshipCardinality.ONE:
+                        related_attribute, related_value = remainder.split("__", maxsplit=1)
+                        node_content = field_content.get("node") or {}
+                        related_attribute_content = node_content.get(related_attribute) or {}
+                        response.variables[variable] = related_attribute_content.get(related_value)
+
+        return response

backend/infrahub/computed_attribute/tasks.py

@@ -2,10 +2,7 @@
 
 from typing import TYPE_CHECKING
 
-from infrahub_sdk.protocols import (
-    CoreNode,  # noqa: TC002
-    CoreTransformPython,
-)
+from infrahub_sdk.protocols import CoreTransformPython
 from infrahub_sdk.template import Jinja2Template
 from prefect import flow
 from prefect.client.orchestration import get_client
@@ -28,9 +25,7 @@
 from infrahub.workflows.utils import add_tags, wait_for_schema_to_converge
 
 from .gather import gather_trigger_computed_attribute_jinja2, gather_trigger_computed_attribute_python
-from .models import (
-    PythonTransformTarget,
-)
+from .models import ComputedAttrJinja2GraphQL, ComputedAttrJinja2GraphQLResponse, PythonTransformTarget
 
 if TYPE_CHECKING:
     from infrahub.core.schema.computed_attribute import ComputedAttribute
@@ -167,49 +162,33 @@ async def trigger_update_python_computed_attributes(
     flow_run_name="Update value for computed attribute {attribute_name}",
 )
 async def update_computed_attribute_value_jinja2(
-    branch_name: str, obj: CoreNode, attribute_name: str, template_value: str, service: InfrahubServices
+    branch_name: str,
+    obj: ComputedAttrJinja2GraphQLResponse,
+    node_kind: str,
+    attribute_name: str,
+    template: Jinja2Template,
+    service: InfrahubServices,
 ) -> None:
     log = get_run_logger()
 
-    await add_tags(branches=[branch_name], nodes=[obj.id], db_change=True)
-
-    jinja_template = Jinja2Template(template=template_value)
-    variables = {}
-    for variable in jinja_template.get_variables():
-        components = variable.split("__")
-        if len(components) == 2:
-            property_name = components[0]
-            property_value = components[1]
-            attribute_property = getattr(obj, property_name)
-            variables[variable] = getattr(attribute_property, property_value)
-        elif len(components) == 3:
-            relationship_name = components[0]
-            property_name = components[1]
-            property_value = components[2]
-            relationship = getattr(obj, relationship_name)
-            try:
-                attribute_property = getattr(relationship.peer, property_name)
-                variables[variable] = getattr(attribute_property, property_value)
-            except ValueError:
-                variables[variable] = ""
-
-    value = await jinja_template.render(variables=variables)
-    existing_value = getattr(obj, attribute_name).value
-    if value == existing_value:
+    await add_tags(branches=[branch_name], nodes=[obj.node_id], db_change=True)
+
+    value = await template.render(variables=obj.variables)
+    if value == obj.computed_attribute_value:
         log.debug(f"Ignoring to update {obj} with existing value on {attribute_name}={value}")
         return
 
     await service.client.execute_graphql(
         query=UPDATE_ATTRIBUTE,
         variables={
-            "id": obj.id,
-            "kind": obj.get_kind(),
+            "id": obj.node_id,
+            "kind": node_kind,
             "attribute": attribute_name,
             "value": value,
         },
         branch_name=branch_name,
     )
-    log.info(f"Updating computed attribute {obj.get_kind()}.{attribute_name}='{value}' ({obj.id})")
+    log.info(f"Updating computed attribute {node_kind}.{attribute_name}='{value}' ({obj.node_id})")
 
 
 @flow(
@@ -235,41 +214,43 @@ async def process_jinja2(
         branch_name if branch_name in registry.get_altered_schema_branches() else registry.default_branch
     )
     schema_branch = registry.schema.get_schema_branch(name=target_branch_schema)
-    await service.client.schema.all(branch=branch_name, refresh=True, schema_hash=schema_branch.get_hash())
-
+    node_schema = schema_branch.get_node(name=computed_attribute_kind, duplicate=False)
     computed_macros = [
         attrib
         for attrib in schema_branch.computed_attributes.get_impacted_jinja2_targets(kind=node_kind, updates=updates)
         if attrib.kind == computed_attribute_kind and attrib.attribute.name == computed_attribute_name
     ]
     for computed_macro in computed_macros:
-        found: list[CoreNode] = []
+        found: list[ComputedAttrJinja2GraphQLResponse] = []
+        template_string = "n/a"
+        if computed_macro.attribute.computed_attribute and computed_macro.attribute.computed_attribute.jinja2_template:
+            template_string = computed_macro.attribute.computed_attribute.jinja2_template
+
+        jinja_template = Jinja2Template(template=template_string)
+        variables = jinja_template.get_variables()
+
+        attribute_graphql = ComputedAttrJinja2GraphQL(
+            node_schema=node_schema, attribute_schema=computed_macro.attribute, variables=variables
+        )
+
         for id_filter in computed_macro.node_filters:
-            filters = {id_filter: object_id}
-            nodes: list[CoreNode] = await service.client.filters(
-                kind=computed_macro.kind,
-                branch=branch_name,
-                prefetch_relationships=True,
-                populate_store=True,
-                **filters,
-            )
-            found.extend(nodes)
+            query = attribute_graphql.render_graphql_query(query_filter=id_filter, filter_id=object_id)
+            response = await service.client.execute_graphql(query=query, branch_name=branch_name)
+            output = attribute_graphql.parse_response(response=response)
+            found.extend(output)
 
         if not found:
             log.debug("No nodes found that requires updates")
 
-        template_string = "n/a"
-        if computed_macro.attribute.computed_attribute and computed_macro.attribute.computed_attribute.jinja2_template:
-            template_string = computed_macro.attribute.computed_attribute.jinja2_template
-
         batch = await service.client.create_batch()
         for node in found:
             batch.add(
                 task=update_computed_attribute_value_jinja2,
                 branch_name=branch_name,
                 obj=node,
+                node_kind=node_schema.kind,
                 attribute_name=computed_macro.attribute.name,
-                template_value=template_string,
+                template=jinja_template,
                 service=service,
             )
 

changelog/6351.fixed.md

@@ -0,0 +1 @@
+Improve performance of computed attributes when updating a large number of objects at once. Replaced client.filter call in Jinja2 based computed attributes.