perf(computed-attribute): batch Python transform recompute

backend/infrahub/computed_attribute/tasks.py

@@ -3,8 +3,7 @@
 from typing import TYPE_CHECKING
 
 from infrahub_sdk.exceptions import URLNotFoundError
-from prefect import flow, task
-from prefect.cache_policies import NONE
+from prefect import flow
 from prefect.client.orchestration import get_client as get_prefect_client
 from prefect.logging import get_run_logger
 
@@ -29,6 +28,7 @@
     TRIGGER_UPDATE_JINJA_COMPUTED_ATTRIBUTES,
     TRIGGER_UPDATE_PYTHON_COMPUTED_ATTRIBUTES,
 )
+from infrahub.workflows.constants import WorkflowTag
 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
@@ -52,6 +52,7 @@
 if TYPE_CHECKING:
     from infrahub.core.schema.computed_attribute import ComputedAttribute
     from infrahub.database import InfrahubDatabase
+    from infrahub.git.repository import InfrahubReadOnlyRepository, InfrahubRepository
     from infrahub.graphql.analyzer import GraphQLQueryReport
 
 
@@ -78,24 +79,6 @@ async def _reconcile_python_computed_attribute_automations(db: InfrahubDatabase)
     log.debug("Reconciled Python computed-attribute node-input automations")
 
 
-UPDATE_ATTRIBUTE = """
-mutation UpdateAttribute(
-    $id: String!,
-    $kind: String!,
-    $attribute: String!,
-    $value: String!
-    $context_account_id: String!
-  ) {
-  InfrahubUpdateComputedAttribute(
-    context: {account: {id: $context_account_id}},
-    data: {id: $id, attribute: $attribute, value: $value, kind: $kind}
-  ) {
-    ok
-  }
-}
-"""
-
-
 def _resolve_changed_elements(
     changed_elements: ChangedElementsPayload | None,
 ) -> ChangedElementSet | None:
@@ -116,34 +99,29 @@ def _transform_read_set_from_query_report(report: GraphQLQueryReport) -> Transfo
     return TransformReadSet.from_read_fields({kind: access.fields for kind, access in report.requested_read.items()})
 
 
-@task(name="computed-attribute-process-transform-for-node", cache_policy=NONE)
-async def process_transform_for_node(
+async def _transform_value_for_node(
+    *,
     branch_name: str,
     object_id: str,
-    node_kind: str,
     attribute_name: str,
     query_id: str,
     transform_timeout: int | None,
-    repository_id: str,
-    repository_name: str,
-    repository_kind: str,
     commit: str | None,
     file_path: str,
     class_name: str,
     convert_query_response: bool,
     context: EventContext,
-) -> None:
+    repo: InfrahubReadOnlyRepository | InfrahubRepository,
+) -> AttributeValueWrite:
+    """Compute one node's value against a shared, pre-initialized repository.
+
+    ``update_group`` keeps the node subscribed to the transform's query group, which is
+    what routes future source changes back to this node. ``context`` is reapplied here
+    because ``get_client()`` builds a fresh client per call rather than inheriting one.
+    """
     client = get_client()
     client.request_context = context.to_request_context()
 
