diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/task_arg_binding.py b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/task_arg_binding.py new file mode 100644 index 0000000000000..982c199868394 --- /dev/null +++ b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/task_arg_binding.py @@ -0,0 +1,104 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +Positional-argument binding spec for stub (foreign-runtime) tasks. + +Captured at parse time from the ``@task.stub`` TaskFlow call, stored in the serialized +Dag, and delivered to the lang-SDK runtime via ``TIRunContext.arg_bindings``. +""" + +from __future__ import annotations + +from functools import cache +from typing import Annotated, Literal + +from pydantic import Field, JsonValue, TypeAdapter +from typing_extensions import TypeAliasType + +from airflow.api_fastapi.core_api.base import BaseModel + +# A named, titled alias (like TaskArgBinding below) kept as free-form JSON rather than a +# typed model, so unknown JSON-schema keywords survive re-serialization along the way. +ArgValueSchema = TypeAliasType( + "ArgValueSchema", Annotated[dict[str, JsonValue], Field(title="ArgValueSchema")] +) +"""JSON-schema fragment constraining the value a stub-task argument binds to; generated +by pydantic from the stub annotation, carried verbatim, unknown keywords ignored.""" + + +class XComArgBinding(BaseModel): + """One positional stub-task argument pulled from an upstream task's XCom.""" + + # No default: it would drop ``kind`` from ``required``, and the generated task-sdk + # client then types it ``Literal | None``, invalid as a tagged-union discriminator. + kind: Literal["xcom"] + + name: str + """The stub function's parameter name this binding fills, in declaration order.""" + + value_schema: ArgValueSchema | None = None + """Schema fragment from the stub function's annotation; omitted when unconstrained.""" + + task_id: str + """Upstream task id whose ``return_value`` XCom is pulled.""" + + map_index: int = -1 + """Map index of the upstream XCom row to pull; -1 is the unmapped row.""" + + element_index: int | None = None + """When set, the pulled value is a sequence and this binding takes the element at + this index (the stub was expanded over an unmapped upstream's output). The lang-SDK + side will GET the single row ``(task_id, map_index=-1)``, decode it, and take + ``value[element_index]``.""" + + +class LiteralArgBinding(BaseModel): + """One positional stub-task argument carrying an inline literal from the Dag file.""" + + kind: Literal["literal"] + """No default, for the same generated-client reason as ``XComArgBinding.kind``.""" + + name: str + """The stub function's parameter name this binding fills, in declaration order.""" + + value_schema: ArgValueSchema | None = None + """Schema fragment from the stub function's annotation; omitted when unconstrained.""" + + value: JsonValue | None = None + """The literal value from the Dag file.""" + + from_default: bool = False + """True when the value was filled from the stub signature's default rather than passed in the call.""" + + +# A named alias with an explicit title so the union lands in every schema as its own +# named definition, which the supervisor-schema dump dedups with its task-sdk twin by title. +TaskArgBinding = TypeAliasType( + "TaskArgBinding", + Annotated[XComArgBinding | LiteralArgBinding, Field(discriminator="kind", title="TaskArgBinding")], +) +"""One positional argument of a stub (foreign-runtime) task, in declaration order.""" + + +@cache +def get_arg_bindings_adapter() -> TypeAdapter[list[TaskArgBinding]]: + """ + Build (lazily, then cache) the adapter validating serialized dicts into ``TaskArgBinding``. + + Only the stub-task path in the execution API needs it, so regular runs never pay for it. + """ + return TypeAdapter(list[TaskArgBinding]) diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py index ad051b3e6d340..5e09e0ac06619 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/taskinstance.py @@ -36,6 +36,7 @@ from airflow.api_fastapi.core_api.base import BaseModel, StrictBaseModel from airflow.api_fastapi.execution_api.datamodels.asset import AssetProfile from airflow.api_fastapi.execution_api.datamodels.connection import ConnectionResponse +from airflow.api_fastapi.execution_api.datamodels.task_arg_binding import TaskArgBinding from airflow.api_fastapi.execution_api.datamodels.variable import VariableResponse from airflow.utils.state import ( DagRunState, @@ -435,6 +436,13 @@ class TIRunContext(BaseModel): always reflects when the task *first* started, not when it was rescheduled/resumed. """ + arg_bindings: list[TaskArgBinding] | None = None + """ + Ordered positional-argument binding spec for stub (foreign-runtime) tasks. + + ``None`` for regular tasks and for stub tasks that declare no parameters. + """ + class PrevSuccessfulDagRunResponse(BaseModel): """Schema for response with previous successful DagRun information for Task Template Context.""" diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py index 34c3dc35406f8..694a4f4729cf5 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py @@ -32,7 +32,7 @@ from opentelemetry import trace from opentelemetry.trace import StatusCode from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator -from pydantic import JsonValue +from pydantic import JsonValue, ValidationError from sqlalchemy import and_, func, or_, tuple_, update from sqlalchemy.engine import CursorResult from sqlalchemy.exc import DataError, NoResultFound, SQLAlchemyError @@ -49,6 +49,7 @@ from airflow.api_fastapi.common.types import UtcDateTime from airflow.api_fastapi.compat import HTTP_422_UNPROCESSABLE_CONTENT from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc +from airflow.api_fastapi.execution_api.datamodels.task_arg_binding import get_arg_bindings_adapter from airflow.api_fastapi.execution_api.datamodels.taskinstance import ( InactiveAssetsResponse, PreviousTIResponse, @@ -75,6 +76,8 @@ get_team_name_for_ti, require_auth, ) +from airflow.api_fastapi.execution_api.services.task_instances import STUB_TASK_TYPE, get_arg_bindings +from airflow.api_fastapi.execution_api.versions import bundle from airflow.configuration import conf from airflow.exceptions import InvalidPartitionKeyError, TaskNotFound from airflow.models.asset import AssetActive @@ -110,6 +113,22 @@ log = structlog.get_logger(__name__) tracer = trace.get_tracer(__name__) +# The first execution API version whose TIRunContext carries ``arg_bindings``. +ARG_BINDINGS_API_VERSION = "2026-10-30" + + +def _client_supports_arg_bindings() -> bool: + """ + Whether the request's negotiated API version can receive ``arg_bindings``. + + Clients on older versions never see the field (the version migration strips it from + the response), so the derivation -- and the structured failures it raises for + undeliverable mapped-stub specs -- must not run for them: a stub Dag that ran before + arg bindings existed keeps running against those clients. + """ + version = bundle.api_version_var.get(None) + return version is None or str(version) >= ARG_BINDINGS_API_VERSION + @ti_id_router.patch( "/{task_instance_id}/run", @@ -163,6 +182,8 @@ def ti_run( TI.hostname, TI.unixname, TI.pid, + TI.operator, + TI.dag_version_id, # This selects the raw JSON value, bypassing the deserialization -- we want that to happen on the # client column("next_kwargs", JSON), @@ -310,6 +331,30 @@ def ti_run( should_retry=_is_eligible_to_retry(previous_state, ti.try_number, ti.max_tries), ) + # Only set for stub (foreign-runtime) tasks with a captured TaskFlow arg + # spec; the route excludes unset fields, keeping regular responses lean. + if ( + ti.operator == STUB_TASK_TYPE + and _client_supports_arg_bindings() + and (arg_bindings := get_arg_bindings(dag_bag, ti, session=session)) + ): + try: + context.arg_bindings = get_arg_bindings_adapter().validate_python(arg_bindings) + except ValidationError: + log.exception( + "Serialized arg_bindings spec failed validation", + dag_id=ti.dag_id, + task_id=ti.task_id, + dag_version_id=ti.dag_version_id, + ) + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail={ + "reason": "invalid_arg_bindings", + "message": "The serialized TaskFlow arg spec for this stub task is not valid.", + }, + ) + # Only set if they are non-null if ti.next_method: context.next_method = ti.next_method diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/services/__init__.py b/airflow-core/src/airflow/api_fastapi/execution_api/services/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/airflow-core/src/airflow/api_fastapi/execution_api/services/__init__.py @@ -0,0 +1,16 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py b/airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py new file mode 100644 index 0000000000000..4a8c2589438d8 --- /dev/null +++ b/airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py @@ -0,0 +1,167 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""Business logic backing the task-instance execution routes.""" + +from __future__ import annotations + +import json +from typing import TYPE_CHECKING, Any, NoReturn + +from fastapi import HTTPException, status + +from airflow.models.expandinput import NotFullyPopulated, SchedulerDictOfListsExpandInput +from airflow.models.xcom import XCOM_RETURN_KEY +from airflow.serialization.definitions.mappedoperator import is_mapped +from airflow.serialization.definitions.xcom_arg import SchedulerPlainXComArg, SchedulerXComArg +from airflow.serialization.serialized_objects import _XComRef + +if TYPE_CHECKING: + from sqlalchemy.orm import Session + + from airflow.models.dagbag import DBDagBag + from airflow.serialization.definitions.dag import SerializedDAG + from airflow.serialization.definitions.mappedoperator import SerializedMappedOperator + +# Task type recorded on the TI row (``TaskInstance.operator``) for +# ``airflow.providers.standard.decorators.stub._StubOperator``. Used to gate the +# serialized-Dag lookup for ``arg_bindings`` so regular tasks never pay for it. +# The gate matches the exact class name; a subclass would need its own entry here. +STUB_TASK_TYPE = "_StubOperator" + + +def get_arg_bindings(dag_bag: DBDagBag, ti: Any, *, session: Session) -> list | None: + """Extract or derive the stub task's TaskFlow arg spec from its Dag version.""" + if ti.dag_version_id is None: + return None + if (dag := dag_bag.get_dag(ti.dag_version_id, session=session)) is None: + return None + if (task := dag.task_dict.get(ti.task_id)) is None: + return None + if is_mapped(task): + return _resolve_mapped_stub_arg_bindings(task, ti, dag=dag, session=session) + return getattr(task, "_arg_bindings", None) + + +def _unsupported_arg_bindings(detail: str) -> NoReturn: + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail={ + "reason": "invalid_arg_bindings", + "message": f"The stub task's TaskFlow arguments cannot be delivered: {detail}.", + }, + ) + + +def _resolve_mapped_stub_arg_bindings( + task: SerializedMappedOperator, ti: Any, *, dag: SerializedDAG, session: Session +) -> list[dict[str, Any]] | None: + """ + Build the per-map-index arg spec for a mapped (``.expand()``) stub task. + + A mapped stub never instantiates at parse time; the Dag serializer captures its + per-parameter metadata (declaration order, defaults, value schemas) from the stub + signature via ``get_mapped_serialized_fields``, and the map-index decomposition is + delegated to ``SchedulerDictOfListsExpandInput.resolve_expansion_sub_indexes``. + Dags serialized without the metadata (an older provider) resolve to ``None``: their + args were never deliverable, so they keep the legacy ignored-args behavior rather + than receive bindings whose order the server cannot know. + """ + metadata = getattr(task, "_mapped_arg_binding_params", None) + if metadata is None: + return None + # The isinstance/map_index/unclaimed checks below re-reject what the provider now + # fails at parse time, for serialized Dags produced by other provider versions. + expand_input = task._get_specified_expand_input() + if not isinstance(expand_input, SchedulerDictOfListsExpandInput): + _unsupported_arg_bindings("expand_kwargs() is not supported on stub tasks") + if ti.map_index < 0: + _unsupported_arg_bindings("the task instance has not been expanded to a map index") + + expand_value = expand_input.value + partial_op_kwargs = task.partial_kwargs.get("op_kwargs") or {} + if unclaimed := (set(expand_value) | set(partial_op_kwargs)) - {meta["name"] for meta in metadata}: + _unsupported_arg_bindings(f"kwargs {sorted(unclaimed)} are not in the captured parameter metadata") + try: + sub_indexes = expand_input.resolve_expansion_sub_indexes(ti.map_index, ti.run_id, session=session) + except NotFullyPopulated as e: + # Neither this nor the ValueError below can happen on the happy path: both take + # someone clearing upstream TIs or XComs themselves during the DagRun. + _unsupported_arg_bindings(f"upstream map lengths are not yet known for {sorted(e.missing)}") + except ValueError as e: + _unsupported_arg_bindings(str(e)) + + spec = [] + for meta in metadata: # Declaration order, captured at parse time. + name = meta["name"] + if name in expand_value: + entry = _bind_mapped_stub_arg(name, expand_value[name], sub_index=sub_indexes[name]) + elif name in partial_op_kwargs: + value = partial_op_kwargs[name] + # XComArgs inside partial() op_kwargs deserialize to _XComRef and are never + # dereferenced (set_task_dag_references only derefs the expand inputs). + if isinstance(value, _XComRef): + value = value.deref(dag) + entry = _bind_mapped_stub_arg(name, value, sub_index=None) + elif "default" in meta: + entry = {"name": name, "kind": "literal", "value": meta["default"], "from_default": True} + else: + _unsupported_arg_bindings(f"parameter {name!r} has no expanded, partial, or default value") + if (value_schema := meta.get("value_schema")) is not None: + entry["value_schema"] = value_schema + spec.append(entry) + return spec + + +def _bind_mapped_stub_arg(name: str, value: Any, *, sub_index: int | None) -> dict[str, Any]: + """Build one arg-binding dict; ``sub_index`` is set for expanded kwargs, None for partial ones.""" + if isinstance(value, SchedulerPlainXComArg): + if value.key != XCOM_RETURN_KEY: + _unsupported_arg_bindings(f"parameter {name!r} references the XCom key {value.key!r}") + if sub_index is None and value.operator.is_mapped: + # A partial() kwarg over a mapped upstream would bind the unmapped XCom row + # (map_index=-1), which never exists; the aggregated output is inexpressible. + _unsupported_arg_bindings( + f"parameter {name!r} references the aggregated output of the mapped task" + f" {value.operator.task_id!r}" + ) + entry: dict[str, Any] = {"name": name, "kind": "xcom", "task_id": value.operator.task_id} + if sub_index is not None: + if value.operator.is_mapped: + entry["map_index"] = sub_index + else: + entry["element_index"] = sub_index + return entry + if isinstance(value, SchedulerXComArg): + _unsupported_arg_bindings( + f"parameter {name!r} received a {type(value).__name__}; only direct upstream" + " task outputs and literals are supported" + ) + if sub_index is not None: + # This kwarg was expanded over a literal collection written in the Dag file. + items = list(value.items()) if isinstance(value, dict) else value + try: + value = items[sub_index] + except (IndexError, KeyError, TypeError): + _unsupported_arg_bindings(f"parameter {name!r} has no element at expansion index {sub_index}") + try: + json.dumps(value, allow_nan=False) + except (TypeError, ValueError): + _unsupported_arg_bindings( + f"parameter {name!r} carries a {type(value).__name__} value, which cannot cross" + " the language boundary" + ) + return {"name": name, "kind": "literal", "value": value} diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py index dc7035d31e3c9..d56ec735c8f13 100644 --- a/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/__init__.py @@ -51,9 +51,11 @@ AddTeamNameField, AddVariableKeysEndpoint, ) +from airflow.api_fastapi.execution_api.versions.v2026_10_30 import AddArgBindingsToTIRunContext bundle = VersionBundle( HeadVersion(), + Version("2026-10-30", AddArgBindingsToTIRunContext), Version( "2026-06-30", AddVariableKeysEndpoint, diff --git a/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py new file mode 100644 index 0000000000000..2b456eae4da1f --- /dev/null +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py @@ -0,0 +1,40 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +from cadwyn import ( + ResponseInfo, + VersionChange, + convert_response_to_previous_version_for, + schema, +) + +from airflow.api_fastapi.execution_api.datamodels.taskinstance import TIRunContext + + +class AddArgBindingsToTIRunContext(VersionChange): + """Add the ``arg_bindings`` argument-binding spec for stub (foreign-runtime) tasks.""" + + description = __doc__ + + instructions_to_migrate_to_previous_version = (schema(TIRunContext).field("arg_bindings").didnt_exist,) + + @convert_response_to_previous_version_for(TIRunContext) # type: ignore[arg-type] + def remove_arg_bindings_field(response: ResponseInfo) -> None: # type: ignore[misc] + """Strip ``arg_bindings`` from the run context for older clients.""" + response.body.pop("arg_bindings", None) diff --git a/airflow-core/src/airflow/models/expandinput.py b/airflow-core/src/airflow/models/expandinput.py index 0363bae92620a..ecc1b398eb735 100644 --- a/airflow-core/src/airflow/models/expandinput.py +++ b/airflow-core/src/airflow/models/expandinput.py @@ -151,6 +151,33 @@ def get_total_map_length(self, run_id: str, *, session: Session) -> int: lengths = self._get_map_lengths(run_id, session=session) return functools.reduce(operator.mul, (lengths[name] for name in self.value), 1) + def resolve_expansion_sub_indexes( + self, map_index: int, run_id: str, *, session: Session + ) -> dict[str, int]: + """ + Decompose a task instance's map index into one index per expanded kwarg. + + Server-side counterpart of the index decomposition in the SDK's + ``DictOfListsExpandInput._expand_mapped_field``: the cross-product of the + expanded kwargs is ordered with the last kwarg varying fastest. A single + expanded kwarg maps one-to-one, skipping the upstream length lookups. + + :raises NotFullyPopulated: if upstream map lengths are not all known yet. + :raises ValueError: if an expanded kwarg's recorded length is zero, e.g. an + upstream was cleared and re-ran to an empty list after this task instance + was expanded (the SDK twin guards the same case). + """ + if len(self.value) == 1: + return dict.fromkeys(self.value, map_index) + lengths = self._get_map_lengths(run_id, session=session) + sub_indexes = {} + for key in reversed(self.value): + if (length := lengths[key]) < 1: + raise ValueError(f"cannot decompose map index over expanded kwarg {key!r} of length 0") + sub_indexes[key] = map_index % length + map_index //= length + return sub_indexes + def iter_references(self) -> Iterable[tuple[Operator, str]]: from airflow.models.referencemixin import ReferenceMixin diff --git a/airflow-core/src/airflow/serialization/schema.json b/airflow-core/src/airflow/serialization/schema.json index 872c3a1331ee3..39a33d1b7757b 100644 --- a/airflow-core/src/airflow/serialization/schema.json +++ b/airflow-core/src/airflow/serialization/schema.json @@ -142,6 +142,72 @@ "description": "A python dictionary containing values of any type", "type": "object" }, + "typed_dict": { + "type": "object", + "properties": { + "__type": { + "type": "string", + "const": "dict" + }, + "__var": { "$ref": "#/definitions/dict" } + }, + "required": [ + "__type", + "__var" + ], + "additionalProperties": false + }, + "arg_binding": { + "$comment": "One captured TaskFlow call argument of a @task.stub task, in dict-encoded form. The inner object stays open so future binding fields keep validating on older cores", + "type": "object", + "properties": { + "__type": { + "type": "string", + "const": "dict" + }, + "__var": { + "type": "object", + "properties": { + "name": { "type": "string" }, + "kind": { "type": "string" }, + "value_schema": { "$ref": "#/definitions/typed_dict" }, + "task_id": { "type": "string" }, + "value": {}, + "from_default": { "type": "boolean" } + }, + "required": [ "name", "kind" ] + } + }, + "required": [ + "__type", + "__var" + ], + "additionalProperties": false + }, + "arg_binding_param": { + "$comment": "Per-parameter metadata of a mapped @task.stub task, in dict-encoded form and declaration order. The inner object stays open so future metadata fields keep validating on older cores", + "type": "object", + "properties": { + "__type": { + "type": "string", + "const": "dict" + }, + "__var": { + "type": "object", + "properties": { + "name": { "type": "string" }, + "value_schema": { "$ref": "#/definitions/typed_dict" }, + "default": {} + }, + "required": [ "name" ] + } + }, + "required": [ + "__type", + "__var" + ], + "additionalProperties": false + }, "color": { "type": "string", "pattern": "^#[a-fA-F0-9]{3,6}$" @@ -345,7 +411,17 @@ "is_teardown": {"type": "boolean", "default": false}, "on_failure_fail_dagrun": {"type": "boolean", "default": false}, "max_active_tis_per_dag": {"type": "integer"}, - "max_active_tis_per_dagrun": {"type": "integer"} + "max_active_tis_per_dagrun": {"type": "integer"}, + "_arg_bindings": { + "$comment": "Only present on @task.stub tasks called with TaskFlow arguments", + "type": "array", + "items": { "$ref": "#/definitions/arg_binding" } + }, + "_mapped_arg_binding_params": { + "$comment": "Only present on mapped @task.stub tasks with parameters; ordered per-parameter binding metadata", + "type": "array", + "items": { "$ref": "#/definitions/arg_binding_param" } + } }, "dependencies": { "expand_input": ["partial_kwargs", "_is_mapped"], diff --git a/airflow-core/src/airflow/serialization/serialized_objects.py b/airflow-core/src/airflow/serialization/serialized_objects.py index 54bc3389c64ce..8f19dd9b8b20a 100644 --- a/airflow-core/src/airflow/serialization/serialized_objects.py +++ b/airflow-core/src/airflow/serialization/serialized_objects.py @@ -996,6 +996,16 @@ def serialize_mapped_operator(cls, op: MappedOperator) -> dict[str, Any]: ) del serialized_op["partial_kwargs"]["python_callable"] + # Optional per-class capability: an operator class may contribute extra serialized + # fields for its mapped form. This is the only point where operator_class is the + # real class and python_callable the real function (neither survives + # serialization), so signature-derived data must be captured here. Used by the + # standard provider's _StubOperator for its TaskFlow arg-binding metadata. + get_extra_fields = getattr(op.operator_class, "get_mapped_serialized_fields", None) + if get_extra_fields is not None: + for key, value in get_extra_fields(op).items(): + serialized_op[key] = cls.serialize(value) + serialized_op["_is_mapped"] = True return serialized_op diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py index 8a152bebe0d3f..0d61d4625f97a 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py @@ -17,6 +17,7 @@ from __future__ import annotations +import itertools from datetime import datetime from types import SimpleNamespace from typing import TYPE_CHECKING @@ -32,6 +33,7 @@ from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from opentelemetry.trace import StatusCode from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator +from pydantic import ValidationError from sqlalchemy import select, update from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.orm import Session @@ -371,6 +373,464 @@ async def workload_token(request: Request) -> TIToken: assert extras["scope"] == "execution" assert extras["sub"] == str(ti.id) + def test_ti_run_returns_arg_bindings_for_stub_task(self, client, dag_maker): + """A stub task's TaskFlow arg spec is extracted from the serialized Dag and returned.""" + with dag_maker("test_arg_bindings_dag", serialized=True): + + @task.stub + def extract(): ... + + @task.stub + def transform(country: str, extracted: dict, limit: int = 10): ... + + transform("uk", extract()) + + dr = dag_maker.create_dagrun() + tis = {ti.task_id: ti for ti in dr.get_task_instances()} + for ti in tis.values(): + ti.set_state(State.QUEUED) + dag_maker.session.flush() + + payload = { + "state": "running", + "hostname": "random-hostname", + "unixname": "random-unixname", + "pid": 100, + "start_date": "2024-09-30T12:00:00Z", + } + + response = client.patch(f"/execution/task-instances/{tis['transform'].id}/run", json=payload) + assert response.status_code == 200 + assert response.json()["arg_bindings"] == [ + {"name": "country", "kind": "literal", "value_schema": {"type": "string"}, "value": "uk"}, + { + "name": "extracted", + "kind": "xcom", + "value_schema": {"type": "object", "additionalProperties": True}, + "task_id": "extract", + }, + { + "name": "limit", + "kind": "literal", + "value_schema": {"type": "integer", "format": "int64"}, + "value": 10, + "from_default": True, + }, + ] + + # An argless stub has no captured spec, so the field stays unset. + response = client.patch(f"/execution/task-instances/{tis['extract'].id}/run", json=payload) + assert response.status_code == 200 + assert "arg_bindings" not in response.json() + + @mock.patch( + "airflow.api_fastapi.execution_api.routes.task_instances.get_arg_bindings", + autospec=True, + return_value=[{"name": "country", "kind": "hologram", "value": "uk"}], + ) + def test_ti_run_reports_invalid_arg_bindings_spec(self, _, client, dag_maker): + """A serialized spec this core version cannot validate fails with a structured error, not a bare 500.""" + with dag_maker("test_invalid_arg_bindings_dag", serialized=True): + + @task.stub + def transform(country: str): ... + + transform("uk") + + dr = dag_maker.create_dagrun() + (ti,) = dr.get_task_instances() + ti.set_state(State.QUEUED) + dag_maker.session.flush() + + response = client.patch( + f"/execution/task-instances/{ti.id}/run", + json={ + "state": "running", + "hostname": "random-hostname", + "unixname": "random-unixname", + "pid": 100, + "start_date": "2024-09-30T12:00:00Z", + }, + ) + + assert response.status_code == 500 + assert response.json()["detail"]["reason"] == "invalid_arg_bindings" + + RUN_PAYLOAD = { + "state": "running", + "hostname": "random-hostname", + "unixname": "random-unixname", + "pid": 100, + "start_date": "2024-09-30T12:00:00Z", + } + + def test_ti_run_resolves_mapped_stub_literal_expand(self, client, dag_maker): + """Expanding a stub over a literal list resolves each map index to its element server-side.""" + with dag_maker("test_mapped_stub_literal", serialized=True): + + @task.stub + def transform(country: str): ... + + transform.expand(country=["uk", "fr", "de"]) + + dr = dag_maker.create_dagrun() + tis = {ti.map_index: ti for ti in dr.get_task_instances()} + assert set(tis) == {0, 1, 2} + for ti in tis.values(): + ti.set_state(State.QUEUED) + dag_maker.session.flush() + + for map_index, country in enumerate(["uk", "fr", "de"]): + response = client.patch( + f"/execution/task-instances/{tis[map_index].id}/run", json=self.RUN_PAYLOAD + ) + assert response.status_code == 200 + assert response.json()["arg_bindings"] == [ + {"name": "country", "kind": "literal", "value_schema": {"type": "string"}, "value": country} + ] + + def test_ti_run_resolves_mapped_stub_over_unmapped_upstream(self, client, dag_maker): + """Expanding over an unmapped upstream's output binds the whole XCom plus an element index.""" + with dag_maker("test_mapped_stub_unmapped_upstream", serialized=True): + + @task.stub + def extract(): ... + + @task.stub + def transform(extracted: dict): ... + + transform.expand(extracted=extract()) + + dr = dag_maker.create_dagrun() + ti = dr.get_task_instance("transform") + ti.map_index = 1 + ti.set_state(State.QUEUED) + dag_maker.session.flush() + + response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD) + assert response.status_code == 200 + assert response.json()["arg_bindings"] == [ + { + "name": "extracted", + "kind": "xcom", + "value_schema": {"type": "object", "additionalProperties": True}, + "task_id": "extract", + "element_index": 1, + } + ] + + def test_ti_run_resolves_mapped_stub_over_mapped_upstream(self, client, dag_maker): + """Expanding over a mapped upstream binds the upstream XCom row at the same map index.""" + with dag_maker("test_mapped_stub_mapped_upstream", serialized=True): + + @task.stub + def seed(n: int): ... + + @task.stub + def transform(extracted: dict): ... + + transform.expand(extracted=seed.expand(n=[1, 2])) + + dr = dag_maker.create_dagrun() + ti = dr.get_task_instance("transform") + ti.map_index = 1 + ti.set_state(State.QUEUED) + dag_maker.session.flush() + + response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD) + assert response.status_code == 200 + assert response.json()["arg_bindings"] == [ + { + "name": "extracted", + "kind": "xcom", + "value_schema": {"type": "object", "additionalProperties": True}, + "task_id": "seed", + "map_index": 1, + } + ] + + def test_ti_run_decomposes_multi_kwarg_mapped_stub(self, client, dag_maker): + """Cross-product expansion decomposes the map index per kwarg like the task-sdk does.""" + with dag_maker("test_mapped_stub_multi_kwarg", serialized=True): + + @task.stub + def combine(a: str, b: int): ... + + combine.expand(a=["x", "y"], b=[1, 2, 3]) + + dr = dag_maker.create_dagrun() + tis = {ti.map_index: ti for ti in dr.get_task_instances()} + assert set(tis) == set(range(6)) + for ti in tis.values(): + ti.set_state(State.QUEUED) + dag_maker.session.flush() + + for map_index, (a, b) in enumerate(itertools.product(["x", "y"], [1, 2, 3])): + response = client.patch( + f"/execution/task-instances/{tis[map_index].id}/run", json=self.RUN_PAYLOAD + ) + assert response.status_code == 200 + assert response.json()["arg_bindings"] == [ + {"name": "a", "kind": "literal", "value_schema": {"type": "string"}, "value": a}, + { + "name": "b", + "kind": "literal", + "value_schema": {"type": "integer", "format": "int64"}, + "value": b, + }, + ] + + def test_ti_run_binds_partial_kwargs_of_mapped_stub(self, client, dag_maker): + """partial() kwargs bind like an unmapped TaskFlow call alongside the expanded ones.""" + with dag_maker("test_mapped_stub_partial", serialized=True): + + @task.stub + def transform(country: str, extracted: dict): ... + + transform.partial(country="uk").expand(extracted=[{"a": 1}, {"b": 2}]) + + dr = dag_maker.create_dagrun() + tis = {ti.map_index: ti for ti in dr.get_task_instances()} + for ti in tis.values(): + ti.set_state(State.QUEUED) + dag_maker.session.flush() + + for map_index, extracted in enumerate([{"a": 1}, {"b": 2}]): + response = client.patch( + f"/execution/task-instances/{tis[map_index].id}/run", json=self.RUN_PAYLOAD + ) + assert response.status_code == 200 + assert response.json()["arg_bindings"] == [ + {"name": "country", "kind": "literal", "value_schema": {"type": "string"}, "value": "uk"}, + { + "name": "extracted", + "kind": "literal", + "value_schema": {"type": "object", "additionalProperties": True}, + "value": extracted, + }, + ] + + def test_ti_run_binds_partial_xcom_kwarg_over_unmapped_upstream(self, client, dag_maker): + """A partial() kwarg carrying an unmapped upstream's output binds that XCom for every index.""" + with dag_maker("test_mapped_stub_partial_xcom", serialized=True): + + @task.stub + def extract(): ... + + @task.stub + def transform(extracted: dict, country: str): ... + + transform.partial(extracted=extract()).expand(country=["uk", "fr"]) + + dr = dag_maker.create_dagrun() + ti = next(ti for ti in dr.get_task_instances() if ti.task_id == "transform" and ti.map_index == 1) + ti.set_state(State.QUEUED) + dag_maker.session.flush() + + response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD) + assert response.status_code == 200 + assert response.json()["arg_bindings"] == [ + { + "name": "extracted", + "kind": "xcom", + "value_schema": {"type": "object", "additionalProperties": True}, + "task_id": "extract", + }, + {"name": "country", "kind": "literal", "value_schema": {"type": "string"}, "value": "fr"}, + ] + + def test_ti_run_rejects_partial_kwarg_over_mapped_upstream(self, client, dag_maker): + """ + A partial() kwarg over a mapped upstream would bind the nonexistent unmapped XCom row. + + The provider rejects this at parse time now; patching its capture hook simulates a + Dag serialized by another provider version, exercising the server-side backstop. + """ + from airflow.providers.standard.decorators.stub import _StubOperator + + fabricated = {"_mapped_arg_binding_params": [{"name": "extracted"}, {"name": "country"}]} + with mock.patch.object(_StubOperator, "get_mapped_serialized_fields", return_value=fabricated): + with dag_maker("test_mapped_stub_partial_mapped_upstream", serialized=True): + + @task.stub + def seed(n: int): ... + + @task.stub + def transform(extracted: dict, country: str): ... + + transform.partial(extracted=seed.expand(n=[1, 2])).expand(country=["uk", "fr"]) + + dr = dag_maker.create_dagrun() + ti = next(ti for ti in dr.get_task_instances() if ti.task_id == "transform" and ti.map_index == 0) + ti.set_state(State.QUEUED) + dag_maker.session.flush() + + response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD) + assert response.status_code == 500 + assert response.json()["detail"]["reason"] == "invalid_arg_bindings" + assert "aggregated output" in response.json()["detail"]["message"] + + def test_ti_run_orders_mapped_stub_spec_by_declaration_with_defaults(self, client, dag_maker): + """The spec follows the signature, not the call sites, and ships defaulted params.""" + with dag_maker("test_mapped_stub_declaration_order", serialized=True): + + @task.stub + def transform(country: str, extracted: dict, retries_num: int = 3): ... + + # The partial() kwarg is declared after the expanded one on purpose. + transform.partial(extracted={"a": 1}).expand(country=["uk", "fr"]) + + dr = dag_maker.create_dagrun() + ti = next(t for t in dr.get_task_instances() if t.map_index == 1) + ti.set_state(State.QUEUED) + dag_maker.session.flush() + + response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD) + assert response.status_code == 200 + assert response.json()["arg_bindings"] == [ + {"name": "country", "kind": "literal", "value_schema": {"type": "string"}, "value": "fr"}, + { + "name": "extracted", + "kind": "literal", + "value_schema": {"type": "object", "additionalProperties": True}, + "value": {"a": 1}, + }, + { + "name": "retries_num", + "kind": "literal", + "value_schema": {"type": "integer", "format": "int64"}, + "value": 3, + "from_default": True, + }, + ] + + def test_ti_run_ignores_args_for_legacy_serialized_mapped_stub(self, client, dag_maker): + """A mapped stub serialized without parameter metadata keeps the ignored-args behavior.""" + from airflow.providers.standard.decorators.stub import _StubOperator + + with mock.patch.object(_StubOperator, "get_mapped_serialized_fields", return_value={}): + with dag_maker("test_mapped_stub_legacy", serialized=True): + + @task.stub + def transform(country: str): ... + + transform.expand(country=["uk", "fr"]) + + dr = dag_maker.create_dagrun() + ti = next(t for t in dr.get_task_instances() if t.map_index == 0) + ti.set_state(State.QUEUED) + dag_maker.session.flush() + + response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD) + assert response.status_code == 200 + assert "arg_bindings" not in response.json() + + def test_ti_run_rejects_unexpanded_mapped_stub_ti(self, client, dag_maker): + """A mapped stub TI still at map_index=-1 cannot receive per-index bindings.""" + with dag_maker("test_mapped_stub_unexpanded", serialized=True): + + @task.stub + def extract(): ... + + @task.stub + def transform(extracted: dict): ... + + transform.expand(extracted=extract()) + + dr = dag_maker.create_dagrun() + ti = dr.get_task_instance("transform") + assert ti.map_index == -1 + ti.set_state(State.QUEUED) + dag_maker.session.flush() + + response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD) + assert response.status_code == 500 + assert response.json()["detail"]["reason"] == "invalid_arg_bindings" + assert "not been expanded" in response.json()["detail"]["message"] + + def test_ti_run_rejects_zero_length_expansion_on_stub(self, client, dag_maker, session): + """An upstream re-run to an empty list after expansion fails structurally, not with a crash.""" + from airflow.models.taskmap import TaskMap + + with dag_maker("test_mapped_stub_zero_length", serialized=True, session=session): + + @task + def seed(): + return [0, 1] + + @task.stub + def combine(a: int, b: int): ... + + combine.expand(a=seed(), b=[1, 2]) + + dr = dag_maker.create_dagrun() + decision = dr.task_instance_scheduling_decisions(session=session) + (seed_ti,) = decision.schedulable_tis + seed_ti.state = TaskInstanceState.SUCCESS + session.add(TaskMap.from_task_instance_xcom(seed_ti, [0, 1])) + session.flush() + + decision = dr.task_instance_scheduling_decisions(session=session) + ti = next(t for t in decision.schedulable_tis if t.map_index == 0) + ti.set_state(State.QUEUED, session=session) + # Simulate the upstream being cleared and re-run to an empty list while this + # expanded TI is still queued. + session.execute(update(TaskMap).where(TaskMap.task_id == "seed").values(length=0, keys=None)) + session.commit() + + response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD) + assert response.status_code == 500 + assert response.json()["detail"]["reason"] == "invalid_arg_bindings" + assert "length 0" in response.json()["detail"]["message"] + + def test_ti_run_rejects_expand_kwargs_on_stub(self, client, dag_maker): + """ + expand_kwargs() has no per-parameter spec to derive, so delivery fails structurally. + + The provider rejects this at parse time now; patching its capture hook simulates a + Dag serialized by another provider version, exercising the server-side backstop. + """ + from airflow.providers.standard.decorators.stub import _StubOperator + + fabricated = {"_mapped_arg_binding_params": [{"name": "country"}]} + with mock.patch.object(_StubOperator, "get_mapped_serialized_fields", return_value=fabricated): + with dag_maker("test_mapped_stub_expand_kwargs", serialized=True): + + @task.stub + def transform(country: str): ... + + transform.expand_kwargs([{"country": "uk"}]) + + dr = dag_maker.create_dagrun() + (ti,) = dr.get_task_instances() + ti.set_state(State.QUEUED) + dag_maker.session.flush() + + response = client.patch(f"/execution/task-instances/{ti.id}/run", json=self.RUN_PAYLOAD) + assert response.status_code == 500 + assert response.json()["detail"]["reason"] == "invalid_arg_bindings" + assert "expand_kwargs" in response.json()["detail"]["message"] + + def test_arg_bindings_adapter_rejects_unknown_kind(self): + """The discriminated union refuses serialized specs with an unrecognised kind.""" + from airflow.api_fastapi.execution_api.datamodels.task_arg_binding import get_arg_bindings_adapter + + with pytest.raises(ValidationError, match="does not match any of the expected tags"): + get_arg_bindings_adapter().validate_python( + [{"name": "country", "kind": "template", "value": "x"}] + ) + + def test_arg_bindings_adapter_carries_value_schema_fragments_verbatim(self): + """The fragment is free-form JSON schema: every keyword the provider generated must + survive validation untouched -- a typed model would silently strip what it doesn't know.""" + from airflow.api_fastapi.execution_api.datamodels.task_arg_binding import get_arg_bindings_adapter + + fragment = {"anyOf": [{"type": "array", "items": {"type": "string"}}, {"type": "null"}]} + (binding,) = get_arg_bindings_adapter().validate_python( + [{"name": "tags", "kind": "literal", "value_schema": fragment, "value": ["a"]}] + ) + assert binding.value_schema == fragment + def test_dynamic_task_mapping_with_parse_time_value(self, client, dag_maker): """Test that dynamic task mapping works correctly with parse-time values.""" with dag_maker("test_dynamic_task_mapping_with_parse_time_value", serialized=True): diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/__init__.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/__init__.py new file mode 100644 index 0000000000000..13a83393a9124 --- /dev/null +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/__init__.py @@ -0,0 +1,16 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py new file mode 100644 index 0000000000000..b4dd521c760e2 --- /dev/null +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py @@ -0,0 +1,124 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +import pytest + +from airflow.sdk import task +from airflow.utils.state import State + +from tests_common.test_utils.db import clear_db_runs + +pytestmark = pytest.mark.db_test + +TIMESTAMP_STR = "2024-09-30T12:00:00Z" + +RUN_PATCH_BODY = { + "state": "running", + "hostname": "h", + "unixname": "u", + "pid": 1, + "start_date": TIMESTAMP_STR, +} + + +@pytest.fixture +def old_ver_client(client): + """Execution API version immediately before ``arg_bindings`` was added.""" + client.headers["Airflow-API-Version"] = "2026-06-30" + return client + + +class TestArgBindingsFieldBackwardCompat: + @pytest.fixture(autouse=True) + def _freeze_time(self, time_machine): + time_machine.move_to(TIMESTAMP_STR, tick=False) + + def setup_method(self): + clear_db_runs() + + def teardown_method(self): + clear_db_runs() + + @pytest.fixture + def stub_ti(self, dag_maker): + with dag_maker("test_arg_bindings_compat_dag", serialized=True): + + @task.stub + def extract(): ... + + @task.stub + def transform(country: str, extracted: dict, limit: int = 10): ... + + transform("uk", extract()) + + dr = dag_maker.create_dagrun() + tis = {ti.task_id: ti for ti in dr.get_task_instances()} + for ti in tis.values(): + ti.set_state(State.QUEUED) + dag_maker.session.flush() + return tis["transform"] + + def test_old_version_strips_arg_bindings_even_when_set(self, old_ver_client, stub_ti): + response = old_ver_client.patch(f"/execution/task-instances/{stub_ti.id}/run", json=RUN_PATCH_BODY) + assert response.status_code == 200 + assert "arg_bindings" not in response.json() + + def test_old_version_skips_undeliverable_arg_bindings_derivation(self, old_ver_client, dag_maker): + """A stub whose bindings cannot be delivered must keep running for clients that never see them.""" + with dag_maker("test_arg_bindings_compat_unexpanded", serialized=True): + + @task.stub + def extract(): ... + + @task.stub + def transform(extracted: dict): ... + + transform.expand(extracted=extract()) + + dr = dag_maker.create_dagrun() + ti = dr.get_task_instance("transform") + assert ti.map_index == -1 + ti.set_state(State.QUEUED) + dag_maker.session.flush() + + # At head this TI fails ti_run with a structured 500 (it has not been expanded + # to a map index); a pre-arg-bindings client keeps the legacy behavior. + response = old_ver_client.patch(f"/execution/task-instances/{ti.id}/run", json=RUN_PATCH_BODY) + assert response.status_code == 200 + assert "arg_bindings" not in response.json() + + def test_head_version_includes_arg_bindings(self, client, stub_ti): + response = client.patch(f"/execution/task-instances/{stub_ti.id}/run", json=RUN_PATCH_BODY) + assert response.status_code == 200 + assert response.json()["arg_bindings"] == [ + {"name": "country", "kind": "literal", "value_schema": {"type": "string"}, "value": "uk"}, + { + "name": "extracted", + "kind": "xcom", + "value_schema": {"type": "object", "additionalProperties": True}, + "task_id": "extract", + }, + { + "name": "limit", + "kind": "literal", + "value_schema": {"type": "integer", "format": "int64"}, + "value": 10, + "from_default": True, + }, + ] diff --git a/airflow-core/tests/unit/serialization/test_dag_serialization.py b/airflow-core/tests/unit/serialization/test_dag_serialization.py index 4708365e846e2..a7459b204cc5c 100644 --- a/airflow-core/tests/unit/serialization/test_dag_serialization.py +++ b/airflow-core/tests/unit/serialization/test_dag_serialization.py @@ -77,7 +77,7 @@ from airflow.serialization.definitions.param import SerializedParam from airflow.serialization.definitions.xcom_arg import SchedulerPlainXComArg from airflow.serialization.encoders import ensure_serialized_asset -from airflow.serialization.enums import Encoding +from airflow.serialization.enums import DagAttributeTypes as DAT, Encoding from airflow.serialization.json_schema import load_dag_schema_dict from airflow.serialization.serialized_objects import ( BaseSerialization, @@ -3405,6 +3405,127 @@ def inner(): assert serialized3["python_callable_name"] == "empty_function" +def test_stub_task_args_round_trip(): + """The stub task's TaskFlow arg spec (``_arg_bindings``) survives Dag serialization.""" + from airflow.sdk import task + + with DAG(dag_id="arg_bindings_dag", schedule=None) as dag: + + @task.stub + def extract(): ... + + @task.stub + def transform(country: str, extracted: dict): ... + + transform("uk", extract()) + + ser_dag = DagSerialization.to_dict(dag) + encoded_tasks = {t[Encoding.VAR]["task_id"]: t[Encoding.VAR] for t in ser_dag["dag"]["tasks"]} + assert "_arg_bindings" not in encoded_tasks["extract"], "argless stubs must not serialize a spec" + assert encoded_tasks["transform"]["_arg_bindings"] == [ + { + Encoding.TYPE: DAT.DICT, + Encoding.VAR: { + "name": "country", + "kind": "literal", + "value_schema": {Encoding.TYPE: DAT.DICT, Encoding.VAR: {"type": "string"}}, + "value": "uk", + }, + }, + { + Encoding.TYPE: DAT.DICT, + Encoding.VAR: { + "name": "extracted", + "kind": "xcom", + "value_schema": { + Encoding.TYPE: DAT.DICT, + Encoding.VAR: {"type": "object", "additionalProperties": True}, + }, + "task_id": "extract", + }, + }, + ] + + round_tripped = DagSerialization.from_dict(ser_dag) + assert round_tripped.task_dict["transform"]._arg_bindings == [ + {"name": "country", "kind": "literal", "value_schema": {"type": "string"}, "value": "uk"}, + { + "name": "extracted", + "kind": "xcom", + "value_schema": {"type": "object", "additionalProperties": True}, + "task_id": "extract", + }, + ] + assert not hasattr(round_tripped.task_dict["extract"], "_arg_bindings") + + +def test_mapped_stub_param_metadata_round_trip(): + """The serializer collects the mapped stub's parameter metadata and it survives the round trip.""" + from airflow.sdk import task + + with DAG(dag_id="mapped_arg_binding_params_dag", schedule=None) as dag: + + @task.stub + def transform(country: str, extracted, retries_num: int = 3): ... + + @task + def plain(x): ... + + transform.partial(extracted={"a": 1}).expand(country=["uk", "fr"]) + plain.expand(x=[1, 2]) + + ser_dag = DagSerialization.to_dict(dag) + DagSerialization.validate_schema(ser_dag) + encoded_tasks = {t[Encoding.VAR]["task_id"]: t[Encoding.VAR] for t in ser_dag["dag"]["tasks"]} + assert encoded_tasks["transform"]["_mapped_arg_binding_params"] == [ + { + Encoding.TYPE: DAT.DICT, + Encoding.VAR: { + "name": "country", + "value_schema": {Encoding.TYPE: DAT.DICT, Encoding.VAR: {"type": "string"}}, + }, + }, + {Encoding.TYPE: DAT.DICT, Encoding.VAR: {"name": "extracted"}}, + { + Encoding.TYPE: DAT.DICT, + Encoding.VAR: { + "name": "retries_num", + "value_schema": { + Encoding.TYPE: DAT.DICT, + Encoding.VAR: {"type": "integer", "format": "int64"}, + }, + "default": 3, + }, + }, + ], "metadata must serialize in declaration order" + assert "_mapped_arg_binding_params" not in encoded_tasks["plain"], ( + "only operator classes defining the capture hook contribute mapped fields" + ) + + round_tripped = DagSerialization.from_dict(ser_dag) + assert round_tripped.task_dict["transform"]._mapped_arg_binding_params == [ + {"name": "country", "value_schema": {"type": "string"}}, + {"name": "extracted"}, + {"name": "retries_num", "value_schema": {"type": "integer", "format": "int64"}, "default": 3}, + ] + assert not hasattr(round_tripped.task_dict["plain"], "_mapped_arg_binding_params") + + +def test_mapped_stub_capture_error_fails_serialization(): + """A capture-hook rejection surfaces as a Dag serialization (import) error.""" + from airflow.sdk import task + + with DAG(dag_id="mapped_arg_binding_params_invalid_dag", schedule=None) as dag: + + @task.stub + def transform(country: str): ... + + transform.expand_kwargs([{"country": "uk"}]) + + with pytest.raises(SerializationError, match="does not support expand_kwargs"): + DagSerialization.to_dict(dag) + + def test_handle_v1_serdag(): v1 = { "__version": 1, diff --git a/devel-common/src/tests_common/test_utils/version_compat.py b/devel-common/src/tests_common/test_utils/version_compat.py index 7eb25dec2b3cb..d96b9dce07b4d 100644 --- a/devel-common/src/tests_common/test_utils/version_compat.py +++ b/devel-common/src/tests_common/test_utils/version_compat.py @@ -42,6 +42,7 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: AIRFLOW_V_3_2_PLUS = get_base_airflow_version_tuple() >= (3, 2, 0) AIRFLOW_V_3_2_2_PLUS = get_base_airflow_version_tuple() >= (3, 2, 2) AIRFLOW_V_3_3_PLUS = get_base_airflow_version_tuple() >= (3, 3, 0) +AIRFLOW_V_3_4_PLUS = get_base_airflow_version_tuple() >= (3, 4, 0) if AIRFLOW_V_3_1_PLUS: from airflow.sdk import PokeReturnValue, timezone diff --git a/go-sdk/cmd/airflow-go-pack/pack_integration_test.go b/go-sdk/cmd/airflow-go-pack/pack_integration_test.go index 77725b0ac4efc..84e9a1045f4e8 100644 --- a/go-sdk/cmd/airflow-go-pack/pack_integration_test.go +++ b/go-sdk/cmd/airflow-go-pack/pack_integration_test.go @@ -34,6 +34,7 @@ import ( "github.com/apache/airflow/go-sdk/internal/airflowmetadata" "github.com/apache/airflow/go-sdk/internal/bundlefooter" + "github.com/apache/airflow/go-sdk/pkg/execution" ) // crossArchFor returns an architecture different from the host that the Go @@ -142,7 +143,7 @@ func TestPack_CrossArchExecutableWithMetadataFile(t *testing.T) { sdk: language: "go" version: "` + sdkVersion + `" - supervisor_schema_version: "2026-06-16" + supervisor_schema_version: "` + execution.SupervisorSchemaVersion + `" source: "main.go" dags: concurrent_xcom_dag: diff --git a/go-sdk/pkg/execution/messages.go b/go-sdk/pkg/execution/messages.go index 72d451a866632..bb81d60c0a4ff 100644 --- a/go-sdk/pkg/execution/messages.go +++ b/go-sdk/pkg/execution/messages.go @@ -32,7 +32,7 @@ import ( // reported in a bundle's airflow-metadata manifest as // sdk.supervisor_schema_version so the supervisor can down/upgrade messages to // a shape the bundle understands. -const SupervisorSchemaVersion = "2026-06-16" +const SupervisorSchemaVersion = "2026-10-30" // The message-type discriminator strings (genmodels.Type*) are generated from the // schema's "type" consts in discriminators.gen.go; outbound messages stamp the diff --git a/providers/common/compat/src/airflow/providers/common/compat/sdk.py b/providers/common/compat/src/airflow/providers/common/compat/sdk.py index 93174df7b2a28..772650f5499e5 100644 --- a/providers/common/compat/src/airflow/providers/common/compat/sdk.py +++ b/providers/common/compat/src/airflow/providers/common/compat/sdk.py @@ -83,9 +83,13 @@ from airflow.sdk.bases.sensor import poke_mode_only as poke_mode_only from airflow.sdk.bases.skipmixin import SkipMixin as SkipMixin from airflow.sdk.configuration import conf as conf - from airflow.sdk.definitions.context import context_merge as context_merge + from airflow.sdk.definitions.context import ( + KNOWN_CONTEXT_KEYS as KNOWN_CONTEXT_KEYS, + context_merge as context_merge, + ) from airflow.sdk.definitions.mappedoperator import MappedOperator as MappedOperator from airflow.sdk.definitions.template import literal as literal + from airflow.sdk.definitions.xcom_arg import PlainXComArg as PlainXComArg from airflow.sdk.exceptions import ( AirflowConfigException as AirflowConfigException, AirflowException as AirflowException, @@ -192,6 +196,7 @@ "DAG": ("airflow.sdk", "airflow.models.dag"), "Param": ("airflow.sdk", "airflow.models.param"), "XComArg": ("airflow.sdk", "airflow.models.xcom_arg"), + "PlainXComArg": ("airflow.sdk.definitions.xcom_arg", "airflow.models.xcom_arg"), "DecoratedOperator": ("airflow.sdk.bases.decorator", "airflow.decorators.base"), "DecoratedMappedOperator": ("airflow.sdk.bases.decorator", "airflow.decorators.base"), "MappedOperator": ("airflow.sdk.definitions.mappedoperator", "airflow.models.mappedoperator"), @@ -246,6 +251,7 @@ # ============================================================================ "Context": ("airflow.sdk", "airflow.utils.context"), "context_merge": ("airflow.sdk.definitions.context", "airflow.utils.context"), + "KNOWN_CONTEXT_KEYS": ("airflow.sdk.definitions.context", "airflow.utils.context"), "context_to_airflow_vars": ("airflow.sdk.execution_time.context", "airflow.utils.operator_helpers"), "AIRFLOW_VAR_NAME_FORMAT_MAPPING": ( "airflow.sdk.execution_time.context", diff --git a/providers/standard/pyproject.toml b/providers/standard/pyproject.toml index b385a9e9eb7f6..373d16bba6c4e 100644 --- a/providers/standard/pyproject.toml +++ b/providers/standard/pyproject.toml @@ -60,7 +60,7 @@ requires-python = ">=3.10" # After you modify the dependencies, and rebuild your Breeze CI image with ``breeze ci-image build`` dependencies = [ "apache-airflow>=2.11.0", - "apache-airflow-providers-common-compat>=1.14.1", + "apache-airflow-providers-common-compat>=1.14.1", # use next version ] # The optional dependencies should be modified in place in the generated file diff --git a/providers/standard/src/airflow/providers/standard/decorators/stub.py b/providers/standard/src/airflow/providers/standard/decorators/stub.py index 08bcf163a56ad..adf554095bac5 100644 --- a/providers/standard/src/airflow/providers/standard/decorators/stub.py +++ b/providers/standard/src/airflow/providers/standard/decorators/stub.py @@ -18,12 +18,33 @@ from __future__ import annotations import ast -from collections.abc import Callable +import copy +import datetime +import inspect +import json +import types +import typing +from collections.abc import Callable, Collection, Mapping +from functools import cache from typing import TYPE_CHECKING, Any +try: + from pydantic import PydanticInvalidForJsonSchema, PydanticSchemaGenerationError, TypeAdapter + from pydantic.json_schema import GenerateJsonSchema +except ImportError: + # Airflow 3 always ships pydantic but Airflow 2.x base installs do not; without it, + # stub args carry no value schemas and runtimes keep their decode-only fallback. + GenerateJsonSchema = object # type: ignore[assignment,misc] + TypeAdapter = None # type: ignore[assignment,misc] + PydanticInvalidForJsonSchema = PydanticSchemaGenerationError = None # type: ignore[assignment,misc] + from airflow.providers.common.compat.sdk import ( + KNOWN_CONTEXT_KEYS, + XCOM_RETURN_KEY, DecoratedOperator, + PlainXComArg, TaskDecorator, + XComArg, task_decorator_factory, ) @@ -31,6 +52,278 @@ from airflow.providers.common.compat.sdk import Context +class _ValueSchemaGenerator(GenerateJsonSchema): + """ + Pydantic's stock JSON-schema generation plus OpenAPI's fixed-width numeric formats. + + A foreign runtime decodes numbers into machine types, which the bare + ``integer``/``number`` type names cannot convey; ``format`` is an annotation per + JSON schema, so runtimes that don't know these names simply skip them. + """ + + def int_schema(self, schema): + return {**super().int_schema(schema), "format": "int64"} + + def float_schema(self, schema): + return {**super().float_schema(schema), "format": "double"} + + +# Most-derived first: datetime subclasses date, so it must be matched before date. +_TEMPORAL_BASES = (datetime.datetime, datetime.date, datetime.time, datetime.timedelta) + + +def _normalize_temporal_annotation(annotation: Any) -> Any: + """ + Map temporal subclasses (e.g. ``pendulum.DateTime``) to their stdlib base. + + Applied recursively through unions and containers, and only as a retry when direct + schema generation fails, so temporal types carrying their own pydantic schema keep it. + """ + # Parametrized generics must be detected before the plain-class branch: on Python + # 3.10, isinstance(list[X], type) is True and issubclass silently consults the + # origin, so the class branch would return list[X] unnormalized. + origin = typing.get_origin(annotation) + args = typing.get_args(annotation) + if origin is not None and args: + normalized = tuple(_normalize_temporal_annotation(arg) for arg in args) + if normalized == args: + return annotation + if origin in (typing.Union, types.UnionType): + return typing.Union[normalized] # noqa: UP007 -- runtime construction from a tuple + return origin[normalized] + if isinstance(annotation, type): + return next((base for base in _TEMPORAL_BASES if issubclass(annotation, base)), annotation) + return annotation + + +def _infer_value_schema(annotation: Any) -> dict[str, Any] | None: + """ + Build the JSON-schema fragment for one stub parameter annotation, via pydantic. + + The pydantic-generated schema ships verbatim, so runtimes must treat it as + open-vocabulary JSON schema. Returns ``None`` when the annotation constrains nothing + (missing, ``Any``, bare ``None``) or pydantic cannot generate a schema for it; the + binding then omits ``value_schema`` and the foreign runtime falls back to a + decode-only check. + """ + if TypeAdapter is None: + return None + if annotation is inspect.Parameter.empty or annotation is None or annotation is Any: + return None + if annotation is type(None): + # get_type_hints normalizes a bare ``None`` annotation to NoneType; a parameter + # that can only ever be None constrains nothing worth shipping. + return None + try: + schema = _generate_value_schema(annotation) + except TypeError: + # Unhashable annotations cannot key the cache; generate directly. + schema = _generate_value_schema.__wrapped__(annotation) + # Deep-copy so callers embedding the fragment never alias the cached dict. + return copy.deepcopy(schema) if schema else None + + +@cache +def _generate_value_schema(annotation: Any) -> dict[str, Any] | None: + """ + Generate the schema for one annotation, cached for the process lifetime. + + TypeAdapter construction is one of pydantic's most expensive operations and + annotations are static, so re-parses of the same Dag file must not re-pay it. + """ + try: + return TypeAdapter(annotation).json_schema(schema_generator=_ValueSchemaGenerator) + except (PydanticSchemaGenerationError, PydanticInvalidForJsonSchema): + normalized = _normalize_temporal_annotation(annotation) + if normalized is annotation: + return None + try: + return TypeAdapter(normalized).json_schema(schema_generator=_ValueSchemaGenerator) + except (PydanticSchemaGenerationError, PydanticInvalidForJsonSchema): + return None + + +def _validate_stub_signature(signature: inspect.Signature, task_id: str) -> None: + for param in signature.parameters.values(): + if param.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD): + raise ValueError( + f"@task.stub task {task_id!r} must declare a fixed number of parameters for the " + f"foreign runtime to bind against; *{param.name} is not supported" + ) + if param.name in KNOWN_CONTEXT_KEYS: + raise ValueError( + f"@task.stub task {task_id!r} parameter {param.name!r} is an Airflow context key; " + "stub signatures declare only data parameters -- the lang-SDK runtime injects its " + "own task context natively (e.g. the Go SDK's sdk.TIRunContext parameter)" + ) + + +def _resolve_param_annotations(python_callable: Callable, signature: inspect.Signature) -> dict[str, Any]: + """Map each parameter to its parse-time-resolvable annotation (``Parameter.empty`` when not).""" + try: + hints = typing.get_type_hints(python_callable) + except (NameError, TypeError): + # Annotations that cannot be resolved at parse time (e.g. names behind + # TYPE_CHECKING with ``from __future__ import annotations``) degrade to "any". + hints = {} + + def resolve(name: str, param: inspect.Parameter) -> Any: + if name in hints: + return hints[name] + if isinstance(param.annotation, str): + return inspect.Parameter.empty + return param.annotation + + return {name: resolve(name, param) for name, param in signature.parameters.items()} + + +def _ensure_json_literal(value: Any, task_id: str, name: str) -> None: + try: + json.dumps(value, allow_nan=False) + except (TypeError, ValueError): + raise ValueError( + f"@task.stub task {task_id!r} parameter {name!r} received a literal of type " + f"{type(value).__name__} that is not JSON-serializable, so it cannot be passed " + "to the foreign runtime" + ) + + +def _validate_xcom_value(value: Any, task_id: str, name: str, *, allow_mapped_upstream: bool = False) -> bool: + """Validate an XComArg argument, returning True when it is a bindable direct upstream output.""" + if isinstance(value, PlainXComArg): + if value.key != XCOM_RETURN_KEY: + raise ValueError( + f"@task.stub task {task_id!r} parameter {name!r} references the XCom key " + f"{value.key!r}; only an upstream task's return value can cross the language " + "boundary -- indexing an output by a custom key is not supported" + ) + if value.operator.is_mapped and not allow_mapped_upstream: + raise ValueError( + f"@task.stub task {task_id!r} parameter {name!r} references the aggregated " + f"output of the mapped task {value.operator.task_id!r}; a foreign runtime " + "pulls single XCom rows, so a mapped upstream's combined output is not " + "supported -- use .expand() on the stub to consume it per element" + ) + return True + if isinstance(value, XComArg): + raise ValueError( + f"@task.stub task {task_id!r} parameter {name!r} received a " + f"{type(value).__name__}; only direct upstream task outputs can cross the " + "language boundary -- .map()/.zip()/.concat() results are not supported" + ) + return False + + +def _build_arg_bindings( + python_callable: Callable, + op_args: Collection[Any], + op_kwargs: Mapping[str, Any], + task_id: str, +) -> list[dict[str, Any]] | None: + """ + Bind the TaskFlow call arguments to the stub signature and build the ordered arg spec. + + Each spec entry is a plain dict matching one variant of the execution API's + ``TaskArgBinding`` union: an ``XComArgBinding`` (``kind="xcom"``) for upstream TaskFlow + outputs, or a ``LiteralArgBinding`` (``kind="literal"``) for everything else. ``name`` is + always the stub function's parameter name, so a foreign runtime can bind by name (e.g. the + Go SDK's ``sdk.TaskInput`` struct fields) in addition to the existing positional order. + Returns ``None`` for argless calls: the binding contract (including the signature checks + below) applies only once a TaskFlow call actually passes arguments, so pre-TaskFlow stub + Dags whose call arguments were always ignored keep parsing. + """ + if not op_args and not op_kwargs: + return None + + signature = inspect.signature(python_callable) + _validate_stub_signature(signature, task_id) + + bound = signature.bind(*op_args, **op_kwargs) + explicitly_bound = set(bound.arguments) + bound.apply_defaults() + + annotations = _resolve_param_annotations(python_callable, signature) + + spec: list[dict[str, Any]] = [] + for name in signature.parameters: + value = bound.arguments[name] + value_schema = _infer_value_schema(annotations[name]) + if _validate_xcom_value(value, task_id, name): + xcom_entry: dict[str, Any] = {"name": name, "kind": "xcom", "task_id": value.operator.task_id} + if value_schema is not None: + xcom_entry["value_schema"] = value_schema + spec.append(xcom_entry) + continue + _ensure_json_literal(value, task_id, name) + entry: dict[str, Any] = {"name": name, "kind": "literal", "value": value} + if value_schema is not None: + # Key omission (never ``None``) is the wire contract for "unconstrained": + # ti_run responds with ``exclude_unset``, so an absent key stays absent. + entry["value_schema"] = value_schema + if name not in explicitly_bound: + entry["from_default"] = True + spec.append(entry) + return spec + + +def _build_mapped_arg_binding_params( + python_callable: Callable, + *, + partial_op_kwargs: Mapping[str, Any], + expand_input: Any, + task_id: str, +) -> list[dict[str, Any]] | None: + """ + Build the ordered per-parameter metadata for a mapped (``.expand()``) stub task. + + Per-map-index values only resolve at run time, so unlike ``_build_arg_bindings`` this + captures what the server-side derivation cannot recover from the serialized Dag alone: + the declaration order the wire contract promises, defaults for parameters no kwarg + covers, and each parameter's value schema. Returns ``None`` for parameterless stubs + (the legacy fan-out shape whose call args were always ignored keeps parsing). + """ + signature = inspect.signature(python_callable) + if not signature.parameters: + return None + _validate_stub_signature(signature, task_id) + if not isinstance(expand_input.value, Mapping): + # expand_kwargs() carries a list (or upstream XCom) of kwarg dicts whose + # parameter names are unknowable at parse time. + raise ValueError( + f"@task.stub task {task_id!r} does not support expand_kwargs(); the parameter " + "binding must be derivable at parse time, so use .expand() with explicit kwargs" + ) + expand_kwargs = expand_input.value + + try: + bound = signature.bind(**{**partial_op_kwargs, **expand_kwargs}) + except TypeError as e: + raise ValueError(f"@task.stub task {task_id!r} TaskFlow mapping does not bind to its signature: {e}") + bound.apply_defaults() + + annotations = _resolve_param_annotations(python_callable, signature) + + params: list[dict[str, Any]] = [] + for name in signature.parameters: + entry: dict[str, Any] = {"name": name} + if (value_schema := _infer_value_schema(annotations[name])) is not None: + entry["value_schema"] = value_schema + if name in expand_kwargs: + # The whole expanded collection ships through XCom/serialization; an upstream + # output is consumed per element, so a mapped upstream is fine here. + if not _validate_xcom_value(expand_kwargs[name], task_id, name, allow_mapped_upstream=True): + _ensure_json_literal(expand_kwargs[name], task_id, name) + elif name in partial_op_kwargs: + if not _validate_xcom_value(partial_op_kwargs[name], task_id, name): + _ensure_json_literal(partial_op_kwargs[name], task_id, name) + else: + default = bound.arguments[name] + _ensure_json_literal(default, task_id, name) + entry["default"] = default + params.append(entry) + return params + + class _StubOperator(DecoratedOperator): custom_operator_name: str = "@task.stub" @@ -60,10 +353,10 @@ def __init__( module = ast.parse(self.get_python_source()) if len(module.body) != 1: - raise RuntimeError("Expected a single statement") + raise ValueError("Expected a single statement") fn = module.body[0] if not isinstance(fn, ast.FunctionDef): - raise RuntimeError("Expected a single sync function") + raise ValueError("Expected a single sync function") for stmt in fn.body: if isinstance(stmt, ast.Pass): continue @@ -75,7 +368,46 @@ def __init__( f"Functions passed to @task.stub must be an empty function (`pass`, or `...` only) (got {stmt})" ) - ... + # Bind the TaskFlow call to the *original* signature (DecoratedOperator mangles context + # key defaults, which stubs reject anyway) and persist the ordered arg spec so the + # execution API can hand it to the foreign runtime via StartupDetails. + self._arg_bindings = _build_arg_bindings(python_callable, self.op_args, self.op_kwargs, self.task_id) + + # Direct .expand() on the stub needs no parse-time spec (ti_run derives per-map-index + # bindings from the serialized expand input), but a mapped task group creates + # per-map-index instances of the tasks inside it with no expand input of their own, + # so their arg values are unresolvable both here and server-side. + in_mapped_group = getattr(self, "get_closest_mapped_task_group", lambda: None)() is not None + if self._arg_bindings is not None and in_mapped_group: + raise ValueError( + f"@task.stub task {self.task_id!r} passes TaskFlow call arguments inside a mapped " + "task group; the captured spec cannot carry values that resolve per map index at " + "runtime, so stub tasks with arguments are not supported under a task group's " + ".expand()" + ) + + @classmethod + def get_serialized_fields(cls): + return super().get_serialized_fields() | {"_arg_bindings"} + + @classmethod + def get_mapped_serialized_fields(cls, mapped_op: Any) -> dict[str, Any]: + """ + Extra serialized fields for the mapped (``.expand()``) form of this operator. + + Called by the core Dag serializer (Airflow 3.4+) while ``python_callable`` is + still the real function; older cores never call it, so mapped stubs there keep + the legacy ignored-args behavior. + """ + params = _build_mapped_arg_binding_params( + mapped_op.python_callable, + partial_op_kwargs=mapped_op.partial_kwargs.get("op_kwargs") or {}, + expand_input=mapped_op._get_specified_expand_input(), + task_id=mapped_op.task_id, + ) + if params is None: + return {} + return {"_mapped_arg_binding_params": params} def execute(self, context: Context) -> Any: raise RuntimeError( @@ -96,6 +428,10 @@ def stub( Stub tasks exist in the Dag graph only, but the execution must happen in an external environment via the Task Execution Interface. + Stub functions may declare parameters and be called TaskFlow-style with upstream task + outputs or JSON-serializable literals; the resulting argument-binding spec (parameter + names, value schemas, and values, in declaration order) is delivered to the foreign + runtime, which binds the values onto the native task function. """ return task_decorator_factory( decorated_operator_class=_StubOperator, diff --git a/providers/standard/tests/unit/standard/decorators/test_stub.py b/providers/standard/tests/unit/standard/decorators/test_stub.py index 2a17c3fdd82c1..b73c82f38d2f6 100644 --- a/providers/standard/tests/unit/standard/decorators/test_stub.py +++ b/providers/standard/tests/unit/standard/decorators/test_stub.py @@ -17,10 +17,16 @@ from __future__ import annotations import contextlib +import datetime +import typing +from typing import Any +from unittest import mock +import pendulum import pytest -from airflow.providers.standard.decorators.stub import stub +from airflow.providers.common.compat.sdk import DAG, task_group +from airflow.providers.standard.decorators.stub import _infer_value_schema, _StubOperator, stub from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS @@ -69,3 +75,456 @@ def test_stub_rejects_retry_policy(): def test_stub_allows_retries(): stub(fn_pass, retries=5)() + + +def fn_extract(): ... + + +def fn_transform(country: str, extracted: dict, retries_num: int = 3): ... + + +def fn_untyped(a, b): ... + + +def fn_varargs(*args): ... + + +def fn_kwonly_varkw(**kwargs): ... + + +def fn_context_key(ti): ... + + +class TestStubTaskflowArgs: + """The TaskFlow call on a stub captures the ordered positional-arg spec (``_arg_bindings``).""" + + def test_literal_and_xcom_spec(self): + with DAG(dag_id="d"): + extracted = stub(fn_extract)() + result = stub(fn_transform)("uk", extracted) + + op = result.operator + assert op._arg_bindings == [ + {"name": "country", "kind": "literal", "value_schema": {"type": "string"}, "value": "uk"}, + { + "name": "extracted", + "kind": "xcom", + "value_schema": {"type": "object", "additionalProperties": True}, + "task_id": "fn_extract", + }, + { + "name": "retries_num", + "kind": "literal", + "value_schema": {"type": "integer", "format": "int64"}, + "value": 3, + "from_default": True, + }, + ] + assert op.upstream_task_ids == {"fn_extract"} + + def test_kwargs_normalize_to_declaration_order(self): + with DAG(dag_id="d"): + extracted = stub(fn_extract)() + result = stub(fn_transform)(extracted=extracted, country="fr", retries_num=7) + + assert result.operator._arg_bindings == [ + {"name": "country", "kind": "literal", "value_schema": {"type": "string"}, "value": "fr"}, + { + "name": "extracted", + "kind": "xcom", + "value_schema": {"type": "object", "additionalProperties": True}, + "task_id": "fn_extract", + }, + { + "name": "retries_num", + "kind": "literal", + "value_schema": {"type": "integer", "format": "int64"}, + "value": 7, + }, + ] + + def test_explicitly_passing_the_default_value_is_not_from_default(self): + """The flag tracks provenance, not value equality: an author-passed argument is explicit + even when it equals the signature default, so keyword-style consumers must still claim it.""" + with DAG(dag_id="d"): + extracted = stub(fn_extract)() + result = stub(fn_transform)("uk", extracted, retries_num=3) + + assert result.operator._arg_bindings[2] == { + "name": "retries_num", + "kind": "literal", + "value_schema": {"type": "integer", "format": "int64"}, + "value": 3, + } + + def test_custom_xcom_key_rejected(self): + with DAG(dag_id="d"): + extracted = stub(fn_extract)() + with pytest.raises(ValueError, match="indexing an output by a custom key"): + stub(fn_transform)("uk", extracted["part"]) + + def test_zero_param_stub_has_no_spec(self): + assert stub(fn_pass)().operator._arg_bindings is None + + def test_untyped_params_omit_value_schema(self): + """Key absence (never ``None``) is the wire contract for an unconstrained argument.""" + with DAG(dag_id="d"): + result = stub(fn_untyped)(1, "x") + + assert result.operator._arg_bindings == [ + {"name": "a", "kind": "literal", "value": 1}, + {"name": "b", "kind": "literal", "value": "x"}, + ] + + def test_unresolvable_annotation_omits_value_schema(self): + def fn(x): ... + + fn.__annotations__ = {"x": "NotARealType"} + with DAG(dag_id="d"): + result = stub(fn)("v") + + assert result.operator._arg_bindings == [{"name": "x", "kind": "literal", "value": "v"}] + + def test_varargs_rejected(self): + with pytest.raises(ValueError, match="fixed number of parameters"): + stub(fn_varargs)(1, 2) + + def test_varkw_rejected(self): + with pytest.raises(ValueError, match="fixed number of parameters"): + stub(fn_kwonly_varkw)(x=1) + + def test_context_key_param_rejected(self): + with pytest.raises(ValueError, match="is an Airflow context key"): + stub(fn_context_key)(1) + + @pytest.mark.parametrize("fn", [fn_varargs, fn_kwonly_varkw, fn_context_key], ids=lambda f: f.__name__) + def test_argless_call_skips_signature_checks(self, fn): + """Pre-TaskFlow stub Dags never passed arguments; their signatures must keep parsing.""" + assert stub(fn)().operator._arg_bindings is None + + def test_argless_call_captures_no_spec_for_defaulted_params(self): + def fn(limit: int = 10): ... + + assert stub(fn)().operator._arg_bindings is None + + def test_non_json_literal_rejected(self): + with DAG(dag_id="d"), pytest.raises(ValueError, match="not JSON-serializable"): + stub(fn_transform)("uk", object()) + + def test_nan_literal_rejected(self): + with DAG(dag_id="d"), pytest.raises(ValueError, match="not JSON-serializable"): + stub(fn_transform)("uk", {"ratio": float("nan")}) + + def test_mapped_xcom_arg_rejected(self): + with DAG(dag_id="d"): + extracted = stub(fn_extract)() + with pytest.raises(ValueError, match="only direct upstream task outputs"): + stub(fn_transform)("uk", extracted.map(lambda v: v)) + + def test_mapped_upstream_aggregated_output_rejected(self): + def fn_produce(n: int): ... + + with DAG(dag_id="d"): + vals = stub(fn_produce).expand(n=[1, 2]) + with pytest.raises(ValueError, match="aggregated output of the mapped task"): + stub(fn_transform)("uk", vals) + + def test_arg_bindings_survive_dag_serialization_round_trip(self): + """The captured spec must survive whichever core serializer the provider runs against.""" + try: + from airflow.serialization.serialized_objects import DagSerialization + except ImportError: # Airflow 2 exposes the round-trip API on SerializedDAG + from airflow.serialization.serialized_objects import SerializedDAG as DagSerialization + + with DAG(dag_id="d") as dag: + extracted = stub(fn_extract)() + stub(fn_transform)("uk", extracted) + + round_tripped = DagSerialization.from_dict(DagSerialization.to_dict(dag)) + assert round_tripped.task_dict["fn_transform"]._arg_bindings == [ + {"name": "country", "kind": "literal", "value_schema": {"type": "string"}, "value": "uk"}, + { + "name": "extracted", + "kind": "xcom", + "value_schema": {"type": "object", "additionalProperties": True}, + "task_id": "fn_extract", + }, + { + "name": "retries_num", + "kind": "literal", + "value_schema": {"type": "integer", "format": "int64"}, + "value": 3, + "from_default": True, + }, + ] + + def test_expand_builds_mapped_stub_without_parse_time_bindings(self): + with DAG(dag_id="d"): + result = stub(fn_transform).expand(country=["uk", "fr"], extracted=[{}, {}]) + # op_kwargs_expand_input/partial_kwargs (not is_mapped) so the assertions also + # hold on the Airflow 2.x MappedOperator, which the provider still supports. + assert result.operator.op_kwargs_expand_input.value == { + "country": ["uk", "fr"], + "extracted": [{}, {}], + } + assert "_arg_bindings" not in result.operator.partial_kwargs + + def test_stub_with_args_inside_mapped_task_group_rejected(self): + @task_group + def group(n): + stub(fn_transform)("uk", {}) + + with DAG(dag_id="d"): + with pytest.raises(ValueError, match="mapped task group"): + group.expand(n=[1, 2]) + + def test_argless_stub_inside_mapped_task_group_allowed(self): + @task_group + def group(n): + stub(fn_extract)() + + with DAG(dag_id="d"): + group.expand(n=[1, 2]) + + +class TestMappedStubArgBindingParams: + """The serializer hook captures ordered per-parameter metadata for mapped stubs.""" + + def get_hook_fields(self, operator): + return _StubOperator.get_mapped_serialized_fields(operator) + + def test_params_follow_declaration_order_with_defaults_and_schemas(self): + with DAG(dag_id="d"): + # The partial() kwarg is declared *after* the expanded one: the captured + # order must come from the signature, not from the call sites. + result = stub(fn_transform).partial(extracted={"a": 1}).expand(country=["uk", "fr"]) + + assert self.get_hook_fields(result.operator) == { + "_mapped_arg_binding_params": [ + {"name": "country", "value_schema": {"type": "string"}}, + {"name": "extracted", "value_schema": {"type": "object", "additionalProperties": True}}, + { + "name": "retries_num", + "value_schema": {"type": "integer", "format": "int64"}, + "default": 3, + }, + ] + } + + def test_none_default_is_captured_by_key_presence(self): + def fn(x: str, y=None): ... + + with DAG(dag_id="d"): + result = stub(fn).expand(x=["a"]) + + params = self.get_hook_fields(result.operator)["_mapped_arg_binding_params"] + assert params[1] == {"name": "y", "default": None} + + def test_untyped_params_omit_value_schema(self): + with DAG(dag_id="d"): + result = stub(fn_untyped).expand(a=[1], b=[2]) + + assert self.get_hook_fields(result.operator) == { + "_mapped_arg_binding_params": [{"name": "a"}, {"name": "b"}] + } + + def test_parameterless_stub_captures_nothing(self): + with DAG(dag_id="d"): + result = stub(fn_extract).expand_kwargs([{}]) + + assert self.get_hook_fields(result.operator) == {} + + def test_expand_kwargs_rejected_for_parameterful_stub(self): + with DAG(dag_id="d"): + result = stub(fn_transform).expand_kwargs([{"country": "uk", "extracted": {}}]) + + with pytest.raises(ValueError, match="does not support expand_kwargs"): + self.get_hook_fields(result.operator) + + def test_missing_required_parameter_rejected(self): + with DAG(dag_id="d"): + result = stub(fn_transform).expand(country=["uk"]) + + with pytest.raises(ValueError, match="does not bind to its signature"): + self.get_hook_fields(result.operator) + + def test_partial_kwarg_over_mapped_upstream_rejected(self): + def fn_produce(n: int): ... + + with DAG(dag_id="d"): + vals = stub(fn_produce).expand(n=[1, 2]) + result = stub(fn_transform).partial(extracted=vals).expand(country=["uk"]) + + with pytest.raises(ValueError, match="aggregated output of the mapped task"): + self.get_hook_fields(result.operator) + + def test_expand_over_mapped_upstream_allowed(self): + def fn_produce(n: int): ... + + with DAG(dag_id="d"): + vals = stub(fn_produce).expand(n=[1, 2]) + result = stub(fn_transform).partial(country="uk").expand(extracted=vals) + + params = self.get_hook_fields(result.operator)["_mapped_arg_binding_params"] + assert [p["name"] for p in params] == ["country", "extracted", "retries_num"] + + def test_non_json_expand_literal_rejected(self): + with DAG(dag_id="d"): + result = stub(fn_transform).partial(country="uk").expand(extracted=[object()]) + + with pytest.raises(ValueError, match="not JSON-serializable"): + self.get_hook_fields(result.operator) + + def test_non_json_needed_default_rejected(self): + not_jsonable = object() + + def fn(x: str, y=not_jsonable): ... + + with DAG(dag_id="d"): + result = stub(fn).expand(x=["a"]) + + with pytest.raises(ValueError, match="not JSON-serializable"): + self.get_hook_fields(result.operator) + + def test_mapped_stub_inside_mapped_task_group_unconstructible(self): + """The SDK bans expansion inside an expanded group outright, so no hook guard is needed.""" + + @task_group + def group(n): + stub(fn_transform).partial(country="uk").expand(extracted=[{}]) + + with DAG(dag_id="d"): + with pytest.raises(NotImplementedError, match="expansion in an expanded task group"): + group.expand(n=[1, 2]) + + +@pytest.mark.parametrize( + ("annotation", "expected"), + [ + pytest.param(str, {"type": "string"}, id="str"), + pytest.param(bool, {"type": "boolean"}, id="bool"), + pytest.param(int, {"type": "integer", "format": "int64"}, id="int"), + pytest.param(float, {"type": "number", "format": "double"}, id="float"), + pytest.param(dict, {"type": "object", "additionalProperties": True}, id="dict"), + pytest.param( + dict[str, int], + {"type": "object", "additionalProperties": {"type": "integer", "format": "int64"}}, + id="dict-parameterized", + ), + pytest.param( + typing.Mapping[str, int], + {"type": "object", "additionalProperties": {"type": "integer", "format": "int64"}}, + id="mapping", + ), + pytest.param(list, {"type": "array", "items": {}}, id="list"), + pytest.param( + list[int], + {"type": "array", "items": {"type": "integer", "format": "int64"}}, + id="list-parameterized", + ), + pytest.param(tuple, {"type": "array", "items": {}}, id="tuple"), + pytest.param(set, {"type": "array", "items": {}, "uniqueItems": True}, id="set"), + pytest.param( + typing.Sequence[int], + {"type": "array", "items": {"type": "integer", "format": "int64"}}, + id="sequence", + ), + pytest.param(datetime.datetime, {"type": "string", "format": "date-time"}, id="datetime"), + pytest.param(datetime.date, {"type": "string", "format": "date"}, id="date"), + pytest.param(datetime.time, {"type": "string", "format": "time"}, id="time"), + pytest.param(datetime.timedelta, {"type": "string", "format": "duration"}, id="timedelta"), + pytest.param(bytes, {"type": "string", "format": "binary"}, id="bytes"), + pytest.param( + typing.Literal["a", "b"], + {"type": "string", "enum": ["a", "b"]}, + id="literal", + ), + pytest.param(Any, None, id="any"), + pytest.param(None, None, id="none"), + pytest.param(type(None), None, id="nonetype"), + pytest.param( + pendulum.DateTime, + {"type": "string", "format": "date-time"}, + id="pendulum-datetime", + ), + pytest.param( + pendulum.DateTime | None, + {"anyOf": [{"type": "string", "format": "date-time"}, {"type": "null"}]}, + id="optional-pendulum-datetime", + ), + pytest.param( + list[pendulum.DateTime], + {"type": "array", "items": {"type": "string", "format": "date-time"}}, + id="list-pendulum-datetime", + ), + pytest.param(pendulum.Duration, {"type": "string", "format": "duration"}, id="pendulum-duration"), + pytest.param( + typing.Optional[str], # noqa: UP045 -- legacy form on purpose + {"anyOf": [{"type": "string"}, {"type": "null"}]}, + id="optional-str", + ), + pytest.param( + typing.Union[int, str], # noqa: UP007 -- legacy form on purpose + {"anyOf": [{"type": "integer", "format": "int64"}, {"type": "string"}]}, + id="union", + ), + pytest.param(str | None, {"anyOf": [{"type": "string"}, {"type": "null"}]}, id="pep604-optional"), + pytest.param( + int | None, + {"anyOf": [{"type": "integer", "format": "int64"}, {"type": "null"}]}, + id="optional-int", + ), + pytest.param( + datetime.datetime | None, + {"anyOf": [{"type": "string", "format": "date-time"}, {"type": "null"}]}, + id="optional-datetime", + ), + pytest.param( + dict | bool, + {"anyOf": [{"type": "object", "additionalProperties": True}, {"type": "boolean"}]}, + id="union-dict-bool", + ), + pytest.param( + str | int | None, + {"anyOf": [{"type": "string"}, {"type": "integer", "format": "int64"}, {"type": "null"}]}, + id="union-with-null", + ), + pytest.param(list | tuple, {"type": "array", "items": {}}, id="union-dedupes-equal-members"), + pytest.param( + datetime.datetime | str, + {"anyOf": [{"type": "string", "format": "date-time"}, {"type": "string"}]}, + id="mixed-format-union-keeps-both", + ), + pytest.param( + str | contextlib.AbstractContextManager, + None, + id="union-unclassifiable-member", + ), + pytest.param(contextlib.AbstractContextManager, None, id="custom-class"), + pytest.param(typing.Callable[[int], str], None, id="callable-invalid-for-json-schema"), + pytest.param( + pendulum.DateTime | contextlib.AbstractContextManager, + None, + id="union-temporal-and-unclassifiable", + ), + ], +) +def test_infer_value_schema(annotation, expected): + assert _infer_value_schema(annotation) == expected + + +@mock.patch("airflow.providers.standard.decorators.stub.TypeAdapter", None) +def test_infer_value_schema_without_pydantic(): + assert _infer_value_schema(str) is None + + +def test_infer_value_schema_cache_returns_isolated_copies(): + first = _infer_value_schema(dict) + second = _infer_value_schema(dict) + assert first == second + assert first is not second, "callers embed and serialize the fragment, so it must not alias the cache" + + +def test_infer_value_schema_unhashable_annotation_generates_uncached(): + annotation = typing.Annotated[int, {"unhashable": True}] + assert _infer_value_schema(annotation) == {"type": "integer", "format": "int64"} diff --git a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index 8e5bfc1d076e5..739ac5ef66595 100644 --- a/task-sdk/src/airflow/sdk/api/datamodels/_generated.py +++ b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py @@ -27,7 +27,7 @@ from pydantic import AwareDatetime, BaseModel, ConfigDict, Field, JsonValue, RootModel -API_VERSION: Final[str] = "2026-06-30" +API_VERSION: Final[str] = "2026-10-30" class AssetAliasReferenceAssetEventDagRun(BaseModel): @@ -608,6 +608,10 @@ class DagAttributeTypes(str, Enum): TASK_GROUP = "taskgroup" +class ArgValueSchema(RootModel[dict[str, JsonValue]]): + root: dict[str, JsonValue] + + class AssetReferenceAssetEventDagRun(BaseModel): """ Schema for AssetModel used in AssetEventDagRunReference. @@ -697,6 +701,18 @@ class HTTPValidationError(BaseModel): detail: Annotated[list[ValidationError] | None, Field(title="Detail")] = None +class LiteralArgBinding(BaseModel): + """ + One positional stub-task argument carrying an inline literal from the Dag file. + """ + + kind: Annotated[Literal["literal"], Field(title="Kind")] + name: Annotated[str, Field(title="Name")] + value_schema: ArgValueSchema | None = None + value: JsonValue | None = None + from_default: Annotated[bool | None, Field(title="From Default")] = False + + class TITerminalStatePayload(BaseModel): """ Schema for updating TaskInstance to a terminal state except SUCCESS state. @@ -710,6 +726,19 @@ class TITerminalStatePayload(BaseModel): rendered_map_index: Annotated[str | None, Field(title="Rendered Map Index")] = None +class XComArgBinding(BaseModel): + """ + One positional stub-task argument pulled from an upstream task's XCom. + """ + + kind: Annotated[Literal["xcom"], Field(title="Kind")] + name: Annotated[str, Field(title="Name")] + value_schema: ArgValueSchema | None = None + task_id: Annotated[str, Field(title="Task Id")] + map_index: Annotated[int | None, Field(title="Map Index")] = -1 + element_index: Annotated[int | None, Field(title="Element Index")] = None + + class AssetEventDagRunReference(BaseModel): """ Schema for AssetEvent model used in DagRun. @@ -782,6 +811,10 @@ class DagRun(BaseModel): team_name: Annotated[str | None, Field(title="Team Name")] = None +class TaskArgBinding(RootModel[XComArgBinding | LiteralArgBinding]): + root: Annotated[XComArgBinding | LiteralArgBinding, Field(discriminator="kind", title="TaskArgBinding")] + + class TIRunContext(BaseModel): """ Response schema for TaskInstance run context. @@ -797,3 +830,4 @@ class TIRunContext(BaseModel): xcom_keys_to_clear: Annotated[list[str] | None, Field(title="Xcom Keys To Clear")] = None should_retry: Annotated[bool | None, Field(title="Should Retry")] = False start_date: Annotated[AwareDatetime | None, Field(title="Start Date")] = None + arg_bindings: Annotated[list[TaskArgBinding] | None, Field(title="Arg Bindings")] = None diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json index 0ec8fe4e49a1c..4c2f503dc828e 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -1,8 +1,15 @@ { "$schema": "https://json-schema.org/draft/2020-12/schema", - "api_version": "2026-06-16", + "api_version": "2026-10-30", "description": "Apache Airflow SDK Supervisor Schema", "$defs": { + "ArgValueSchema": { + "additionalProperties": { + "$ref": "#/$defs/JsonValue" + }, + "title": "ArgValueSchema", + "type": "object" + }, "AssetAliasReferenceAssetEventDagRun": { "additionalProperties": false, "description": "Schema for AssetAliasModel used in AssetEventDagRunReference.", @@ -4563,6 +4570,124 @@ "title": "ConnectionResponse", "type": "object" }, + "LiteralArgBinding": { + "description": "One positional stub-task argument carrying an inline literal from the Dag file.", + "properties": { + "kind": { + "const": "literal", + "title": "Kind", + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + }, + "value_schema": { + "anyOf": [ + { + "$ref": "#/$defs/ArgValueSchema" + }, + { + "type": "null" + } + ], + "default": null + }, + "value": { + "anyOf": [ + { + "$ref": "#/$defs/JsonValue" + }, + { + "type": "null" + } + ], + "default": null + }, + "from_default": { + "default": false, + "title": "From Default", + "type": "boolean" + } + }, + "required": [ + "kind", + "name" + ], + "title": "LiteralArgBinding", + "type": "object" + }, + "TaskArgBinding": { + "discriminator": { + "mapping": { + "literal": "#/$defs/LiteralArgBinding", + "xcom": "#/$defs/XComArgBinding" + }, + "propertyName": "kind" + }, + "oneOf": [ + { + "$ref": "#/$defs/XComArgBinding" + }, + { + "$ref": "#/$defs/LiteralArgBinding" + } + ], + "title": "TaskArgBinding" + }, + "XComArgBinding": { + "description": "One positional stub-task argument pulled from an upstream task's XCom.", + "properties": { + "kind": { + "const": "xcom", + "title": "Kind", + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + }, + "value_schema": { + "anyOf": [ + { + "$ref": "#/$defs/ArgValueSchema" + }, + { + "type": "null" + } + ], + "default": null + }, + "task_id": { + "title": "Task Id", + "type": "string" + }, + "map_index": { + "default": -1, + "title": "Map Index", + "type": "integer" + }, + "element_index": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Element Index" + } + }, + "required": [ + "kind", + "name", + "task_id" + ], + "title": "XComArgBinding", + "type": "object" + }, "AssetEventDagRunReference": { "additionalProperties": false, "description": "Schema for AssetEvent model used in DagRun.", @@ -4926,6 +5051,21 @@ ], "default": null, "title": "Start Date" + }, + "arg_bindings": { + "anyOf": [ + { + "items": { + "$ref": "#/$defs/TaskArgBinding" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Arg Bindings" } }, "required": [ diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py b/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py index 9491a8993fdc3..7e5ce93f86bdc 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py +++ b/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py @@ -37,8 +37,13 @@ def get_bundle() -> VersionBundle: """ from cadwyn import HeadVersion, Version, VersionBundle + from airflow.sdk.execution_time.schema.versions.v2026_10_30 import ( + AddArgBindingsToSupervisorTIRunContext, + ) + return VersionBundle( HeadVersion(), + Version("2026-10-30", AddArgBindingsToSupervisorTIRunContext), Version("2026-06-16"), ) diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py b/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py new file mode 100644 index 0000000000000..e6b93f5dea805 --- /dev/null +++ b/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py @@ -0,0 +1,36 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +from cadwyn import VersionChange, schema + +from airflow.sdk.api.datamodels._generated import TIRunContext + + +class AddArgBindingsToSupervisorTIRunContext(VersionChange): + """ + Add the ``arg_bindings`` argument-binding spec for stub (foreign-runtime) tasks. + + Each entry is a discriminated union of ``XComArgBinding`` and ``LiteralArgBinding`` + keyed on ``kind``. The supervisor-schema mirror of the execution API's + ``AddArgBindingsToTIRunContext``, named apart so the two migrations are not confused. + """ + + description = __doc__ + + instructions_to_migrate_to_previous_version = (schema(TIRunContext).field("arg_bindings").didnt_exist,) diff --git a/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py index 05218aded3d3b..8e2379e2fdbba 100644 --- a/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py +++ b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py @@ -105,12 +105,9 @@ def _backfill_sentry_trace(request): class TestSchemaVersionMigratorDowngrade: """ Drive the downgrade direction against a mock bundle so we can pin - *field-level* migration behaviour. The real supervisor bundle has - no schema-level migrations on the IPC bodies yet, so it would no-op - every version -- which proves nothing about the migration chain. - The mock bundle's mechanism is identical to the real one, so what - we prove about it applies to the real bundle the moment a - ``schema(...)`` instruction lands. + *field-level* migration behaviour independent of the real bundle's + contents. The real bundle's ``arg_bindings`` migration is covered by + :class:`TestRealBundleArgBindingsDowngrade` below. """ @pytest.fixture @@ -369,3 +366,103 @@ def test_accessing_bundle_loads_cadwyn(self): "assert 'cadwyn' in sys.modules, 'cadwyn should load when the bundle is accessed'" ) subprocess.run([sys.executable, "-c", code], check=True, capture_output=True, text=True) + + +class TestRealBundleArgBindingsDowngrade: + """ + Drive the *real* supervisor bundle through the ``arg_bindings`` migration. + + ``AddArgBindingsToSupervisorTIRunContext`` is the bundle's first ``schema(...)`` + instruction on a model *nested* inside a registered body + (``StartupDetails.ti_context``); this pins that the downgrade + re-validation strips the nested field on the wire for a runtime + pinned to the previous version, and keeps it at head. + """ + + @pytest.fixture + def startup_details(self): + import datetime + import uuid + + from airflow.sdk.api.datamodels._generated import ( + BundleInfo, + DagRun, + DagRunState, + DagRunType, + TaskInstance, + TIRunContext, + ) + from airflow.sdk.execution_time.comms import StartupDetails + + now = datetime.datetime.now(datetime.timezone.utc) + return StartupDetails( + ti=TaskInstance( + id=uuid.uuid4(), + task_id="transform", + dag_id="d", + run_id="r", + try_number=1, + dag_version_id=uuid.uuid4(), + ), + dag_rel_path="d.py", + bundle_info=BundleInfo(name="b", version=None), + start_date=now, + ti_context=TIRunContext( + dag_run=DagRun( + dag_id="d", + run_id="r", + logical_date=now, + start_date=now, + run_type=DagRunType.MANUAL, + state=DagRunState.RUNNING, + run_after=now, + consumed_asset_events=[], + ), + max_tries=1, + arg_bindings=[ + # No value_schema: the unconstrained ("any") case rides through the migrator too. + {"name": "country", "kind": "literal", "value": "uk"}, + { + "name": "extracted", + "kind": "xcom", + "value_schema": {"type": "object"}, + "task_id": "extract", + }, + { + "name": "limit", + "kind": "literal", + "value_schema": {"type": "integer", "format": "int64"}, + "value": 10, + "from_default": True, + }, + ], + ), + sentry_integration="", + ) + + @pytest.fixture + def real_migrator(self) -> SchemaVersionMigrator: + return get_schema_version_migrator() + + def test_downgrade_strips_arg_bindings_for_previous_version(self, real_migrator, startup_details): + out = real_migrator.downgrade(startup_details, "2026-06-16").model_dump() + assert "arg_bindings" not in out["ti_context"] + + def test_head_version_keeps_arg_bindings(self, real_migrator, startup_details): + from airflow.sdk.api.datamodels._generated import LiteralArgBinding, XComArgBinding + + out = real_migrator.downgrade(startup_details, "2026-10-30") + assert out.ti_context.arg_bindings is not None + literal, xcom, defaulted = (a.root for a in out.ti_context.arg_bindings) + assert isinstance(literal, LiteralArgBinding) + assert literal.value == "uk" + assert literal.name == "country" + assert literal.from_default is False + assert literal.value_schema is None + assert isinstance(xcom, XComArgBinding) + assert xcom.task_id == "extract" + assert xcom.name == "extracted" + assert xcom.value_schema.root == {"type": "object"} + assert isinstance(defaulted, LiteralArgBinding) + assert defaulted.from_default is True + assert defaulted.value_schema.root == {"type": "integer", "format": "int64"} diff --git a/ts-sdk/src/generated/supervisor.ts b/ts-sdk/src/generated/supervisor.ts index ab2632831ab95..31d44a52b18c5 100644 --- a/ts-sdk/src/generated/supervisor.ts +++ b/ts-sdk/src/generated/supervisor.ts @@ -22,17 +22,17 @@ // // Re-run with: pnpm run generate:supervisor +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "JsonValue". + */ +export type JsonValue = unknown; export type Name = string; export type Id = number; export type Timestamp = string; export type Extra = { [k: string]: JsonValue; } | null; -/** - * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema - * via the `definition` "JsonValue". - */ -export type JsonValue = unknown; export type Name1 = string; export type Uri = string; export type Group = string; @@ -245,6 +245,20 @@ export type NextKwargs1 = export type XcomKeysToClear = string[]; export type ShouldRetry = boolean; export type StartDate2 = string | null; +export type ArgBindings = TaskArgBinding[] | null; +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "TaskArgBinding". + */ +export type TaskArgBinding = XComArgBinding | LiteralArgBinding; +export type Kind = "xcom"; +export type Name8 = string; +export type TaskId1 = string; +export type MapIndex1 = number; +export type ElementIndex = number | null; +export type Kind1 = "literal"; +export type Name9 = string; +export type FromDefault = boolean; export type Type13 = "TaskCallbackRequest"; export type Filepath2 = string; export type BundleName3 = string; @@ -310,7 +324,7 @@ export type NextKwargs2 = { } | null; export type RenderedMapIndex1 = string | null; export type Type20 = "DeferTask"; -export type Name8 = string; +export type Name10 = string; export type Key1 = string; export type Type21 = "DeleteAssetStateStoreByName"; export type Uri5 = string; @@ -324,8 +338,8 @@ export type Type24 = "DeleteVariable"; export type Key5 = string; export type DagId6 = string; export type RunId5 = string; -export type TaskId1 = string; -export type MapIndex1 = number | null; +export type TaskId2 = string; +export type MapIndex2 = number | null; export type Type25 = "DeleteXCom"; /** * Error types used in the API client. @@ -362,11 +376,11 @@ export type ErrorType1 = | "PERMISSION_DENIED" | "GENERIC_ERROR" | "API_SERVER_ERROR"; -export type Name9 = string; +export type Name11 = string; export type Type27 = "GetAssetByName"; export type Uri6 = string; export type Type28 = "GetAssetByUri"; -export type Name10 = string | null; +export type Name12 = string | null; export type Uri7 = string | null; export type After = string | null; export type Before = string | null; @@ -389,7 +403,7 @@ export type Extra8 = { [k: string]: string; } | null; export type Type30 = "GetAssetEventByAssetAlias"; -export type Name11 = string; +export type Name13 = string; export type Key6 = string; export type Type31 = "GetAssetStateStoreByName"; export type Uri8 = string; @@ -421,12 +435,12 @@ export type LogicalDate3 = string; export type State3 = string | null; export type Type41 = "GetPreviousDagRun"; export type DagId12 = string; -export type TaskId2 = string; +export type TaskId3 = string; export type LogicalDate4 = string | null; -export type MapIndex2 = number; +export type MapIndex3 = number; export type Type42 = "GetPreviousTI"; export type DagId13 = string; -export type MapIndex3 = number | null; +export type MapIndex4 = number | null; export type TaskIds = string[] | null; export type TaskGroupId = string | null; export type LogicalDates1 = string[] | null; @@ -443,7 +457,7 @@ export type TiId6 = string; export type Key8 = string; export type Type46 = "GetTaskStateStore"; export type DagId15 = string; -export type MapIndex4 = number | null; +export type MapIndex5 = number | null; export type TaskIds1 = string[] | null; export type TaskGroupId1 = string | null; export type LogicalDates2 = string[] | null; @@ -458,25 +472,25 @@ export type Type49 = "GetVariableKeys"; export type Key10 = string; export type DagId16 = string; export type RunId9 = string; -export type TaskId3 = string; -export type MapIndex5 = number | null; +export type TaskId4 = string; +export type MapIndex6 = number | null; export type IncludePriorDates = boolean; export type Type50 = "GetXCom"; export type Key11 = string; export type DagId17 = string; export type RunId10 = string; -export type TaskId4 = string; +export type TaskId5 = string; export type Type51 = "GetXComCount"; export type Key12 = string; export type DagId18 = string; export type RunId11 = string; -export type TaskId5 = string; +export type TaskId6 = string; export type Offset1 = number; export type Type52 = "GetXComSequenceItem"; export type Key13 = string; export type DagId19 = string; export type RunId12 = string; -export type TaskId6 = string; +export type TaskId7 = string; export type Start = number | null; export type Stop = number | null; export type Step = number | null; @@ -498,7 +512,7 @@ export type AssignedUsers1 = HITLUser[] | null; export type Type54 = "HITLDetailRequestResult"; export type InactiveAssets = AssetProfile[] | null; export type Type55 = "InactiveAssetsResult"; -export type Name12 = string | null; +export type Name14 = string | null; export type Type56 = "MaskSecret"; export type Ok = boolean; export type Type57 = "OKResponse"; @@ -508,7 +522,7 @@ export type StartDate4 = string | null; export type EndDate3 = string | null; export type Type58 = "PrevSuccessfulDagRunResult"; export type Type59 = "PreviousDagRunResult"; -export type TaskId7 = string; +export type TaskId8 = string; export type DagId20 = string; export type RunId13 = string; export type LogicalDate5 = string | null; @@ -516,7 +530,7 @@ export type StartDate5 = string | null; export type EndDate4 = string | null; export type State4 = string | null; export type TryNumber2 = number; -export type MapIndex6 = number | null; +export type MapIndex7 = number | null; export type Duration = number | null; export type Type60 = "PreviousTIResult"; export type Key14 = string; @@ -536,7 +550,7 @@ export type RetryReason = string | null; export type Type64 = "RetryTask"; export type Type65 = "SentFDs"; export type Fds = number[]; -export type Name13 = string; +export type Name15 = string; export type Key15 = string; export type Type66 = "SetAssetStateStoreByName"; export type Uri9 = string; @@ -552,8 +566,8 @@ export type Type70 = "SetTaskStateStore"; export type Key18 = string; export type DagId21 = string; export type RunId14 = string; -export type TaskId8 = string; -export type MapIndex7 = number | null; +export type TaskId9 = string; +export type MapIndex8 = number | null; export type DagResult1 = boolean; export type MappedLength = number | null; export type Type71 = "SetXCom"; @@ -624,6 +638,13 @@ export type Root = JsonValue[]; export type Type89 = "XComSequenceSliceResult"; export interface SupervisorWireSchema {} +/** + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "ArgValueSchema". + */ +export interface ArgValueSchema { + [k: string]: JsonValue; +} /** * Schema for AssetAliasModel used in AssetEventDagRunReference. * @@ -1003,6 +1024,7 @@ export interface TIRunContext { xcom_keys_to_clear?: XcomKeysToClear; should_retry?: ShouldRetry; start_date?: StartDate2; + arg_bindings?: ArgBindings; } /** * Variable schema for responses with fields that are needed for Runtime. @@ -1030,6 +1052,33 @@ export interface ConnectionResponse { port: Port1; extra: Extra6; } +/** + * One positional stub-task argument pulled from an upstream task's XCom. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "XComArgBinding". + */ +export interface XComArgBinding { + kind: Kind; + name: Name8; + value_schema?: ArgValueSchema | null; + task_id: TaskId1; + map_index?: MapIndex1; + element_index?: ElementIndex; +} +/** + * One positional stub-task argument carrying an inline literal from the Dag file. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "LiteralArgBinding". + */ +export interface LiteralArgBinding { + kind: Kind1; + name: Name9; + value_schema?: ArgValueSchema | null; + value?: unknown; + from_default?: FromDefault; +} /** * Email notification request for task failures/retries. * @@ -1149,7 +1198,7 @@ export interface DeferTask { * via the `definition` "DeleteAssetStateStoreByName". */ export interface DeleteAssetStateStoreByName { - name: Name8; + name: Name10; key: Key1; type?: Type21; } @@ -1187,8 +1236,8 @@ export interface DeleteXCom { key: Key5; dag_id: DagId6; run_id: RunId5; - task_id: TaskId1; - map_index?: MapIndex1; + task_id: TaskId2; + map_index?: MapIndex2; type?: Type25; } /** @@ -1205,7 +1254,7 @@ export interface ErrorResponse { * via the `definition` "GetAssetByName". */ export interface GetAssetByName { - name: Name9; + name: Name11; type?: Type27; } /** @@ -1221,7 +1270,7 @@ export interface GetAssetByUri { * via the `definition` "GetAssetEventByAsset". */ export interface GetAssetEventByAsset { - name: Name10; + name: Name12; uri: Uri7; after?: After; before?: Before; @@ -1252,7 +1301,7 @@ export interface GetAssetEventByAssetAlias { * via the `definition` "GetAssetStateStoreByName". */ export interface GetAssetStateStoreByName { - name: Name11; + name: Name13; key: Key6; type?: Type31; } @@ -1354,9 +1403,9 @@ export interface GetPreviousDagRun { */ export interface GetPreviousTI { dag_id: DagId12; - task_id: TaskId2; + task_id: TaskId3; logical_date?: LogicalDate4; - map_index?: MapIndex2; + map_index?: MapIndex3; state?: TaskInstanceState | null; type?: Type42; } @@ -1366,7 +1415,7 @@ export interface GetPreviousTI { */ export interface GetTICount { dag_id: DagId13; - map_index?: MapIndex3; + map_index?: MapIndex4; task_ids?: TaskIds; task_group_id?: TaskGroupId; logical_dates?: LogicalDates1; @@ -1407,7 +1456,7 @@ export interface GetTaskStateStore { */ export interface GetTaskStates { dag_id: DagId15; - map_index?: MapIndex4; + map_index?: MapIndex5; task_ids?: TaskIds1; task_group_id?: TaskGroupId1; logical_dates?: LogicalDates2; @@ -1440,8 +1489,8 @@ export interface GetXCom { key: Key10; dag_id: DagId16; run_id: RunId9; - task_id: TaskId3; - map_index?: MapIndex5; + task_id: TaskId4; + map_index?: MapIndex6; include_prior_dates?: IncludePriorDates; type?: Type50; } @@ -1455,7 +1504,7 @@ export interface GetXComCount { key: Key11; dag_id: DagId17; run_id: RunId10; - task_id: TaskId4; + task_id: TaskId5; type?: Type51; } /** @@ -1466,7 +1515,7 @@ export interface GetXComSequenceItem { key: Key12; dag_id: DagId18; run_id: RunId11; - task_id: TaskId5; + task_id: TaskId6; offset: Offset1; type?: Type52; } @@ -1478,7 +1527,7 @@ export interface GetXComSequenceSlice { key: Key13; dag_id: DagId19; run_id: RunId12; - task_id: TaskId6; + task_id: TaskId7; start: Start; stop: Stop; step: Step; @@ -1520,7 +1569,7 @@ export interface InactiveAssetsResult { */ export interface MaskSecret { value: JsonValue; - name?: Name12; + name?: Name14; type?: Type56; } /** @@ -1559,7 +1608,7 @@ export interface PreviousDagRunResult { * via the `definition` "PreviousTIResponse". */ export interface PreviousTIResponse { - task_id: TaskId7; + task_id: TaskId8; dag_id: DagId20; run_id: RunId13; logical_date?: LogicalDate5; @@ -1567,7 +1616,7 @@ export interface PreviousTIResponse { end_date?: EndDate4; state?: State4; try_number: TryNumber2; - map_index?: MapIndex6; + map_index?: MapIndex7; duration?: Duration; } /** @@ -1636,7 +1685,7 @@ export interface SentFDs { * via the `definition` "SetAssetStateStoreByName". */ export interface SetAssetStateStoreByName { - name: Name13; + name: Name15; key: Key15; value: JsonValue; type?: Type66; @@ -1694,8 +1743,8 @@ export interface SetXCom { value: JsonValue; dag_id: DagId21; run_id: RunId14; - task_id: TaskId8; - map_index?: MapIndex7; + task_id: TaskId9; + map_index?: MapIndex8; dag_result?: DagResult1; mapped_length?: MappedLength; type?: Type71; @@ -1896,4 +1945,4 @@ export interface XComSequenceSliceResult { * (e.g. bundle metadata) and runs the migrator accordingly. * Exposed so the SDK author / operator can confirm which schema * version their build is pinned to. */ -export const SUPERVISOR_API_VERSION = "2026-06-16" as const; +export const SUPERVISOR_API_VERSION = "2026-10-30" as const;