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/7] 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/7] 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/7] 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/7] 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/7] 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/7] 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/7] 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()); }