-    repo = await get_initialized_repo(
-        client=client,
-        repository_id=repository_id,
-        name=repository_name,
-        repository_kind=repository_kind,
-        commit=commit,
-    )
-
     data = await client.query_gql_query(
         name=query_id,
         branch_name=branch_name,
@@ -152,7 +130,7 @@ async def process_transform_for_node(
         subscribers=[object_id],
     )
 
-    transformed_data = await repo.execute_python_transform.with_options(timeout_seconds=transform_timeout)(
+    value = await repo.execute_python_transform.with_options(timeout_seconds=transform_timeout)(
         client=client,
         branch_name=branch_name,
         commit=commit,
@@ -161,17 +139,27 @@ async def process_transform_for_node(
         convert_query_response=convert_query_response,
     )  # type: ignore[call-overload]
 
-    await client.execute_graphql(
-        query=UPDATE_ATTRIBUTE,
-        variables={
-            "id": object_id,
-            "kind": node_kind,
-            "attribute": attribute_name,
-            "value": transformed_data,
-            "context_account_id": context.account_id,
-        },
-        branch_name=branch_name,
-    )
+    return AttributeValueWrite(node_id=object_id, field=attribute_name, value=value)
+
+
+def _partition_transform_results(
+    results: list[tuple[str, AttributeValueWrite | Exception]],
+) -> tuple[list[AttributeValueWrite], list[tuple[str, str]]]:
+    """Split results into values to persist and ``(node_id, reason)`` skips.
+
+    A raised or non-string result is skipped so one bad node cannot block its siblings,
+    and a failure never overwrites the last good value with null.
+    """
+    writes: list[AttributeValueWrite] = []
+    skipped: list[tuple[str, str]] = []
+    for object_id, result in results:
+        if isinstance(result, Exception):
+            skipped.append((object_id, f"transform raised {result!r}"))
+        elif not isinstance(result.value, str):
+            skipped.append((object_id, f"transform returned {type(result.value).__name__}, expected a string"))
+        else:
+            writes.append(result)
+    return writes, skipped
 
 
 @flow(
@@ -188,8 +176,21 @@ async def process_transform(
     object_ids: list[str] | None = None,
     updated_fields: list[str] | None = None,  # noqa: ARG001
 ) -> None:
+    """Recompute Python computed attributes for a batch of nodes.
+
+    One repository init and one bulk write per batch; unchanged values emit no events
+    and fan out no further recompute; a failing node keeps its previous value without
+    blocking its siblings.
+
+    Raises:
+        ValueError: if a computed attribute has no transform configured or the transform cannot be fetched.
+
+    """
+    log = get_run_logger()
     all_ids = list({*([object_id] if object_id else []), *(object_ids or [])})
     await add_tags(branches=[branch_name], nodes=all_ids)
+    if not all_ids:
+        return
     client = get_client()
     client.request_context = context.to_request_context()
 
@@ -219,26 +220,51 @@ async def process_transform(
                 f"Unable to fetch transform '{transform_attribute.transform}' for computed attribute '{attribute_name}'"
             )
 
-        batch = await client.create_batch()
+        repo = await get_initialized_repo(
+            client=client,
+            repository_id=transform.repository_id,
+            name=transform.repository_name,
+            repository_kind=transform.repository_typename,
+            commit=transform.repository_commit,
+        )
+
+        # One failing node must not abort the batch and drop its siblings' writes.
+        batch = await client.create_batch(return_exceptions=True)
         for oid in all_ids:
             batch.add(
-                task=process_transform_for_node,
+                task=_transform_value_for_node,
+                node=oid,
                 branch_name=branch_name,
                 object_id=oid,
-                node_kind=node_kind,
                 attribute_name=attribute_name,
                 query_id=transform.query_name,
                 transform_timeout=transform.timeout,
-                repository_id=transform.repository_id,
-                repository_name=transform.repository_name,
-                repository_kind=transform.repository_typename,
                 commit=transform.repository_commit,
                 file_path=transform.file_path,
                 class_name=transform.class_name,
                 convert_query_response=transform.convert_query_response,
                 context=context,
+                repo=repo,
             )
-        _ = [r async for _, r in batch.execute()]
+        results: list[tuple[str, AttributeValueWrite | Exception]] = [
+            (oid, result) async for oid, result in batch.execute()
+        ]
+        writes, skipped = _partition_transform_results(results)
+        for skipped_id, reason in skipped:
+            log.warning(f"Skipping recompute of '{attribute_name}' for node {skipped_id}: {reason}")
+
+        dispatcher = await build_bulk_recompute_dispatcher(schema_branch=schema_branch)
+        await dispatcher.dispatch(
+            writes=writes,
+            branch_name=branch_name,
+            context=context,
+            coalesced=False,
+            recompute_depth=0,
+        )
+        log.info(
+            f"Recompute of '{attribute_name}' complete: submitted={len(results)} "
+            f"written={len(writes)} skipped={len(skipped)}"
+        )
 
 
 @flow(
@@ -274,6 +300,8 @@ async def trigger_update_python_computed_attributes(
                 "computed_attribute_kind": computed_attribute_kind,
                 "context": context,
             },
+            # Must be a creation tag: in-flow tag updates drop tags added mid-run.
+            tags=[WorkflowTag.BRANCH.render(identifier=branch_name)],
         )
 
 
@@ -675,6 +703,8 @@ async def query_transform_targets(
                     "computed_attribute_kind": kind,
                     "context": context,
                 },
+                # Must be a creation tag: in-flow tag updates drop tags added mid-run.
+                tags=[WorkflowTag.BRANCH.render(identifier=branch_name)],
             )
 
 

backend/infrahub/computed_attribute/transform_recompute.py

@@ -6,6 +6,7 @@
 from infrahub.core.registry import registry
 from infrahub.log import get_logger
 from infrahub.workflows.catalogue import TRIGGER_UPDATE_PYTHON_COMPUTED_ATTRIBUTES
+from infrahub.workflows.constants import WorkflowTag
 
 from .recompute_resolution import RecomputeResolver
 
@@ -66,5 +67,7 @@ async def submit(self, *, branch_name: str, transform_id: str, context: EventCon
                     "computed_attribute_kind": definition.kind,
                     "context": context,
                 },
+                # Must be a creation tag: in-flow tag updates drop tags added mid-run.
+                tags=[WorkflowTag.BRANCH.render(identifier=branch_name)],
             )
         return len(definitions)

backend/tests/adapters/workflow.py

@@ -54,7 +54,7 @@ async def submit_workflow(
         tags: list[str] | None = None,
         priority: WorkflowPriority | None = None,
     ) -> WorkflowInfo:
-        self.calls.append({"kind": "submit", "workflow": workflow, "parameters": parameters or {}})
+        self.calls.append({"kind": "submit", "workflow": workflow, "parameters": parameters or {}, "tags": tags or []})
         return WorkflowInfo(id=uuid.uuid4())
 
     def get_execute_calls_for(self, workflow: WorkflowDefinition) -> list[dict[str, Any]]:

backend/tests/component/computed_attribute/_base.py

@@ -11,7 +11,8 @@
 from infrahub.auth.session import AccountSession
 from infrahub.auth.types import AuthType
 from infrahub.context import InfrahubContext
-from infrahub.core.constants import RelationshipCardinality
+from infrahub.core.constants import InfrahubKind, RelationshipCardinality
+from infrahub.core.node import Node
 from infrahub.core.schema import AttributeSchema, NodeSchema, RelationshipSchema, SchemaRoot
 from infrahub.core.schema.computed_attribute import ComputedAttribute, ComputedAttributeKind
 from infrahub.events.schema_action import ChangedElementsPayload  # noqa: TC001  used in dataclass field
@@ -26,6 +27,7 @@
 
     from infrahub.core.branch import Branch
     from infrahub.core.protocols import CoreAccount
+    from infrahub.database import InfrahubDatabase
     from infrahub.events.models import EventContext
     from infrahub.services import InfrahubServices
     from infrahub.workflows.models import WorkflowDefinition
@@ -83,6 +85,35 @@
 )
 
 
+async def create_transform01(db: InfrahubDatabase, branch_name: str) -> Node:
+    """query01/repo01/transform01: a Python transform whose query reads only TestCar.name.
+
+    The edges/node query structure is required for the analyzer to record the field as
+    a read, making ``name`` the single "related" field. Returns the repository so a
+    test can attach further transforms to it.
+    """
+    query = await Node.init(db=db, schema=InfrahubKind.GRAPHQLQUERY)
+    await query.new(
+        db=db,
+        name="query01",
+        query="query { TestCar { edges { node { name { value } } } } }",
+        models=["TestCar", "TestPerson"],
+    )
+    await query.save(db=db)
+
+    repo = await Node.init(db=db, schema=InfrahubKind.READONLYREPOSITORY)
+    await repo.new(db=db, name="repo01", ref=branch_name, commit="commit01", location="location01", queries=[query])
+    await repo.save(db=db)
+
+    transform = await Node.init(db=db, schema=InfrahubKind.TRANSFORMPYTHON)
+    await transform.new(
+        db=db, name="transform01", file_path="transform.py", class_name="Transform", query=query, repository=repo
+    )
+    await transform.save(db=db)
+
+    return repo
+
+
 @dataclass
 class ScopedRecomputeCase:
     """A single ``(changed_elements -> expected submitted set)`` parametrize case."""