From d348817788d66c278b1b4b43bccfa87d86dfc793 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Tue, 4 Aug 2026 13:02:10 +0000 Subject: [PATCH 1/8] Cut common-compat 1.19.0 with the SDK surface the stub decorator needs The @task.stub TaskFlow support in providers-standard imports KNOWN_CONTEXT_KEYS, PlainXComArg, MappedOperator and the decorator base classes through the compat layer so the provider keeps working down to Airflow 2.11. Those symbols first ship in common-compat 1.19.0 (1.18.0 was released from main in the meantime without them), so the version is cut here for the standard provider's pin to resolve. --- providers/common/compat/docs/changelog.rst | 8 ++++++++ providers/common/compat/docs/index.rst | 6 +++--- providers/common/compat/provider.yaml | 1 + providers/common/compat/pyproject.toml | 6 +++--- .../src/airflow/providers/common/compat/__init__.py | 2 +- .../compat/src/airflow/providers/common/compat/sdk.py | 8 +++++++- uv.lock | 2 +- 7 files changed, 24 insertions(+), 9 deletions(-) diff --git a/providers/common/compat/docs/changelog.rst b/providers/common/compat/docs/changelog.rst index 919c2b980ecfd..ae758d1062ac4 100644 --- a/providers/common/compat/docs/changelog.rst +++ b/providers/common/compat/docs/changelog.rst @@ -25,6 +25,14 @@ Changelog --------- +1.19.0 +...... + +Features +~~~~~~~~ + +* ``Expose KNOWN_CONTEXT_KEYS and PlainXComArg through the common.compat SDK seam`` + 1.18.0 ...... diff --git a/providers/common/compat/docs/index.rst b/providers/common/compat/docs/index.rst index 1f4a78a79c61d..7e405381a81e7 100644 --- a/providers/common/compat/docs/index.rst +++ b/providers/common/compat/docs/index.rst @@ -62,7 +62,7 @@ apache-airflow-providers-common-compat package Common Compatibility Provider - providing compatibility code for previous Airflow versions -Release: 1.18.0 +Release: 1.19.0 Provider package ---------------- @@ -133,5 +133,5 @@ Downloading official packages You can download officially released packages and verify their checksums and signatures from the `Official Apache Download site `_ -* `The apache-airflow-providers-common-compat 1.18.0 sdist package `_ (`asc `__, `sha512 `__) -* `The apache-airflow-providers-common-compat 1.18.0 wheel package `_ (`asc `__, `sha512 `__) +* `The apache-airflow-providers-common-compat 1.19.0 sdist package `_ (`asc `__, `sha512 `__) +* `The apache-airflow-providers-common-compat 1.19.0 wheel package `_ (`asc `__, `sha512 `__) diff --git a/providers/common/compat/provider.yaml b/providers/common/compat/provider.yaml index 102a7f5559df4..0da08cfe5e1f2 100644 --- a/providers/common/compat/provider.yaml +++ b/providers/common/compat/provider.yaml @@ -29,6 +29,7 @@ source-date-epoch: 1785633505 # In such case adding >= NEW_VERSION and bumping to NEW_VERSION in a provider have # to be done in the same PR versions: + - 1.19.0 - 1.18.0 - 1.17.0 - 1.16.0 diff --git a/providers/common/compat/pyproject.toml b/providers/common/compat/pyproject.toml index ded1b6fcbe447..1ef143d64b5cb 100644 --- a/providers/common/compat/pyproject.toml +++ b/providers/common/compat/pyproject.toml @@ -25,7 +25,7 @@ build-backend = "flit_core.buildapi" [project] name = "apache-airflow-providers-common-compat" -version = "1.18.0" +version = "1.19.0" description = "Provider package apache-airflow-providers-common-compat for Apache Airflow" readme = "README.rst" license = "Apache-2.0" @@ -109,8 +109,8 @@ apache-airflow-providers-common-sql = {workspace = true} apache-airflow-providers-standard = {workspace = true} [project.urls] -"Documentation" = "https://airflow.apache.org/docs/apache-airflow-providers-common-compat/1.18.0" -"Changelog" = "https://airflow.apache.org/docs/apache-airflow-providers-common-compat/1.18.0/changelog.html" +"Documentation" = "https://airflow.apache.org/docs/apache-airflow-providers-common-compat/1.19.0" +"Changelog" = "https://airflow.apache.org/docs/apache-airflow-providers-common-compat/1.19.0/changelog.html" "Bug Tracker" = "https://github.com/apache/airflow/issues" "Source Code" = "https://github.com/apache/airflow" "Slack Chat" = "https://s.apache.org/airflow-slack" diff --git a/providers/common/compat/src/airflow/providers/common/compat/__init__.py b/providers/common/compat/src/airflow/providers/common/compat/__init__.py index fa614ba20ea89..cd2ec579d0697 100644 --- a/providers/common/compat/src/airflow/providers/common/compat/__init__.py +++ b/providers/common/compat/src/airflow/providers/common/compat/__init__.py @@ -29,7 +29,7 @@ __all__ = ["__version__"] -__version__ = "1.18.0" +__version__ = "1.19.0" if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse( "2.11.0" 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/uv.lock b/uv.lock index f039ed5c273fd..2cddd02b5b5f5 100644 --- a/uv.lock +++ b/uv.lock @@ -4529,7 +4529,7 @@ docs = [{ name = "apache-airflow-devel-common", extras = ["docs"], editable = "d [[package]] name = "apache-airflow-providers-common-compat" -version = "1.18.0" +version = "1.19.0" source = { editable = "providers/common/compat" } dependencies = [ { name = "apache-airflow" }, From f31625eaf8971f3cafe81765fbd8afe078923bad Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Tue, 28 Jul 2026 12:21:19 +0000 Subject: [PATCH 2/8] Support TaskFlow call syntax on @task.stub tasks Stub tasks silently ignored TaskFlow call arguments, so a Dag author could not hand literals or upstream XCom results to a lang-SDK runtime. The decorator now binds the call to the stub's signature at parse time and captures an ordered arg spec (literal values and direct upstream XCom references, with pydantic-derived JSON value schemas) that serializes with the Dag, while rejecting what cannot cross the language boundary: custom XCom keys, aggregated mapped outputs, non-JSON literals, and stubs with arguments inside mapped task groups. Mapped (.expand()) stubs capture no spec and keep the legacy behavior until a follow-up delivers per-map-index bindings. --- .../src/airflow/serialization/schema.json | 49 ++- .../serialization/test_dag_serialization.py | 98 +++++ providers/standard/README.rst | 2 +- providers/standard/docs/index.rst | 2 +- providers/standard/pyproject.toml | 2 +- .../providers/standard/decorators/stub.py | 290 +++++++++++++- .../unit/standard/decorators/test_stub.py | 376 +++++++++++++++++- 7 files changed, 810 insertions(+), 9 deletions(-) diff --git a/airflow-core/src/airflow/serialization/schema.json b/airflow-core/src/airflow/serialization/schema.json index 872c3a1331ee3..b860ca5e1bf55 100644 --- a/airflow-core/src/airflow/serialization/schema.json +++ b/airflow-core/src/airflow/serialization/schema.json @@ -142,6 +142,48 @@ "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", "enum": [ "xcom", "literal" ] }, + "value_schema": { "$ref": "#/definitions/typed_dict" }, + "task_id": { "type": "string" }, + "value": {}, + "from_default": { "type": "boolean" } + }, + "required": [ "name", "kind" ] + } + }, + "required": [ + "__type", + "__var" + ], + "additionalProperties": false + }, "color": { "type": "string", "pattern": "^#[a-fA-F0-9]{3,6}$" @@ -345,7 +387,12 @@ "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" } + } }, "dependencies": { "expand_input": ["partial_kwargs", "_is_mapped"], diff --git a/airflow-core/tests/unit/serialization/test_dag_serialization.py b/airflow-core/tests/unit/serialization/test_dag_serialization.py index 7852c25dc5ee8..575a7c5dcab6d 100644 --- a/airflow-core/tests/unit/serialization/test_dag_serialization.py +++ b/airflow-core/tests/unit/serialization/test_dag_serialization.py @@ -3524,6 +3524,104 @@ 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): ... + + # Nested value_schema (dict[str, int] re-encodes its additionalProperties) plus + # dict/list literal values, whose contents must not collide with the {__type,__var} + # encoding during round-trip. + @task.stub + def aggregate(counts: dict[str, int], tags: list, config: dict): ... + + data = extract() + transform("uk", data) + aggregate(data, ["metrics", "hourly"], {"threshold": {"warn": 1}}) + + ser_dag = DagSerialization.to_dict(dag) + # The serialized form must satisfy schema.json (arg_binding / typed_dict definitions). + DagSerialization.validate_schema(ser_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: DagAttributeTypes.DICT, + Encoding.VAR: { + "name": "country", + "kind": "literal", + "value_schema": {Encoding.TYPE: DagAttributeTypes.DICT, Encoding.VAR: {"type": "string"}}, + "value": "uk", + }, + }, + { + Encoding.TYPE: DagAttributeTypes.DICT, + Encoding.VAR: { + "name": "extracted", + "kind": "xcom", + "value_schema": { + Encoding.TYPE: DagAttributeTypes.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", + }, + ] + # The nested value_schema and dict/list literal values survive the round-trip intact. + assert round_tripped.task_dict["aggregate"]._arg_bindings == dag.task_dict["aggregate"]._arg_bindings + assert round_tripped.task_dict["aggregate"]._arg_bindings == [ + { + "name": "counts", + "kind": "xcom", + "value_schema": { + "type": "object", + "additionalProperties": {"type": "integer", "format": "int64"}, + }, + "task_id": "extract", + }, + { + "name": "tags", + "kind": "literal", + "value_schema": {"type": "array", "items": {}}, + "value": ["metrics", "hourly"], + }, + { + "name": "config", + "kind": "literal", + "value_schema": {"type": "object", "additionalProperties": True}, + "value": {"threshold": {"warn": 1}}, + }, + ] + assert not hasattr(round_tripped.task_dict["extract"], "_arg_bindings") + + # The deserialized spec must be plain JSON (no {__type, __var} encoding sentinels) so the + # execution API can validate it straight off the serialized Dag -- this is the contract + # ti_run relies on when it feeds get_arg_bindings() into the TaskArgBinding adapter. + from airflow.api_fastapi.execution_api.datamodels.task_arg_binding import get_arg_bindings_adapter + + for task_id in ("transform", "aggregate"): + get_arg_bindings_adapter().validate_python(round_tripped.task_dict[task_id]._arg_bindings) + + def test_handle_v1_serdag(): v1 = { "__version": 1, diff --git a/providers/standard/README.rst b/providers/standard/README.rst index f3e9502574084..19acc04cdacba 100644 --- a/providers/standard/README.rst +++ b/providers/standard/README.rst @@ -54,7 +54,7 @@ Requirements PIP package Version required ========================================== ================== ``apache-airflow`` ``>=2.11.0`` -``apache-airflow-providers-common-compat`` ``>=1.14.1`` +``apache-airflow-providers-common-compat`` ``>=1.19.0`` ========================================== ================== Optional cross provider package dependencies diff --git a/providers/standard/docs/index.rst b/providers/standard/docs/index.rst index a1d2b35831646..f621160a04878 100644 --- a/providers/standard/docs/index.rst +++ b/providers/standard/docs/index.rst @@ -90,7 +90,7 @@ The minimum Apache Airflow version supported by this provider distribution is `` PIP package Version required ========================================== ================== ``apache-airflow`` ``>=2.11.0`` -``apache-airflow-providers-common-compat`` ``>=1.14.1`` +``apache-airflow-providers-common-compat`` ``>=1.19.0`` ========================================== ================== Optional cross provider package dependencies diff --git a/providers/standard/pyproject.toml b/providers/standard/pyproject.toml index b1d2faf0f4927..817e09db8be8a 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.19.0", # 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..67952d67c3742 100644 --- a/providers/standard/src/airflow/providers/standard/decorators/stub.py +++ b/providers/standard/src/airflow/providers/standard/decorators/stub.py @@ -18,12 +18,34 @@ 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 PydanticUserError, 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] + PydanticUserError = None # type: ignore[assignment,misc] + from airflow.providers.common.compat.sdk import ( + KNOWN_CONTEXT_KEYS, + XCOM_RETURN_KEY, DecoratedOperator, + MappedOperator, + PlainXComArg, TaskDecorator, + XComArg, task_decorator_factory, ) @@ -31,6 +53,242 @@ 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. Any pydantic + # failure inside the body degrades to None there, so this retry never re-raises. + 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. + """ + # Reached only when pydantic is installed (``_infer_value_schema`` guards on + # ``TypeAdapter is None``), so ``PydanticUserError`` is a real exception class here. + # It is the base of PydanticSchemaGenerationError and PydanticInvalidForJsonSchema and + # covers annotations pydantic rejects outright (e.g. bare ClassVar); TypeError catches + # the exotic generics pydantic chokes on with a plain TypeError. Either way, "pydantic + # cannot schema this" degrades to no schema rather than failing Dag parsing. + try: + return TypeAdapter(annotation).json_schema(schema_generator=_ValueSchemaGenerator) + except (PydanticUserError, TypeError): + normalized = _normalize_temporal_annotation(annotation) + if normalized is annotation: + return None + try: + return TypeAdapter(normalized).json_schema(schema_generator=_ValueSchemaGenerator) + except (PydanticUserError, TypeError): + 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; pass it in its JSON form instead" + ) + + +def _validate_xcom_value(value: Any, task_id: str, name: str) -> 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" + ) + # isinstance, not .is_mapped: Airflow 2.11 operators have no is_mapped attribute. + if isinstance(value.operator, MappedOperator): + 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" + ) + 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, + *, + in_mapped_group: bool, +) -> 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 + + # 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. + if in_mapped_group: + raise ValueError( + f"@task.stub task {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()" + ) + + 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 + + class _StubOperator(DecoratedOperator): custom_operator_name: str = "@task.stub" @@ -60,10 +318,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 +333,23 @@ 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, + in_mapped_group=self.get_closest_mapped_task_group() is not None, + ) + + @classmethod + def get_serialized_fields(cls): + # _arg_bindings must round-trip back to plain JSON (not {__type, __var}-encoded) so the + # execution API can validate it straight off the serialized Dag: it deserializes fully + # only while it stays out of SerializedBaseOperator's static serialized-field set. + return super().get_serialized_fields() | {"_arg_bindings"} def execute(self, context: Context) -> Any: raise RuntimeError( @@ -96,6 +370,14 @@ 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. + + Mapped (``.expand()``) stubs do not receive TaskFlow arguments yet -- their call args + keep the legacy ignored behavior; per-map-index delivery is part of + https://github.com/apache/airflow/issues/66937 and lands in a follow-up. """ 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..7dee1b59a81ef 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, stub from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS @@ -69,3 +75,371 @@ 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_temporal_literal_rejected(self): + def fn(when: datetime.datetime): ... + + with DAG(dag_id="d"), pytest.raises(ValueError, match="not JSON-serializable"): + stub(fn)(datetime.datetime(2020, 1, 1)) + + 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): + """Mapped stubs capture no spec: their call args keep the legacy ignored behavior for now.""" + 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]) + + +@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"), + # pydantic raises PydanticUserError (not the JSON-schema subclasses) for these; they + # must still degrade to no schema rather than crash Dag parsing. + pytest.param(typing.ClassVar, None, id="pydantic-user-error"), + 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"} + + +def test_infer_value_schema_degrades_on_pydantic_typeerror(monkeypatch): + """A bare TypeError from pydantic degrades to no schema rather than crashing Dag parsing.""" + from airflow.providers.standard.decorators import stub as stub_module + + def _raise_type_error(_annotation): + raise TypeError("pydantic cannot build a schema for this") + + monkeypatch.setattr(stub_module, "TypeAdapter", _raise_type_error) + + # A fresh class dodges the process-lifetime schema cache and exercises the hashable-but- + # unschemable path, where a naive ``except TypeError`` retry would re-raise and crash. + class _Unschemable: ... + + assert _infer_value_schema(_Unschemable) is None From b9f1951ddc4cb566b9dac10c2b93f18641758ace Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Tue, 28 Jul 2026 12:25:41 +0000 Subject: [PATCH 3/8] Ship stub arg_bindings in a new execution API version TIRunContext gains an arg_bindings field so a lang-SDK runtime receives the stub task's TaskFlow arg spec at startup. ti_run derives it from the serialized Dag only for stub operators, so regular tasks never pay for the lookup, and only for clients on the new API version -- gated on the Cadwyn VersionChangeWithSideEffects.is_applied check rather than a date comparison -- so stub Dags that predate arg bindings keep running against older clients, for which the version migration strips the field. --- .../datamodels/task_arg_binding.py | 93 ++++++++++++++ .../execution_api/datamodels/taskinstance.py | 8 ++ .../execution_api/routes/task_instances.py | 34 +++++- .../execution_api/services/__init__.py | 16 +++ .../execution_api/services/task_instances.py | 67 +++++++++++ .../execution_api/versions/__init__.py | 2 + .../execution_api/versions/v2026_10_30.py | 42 +++++++ .../versions/head/test_task_instances.py | 113 ++++++++++++++++++ .../versions/v2026_10_30/__init__.py | 16 +++ .../v2026_10_30/test_task_instances.py | 100 ++++++++++++++++ .../airflow/sdk/api/datamodels/_generated.py | 34 +++++- 11 files changed, 523 insertions(+), 2 deletions(-) create mode 100644 airflow-core/src/airflow/api_fastapi/execution_api/datamodels/task_arg_binding.py create mode 100644 airflow-core/src/airflow/api_fastapi/execution_api/services/__init__.py create mode 100644 airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py create mode 100644 airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py create mode 100644 airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/__init__.py create mode 100644 airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py 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..94c653d577ca8 --- /dev/null +++ b/airflow-core/src/airflow/api_fastapi/execution_api/datamodels/task_arg_binding.py @@ -0,0 +1,93 @@ +# 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 _ArgBindingBase(BaseModel): + """Fields every :class:`TaskArgBinding` variant carries, regardless of ``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.""" + + +class XComArgBinding(_ArgBindingBase): + """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"] + + task_id: str + """Upstream task id whose ``return_value`` XCom is pulled.""" + + +class LiteralArgBinding(_ArgBindingBase): + """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``.""" + + 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 41ecf49b053fb..c713d505e3551 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,11 @@ get_team_name_for_ti, require_auth, ) +from airflow.api_fastapi.execution_api.services.task_instances import ( + LANG_SDK_OPERATORS, + client_supports_arg_bindings, + get_arg_bindings, +) from airflow.configuration import conf from airflow.exceptions import InvalidPartitionKeyError, TaskNotFound from airflow.models.asset import AssetActive @@ -163,6 +169,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 +318,30 @@ def ti_run( should_retry=_is_eligible_to_retry(previous_state, ti.try_number, ti.max_tries), ) + # Only set for lang-SDK (foreign-runtime) tasks with a captured TaskFlow arg + # spec; the route excludes unset fields, keeping regular responses lean. + if ( + ti.operator in LANG_SDK_OPERATORS + 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..2d90bd52bd722 --- /dev/null +++ b/airflow-core/src/airflow/api_fastapi/execution_api/services/task_instances.py @@ -0,0 +1,67 @@ +# 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 + +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from sqlalchemy.orm import Session + + from airflow.models.dagbag import DBDagBag + +# Task types (``TaskInstance.operator``, the operator class name) whose tasks carry a +# lang-SDK ``arg_bindings`` spec. Used to gate the serialized-Dag lookup so regular tasks +# never pay for it. The gate matches exact class names; a new lang-SDK operator adds its +# name here. +LANG_SDK_OPERATORS = frozenset({"_StubOperator"}) + + +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 must not run for them. + + Rather than comparing the negotiated version by date, we check the + ``VersionChangeWithSideEffects`` subclass's ``is_applied`` flag; see + https://docs.cadwyn.dev/concepts/version_changes/#version-changes-with-side-effects + """ + # Imported locally: the versions package transitively imports the routes, which import + # this module, so a top-level import here would be circular. + from airflow.api_fastapi.execution_api.versions.v2026_10_30 import AddArgBindingsToTIRunContext + + return AddArgBindingsToTIRunContext.is_applied + + +def get_arg_bindings(dag_bag: DBDagBag, ti: Any, *, session: Session) -> list | None: + """ + Extract the stub task's TaskFlow arg spec from its Dag version. + + Mapped (``.expand()``) stubs never capture a parse-time spec, so they resolve to + ``None`` here and keep the legacy ignored-args behavior; per-map-index delivery + lands in a follow-up. + """ + 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 + return getattr(task, "_arg_bindings", None) 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..1c85aed252c06 --- /dev/null +++ b/airflow-core/src/airflow/api_fastapi/execution_api/versions/v2026_10_30.py @@ -0,0 +1,42 @@ +# 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, + VersionChangeWithSideEffects, + convert_response_to_previous_version_for, + schema, +) + +from airflow.api_fastapi.execution_api.datamodels.taskinstance import TIRunContext + + +class AddArgBindingsToTIRunContext(VersionChangeWithSideEffects): + """Add the ``arg_bindings`` argument-binding spec for stub (foreign-runtime) tasks.""" + + description = __doc__ + + # A side-effect change, not just a schema one, so ti_run can gate the server-side spec + # derivation on ``is_applied``: clients older than this version never receive the field. + 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/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 bb3c0f7e5a785..4064c66078678 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 @@ -32,6 +32,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 @@ -161,6 +162,14 @@ def test_id_matches_sub_claim(client, session, create_task_instance): class TestTIRunState: + RUN_PAYLOAD = { + "state": "running", + "hostname": "random-hostname", + "unixname": "random-unixname", + "pid": 100, + "start_date": "2024-09-30T12:00:00Z", + } + def setup_method(self): clear_db_logs() clear_db_runs() @@ -372,6 +381,110 @@ 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() + + response = client.patch(f"/execution/task-instances/{tis['transform'].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": "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=self.RUN_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=self.RUN_PAYLOAD) + + assert response.status_code == 500 + assert response.json()["detail"]["reason"] == "invalid_arg_bindings" + + def test_ti_run_returns_no_arg_bindings_for_mapped_stub(self, client, dag_maker): + """Mapped stubs keep the legacy ignored-args behavior until per-map-index delivery lands.""" + with dag_maker("test_mapped_stub_ignored_args", 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_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..a4b98bd10206e --- /dev/null +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/v2026_10_30/test_task_instances.py @@ -0,0 +1,100 @@ +# 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_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/task-sdk/src/airflow/sdk/api/datamodels/_generated.py b/task-sdk/src/airflow/sdk/api/datamodels/_generated.py index cc3c7eb0a8f20..201f218c3c973 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 | None]]): + root: dict[str, JsonValue | None] + + 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. + """ + + name: Annotated[str, Field(title="Name")] + value_schema: ArgValueSchema | None = None + kind: Annotated[Literal["literal"], Field(title="Kind")] + 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,17 @@ 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. + """ + + name: Annotated[str, Field(title="Name")] + value_schema: ArgValueSchema | None = None + kind: Annotated[Literal["xcom"], Field(title="Kind")] + task_id: Annotated[str, Field(title="Task Id")] + + class AssetEventDagRunReference(BaseModel): """ Schema for AssetEvent model used in DagRun. @@ -782,6 +809,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 +828,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 From 685da024852149c094201cf8e7c87ba9df125d3b Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Tue, 28 Jul 2026 12:26:34 +0000 Subject: [PATCH 4/8] Deliver stub arg bindings to SDK runtimes via the supervisor schema StartupDetails in the supervisor wire schema carries the new arg_bindings so foreign runtimes receive the spec at task startup, with a version migration that strips it for runtimes pinned to the previous schema. The Go and TS SDKs regenerate against the new schema version; the Go arg-binding runtime itself lands in a stacked follow-up PR. --- .../airflow-go-pack/pack_integration_test.go | 3 +- go-sdk/pkg/execution/messages.go | 2 +- .../sdk/execution_time/schema/schema.json | 125 +++++++++++++++++- .../schema/versions/__init__.py | 5 + .../schema/versions/v2026_10_30.py | 36 +++++ .../execution_time/schema/test_migrator.py | 113 +++++++++++++++- ts-sdk/src/generated/supervisor.ts | 113 +++++++++++----- 7 files changed, 354 insertions(+), 43 deletions(-) create mode 100644 task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py 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/task-sdk/src/airflow/sdk/execution_time/schema/schema.json b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json index 8d606cf968043..4524c74ff794a 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -1,6 +1,6 @@ { "$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": { "AssetAliasReferenceAssetEventDagRun": { @@ -4590,6 +4590,114 @@ "title": "XComSequenceSliceResult", "type": "object" }, + "ArgValueSchema": { + "additionalProperties": { + "$ref": "#/$defs/JsonValue" + }, + "title": "ArgValueSchema", + "type": "object" + }, + "LiteralArgBinding": { + "description": "One positional stub-task argument carrying an inline literal from the Dag file.", + "properties": { + "name": { + "title": "Name", + "type": "string" + }, + "value_schema": { + "anyOf": [ + { + "$ref": "#/$defs/ArgValueSchema" + }, + { + "type": "null" + } + ], + "default": null + }, + "kind": { + "const": "literal", + "title": "Kind", + "type": "string" + }, + "value": { + "anyOf": [ + { + "$ref": "#/$defs/JsonValue" + }, + { + "type": "null" + } + ], + "default": null + }, + "from_default": { + "default": false, + "title": "From Default", + "type": "boolean" + } + }, + "required": [ + "name", + "kind" + ], + "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": { + "name": { + "title": "Name", + "type": "string" + }, + "value_schema": { + "anyOf": [ + { + "$ref": "#/$defs/ArgValueSchema" + }, + { + "type": "null" + } + ], + "default": null + }, + "kind": { + "const": "xcom", + "title": "Kind", + "type": "string" + }, + "task_id": { + "title": "Task Id", + "type": "string" + } + }, + "required": [ + "name", + "kind", + "task_id" + ], + "title": "XComArgBinding", + "type": "object" + }, "AssetEventDagRunReference": { "additionalProperties": false, "description": "Schema for AssetEvent model used in DagRun.", @@ -4981,6 +5089,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..cd5f5fff5fb56 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,107 @@ 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, + data_interval_start=None, + data_interval_end=None, + start_date=now, + end_date=None, + run_type=DagRunType.MANUAL, + state=DagRunState.RUNNING, + run_after=now, + consumed_asset_events=[], + partition_key=None, + ), + 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 049b0c1ce92f9..83170516900a4 100644 --- a/ts-sdk/src/generated/supervisor.ts +++ b/ts-sdk/src/generated/supervisor.ts @@ -22,6 +22,11 @@ // // 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; @@ -166,11 +171,6 @@ export type Conf = { export type TriggeringUserName = string | null; export type Name7 = string; export type Uri4 = string; -/** - * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema - * via the `definition` "JsonValue". - */ -export type JsonValue = unknown; export type SourceTaskId1 = string | null; export type SourceDagId1 = string | null; export type SourceRunId1 = string | null; @@ -245,6 +245,18 @@ 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 Name8 = string; +export type Kind = "xcom"; +export type TaskId1 = string; +export type Name9 = string; +export type Kind1 = "literal"; +export type FromDefault = boolean; export type Type13 = "TaskCallbackRequest"; export type Filepath2 = string; export type BundleName3 = string; @@ -310,7 +322,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,7 +336,7 @@ export type Type24 = "DeleteVariable"; export type Key5 = string; export type DagId6 = string; export type RunId5 = string; -export type TaskId1 = string; +export type TaskId2 = string; export type MapIndex1 = number | null; export type Type25 = "DeleteXCom"; /** @@ -362,11 +374,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 +401,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,7 +433,7 @@ 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 Type42 = "GetPreviousTI"; @@ -458,25 +470,25 @@ export type Type49 = "GetVariableKeys"; export type Key10 = string; export type DagId16 = string; export type RunId9 = string; -export type TaskId3 = string; +export type TaskId4 = string; export type MapIndex5 = 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 +510,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 +520,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; @@ -536,7 +548,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,7 +564,7 @@ export type Type70 = "SetTaskStateStore"; export type Key18 = string; export type DagId21 = string; export type RunId14 = string; -export type TaskId8 = string; +export type TaskId9 = string; export type MapIndex7 = number | null; export type DagResult1 = boolean; export type MappedLength = number | null; @@ -624,6 +636,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. * @@ -1019,6 +1038,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 +1050,31 @@ export interface VariableResponse { key: Key; value: Value; } +/** + * 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 { + name: Name8; + value_schema?: ArgValueSchema | null; + kind: Kind; + task_id: TaskId1; +} +/** + * 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 { + name: Name9; + value_schema?: ArgValueSchema | null; + kind: Kind1; + value?: unknown; + from_default?: FromDefault; +} /** * Email notification request for task failures/retries. * @@ -1149,7 +1194,7 @@ export interface DeferTask { * via the `definition` "DeleteAssetStateStoreByName". */ export interface DeleteAssetStateStoreByName { - name: Name8; + name: Name10; key: Key1; type?: Type21; } @@ -1187,7 +1232,7 @@ export interface DeleteXCom { key: Key5; dag_id: DagId6; run_id: RunId5; - task_id: TaskId1; + task_id: TaskId2; map_index?: MapIndex1; type?: Type25; } @@ -1205,7 +1250,7 @@ export interface ErrorResponse { * via the `definition` "GetAssetByName". */ export interface GetAssetByName { - name: Name9; + name: Name11; type?: Type27; } /** @@ -1221,7 +1266,7 @@ export interface GetAssetByUri { * via the `definition` "GetAssetEventByAsset". */ export interface GetAssetEventByAsset { - name: Name10; + name: Name12; uri: Uri7; after?: After; before?: Before; @@ -1252,7 +1297,7 @@ export interface GetAssetEventByAssetAlias { * via the `definition` "GetAssetStateStoreByName". */ export interface GetAssetStateStoreByName { - name: Name11; + name: Name13; key: Key6; type?: Type31; } @@ -1354,7 +1399,7 @@ export interface GetPreviousDagRun { */ export interface GetPreviousTI { dag_id: DagId12; - task_id: TaskId2; + task_id: TaskId3; logical_date?: LogicalDate4; map_index?: MapIndex2; state?: TaskInstanceState | null; @@ -1440,7 +1485,7 @@ export interface GetXCom { key: Key10; dag_id: DagId16; run_id: RunId9; - task_id: TaskId3; + task_id: TaskId4; map_index?: MapIndex5; include_prior_dates?: IncludePriorDates; type?: Type50; @@ -1455,7 +1500,7 @@ export interface GetXComCount { key: Key11; dag_id: DagId17; run_id: RunId10; - task_id: TaskId4; + task_id: TaskId5; type?: Type51; } /** @@ -1466,7 +1511,7 @@ export interface GetXComSequenceItem { key: Key12; dag_id: DagId18; run_id: RunId11; - task_id: TaskId5; + task_id: TaskId6; offset: Offset1; type?: Type52; } @@ -1478,7 +1523,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 +1565,7 @@ export interface InactiveAssetsResult { */ export interface MaskSecret { value: JsonValue; - name?: Name12; + name?: Name14; type?: Type56; } /** @@ -1559,7 +1604,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; @@ -1636,7 +1681,7 @@ export interface SentFDs { * via the `definition` "SetAssetStateStoreByName". */ export interface SetAssetStateStoreByName { - name: Name13; + name: Name15; key: Key15; value: JsonValue; type?: Type66; @@ -1694,7 +1739,7 @@ export interface SetXCom { value: JsonValue; dag_id: DagId21; run_id: RunId14; - task_id: TaskId8; + task_id: TaskId9; map_index?: MapIndex7; dag_result?: DagResult1; mapped_length?: MappedLength; @@ -1896,4 +1941,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; From 08b9cf87db9dad39da158facd8b34afed365e063 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Thu, 30 Jul 2026 00:34:21 +0000 Subject: [PATCH 5/8] Reject upstream outputs nested inside stub literal collections An XComArg buried in a list or dict literal fell through to the JSON check, whose "pass it in its JSON form instead" advice is impossible to follow for a task output. Detect nested references up front and point the author at the working alternative: pass the upstream output as its own argument. --- .../src/airflow/providers/standard/decorators/stub.py | 6 ++++++ .../standard/tests/unit/standard/decorators/test_stub.py | 7 +++++++ 2 files changed, 13 insertions(+) diff --git a/providers/standard/src/airflow/providers/standard/decorators/stub.py b/providers/standard/src/airflow/providers/standard/decorators/stub.py index 67952d67c3742..d40f6bd4c7587 100644 --- a/providers/standard/src/airflow/providers/standard/decorators/stub.py +++ b/providers/standard/src/airflow/providers/standard/decorators/stub.py @@ -186,6 +186,12 @@ def resolve(name: str, param: inspect.Parameter) -> Any: def _ensure_json_literal(value: Any, task_id: str, name: str) -> None: + if next(XComArg.iter_xcom_references(value), None) is not None: + raise ValueError( + f"@task.stub task {task_id!r} parameter {name!r} received a collection with an " + "upstream task output nested inside it; only a direct XComArg argument can cross " + "the language boundary -- pass the upstream output as its own argument" + ) try: json.dumps(value, allow_nan=False) except (TypeError, ValueError): diff --git a/providers/standard/tests/unit/standard/decorators/test_stub.py b/providers/standard/tests/unit/standard/decorators/test_stub.py index 7dee1b59a81ef..95b029aeac7b2 100644 --- a/providers/standard/tests/unit/standard/decorators/test_stub.py +++ b/providers/standard/tests/unit/standard/decorators/test_stub.py @@ -221,6 +221,13 @@ def fn(when: datetime.datetime): ... with DAG(dag_id="d"), pytest.raises(ValueError, match="not JSON-serializable"): stub(fn)(datetime.datetime(2020, 1, 1)) + @pytest.mark.parametrize("wrap", [lambda x: [x], lambda x: {"data": x}], ids=["list", "dict"]) + def test_xcom_nested_in_collection_literal_rejected(self, wrap): + with DAG(dag_id="d"): + extracted = stub(fn_extract)() + with pytest.raises(ValueError, match="nested inside"): + stub(fn_transform)("uk", wrap(extracted)) + def test_mapped_xcom_arg_rejected(self): with DAG(dag_id="d"): extracted = stub(fn_extract)() From a2d48e6424c6d12792a738899e9f23b9d032562a Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 5 Aug 2026 02:29:33 +0000 Subject: [PATCH 6/8] Prepare provider documentation against rc tags during a release vote When a PR cuts a new provider version while the previous version is still being voted on, only the rcN tags exist on the apache remote - the final tag is pushed after the vote passes. The changes-table walk in _get_all_changes_for_package assumed every past version has a final tag and crashed with git exit 128 in that window, breaking CI for any PR that bumps a provider version during a release wave. --- .../provider_documentation.py | 35 ++++++++++++++++++- .../tests/test_provider_documentation.py | 21 +++++++++++ 2 files changed, 55 insertions(+), 1 deletion(-) diff --git a/dev/breeze/src/airflow_breeze/prepare_providers/provider_documentation.py b/dev/breeze/src/airflow_breeze/prepare_providers/provider_documentation.py index 0a3053f11842d..171ed209ee8ad 100644 --- a/dev/breeze/src/airflow_breeze/prepare_providers/provider_documentation.py +++ b/dev/breeze/src/airflow_breeze/prepare_providers/provider_documentation.py @@ -402,6 +402,39 @@ def _print_changes_table(changes_table): console_print(syntax) +def _resolve_existing_version_tag(version_tag: str) -> str: + """Return the tag to diff a released version against. + + While a provider release vote is in progress only the ``rcN`` tags exist; the + final tag is pushed once the vote passes. Fall back to the newest rc tag in + that window so documentation preparation keeps working. + """ + result = run_command( + ["git", "rev-parse", version_tag], + cwd=AIRFLOW_ROOT_PATH, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + check=False, + ) + if result.returncode == 0: + return version_tag + result = run_command( + ["git", "tag", "--list", f"{version_tag}rc*", "--sort=-version:refname"], + cwd=AIRFLOW_ROOT_PATH, + capture_output=True, + text=True, + check=True, + ) + rc_tags = result.stdout.split() + if not rc_tags: + return version_tag + console_print( + f"[warning]The tag {version_tag} does not exist yet (release vote likely in progress). " + f"Using {rc_tags[0]} instead.[/]" + ) + return rc_tags[0] + + def _get_all_changes_for_package( provider_id: str, base_branch: str, @@ -511,7 +544,7 @@ def _get_all_changes_for_package( current_version = provider_details.versions[0] list_of_list_of_changes: list[list[Change]] = [] for version in provider_details.versions[1:]: - version_tag = get_version_tag(version, provider_id) + version_tag = _resolve_existing_version_tag(get_version_tag(version, provider_id)) result = run_command( _get_git_log_command( providers_folder_paths_for_git_commit_retrieval, next_version_tag, version_tag diff --git a/dev/breeze/tests/test_provider_documentation.py b/dev/breeze/tests/test_provider_documentation.py index b10484520711c..0d92d43713619 100644 --- a/dev/breeze/tests/test_provider_documentation.py +++ b/dev/breeze/tests/test_provider_documentation.py @@ -19,6 +19,7 @@ import random import string from pathlib import Path +from unittest import mock import pytest @@ -34,6 +35,7 @@ _get_change_from_line, _get_changes_classified, _get_git_log_command, + _resolve_existing_version_tag, classification_result, classify_change_deterministically, get_most_impactful_change, @@ -102,6 +104,25 @@ def test_get_version_tag(version: str, provider_id: str, suffix: str, tag: str): assert get_version_tag(version, provider_id, suffix) == tag +@pytest.mark.parametrize( + ("rev_parse_returncode", "rc_tags_output", "expected_tag"), + [ + (0, "", "providers-asana/1.0.1"), + (128, "providers-asana/1.0.1rc2\nproviders-asana/1.0.1rc1\n", "providers-asana/1.0.1rc2"), + (128, "", "providers-asana/1.0.1"), + ], +) +@mock.patch("airflow_breeze.prepare_providers.provider_documentation.run_command") +def test_resolve_existing_version_tag( + mock_run_command, rev_parse_returncode: int, rc_tags_output: str, expected_tag: str +): + mock_run_command.side_effect = [ + mock.Mock(returncode=rev_parse_returncode), + mock.Mock(returncode=0, stdout=rc_tags_output), + ] + assert _resolve_existing_version_tag("providers-asana/1.0.1") == expected_tag + + @pytest.mark.parametrize( ("folder_paths", "from_commit", "to_commit", "git_command"), [ From fecbcff0c543461eb4da447819bd7d0125961abc Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 5 Aug 2026 08:06:58 +0000 Subject: [PATCH 7/8] Java SDK: Register tasks as first-class TaskDef objects `dag.addTask("extract", Extract.class)` stored tasks as a plain `Map>`, which leaves nowhere to hang anything else a task needs: dependency edges, task-level configuration, and argument wiring all have to attach to a per-task object, and a map of classes cannot carry them. Introducing that object now keeps those follow-ups additive instead of forcing another break of the registration API later. The annotation surface keeps `Builder.Dag` / `Builder.Task`, and the interface users implement keeps the `Task` name, so the definition objects are `DagDef` and `TaskDef` -- a pairing that stays unambiguous next to `Task` at a use site. The SDK is pre-1.0, so the old string-keyed overload is removed outright rather than deprecated. --- .../language-sdks/java.rst | 12 +-- .../airflow/example/ExampleBundleBuilder.java | 2 +- .../example/InterfaceExampleBuilder.java | 11 ++- .../apache/airflow/sdk/BuilderProcessor.kt | 11 +-- .../org/apache/airflow/sdk/BuilderTest.kt | 48 ++++++----- .../airflow/example/ScalaSparkExample.scala | 14 +-- .../kotlin/org/apache/airflow/sdk/Builder.kt | 4 +- .../kotlin/org/apache/airflow/sdk/Bundle.kt | 20 ++--- .../apache/airflow/sdk/{Dag.kt => DagDef.kt} | 60 +++++++++---- .../org/apache/airflow/sdk/execution/Task.kt | 6 +- .../org/apache/airflow/sdk/BundleTest.kt | 4 +- .../org/apache/airflow/sdk/DagDefTest.kt | 86 +++++++++++++++++++ .../apache/airflow/sdk/execution/TaskTest.kt | 6 +- .../airflow/k8sexample/K8sBundleBuilder.java | 2 +- 14 files changed, 202 insertions(+), 84 deletions(-) rename java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/{Dag.kt => DagDef.kt} (63%) create mode 100644 java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/DagDefTest.kt diff --git a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst index c8b73b4626a28..f84aa91094b10 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst @@ -115,7 +115,7 @@ Java entry point public class Main implements BundleBuilder { @Override - public Iterable getDags() { + public Iterable getDags() { return List.of(SalesPipelineBuilder.build()); // SalesPipelineBuilder generated at compile time } @@ -203,7 +203,7 @@ Interface-based API ~~~~~~~~~~~~~~~~~~~ Implement the ``Task`` interface directly for full control over how tasks are registered and how XComs are -read. +read. Each task is registered as a ``TaskDef`` on a ``DagDef``. .. code-block:: java @@ -224,10 +224,10 @@ Register tasks manually in a ``BundleBuilder``: public class MyBundle implements BundleBuilder { @Override - public Iterable getDags() { - var dag = new Dag("my_dag"); - dag.addTask("fetch", FetchTask.class); - dag.addTask("process", ProcessTask.class); + public Iterable getDags() { + var dag = new DagDef("my_dag") + .addTask(new TaskDef("fetch", FetchTask.class)) + .addTask(new TaskDef("process", ProcessTask.class)); return List.of(dag); } } diff --git a/java-sdk/example/src/java/org/apache/airflow/example/ExampleBundleBuilder.java b/java-sdk/example/src/java/org/apache/airflow/example/ExampleBundleBuilder.java index f63d6c4d74337..fa1a860755863 100644 --- a/java-sdk/example/src/java/org/apache/airflow/example/ExampleBundleBuilder.java +++ b/java-sdk/example/src/java/org/apache/airflow/example/ExampleBundleBuilder.java @@ -26,7 +26,7 @@ public class ExampleBundleBuilder implements BundleBuilder { @NotNull @Override - public Iterable getDags() { + public Iterable getDags() { return List.of( InterfaceExampleBuilder.build(), AnnotationExampleBuilder.build(), diff --git a/java-sdk/example/src/java/org/apache/airflow/example/InterfaceExampleBuilder.java b/java-sdk/example/src/java/org/apache/airflow/example/InterfaceExampleBuilder.java index 1c536c3cbf2c7..74eb4f5d7533a 100644 --- a/java-sdk/example/src/java/org/apache/airflow/example/InterfaceExampleBuilder.java +++ b/java-sdk/example/src/java/org/apache/airflow/example/InterfaceExampleBuilder.java @@ -71,11 +71,10 @@ public void execute(@NotNull Context context, Client client) { } } - public static Dag build() { - var dag = new Dag("java_interface_example"); - dag.addTask("extract", Extract.class); - dag.addTask("transform", Transform.class); - dag.addTask("load", Load.class); - return dag; + public static DagDef build() { + return new DagDef("java_interface_example") + .addTask(new TaskDef("extract", Extract.class)) + .addTask(new TaskDef("transform", Transform.class)) + .addTask(new TaskDef("load", Load.class)); } } diff --git a/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt b/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt index 11202ffcdb455..56cbf1e76ad9f 100644 --- a/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt +++ b/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt @@ -54,8 +54,8 @@ import javax.tools.Diagnostic * containing: * * - One inner class per [Builder.Task]-annotated method, implementing [Task]. - * - A static `build()` method that constructs the [Dag] and registers those - * inner classes as tasks. + * - A static `build()` method that constructs the [DagDef] and registers those + * inner classes as [TaskDef]s. * * [Builder.XCom]-annotated parameters are resolved via `client.getXCom` in the * generated `execute` body, with the result cast to the parameter's declared @@ -102,8 +102,8 @@ class BuilderProcessor : AbstractProcessor() { MethodSpec .methodBuilder("build") .addModifiers(Modifier.PUBLIC, Modifier.STATIC) - .returns(ClassName.get(Dag::class.java)) - .addStatement($$"var dag = new $T($S)", ClassName.get(Dag::class.java), ann.id.ifBlank { el.simpleName }) + .returns(ClassName.get(DagDef::class.java)) + .addStatement($$"var dag = new $T($S)", ClassName.get(DagDef::class.java), ann.id.ifBlank { el.simpleName }) for (inner in el.enclosedElements) { if (inner !is ExecutableElement) continue @@ -116,7 +116,8 @@ class BuilderProcessor : AbstractProcessor() { builderClass.addType(task.spec) buildMethod.addStatement( - $$"dag.addTask($S, $L.class)", + $$"dag.addTask(new $T($S, $L.class))", + ClassName.get(TaskDef::class.java), ann.id.ifBlank { inner.simpleName }, innerName, ) diff --git a/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt b/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt index 3e28e1b009afd..6a08979c56163 100644 --- a/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt +++ b/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt @@ -83,16 +83,17 @@ class BuilderTest { import java.util.Optional; import org.apache.airflow.sdk.Client; import org.apache.airflow.sdk.Context; - import org.apache.airflow.sdk.Dag; + import org.apache.airflow.sdk.DagDef; import org.apache.airflow.sdk.MissingXComException; import org.apache.airflow.sdk.Task; + import org.apache.airflow.sdk.TaskDef; public final class TestExampleBuilder { - public static Dag build() { - var dag = new Dag("TestExample"); - dag.addTask("t1", T1.class); - dag.addTask("t2", T2.class); - dag.addTask("t3", T3.class); + public static DagDef build() { + var dag = new DagDef("TestExample"); + dag.addTask(new TaskDef("t1", T1.class)); + dag.addTask(new TaskDef("t2", T2.class)); + dag.addTask(new TaskDef("t3", T3.class)); return dag; } public static final class T1 implements Task { @@ -157,14 +158,15 @@ class BuilderTest { import java.util.Optional; import org.apache.airflow.sdk.Client; import org.apache.airflow.sdk.Context; - import org.apache.airflow.sdk.Dag; + import org.apache.airflow.sdk.DagDef; import org.apache.airflow.sdk.MissingXComException; import org.apache.airflow.sdk.Task; + import org.apache.airflow.sdk.TaskDef; public final class TestExampleBuilder { - public static Dag build() { - var dag = new Dag("TestExample"); - dag.addTask("t", T.class); + public static DagDef build() { + var dag = new DagDef("TestExample"); + dag.addTask(new TaskDef("t", T.class)); return dag; } public static final class T implements Task { @@ -220,14 +222,15 @@ class BuilderTest { import java.util.Optional; import org.apache.airflow.sdk.Client; import org.apache.airflow.sdk.Context; - import org.apache.airflow.sdk.Dag; + import org.apache.airflow.sdk.DagDef; import org.apache.airflow.sdk.MissingXComException; import org.apache.airflow.sdk.Task; + import org.apache.airflow.sdk.TaskDef; public final class TestExampleBuilder { - public static Dag build() { - var dag = new Dag("TestExample"); - dag.addTask("t", T.class); + public static DagDef build() { + var dag = new DagDef("TestExample"); + dag.addTask(new TaskDef("t", T.class)); return dag; } public static final class T implements Task { @@ -261,8 +264,8 @@ class BuilderTest { "org.apache.airflow.example.TestExampleBuilder", """ package org.apache.airflow.example; - import org.apache.airflow.sdk.Dag; - public final class TestExampleBuilder { public static Dag build() { var dag = new Dag("foo"); return dag; } } + import org.apache.airflow.sdk.DagDef; + public final class TestExampleBuilder { public static DagDef build() { var dag = new DagDef("foo"); return dag; } } """, ) } @@ -284,8 +287,8 @@ class BuilderTest { "org.apache.airflow.example.Foo", """ package org.apache.airflow.example; - import org.apache.airflow.sdk.Dag; - public final class Foo { public static Dag build() { var dag = new Dag("TestExample"); return dag; } } + import org.apache.airflow.sdk.DagDef; + public final class Foo { public static DagDef build() { var dag = new DagDef("TestExample"); return dag; } } """, ) } @@ -313,12 +316,13 @@ class BuilderTest { import java.lang.Override; import org.apache.airflow.sdk.Client; import org.apache.airflow.sdk.Context; - import org.apache.airflow.sdk.Dag; + import org.apache.airflow.sdk.DagDef; import org.apache.airflow.sdk.Task; + import org.apache.airflow.sdk.TaskDef; public final class TestExampleBuilder { - public static Dag build() { - var dag = new Dag("TestExample"); - dag.addTask("foo", T1.class); + public static DagDef build() { + var dag = new DagDef("TestExample"); + dag.addTask(new TaskDef("foo", T1.class)); return dag; } public static final class T1 implements Task { diff --git a/java-sdk/scala_spark_example/src/main/scala/org/apache/airflow/example/ScalaSparkExample.scala b/java-sdk/scala_spark_example/src/main/scala/org/apache/airflow/example/ScalaSparkExample.scala index ae779bf02135a..9ecebafb0cfe0 100644 --- a/java-sdk/scala_spark_example/src/main/scala/org/apache/airflow/example/ScalaSparkExample.scala +++ b/java-sdk/scala_spark_example/src/main/scala/org/apache/airflow/example/ScalaSparkExample.scala @@ -19,7 +19,7 @@ package org.apache.airflow.example -import org.apache.airflow.sdk.{Bundle, BundleBuilder, Client, Context, Dag, Server, Task} +import org.apache.airflow.sdk.{Bundle, BundleBuilder, Client, Context, DagDef, Server, Task, TaskDef} import org.apache.logging.log4j.{LogManager, Logger} import org.apache.spark.sql.SparkSession import org.apache.spark.sql.functions.sum @@ -133,16 +133,16 @@ class SparkLoad extends Task { } object ScalaSparkExample { - def build(): Dag = - new Dag(SparkEtl.DagId) - .addTask(SparkEtl.ExtractTaskId, classOf[SparkExtract]) - .addTask(SparkEtl.TransformTaskId, classOf[SparkTransform]) - .addTask(SparkEtl.LoadTaskId, classOf[SparkLoad]) + def build(): DagDef = + new DagDef(SparkEtl.DagId) + .addTask(new TaskDef(SparkEtl.ExtractTaskId, classOf[SparkExtract])) + .addTask(new TaskDef(SparkEtl.TransformTaskId, classOf[SparkTransform])) + .addTask(new TaskDef(SparkEtl.LoadTaskId, classOf[SparkLoad])) } /** Bundle entry point served to Airflow's Java coordinator. */ object ScalaSparkBundleBuilder extends BundleBuilder { - override def getDags(): java.lang.Iterable[Dag] = java.util.List.of(ScalaSparkExample.build()) + override def getDags(): java.lang.Iterable[DagDef] = java.util.List.of(ScalaSparkExample.build()) def main(args: Array[String]): Unit = Server.create(args).serve(new Bundle(getDags())) diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Builder.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Builder.kt index 3a5b84d2daf84..9dbfbfbefc390 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Builder.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Builder.kt @@ -24,7 +24,7 @@ package org.apache.airflow.sdk * * This class is not instantiated directly. Its nested annotations drive the * `BuilderProcessor` annotation processor in the :processor project, - * which generates a `*Builder` class for each class annotated with [Dag]. + * which generates a `*Builder` class for each class annotated with [Builder.Dag]. * * Example: * @@ -41,7 +41,7 @@ package org.apache.airflow.sdk * ``` * * The processor generates `MyPipelineBuilder.build()`, which returns a - * fully wired-up [Dag] ready to add to a [Bundle]. + * fully wired-up [DagDef] ready to add to a [Bundle]. */ class Builder internal constructor() { /** diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Bundle.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Bundle.kt index 677ec48eb9311..6cd549270a8a4 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Bundle.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Bundle.kt @@ -20,22 +20,22 @@ package org.apache.airflow.sdk /** - * An immutable snapshot of all [Dag]s that this JVM process can execute. + * An immutable snapshot of all [DagDef]s that this JVM process can execute. * * Build a [Bundle] by implementing [BundleBuilder], then pass it to * [Server.serve] to start accepting task-execution requests. * - * @property dags All registered Dags keyed by [Dag.id]. + * @property dags All registered Dags keyed by [DagDef.id]. * @throws IllegalArgumentException if any two Dags share the same ID. */ class Bundle( - dags: Iterable, + dags: Iterable, ) { - internal val dags: Map = dags.associateByDagId() + internal val dags: Map = dags.associateByDagId() } -private fun Iterable.associateByDagId(): Map { - val dagMap = linkedMapOf() +private fun Iterable.associateByDagId(): Map { + val dagMap = linkedMapOf() for (dag in this) { require(dagMap.putIfAbsent(dag.id, dag) == null) { "Dags in bundle have duplicate ID: ${dag.id}" @@ -45,14 +45,14 @@ private fun Iterable.associateByDagId(): Map { } /** - * Entry point for declaring the [Dag]s that this bundle contains. + * Entry point for declaring the [DagDef]s that this bundle contains. * * Implement this interface to create a Dag bundle to be served by [Server]. * * ```java * public class MyBundleBuilder implements BundleBuilder { * @Override - * public Iterable getDags() { + * public Iterable getDags() { * return List.of(MyDagBuilder.build()); * } * @@ -64,14 +64,14 @@ private fun Iterable.associateByDagId(): Map { */ interface BundleBuilder { /** - * Returns all [Dag]s that belong to this bundle. + * Returns all [DagDef]s that belong to this bundle. * * Called once during [build]; Dag IDs must be unique across the returned * collection. * * @throws IllegalArgumentException if any two Dags share the same ID. */ - fun getDags(): Iterable + fun getDags(): Iterable /** * Constructs a [Bundle] from the Dags returned by [getDags]. diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Dag.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/DagDef.kt similarity index 63% rename from java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Dag.kt rename to java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/DagDef.kt index c998580374169..a6bd50e949a1c 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Dag.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/DagDef.kt @@ -24,46 +24,72 @@ import kotlin.Throws /** * A collection of tasks with directional dependencies. * - * Create a [Dag] directly and register tasks with [addTask]. + * Create a [DagDef] directly and register [TaskDef]s with [addTask]. * * The [Builder.Dag] annotation should generally be preferred in user code, * where the annotation processor generates the wiring for you. Only use this - * class directly if you need to do low-level plumbing. + * class directly if you need to do low-level plumbing: + * + * ```java + * var dag = new DagDef("java_etl") + * .addTask(new TaskDef("extract", Extract.class)) + * .addTask(new TaskDef("load", Load.class)); + * ``` * * @param id Dag identifier. Must contain only ASCII alphanumeric characters, * dashes, dots, or underscores; must be unique within a [Bundle]. * * @see Builder.Dag */ -class Dag( +class DagDef( val id: String, // TODO: charset check? ) { - internal var tasks = mutableMapOf>() + internal val tasks = linkedMapOf() /** * Registers a task with this Dag. * - * The class must have a public no-argument constructor and implement [Task]. - * Task IDs must be unique within a Dag. + * A [TaskDef] belongs to at most one [DagDef]; registering the same instance + * with a second Dag, or twice with the same one, fails. Task IDs must be + * unique within a Dag. * - * @param id Task identifier, unique within this Dag. - * @param definition Class that implements [Task]. Must have a public no-arg - * constructor. + * @param task Task definition to register. * @return This Dag, for chaining. - * @throws IllegalArgumentException if a task already exists in the Dag with - * the same ID. + * @throws IllegalArgumentException if the task already belongs to a Dag or a + * task with the same ID is already registered. */ - fun addTask( - id: String, - definition: Class, - ): Dag { - require(tasks.putIfAbsent(id, definition) == null) { - "Tasks in Dag have duplicate ID: $id" + fun addTask(task: TaskDef): DagDef { + require(task.owner == null) { + "Task '${task.id}' already belongs to Dag '${task.owner?.id}'" + } + require(tasks.putIfAbsent(task.id, task) == null) { + "Tasks in Dag have duplicate ID: ${task.id}" } + task.owner = this return this } } +/** + * One task definition: its ID and the class that implements it. + * + * ```java + * var extract = new TaskDef("extract", Extract.class); + * ``` + * + * @param id Task identifier, unique within a [DagDef]. + * @param definition Class that implements [Task]. Must have a public no-arg + * constructor. + * + * @see Builder.Task + */ +class TaskDef( + val id: String, + val definition: Class, +) { + internal var owner: DagDef? = null +} + /** * A single unit of work executed by Airflow. * diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Task.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Task.kt index 60258dacc628d..1f1a01e674a58 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Task.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Task.kt @@ -70,9 +70,11 @@ internal object TaskRunner { request: StartupDetails, client: Client, ): Any { - val task = bundle.dags[request.ti.dagId]?.tasks[request.ti.taskId] ?: return TaskResult.of(TaskState.State.REMOVED) + val definition = + bundle.dags[request.ti.dagId]?.tasks[request.ti.taskId]?.definition + ?: return TaskResult.of(TaskState.State.REMOVED) return try { - task.getDeclaredConstructor().newInstance().execute(Context.from(request), client) + definition.getDeclaredConstructor().newInstance().execute(Context.from(request), client) TaskResult.success() } catch (e: Throwable) { logger.error("Error executing task", mapOf("ti" to request.ti, "error" to e, "trace" to e.stackTraceToString())) diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/BundleTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/BundleTest.kt index 0e4afb1894a7f..57050754e887e 100644 --- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/BundleTest.kt +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/BundleTest.kt @@ -27,7 +27,7 @@ internal class BundleTest { @Test @DisplayName("Should index dags by dagId") fun shouldIndexDagsByDagId() { - val dag = Dag("dag") + val dag = DagDef("dag") val bundle = Bundle(listOf(dag)) @@ -39,7 +39,7 @@ internal class BundleTest { fun shouldRejectDuplicateDagIds() { val error = Assertions.assertThrows(IllegalArgumentException::class.java) { - Bundle(listOf(Dag("dag"), Dag("dag"))) + Bundle(listOf(DagDef("dag"), DagDef("dag"))) } Assertions.assertEquals("Dags in bundle have duplicate ID: dag", error.message) diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/DagDefTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/DagDefTest.kt new file mode 100644 index 0000000000000..f22c26432c30c --- /dev/null +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/DagDefTest.kt @@ -0,0 +1,86 @@ +/* + * 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. + */ + +package org.apache.airflow.sdk + +import org.junit.jupiter.api.Assertions +import org.junit.jupiter.api.DisplayName +import org.junit.jupiter.api.Test + +internal class DagDefTest { + private class NoOp : Task { + override fun execute( + context: Context, + client: Client, + ) = Unit + } + + @Test + @DisplayName("Should index tasks by taskId in registration order") + fun shouldIndexTasksByTaskId() { + val extract = TaskDef("extract", NoOp::class.java) + val load = TaskDef("load", NoOp::class.java) + + val dag = DagDef("dag").addTask(extract).addTask(load) + + Assertions.assertEquals(listOf("extract", "load"), dag.tasks.keys.toList()) + Assertions.assertEquals(mapOf("extract" to extract, "load" to load), dag.tasks) + } + + @Test + @DisplayName("Should reject duplicate task ids") + fun shouldRejectDuplicateTaskIds() { + val dag = DagDef("dag").addTask(TaskDef("extract", NoOp::class.java)) + + val error = + Assertions.assertThrows(IllegalArgumentException::class.java) { + dag.addTask(TaskDef("extract", NoOp::class.java)) + } + + Assertions.assertEquals("Tasks in Dag have duplicate ID: extract", error.message) + } + + @Test + @DisplayName("Should reject a task already registered with another dag") + fun shouldRejectTaskOwnedByAnotherDag() { + val extract = TaskDef("extract", NoOp::class.java) + DagDef("first").addTask(extract) + + val error = + Assertions.assertThrows(IllegalArgumentException::class.java) { + DagDef("second").addTask(extract) + } + + Assertions.assertEquals("Task 'extract' already belongs to Dag 'first'", error.message) + } + + @Test + @DisplayName("Should reject the same task registered twice with one dag") + fun shouldRejectTaskRegisteredTwice() { + val extract = TaskDef("extract", NoOp::class.java) + val dag = DagDef("dag").addTask(extract) + + val error = + Assertions.assertThrows(IllegalArgumentException::class.java) { + dag.addTask(extract) + } + + Assertions.assertEquals("Task 'extract' already belongs to Dag 'dag'", error.message) + } +} diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskTest.kt index 42a083b94aacb..5d303a3dd814a 100644 --- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskTest.kt +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskTest.kt @@ -22,8 +22,9 @@ package org.apache.airflow.sdk.execution import org.apache.airflow.sdk.Bundle import org.apache.airflow.sdk.Client import org.apache.airflow.sdk.Context -import org.apache.airflow.sdk.Dag +import org.apache.airflow.sdk.DagDef import org.apache.airflow.sdk.Task +import org.apache.airflow.sdk.TaskDef import org.apache.airflow.sdk.execution.comm.BundleInfo import org.apache.airflow.sdk.execution.comm.DagRun import org.apache.airflow.sdk.execution.comm.RetryTask @@ -88,8 +89,7 @@ class TaskTest { taskId: String, taskClass: Class, ): Bundle { - val dag = Dag("test_dag") - dag.addTask(taskId, taskClass) + val dag = DagDef("test_dag").addTask(TaskDef(taskId, taskClass)) return Bundle(listOf(dag)) } diff --git a/kubernetes-tests/lang_sdk/java_example/src/java/org/apache/airflow/k8sexample/K8sBundleBuilder.java b/kubernetes-tests/lang_sdk/java_example/src/java/org/apache/airflow/k8sexample/K8sBundleBuilder.java index 8333567da83ed..c3b511b898ced 100644 --- a/kubernetes-tests/lang_sdk/java_example/src/java/org/apache/airflow/k8sexample/K8sBundleBuilder.java +++ b/kubernetes-tests/lang_sdk/java_example/src/java/org/apache/airflow/k8sexample/K8sBundleBuilder.java @@ -28,7 +28,7 @@ public class K8sBundleBuilder implements BundleBuilder { @NotNull @Override - public Iterable getDags() { + public Iterable getDags() { return List.of(CombinedExampleBuilder.build()); } From 9e7cab73e8f0a6acb2fa5ce95ac0ab46bdb7b9cc Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 5 Aug 2026 08:28:37 +0000 Subject: [PATCH 8/8] Java SDK: Honor TaskFlow arg bindings sent by the supervisor For a stub-backed Dag the Python file's `@task.stub` call site is the graph the scheduler actually orders the run by, so it must also be what feeds the Java task its inputs. The Java side previously re-declared that data flow with `@Builder.XCom(task = "...")`, duplicating the Dag file's wiring in a second place that nothing keeps honest: rename or re-wire a task in Python and the Java annotation silently keeps pulling the old upstream. The 2026-10-30 supervisor schema delivers the call site's bindings with every task run, so the runtime can read them instead of guessing. Binding is positional, matching the Go SDK's flat-parameter contract: Java parameter names are not API, so an IDE rename must never rebind an input. Keyword-style calls bind by name only through an explicit `TaskInput` bundle whose public fields declare their wire names -- the deliberate, tagged boundary for snake_case-to-camelCase crossings. A task declares flat data parameters or one bundle, never both, so field names and positions cannot shift each other. jsonSchema2Pojo cannot express the kind-discriminated binding union, so the generated `TIRunContext` carries the raw payload and a small hand-written decoder materializes the typed view. --- .../language-sdks/java.rst | 131 +++++++- .../airflow/example/AnnotationExample.java | 25 +- .../example/InterfaceExampleBuilder.java | 11 +- .../airflow/example/XComCastingExample.java | 11 +- .../src/resources/dags/java_examples.py | 36 ++- java-sdk/gradle.properties | 2 +- .../apache/airflow/sdk/BuilderProcessor.kt | 268 ++++++++++------ .../org/apache/airflow/sdk/BuilderTest.kt | 264 ++++++++++++---- java-sdk/sdk/schema/schema.json | 259 ++++++++++++++- .../kotlin/org/apache/airflow/sdk/ArgName.kt | 40 +++ .../kotlin/org/apache/airflow/sdk/Builder.kt | 19 +- .../kotlin/org/apache/airflow/sdk/Client.kt | 94 ++++++ .../org/apache/airflow/sdk/TaskInput.kt | 45 +++ .../airflow/sdk/execution/ArgBinding.kt | 77 +++++ .../apache/airflow/sdk/internal/ArgValues.kt | 175 +++++++++++ .../org/apache/airflow/sdk/ClientArgTest.kt | 297 ++++++++++++++++++ 16 files changed, 1550 insertions(+), 204 deletions(-) create mode 100644 java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/ArgName.kt create mode 100644 java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/TaskInput.kt create mode 100644 java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/ArgBinding.kt create mode 100644 java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/ArgValues.kt create mode 100644 java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ClientArgTest.kt diff --git a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst index f84aa91094b10..b3d1f2792fc78 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst @@ -93,10 +93,7 @@ Java implementation } @Builder.Task(id = "transform") - public long transform( - Client client, - @Builder.XCom(task = "extract") long recordCount - ) { + public long transform(Client client, long recordCount) { var threshold = (String) client.getVariable("transform_threshold"); // ... process data ... return transformedCount; @@ -105,8 +102,11 @@ Java implementation .. note:: - See how both ``transform`` in Python and Java need to have an argument to accept upstream XCom. The - Python one is needed to declare dependency, and the Java one is needed to actually retrieve the value. + The graph is declared once, in the Python Dag file: ``transform(extract())`` feeds the upstream's + return value into the downstream's parameter by calling tasks like functions. The supervisor sends + the resulting *argument bindings* to the Java runtime, and each Java data parameter receives + whatever the Python call site bound at its position — an upstream task's XCom or an inline + literal. See :ref:`java-sdk/arg-binding`. Java entry point ~~~~~~~~~~~~~~~~ @@ -164,12 +164,17 @@ Annotate a plain Java class and let the SDK generate the boilerplate at compile * - ``@Builder.Task(id = "...")`` - Marks a method as a task implementation. The ``id`` must match the ``@task.stub`` function name in the Python Dag. If ``id`` is omitted the method name is used. - * - ``@Builder.XCom(task = "...")`` - - Injects the ``return_value`` XCom from the named upstream task as a method parameter. - The parameter type must be compatible with the stored value (see :ref:`java-sdk/types`). + * - ``TaskInput`` / ``@ArgName("...")`` + - Marks a class as a task's input bundle, so keyword arguments bind by name instead of by + position: each public field receives the binding whose name matches it (the ``@ArgName`` + value, or the verbatim field name). See :ref:`java-sdk/arg-binding`. + +Besides the annotations, a task method may declare a ``Client`` and a ``Context`` parameter in any +position; the SDK injects both. Every other parameter is a *data parameter* and receives an +argument bound by the Python ``@task.stub`` call site. The annotation processor generates a ``Builder`` class that wires up the task -registry and handles XCom injection automatically. +registry and resolves data parameters and XCom pushes automatically. .. code-block:: java @@ -184,10 +189,7 @@ registry and handles XCom injection automatically. } @Builder.Task(id = "process") - public long process( - Client client, - @Builder.XCom(task = "fetch") String fetched - ) { + public long process(Client client, String fetched) { var threshold = (String) client.getVariable("process_threshold"); // implement task logic return count; @@ -234,6 +236,105 @@ Register tasks manually in a ``BundleBuilder``: See the `Java SDK API Reference `__ for more details. +.. _java-sdk/arg-binding: + +Binding stub arguments +~~~~~~~~~~~~~~~~~~~~~~ + +Calling a ``@task.stub`` TaskFlow-style in the Python Dag is what declares the graph, and the +supervisor delivers the resulting argument bindings to the Java runtime with every task run. A +binding carries either an upstream task's ``return_value`` XCom or an inline literal written at the +call site. + +Positional binding +^^^^^^^^^^^^^^^^^^ + +A task method's data parameters bind **by position**, in declaration order — the injected ``Client`` +and ``Context`` parameters do not take up a position. Java parameter names are not part of the API, +so renaming one in an IDE never rebinds an input. + +.. code-block:: python + + @task.stub(queue="java") + def score(rows, threshold): ... + + + score(load_rows(), 0.75) + +.. code-block:: java + + @Builder.Task(id = "score") + public long score(Client client, long rows, double threshold) { + // rows <- the load_rows XCom (position 0) + // threshold <- the literal 0.75 (position 1) + } + +A primitive parameter cannot hold ``null``, so the task fails with ``MissingXComException`` when its +binding resolves to nothing; declare a boxed type (``Long``, ``Double``, …) to receive ``null`` +instead. Declaring more data parameters than the call site bound also fails the task, rather than +running it with missing inputs. + +Named binding with a ``TaskInput`` bundle +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +To bind keyword arguments by name, declare a single parameter whose class implements ``TaskInput``. +Each public non-final field receives the binding named by its ``@ArgName`` value, or by the verbatim +field name — the deliberate, tagged boundary where the stub's ``snake_case`` argument names cross +into ``camelCase`` Java fields. The class needs a public no-argument constructor. + +.. code-block:: python + + @task.stub(queue="java") + def score(region_code, threshold): ... + + + score(region_code="emea", threshold=load_threshold()) + +.. code-block:: java + + public static class ScoreInput implements TaskInput { + @ArgName("region_code") + public String region; + + public double threshold; + } + + @Builder.Task(id = "score") + public long score(Client client, ScoreInput input) { ... } + +A task declares flat data parameters **or** one ``TaskInput`` bundle, never both, so field names and +flat positions cannot shift each other. Mixing them, or declaring two bundles, fails the build. + +Reading bindings from the interface API +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Tasks written against the ``Task`` interface read the same bindings imperatively: + +.. list-table:: + :header-rows: 1 + :widths: 40 60 + + * - ``Client`` method + - Returns + * - ``hasArgs()`` + - Whether the Python Dag called this stub with any TaskFlow arguments at all. + * - ``hasArg(int position)`` / ``hasArg(String name)`` + - Whether an argument was bound at that position, or with that name. + * - ``getArg(int position)`` / ``getArg(String name)`` + - The bound value — the inline literal, or the bound upstream's XCom. Throws + ``IllegalArgumentException`` when nothing was bound there; probe with ``hasArg`` first. + +.. code-block:: java + + public class ScoreTask implements Task { + @Override + public void execute(Context context, Client client) throws Exception { + var rows = client.getArg(0); + var threshold = client.hasArg("threshold") ? client.getArg("threshold") : 0.5; + // implement task logic + } + } + .. _java-sdk/logging: Logging @@ -428,7 +529,7 @@ represented as Java objects when read back via ``getXCom``. .. note:: - An ``@Builder.XCom`` parameter that reads a value which was never pushed resolves to + A data parameter whose binding resolves to a value that was never pushed receives ``null``. A boxed parameter (``Integer``, ``Long``, ``Boolean``, …) receives ``null`` safely, but a primitive parameter (``int``, ``long``, ``boolean``, …) cannot represent ``null`` and the task fails with ``MissingXComException``. Declare the parameter with a diff --git a/java-sdk/example/src/java/org/apache/airflow/example/AnnotationExample.java b/java-sdk/example/src/java/org/apache/airflow/example/AnnotationExample.java index bb715a73cb502..42ee68a745ce4 100644 --- a/java-sdk/example/src/java/org/apache/airflow/example/AnnotationExample.java +++ b/java-sdk/example/src/java/org/apache/airflow/example/AnnotationExample.java @@ -27,6 +27,9 @@ import java.util.concurrent.Executors; import org.apache.airflow.sdk.*; +// The Python Dag file (src/resources/dags/java_examples.py) owns the graph: each +// data parameter below receives whatever the @task.stub call site bound at its +// position, so this class registers task implementations only. @SuppressWarnings("DuplicatedCode") @Builder.Dag(id = "java_annotation_example") public class AnnotationExample { @@ -52,7 +55,7 @@ public long extractValue(Client client) throws InterruptedException { } @Builder.Task(id = "transform") - public long transformValue(Client client, @Builder.XCom(task = "extract") long extracted) { + public long transformValue(Client client, long extracted) { log.log(INFO, "Got XCom from extract: {0}", extracted); var variable = client.getVariable("my_variable"); @@ -68,7 +71,7 @@ public long transformValue(Client client, @Builder.XCom(task = "extract") long e // RetryTask (instead of a terminal FAILED) when ti_context.should_retry is // set. The retry then runs this task again and it returns normally. @Builder.Task - public void load(Context context, @Builder.XCom(task = "transform") long transformed) { + public void load(Context context, long transformed) { log.log(INFO, "Got XCom from transform: {0}", transformed); if (context.ti.tryNumber == 1) { throw new RuntimeException("I failed"); @@ -76,6 +79,24 @@ public void load(Context context, @Builder.XCom(task = "transform") long transfo log.log(INFO, "Recovered on retry, try number {0}", context.ti.tryNumber); } + // Keyword arguments bind by name instead of by position, through a TaskInput + // bundle: the tagged boundary where the stub's snake_case argument names + // cross into camelCase Java fields. + public static class ReportInput implements TaskInput { + @ArgName("run_label") + public String runLabel; + + public long transformed; + } + + @Builder.Task(id = "report") + public void report(ReportInput input) { + log.log(INFO, "Report {0} for transformed value {1}", input.runLabel, input.transformed); + if (!"nightly".equals(input.runLabel)) { + throw new RuntimeException("expected run label 'nightly' but got " + input.runLabel); + } + } + // Verify one supervisor channel can handle client calls across threads. @Builder.Task(id = "concurrent") public void concurrentClientCalls(Client client) throws Exception { diff --git a/java-sdk/example/src/java/org/apache/airflow/example/InterfaceExampleBuilder.java b/java-sdk/example/src/java/org/apache/airflow/example/InterfaceExampleBuilder.java index 74eb4f5d7533a..1ff317959d43f 100644 --- a/java-sdk/example/src/java/org/apache/airflow/example/InterfaceExampleBuilder.java +++ b/java-sdk/example/src/java/org/apache/airflow/example/InterfaceExampleBuilder.java @@ -52,8 +52,10 @@ public void execute(@NotNull Context context, Client client) throws Exception { public static class Transform implements Task { public void execute(@NotNull Context context, Client client) { - var extracted = client.getXCom("extract"); - log.log(INFO, "Got XCom from extract: {0}", extracted); + // The Python Dag file calls transform(extracted): an interface-API task + // reads the same binding imperatively, by position or by argument name. + var extracted = client.getArg(0); + log.log(INFO, "Got extracted value from the bound argument: {0}", extracted); var variable = client.getVariable("my_variable"); log.log(INFO, "Got variable: {0}", variable); @@ -65,7 +67,10 @@ public void execute(@NotNull Context context, Client client) { public static class Load implements Task { public void execute(@NotNull Context context, Client client) { - var transformed = client.getXCom("transform"); + // hasArg probes whether the Dag file bound an argument at all, so a task + // can still fall back to reading an upstream XCom by task id. + var transformed = + client.hasArg("transformed") ? client.getArg("transformed") : client.getXCom("transform"); log.log(INFO, "Got XCom from transform: {0}", transformed); throw new RuntimeException("I failed"); } diff --git a/java-sdk/example/src/java/org/apache/airflow/example/XComCastingExample.java b/java-sdk/example/src/java/org/apache/airflow/example/XComCastingExample.java index c4945af865df1..092d2e168f9e8 100644 --- a/java-sdk/example/src/java/org/apache/airflow/example/XComCastingExample.java +++ b/java-sdk/example/src/java/org/apache/airflow/example/XComCastingExample.java @@ -23,6 +23,9 @@ import org.apache.airflow.sdk.*; +// Stub-backed tasks wired by the Python Dag file: each parameter receives the +// value the stub call bound at its position, widening or narrowing to the +// declared type at run time. @Builder.Dag(id = "java_xcom_casting_example") public class XComCastingExample { private static final System.Logger log = System.getLogger(XComCastingExample.class.getName()); @@ -35,13 +38,13 @@ public int produceNumber() { // Any primitive numeric type (byte, short, int, long, float, double) and its boxed form works the same way. @Builder.Task(id = "widen_to_long") - public long widenToLong(@Builder.XCom(task = "produce_number") long value) { + public long widenToLong(long value) { log.log(INFO, "Got long {0}", value); return value + 1; } @Builder.Task(id = "widen_to_double") - public void widenToDouble(@Builder.XCom(task = "widen_to_long") double value) { + public void widenToDouble(double value) { log.log(INFO, "Got double {0}", value); if (value != 8.0) { throw new RuntimeException("expected 8.0 but got " + value); @@ -54,7 +57,7 @@ public void produceNothing() { } @Builder.Task(id = "consume_nullable") - public void consumeNullable(@Builder.XCom(task = "produce_nothing") Integer value) { + public void consumeNullable(Integer value) { log.log(INFO, "Got nullable int {0}", value); if (value != null) { throw new RuntimeException("expected null but got " + value); @@ -68,7 +71,7 @@ public double produceFraction() { } @Builder.Task(id = "consume_float") - public void consumeFloat(@Builder.XCom(task = "produce_fraction") float value) { + public void consumeFloat(float value) { log.log(INFO, "Got float {0}", value); if (value != 1.5f) { throw new RuntimeException("expected 1.5 but got " + value); diff --git a/java-sdk/example/src/resources/dags/java_examples.py b/java-sdk/example/src/resources/dags/java_examples.py index 5426911b0faee..49b79b892ad65 100644 --- a/java-sdk/example/src/resources/dags/java_examples.py +++ b/java-sdk/example/src/resources/dags/java_examples.py @@ -34,27 +34,32 @@ def extract(): ... @task.stub(queue="java") -def transform(): ... +def transform(extracted): ... @task.stub(queue="java", retries=1, retry_delay=timedelta(seconds=5)) -def load(): ... +def load(transformed): ... @task.stub(queue="java") def concurrent(): ... +# Keyword arguments bind to the public fields of the Java task's TaskInput bundle. +@task.stub(queue="java") +def report(run_label, transformed): ... + + @task.stub(queue="java") def produce_number(): ... @task.stub(queue="java") -def widen_to_long(): ... +def widen_to_long(value): ... @task.stub(queue="java") -def widen_to_double(): ... +def widen_to_double(value): ... @task.stub(queue="java") @@ -62,7 +67,7 @@ def produce_nothing(): ... @task.stub(queue="java") -def consume_nullable(): ... +def consume_nullable(value): ... @task.stub(queue="java") @@ -70,7 +75,7 @@ def produce_fraction(): ... @task.stub(queue="java") -def consume_float(): ... +def consume_float(value): ... @task() @@ -82,25 +87,28 @@ def python_task_2(transformed): @dag(dag_id="java_interface_example") def java_interface_example(): - transformed = transform() - python_task_1() >> extract() >> transformed + extracted = extract() + python_task_1() >> extracted + transformed = transform(extracted) python_task_2(transformed) @dag(dag_id="java_annotation_example") def java_annotation_example(): - transformed = transform() - python_task_1() >> extract() >> transformed + extracted = extract() + python_task_1() >> extracted + transformed = transform(extracted) python_task_2(transformed) - transformed >> load() + load(transformed) + report(run_label="nightly", transformed=transformed) concurrent() @dag(dag_id="java_xcom_casting_example") def java_xcom_casting_example(): - produce_number() >> widen_to_long() >> widen_to_double() - produce_nothing() >> consume_nullable() - produce_fraction() >> consume_float() + widen_to_double(widen_to_long(produce_number())) + consume_nullable(produce_nothing()) + consume_float(produce_fraction()) java_interface_example() diff --git a/java-sdk/gradle.properties b/java-sdk/gradle.properties index 9438ba6435532..477b31da386c4 100644 --- a/java-sdk/gradle.properties +++ b/java-sdk/gradle.properties @@ -17,7 +17,7 @@ org.gradle.configuration-cache=true -airflowSupervisorSchemaVersion=2026-06-16 +airflowSupervisorSchemaVersion=2026-10-30 projectVersion=1.0.0-SNAPSHOT diff --git a/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt b/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt index 56cbf1e76ad9f..3bd8eb336b95d 100644 --- a/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt +++ b/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt @@ -25,18 +25,21 @@ import com.squareup.javapoet.ClassName import com.squareup.javapoet.CodeBlock import com.squareup.javapoet.JavaFile import com.squareup.javapoet.MethodSpec +import com.squareup.javapoet.ParameterizedTypeName import com.squareup.javapoet.TypeName import com.squareup.javapoet.TypeSpec -import java.util.Optional +import org.apache.airflow.sdk.internal.ArgValues import javax.annotation.processing.AbstractProcessor import javax.annotation.processing.ProcessingEnvironment import javax.annotation.processing.RoundEnvironment import javax.annotation.processing.SupportedAnnotationTypes import javax.annotation.processing.SupportedSourceVersion import javax.lang.model.SourceVersion +import javax.lang.model.element.ElementKind import javax.lang.model.element.ExecutableElement import javax.lang.model.element.Modifier import javax.lang.model.element.TypeElement +import javax.lang.model.element.VariableElement import javax.lang.model.type.TypeKind import javax.lang.model.type.TypeMirror import javax.tools.Diagnostic @@ -57,9 +60,11 @@ import javax.tools.Diagnostic * - A static `build()` method that constructs the [DagDef] and registers those * inner classes as [TaskDef]s. * - * [Builder.XCom]-annotated parameters are resolved via `client.getXCom` in the - * generated `execute` body, with the result cast to the parameter's declared - * type. Non-`void` return values are forwarded to `client.setXCom`. + * In the generated `execute` body, a task's data parameters resolve through + * [ArgValues] against the arg bindings the supervisor delivered for the run: + * flat parameters by their position among the data parameters, [TaskInput] + * bundle fields by wire name. Non-`void` return values are forwarded to + * `client.setXCom`. */ @SupportedAnnotationTypes("org.apache.airflow.sdk.Builder.Dag") @SupportedSourceVersion(SourceVersion.RELEASE_11) @@ -102,23 +107,22 @@ class BuilderProcessor : AbstractProcessor() { MethodSpec .methodBuilder("build") .addModifiers(Modifier.PUBLIC, Modifier.STATIC) - .returns(ClassName.get(DagDef::class.java)) - .addStatement($$"var dag = new $T($S)", ClassName.get(DagDef::class.java), ann.id.ifBlank { el.simpleName }) + .returns(DAG_DEF_TYPE) + .addStatement($$"var dag = new $T($S)", DAG_DEF_TYPE, ann.id.ifBlank { el.simpleName }) for (inner in el.enclosedElements) { if (inner !is ExecutableElement) continue if (inner.isVarArgs) throw IllegalArgumentException("Cannot create task from vararg function ${inner.simpleName}") - val ann = inner.getAnnotation(Builder.Task::class.java) ?: continue + val taskAnn = inner.getAnnotation(Builder.Task::class.java) ?: continue val innerName = inner.simpleName.toString().replaceFirstChar(Char::uppercase) - val task = buildTask(innerName, inner, el) - builderClass.addType(task.spec) + builderClass.addType(buildTask(innerName, inner, el)) buildMethod.addStatement( $$"dag.addTask(new $T($S, $L.class))", - ClassName.get(TaskDef::class.java), - ann.id.ifBlank { inner.simpleName }, + TASK_DEF_TYPE, + taskAnn.id.ifBlank { inner.simpleName }, innerName, ) } @@ -132,40 +136,44 @@ class BuilderProcessor : AbstractProcessor() { name: String, inner: ExecutableElement, parent: TypeElement, - ): BuildTaskResult { - val clientType = ClassName.get(Client::class.java) - val contextType = ClassName.get(Context::class.java) - + ): TypeSpec { val executeSpec = MethodSpec .methodBuilder("execute") .addAnnotation(Override::class.java) .addModifiers(Modifier.PUBLIC) .returns(TypeName.VOID) - .addParameter(contextType, "context") - .addParameter(clientType, "client") + .addParameter(CONTEXT_TYPE, "context") + .addParameter(CLIENT_TYPE, "client") .addException(Exception::class.java) - val required = mutableListOf() + val dataParams = collectDataParams(inner) + val dataByName = dataParams.associateBy { it.name } val innerArgs = with(processingEnv) { inner.parameters.joinToString { param -> - val anno = param.getAnnotation(Builder.XCom::class.java) val type = param.asType() when { - anno != null -> - param.simpleName.toString().also { - required += RequiredXCom(type, it, anno.task.ifBlank { it }) - } - isType(type, clientType) -> "client" - isType(type, contextType) -> "context" - else -> throw IllegalArgumentException("Unsupported task parameter '${param.simpleName}' with type: $type") + isType(type, CLIENT_TYPE) -> "client" + isType(type, CONTEXT_TYPE) -> "context" + else -> dataByName.getValue(param.simpleName.toString()).name } } } - required.forEach { - executeSpec.addStatement($$"var $L = $L", it.paramName, xcomAccess(it)) + + dataParams.forEach { param -> + val paramType = TypeName.get(param.type) + val fields = param.bundleFields + if (fields == null) { + executeSpec.addStatement($$"$T $L = $L", paramType, param.name, positionalAccess(param)) + } else { + executeSpec.addStatement($$"$T $L = new $T()", paramType, param.name, paramType) + fields.forEach { field -> + executeSpec.addStatement($$"$L.$L = $L", param.name, field.fieldName, namedAccess(field)) + } + } } + if (inner.returnType.kind == TypeKind.VOID) { $$"new $T().$L($L)" } else { @@ -179,80 +187,154 @@ class BuilderProcessor : AbstractProcessor() { ) } - val spec = - TypeSpec - .classBuilder(name) - .addSuperinterface(Task::class.java) - .addModifiers(Modifier.PUBLIC, Modifier.FINAL, Modifier.STATIC) - .addMethod(executeSpec.build()) - .build() - return BuildTaskResult(spec) + return TypeSpec + .classBuilder(name) + .addSuperinterface(Task::class.java) + .addModifiers(Modifier.PUBLIC, Modifier.FINAL, Modifier.STATIC) + .addMethod(executeSpec.build()) + .build() + } + + /** + * Collects the task method's data parameters — every parameter the SDK does + * not inject — in declaration order. A parameter's index in the returned + * list is the position it binds at: Java parameter names are not API, so + * renaming one must never rebind an input. + */ + private fun collectDataParams(method: ExecutableElement): List { + val params = mutableListOf() + with(processingEnv) { + for (param in method.parameters) { + val type = param.asType() + if (isType(type, CLIENT_TYPE) || isType(type, CONTEXT_TYPE)) continue + val bundleFields = if (isTaskInput(type)) collectBundleFields(method, param) else null + params += DataParam(type, param.simpleName.toString(), params.size, bundleFields) + } + } + val bundles = params.filter { it.bundleFields != null } + require(bundles.size <= 1) { + "Task method '${method.simpleName}' declares more than one TaskInput parameter: " + + bundles.joinToString { "'${it.name}'" } + } + bundles.singleOrNull()?.let { bundle -> + require(params.size == 1) { + "Task method '${method.simpleName}' declares TaskInput parameter '${bundle.name}' and other data " + + "parameters; a TaskInput bundle owns the whole named-argument surface, so it must be the only one" + } + } + return params + } + + private fun ProcessingEnvironment.isTaskInput(type: TypeMirror): Boolean { + val marker = elementUtils.getTypeElement(TASK_INPUT_TYPE.canonicalName()) ?: return false + return !type.kind.isPrimitive && typeUtils.isAssignable(type, marker.asType()) + } + + /** + * Introspects a [TaskInput] bundle class: every public non-static non-final + * field receives the binding named by its [ArgName] value, or by its + * verbatim field name. + */ + private fun ProcessingEnvironment.collectBundleFields( + method: ExecutableElement, + param: VariableElement, + ): List { + val bundleType = + typeUtils.asElement(param.asType()) as? TypeElement + ?: throw IllegalArgumentException( + "TaskInput parameter '${param.simpleName}' of task method '${method.simpleName}' has no class type", + ) + val hasNoArgConstructor = + bundleType.enclosedElements + .filterIsInstance() + .any { it.kind == ElementKind.CONSTRUCTOR && it.parameters.isEmpty() && Modifier.PUBLIC in it.modifiers } + require(hasNoArgConstructor) { + "TaskInput class ${bundleType.simpleName} needs a public no-argument constructor" + } + return bundleType.enclosedElements + .filterIsInstance() + .filter { it.kind == ElementKind.FIELD && Modifier.STATIC !in it.modifiers } + .map { field -> + require(Modifier.PUBLIC in field.modifiers && Modifier.FINAL !in field.modifiers) { + "TaskInput field ${bundleType.simpleName}.${field.simpleName} must be public and non-final " + + "so the generated code can assign its binding" + } + BundleField( + type = field.asType(), + fieldName = field.simpleName.toString(), + wireName = field.getAnnotation(ArgName::class.java)?.value ?: field.simpleName.toString(), + ) + } } } +/** + * One data parameter of a task method, positioned among its peers. + * [bundleFields] is non-null for a [TaskInput] bundle parameter. + */ +private class DataParam( + val type: TypeMirror, + val name: String, + val position: Int, + val bundleFields: List?, +) + +/** One public field of a [TaskInput] bundle class, with its wire name. */ +private class BundleField( + val type: TypeMirror, + val fieldName: String, + val wireName: String, +) + +private val DAG_DEF_TYPE = ClassName.get(DagDef::class.java) +private val TASK_DEF_TYPE = ClassName.get(TaskDef::class.java) +private val CLIENT_TYPE = ClassName.get(Client::class.java) +private val CONTEXT_TYPE = ClassName.get(Context::class.java) +private val TASK_INPUT_TYPE = ClassName.get(TaskInput::class.java) +private val ARG_VALUES_TYPE = ClassName.get(ArgValues::class.java) + private fun ProcessingEnvironment.isType( t: TypeMirror, c: ClassName, ): Boolean = typeUtils.isSameType(t, elementUtils.getTypeElement(c.canonicalName()).asType()) -private data class RequiredXCom( - val paramType: TypeMirror, - val paramName: String, - val taskId: String, -) - -private val NUMBER_ACCESSORS: Map = - buildMap { - mapOf( - TypeName.BYTE to "byteValue", - TypeName.SHORT to "shortValue", - TypeName.INT to "intValue", - TypeName.LONG to "longValue", - TypeName.FLOAT to "floatValue", - TypeName.DOUBLE to "doubleValue", - ).forEach { (primitive, accessor) -> - put(primitive, accessor) - put(primitive.box(), accessor) - } +/** + * Emits the resolve-and-decode expression for one flat data parameter, bound + * at its position. A primitive parameter cannot hold null, so it fails with a + * clear [MissingXComException] when the binding resolves to nothing; boxed and + * reference parameters receive null instead. + */ +private fun positionalAccess(param: DataParam): CodeBlock { + val type = TypeName.get(param.type) + return if (type.isPrimitive) { + CodeBlock.of( + $$"$T.requiredInput(context, client, $L, $T.class, $S)", + ARG_VALUES_TYPE, + param.position, + type.box(), + param.name, + ) + } else { + val raw = (type as? ParameterizedTypeName)?.rawType ?: type + val call = CodeBlock.of($$"$T.optionalInput(context, client, $L, $T.class)", ARG_VALUES_TYPE, param.position, raw) + if (type is ParameterizedTypeName) CodeBlock.of($$"($T) $L", type, call) else call } +} -private fun xcomAccess(xcom: RequiredXCom): CodeBlock { - val type = TypeName.get(xcom.paramType) - val accessor = NUMBER_ACCESSORS[type] - val number = ClassName.get(Number::class.java) - val optional = ClassName.get(Optional::class.java) - // A primitive parameter cannot hold null, so fail with a clear error instead of an - // opaque NullPointerException while unboxing when the XCom is absent. - val value = - if (type.isPrimitive) { - CodeBlock.of( - $$"$T.ofNullable(client.getXCom($S)).orElseThrow(() -> new $T($S, $S))", - optional, - xcom.taskId, - ClassName.get(MissingXComException::class.java), - xcom.taskId, - xcom.paramName, - ) - } else { - CodeBlock.of($$"client.getXCom($S)", xcom.taskId) - } - // Wire integers decode to Long and floats to Double, so a direct (Integer)/(Float) - // cast throws ClassCastException; widen via Number instead. - return when { - accessor == null -> CodeBlock.of($$"($T) $L", if (type.isPrimitive) type.box() else type, value) - type.isPrimitive -> CodeBlock.of($$"(($T) $L).$L()", number, value, accessor) - else -> - CodeBlock.of( - $$"$T.ofNullable(($T) $L).map($T::$L).orElse(null)", - optional, - number, - value, - number, - accessor, - ) +/** Emits the resolve-and-decode expression for one bundle field, bound by wire name. */ +private fun namedAccess(field: BundleField): CodeBlock { + val type = TypeName.get(field.type) + return if (type.isPrimitive) { + CodeBlock.of( + $$"$T.requiredNamed(client, $S, $T.class, $S)", + ARG_VALUES_TYPE, + field.wireName, + type.box(), + field.fieldName, + ) + } else { + val raw = (type as? ParameterizedTypeName)?.rawType ?: type + val call = CodeBlock.of($$"$T.optionalNamed(client, $S, $T.class)", ARG_VALUES_TYPE, field.wireName, raw) + if (type is ParameterizedTypeName) CodeBlock.of($$"($T) $L", type, call) else call } } - -private data class BuildTaskResult( - val spec: TypeSpec, -) diff --git a/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt b/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt index 6a08979c56163..679aeea4338dd 100644 --- a/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt +++ b/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt @@ -62,7 +62,7 @@ class BuilderTest { } @Builder.Task - public void t3(Context ctx, @Builder.XCom(task = "t2") int value) { + public void t3(Context ctx, int value) { System.out.println(String.format("%s %s", ctx.ti, value)); } } @@ -78,15 +78,14 @@ class BuilderTest { package org.apache.airflow.example; import java.lang.Exception; - import java.lang.Number; + import java.lang.Integer; import java.lang.Override; - import java.util.Optional; import org.apache.airflow.sdk.Client; import org.apache.airflow.sdk.Context; import org.apache.airflow.sdk.DagDef; - import org.apache.airflow.sdk.MissingXComException; import org.apache.airflow.sdk.Task; import org.apache.airflow.sdk.TaskDef; + import org.apache.airflow.sdk.internal.ArgValues; public final class TestExampleBuilder { public static DagDef build() { @@ -111,7 +110,7 @@ class BuilderTest { public static final class T3 implements Task { @Override public void execute(Context context, Client client) throws Exception { - var value = ((Number) Optional.ofNullable(client.getXCom("t2")).orElseThrow(() -> new MissingXComException("t2", "value"))).intValue(); + int value = ArgValues.requiredInput(context, client, 0, Integer.class, "value"); new TestExample().t3(context, value); } } @@ -121,25 +120,19 @@ class BuilderTest { } @Test - @DisplayName("widen primitive numerics directly and boxed numerics null-safely") - fun generateBuilderWidensNumericXCom() { + @DisplayName("bind data parameters by position, skipping the injected Client and Context") + fun generateBuilderBindsDataParametersByPosition() { val compilation = compile( """ package org.apache.airflow.example; import org.apache.airflow.sdk.Builder; + import org.apache.airflow.sdk.Client; + import org.apache.airflow.sdk.Context; @Builder.Dag public class TestExample { @Builder.Task - public void t( - @Builder.XCom(task = "a") int i, - @Builder.XCom(task = "b") long l, - @Builder.XCom(task = "c") double d, - @Builder.XCom(task = "f") float fl, - @Builder.XCom(task = "e") Integer boxedInteger, - @Builder.XCom(task = "g") Long boxedLong, - @Builder.XCom(task = "h") Double boxedDouble, - @Builder.XCom(task = "j") Float boxedFloat) {} + public void t(long first, Client client, String second, Context ctx, Integer third) {} } """, ) @@ -153,15 +146,16 @@ class BuilderTest { package org.apache.airflow.example; import java.lang.Exception; - import java.lang.Number; + import java.lang.Integer; + import java.lang.Long; import java.lang.Override; - import java.util.Optional; + import java.lang.String; import org.apache.airflow.sdk.Client; import org.apache.airflow.sdk.Context; import org.apache.airflow.sdk.DagDef; - import org.apache.airflow.sdk.MissingXComException; import org.apache.airflow.sdk.Task; import org.apache.airflow.sdk.TaskDef; + import org.apache.airflow.sdk.internal.ArgValues; public final class TestExampleBuilder { public static DagDef build() { @@ -172,15 +166,10 @@ class BuilderTest { public static final class T implements Task { @Override public void execute(Context context, Client client) throws Exception { - var i = ((Number) Optional.ofNullable(client.getXCom("a")).orElseThrow(() -> new MissingXComException("a", "i"))).intValue(); - var l = ((Number) Optional.ofNullable(client.getXCom("b")).orElseThrow(() -> new MissingXComException("b", "l"))).longValue(); - var d = ((Number) Optional.ofNullable(client.getXCom("c")).orElseThrow(() -> new MissingXComException("c", "d"))).doubleValue(); - var fl = ((Number) Optional.ofNullable(client.getXCom("f")).orElseThrow(() -> new MissingXComException("f", "fl"))).floatValue(); - var boxedInteger = Optional.ofNullable((Number) client.getXCom("e")).map(Number::intValue).orElse(null); - var boxedLong = Optional.ofNullable((Number) client.getXCom("g")).map(Number::longValue).orElse(null); - var boxedDouble = Optional.ofNullable((Number) client.getXCom("h")).map(Number::doubleValue).orElse(null); - var boxedFloat = Optional.ofNullable((Number) client.getXCom("j")).map(Number::floatValue).orElse(null); - new TestExample().t(i, l, d, fl, boxedInteger, boxedLong, boxedDouble, boxedFloat); + long first = ArgValues.requiredInput(context, client, 0, Long.class, "first"); + String second = ArgValues.optionalInput(context, client, 1, String.class); + Integer third = ArgValues.optionalInput(context, client, 2, Integer.class); + new TestExample().t(first, client, second, context, third); } } } @@ -189,20 +178,19 @@ class BuilderTest { } @Test - @DisplayName("guard non-numeric primitives, leave objects and boxed types nullable") - fun generateBuilderGuardsNonNumericPrimitiveXCom() { + @DisplayName("require primitive parameters, leave boxed and parameterized types nullable") + fun generateBuilderRequiresPrimitivesOnly() { val compilation = compile( """ package org.apache.airflow.example; + import java.util.List; + import java.util.Map; import org.apache.airflow.sdk.Builder; @Builder.Dag public class TestExample { @Builder.Task - public void t( - @Builder.XCom(task = "a") boolean flag, - @Builder.XCom(task = "b") String text, - @Builder.XCom(task = "c") Boolean boxed) {} + public void t(boolean flag, float fraction, Double boxed, List tags, Map raw) {} } """, ) @@ -216,16 +204,19 @@ class BuilderTest { package org.apache.airflow.example; import java.lang.Boolean; + import java.lang.Double; import java.lang.Exception; + import java.lang.Float; import java.lang.Override; import java.lang.String; - import java.util.Optional; + import java.util.List; + import java.util.Map; import org.apache.airflow.sdk.Client; import org.apache.airflow.sdk.Context; import org.apache.airflow.sdk.DagDef; - import org.apache.airflow.sdk.MissingXComException; import org.apache.airflow.sdk.Task; import org.apache.airflow.sdk.TaskDef; + import org.apache.airflow.sdk.internal.ArgValues; public final class TestExampleBuilder { public static DagDef build() { @@ -236,10 +227,12 @@ class BuilderTest { public static final class T implements Task { @Override public void execute(Context context, Client client) throws Exception { - var flag = (Boolean) Optional.ofNullable(client.getXCom("a")).orElseThrow(() -> new MissingXComException("a", "flag")); - var text = (String) client.getXCom("b"); - var boxed = (Boolean) client.getXCom("c"); - new TestExample().t(flag, text, boxed); + boolean flag = ArgValues.requiredInput(context, client, 0, Boolean.class, "flag"); + float fraction = ArgValues.requiredInput(context, client, 1, Float.class, "fraction"); + Double boxed = ArgValues.optionalInput(context, client, 2, Double.class); + List tags = (List) ArgValues.optionalInput(context, client, 3, List.class); + Map raw = ArgValues.optionalInput(context, client, 4, Map.class); + new TestExample().t(flag, fraction, boxed, tags, raw); } } } @@ -247,6 +240,179 @@ class BuilderTest { ) } + @Test + @DisplayName("bind input-bundle fields by wire name") + fun generateBuilderBindsInputBundleFields() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import java.util.List; + import org.apache.airflow.sdk.ArgName; + import org.apache.airflow.sdk.Builder; + import org.apache.airflow.sdk.Client; + import org.apache.airflow.sdk.TaskInput; + @Builder.Dag + public class TestExample { + public static class ScoreInput implements TaskInput { + @ArgName("region_code") public String region; + public double threshold; + public List tags; + } + + @Builder.Task + public double score(Client client, ScoreInput input) { return input.threshold; } + } + """, + ) + + assertThat(compilation).succeeded() + assertThat(compilation) + .generatedSourceFile("org.apache.airflow.example.TestExampleBuilder") + .hasSourceEquivalentTo( + "org.apache.airflow.example.TestExampleBuilder", + """ + package org.apache.airflow.example; + + import java.lang.Double; + import java.lang.Exception; + import java.lang.Override; + import java.lang.String; + import java.util.List; + import org.apache.airflow.sdk.Client; + import org.apache.airflow.sdk.Context; + import org.apache.airflow.sdk.DagDef; + import org.apache.airflow.sdk.Task; + import org.apache.airflow.sdk.TaskDef; + import org.apache.airflow.sdk.internal.ArgValues; + + public final class TestExampleBuilder { + public static DagDef build() { + var dag = new DagDef("TestExample"); + dag.addTask(new TaskDef("score", Score.class)); + return dag; + } + public static final class Score implements Task { + @Override + public void execute(Context context, Client client) throws Exception { + TestExample.ScoreInput input = new TestExample.ScoreInput(); + input.region = ArgValues.optionalNamed(client, "region_code", String.class); + input.threshold = ArgValues.requiredNamed(client, "threshold", Double.class, "threshold"); + input.tags = (List) ArgValues.optionalNamed(client, "tags", List.class); + client.setXCom(new TestExample().score(client, input)); + } + } + } + """, + ) + } + + @Test + @DisplayName("reject an input bundle mixed with flat data parameters") + fun rejectBundleMixedWithFlatParams() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + import org.apache.airflow.sdk.TaskInput; + @Builder.Dag + public class TestExample { + public static class ScoreInput implements TaskInput { + public double threshold; + } + + @Builder.Task + public void t(ScoreInput input, int extra) {} + } + """, + ) + assertThat(compilation).failed() + assertThat(compilation).hadErrorContaining( + "Task method 't' declares TaskInput parameter 'input' and other data parameters", + ) + } + + @Test + @DisplayName("reject a task declaring more than one input bundle") + fun rejectMultipleBundles() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + import org.apache.airflow.sdk.TaskInput; + @Builder.Dag + public class TestExample { + public static class ScoreInput implements TaskInput { + public double threshold; + } + + @Builder.Task + public void t(ScoreInput first, ScoreInput second) {} + } + """, + ) + assertThat(compilation).failed() + assertThat(compilation).hadErrorContaining( + "Task method 't' declares more than one TaskInput parameter: 'first', 'second'", + ) + } + + @Test + @DisplayName("reject an input bundle with a non-public field") + fun rejectBundleWithNonPublicField() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + import org.apache.airflow.sdk.TaskInput; + @Builder.Dag + public class TestExample { + public static class ScoreInput implements TaskInput { + double threshold; + } + + @Builder.Task + public void t(ScoreInput input) {} + } + """, + ) + assertThat(compilation).failed() + assertThat(compilation).hadErrorContaining( + "TaskInput field ScoreInput.threshold must be public and non-final", + ) + } + + @Test + @DisplayName("reject an input bundle without a public no-argument constructor") + fun rejectBundleWithoutNoArgConstructor() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + import org.apache.airflow.sdk.TaskInput; + @Builder.Dag + public class TestExample { + public static class ScoreInput implements TaskInput { + public double threshold; + + public ScoreInput(double threshold) { this.threshold = threshold; } + } + + @Builder.Task + public void t(ScoreInput input) {} + } + """, + ) + assertThat(compilation).failed() + assertThat(compilation).hadErrorContaining( + "TaskInput class ScoreInput needs a public no-argument constructor", + ) + } + @Test @DisplayName("generate builder for dag class with custom dag id") fun generateBuilderWithCustomDagId() { @@ -333,24 +499,6 @@ class BuilderTest { ) } - @Test - @DisplayName("generate builder for dag class with invalid task parameter") - fun generateBuilderForDagClassWithInvalidTaskParameter() { - val compilation = - compile( - """ - package org.apache.airflow.example; - import org.apache.airflow.sdk.Builder; - @Builder.Dag - public class TestExample { @Builder.Task(id = "foo") public void t1(String client) {} } - """, - ) - assertThat(compilation).failed() - assertThat(compilation).hadErrorContaining( - "Unsupported task parameter 'client' with type: java.lang.String", - ) - } - @Test @DisplayName("generate builder for dag class with varargs task parameter") fun generateBuilderForDagClassWithVarArgsTaskParameter() { diff --git a/java-sdk/sdk/schema/schema.json b/java-sdk/sdk/schema/schema.json index e6ce8aa3d066e..b671959c50a00 100644 --- a/java-sdk/sdk/schema/schema.json +++ b/java-sdk/sdk/schema/schema.json @@ -1,6 +1,6 @@ { "$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": { "AssetAliasReferenceAssetEventDagRun": { @@ -753,6 +753,19 @@ ], "title": "Bundle Version" }, + "version_data": { + "anyOf": [ + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Version Data" + }, "msg": { "anyOf": [ { @@ -1651,6 +1664,19 @@ ], "title": "Bundle Version" }, + "version_data": { + "anyOf": [ + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Version Data" + }, "msg": { "anyOf": [ { @@ -1846,6 +1872,45 @@ "title": "Ascending", "type": "boolean" }, + "partition_key": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Partition Key" + }, + "partition_key_regexp_pattern": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Partition Key Regexp Pattern" + }, + "extra": { + "anyOf": [ + { + "additionalProperties": { + "type": "string" + }, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Extra" + }, "type": { "const": "GetAssetEventByAsset", "default": "GetAssetEventByAsset", @@ -1909,6 +1974,45 @@ "title": "Ascending", "type": "boolean" }, + "partition_key": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Partition Key" + }, + "partition_key_regexp_pattern": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Partition Key Regexp Pattern" + }, + "extra": { + "anyOf": [ + { + "additionalProperties": { + "type": "string" + }, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Extra" + }, "type": { "const": "GetAssetEventByAssetAlias", "default": "GetAssetEventByAssetAlias", @@ -3859,6 +3963,19 @@ ], "title": "Bundle Version" }, + "version_data": { + "anyOf": [ + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Version Data" + }, "msg": { "anyOf": [ { @@ -4446,6 +4563,131 @@ "title": "ConnectionResponse", "type": "object" }, + "ArgValueSchema": { + "additionalProperties": { + "$ref": "#/$defs/JsonValue" + }, + "title": "ArgValueSchema", + "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.", @@ -4809,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/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/ArgName.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/ArgName.kt new file mode 100644 index 0000000000000..3af778d7cd71b --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/ArgName.kt @@ -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. + */ + +package org.apache.airflow.sdk + +/** + * Declares the wire name of a [TaskInput] field explicitly, for when the + * Python stub signature's argument name is not a valid (or desirable) Java + * field name — typically `snake_case` arguments crossing into `camelCase` + * fields. + * + * ```java + * @ArgName("region_code") public String region; + * ``` + * + * Fields without the annotation bind their verbatim field name. + * + * @param value Argument name as declared in the stub task's signature. + */ +@Target(AnnotationTarget.FIELD) +@MustBeDocumented +annotation class ArgName( + val value: String, +) diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Builder.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Builder.kt index 9dbfbfbefc390..1e8ca63670838 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Builder.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Builder.kt @@ -36,10 +36,15 @@ package org.apache.airflow.sdk * public long extract(Client client) { ... } * * @Builder.Task(id = "transform") - * public long transform(Client client, @Builder.XCom(task = "extract") long extracted) { ... } + * public long transform(Client client, long extracted) { ... } * } * ``` * + * A task method's data parameters — everything other than the injected + * [Client] and [Context] — receive the arguments the Python `@task.stub` call + * site bound, by position. Keyword arguments bind by name instead through a + * single [TaskInput] bundle parameter. + * * The processor generates `MyPipelineBuilder.build()`, which returns a * fully wired-up [DagDef] ready to add to a [Bundle]. */ @@ -73,16 +78,4 @@ class Builder internal constructor() { annotation class Task( val id: String = "", ) - - /** - * Annotation to mark a task definition's method parameter as an XCom input. - * - * @param task The task ID to pull. If empty or not given, the annotated - * parameter's name is used by default. - */ - @Target(AnnotationTarget.VALUE_PARAMETER) - @MustBeDocumented - annotation class XCom( - val task: String = "", - ) } diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Client.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Client.kt index 59aa7832a3756..941ac86aad7ac 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Client.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Client.kt @@ -19,8 +19,10 @@ package org.apache.airflow.sdk +import org.apache.airflow.sdk.execution.ArgBinding import org.apache.airflow.sdk.execution.Client import org.apache.airflow.sdk.execution.comm.StartupDetails +import org.apache.airflow.sdk.execution.decodeArgBindings /** * A connection registered in Airflow's connection store. @@ -152,6 +154,98 @@ class Client internal constructor( runId = details.ti.runId, mapIndex = details.ti.mapIndex ?: -1, ) + + internal val argBindings: List by lazy { + decodeArgBindings(details.tiContext?.argBindings) + } + + // A literal binding carries the inline value from the Dag file; an XCom + // binding pulls the bound upstream task's return-value XCom, honouring the + // bound map index and element index. + internal fun resolveBinding(binding: ArgBinding): Any? = + when (binding) { + is ArgBinding.Literal -> binding.value + is ArgBinding.XCom -> { + val value = getXCom(taskId = binding.taskId, mapIndex = binding.mapIndex.takeIf { it >= 0 }) + when { + binding.elementIndex == null -> value + value is List<*> -> value[binding.elementIndex] + else -> + error( + "Argument '${binding.name}' binds element ${binding.elementIndex} of task '${binding.taskId}', " + + "but its XCom is not a list", + ) + } + } + } + + /** + * Whether the supervisor delivered TaskFlow arg bindings for this run — + * that is, the Python Dag file called this stub task with TaskFlow + * arguments. + */ + fun hasArgs(): Boolean = argBindings.isNotEmpty() + + /** + * Whether a TaskFlow argument was bound at [position] of the stub task's + * signature. + * + * @param position Zero-based position in the stub call's argument list. + */ + fun hasArg(position: Int): Boolean = position in argBindings.indices + + /** + * Whether the Python Dag file bound a TaskFlow argument with this name to + * the current stub task. + * + * @param name Argument name as declared in the stub task's signature. + */ + fun hasArg(name: String): Boolean = argBindings.any { it.name == name } + + /** + * Resolves the TaskFlow argument bound at [position] of the stub task's + * signature. + * + * A literal binding returns the inline value from the Dag file; an XCom + * binding pulls the bound upstream task's return-value XCom, honouring the + * bound map index and element index. + * + * @param position Zero-based position in the stub call's argument list. + * @return The bound value, or `null` when the bound value is null or the + * upstream pushed no value. + * @throws IllegalArgumentException if no argument was bound at this + * position; use [hasArg] to probe. + * @throws ApiError if the underlying XCom read fails. + */ + fun getArg(position: Int): Any? { + require(position in argBindings.indices) { + "No TaskFlow argument bound at position: $position" + } + return resolveBinding(argBindings[position]) + } + + /** + * Resolves the TaskFlow argument bound with [name] at the `@task.stub` + * call site in the Python Dag file. + * + * A literal binding returns the inline value from the Dag file; an XCom + * binding pulls the bound upstream task's return-value XCom, honouring the + * bound map index and element index. + * + * @param name Argument name as declared in the stub task's signature. + * @return The bound value, or `null` when the bound value is null or the + * upstream pushed no value. + * @throws IllegalArgumentException if no argument with this name was bound; + * use [hasArg] to probe. + * @throws ApiError if the underlying XCom read fails. + */ + fun getArg(name: String): Any? { + val binding = + requireNotNull(argBindings.firstOrNull { it.name == name }) { + "No TaskFlow argument bound with name: '$name'" + } + return resolveBinding(binding) + } } /** diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/TaskInput.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/TaskInput.kt new file mode 100644 index 0000000000000..306600e0ff443 --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/TaskInput.kt @@ -0,0 +1,45 @@ +/* + * 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. + */ + +package org.apache.airflow.sdk + +/** + * Marks a class as a task's input bundle: when the Python Dag file calls the + * stub task with keyword arguments, each public field receives the runtime + * binding whose name matches the field ([ArgName] value, or the verbatim + * field name). + * + * A task method may declare at most one `TaskInput` parameter and, if it + * does, no other data parameters — the bundle owns the whole named-argument + * surface, so field names and flat positions cannot shift each other. + * + * ```java + * public static class ScoreInput implements TaskInput { + * @ArgName("region_code") public String region; // explicit wire name + * public double threshold; // binds "threshold" + * } + * + * @Builder.Task + * public Result score(Client client, ScoreInput input) { ... } + * ``` + * + * The class needs a public no-argument constructor and public non-final + * fields. + */ +interface TaskInput diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/ArgBinding.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/ArgBinding.kt new file mode 100644 index 0000000000000..808eff687d51e --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/ArgBinding.kt @@ -0,0 +1,77 @@ +/* + * 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. + */ + +package org.apache.airflow.sdk.execution + +/** + * One stub-task argument bound at the `@task.stub` TaskFlow call site in the + * Python Dag file, delivered via `TIRunContext.arg_bindings`. + * + * The supervisor schema models this as a `kind`-discriminated union + * (`XComArgBinding` / `LiteralArgBinding`), which jsonSchema2Pojo cannot + * express as a typed field — the generated `TIRunContext.argBindings` is a + * plain `Object` holding the msgpack-decoded list of maps — so this hand- + * written decoder materializes the typed view. + */ +internal sealed class ArgBinding { + abstract val name: String + + internal data class XCom( + override val name: String, + val taskId: String, + val mapIndex: Int, + val elementIndex: Int?, + ) : ArgBinding() + + internal data class Literal( + override val name: String, + val value: Any?, + ) : ArgBinding() +} + +/** + * Decodes the raw `TIRunContext.argBindings` payload into a list of bindings + * preserving the stub signature's parameter order — flat task parameters + * bind by that position, input-bundle fields by [ArgBinding.name]. + * + * @throws IllegalStateException on a malformed payload, an unsupported + * binding kind, or a duplicate argument name; the task cannot bind its + * arguments correctly, so it must fail rather than run with wrong inputs. + */ +internal fun decodeArgBindings(raw: Any?): List { + if (raw == null) return emptyList() + check(raw is List<*>) { "arg_bindings payload is not a list: ${raw.javaClass.name}" } + val seen = mutableSetOf() + return raw.map { entry -> + check(entry is Map<*, *>) { "arg_bindings entry is not a map: $entry" } + val name = checkNotNull(entry["name"] as? String) { "arg_bindings entry has no name: $entry" } + check(seen.add(name)) { "arg_bindings entries have duplicate name: '$name'" } + when (val kind = entry["kind"]) { + "literal" -> ArgBinding.Literal(name = name, value = entry["value"]) + "xcom" -> + ArgBinding.XCom( + name = name, + taskId = checkNotNull(entry["task_id"] as? String) { "xcom arg binding '$name' has no task_id" }, + mapIndex = (entry["map_index"] as? Number)?.toInt() ?: -1, + elementIndex = (entry["element_index"] as? Number)?.toInt(), + ) + else -> error("Unsupported arg binding kind '$kind' for argument '$name'") + } + } +} diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/ArgValues.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/ArgValues.kt new file mode 100644 index 0000000000000..cd56f38470668 --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/ArgValues.kt @@ -0,0 +1,175 @@ +/* + * 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. + */ + +@file:Suppress("PLATFORM_CLASS_MAPPED_TO_KOTLIN") + +package org.apache.airflow.sdk.internal + +import com.fasterxml.jackson.databind.ObjectMapper +import com.fasterxml.jackson.databind.json.JsonMapper +import org.apache.airflow.sdk.Client +import org.apache.airflow.sdk.Context +import org.apache.airflow.sdk.MissingXComException +import org.apache.airflow.sdk.execution.ArgBinding + +/** + * Resolves a task's data parameters from the arg bindings the supervisor + * delivered, and decodes their raw wire values into the declared parameter + * types. Public so that processor-generated task classes can call it; not + * user-facing API. + * + * The bindings come from the Python `@task.stub` call site, which is also the + * graph the scheduler ordered the run by. Flat data parameters resolve the + * binding at their position; input-bundle fields resolve bindings by name. + */ +object ArgValues { + private val mapper: ObjectMapper = JsonMapper.builder().build().findAndRegisterModules() + + /** + * Resolves the data parameter at [position] into [type] for a parameter that + * cannot be null. + * + * @param position Zero-based index among the task's data parameters, in + * declaration order. + * @throws MissingXComException if the resolved value is null. + */ + @JvmStatic + fun requiredInput( + context: Context, + client: Client, + position: Int, + type: Class, + paramName: String, + ): T { + val binding = bindingAt(context, client, position) + return decode(client.resolveBinding(binding), type) ?: throw missing(binding, paramName) + } + + /** + * Resolves the data parameter at [position] into [type], passing null + * through. + * + * @param position Zero-based index among the task's data parameters, in + * declaration order. + */ + @JvmStatic + fun optionalInput( + context: Context, + client: Client, + position: Int, + type: Class, + ): T? = decode(client.resolveBinding(bindingAt(context, client, position)), type) + + /** + * Resolves the runtime binding named [name] into [type] for an input-bundle + * field that cannot be null. + * + * @param name Wire name of the argument (`@ArgName` value or the verbatim + * field name). + * @throws IllegalStateException if the stub call bound no argument named + * [name]. + * @throws MissingXComException if the resolved value is null. + */ + @JvmStatic + fun requiredNamed( + client: Client, + name: String, + type: Class, + fieldName: String, + ): T { + val binding = + checkNotNull(client.argBindings.firstOrNull { it.name == name }) { + "The stub call bound no argument named '$name', required by input field '$fieldName'" + } + return decode(client.resolveBinding(binding), type) ?: throw missing(binding, fieldName, name) + } + + /** + * Resolves the runtime binding named [name] into [type], passing null + * through. An absent binding resolves to null. + * + * @param name Wire name of the argument (`@ArgName` value or the verbatim + * field name). + */ + @JvmStatic + fun optionalNamed( + client: Client, + name: String, + type: Class, + ): T? { + val binding = client.argBindings.firstOrNull { it.name == name } ?: return null + return decode(client.resolveBinding(binding), type) + } + + private fun bindingAt( + context: Context, + client: Client, + position: Int, + ): ArgBinding { + val bindings = client.argBindings + check(position < bindings.size) { + "Task '${context.ti.taskId}' declares a data parameter at position $position " + + "but the stub call bound only ${bindings.size} argument(s)" + } + return bindings[position] + } + + private fun missing( + binding: ArgBinding, + target: String, + argName: String? = null, + ): MissingXComException = + when (binding) { + is ArgBinding.XCom -> MissingXComException(binding.taskId, target) + is ArgBinding.Literal -> + MissingXComException( + "'$target' has a primitive type but the stub call bound a null literal" + + (argName?.let { " for argument '$it'" } ?: "") + + "; declare a boxed type (e.g. Integer instead of int) to receive null.", + ) + } + + internal fun decode( + value: Any?, + type: Class, + ): T? { + if (value == null) return null + if (type.isInstance(value)) return type.cast(value) + // The msgpack decoder yields Long for wire integers and Double for wire + // floats, so widen numerics via Number instead of casting. + if (value is Number) { + numberConverter(type)?.let { return type.cast(it(value)) } + } + // Structured wire values (maps, lists) convert into the declared POJO or + // collection type; unknown fields fail the task, mirroring the Go SDK's + // strict decode of task inputs. + return mapper.convertValue(value, type) + } + + private fun numberConverter(type: Class<*>): ((Number) -> Any)? = + when (type) { + java.lang.Byte::class.java -> { n -> n.toByte() } + java.lang.Short::class.java -> { n -> n.toShort() } + java.lang.Integer::class.java -> { n -> n.toInt() } + java.lang.Long::class.java -> { n -> n.toLong() } + java.lang.Float::class.java -> { n -> n.toFloat() } + java.lang.Double::class.java -> { n -> n.toDouble() } + else -> null + } +} diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ClientArgTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ClientArgTest.kt new file mode 100644 index 0000000000000..e6b68629212b3 --- /dev/null +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ClientArgTest.kt @@ -0,0 +1,297 @@ +/* + * 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. + */ + +@file:Suppress("PLATFORM_CLASS_MAPPED_TO_KOTLIN") + +package org.apache.airflow.sdk + +import org.apache.airflow.sdk.execution.comm.ConnectionResult +import org.apache.airflow.sdk.execution.comm.StartupDetails +import org.apache.airflow.sdk.execution.comm.TIRunContext +import org.apache.airflow.sdk.execution.comm.VariableResult +import org.apache.airflow.sdk.execution.comm.XComResult +import org.apache.airflow.sdk.internal.ArgValues +import org.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.Assertions.assertFalse +import org.junit.jupiter.api.Assertions.assertNull +import org.junit.jupiter.api.Assertions.assertThrows +import org.junit.jupiter.api.Assertions.assertTrue +import org.junit.jupiter.api.DisplayName +import org.junit.jupiter.api.Test +import org.apache.airflow.sdk.execution.comm.TaskInstance as CommTaskInstance + +/** Records getXCom calls and serves canned values keyed by task id. */ +private class FakeXComTransport( + val xcoms: Map = emptyMap(), +) : org.apache.airflow.sdk.execution.Client { + val pulls = mutableListOf>() + + override fun getConnection(id: String): ConnectionResult = throw NotImplementedError() + + override fun getVariable(key: String): VariableResult = throw NotImplementedError() + + override fun getXCom( + key: String, + dagId: String, + taskId: String, + runId: String, + mapIndex: Int?, + includePriorDates: Boolean, + ): XComResult { + pulls += taskId to mapIndex + return XComResult().also { + it.key = key + it.value = xcoms[taskId] + } + } + + override fun setXCom( + key: String, + value: Any, + dagId: String, + taskId: String, + runId: String, + mapIndex: Int, + ) = throw NotImplementedError() +} + +private fun startupDetails(argBindings: List>?): StartupDetails = + StartupDetails().also { details -> + details.ti = + CommTaskInstance().also { + it.dagId = "d" + it.runId = "r" + it.taskId = "t" + it.tryNumber = 1 + } + details.tiContext = TIRunContext().also { it.argBindings = argBindings } + } + +private fun clientWith( + argBindings: List>?, + xcoms: Map = emptyMap(), +): Pair { + val transport = FakeXComTransport(xcoms) + return Client(startupDetails(argBindings), transport) to transport +} + +private fun taskContext(): Context = + Context( + dagRun = DagRun("d", "r", null, null, null, null, null, emptyMap()), + ti = TaskInstance("d", "r", "t", null, 1), + ) + +internal class ClientArgTest { + @Test + @DisplayName("Should resolve a literal binding to its inline value, by position and by name") + fun shouldResolveLiteralBinding() { + val (client, transport) = clientWith(listOf(mapOf("kind" to "literal", "name" to "x", "value" to 42L))) + + assertTrue(client.hasArgs()) + assertTrue(client.hasArg(0)) + assertTrue(client.hasArg("x")) + assertEquals(42L, client.getArg(0)) + assertEquals(42L, client.getArg("x")) + assertEquals(emptyList>(), transport.pulls) + } + + @Test + @DisplayName("Should resolve an xcom binding by pulling the bound task's return value") + fun shouldResolveXComBinding() { + val (client, transport) = + clientWith( + listOf(mapOf("kind" to "xcom", "name" to "x", "task_id" to "upstream", "map_index" to -1L)), + xcoms = mapOf("upstream" to 7L), + ) + + assertEquals(7L, client.getArg(0)) + assertEquals(listOf("upstream" to null), transport.pulls) + } + + @Test + @DisplayName("Should keep bindings in stub-signature order") + fun shouldKeepBindingOrder() { + val (client, _) = + clientWith( + listOf( + mapOf("kind" to "literal", "name" to "b", "value" to 2L), + mapOf("kind" to "literal", "name" to "a", "value" to 1L), + ), + ) + + assertEquals(2L, client.getArg(0)) + assertEquals(1L, client.getArg(1)) + } + + @Test + @DisplayName("Should pass a non-negative bound map index to the XCom read") + fun shouldPassBoundMapIndex() { + val (client, transport) = + clientWith( + listOf(mapOf("kind" to "xcom", "name" to "x", "task_id" to "upstream", "map_index" to 2L)), + xcoms = mapOf("upstream" to 7L), + ) + + client.getArg(0) + + assertEquals(listOf("upstream" to 2), transport.pulls) + } + + @Test + @DisplayName("Should index into a list XCom when the binding has an element index") + fun shouldResolveElementIndex() { + val (client, _) = + clientWith( + listOf(mapOf("kind" to "xcom", "name" to "x", "task_id" to "upstream", "element_index" to 1L)), + xcoms = mapOf("upstream" to listOf("a", "b", "c")), + ) + + assertEquals("b", client.getArg(0)) + } + + @Test + @DisplayName("Should fail when an element index points into a non-list XCom") + fun shouldRejectElementIndexOnNonList() { + val (client, _) = + clientWith( + listOf(mapOf("kind" to "xcom", "name" to "x", "task_id" to "upstream", "element_index" to 1L)), + xcoms = mapOf("upstream" to "scalar"), + ) + + assertThrows(IllegalStateException::class.java) { client.getArg(0) } + } + + @Test + @DisplayName("Should reject reading an argument that was never bound") + fun shouldRejectUnknownArg() { + val (client, _) = clientWith(listOf(mapOf("kind" to "literal", "name" to "x", "value" to 1L))) + + assertFalse(client.hasArg("y")) + assertFalse(client.hasArg(1)) + assertThrows(IllegalArgumentException::class.java) { client.getArg("y") } + assertThrows(IllegalArgumentException::class.java) { client.getArg(1) } + } + + @Test + @DisplayName("Should report no bound arguments when the supervisor sent none") + fun shouldHandleAbsentBindings() { + val (client, _) = clientWith(null) + + assertFalse(client.hasArgs()) + assertFalse(client.hasArg("x")) + assertFalse(client.hasArg(0)) + } + + @Test + @DisplayName("Should fail on an unsupported binding kind") + fun shouldRejectUnknownBindingKind() { + val (client, _) = clientWith(listOf(mapOf("kind" to "mystery", "name" to "x"))) + + assertThrows(IllegalStateException::class.java) { client.hasArg("x") } + } + + @Test + @DisplayName("Should fail on duplicate binding names") + fun shouldRejectDuplicateBindingNames() { + val (client, _) = + clientWith( + listOf( + mapOf("kind" to "literal", "name" to "x", "value" to 1L), + mapOf("kind" to "literal", "name" to "x", "value" to 2L), + ), + ) + + assertThrows(IllegalStateException::class.java) { client.hasArgs() } + } + + @Test + @DisplayName("Should resolve a flat data parameter from the binding at its position") + fun shouldResolvePositionalBinding() { + val (client, transport) = + clientWith( + listOf( + mapOf("kind" to "literal", "name" to "first", "value" to 5L), + mapOf("kind" to "xcom", "name" to "second", "task_id" to "upstream"), + ), + xcoms = mapOf("upstream" to "pulled"), + ) + + assertEquals(5, ArgValues.requiredInput(taskContext(), client, 0, Integer::class.java, "first").toInt()) + assertEquals("pulled", ArgValues.optionalInput(taskContext(), client, 1, String::class.java)) + assertEquals(listOf("upstream" to null), transport.pulls) + } + + @Test + @DisplayName("Should fail fast when the stub call bound fewer arguments than declared") + fun shouldFailOnArityMismatch() { + val (client, _) = clientWith(listOf(mapOf("kind" to "literal", "name" to "only", "value" to 1L))) + + val error = + assertThrows(IllegalStateException::class.java) { + ArgValues.optionalInput(taskContext(), client, 1, Integer::class.java) + } + + assertEquals( + "Task 't' declares a data parameter at position 1 but the stub call bound only 1 argument(s)", + error.message, + ) + } + + @Test + @DisplayName("Should throw MissingXComException for a required argument bound to a null literal") + fun shouldThrowForNullLiteralOnRequired() { + val (client, _) = clientWith(listOf(mapOf("kind" to "literal", "name" to "value", "value" to null))) + + assertThrows(MissingXComException::class.java) { + ArgValues.requiredInput(taskContext(), client, 0, Integer::class.java, "value") + } + } + + @Test + @DisplayName("Should resolve input-bundle fields by wire name") + fun shouldResolveNamedBindings() { + val (client, _) = + clientWith( + listOf( + mapOf("kind" to "literal", "name" to "region_code", "value" to "emea"), + mapOf("kind" to "xcom", "name" to "threshold", "task_id" to "upstream"), + ), + xcoms = mapOf("upstream" to 0.5), + ) + + assertEquals("emea", ArgValues.optionalNamed(client, "region_code", String::class.java)) + assertEquals(0.5, ArgValues.requiredNamed(client, "threshold", java.lang.Double::class.java, "threshold")) + } + + @Test + @DisplayName("Should resolve an absent named binding to null for optional fields and fail for required ones") + fun shouldHandleAbsentNamedBinding() { + val (client, _) = clientWith(listOf(mapOf("kind" to "literal", "name" to "other", "value" to 1L))) + + assertNull(ArgValues.optionalNamed(client, "missing", String::class.java)) + val error = + assertThrows(IllegalStateException::class.java) { + ArgValues.requiredNamed(client, "missing", Integer::class.java, "field") + } + assertEquals( + "The stub call bound no argument named 'missing', required by input field 'field'", + error.message, + ) + } +}