From d348817788d66c278b1b4b43bccfa87d86dfc793 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Tue, 4 Aug 2026 13:02:10 +0000 Subject: [PATCH 01/10] 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 02/10] 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 03/10] 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 04/10] 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 05/10] 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 06/10] 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 07/10] Java SDK: Register tasks as first-class TaskDef objects `dag.addTask("extract", Extract.class)` stored tasks as a plain `Map>`, which leaves nowhere to hang anything else a task needs: dependency edges, task-level configuration, and argument wiring all have to attach to a per-task object, and a map of classes cannot carry them. Introducing that object now keeps those follow-ups additive instead of forcing another break of the registration API later. The annotation surface keeps `Builder.Dag` / `Builder.Task`, and the interface users implement keeps the `Task` name, so the definition objects are `DagDef` and `TaskDef` -- a pairing that stays unambiguous next to `Task` at a use site. The SDK is pre-1.0, so the old string-keyed overload is removed outright rather than deprecated. --- .../language-sdks/java.rst | 12 +-- .../airflow/example/ExampleBundleBuilder.java | 2 +- .../example/InterfaceExampleBuilder.java | 11 ++- .../apache/airflow/sdk/BuilderProcessor.kt | 11 +-- .../org/apache/airflow/sdk/BuilderTest.kt | 48 ++++++----- .../airflow/example/ScalaSparkExample.scala | 14 +-- .../kotlin/org/apache/airflow/sdk/Builder.kt | 4 +- .../kotlin/org/apache/airflow/sdk/Bundle.kt | 20 ++--- .../apache/airflow/sdk/{Dag.kt => DagDef.kt} | 60 +++++++++---- .../org/apache/airflow/sdk/execution/Task.kt | 6 +- .../org/apache/airflow/sdk/BundleTest.kt | 4 +- .../org/apache/airflow/sdk/DagDefTest.kt | 86 +++++++++++++++++++ .../apache/airflow/sdk/execution/TaskTest.kt | 6 +- .../airflow/k8sexample/K8sBundleBuilder.java | 2 +- 14 files changed, 202 insertions(+), 84 deletions(-) rename java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/{Dag.kt => DagDef.kt} (63%) create mode 100644 java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/DagDefTest.kt diff --git a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst index c8b73b4626a28..f84aa91094b10 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst @@ -115,7 +115,7 @@ Java entry point public class Main implements BundleBuilder { @Override - public Iterable getDags() { + public Iterable getDags() { return List.of(SalesPipelineBuilder.build()); // SalesPipelineBuilder generated at compile time } @@ -203,7 +203,7 @@ Interface-based API ~~~~~~~~~~~~~~~~~~~ Implement the ``Task`` interface directly for full control over how tasks are registered and how XComs are -read. +read. Each task is registered as a ``TaskDef`` on a ``DagDef``. .. code-block:: java @@ -224,10 +224,10 @@ Register tasks manually in a ``BundleBuilder``: public class MyBundle implements BundleBuilder { @Override - public Iterable getDags() { - var dag = new Dag("my_dag"); - dag.addTask("fetch", FetchTask.class); - dag.addTask("process", ProcessTask.class); + public Iterable getDags() { + var dag = new DagDef("my_dag") + .addTask(new TaskDef("fetch", FetchTask.class)) + .addTask(new TaskDef("process", ProcessTask.class)); return List.of(dag); } } diff --git a/java-sdk/example/src/java/org/apache/airflow/example/ExampleBundleBuilder.java b/java-sdk/example/src/java/org/apache/airflow/example/ExampleBundleBuilder.java index f63d6c4d74337..fa1a860755863 100644 --- a/java-sdk/example/src/java/org/apache/airflow/example/ExampleBundleBuilder.java +++ b/java-sdk/example/src/java/org/apache/airflow/example/ExampleBundleBuilder.java @@ -26,7 +26,7 @@ public class ExampleBundleBuilder implements BundleBuilder { @NotNull @Override - public Iterable getDags() { + public Iterable getDags() { return List.of( InterfaceExampleBuilder.build(), AnnotationExampleBuilder.build(), diff --git a/java-sdk/example/src/java/org/apache/airflow/example/InterfaceExampleBuilder.java b/java-sdk/example/src/java/org/apache/airflow/example/InterfaceExampleBuilder.java index 1c536c3cbf2c7..74eb4f5d7533a 100644 --- a/java-sdk/example/src/java/org/apache/airflow/example/InterfaceExampleBuilder.java +++ b/java-sdk/example/src/java/org/apache/airflow/example/InterfaceExampleBuilder.java @@ -71,11 +71,10 @@ public void execute(@NotNull Context context, Client client) { } } - public static Dag build() { - var dag = new Dag("java_interface_example"); - dag.addTask("extract", Extract.class); - dag.addTask("transform", Transform.class); - dag.addTask("load", Load.class); - return dag; + public static DagDef build() { + return new DagDef("java_interface_example") + .addTask(new TaskDef("extract", Extract.class)) + .addTask(new TaskDef("transform", Transform.class)) + .addTask(new TaskDef("load", Load.class)); } } diff --git a/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt b/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt index 11202ffcdb455..56cbf1e76ad9f 100644 --- a/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt +++ b/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt @@ -54,8 +54,8 @@ import javax.tools.Diagnostic * containing: * * - One inner class per [Builder.Task]-annotated method, implementing [Task]. - * - A static `build()` method that constructs the [Dag] and registers those - * inner classes as tasks. + * - A static `build()` method that constructs the [DagDef] and registers those + * inner classes as [TaskDef]s. * * [Builder.XCom]-annotated parameters are resolved via `client.getXCom` in the * generated `execute` body, with the result cast to the parameter's declared @@ -102,8 +102,8 @@ class BuilderProcessor : AbstractProcessor() { MethodSpec .methodBuilder("build") .addModifiers(Modifier.PUBLIC, Modifier.STATIC) - .returns(ClassName.get(Dag::class.java)) - .addStatement($$"var dag = new $T($S)", ClassName.get(Dag::class.java), ann.id.ifBlank { el.simpleName }) + .returns(ClassName.get(DagDef::class.java)) + .addStatement($$"var dag = new $T($S)", ClassName.get(DagDef::class.java), ann.id.ifBlank { el.simpleName }) for (inner in el.enclosedElements) { if (inner !is ExecutableElement) continue @@ -116,7 +116,8 @@ class BuilderProcessor : AbstractProcessor() { builderClass.addType(task.spec) buildMethod.addStatement( - $$"dag.addTask($S, $L.class)", + $$"dag.addTask(new $T($S, $L.class))", + ClassName.get(TaskDef::class.java), ann.id.ifBlank { inner.simpleName }, innerName, ) diff --git a/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt b/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt index 3e28e1b009afd..6a08979c56163 100644 --- a/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt +++ b/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt @@ -83,16 +83,17 @@ class BuilderTest { import java.util.Optional; import org.apache.airflow.sdk.Client; import org.apache.airflow.sdk.Context; - import org.apache.airflow.sdk.Dag; + import org.apache.airflow.sdk.DagDef; import org.apache.airflow.sdk.MissingXComException; import org.apache.airflow.sdk.Task; + import org.apache.airflow.sdk.TaskDef; public final class TestExampleBuilder { - public static Dag build() { - var dag = new Dag("TestExample"); - dag.addTask("t1", T1.class); - dag.addTask("t2", T2.class); - dag.addTask("t3", T3.class); + public static DagDef build() { + var dag = new DagDef("TestExample"); + dag.addTask(new TaskDef("t1", T1.class)); + dag.addTask(new TaskDef("t2", T2.class)); + dag.addTask(new TaskDef("t3", T3.class)); return dag; } public static final class T1 implements Task { @@ -157,14 +158,15 @@ class BuilderTest { import java.util.Optional; import org.apache.airflow.sdk.Client; import org.apache.airflow.sdk.Context; - import org.apache.airflow.sdk.Dag; + import org.apache.airflow.sdk.DagDef; import org.apache.airflow.sdk.MissingXComException; import org.apache.airflow.sdk.Task; + import org.apache.airflow.sdk.TaskDef; public final class TestExampleBuilder { - public static Dag build() { - var dag = new Dag("TestExample"); - dag.addTask("t", T.class); + public static DagDef build() { + var dag = new DagDef("TestExample"); + dag.addTask(new TaskDef("t", T.class)); return dag; } public static final class T implements Task { @@ -220,14 +222,15 @@ class BuilderTest { import java.util.Optional; import org.apache.airflow.sdk.Client; import org.apache.airflow.sdk.Context; - import org.apache.airflow.sdk.Dag; + import org.apache.airflow.sdk.DagDef; import org.apache.airflow.sdk.MissingXComException; import org.apache.airflow.sdk.Task; + import org.apache.airflow.sdk.TaskDef; public final class TestExampleBuilder { - public static Dag build() { - var dag = new Dag("TestExample"); - dag.addTask("t", T.class); + public static DagDef build() { + var dag = new DagDef("TestExample"); + dag.addTask(new TaskDef("t", T.class)); return dag; } public static final class T implements Task { @@ -261,8 +264,8 @@ class BuilderTest { "org.apache.airflow.example.TestExampleBuilder", """ package org.apache.airflow.example; - import org.apache.airflow.sdk.Dag; - public final class TestExampleBuilder { public static Dag build() { var dag = new Dag("foo"); return dag; } } + import org.apache.airflow.sdk.DagDef; + public final class TestExampleBuilder { public static DagDef build() { var dag = new DagDef("foo"); return dag; } } """, ) } @@ -284,8 +287,8 @@ class BuilderTest { "org.apache.airflow.example.Foo", """ package org.apache.airflow.example; - import org.apache.airflow.sdk.Dag; - public final class Foo { public static Dag build() { var dag = new Dag("TestExample"); return dag; } } + import org.apache.airflow.sdk.DagDef; + public final class Foo { public static DagDef build() { var dag = new DagDef("TestExample"); return dag; } } """, ) } @@ -313,12 +316,13 @@ class BuilderTest { import java.lang.Override; import org.apache.airflow.sdk.Client; import org.apache.airflow.sdk.Context; - import org.apache.airflow.sdk.Dag; + import org.apache.airflow.sdk.DagDef; import org.apache.airflow.sdk.Task; + import org.apache.airflow.sdk.TaskDef; public final class TestExampleBuilder { - public static Dag build() { - var dag = new Dag("TestExample"); - dag.addTask("foo", T1.class); + public static DagDef build() { + var dag = new DagDef("TestExample"); + dag.addTask(new TaskDef("foo", T1.class)); return dag; } public static final class T1 implements Task { diff --git a/java-sdk/scala_spark_example/src/main/scala/org/apache/airflow/example/ScalaSparkExample.scala b/java-sdk/scala_spark_example/src/main/scala/org/apache/airflow/example/ScalaSparkExample.scala index ae779bf02135a..9ecebafb0cfe0 100644 --- a/java-sdk/scala_spark_example/src/main/scala/org/apache/airflow/example/ScalaSparkExample.scala +++ b/java-sdk/scala_spark_example/src/main/scala/org/apache/airflow/example/ScalaSparkExample.scala @@ -19,7 +19,7 @@ package org.apache.airflow.example -import org.apache.airflow.sdk.{Bundle, BundleBuilder, Client, Context, Dag, Server, Task} +import org.apache.airflow.sdk.{Bundle, BundleBuilder, Client, Context, DagDef, Server, Task, TaskDef} import org.apache.logging.log4j.{LogManager, Logger} import org.apache.spark.sql.SparkSession import org.apache.spark.sql.functions.sum @@ -133,16 +133,16 @@ class SparkLoad extends Task { } object ScalaSparkExample { - def build(): Dag = - new Dag(SparkEtl.DagId) - .addTask(SparkEtl.ExtractTaskId, classOf[SparkExtract]) - .addTask(SparkEtl.TransformTaskId, classOf[SparkTransform]) - .addTask(SparkEtl.LoadTaskId, classOf[SparkLoad]) + def build(): DagDef = + new DagDef(SparkEtl.DagId) + .addTask(new TaskDef(SparkEtl.ExtractTaskId, classOf[SparkExtract])) + .addTask(new TaskDef(SparkEtl.TransformTaskId, classOf[SparkTransform])) + .addTask(new TaskDef(SparkEtl.LoadTaskId, classOf[SparkLoad])) } /** Bundle entry point served to Airflow's Java coordinator. */ object ScalaSparkBundleBuilder extends BundleBuilder { - override def getDags(): java.lang.Iterable[Dag] = java.util.List.of(ScalaSparkExample.build()) + override def getDags(): java.lang.Iterable[DagDef] = java.util.List.of(ScalaSparkExample.build()) def main(args: Array[String]): Unit = Server.create(args).serve(new Bundle(getDags())) diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Builder.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Builder.kt index 3a5b84d2daf84..9dbfbfbefc390 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Builder.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Builder.kt @@ -24,7 +24,7 @@ package org.apache.airflow.sdk * * This class is not instantiated directly. Its nested annotations drive the * `BuilderProcessor` annotation processor in the :processor project, - * which generates a `*Builder` class for each class annotated with [Dag]. + * which generates a `*Builder` class for each class annotated with [Builder.Dag]. * * Example: * @@ -41,7 +41,7 @@ package org.apache.airflow.sdk * ``` * * The processor generates `MyPipelineBuilder.build()`, which returns a - * fully wired-up [Dag] ready to add to a [Bundle]. + * fully wired-up [DagDef] ready to add to a [Bundle]. */ class Builder internal constructor() { /** diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Bundle.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Bundle.kt index 677ec48eb9311..6cd549270a8a4 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Bundle.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Bundle.kt @@ -20,22 +20,22 @@ package org.apache.airflow.sdk /** - * An immutable snapshot of all [Dag]s that this JVM process can execute. + * An immutable snapshot of all [DagDef]s that this JVM process can execute. * * Build a [Bundle] by implementing [BundleBuilder], then pass it to * [Server.serve] to start accepting task-execution requests. * - * @property dags All registered Dags keyed by [Dag.id]. + * @property dags All registered Dags keyed by [DagDef.id]. * @throws IllegalArgumentException if any two Dags share the same ID. */ class Bundle( - dags: Iterable, + dags: Iterable, ) { - internal val dags: Map = dags.associateByDagId() + internal val dags: Map = dags.associateByDagId() } -private fun Iterable.associateByDagId(): Map { - val dagMap = linkedMapOf() +private fun Iterable.associateByDagId(): Map { + val dagMap = linkedMapOf() for (dag in this) { require(dagMap.putIfAbsent(dag.id, dag) == null) { "Dags in bundle have duplicate ID: ${dag.id}" @@ -45,14 +45,14 @@ private fun Iterable.associateByDagId(): Map { } /** - * Entry point for declaring the [Dag]s that this bundle contains. + * Entry point for declaring the [DagDef]s that this bundle contains. * * Implement this interface to create a Dag bundle to be served by [Server]. * * ```java * public class MyBundleBuilder implements BundleBuilder { * @Override - * public Iterable getDags() { + * public Iterable getDags() { * return List.of(MyDagBuilder.build()); * } * @@ -64,14 +64,14 @@ private fun Iterable.associateByDagId(): Map { */ interface BundleBuilder { /** - * Returns all [Dag]s that belong to this bundle. + * Returns all [DagDef]s that belong to this bundle. * * Called once during [build]; Dag IDs must be unique across the returned * collection. * * @throws IllegalArgumentException if any two Dags share the same ID. */ - fun getDags(): Iterable + fun getDags(): Iterable /** * Constructs a [Bundle] from the Dags returned by [getDags]. diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Dag.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/DagDef.kt similarity index 63% rename from java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Dag.kt rename to java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/DagDef.kt index c998580374169..a6bd50e949a1c 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Dag.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/DagDef.kt @@ -24,46 +24,72 @@ import kotlin.Throws /** * A collection of tasks with directional dependencies. * - * Create a [Dag] directly and register tasks with [addTask]. + * Create a [DagDef] directly and register [TaskDef]s with [addTask]. * * The [Builder.Dag] annotation should generally be preferred in user code, * where the annotation processor generates the wiring for you. Only use this - * class directly if you need to do low-level plumbing. + * class directly if you need to do low-level plumbing: + * + * ```java + * var dag = new DagDef("java_etl") + * .addTask(new TaskDef("extract", Extract.class)) + * .addTask(new TaskDef("load", Load.class)); + * ``` * * @param id Dag identifier. Must contain only ASCII alphanumeric characters, * dashes, dots, or underscores; must be unique within a [Bundle]. * * @see Builder.Dag */ -class Dag( +class DagDef( val id: String, // TODO: charset check? ) { - internal var tasks = mutableMapOf>() + internal val tasks = linkedMapOf() /** * Registers a task with this Dag. * - * The class must have a public no-argument constructor and implement [Task]. - * Task IDs must be unique within a Dag. + * A [TaskDef] belongs to at most one [DagDef]; registering the same instance + * with a second Dag, or twice with the same one, fails. Task IDs must be + * unique within a Dag. * - * @param id Task identifier, unique within this Dag. - * @param definition Class that implements [Task]. Must have a public no-arg - * constructor. + * @param task Task definition to register. * @return This Dag, for chaining. - * @throws IllegalArgumentException if a task already exists in the Dag with - * the same ID. + * @throws IllegalArgumentException if the task already belongs to a Dag or a + * task with the same ID is already registered. */ - fun addTask( - id: String, - definition: Class, - ): Dag { - require(tasks.putIfAbsent(id, definition) == null) { - "Tasks in Dag have duplicate ID: $id" + fun addTask(task: TaskDef): DagDef { + require(task.owner == null) { + "Task '${task.id}' already belongs to Dag '${task.owner?.id}'" + } + require(tasks.putIfAbsent(task.id, task) == null) { + "Tasks in Dag have duplicate ID: ${task.id}" } + task.owner = this return this } } +/** + * One task definition: its ID and the class that implements it. + * + * ```java + * var extract = new TaskDef("extract", Extract.class); + * ``` + * + * @param id Task identifier, unique within a [DagDef]. + * @param definition Class that implements [Task]. Must have a public no-arg + * constructor. + * + * @see Builder.Task + */ +class TaskDef( + val id: String, + val definition: Class, +) { + internal var owner: DagDef? = null +} + /** * A single unit of work executed by Airflow. * diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Task.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Task.kt index 60258dacc628d..1f1a01e674a58 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Task.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Task.kt @@ -70,9 +70,11 @@ internal object TaskRunner { request: StartupDetails, client: Client, ): Any { - val task = bundle.dags[request.ti.dagId]?.tasks[request.ti.taskId] ?: return TaskResult.of(TaskState.State.REMOVED) + val definition = + bundle.dags[request.ti.dagId]?.tasks[request.ti.taskId]?.definition + ?: return TaskResult.of(TaskState.State.REMOVED) return try { - task.getDeclaredConstructor().newInstance().execute(Context.from(request), client) + definition.getDeclaredConstructor().newInstance().execute(Context.from(request), client) TaskResult.success() } catch (e: Throwable) { logger.error("Error executing task", mapOf("ti" to request.ti, "error" to e, "trace" to e.stackTraceToString())) diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/BundleTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/BundleTest.kt index 0e4afb1894a7f..57050754e887e 100644 --- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/BundleTest.kt +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/BundleTest.kt @@ -27,7 +27,7 @@ internal class BundleTest { @Test @DisplayName("Should index dags by dagId") fun shouldIndexDagsByDagId() { - val dag = Dag("dag") + val dag = DagDef("dag") val bundle = Bundle(listOf(dag)) @@ -39,7 +39,7 @@ internal class BundleTest { fun shouldRejectDuplicateDagIds() { val error = Assertions.assertThrows(IllegalArgumentException::class.java) { - Bundle(listOf(Dag("dag"), Dag("dag"))) + Bundle(listOf(DagDef("dag"), DagDef("dag"))) } Assertions.assertEquals("Dags in bundle have duplicate ID: dag", error.message) diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/DagDefTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/DagDefTest.kt new file mode 100644 index 0000000000000..f22c26432c30c --- /dev/null +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/DagDefTest.kt @@ -0,0 +1,86 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.airflow.sdk + +import org.junit.jupiter.api.Assertions +import org.junit.jupiter.api.DisplayName +import org.junit.jupiter.api.Test + +internal class DagDefTest { + private class NoOp : Task { + override fun execute( + context: Context, + client: Client, + ) = Unit + } + + @Test + @DisplayName("Should index tasks by taskId in registration order") + fun shouldIndexTasksByTaskId() { + val extract = TaskDef("extract", NoOp::class.java) + val load = TaskDef("load", NoOp::class.java) + + val dag = DagDef("dag").addTask(extract).addTask(load) + + Assertions.assertEquals(listOf("extract", "load"), dag.tasks.keys.toList()) + Assertions.assertEquals(mapOf("extract" to extract, "load" to load), dag.tasks) + } + + @Test + @DisplayName("Should reject duplicate task ids") + fun shouldRejectDuplicateTaskIds() { + val dag = DagDef("dag").addTask(TaskDef("extract", NoOp::class.java)) + + val error = + Assertions.assertThrows(IllegalArgumentException::class.java) { + dag.addTask(TaskDef("extract", NoOp::class.java)) + } + + Assertions.assertEquals("Tasks in Dag have duplicate ID: extract", error.message) + } + + @Test + @DisplayName("Should reject a task already registered with another dag") + fun shouldRejectTaskOwnedByAnotherDag() { + val extract = TaskDef("extract", NoOp::class.java) + DagDef("first").addTask(extract) + + val error = + Assertions.assertThrows(IllegalArgumentException::class.java) { + DagDef("second").addTask(extract) + } + + Assertions.assertEquals("Task 'extract' already belongs to Dag 'first'", error.message) + } + + @Test + @DisplayName("Should reject the same task registered twice with one dag") + fun shouldRejectTaskRegisteredTwice() { + val extract = TaskDef("extract", NoOp::class.java) + val dag = DagDef("dag").addTask(extract) + + val error = + Assertions.assertThrows(IllegalArgumentException::class.java) { + dag.addTask(extract) + } + + Assertions.assertEquals("Task 'extract' already belongs to Dag 'dag'", error.message) + } +} diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskTest.kt index 42a083b94aacb..5d303a3dd814a 100644 --- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskTest.kt +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskTest.kt @@ -22,8 +22,9 @@ package org.apache.airflow.sdk.execution import org.apache.airflow.sdk.Bundle import org.apache.airflow.sdk.Client import org.apache.airflow.sdk.Context -import org.apache.airflow.sdk.Dag +import org.apache.airflow.sdk.DagDef import org.apache.airflow.sdk.Task +import org.apache.airflow.sdk.TaskDef import org.apache.airflow.sdk.execution.comm.BundleInfo import org.apache.airflow.sdk.execution.comm.DagRun import org.apache.airflow.sdk.execution.comm.RetryTask @@ -88,8 +89,7 @@ class TaskTest { taskId: String, taskClass: Class, ): Bundle { - val dag = Dag("test_dag") - dag.addTask(taskId, taskClass) + val dag = DagDef("test_dag").addTask(TaskDef(taskId, taskClass)) return Bundle(listOf(dag)) } diff --git a/kubernetes-tests/lang_sdk/java_example/src/java/org/apache/airflow/k8sexample/K8sBundleBuilder.java b/kubernetes-tests/lang_sdk/java_example/src/java/org/apache/airflow/k8sexample/K8sBundleBuilder.java index 8333567da83ed..c3b511b898ced 100644 --- a/kubernetes-tests/lang_sdk/java_example/src/java/org/apache/airflow/k8sexample/K8sBundleBuilder.java +++ b/kubernetes-tests/lang_sdk/java_example/src/java/org/apache/airflow/k8sexample/K8sBundleBuilder.java @@ -28,7 +28,7 @@ public class K8sBundleBuilder implements BundleBuilder { @NotNull @Override - public Iterable getDags() { + public Iterable getDags() { return List.of(CombinedExampleBuilder.build()); } From 9e7cab73e8f0a6acb2fa5ce95ac0ab46bdb7b9cc Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 5 Aug 2026 08:28:37 +0000 Subject: [PATCH 08/10] Java SDK: Honor TaskFlow arg bindings sent by the supervisor For a stub-backed Dag the Python file's `@task.stub` call site is the graph the scheduler actually orders the run by, so it must also be what feeds the Java task its inputs. The Java side previously re-declared that data flow with `@Builder.XCom(task = "...")`, duplicating the Dag file's wiring in a second place that nothing keeps honest: rename or re-wire a task in Python and the Java annotation silently keeps pulling the old upstream. The 2026-10-30 supervisor schema delivers the call site's bindings with every task run, so the runtime can read them instead of guessing. Binding is positional, matching the Go SDK's flat-parameter contract: Java parameter names are not API, so an IDE rename must never rebind an input. Keyword-style calls bind by name only through an explicit `TaskInput` bundle whose public fields declare their wire names -- the deliberate, tagged boundary for snake_case-to-camelCase crossings. A task declares flat data parameters or one bundle, never both, so field names and positions cannot shift each other. jsonSchema2Pojo cannot express the kind-discriminated binding union, so the generated `TIRunContext` carries the raw payload and a small hand-written decoder materializes the typed view. --- .../language-sdks/java.rst | 131 +++++++- .../airflow/example/AnnotationExample.java | 25 +- .../example/InterfaceExampleBuilder.java | 11 +- .../airflow/example/XComCastingExample.java | 11 +- .../src/resources/dags/java_examples.py | 36 ++- java-sdk/gradle.properties | 2 +- .../apache/airflow/sdk/BuilderProcessor.kt | 268 ++++++++++------ .../org/apache/airflow/sdk/BuilderTest.kt | 264 ++++++++++++---- java-sdk/sdk/schema/schema.json | 259 ++++++++++++++- .../kotlin/org/apache/airflow/sdk/ArgName.kt | 40 +++ .../kotlin/org/apache/airflow/sdk/Builder.kt | 19 +- .../kotlin/org/apache/airflow/sdk/Client.kt | 94 ++++++ .../org/apache/airflow/sdk/TaskInput.kt | 45 +++ .../airflow/sdk/execution/ArgBinding.kt | 77 +++++ .../apache/airflow/sdk/internal/ArgValues.kt | 175 +++++++++++ .../org/apache/airflow/sdk/ClientArgTest.kt | 297 ++++++++++++++++++ 16 files changed, 1550 insertions(+), 204 deletions(-) create mode 100644 java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/ArgName.kt create mode 100644 java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/TaskInput.kt create mode 100644 java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/ArgBinding.kt create mode 100644 java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/ArgValues.kt create mode 100644 java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ClientArgTest.kt diff --git a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst index f84aa91094b10..b3d1f2792fc78 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst @@ -93,10 +93,7 @@ Java implementation } @Builder.Task(id = "transform") - public long transform( - Client client, - @Builder.XCom(task = "extract") long recordCount - ) { + public long transform(Client client, long recordCount) { var threshold = (String) client.getVariable("transform_threshold"); // ... process data ... return transformedCount; @@ -105,8 +102,11 @@ Java implementation .. note:: - See how both ``transform`` in Python and Java need to have an argument to accept upstream XCom. The - Python one is needed to declare dependency, and the Java one is needed to actually retrieve the value. + The graph is declared once, in the Python Dag file: ``transform(extract())`` feeds the upstream's + return value into the downstream's parameter by calling tasks like functions. The supervisor sends + the resulting *argument bindings* to the Java runtime, and each Java data parameter receives + whatever the Python call site bound at its position — an upstream task's XCom or an inline + literal. See :ref:`java-sdk/arg-binding`. Java entry point ~~~~~~~~~~~~~~~~ @@ -164,12 +164,17 @@ Annotate a plain Java class and let the SDK generate the boilerplate at compile * - ``@Builder.Task(id = "...")`` - Marks a method as a task implementation. The ``id`` must match the ``@task.stub`` function name in the Python Dag. If ``id`` is omitted the method name is used. - * - ``@Builder.XCom(task = "...")`` - - Injects the ``return_value`` XCom from the named upstream task as a method parameter. - The parameter type must be compatible with the stored value (see :ref:`java-sdk/types`). + * - ``TaskInput`` / ``@ArgName("...")`` + - Marks a class as a task's input bundle, so keyword arguments bind by name instead of by + position: each public field receives the binding whose name matches it (the ``@ArgName`` + value, or the verbatim field name). See :ref:`java-sdk/arg-binding`. + +Besides the annotations, a task method may declare a ``Client`` and a ``Context`` parameter in any +position; the SDK injects both. Every other parameter is a *data parameter* and receives an +argument bound by the Python ``@task.stub`` call site. The annotation processor generates a ``Builder`` class that wires up the task -registry and handles XCom injection automatically. +registry and resolves data parameters and XCom pushes automatically. .. code-block:: java @@ -184,10 +189,7 @@ registry and handles XCom injection automatically. } @Builder.Task(id = "process") - public long process( - Client client, - @Builder.XCom(task = "fetch") String fetched - ) { + public long process(Client client, String fetched) { var threshold = (String) client.getVariable("process_threshold"); // implement task logic return count; @@ -234,6 +236,105 @@ Register tasks manually in a ``BundleBuilder``: See the `Java SDK API Reference `__ for more details. +.. _java-sdk/arg-binding: + +Binding stub arguments +~~~~~~~~~~~~~~~~~~~~~~ + +Calling a ``@task.stub`` TaskFlow-style in the Python Dag is what declares the graph, and the +supervisor delivers the resulting argument bindings to the Java runtime with every task run. A +binding carries either an upstream task's ``return_value`` XCom or an inline literal written at the +call site. + +Positional binding +^^^^^^^^^^^^^^^^^^ + +A task method's data parameters bind **by position**, in declaration order — the injected ``Client`` +and ``Context`` parameters do not take up a position. Java parameter names are not part of the API, +so renaming one in an IDE never rebinds an input. + +.. code-block:: python + + @task.stub(queue="java") + def score(rows, threshold): ... + + + score(load_rows(), 0.75) + +.. code-block:: java + + @Builder.Task(id = "score") + public long score(Client client, long rows, double threshold) { + // rows <- the load_rows XCom (position 0) + // threshold <- the literal 0.75 (position 1) + } + +A primitive parameter cannot hold ``null``, so the task fails with ``MissingXComException`` when its +binding resolves to nothing; declare a boxed type (``Long``, ``Double``, …) to receive ``null`` +instead. Declaring more data parameters than the call site bound also fails the task, rather than +running it with missing inputs. + +Named binding with a ``TaskInput`` bundle +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +To bind keyword arguments by name, declare a single parameter whose class implements ``TaskInput``. +Each public non-final field receives the binding named by its ``@ArgName`` value, or by the verbatim +field name — the deliberate, tagged boundary where the stub's ``snake_case`` argument names cross +into ``camelCase`` Java fields. The class needs a public no-argument constructor. + +.. code-block:: python + + @task.stub(queue="java") + def score(region_code, threshold): ... + + + score(region_code="emea", threshold=load_threshold()) + +.. code-block:: java + + public static class ScoreInput implements TaskInput { + @ArgName("region_code") + public String region; + + public double threshold; + } + + @Builder.Task(id = "score") + public long score(Client client, ScoreInput input) { ... } + +A task declares flat data parameters **or** one ``TaskInput`` bundle, never both, so field names and +flat positions cannot shift each other. Mixing them, or declaring two bundles, fails the build. + +Reading bindings from the interface API +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Tasks written against the ``Task`` interface read the same bindings imperatively: + +.. list-table:: + :header-rows: 1 + :widths: 40 60 + + * - ``Client`` method + - Returns + * - ``hasArgs()`` + - Whether the Python Dag called this stub with any TaskFlow arguments at all. + * - ``hasArg(int position)`` / ``hasArg(String name)`` + - Whether an argument was bound at that position, or with that name. + * - ``getArg(int position)`` / ``getArg(String name)`` + - The bound value — the inline literal, or the bound upstream's XCom. Throws + ``IllegalArgumentException`` when nothing was bound there; probe with ``hasArg`` first. + +.. code-block:: java + + public class ScoreTask implements Task { + @Override + public void execute(Context context, Client client) throws Exception { + var rows = client.getArg(0); + var threshold = client.hasArg("threshold") ? client.getArg("threshold") : 0.5; + // implement task logic + } + } + .. _java-sdk/logging: Logging @@ -428,7 +529,7 @@ represented as Java objects when read back via ``getXCom``. .. note:: - An ``@Builder.XCom`` parameter that reads a value which was never pushed resolves to + A data parameter whose binding resolves to a value that was never pushed receives ``null``. A boxed parameter (``Integer``, ``Long``, ``Boolean``, …) receives ``null`` safely, but a primitive parameter (``int``, ``long``, ``boolean``, …) cannot represent ``null`` and the task fails with ``MissingXComException``. Declare the parameter with a diff --git a/java-sdk/example/src/java/org/apache/airflow/example/AnnotationExample.java b/java-sdk/example/src/java/org/apache/airflow/example/AnnotationExample.java index bb715a73cb502..42ee68a745ce4 100644 --- a/java-sdk/example/src/java/org/apache/airflow/example/AnnotationExample.java +++ b/java-sdk/example/src/java/org/apache/airflow/example/AnnotationExample.java @@ -27,6 +27,9 @@ import java.util.concurrent.Executors; import org.apache.airflow.sdk.*; +// The Python Dag file (src/resources/dags/java_examples.py) owns the graph: each +// data parameter below receives whatever the @task.stub call site bound at its +// position, so this class registers task implementations only. @SuppressWarnings("DuplicatedCode") @Builder.Dag(id = "java_annotation_example") public class AnnotationExample { @@ -52,7 +55,7 @@ public long extractValue(Client client) throws InterruptedException { } @Builder.Task(id = "transform") - public long transformValue(Client client, @Builder.XCom(task = "extract") long extracted) { + public long transformValue(Client client, long extracted) { log.log(INFO, "Got XCom from extract: {0}", extracted); var variable = client.getVariable("my_variable"); @@ -68,7 +71,7 @@ public long transformValue(Client client, @Builder.XCom(task = "extract") long e // RetryTask (instead of a terminal FAILED) when ti_context.should_retry is // set. The retry then runs this task again and it returns normally. @Builder.Task - public void load(Context context, @Builder.XCom(task = "transform") long transformed) { + public void load(Context context, long transformed) { log.log(INFO, "Got XCom from transform: {0}", transformed); if (context.ti.tryNumber == 1) { throw new RuntimeException("I failed"); @@ -76,6 +79,24 @@ public void load(Context context, @Builder.XCom(task = "transform") long transfo log.log(INFO, "Recovered on retry, try number {0}", context.ti.tryNumber); } + // Keyword arguments bind by name instead of by position, through a TaskInput + // bundle: the tagged boundary where the stub's snake_case argument names + // cross into camelCase Java fields. + public static class ReportInput implements TaskInput { + @ArgName("run_label") + public String runLabel; + + public long transformed; + } + + @Builder.Task(id = "report") + public void report(ReportInput input) { + log.log(INFO, "Report {0} for transformed value {1}", input.runLabel, input.transformed); + if (!"nightly".equals(input.runLabel)) { + throw new RuntimeException("expected run label 'nightly' but got " + input.runLabel); + } + } + // Verify one supervisor channel can handle client calls across threads. @Builder.Task(id = "concurrent") public void concurrentClientCalls(Client client) throws Exception { diff --git a/java-sdk/example/src/java/org/apache/airflow/example/InterfaceExampleBuilder.java b/java-sdk/example/src/java/org/apache/airflow/example/InterfaceExampleBuilder.java index 74eb4f5d7533a..1ff317959d43f 100644 --- a/java-sdk/example/src/java/org/apache/airflow/example/InterfaceExampleBuilder.java +++ b/java-sdk/example/src/java/org/apache/airflow/example/InterfaceExampleBuilder.java @@ -52,8 +52,10 @@ public void execute(@NotNull Context context, Client client) throws Exception { public static class Transform implements Task { public void execute(@NotNull Context context, Client client) { - var extracted = client.getXCom("extract"); - log.log(INFO, "Got XCom from extract: {0}", extracted); + // The Python Dag file calls transform(extracted): an interface-API task + // reads the same binding imperatively, by position or by argument name. + var extracted = client.getArg(0); + log.log(INFO, "Got extracted value from the bound argument: {0}", extracted); var variable = client.getVariable("my_variable"); log.log(INFO, "Got variable: {0}", variable); @@ -65,7 +67,10 @@ public void execute(@NotNull Context context, Client client) { public static class Load implements Task { public void execute(@NotNull Context context, Client client) { - var transformed = client.getXCom("transform"); + // hasArg probes whether the Dag file bound an argument at all, so a task + // can still fall back to reading an upstream XCom by task id. + var transformed = + client.hasArg("transformed") ? client.getArg("transformed") : client.getXCom("transform"); log.log(INFO, "Got XCom from transform: {0}", transformed); throw new RuntimeException("I failed"); } diff --git a/java-sdk/example/src/java/org/apache/airflow/example/XComCastingExample.java b/java-sdk/example/src/java/org/apache/airflow/example/XComCastingExample.java index c4945af865df1..092d2e168f9e8 100644 --- a/java-sdk/example/src/java/org/apache/airflow/example/XComCastingExample.java +++ b/java-sdk/example/src/java/org/apache/airflow/example/XComCastingExample.java @@ -23,6 +23,9 @@ import org.apache.airflow.sdk.*; +// Stub-backed tasks wired by the Python Dag file: each parameter receives the +// value the stub call bound at its position, widening or narrowing to the +// declared type at run time. @Builder.Dag(id = "java_xcom_casting_example") public class XComCastingExample { private static final System.Logger log = System.getLogger(XComCastingExample.class.getName()); @@ -35,13 +38,13 @@ public int produceNumber() { // Any primitive numeric type (byte, short, int, long, float, double) and its boxed form works the same way. @Builder.Task(id = "widen_to_long") - public long widenToLong(@Builder.XCom(task = "produce_number") long value) { + public long widenToLong(long value) { log.log(INFO, "Got long {0}", value); return value + 1; } @Builder.Task(id = "widen_to_double") - public void widenToDouble(@Builder.XCom(task = "widen_to_long") double value) { + public void widenToDouble(double value) { log.log(INFO, "Got double {0}", value); if (value != 8.0) { throw new RuntimeException("expected 8.0 but got " + value); @@ -54,7 +57,7 @@ public void produceNothing() { } @Builder.Task(id = "consume_nullable") - public void consumeNullable(@Builder.XCom(task = "produce_nothing") Integer value) { + public void consumeNullable(Integer value) { log.log(INFO, "Got nullable int {0}", value); if (value != null) { throw new RuntimeException("expected null but got " + value); @@ -68,7 +71,7 @@ public double produceFraction() { } @Builder.Task(id = "consume_float") - public void consumeFloat(@Builder.XCom(task = "produce_fraction") float value) { + public void consumeFloat(float value) { log.log(INFO, "Got float {0}", value); if (value != 1.5f) { throw new RuntimeException("expected 1.5 but got " + value); diff --git a/java-sdk/example/src/resources/dags/java_examples.py b/java-sdk/example/src/resources/dags/java_examples.py index 5426911b0faee..49b79b892ad65 100644 --- a/java-sdk/example/src/resources/dags/java_examples.py +++ b/java-sdk/example/src/resources/dags/java_examples.py @@ -34,27 +34,32 @@ def extract(): ... @task.stub(queue="java") -def transform(): ... +def transform(extracted): ... @task.stub(queue="java", retries=1, retry_delay=timedelta(seconds=5)) -def load(): ... +def load(transformed): ... @task.stub(queue="java") def concurrent(): ... +# Keyword arguments bind to the public fields of the Java task's TaskInput bundle. +@task.stub(queue="java") +def report(run_label, transformed): ... + + @task.stub(queue="java") def produce_number(): ... @task.stub(queue="java") -def widen_to_long(): ... +def widen_to_long(value): ... @task.stub(queue="java") -def widen_to_double(): ... +def widen_to_double(value): ... @task.stub(queue="java") @@ -62,7 +67,7 @@ def produce_nothing(): ... @task.stub(queue="java") -def consume_nullable(): ... +def consume_nullable(value): ... @task.stub(queue="java") @@ -70,7 +75,7 @@ def produce_fraction(): ... @task.stub(queue="java") -def consume_float(): ... +def consume_float(value): ... @task() @@ -82,25 +87,28 @@ def python_task_2(transformed): @dag(dag_id="java_interface_example") def java_interface_example(): - transformed = transform() - python_task_1() >> extract() >> transformed + extracted = extract() + python_task_1() >> extracted + transformed = transform(extracted) python_task_2(transformed) @dag(dag_id="java_annotation_example") def java_annotation_example(): - transformed = transform() - python_task_1() >> extract() >> transformed + extracted = extract() + python_task_1() >> extracted + transformed = transform(extracted) python_task_2(transformed) - transformed >> load() + load(transformed) + report(run_label="nightly", transformed=transformed) concurrent() @dag(dag_id="java_xcom_casting_example") def java_xcom_casting_example(): - produce_number() >> widen_to_long() >> widen_to_double() - produce_nothing() >> consume_nullable() - produce_fraction() >> consume_float() + widen_to_double(widen_to_long(produce_number())) + consume_nullable(produce_nothing()) + consume_float(produce_fraction()) java_interface_example() diff --git a/java-sdk/gradle.properties b/java-sdk/gradle.properties index 9438ba6435532..477b31da386c4 100644 --- a/java-sdk/gradle.properties +++ b/java-sdk/gradle.properties @@ -17,7 +17,7 @@ org.gradle.configuration-cache=true -airflowSupervisorSchemaVersion=2026-06-16 +airflowSupervisorSchemaVersion=2026-10-30 projectVersion=1.0.0-SNAPSHOT diff --git a/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt b/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt index 56cbf1e76ad9f..3bd8eb336b95d 100644 --- a/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt +++ b/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt @@ -25,18 +25,21 @@ import com.squareup.javapoet.ClassName import com.squareup.javapoet.CodeBlock import com.squareup.javapoet.JavaFile import com.squareup.javapoet.MethodSpec +import com.squareup.javapoet.ParameterizedTypeName import com.squareup.javapoet.TypeName import com.squareup.javapoet.TypeSpec -import java.util.Optional +import org.apache.airflow.sdk.internal.ArgValues import javax.annotation.processing.AbstractProcessor import javax.annotation.processing.ProcessingEnvironment import javax.annotation.processing.RoundEnvironment import javax.annotation.processing.SupportedAnnotationTypes import javax.annotation.processing.SupportedSourceVersion import javax.lang.model.SourceVersion +import javax.lang.model.element.ElementKind import javax.lang.model.element.ExecutableElement import javax.lang.model.element.Modifier import javax.lang.model.element.TypeElement +import javax.lang.model.element.VariableElement import javax.lang.model.type.TypeKind import javax.lang.model.type.TypeMirror import javax.tools.Diagnostic @@ -57,9 +60,11 @@ import javax.tools.Diagnostic * - A static `build()` method that constructs the [DagDef] and registers those * inner classes as [TaskDef]s. * - * [Builder.XCom]-annotated parameters are resolved via `client.getXCom` in the - * generated `execute` body, with the result cast to the parameter's declared - * type. Non-`void` return values are forwarded to `client.setXCom`. + * In the generated `execute` body, a task's data parameters resolve through + * [ArgValues] against the arg bindings the supervisor delivered for the run: + * flat parameters by their position among the data parameters, [TaskInput] + * bundle fields by wire name. Non-`void` return values are forwarded to + * `client.setXCom`. */ @SupportedAnnotationTypes("org.apache.airflow.sdk.Builder.Dag") @SupportedSourceVersion(SourceVersion.RELEASE_11) @@ -102,23 +107,22 @@ class BuilderProcessor : AbstractProcessor() { MethodSpec .methodBuilder("build") .addModifiers(Modifier.PUBLIC, Modifier.STATIC) - .returns(ClassName.get(DagDef::class.java)) - .addStatement($$"var dag = new $T($S)", ClassName.get(DagDef::class.java), ann.id.ifBlank { el.simpleName }) + .returns(DAG_DEF_TYPE) + .addStatement($$"var dag = new $T($S)", DAG_DEF_TYPE, ann.id.ifBlank { el.simpleName }) for (inner in el.enclosedElements) { if (inner !is ExecutableElement) continue if (inner.isVarArgs) throw IllegalArgumentException("Cannot create task from vararg function ${inner.simpleName}") - val ann = inner.getAnnotation(Builder.Task::class.java) ?: continue + val taskAnn = inner.getAnnotation(Builder.Task::class.java) ?: continue val innerName = inner.simpleName.toString().replaceFirstChar(Char::uppercase) - val task = buildTask(innerName, inner, el) - builderClass.addType(task.spec) + builderClass.addType(buildTask(innerName, inner, el)) buildMethod.addStatement( $$"dag.addTask(new $T($S, $L.class))", - ClassName.get(TaskDef::class.java), - ann.id.ifBlank { inner.simpleName }, + TASK_DEF_TYPE, + taskAnn.id.ifBlank { inner.simpleName }, innerName, ) } @@ -132,40 +136,44 @@ class BuilderProcessor : AbstractProcessor() { name: String, inner: ExecutableElement, parent: TypeElement, - ): BuildTaskResult { - val clientType = ClassName.get(Client::class.java) - val contextType = ClassName.get(Context::class.java) - + ): TypeSpec { val executeSpec = MethodSpec .methodBuilder("execute") .addAnnotation(Override::class.java) .addModifiers(Modifier.PUBLIC) .returns(TypeName.VOID) - .addParameter(contextType, "context") - .addParameter(clientType, "client") + .addParameter(CONTEXT_TYPE, "context") + .addParameter(CLIENT_TYPE, "client") .addException(Exception::class.java) - val required = mutableListOf() + val dataParams = collectDataParams(inner) + val dataByName = dataParams.associateBy { it.name } val innerArgs = with(processingEnv) { inner.parameters.joinToString { param -> - val anno = param.getAnnotation(Builder.XCom::class.java) val type = param.asType() when { - anno != null -> - param.simpleName.toString().also { - required += RequiredXCom(type, it, anno.task.ifBlank { it }) - } - isType(type, clientType) -> "client" - isType(type, contextType) -> "context" - else -> throw IllegalArgumentException("Unsupported task parameter '${param.simpleName}' with type: $type") + isType(type, CLIENT_TYPE) -> "client" + isType(type, CONTEXT_TYPE) -> "context" + else -> dataByName.getValue(param.simpleName.toString()).name } } } - required.forEach { - executeSpec.addStatement($$"var $L = $L", it.paramName, xcomAccess(it)) + + dataParams.forEach { param -> + val paramType = TypeName.get(param.type) + val fields = param.bundleFields + if (fields == null) { + executeSpec.addStatement($$"$T $L = $L", paramType, param.name, positionalAccess(param)) + } else { + executeSpec.addStatement($$"$T $L = new $T()", paramType, param.name, paramType) + fields.forEach { field -> + executeSpec.addStatement($$"$L.$L = $L", param.name, field.fieldName, namedAccess(field)) + } + } } + if (inner.returnType.kind == TypeKind.VOID) { $$"new $T().$L($L)" } else { @@ -179,80 +187,154 @@ class BuilderProcessor : AbstractProcessor() { ) } - val spec = - TypeSpec - .classBuilder(name) - .addSuperinterface(Task::class.java) - .addModifiers(Modifier.PUBLIC, Modifier.FINAL, Modifier.STATIC) - .addMethod(executeSpec.build()) - .build() - return BuildTaskResult(spec) + return TypeSpec + .classBuilder(name) + .addSuperinterface(Task::class.java) + .addModifiers(Modifier.PUBLIC, Modifier.FINAL, Modifier.STATIC) + .addMethod(executeSpec.build()) + .build() + } + + /** + * Collects the task method's data parameters — every parameter the SDK does + * not inject — in declaration order. A parameter's index in the returned + * list is the position it binds at: Java parameter names are not API, so + * renaming one must never rebind an input. + */ + private fun collectDataParams(method: ExecutableElement): List { + val params = mutableListOf() + with(processingEnv) { + for (param in method.parameters) { + val type = param.asType() + if (isType(type, CLIENT_TYPE) || isType(type, CONTEXT_TYPE)) continue + val bundleFields = if (isTaskInput(type)) collectBundleFields(method, param) else null + params += DataParam(type, param.simpleName.toString(), params.size, bundleFields) + } + } + val bundles = params.filter { it.bundleFields != null } + require(bundles.size <= 1) { + "Task method '${method.simpleName}' declares more than one TaskInput parameter: " + + bundles.joinToString { "'${it.name}'" } + } + bundles.singleOrNull()?.let { bundle -> + require(params.size == 1) { + "Task method '${method.simpleName}' declares TaskInput parameter '${bundle.name}' and other data " + + "parameters; a TaskInput bundle owns the whole named-argument surface, so it must be the only one" + } + } + return params + } + + private fun ProcessingEnvironment.isTaskInput(type: TypeMirror): Boolean { + val marker = elementUtils.getTypeElement(TASK_INPUT_TYPE.canonicalName()) ?: return false + return !type.kind.isPrimitive && typeUtils.isAssignable(type, marker.asType()) + } + + /** + * Introspects a [TaskInput] bundle class: every public non-static non-final + * field receives the binding named by its [ArgName] value, or by its + * verbatim field name. + */ + private fun ProcessingEnvironment.collectBundleFields( + method: ExecutableElement, + param: VariableElement, + ): List { + val bundleType = + typeUtils.asElement(param.asType()) as? TypeElement + ?: throw IllegalArgumentException( + "TaskInput parameter '${param.simpleName}' of task method '${method.simpleName}' has no class type", + ) + val hasNoArgConstructor = + bundleType.enclosedElements + .filterIsInstance() + .any { it.kind == ElementKind.CONSTRUCTOR && it.parameters.isEmpty() && Modifier.PUBLIC in it.modifiers } + require(hasNoArgConstructor) { + "TaskInput class ${bundleType.simpleName} needs a public no-argument constructor" + } + return bundleType.enclosedElements + .filterIsInstance() + .filter { it.kind == ElementKind.FIELD && Modifier.STATIC !in it.modifiers } + .map { field -> + require(Modifier.PUBLIC in field.modifiers && Modifier.FINAL !in field.modifiers) { + "TaskInput field ${bundleType.simpleName}.${field.simpleName} must be public and non-final " + + "so the generated code can assign its binding" + } + BundleField( + type = field.asType(), + fieldName = field.simpleName.toString(), + wireName = field.getAnnotation(ArgName::class.java)?.value ?: field.simpleName.toString(), + ) + } } } +/** + * One data parameter of a task method, positioned among its peers. + * [bundleFields] is non-null for a [TaskInput] bundle parameter. + */ +private class DataParam( + val type: TypeMirror, + val name: String, + val position: Int, + val bundleFields: List?, +) + +/** One public field of a [TaskInput] bundle class, with its wire name. */ +private class BundleField( + val type: TypeMirror, + val fieldName: String, + val wireName: String, +) + +private val DAG_DEF_TYPE = ClassName.get(DagDef::class.java) +private val TASK_DEF_TYPE = ClassName.get(TaskDef::class.java) +private val CLIENT_TYPE = ClassName.get(Client::class.java) +private val CONTEXT_TYPE = ClassName.get(Context::class.java) +private val TASK_INPUT_TYPE = ClassName.get(TaskInput::class.java) +private val ARG_VALUES_TYPE = ClassName.get(ArgValues::class.java) + private fun ProcessingEnvironment.isType( t: TypeMirror, c: ClassName, ): Boolean = typeUtils.isSameType(t, elementUtils.getTypeElement(c.canonicalName()).asType()) -private data class RequiredXCom( - val paramType: TypeMirror, - val paramName: String, - val taskId: String, -) - -private val NUMBER_ACCESSORS: Map = - buildMap { - mapOf( - TypeName.BYTE to "byteValue", - TypeName.SHORT to "shortValue", - TypeName.INT to "intValue", - TypeName.LONG to "longValue", - TypeName.FLOAT to "floatValue", - TypeName.DOUBLE to "doubleValue", - ).forEach { (primitive, accessor) -> - put(primitive, accessor) - put(primitive.box(), accessor) - } +/** + * Emits the resolve-and-decode expression for one flat data parameter, bound + * at its position. A primitive parameter cannot hold null, so it fails with a + * clear [MissingXComException] when the binding resolves to nothing; boxed and + * reference parameters receive null instead. + */ +private fun positionalAccess(param: DataParam): CodeBlock { + val type = TypeName.get(param.type) + return if (type.isPrimitive) { + CodeBlock.of( + $$"$T.requiredInput(context, client, $L, $T.class, $S)", + ARG_VALUES_TYPE, + param.position, + type.box(), + param.name, + ) + } else { + val raw = (type as? ParameterizedTypeName)?.rawType ?: type + val call = CodeBlock.of($$"$T.optionalInput(context, client, $L, $T.class)", ARG_VALUES_TYPE, param.position, raw) + if (type is ParameterizedTypeName) CodeBlock.of($$"($T) $L", type, call) else call } +} -private fun xcomAccess(xcom: RequiredXCom): CodeBlock { - val type = TypeName.get(xcom.paramType) - val accessor = NUMBER_ACCESSORS[type] - val number = ClassName.get(Number::class.java) - val optional = ClassName.get(Optional::class.java) - // A primitive parameter cannot hold null, so fail with a clear error instead of an - // opaque NullPointerException while unboxing when the XCom is absent. - val value = - if (type.isPrimitive) { - CodeBlock.of( - $$"$T.ofNullable(client.getXCom($S)).orElseThrow(() -> new $T($S, $S))", - optional, - xcom.taskId, - ClassName.get(MissingXComException::class.java), - xcom.taskId, - xcom.paramName, - ) - } else { - CodeBlock.of($$"client.getXCom($S)", xcom.taskId) - } - // Wire integers decode to Long and floats to Double, so a direct (Integer)/(Float) - // cast throws ClassCastException; widen via Number instead. - return when { - accessor == null -> CodeBlock.of($$"($T) $L", if (type.isPrimitive) type.box() else type, value) - type.isPrimitive -> CodeBlock.of($$"(($T) $L).$L()", number, value, accessor) - else -> - CodeBlock.of( - $$"$T.ofNullable(($T) $L).map($T::$L).orElse(null)", - optional, - number, - value, - number, - accessor, - ) +/** Emits the resolve-and-decode expression for one bundle field, bound by wire name. */ +private fun namedAccess(field: BundleField): CodeBlock { + val type = TypeName.get(field.type) + return if (type.isPrimitive) { + CodeBlock.of( + $$"$T.requiredNamed(client, $S, $T.class, $S)", + ARG_VALUES_TYPE, + field.wireName, + type.box(), + field.fieldName, + ) + } else { + val raw = (type as? ParameterizedTypeName)?.rawType ?: type + val call = CodeBlock.of($$"$T.optionalNamed(client, $S, $T.class)", ARG_VALUES_TYPE, field.wireName, raw) + if (type is ParameterizedTypeName) CodeBlock.of($$"($T) $L", type, call) else call } } - -private data class BuildTaskResult( - val spec: TypeSpec, -) diff --git a/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt b/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt index 6a08979c56163..679aeea4338dd 100644 --- a/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt +++ b/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt @@ -62,7 +62,7 @@ class BuilderTest { } @Builder.Task - public void t3(Context ctx, @Builder.XCom(task = "t2") int value) { + public void t3(Context ctx, int value) { System.out.println(String.format("%s %s", ctx.ti, value)); } } @@ -78,15 +78,14 @@ class BuilderTest { package org.apache.airflow.example; import java.lang.Exception; - import java.lang.Number; + import java.lang.Integer; import java.lang.Override; - import java.util.Optional; import org.apache.airflow.sdk.Client; import org.apache.airflow.sdk.Context; import org.apache.airflow.sdk.DagDef; - import org.apache.airflow.sdk.MissingXComException; import org.apache.airflow.sdk.Task; import org.apache.airflow.sdk.TaskDef; + import org.apache.airflow.sdk.internal.ArgValues; public final class TestExampleBuilder { public static DagDef build() { @@ -111,7 +110,7 @@ class BuilderTest { public static final class T3 implements Task { @Override public void execute(Context context, Client client) throws Exception { - var value = ((Number) Optional.ofNullable(client.getXCom("t2")).orElseThrow(() -> new MissingXComException("t2", "value"))).intValue(); + int value = ArgValues.requiredInput(context, client, 0, Integer.class, "value"); new TestExample().t3(context, value); } } @@ -121,25 +120,19 @@ class BuilderTest { } @Test - @DisplayName("widen primitive numerics directly and boxed numerics null-safely") - fun generateBuilderWidensNumericXCom() { + @DisplayName("bind data parameters by position, skipping the injected Client and Context") + fun generateBuilderBindsDataParametersByPosition() { val compilation = compile( """ package org.apache.airflow.example; import org.apache.airflow.sdk.Builder; + import org.apache.airflow.sdk.Client; + import org.apache.airflow.sdk.Context; @Builder.Dag public class TestExample { @Builder.Task - public void t( - @Builder.XCom(task = "a") int i, - @Builder.XCom(task = "b") long l, - @Builder.XCom(task = "c") double d, - @Builder.XCom(task = "f") float fl, - @Builder.XCom(task = "e") Integer boxedInteger, - @Builder.XCom(task = "g") Long boxedLong, - @Builder.XCom(task = "h") Double boxedDouble, - @Builder.XCom(task = "j") Float boxedFloat) {} + public void t(long first, Client client, String second, Context ctx, Integer third) {} } """, ) @@ -153,15 +146,16 @@ class BuilderTest { package org.apache.airflow.example; import java.lang.Exception; - import java.lang.Number; + import java.lang.Integer; + import java.lang.Long; import java.lang.Override; - import java.util.Optional; + import java.lang.String; import org.apache.airflow.sdk.Client; import org.apache.airflow.sdk.Context; import org.apache.airflow.sdk.DagDef; - import org.apache.airflow.sdk.MissingXComException; import org.apache.airflow.sdk.Task; import org.apache.airflow.sdk.TaskDef; + import org.apache.airflow.sdk.internal.ArgValues; public final class TestExampleBuilder { public static DagDef build() { @@ -172,15 +166,10 @@ class BuilderTest { public static final class T implements Task { @Override public void execute(Context context, Client client) throws Exception { - var i = ((Number) Optional.ofNullable(client.getXCom("a")).orElseThrow(() -> new MissingXComException("a", "i"))).intValue(); - var l = ((Number) Optional.ofNullable(client.getXCom("b")).orElseThrow(() -> new MissingXComException("b", "l"))).longValue(); - var d = ((Number) Optional.ofNullable(client.getXCom("c")).orElseThrow(() -> new MissingXComException("c", "d"))).doubleValue(); - var fl = ((Number) Optional.ofNullable(client.getXCom("f")).orElseThrow(() -> new MissingXComException("f", "fl"))).floatValue(); - var boxedInteger = Optional.ofNullable((Number) client.getXCom("e")).map(Number::intValue).orElse(null); - var boxedLong = Optional.ofNullable((Number) client.getXCom("g")).map(Number::longValue).orElse(null); - var boxedDouble = Optional.ofNullable((Number) client.getXCom("h")).map(Number::doubleValue).orElse(null); - var boxedFloat = Optional.ofNullable((Number) client.getXCom("j")).map(Number::floatValue).orElse(null); - new TestExample().t(i, l, d, fl, boxedInteger, boxedLong, boxedDouble, boxedFloat); + long first = ArgValues.requiredInput(context, client, 0, Long.class, "first"); + String second = ArgValues.optionalInput(context, client, 1, String.class); + Integer third = ArgValues.optionalInput(context, client, 2, Integer.class); + new TestExample().t(first, client, second, context, third); } } } @@ -189,20 +178,19 @@ class BuilderTest { } @Test - @DisplayName("guard non-numeric primitives, leave objects and boxed types nullable") - fun generateBuilderGuardsNonNumericPrimitiveXCom() { + @DisplayName("require primitive parameters, leave boxed and parameterized types nullable") + fun generateBuilderRequiresPrimitivesOnly() { val compilation = compile( """ package org.apache.airflow.example; + import java.util.List; + import java.util.Map; import org.apache.airflow.sdk.Builder; @Builder.Dag public class TestExample { @Builder.Task - public void t( - @Builder.XCom(task = "a") boolean flag, - @Builder.XCom(task = "b") String text, - @Builder.XCom(task = "c") Boolean boxed) {} + public void t(boolean flag, float fraction, Double boxed, List tags, Map raw) {} } """, ) @@ -216,16 +204,19 @@ class BuilderTest { package org.apache.airflow.example; import java.lang.Boolean; + import java.lang.Double; import java.lang.Exception; + import java.lang.Float; import java.lang.Override; import java.lang.String; - import java.util.Optional; + import java.util.List; + import java.util.Map; import org.apache.airflow.sdk.Client; import org.apache.airflow.sdk.Context; import org.apache.airflow.sdk.DagDef; - import org.apache.airflow.sdk.MissingXComException; import org.apache.airflow.sdk.Task; import org.apache.airflow.sdk.TaskDef; + import org.apache.airflow.sdk.internal.ArgValues; public final class TestExampleBuilder { public static DagDef build() { @@ -236,10 +227,12 @@ class BuilderTest { public static final class T implements Task { @Override public void execute(Context context, Client client) throws Exception { - var flag = (Boolean) Optional.ofNullable(client.getXCom("a")).orElseThrow(() -> new MissingXComException("a", "flag")); - var text = (String) client.getXCom("b"); - var boxed = (Boolean) client.getXCom("c"); - new TestExample().t(flag, text, boxed); + boolean flag = ArgValues.requiredInput(context, client, 0, Boolean.class, "flag"); + float fraction = ArgValues.requiredInput(context, client, 1, Float.class, "fraction"); + Double boxed = ArgValues.optionalInput(context, client, 2, Double.class); + List tags = (List) ArgValues.optionalInput(context, client, 3, List.class); + Map raw = ArgValues.optionalInput(context, client, 4, Map.class); + new TestExample().t(flag, fraction, boxed, tags, raw); } } } @@ -247,6 +240,179 @@ class BuilderTest { ) } + @Test + @DisplayName("bind input-bundle fields by wire name") + fun generateBuilderBindsInputBundleFields() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import java.util.List; + import org.apache.airflow.sdk.ArgName; + import org.apache.airflow.sdk.Builder; + import org.apache.airflow.sdk.Client; + import org.apache.airflow.sdk.TaskInput; + @Builder.Dag + public class TestExample { + public static class ScoreInput implements TaskInput { + @ArgName("region_code") public String region; + public double threshold; + public List tags; + } + + @Builder.Task + public double score(Client client, ScoreInput input) { return input.threshold; } + } + """, + ) + + assertThat(compilation).succeeded() + assertThat(compilation) + .generatedSourceFile("org.apache.airflow.example.TestExampleBuilder") + .hasSourceEquivalentTo( + "org.apache.airflow.example.TestExampleBuilder", + """ + package org.apache.airflow.example; + + import java.lang.Double; + import java.lang.Exception; + import java.lang.Override; + import java.lang.String; + import java.util.List; + import org.apache.airflow.sdk.Client; + import org.apache.airflow.sdk.Context; + import org.apache.airflow.sdk.DagDef; + import org.apache.airflow.sdk.Task; + import org.apache.airflow.sdk.TaskDef; + import org.apache.airflow.sdk.internal.ArgValues; + + public final class TestExampleBuilder { + public static DagDef build() { + var dag = new DagDef("TestExample"); + dag.addTask(new TaskDef("score", Score.class)); + return dag; + } + public static final class Score implements Task { + @Override + public void execute(Context context, Client client) throws Exception { + TestExample.ScoreInput input = new TestExample.ScoreInput(); + input.region = ArgValues.optionalNamed(client, "region_code", String.class); + input.threshold = ArgValues.requiredNamed(client, "threshold", Double.class, "threshold"); + input.tags = (List) ArgValues.optionalNamed(client, "tags", List.class); + client.setXCom(new TestExample().score(client, input)); + } + } + } + """, + ) + } + + @Test + @DisplayName("reject an input bundle mixed with flat data parameters") + fun rejectBundleMixedWithFlatParams() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + import org.apache.airflow.sdk.TaskInput; + @Builder.Dag + public class TestExample { + public static class ScoreInput implements TaskInput { + public double threshold; + } + + @Builder.Task + public void t(ScoreInput input, int extra) {} + } + """, + ) + assertThat(compilation).failed() + assertThat(compilation).hadErrorContaining( + "Task method 't' declares TaskInput parameter 'input' and other data parameters", + ) + } + + @Test + @DisplayName("reject a task declaring more than one input bundle") + fun rejectMultipleBundles() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + import org.apache.airflow.sdk.TaskInput; + @Builder.Dag + public class TestExample { + public static class ScoreInput implements TaskInput { + public double threshold; + } + + @Builder.Task + public void t(ScoreInput first, ScoreInput second) {} + } + """, + ) + assertThat(compilation).failed() + assertThat(compilation).hadErrorContaining( + "Task method 't' declares more than one TaskInput parameter: 'first', 'second'", + ) + } + + @Test + @DisplayName("reject an input bundle with a non-public field") + fun rejectBundleWithNonPublicField() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + import org.apache.airflow.sdk.TaskInput; + @Builder.Dag + public class TestExample { + public static class ScoreInput implements TaskInput { + double threshold; + } + + @Builder.Task + public void t(ScoreInput input) {} + } + """, + ) + assertThat(compilation).failed() + assertThat(compilation).hadErrorContaining( + "TaskInput field ScoreInput.threshold must be public and non-final", + ) + } + + @Test + @DisplayName("reject an input bundle without a public no-argument constructor") + fun rejectBundleWithoutNoArgConstructor() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + import org.apache.airflow.sdk.TaskInput; + @Builder.Dag + public class TestExample { + public static class ScoreInput implements TaskInput { + public double threshold; + + public ScoreInput(double threshold) { this.threshold = threshold; } + } + + @Builder.Task + public void t(ScoreInput input) {} + } + """, + ) + assertThat(compilation).failed() + assertThat(compilation).hadErrorContaining( + "TaskInput class ScoreInput needs a public no-argument constructor", + ) + } + @Test @DisplayName("generate builder for dag class with custom dag id") fun generateBuilderWithCustomDagId() { @@ -333,24 +499,6 @@ class BuilderTest { ) } - @Test - @DisplayName("generate builder for dag class with invalid task parameter") - fun generateBuilderForDagClassWithInvalidTaskParameter() { - val compilation = - compile( - """ - package org.apache.airflow.example; - import org.apache.airflow.sdk.Builder; - @Builder.Dag - public class TestExample { @Builder.Task(id = "foo") public void t1(String client) {} } - """, - ) - assertThat(compilation).failed() - assertThat(compilation).hadErrorContaining( - "Unsupported task parameter 'client' with type: java.lang.String", - ) - } - @Test @DisplayName("generate builder for dag class with varargs task parameter") fun generateBuilderForDagClassWithVarArgsTaskParameter() { diff --git a/java-sdk/sdk/schema/schema.json b/java-sdk/sdk/schema/schema.json index e6ce8aa3d066e..b671959c50a00 100644 --- a/java-sdk/sdk/schema/schema.json +++ b/java-sdk/sdk/schema/schema.json @@ -1,6 +1,6 @@ { "$schema": "https://json-schema.org/draft/2020-12/schema", - "api_version": "2026-06-16", + "api_version": "2026-10-30", "description": "Apache Airflow SDK Supervisor Schema", "$defs": { "AssetAliasReferenceAssetEventDagRun": { @@ -753,6 +753,19 @@ ], "title": "Bundle Version" }, + "version_data": { + "anyOf": [ + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Version Data" + }, "msg": { "anyOf": [ { @@ -1651,6 +1664,19 @@ ], "title": "Bundle Version" }, + "version_data": { + "anyOf": [ + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Version Data" + }, "msg": { "anyOf": [ { @@ -1846,6 +1872,45 @@ "title": "Ascending", "type": "boolean" }, + "partition_key": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Partition Key" + }, + "partition_key_regexp_pattern": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Partition Key Regexp Pattern" + }, + "extra": { + "anyOf": [ + { + "additionalProperties": { + "type": "string" + }, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Extra" + }, "type": { "const": "GetAssetEventByAsset", "default": "GetAssetEventByAsset", @@ -1909,6 +1974,45 @@ "title": "Ascending", "type": "boolean" }, + "partition_key": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Partition Key" + }, + "partition_key_regexp_pattern": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Partition Key Regexp Pattern" + }, + "extra": { + "anyOf": [ + { + "additionalProperties": { + "type": "string" + }, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Extra" + }, "type": { "const": "GetAssetEventByAssetAlias", "default": "GetAssetEventByAssetAlias", @@ -3859,6 +3963,19 @@ ], "title": "Bundle Version" }, + "version_data": { + "anyOf": [ + { + "additionalProperties": true, + "type": "object" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Version Data" + }, "msg": { "anyOf": [ { @@ -4446,6 +4563,131 @@ "title": "ConnectionResponse", "type": "object" }, + "ArgValueSchema": { + "additionalProperties": { + "$ref": "#/$defs/JsonValue" + }, + "title": "ArgValueSchema", + "type": "object" + }, + "LiteralArgBinding": { + "description": "One positional stub-task argument carrying an inline literal from the Dag file.", + "properties": { + "kind": { + "const": "literal", + "title": "Kind", + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + }, + "value_schema": { + "anyOf": [ + { + "$ref": "#/$defs/ArgValueSchema" + }, + { + "type": "null" + } + ], + "default": null + }, + "value": { + "anyOf": [ + { + "$ref": "#/$defs/JsonValue" + }, + { + "type": "null" + } + ], + "default": null + }, + "from_default": { + "default": false, + "title": "From Default", + "type": "boolean" + } + }, + "required": [ + "kind", + "name" + ], + "title": "LiteralArgBinding", + "type": "object" + }, + "TaskArgBinding": { + "discriminator": { + "mapping": { + "literal": "#/$defs/LiteralArgBinding", + "xcom": "#/$defs/XComArgBinding" + }, + "propertyName": "kind" + }, + "oneOf": [ + { + "$ref": "#/$defs/XComArgBinding" + }, + { + "$ref": "#/$defs/LiteralArgBinding" + } + ], + "title": "TaskArgBinding" + }, + "XComArgBinding": { + "description": "One positional stub-task argument pulled from an upstream task's XCom.", + "properties": { + "kind": { + "const": "xcom", + "title": "Kind", + "type": "string" + }, + "name": { + "title": "Name", + "type": "string" + }, + "value_schema": { + "anyOf": [ + { + "$ref": "#/$defs/ArgValueSchema" + }, + { + "type": "null" + } + ], + "default": null + }, + "task_id": { + "title": "Task Id", + "type": "string" + }, + "map_index": { + "default": -1, + "title": "Map Index", + "type": "integer" + }, + "element_index": { + "anyOf": [ + { + "type": "integer" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Element Index" + } + }, + "required": [ + "kind", + "name", + "task_id" + ], + "title": "XComArgBinding", + "type": "object" + }, "AssetEventDagRunReference": { "additionalProperties": false, "description": "Schema for AssetEvent model used in DagRun.", @@ -4809,6 +5051,21 @@ ], "default": null, "title": "Start Date" + }, + "arg_bindings": { + "anyOf": [ + { + "items": { + "$ref": "#/$defs/TaskArgBinding" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Arg Bindings" } }, "required": [ diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/ArgName.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/ArgName.kt new file mode 100644 index 0000000000000..3af778d7cd71b --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/ArgName.kt @@ -0,0 +1,40 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.airflow.sdk + +/** + * Declares the wire name of a [TaskInput] field explicitly, for when the + * Python stub signature's argument name is not a valid (or desirable) Java + * field name — typically `snake_case` arguments crossing into `camelCase` + * fields. + * + * ```java + * @ArgName("region_code") public String region; + * ``` + * + * Fields without the annotation bind their verbatim field name. + * + * @param value Argument name as declared in the stub task's signature. + */ +@Target(AnnotationTarget.FIELD) +@MustBeDocumented +annotation class ArgName( + val value: String, +) diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Builder.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Builder.kt index 9dbfbfbefc390..1e8ca63670838 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Builder.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Builder.kt @@ -36,10 +36,15 @@ package org.apache.airflow.sdk * public long extract(Client client) { ... } * * @Builder.Task(id = "transform") - * public long transform(Client client, @Builder.XCom(task = "extract") long extracted) { ... } + * public long transform(Client client, long extracted) { ... } * } * ``` * + * A task method's data parameters — everything other than the injected + * [Client] and [Context] — receive the arguments the Python `@task.stub` call + * site bound, by position. Keyword arguments bind by name instead through a + * single [TaskInput] bundle parameter. + * * The processor generates `MyPipelineBuilder.build()`, which returns a * fully wired-up [DagDef] ready to add to a [Bundle]. */ @@ -73,16 +78,4 @@ class Builder internal constructor() { annotation class Task( val id: String = "", ) - - /** - * Annotation to mark a task definition's method parameter as an XCom input. - * - * @param task The task ID to pull. If empty or not given, the annotated - * parameter's name is used by default. - */ - @Target(AnnotationTarget.VALUE_PARAMETER) - @MustBeDocumented - annotation class XCom( - val task: String = "", - ) } diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Client.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Client.kt index 59aa7832a3756..941ac86aad7ac 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Client.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Client.kt @@ -19,8 +19,10 @@ package org.apache.airflow.sdk +import org.apache.airflow.sdk.execution.ArgBinding import org.apache.airflow.sdk.execution.Client import org.apache.airflow.sdk.execution.comm.StartupDetails +import org.apache.airflow.sdk.execution.decodeArgBindings /** * A connection registered in Airflow's connection store. @@ -152,6 +154,98 @@ class Client internal constructor( runId = details.ti.runId, mapIndex = details.ti.mapIndex ?: -1, ) + + internal val argBindings: List by lazy { + decodeArgBindings(details.tiContext?.argBindings) + } + + // A literal binding carries the inline value from the Dag file; an XCom + // binding pulls the bound upstream task's return-value XCom, honouring the + // bound map index and element index. + internal fun resolveBinding(binding: ArgBinding): Any? = + when (binding) { + is ArgBinding.Literal -> binding.value + is ArgBinding.XCom -> { + val value = getXCom(taskId = binding.taskId, mapIndex = binding.mapIndex.takeIf { it >= 0 }) + when { + binding.elementIndex == null -> value + value is List<*> -> value[binding.elementIndex] + else -> + error( + "Argument '${binding.name}' binds element ${binding.elementIndex} of task '${binding.taskId}', " + + "but its XCom is not a list", + ) + } + } + } + + /** + * Whether the supervisor delivered TaskFlow arg bindings for this run — + * that is, the Python Dag file called this stub task with TaskFlow + * arguments. + */ + fun hasArgs(): Boolean = argBindings.isNotEmpty() + + /** + * Whether a TaskFlow argument was bound at [position] of the stub task's + * signature. + * + * @param position Zero-based position in the stub call's argument list. + */ + fun hasArg(position: Int): Boolean = position in argBindings.indices + + /** + * Whether the Python Dag file bound a TaskFlow argument with this name to + * the current stub task. + * + * @param name Argument name as declared in the stub task's signature. + */ + fun hasArg(name: String): Boolean = argBindings.any { it.name == name } + + /** + * Resolves the TaskFlow argument bound at [position] of the stub task's + * signature. + * + * A literal binding returns the inline value from the Dag file; an XCom + * binding pulls the bound upstream task's return-value XCom, honouring the + * bound map index and element index. + * + * @param position Zero-based position in the stub call's argument list. + * @return The bound value, or `null` when the bound value is null or the + * upstream pushed no value. + * @throws IllegalArgumentException if no argument was bound at this + * position; use [hasArg] to probe. + * @throws ApiError if the underlying XCom read fails. + */ + fun getArg(position: Int): Any? { + require(position in argBindings.indices) { + "No TaskFlow argument bound at position: $position" + } + return resolveBinding(argBindings[position]) + } + + /** + * Resolves the TaskFlow argument bound with [name] at the `@task.stub` + * call site in the Python Dag file. + * + * A literal binding returns the inline value from the Dag file; an XCom + * binding pulls the bound upstream task's return-value XCom, honouring the + * bound map index and element index. + * + * @param name Argument name as declared in the stub task's signature. + * @return The bound value, or `null` when the bound value is null or the + * upstream pushed no value. + * @throws IllegalArgumentException if no argument with this name was bound; + * use [hasArg] to probe. + * @throws ApiError if the underlying XCom read fails. + */ + fun getArg(name: String): Any? { + val binding = + requireNotNull(argBindings.firstOrNull { it.name == name }) { + "No TaskFlow argument bound with name: '$name'" + } + return resolveBinding(binding) + } } /** diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/TaskInput.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/TaskInput.kt new file mode 100644 index 0000000000000..306600e0ff443 --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/TaskInput.kt @@ -0,0 +1,45 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.airflow.sdk + +/** + * Marks a class as a task's input bundle: when the Python Dag file calls the + * stub task with keyword arguments, each public field receives the runtime + * binding whose name matches the field ([ArgName] value, or the verbatim + * field name). + * + * A task method may declare at most one `TaskInput` parameter and, if it + * does, no other data parameters — the bundle owns the whole named-argument + * surface, so field names and flat positions cannot shift each other. + * + * ```java + * public static class ScoreInput implements TaskInput { + * @ArgName("region_code") public String region; // explicit wire name + * public double threshold; // binds "threshold" + * } + * + * @Builder.Task + * public Result score(Client client, ScoreInput input) { ... } + * ``` + * + * The class needs a public no-argument constructor and public non-final + * fields. + */ +interface TaskInput diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/ArgBinding.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/ArgBinding.kt new file mode 100644 index 0000000000000..808eff687d51e --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/ArgBinding.kt @@ -0,0 +1,77 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.airflow.sdk.execution + +/** + * One stub-task argument bound at the `@task.stub` TaskFlow call site in the + * Python Dag file, delivered via `TIRunContext.arg_bindings`. + * + * The supervisor schema models this as a `kind`-discriminated union + * (`XComArgBinding` / `LiteralArgBinding`), which jsonSchema2Pojo cannot + * express as a typed field — the generated `TIRunContext.argBindings` is a + * plain `Object` holding the msgpack-decoded list of maps — so this hand- + * written decoder materializes the typed view. + */ +internal sealed class ArgBinding { + abstract val name: String + + internal data class XCom( + override val name: String, + val taskId: String, + val mapIndex: Int, + val elementIndex: Int?, + ) : ArgBinding() + + internal data class Literal( + override val name: String, + val value: Any?, + ) : ArgBinding() +} + +/** + * Decodes the raw `TIRunContext.argBindings` payload into a list of bindings + * preserving the stub signature's parameter order — flat task parameters + * bind by that position, input-bundle fields by [ArgBinding.name]. + * + * @throws IllegalStateException on a malformed payload, an unsupported + * binding kind, or a duplicate argument name; the task cannot bind its + * arguments correctly, so it must fail rather than run with wrong inputs. + */ +internal fun decodeArgBindings(raw: Any?): List { + if (raw == null) return emptyList() + check(raw is List<*>) { "arg_bindings payload is not a list: ${raw.javaClass.name}" } + val seen = mutableSetOf() + return raw.map { entry -> + check(entry is Map<*, *>) { "arg_bindings entry is not a map: $entry" } + val name = checkNotNull(entry["name"] as? String) { "arg_bindings entry has no name: $entry" } + check(seen.add(name)) { "arg_bindings entries have duplicate name: '$name'" } + when (val kind = entry["kind"]) { + "literal" -> ArgBinding.Literal(name = name, value = entry["value"]) + "xcom" -> + ArgBinding.XCom( + name = name, + taskId = checkNotNull(entry["task_id"] as? String) { "xcom arg binding '$name' has no task_id" }, + mapIndex = (entry["map_index"] as? Number)?.toInt() ?: -1, + elementIndex = (entry["element_index"] as? Number)?.toInt(), + ) + else -> error("Unsupported arg binding kind '$kind' for argument '$name'") + } + } +} diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/ArgValues.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/ArgValues.kt new file mode 100644 index 0000000000000..cd56f38470668 --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/ArgValues.kt @@ -0,0 +1,175 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +@file:Suppress("PLATFORM_CLASS_MAPPED_TO_KOTLIN") + +package org.apache.airflow.sdk.internal + +import com.fasterxml.jackson.databind.ObjectMapper +import com.fasterxml.jackson.databind.json.JsonMapper +import org.apache.airflow.sdk.Client +import org.apache.airflow.sdk.Context +import org.apache.airflow.sdk.MissingXComException +import org.apache.airflow.sdk.execution.ArgBinding + +/** + * Resolves a task's data parameters from the arg bindings the supervisor + * delivered, and decodes their raw wire values into the declared parameter + * types. Public so that processor-generated task classes can call it; not + * user-facing API. + * + * The bindings come from the Python `@task.stub` call site, which is also the + * graph the scheduler ordered the run by. Flat data parameters resolve the + * binding at their position; input-bundle fields resolve bindings by name. + */ +object ArgValues { + private val mapper: ObjectMapper = JsonMapper.builder().build().findAndRegisterModules() + + /** + * Resolves the data parameter at [position] into [type] for a parameter that + * cannot be null. + * + * @param position Zero-based index among the task's data parameters, in + * declaration order. + * @throws MissingXComException if the resolved value is null. + */ + @JvmStatic + fun requiredInput( + context: Context, + client: Client, + position: Int, + type: Class, + paramName: String, + ): T { + val binding = bindingAt(context, client, position) + return decode(client.resolveBinding(binding), type) ?: throw missing(binding, paramName) + } + + /** + * Resolves the data parameter at [position] into [type], passing null + * through. + * + * @param position Zero-based index among the task's data parameters, in + * declaration order. + */ + @JvmStatic + fun optionalInput( + context: Context, + client: Client, + position: Int, + type: Class, + ): T? = decode(client.resolveBinding(bindingAt(context, client, position)), type) + + /** + * Resolves the runtime binding named [name] into [type] for an input-bundle + * field that cannot be null. + * + * @param name Wire name of the argument (`@ArgName` value or the verbatim + * field name). + * @throws IllegalStateException if the stub call bound no argument named + * [name]. + * @throws MissingXComException if the resolved value is null. + */ + @JvmStatic + fun requiredNamed( + client: Client, + name: String, + type: Class, + fieldName: String, + ): T { + val binding = + checkNotNull(client.argBindings.firstOrNull { it.name == name }) { + "The stub call bound no argument named '$name', required by input field '$fieldName'" + } + return decode(client.resolveBinding(binding), type) ?: throw missing(binding, fieldName, name) + } + + /** + * Resolves the runtime binding named [name] into [type], passing null + * through. An absent binding resolves to null. + * + * @param name Wire name of the argument (`@ArgName` value or the verbatim + * field name). + */ + @JvmStatic + fun optionalNamed( + client: Client, + name: String, + type: Class, + ): T? { + val binding = client.argBindings.firstOrNull { it.name == name } ?: return null + return decode(client.resolveBinding(binding), type) + } + + private fun bindingAt( + context: Context, + client: Client, + position: Int, + ): ArgBinding { + val bindings = client.argBindings + check(position < bindings.size) { + "Task '${context.ti.taskId}' declares a data parameter at position $position " + + "but the stub call bound only ${bindings.size} argument(s)" + } + return bindings[position] + } + + private fun missing( + binding: ArgBinding, + target: String, + argName: String? = null, + ): MissingXComException = + when (binding) { + is ArgBinding.XCom -> MissingXComException(binding.taskId, target) + is ArgBinding.Literal -> + MissingXComException( + "'$target' has a primitive type but the stub call bound a null literal" + + (argName?.let { " for argument '$it'" } ?: "") + + "; declare a boxed type (e.g. Integer instead of int) to receive null.", + ) + } + + internal fun decode( + value: Any?, + type: Class, + ): T? { + if (value == null) return null + if (type.isInstance(value)) return type.cast(value) + // The msgpack decoder yields Long for wire integers and Double for wire + // floats, so widen numerics via Number instead of casting. + if (value is Number) { + numberConverter(type)?.let { return type.cast(it(value)) } + } + // Structured wire values (maps, lists) convert into the declared POJO or + // collection type; unknown fields fail the task, mirroring the Go SDK's + // strict decode of task inputs. + return mapper.convertValue(value, type) + } + + private fun numberConverter(type: Class<*>): ((Number) -> Any)? = + when (type) { + java.lang.Byte::class.java -> { n -> n.toByte() } + java.lang.Short::class.java -> { n -> n.toShort() } + java.lang.Integer::class.java -> { n -> n.toInt() } + java.lang.Long::class.java -> { n -> n.toLong() } + java.lang.Float::class.java -> { n -> n.toFloat() } + java.lang.Double::class.java -> { n -> n.toDouble() } + else -> null + } +} diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ClientArgTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ClientArgTest.kt new file mode 100644 index 0000000000000..e6b68629212b3 --- /dev/null +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ClientArgTest.kt @@ -0,0 +1,297 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +@file:Suppress("PLATFORM_CLASS_MAPPED_TO_KOTLIN") + +package org.apache.airflow.sdk + +import org.apache.airflow.sdk.execution.comm.ConnectionResult +import org.apache.airflow.sdk.execution.comm.StartupDetails +import org.apache.airflow.sdk.execution.comm.TIRunContext +import org.apache.airflow.sdk.execution.comm.VariableResult +import org.apache.airflow.sdk.execution.comm.XComResult +import org.apache.airflow.sdk.internal.ArgValues +import org.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.Assertions.assertFalse +import org.junit.jupiter.api.Assertions.assertNull +import org.junit.jupiter.api.Assertions.assertThrows +import org.junit.jupiter.api.Assertions.assertTrue +import org.junit.jupiter.api.DisplayName +import org.junit.jupiter.api.Test +import org.apache.airflow.sdk.execution.comm.TaskInstance as CommTaskInstance + +/** Records getXCom calls and serves canned values keyed by task id. */ +private class FakeXComTransport( + val xcoms: Map = emptyMap(), +) : org.apache.airflow.sdk.execution.Client { + val pulls = mutableListOf>() + + override fun getConnection(id: String): ConnectionResult = throw NotImplementedError() + + override fun getVariable(key: String): VariableResult = throw NotImplementedError() + + override fun getXCom( + key: String, + dagId: String, + taskId: String, + runId: String, + mapIndex: Int?, + includePriorDates: Boolean, + ): XComResult { + pulls += taskId to mapIndex + return XComResult().also { + it.key = key + it.value = xcoms[taskId] + } + } + + override fun setXCom( + key: String, + value: Any, + dagId: String, + taskId: String, + runId: String, + mapIndex: Int, + ) = throw NotImplementedError() +} + +private fun startupDetails(argBindings: List>?): StartupDetails = + StartupDetails().also { details -> + details.ti = + CommTaskInstance().also { + it.dagId = "d" + it.runId = "r" + it.taskId = "t" + it.tryNumber = 1 + } + details.tiContext = TIRunContext().also { it.argBindings = argBindings } + } + +private fun clientWith( + argBindings: List>?, + xcoms: Map = emptyMap(), +): Pair { + val transport = FakeXComTransport(xcoms) + return Client(startupDetails(argBindings), transport) to transport +} + +private fun taskContext(): Context = + Context( + dagRun = DagRun("d", "r", null, null, null, null, null, emptyMap()), + ti = TaskInstance("d", "r", "t", null, 1), + ) + +internal class ClientArgTest { + @Test + @DisplayName("Should resolve a literal binding to its inline value, by position and by name") + fun shouldResolveLiteralBinding() { + val (client, transport) = clientWith(listOf(mapOf("kind" to "literal", "name" to "x", "value" to 42L))) + + assertTrue(client.hasArgs()) + assertTrue(client.hasArg(0)) + assertTrue(client.hasArg("x")) + assertEquals(42L, client.getArg(0)) + assertEquals(42L, client.getArg("x")) + assertEquals(emptyList>(), transport.pulls) + } + + @Test + @DisplayName("Should resolve an xcom binding by pulling the bound task's return value") + fun shouldResolveXComBinding() { + val (client, transport) = + clientWith( + listOf(mapOf("kind" to "xcom", "name" to "x", "task_id" to "upstream", "map_index" to -1L)), + xcoms = mapOf("upstream" to 7L), + ) + + assertEquals(7L, client.getArg(0)) + assertEquals(listOf("upstream" to null), transport.pulls) + } + + @Test + @DisplayName("Should keep bindings in stub-signature order") + fun shouldKeepBindingOrder() { + val (client, _) = + clientWith( + listOf( + mapOf("kind" to "literal", "name" to "b", "value" to 2L), + mapOf("kind" to "literal", "name" to "a", "value" to 1L), + ), + ) + + assertEquals(2L, client.getArg(0)) + assertEquals(1L, client.getArg(1)) + } + + @Test + @DisplayName("Should pass a non-negative bound map index to the XCom read") + fun shouldPassBoundMapIndex() { + val (client, transport) = + clientWith( + listOf(mapOf("kind" to "xcom", "name" to "x", "task_id" to "upstream", "map_index" to 2L)), + xcoms = mapOf("upstream" to 7L), + ) + + client.getArg(0) + + assertEquals(listOf("upstream" to 2), transport.pulls) + } + + @Test + @DisplayName("Should index into a list XCom when the binding has an element index") + fun shouldResolveElementIndex() { + val (client, _) = + clientWith( + listOf(mapOf("kind" to "xcom", "name" to "x", "task_id" to "upstream", "element_index" to 1L)), + xcoms = mapOf("upstream" to listOf("a", "b", "c")), + ) + + assertEquals("b", client.getArg(0)) + } + + @Test + @DisplayName("Should fail when an element index points into a non-list XCom") + fun shouldRejectElementIndexOnNonList() { + val (client, _) = + clientWith( + listOf(mapOf("kind" to "xcom", "name" to "x", "task_id" to "upstream", "element_index" to 1L)), + xcoms = mapOf("upstream" to "scalar"), + ) + + assertThrows(IllegalStateException::class.java) { client.getArg(0) } + } + + @Test + @DisplayName("Should reject reading an argument that was never bound") + fun shouldRejectUnknownArg() { + val (client, _) = clientWith(listOf(mapOf("kind" to "literal", "name" to "x", "value" to 1L))) + + assertFalse(client.hasArg("y")) + assertFalse(client.hasArg(1)) + assertThrows(IllegalArgumentException::class.java) { client.getArg("y") } + assertThrows(IllegalArgumentException::class.java) { client.getArg(1) } + } + + @Test + @DisplayName("Should report no bound arguments when the supervisor sent none") + fun shouldHandleAbsentBindings() { + val (client, _) = clientWith(null) + + assertFalse(client.hasArgs()) + assertFalse(client.hasArg("x")) + assertFalse(client.hasArg(0)) + } + + @Test + @DisplayName("Should fail on an unsupported binding kind") + fun shouldRejectUnknownBindingKind() { + val (client, _) = clientWith(listOf(mapOf("kind" to "mystery", "name" to "x"))) + + assertThrows(IllegalStateException::class.java) { client.hasArg("x") } + } + + @Test + @DisplayName("Should fail on duplicate binding names") + fun shouldRejectDuplicateBindingNames() { + val (client, _) = + clientWith( + listOf( + mapOf("kind" to "literal", "name" to "x", "value" to 1L), + mapOf("kind" to "literal", "name" to "x", "value" to 2L), + ), + ) + + assertThrows(IllegalStateException::class.java) { client.hasArgs() } + } + + @Test + @DisplayName("Should resolve a flat data parameter from the binding at its position") + fun shouldResolvePositionalBinding() { + val (client, transport) = + clientWith( + listOf( + mapOf("kind" to "literal", "name" to "first", "value" to 5L), + mapOf("kind" to "xcom", "name" to "second", "task_id" to "upstream"), + ), + xcoms = mapOf("upstream" to "pulled"), + ) + + assertEquals(5, ArgValues.requiredInput(taskContext(), client, 0, Integer::class.java, "first").toInt()) + assertEquals("pulled", ArgValues.optionalInput(taskContext(), client, 1, String::class.java)) + assertEquals(listOf("upstream" to null), transport.pulls) + } + + @Test + @DisplayName("Should fail fast when the stub call bound fewer arguments than declared") + fun shouldFailOnArityMismatch() { + val (client, _) = clientWith(listOf(mapOf("kind" to "literal", "name" to "only", "value" to 1L))) + + val error = + assertThrows(IllegalStateException::class.java) { + ArgValues.optionalInput(taskContext(), client, 1, Integer::class.java) + } + + assertEquals( + "Task 't' declares a data parameter at position 1 but the stub call bound only 1 argument(s)", + error.message, + ) + } + + @Test + @DisplayName("Should throw MissingXComException for a required argument bound to a null literal") + fun shouldThrowForNullLiteralOnRequired() { + val (client, _) = clientWith(listOf(mapOf("kind" to "literal", "name" to "value", "value" to null))) + + assertThrows(MissingXComException::class.java) { + ArgValues.requiredInput(taskContext(), client, 0, Integer::class.java, "value") + } + } + + @Test + @DisplayName("Should resolve input-bundle fields by wire name") + fun shouldResolveNamedBindings() { + val (client, _) = + clientWith( + listOf( + mapOf("kind" to "literal", "name" to "region_code", "value" to "emea"), + mapOf("kind" to "xcom", "name" to "threshold", "task_id" to "upstream"), + ), + xcoms = mapOf("upstream" to 0.5), + ) + + assertEquals("emea", ArgValues.optionalNamed(client, "region_code", String::class.java)) + assertEquals(0.5, ArgValues.requiredNamed(client, "threshold", java.lang.Double::class.java, "threshold")) + } + + @Test + @DisplayName("Should resolve an absent named binding to null for optional fields and fail for required ones") + fun shouldHandleAbsentNamedBinding() { + val (client, _) = clientWith(listOf(mapOf("kind" to "literal", "name" to "other", "value" to 1L))) + + assertNull(ArgValues.optionalNamed(client, "missing", String::class.java)) + val error = + assertThrows(IllegalStateException::class.java) { + ArgValues.requiredNamed(client, "missing", Integer::class.java, "field") + } + assertEquals( + "The stub call bound no argument named 'missing', required by input field 'field'", + error.message, + ) + } +} From 923b119e7fc2cafd7dbb679cd31084557527462c Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 5 Aug 2026 08:56:10 +0000 Subject: [PATCH 09/10] Java SDK: Author complete Dags in Java without a Python stub file Until now the Java SDK could only supply task bodies: a Python @task.stub Dag had to own the schedule, every task option, and the graph. That splits one pipeline across two languages and two repositories for no reason other than a missing authoring surface, and it left the Java-side model with nothing to describe -- no edges, no configuration -- so there was nothing a native Java Dag could be built from. Java annotations cannot change call semantics the way Python decorators do, so the graph is declared against a compile-time-generated twin class (`Ref`): calling a twin registers the task and passing one twin's handle into another feeds the upstream's output into the downstream's parameter, making the call graph the task graph the way Python TaskFlow does -- but type-checked by javac through the In/TaskRef generics. Keeping the wiring calls the only way to express an edge means there is one graph story to learn instead of two, and the method stays optional so stub-backed classes are unchanged: their graph still lives in the Python Dag file, and runtime arg bindings continue to win over anything Java declares, because for a stub task the Python call site is the graph the scheduler ordered the run by. Dag and task configuration is generated from Airflow's own Dag serialization schema rather than hand-listed, so the Java attributes cannot drift from the Python semantics they mirror and new scalar keys appear after a schema sync. Only attributes written at the use site are applied, leaving Airflow's defaults in charge of everything unset. Design rationale is recorded in ADR-0007. --- .pre-commit-config.yaml | 10 + .../adr/lang-sdk/0007-taskflow-dag-dsl.md | 240 ++++++++ airflow-core/adr/lang-sdk/README.md | 1 + .../language-sdks/java.rst | 100 +++- java-sdk/README.md | 19 +- .../airflow/example/ExampleBundleBuilder.java | 4 +- .../example/nativedag/AnnotationExample.java | 62 ++ .../example/nativedag/InterfaceExample.java | 77 +++ .../apache/airflow/sdk/BuilderProcessor.kt | 354 ++++++++++-- .../org/apache/airflow/sdk/BuilderTest.kt | 542 +++++++++++++++++- java-sdk/sdk/build.gradle.kts | 432 +++++++++++++- java-sdk/sdk/schema/dag-schema.json | 496 ++++++++++++++++ .../kotlin/org/apache/airflow/sdk/Builder.kt | 81 --- .../kotlin/org/apache/airflow/sdk/Bundle.kt | 35 +- .../kotlin/org/apache/airflow/sdk/Context.kt | 3 + .../kotlin/org/apache/airflow/sdk/DagDef.kt | 106 +++- .../main/kotlin/org/apache/airflow/sdk/In.kt | 64 +++ .../kotlin/org/apache/airflow/sdk/Wiring.kt | 49 ++ .../org/apache/airflow/sdk/execution/Task.kt | 10 +- .../apache/airflow/sdk/internal/ArgValues.kt | 92 ++- .../org/apache/airflow/sdk/internal/Fields.kt | 107 ++++ .../org/apache/airflow/sdk/internal/Refs.kt | 68 +++ .../org/apache/airflow/sdk/BundleTest.kt | 75 +++ .../org/apache/airflow/sdk/ClientArgTest.kt | 38 ++ .../org/apache/airflow/sdk/DagDefTest.kt | 132 +++++ .../apache/airflow/sdk/execution/TaskTest.kt | 24 + .../airflow/sdk/internal/ArgValuesTest.kt | 193 +++++++ .../apache/airflow/sdk/internal/RefsTest.kt | 88 +++ 28 files changed, 3321 insertions(+), 181 deletions(-) create mode 100644 airflow-core/adr/lang-sdk/0007-taskflow-dag-dsl.md create mode 100644 java-sdk/example/src/java/org/apache/airflow/example/nativedag/AnnotationExample.java create mode 100644 java-sdk/example/src/java/org/apache/airflow/example/nativedag/InterfaceExample.java create mode 100644 java-sdk/sdk/schema/dag-schema.json delete mode 100644 java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Builder.kt create mode 100644 java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/In.kt create mode 100644 java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Wiring.kt create mode 100644 java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/Fields.kt create mode 100644 java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/Refs.kt create mode 100644 java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/ArgValuesTest.kt create mode 100644 java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/RefsTest.kt diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 038d8ebc80e38..8bf8f7893b02a 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -280,6 +280,16 @@ repos: (?x) ^java-sdk/gradle\.properties$| ^java-sdk/sdk/schema/schema\.json$ + - id: sync-java-sdk-dag-schema + name: Sync Java SDK Dag serialization schema with airflow-core + description: "Copy airflow-core's serialization schema when Java SDK's vendored dag-schema.json drifts" + entry: ./java-sdk/gradlew -p ./java-sdk :sdk:syncDagSchema + language: system + pass_filenames: false + files: > + (?x) + ^airflow-core/src/airflow/serialization/schema\.json$| + ^java-sdk/sdk/schema/dag-schema\.json$ - id: check-go-version-in-sync name: Check Go toolchain version is consistent across build files entry: ./scripts/ci/prek/check_go_version_in_sync.py diff --git a/airflow-core/adr/lang-sdk/0007-taskflow-dag-dsl.md b/airflow-core/adr/lang-sdk/0007-taskflow-dag-dsl.md new file mode 100644 index 0000000000000..8aad4f3ae01a1 --- /dev/null +++ b/airflow-core/adr/lang-sdk/0007-taskflow-dag-dsl.md @@ -0,0 +1,240 @@ + + +# ADR-0007: TaskFlow-Style Dag DSL for Native Java Dags + +## Status + +Accepted + +## Context + +The Go SDK adopted a TaskFlow-style authoring API: `Dag.Task(fn, opts...)` +returns a handle, `Inputs(handle)` both wires the dependency edge and feeds the +upstream's return value into the downstream function's data parameter, and +`TaskSpec` / `DagSpec` configuration structs are generated from the Dag +serialization schema (`airflow-core/src/airflow/serialization/schema.json`). + +The Java SDK could only implement *task bodies*. Everything else about the Dag +had to come from a Python `@task.stub` file: + +- `@Builder.Dag` / `@Builder.Task` carried only identity attributes; there was + no way to express schedule, retries, queue, or any other Dag/task + configuration. +- The Java-side model kept only an `id -> task class` map: no dependency edges + and no configuration, so a Java author could not describe a graph and there + was nothing to serialize for native Dag support. + +Python's TaskFlow shows the shape worth matching: calling tasks like functions +*is* the graph declaration — `load(transform(extract()))`. Java annotations +cannot change call semantics (invoking the real method would execute its body), +so the call syntax has to target compile-time-generated twins. + +## Decision + +### Two authoring surfaces, one model + +**Annotation surface** — task bodies are `@Builder.Task`-annotated methods; +configuration lives in the annotations; the graph is wired by a static +`@Wiring` method that receives a processor-generated *flow twin* class +(`Ref`). Twin methods mirror the task methods: injectable parameters +(`Client`, `Context`) are dropped, data parameters become `In`-typed +inputs, and the return value becomes a `TaskRef` (which extends +`In`). The call graph is the task graph: + +```java +@Builder.Dag(id = "java_etl", schedule = "@daily") +public class EtlPipeline { + + @Builder.Task(id = "extract", retries = 2) + public ExtractResult extract(Client client) { ... } + + @Builder.Task + public TransformResult transform(Client client, ExtractResult in) { ... } + + @Builder.Task + public void load(Context ctx, TransformResult in) { ... } + + @Builder.Task + public void score(double threshold) { ... } + + @Wiring + static void depends(EtlPipelineRef f) { + f.load(f.transform(f.extract())); + f.score(In.value(0.5)); // inline literal input + } +} +``` + +There is deliberately **no id-based wiring mode** and no edge-declaring API +on the handles: the wiring calls, and only they, define the graph. A twin call +registers the task, so a `@Builder.Task` method never invoked in the wiring +method is an error (checked when the generated `build()` runs, at Dag-parse +time). + +The wiring method is **optional**: a `@Builder.Dag` class without one registers +every task with no Java-side edges. That is the shape for stub-backed tasks — +the Python Dag file owns the graph and the supervisor delivers each call-site +argument at run time — so those classes carry no `@Wiring`, no configuration +attributes, and no Java-side dependencies at all. + +**Interface surface** — plain-class task definitions registered as first-class +`TaskDef` objects, with fluent schema-validated `.config(key, value)` calls +and object-reference edges: + +```java +var extract = new TaskDef("extract", Extract.class).config("retries", 2); +var transform = new TaskDef("transform", Transform.class).dependsOn(extract); + +var dag = new DagDef("java_etl") + .config("schedule", "@daily") + .addTask(extract) + .addTask(transform); +``` + +A `TaskDef` belongs to at most one Dag. Both surfaces flow into the same +model: schema-keyed configuration maps, `TaskDef` dependency edges, and +recorded parameter inputs on `DagDef`/`TaskDef`. + +### One package, one wildcard import + +The whole user-facing surface — the `Builder.Dag` / `Builder.Task` +annotations, `@Wiring`, the flow types (`In`, `TaskRef`), and the runtime +types (`DagDef`, `TaskDef`, `Task`, `Client`, `Context`, `Bundle`, ...) — +lives in `org.apache.airflow.sdk`, so every example starts with a single +`import org.apache.airflow.sdk.*`. There is no separate `dsl` package. + +Two top-level types cannot share a fully-qualified name, and `Dag` / `Task` +are the names both the annotations and the runtime types want. Keeping the +annotations **nested** in `Builder` resolves that without renaming anything: +`Builder.Dag` and `Builder.Task` read as the annotations they are (they drive +the `*Builder` codegen), the task-implementation interface keeps the plain +name `Task`, and the Dag model is `DagDef`, symmetric with `TaskDef` — both +are the definition objects of the interface API. `@Wiring` needs no +qualification because nothing else claims that name. + +By convention the `@Wiring` method sits at the end of the Dag class, after +the task methods it wires — read the tasks first, then the graph. + +### Wiring is type-checked by javac itself + +Twin input types make the graph checks ordinary Java type checking rather +than bespoke processor analysis: + +- numeric parameters accept any numeric upstream (`In`, + widened or narrowed at run time by the shared decoder); +- `Object`, raw `Map`, and raw `List` parameters accept any upstream + (`In`, decoded loosely at run time); +- everything else accepts covariant matches of the declared type + (`In`). + +Unknown upstreams are unrepresentable (a handle only exists once its task is +registered), cycles are unconstructible in call syntax, and type mismatches +are javac errors at the twin call site. The interface API's `dependsOn` can +still express a cycle, so `Bundle` construction validates acyclicity (and +that every referenced upstream is registered in the same Dag). + +### Runtime bindings win; wiring is the fallback + +Data parameters resolve by **position**, never by parameter name: Java call +syntax is positional (no kwargs) and Java parameter names are not API — an +IDE rename must not change binding behaviour. This matches the Go SDK's +flat-parameter contract. + +When the supervisor delivered `arg_bindings` for the run, the binding at the +parameter's position wins over anything the `@Wiring` method declared: for a +stub task the Python call site *is* the graph the scheduler ordered the run +by, so the Java class must not be able to disagree with it. The wiring- +recorded inputs are the fallback, used when no bindings arrived — which is +exactly the native-Dag case, where no Python call site exists. A `TaskRef` +input then resolves to the upstream's return-value XCom and a literal input +to its value. Keyword arguments still bind by name through a `TaskInput` +bundle; the generated code branches on whether bindings arrived, because +bindings fill the bundle field by field while the wiring fallback decodes the +bundle wholesale from its single wired input. + +### Generated from the serialization schema + +The `Builder` class — the outer container plus its nested `Dag` and `Task` +annotations, configuration attributes and all — and the `SchemaFields` +validation table are generated at build time from a vendored copy of the Dag +serialization schema (`sdk/schema/dag-schema.json`, kept in sync with +airflow-core by the `sync-java-sdk-dag-schema` prek hook), mirroring the +supervisor-schema → jsonschema2pojo pipeline that already exists. Generating +the whole class rather than merging generated attributes into a hand-written +one keeps a single definition of `Builder`; `id` (and `to` on `Dag`) stay the +leading structural attributes, and generation fails if a schema key ever +camel-cases onto one of them. + +Field selection mirrors the Go SDK's `TaskSpec` generator: scalar properties +only, serializer-owned keys skipped (`_`-prefixed, schema-required, +`has_on_*`), a documented exclusion list for Python-only concerns that fails +generation when it goes stale, and a hand-curated Dag-level allowlist matching +Go's `DagSpec`. `schedule` is a virtual key that the future serializer maps to +the schema's `timetable` object. + +Temporal attributes are ISO-8601 strings in annotations (validated by the +processor at compile time) and `java.time.Duration` / +`java.time.OffsetDateTime` values in `.config` calls. + +### Explicit-only lowering, single-sourced dag id + +The processor lowers **only attributes written at the use site** into +`DagDef.config(...)` / `TaskDef.config(...)` calls (via +`AnnotationMirror.getElementValues`). Annotation defaults mirror schema +defaults but are never emitted, so the serializer's omit-if-default semantics +stay intact and the scheduler's own defaults win for everything unset. + +Because this written-vs-defaulted distinction only exists at compile time, +there is no reflective `new DagDef(SomeClass.class)` constructor. Instead the +generated builder exposes `DAG_ID` and a `dag()` factory (Dag-level config +only, no tasks) alongside `build()`, so the dag id is never restated in user +code. + +### No annotation on data parameters + +Data parameters need no annotation at all: anything that is not an injectable +type (`Client`, `Context`) is a data parameter, bound in declaration order. +Naming an upstream in an annotation would duplicate what the wiring — or the +Python call site — already declares. `Client.getXCom` / `Client.setXCom` +remain for imperative access. + +## Consequences + +- One concept, one name: bindings are *args*/*inputs* across the wire + contract, the Go SDK, and the Java surface; the Java graph is declared the + way Python TaskFlow declares it, but type-checked at compile time. +- Dependency edges, configuration, and parameter inputs now exist in the + Java-side model, which is the prerequisite for emitting + DagSerialization-v3 JSON for native Java Dags. +- Config typos fail at compile time (annotations) or Dag-parse time + (`.config`), never silently. Wiring mistakes fail at compile time (type + mismatch, unknown handle) or Dag-parse time (unregistered task, `dependsOn` + cycle), never at task run time. +- An annotation Dag class that owns its graph writes a `@Wiring` method — + the same posture as the Go SDK, trading a few lines for one wiring story + instead of two. Stub-backed classes omit it and stay registration-only. +- The `@Wiring` method references the generated `Ref` twin, so IDEs + show unresolved symbols until the first successful build (the standard + Java codegen experience, as with Dagger or AutoValue). +- New scalar schema keys show up automatically in the annotations and the + validation table after a schema sync; removals fail the build until the + exclusion rules are updated — the surface cannot drift silently. +- Generated public API (annotation attributes) varies with the vendored + schema version, exactly like the generated supervisor-schema models. diff --git a/airflow-core/adr/lang-sdk/README.md b/airflow-core/adr/lang-sdk/README.md index 7092bc19f3157..a3493cc1818b5 100644 --- a/airflow-core/adr/lang-sdk/README.md +++ b/airflow-core/adr/lang-sdk/README.md @@ -32,6 +32,7 @@ bind core interfaces and apply to every language SDK, not just the Java SDK. - [ADR-0004](0004-dag-parsing.md): language-specific Dag file processing. - [ADR-0005](0005-coordinator-packaging.md): coordinator packaging, module layout, and registration. - [ADR-0006](0006-no-lang-sdk-source-display.md): no Lang-SDK source display for mixed-language (`@task.stub`) Dags. +- [ADR-0007](0007-taskflow-dag-dsl.md): TaskFlow-style Dag DSL for native Java Dags. Decisions specific to a single SDK stay next to that SDK — for example, the Go SDK's bundle-format decisions live in [`go-sdk/adr/`](../../../go-sdk/adr). 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 b3d1f2792fc78..280925dd240ba 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst @@ -160,10 +160,18 @@ Annotate a plain Java class and let the SDK generate the boilerplate at compile * - Annotation - Purpose * - ``@Builder.Dag(id = "...")`` - - Marks the class as a task container. The ``id`` must match the ``dag_id`` in the Python Dag. + - Marks the class as a task container. For a stub-backed Dag the ``id`` must match the + ``dag_id`` in the Python Dag. Further attributes (``schedule``, ``description``, ``tags``, + ``catchup``, …) mirror the Dag serialization schema and configure the Dag itself; only + attributes written explicitly are applied. See :ref:`java-sdk/native-dags`. * - ``@Builder.Task(id = "...")`` - - Marks a method as a task implementation. The ``id`` must match the ``@task.stub`` function - name in the Python Dag. If ``id`` is omitted the method name is used. + - Marks a method as a task implementation. For a stub-backed Dag the ``id`` must match the + ``@task.stub`` function name in the Python Dag. If ``id`` is omitted the method name is + used. Further attributes (``retries``, ``queue``, ``retryDelay``, …) mirror the Dag + serialization schema; only attributes written explicitly are applied. + * - ``@Wiring`` + - Marks the static method that declares the task graph in Java, TaskFlow-style. Only needed + for a Dag that has no Python stub file. See :ref:`java-sdk/native-dags`. * - ``TaskInput`` / ``@ArgName("...")`` - Marks a class as a task's input bundle, so keyword arguments bind by name instead of by position: each public field receives the binding whose name matches it (the ``@ArgName`` @@ -205,7 +213,9 @@ Interface-based API ~~~~~~~~~~~~~~~~~~~ Implement the ``Task`` interface directly for full control over how tasks are registered and how XComs are -read. Each task is registered as a ``TaskDef`` on a ``DagDef``. +read. Each task is registered as a ``TaskDef`` on a ``DagDef``; both carry a fluent +``config(key, value)`` whose keys are Dag serialization schema property names, and ``TaskDef`` also +carries ``dependsOn(...)`` for declaring edges between task definitions. .. code-block:: java @@ -227,9 +237,12 @@ Register tasks manually in a ``BundleBuilder``: public class MyBundle implements BundleBuilder { @Override public Iterable getDags() { + var fetch = new TaskDef("fetch", FetchTask.class).config("retries", 2); + var process = new TaskDef("process", ProcessTask.class).dependsOn(fetch); var dag = new DagDef("my_dag") - .addTask(new TaskDef("fetch", FetchTask.class)) - .addTask(new TaskDef("process", ProcessTask.class)); + .config("schedule", "@daily") + .addTask(fetch) + .addTask(process); return List.of(dag); } } @@ -335,6 +348,81 @@ Tasks written against the ``Task`` interface read the same bindings imperatively } } +.. _java-sdk/native-dags: + +Native Java Dags +---------------- + +A Dag can also be authored entirely in Java, with no Python stub file: the annotations (or the +``TaskDef`` / ``DagDef`` objects) carry the configuration, and Java declares the graph. + +Wiring the graph with ``@Wiring`` +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +The annotation processor generates a ``Ref`` twin class whose methods mirror the +``@Builder.Task`` methods: the injected ``Client`` and ``Context`` parameters are dropped, each data +parameter takes an ``In`` input, and the return value becomes a ``TaskRef``. A static +``@Wiring`` method receives the twin and calls it — calling a twin registers the task, and passing +one twin's result into another feeds the upstream's output into the downstream's parameter *and* +wires the dependency edge. The call graph is the task graph, and ``javac`` type-checks it: + +.. code-block:: java + + @Builder.Dag( + id = "java_etl", + schedule = "@daily", + description = "Pure-Java Dag, no Python stub file", + tags = {"example", "java-sdk"}) + public class EtlPipeline { + + @Builder.Task(id = "extract", retries = 2) + public long extract() { + return 42L; + } + + @Builder.Task(id = "transform") + public long transform(long extracted) { + return extracted * 2; + } + + @Builder.Task(id = "load") + public void load(long transformed) { + // implement task logic + } + + @Wiring + static void depends(EtlPipelineRef f) { + f.load(f.transform(f.extract())); + } + } + +Every ``@Builder.Task`` method must be invoked in the wiring method; a task the wiring missed fails +at Dag-parse time. ``In.value(...)`` wires an inline literal where no upstream feeds a parameter. +The wiring method is optional — a class without one registers every task with no Java-side edges, +which is the shape for stub-backed tasks whose graph the Python Dag file owns. + +.. note:: + + Runtime argument bindings win over Java-declared wiring. When the supervisor delivers bindings + for a run (see :ref:`java-sdk/arg-binding`), the binding at a parameter's position is what the + task receives, because for a stub task the Python call site is the graph the scheduler ordered + the run by. Wired inputs are the fallback, which is what a native Java Dag always uses. + +Configuration attributes +~~~~~~~~~~~~~~~~~~~~~~~~ + +The ``@Builder.Dag`` and ``@Builder.Task`` configuration attributes, and the keys accepted by +``DagDef.config`` and ``TaskDef.config``, are generated from Airflow's Dag serialization schema, so +they carry the same names and types as their Python counterparts. Annotation attributes are +``camelCase`` (``retryDelay``); ``config`` keys are the verbatim schema names (``"retry_delay"``). +Only attributes written explicitly at the use site are applied, so Airflow's own defaults still +apply to everything left out. + +Durations and date-times are ISO-8601 strings in annotations (``retryDelay = "PT5M"``, +``startDate = "2026-01-01T00:00:00Z"``, validated at compile time) and ``java.time.Duration`` / +``java.time.OffsetDateTime`` values in ``config`` calls. An unknown key or a mismatched value type +fails the build (annotations) or Dag parsing (``config``). + .. _java-sdk/logging: Logging diff --git a/java-sdk/README.md b/java-sdk/README.md index 7e3cbeafed1d4..18522a89f49d1 100644 --- a/java-sdk/README.md +++ b/java-sdk/README.md @@ -649,9 +649,13 @@ E2E_TEST_MODE=java_sdk uv run --project airflow-e2e-tests pytest \ not the implementation language. - Keep `sdk/src/main/kotlin/` (the public API surface) free of internal implementation details; those belong in the `execution/` sub-package. -- The annotation processor (`BuilderProcessor.kt`) uses `kapt`. When adding a - new annotation, define it in `Builder.kt`, handle it in - `BuilderProcessor.kt`, and add a golden-output test in +- The annotation processor (`BuilderProcessor.kt`) uses `kapt`. The `Builder` + class holding the `@Builder.Dag` / `@Builder.Task` annotations is generated + from the Dag serialization schema by `:sdk:generateDagDsl` (vendored at + `sdk/schema/dag-schema.json`); `@Wiring` and the `In`/`TaskRef` wiring types + are hand-written next to the rest of the public surface in + `sdk/src/main/kotlin/org/apache/airflow/sdk/`. When adding annotation + behaviour, handle it in `BuilderProcessor.kt` and add a golden-output test in `processor/src/test/kotlin/`. - The Python coordinator subclasses `SubprocessCoordinator`. Do not reach into the JVM process from Python beyond what `_build_execute_task_command` @@ -673,9 +677,14 @@ E2E_TEST_MODE=java_sdk uv run --project airflow-e2e-tests pytest \ 5. Update `airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst` if the change is user-visible. -**Adding a new annotation**: +**Adding a new annotation or configuration attribute**: -1. Define the annotation interface in `Builder.kt`. +1. Hand-written annotations live in `sdk/src/main/kotlin/org/apache/airflow/sdk/` + next to the runtime types (one package, so user code needs a single + `import org.apache.airflow.sdk.*`); the configuration attributes of + `@Builder.Dag` / `@Builder.Task` come from the Dag serialization schema via + `:sdk:generateDagDsl` (adjust its allowlist/exclusion rules in + `sdk/build.gradle.kts` when the exposed field set should change). 2. Handle it in `BuilderProcessor.kt` — generate the appropriate code in the `*Builder` class. 3. Add a test in `BuilderTest.kt` with expected generated output. 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 fa1a860755863..3a0868dafbc2a 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 @@ -30,7 +30,9 @@ public Iterable getDags() { return List.of( InterfaceExampleBuilder.build(), AnnotationExampleBuilder.build(), - XComCastingExampleBuilder.build()); + XComCastingExampleBuilder.build(), + org.apache.airflow.example.nativedag.AnnotationExampleBuilder.build(), + org.apache.airflow.example.nativedag.InterfaceExample.build()); } public static void main(String[] args) { diff --git a/java-sdk/example/src/java/org/apache/airflow/example/nativedag/AnnotationExample.java b/java-sdk/example/src/java/org/apache/airflow/example/nativedag/AnnotationExample.java new file mode 100644 index 0000000000000..67770f44f91f2 --- /dev/null +++ b/java-sdk/example/src/java/org/apache/airflow/example/nativedag/AnnotationExample.java @@ -0,0 +1,62 @@ +/* + * 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. + */ + +// "native" is a Java keyword, so the native-Dag examples live in "nativedag". +package org.apache.airflow.example.nativedag; + +import static java.lang.System.Logger.Level.INFO; + +import org.apache.airflow.sdk.*; + +// A Dag defined entirely in Java, annotation-style: no Python stub file +// describes it. The @Builder.Dag/@Builder.Task attributes carry the +// configuration and the @Wiring method declares the graph -- the single place +// dependencies are defined. +@Builder.Dag( + id = "java_native_annotation_example", + description = "Pure-Java Dag authored with annotations, without a Python stub file", + schedule = "@daily", + startDate = "2026-01-01T00:00:00Z", + catchup = false, + tags = {"example", "java-sdk"}) +public class AnnotationExample { + private static final System.Logger log = System.getLogger(AnnotationExample.class.getName()); + + @Builder.Task(id = "extract", retries = 2) + public long extract() { + log.log(INFO, "Extracting a value"); + return 42L; + } + + @Builder.Task(id = "transform") + public long transform(long extracted) { + log.log(INFO, "Transforming {0}", extracted); + return extracted * 2; + } + + @Builder.Task(id = "load") + public void load(long transformed) { + log.log(INFO, "Loaded {0}", transformed); + } + + @Wiring + static void depends(AnnotationExampleRef f) { + f.load(f.transform(f.extract())); + } +} diff --git a/java-sdk/example/src/java/org/apache/airflow/example/nativedag/InterfaceExample.java b/java-sdk/example/src/java/org/apache/airflow/example/nativedag/InterfaceExample.java new file mode 100644 index 0000000000000..02c91a480c619 --- /dev/null +++ b/java-sdk/example/src/java/org/apache/airflow/example/nativedag/InterfaceExample.java @@ -0,0 +1,77 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +// "native" is a Java keyword, so the native-Dag examples live in "nativedag". +package org.apache.airflow.example.nativedag; + +import static java.lang.System.Logger.Level.INFO; + +import java.util.List; +import org.apache.airflow.sdk.*; + +// A Dag defined entirely in Java, interface-style: no Python stub file +// describes it. TaskDef.config carries the task configuration, dependsOn +// declares the graph, and DagDef.config carries the Dag configuration. +public class InterfaceExample { + private static final System.Logger log = System.getLogger(InterfaceExample.class.getName()); + + public static class Extract implements Task { + @Override + public void execute(Context context, Client client) { + log.log(INFO, "Extracting a value"); + client.setXCom(42L); + } + } + + public static class Transform implements Task { + @Override + public void execute(Context context, Client client) { + var extracted = ((Number) client.getXCom("extract")).longValue(); + log.log(INFO, "Transforming {0}", extracted); + client.setXCom(extracted * 2); + } + } + + public static class Load implements Task { + @Override + public void execute(Context context, Client client) { + var transformed = client.getXCom("transform"); + log.log(INFO, "Loaded {0}", transformed); + } + } + + public static DagDef build() { + var extract = + new TaskDef("extract", Extract.class) + .config("retries", 2) + .config("doc_md", "Extracts a value and pushes it as an XCom."); + var transform = new TaskDef("transform", Transform.class).dependsOn(extract); + var load = new TaskDef("load", Load.class).dependsOn(transform); + return new DagDef("java_native_interface_example") + .config( + "description", + "Pure-Java Dag authored with the interface API, without a Python stub file") + .config("schedule", "@daily") + .config("catchup", false) + .config("tags", List.of("example", "java-sdk")) + .addTask(extract) + .addTask(transform) + .addTask(load); + } +} 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 3bd8eb336b95d..c2769b7d9eb3e 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 @@ -23,18 +23,29 @@ package org.apache.airflow.sdk import com.squareup.javapoet.ClassName import com.squareup.javapoet.CodeBlock +import com.squareup.javapoet.FieldSpec import com.squareup.javapoet.JavaFile import com.squareup.javapoet.MethodSpec import com.squareup.javapoet.ParameterizedTypeName import com.squareup.javapoet.TypeName import com.squareup.javapoet.TypeSpec +import com.squareup.javapoet.WildcardTypeName import org.apache.airflow.sdk.internal.ArgValues +import org.apache.airflow.sdk.internal.Field +import org.apache.airflow.sdk.internal.FieldType +import org.apache.airflow.sdk.internal.Refs +import org.apache.airflow.sdk.internal.SchemaFields +import java.time.Duration +import java.time.OffsetDateTime +import java.time.format.DateTimeParseException import javax.annotation.processing.AbstractProcessor import javax.annotation.processing.ProcessingEnvironment import javax.annotation.processing.RoundEnvironment import javax.annotation.processing.SupportedAnnotationTypes import javax.annotation.processing.SupportedSourceVersion import javax.lang.model.SourceVersion +import javax.lang.model.element.AnnotationValue +import javax.lang.model.element.Element import javax.lang.model.element.ElementKind import javax.lang.model.element.ExecutableElement import javax.lang.model.element.Modifier @@ -53,18 +64,28 @@ import javax.tools.Diagnostic * `META-INF/services/javax.annotation.processing.Processor`; not intended to be * instantiated or referenced directly. * - * For each class annotated with [Builder.Dag], generates a `*Builder` class - * containing: + * For each class annotated with [Builder.Dag], generates: * - * - One inner class per [Builder.Task]-annotated method, implementing [Task]. - * - A static `build()` method that constructs the [DagDef] and registers those - * inner classes as [TaskDef]s. + * - A `*Builder` class containing one inner class per [Builder.Task]-annotated + * method (implementing [Task]), a `DAG_ID` constant, a static `dag()` factory + * that lowers every explicitly-written `@Builder.Dag` attribute into + * `DagDef.config` calls, and a static `build()` that invokes the class's + * [Wiring] method and verifies it registered every task — or, when the class + * has no wiring method, registers every task with no Java-side edges. + * - A `*Ref` twin class (only when a wiring method exists) whose methods + * mirror the task methods: injectable parameters ([Client], [Context]) are + * dropped, data parameters become [In]-typed inputs, and the return value + * becomes a [TaskRef]. Calling a twin registers the task with its + * explicitly-written `@Builder.Task` attributes lowered into `TaskDef.config` + * calls; passing one twin's handle to another wires the dependency edge and + * feeds the upstream's return-value XCom into the downstream's parameter, + * type-checked by javac through the [In] / [TaskRef] generics. * - * In the generated `execute` body, a task's data parameters resolve through - * [ArgValues] against the arg bindings the supervisor delivered for the run: - * flat parameters by their position among the data parameters, [TaskInput] - * bundle fields by wire name. Non-`void` return values are forwarded to - * `client.setXCom`. + * In the generated `execute` bodies, a task's data parameters resolve through + * [ArgValues] against the arg bindings the supervisor delivered for the run, + * falling back to the wired inputs: flat parameters by their position, and + * [TaskInput] bundle fields by wire name. Non-`void` return values are + * forwarded to `client.setXCom`. */ @SupportedAnnotationTypes("org.apache.airflow.sdk.Builder.Dag") @SupportedSourceVersion(SourceVersion.RELEASE_11) @@ -77,12 +98,18 @@ class BuilderProcessor : AbstractProcessor() { roundEnv.getElementsAnnotatedWith(Builder.Dag::class.java).filterIsInstance().forEach { el -> with(processingEnv) { runCatching { + val packageName = elementUtils.getPackageOf(el).qualifiedName.toString() + val declarations = collectTasks(el) + val wiring = findWiring(el) + val builderName = ClassName.get(packageName, dagAnnotation(el).to.ifBlank { "${el.simpleName}Builder" }) + val refName = ClassName.get(packageName, "${el.simpleName}Ref") JavaFile - .builder( - elementUtils.getPackageOf(el).qualifiedName.toString(), - buildDag(el), - ).build() + .builder(packageName, buildBuilder(el, declarations, wiring, builderName, refName)) + .build() .writeTo(filer) + if (wiring != null) { + JavaFile.builder(packageName, buildRef(el, declarations, builderName, refName)).build().writeTo(filer) + } }.onFailure { e -> messager.printMessage( Diagnostic.Kind.ERROR, @@ -95,46 +122,264 @@ class BuilderProcessor : AbstractProcessor() { return true } - private fun buildDag(el: TypeElement): TypeSpec { - val ann = el.getAnnotation(Builder.Dag::class.java)!! + private fun dagAnnotation(el: TypeElement): Builder.Dag = el.getAnnotation(Builder.Dag::class.java)!! + + private fun buildBuilder( + el: TypeElement, + declarations: List, + wiring: ExecutableElement?, + builderName: ClassName, + refName: ClassName, + ): TypeSpec { + val ann = dagAnnotation(el) val builderClass = TypeSpec - .classBuilder(ann.to.ifBlank { "${el.simpleName}Builder" }) + .classBuilder(builderName) .addModifiers(Modifier.PUBLIC, Modifier.FINAL) + .addField( + FieldSpec + .builder(ClassName.get(String::class.java), "DAG_ID", Modifier.PUBLIC, Modifier.STATIC, Modifier.FINAL) + .initializer($$"$S", ann.id.ifBlank { el.simpleName }) + .build(), + ) + + val dagMethod = + MethodSpec + .methodBuilder("dag") + .addModifiers(Modifier.PUBLIC, Modifier.STATIC) + .returns(DAG_DEF_TYPE) + .addJavadoc("Returns a new {@code DagDef} carrying the Dag attributes, with no tasks registered.\n") + .addStatement($$"var dag = new $T(DAG_ID)", DAG_DEF_TYPE) + explicitConfig(el, DAG_ANNOTATION, DAG_STRUCTURAL_ATTRIBUTES, SchemaFields.DAG).forEach { (key, value) -> + dagMethod.addStatement($$"dag.config($S, $L)", key, value) + } + dagMethod.addStatement("return dag") + builderClass.addMethod(dagMethod.build()) val buildMethod = MethodSpec .methodBuilder("build") .addModifiers(Modifier.PUBLIC, Modifier.STATIC) .returns(DAG_DEF_TYPE) - .addStatement($$"var dag = new $T($S)", DAG_DEF_TYPE, ann.id.ifBlank { el.simpleName }) + .addStatement("var dag = dag()") + if (wiring != null) { + buildMethod.addStatement( + $$"$T.$L(new $T(dag))", + ClassName.get(el), + wiring.simpleName, + refName, + ) + buildMethod.addStatement( + $$"$T.requireRegistered(dag, $T.of($L))", + REFS_TYPE, + ClassName.get(List::class.java), + declarations.joinToString { "\"${it.id}\"" }, + ) + } else { + // No wiring method: register every task with no Java-side edges — a + // Python stub Dag defines the graph for these tasks. + declarations.forEach { decl -> + buildMethod.addStatement($$"dag.addTask($L)", taskDefCode(decl, CodeBlock.of($$"$L", decl.className))) + } + } + buildMethod.addStatement("return dag") + builderClass.addMethod(buildMethod.build()) + declarations.forEach { builderClass.addType(buildTask(it, el)) } + return builderClass.build() + } + + private fun buildRef( + el: TypeElement, + declarations: List, + builderName: ClassName, + refName: ClassName, + ): TypeSpec { + val flowClass = + TypeSpec + .classBuilder(refName) + .addModifiers(Modifier.PUBLIC, Modifier.FINAL) + .addJavadoc( + "Task-reference twins of {@link \$T}'s task methods, for wiring its task graph.\n\n" + + "

Calling a twin registers the task with the Dag under construction; passing one\n" + + "twin's return value into another feeds the upstream's output into the downstream's\n" + + "parameter and wires the dependency edge.\n", + ClassName.get(el), + ).addField(DAG_DEF_TYPE, "dag", Modifier.PRIVATE, Modifier.FINAL) + .addMethod( + MethodSpec + .constructorBuilder() + .addModifiers(Modifier.PUBLIC) + .addParameter(DAG_DEF_TYPE, "dag") + .addStatement("this.dag = dag") + .build(), + ) + + for (decl in declarations) { + val twin = + MethodSpec + .methodBuilder(decl.method.simpleName.toString()) + .addModifiers(Modifier.PUBLIC) + .returns(ParameterizedTypeName.get(TASK_HANDLE_TYPE, TypeName.get(decl.method.returnType).boxIfPossible())) + decl.dataParams.forEach { twin.addParameter(inType(it.type), it.name) } + twin.addStatement( + $$"return $T.register(dag, $L, $T.of($L))", + REFS_TYPE, + taskDefCode(decl, CodeBlock.of($$"$T.$L", builderName, decl.className)), + ClassName.get(List::class.java), + decl.dataParams.joinToString { it.name }, + ) + flowClass.addMethod(twin.build()) + } + return flowClass.build() + } + + /** + * Emits `new TaskDef(id, .class)` with the explicitly-written + * `@Builder.Task` attributes lowered into chained `.config` calls. + */ + private fun taskDefCode( + decl: TaskDeclaration, + classRef: CodeBlock, + ): CodeBlock { + val taskDef = + CodeBlock + .builder() + .add($$"new $T($S, $L.class)", TASK_DEF_TYPE, decl.id, classRef) + explicitConfig(decl.method, TASK_ANNOTATION, TASK_STRUCTURAL_ATTRIBUTES, SchemaFields.TASK).forEach { (key, value) -> + taskDef.add($$".config($S, $L)", key, value) + } + return taskDef.build() + } + + /** + * Maps a data parameter's declared type to its twin-input type. Numeric + * parameters accept any numeric upstream (`In`, widened + * at run time); `Object`, raw `Map`, and raw `List` parameters accept any + * upstream (`In`, decoded loosely at run time); everything else accepts + * covariant matches of the declared type (`In`). + */ + private fun inType(paramType: TypeMirror): TypeName { + val boxed = TypeName.get(paramType).boxIfPossible() + val argument = + when { + isNumeric(paramType) -> WildcardTypeName.subtypeOf(TypeName.get(Number::class.java)) + else -> WildcardTypeName.subtypeOf(boxed) + } + return ParameterizedTypeName.get(IN_TYPE, argument) + } + + private fun isNumeric(t: TypeMirror): Boolean = t.kind in NUMERIC_KINDS || TypeName.get(t) in BOXED_NUMERICS + + private fun collectTasks(el: TypeElement): List { + val declarations = mutableListOf() for (inner in el.enclosedElements) { if (inner !is ExecutableElement) continue + val ann = inner.getAnnotation(Builder.Task::class.java) ?: continue if (inner.isVarArgs) throw IllegalArgumentException("Cannot create task from vararg function ${inner.simpleName}") + val id = ann.id.ifBlank { inner.simpleName.toString() } + require(declarations.none { it.id == id }) { "Tasks in Dag have duplicate ID: $id" } + declarations += TaskDeclaration(inner, id, collectDataParams(inner)) + } + return declarations + } - val taskAnn = inner.getAnnotation(Builder.Task::class.java) ?: continue - val innerName = inner.simpleName.toString().replaceFirstChar(Char::uppercase) + /** + * Finds and validates the class's [Wiring] method. The method is optional: + * without one, every task registers with no Java-side edges. + */ + private fun findWiring(el: TypeElement): ExecutableElement? { + val methods = + el.enclosedElements + .filterIsInstance() + .filter { it.getAnnotation(Wiring::class.java) != null } + if (methods.isEmpty()) return null + val wiring = + methods.singleOrNull() + ?: throw IllegalArgumentException( + "Dag class ${el.simpleName} declares more than one @Wiring method: " + + methods.joinToString { it.simpleName.toString() }, + ) + require(Modifier.STATIC in wiring.modifiers && Modifier.PRIVATE !in wiring.modifiers) { + "@Wiring method '${wiring.simpleName}' must be static and non-private" + } + require(wiring.returnType.kind == TypeKind.VOID && wiring.parameters.size == 1) { + "@Wiring method '${wiring.simpleName}' must be void and take the generated ${el.simpleName}Ref as its only parameter" + } + return wiring + } - builderClass.addType(buildTask(innerName, inner, el)) + /** + * Lowers the explicitly-written configuration attributes of [element]'s + * [annotationName] annotation into (schema key, value code) pairs. Only + * attributes present at the use site are lowered, so annotation defaults + * never override the schema's own defaults. + */ + private fun explicitConfig( + element: Element, + annotationName: String, + structural: Set, + table: Map, + ): List> { + val mirror = + element.annotationMirrors.firstOrNull { + (it.annotationType.asElement() as TypeElement).qualifiedName.contentEquals(annotationName) + } ?: return emptyList() + val byAttribute = table.values.associateBy { it.attribute } + return mirror.elementValues.mapNotNull { (attr, value) -> + val name = attr.simpleName.toString() + if (name in structural) return@mapNotNull null + val field = + requireNotNull(byAttribute[name]) { + "Annotation attribute '$name' has no Dag serialization schema key" + } + field.key to configValueCode(field, value) + } + } - buildMethod.addStatement( - $$"dag.addTask(new $T($S, $L.class))", - TASK_DEF_TYPE, - taskAnn.id.ifBlank { inner.simpleName }, - innerName, - ) + private fun configValueCode( + field: Field, + value: AnnotationValue, + ): CodeBlock = + when (field.type) { + FieldType.STRING -> CodeBlock.of($$"$S", value.value) + FieldType.BOOLEAN, FieldType.INTEGER, FieldType.NUMBER -> CodeBlock.of($$"$L", value.value) + FieldType.STRING_ARRAY -> { + @Suppress("UNCHECKED_CAST") + val items = value.value as List + CodeBlock.of( + $$"$T.of($L)", + ClassName.get(List::class.java), + items.joinToString { "\"${it.value}\"" }, + ) + } + FieldType.TIMEDELTA -> { + val text = value.value as String + parseTemporal(field, text) { Duration.parse(text) } + CodeBlock.of($$"$T.parse($S)", ClassName.get(Duration::class.java), text) + } + FieldType.DATETIME -> { + val text = value.value as String + parseTemporal(field, text) { OffsetDateTime.parse(text) } + CodeBlock.of($$"$T.parse($S)", ClassName.get(OffsetDateTime::class.java), text) + } } - buildMethod.addStatement("return dag") - builderClass.addMethod(buildMethod.build()) - return builderClass.build() + private fun parseTemporal( + field: Field, + text: String, + parse: () -> Any, + ) { + try { + parse() + } catch (e: DateTimeParseException) { + throw IllegalArgumentException("Annotation attribute '${field.attribute}' is not valid ISO-8601: '$text'") + } } private fun buildTask( - name: String, - inner: ExecutableElement, + decl: TaskDeclaration, parent: TypeElement, ): TypeSpec { val executeSpec = @@ -147,8 +392,8 @@ class BuilderProcessor : AbstractProcessor() { .addParameter(CLIENT_TYPE, "client") .addException(Exception::class.java) - val dataParams = collectDataParams(inner) - val dataByName = dataParams.associateBy { it.name } + val inner = decl.method + val dataByName = decl.dataParams.associateBy { it.name } val innerArgs = with(processingEnv) { inner.parameters.joinToString { param -> @@ -161,16 +406,24 @@ class BuilderProcessor : AbstractProcessor() { } } - dataParams.forEach { param -> + decl.dataParams.forEach { param -> val paramType = TypeName.get(param.type) val fields = param.bundleFields if (fields == null) { executeSpec.addStatement($$"$T $L = $L", paramType, param.name, positionalAccess(param)) } else { - executeSpec.addStatement($$"$T $L = new $T()", paramType, param.name, paramType) + // Runtime bindings bind the bundle's fields by wire name; the + // Java-wired fallback decodes the bundle wholesale from its single + // wired input. + executeSpec.addStatement($$"$T $L", paramType, param.name) + executeSpec.beginControlFlow($$"if ($T.hasRuntimeBindings(client))", ARG_VALUES_TYPE) + executeSpec.addStatement($$"$L = new $T()", param.name, paramType) fields.forEach { field -> executeSpec.addStatement($$"$L.$L = $L", param.name, field.fieldName, namedAccess(field)) } + executeSpec.nextControlFlow("else") + executeSpec.addStatement($$"$L = $L", param.name, positionalAccess(param)) + executeSpec.endControlFlow() } } @@ -188,7 +441,7 @@ class BuilderProcessor : AbstractProcessor() { } return TypeSpec - .classBuilder(name) + .classBuilder(decl.className) .addSuperinterface(Task::class.java) .addModifiers(Modifier.PUBLIC, Modifier.FINAL, Modifier.STATIC) .addMethod(executeSpec.build()) @@ -268,6 +521,15 @@ class BuilderProcessor : AbstractProcessor() { } } +/** One [Builder.Task]-annotated method with its resolved id and data parameters. */ +private class TaskDeclaration( + val method: ExecutableElement, + val id: String, + val dataParams: List, +) { + val className: String = method.simpleName.toString().replaceFirstChar(Char::uppercase) +} + /** * One data parameter of a task method, positioned among its peers. * [bundleFields] is non-null for a [TaskInput] bundle parameter. @@ -292,6 +554,24 @@ private val CLIENT_TYPE = ClassName.get(Client::class.java) private val CONTEXT_TYPE = ClassName.get(Context::class.java) private val TASK_INPUT_TYPE = ClassName.get(TaskInput::class.java) private val ARG_VALUES_TYPE = ClassName.get(ArgValues::class.java) +private val REFS_TYPE = ClassName.get(Refs::class.java) +private val IN_TYPE = ClassName.get(In::class.java) +private val TASK_HANDLE_TYPE = ClassName.get(TaskRef::class.java) + +private const val DAG_ANNOTATION = "org.apache.airflow.sdk.Builder.Dag" +private const val TASK_ANNOTATION = "org.apache.airflow.sdk.Builder.Task" + +private val DAG_STRUCTURAL_ATTRIBUTES = setOf("id", "to") +private val TASK_STRUCTURAL_ATTRIBUTES = setOf("id") + +private val NUMERIC_KINDS = + setOf(TypeKind.BYTE, TypeKind.SHORT, TypeKind.INT, TypeKind.LONG, TypeKind.FLOAT, TypeKind.DOUBLE) + +private val BOXED_NUMERICS: Set = + setOf(TypeName.BYTE, TypeName.SHORT, TypeName.INT, TypeName.LONG, TypeName.FLOAT, TypeName.DOUBLE) + .mapTo(mutableSetOf()) { it.box() } + +private fun TypeName.boxIfPossible(): TypeName = if (this == TypeName.VOID || isPrimitive) box() else this private fun ProcessingEnvironment.isType( t: TypeMirror, 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 679aeea4338dd..8280955f0a910 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 @@ -40,8 +40,8 @@ private fun JavaFileObjectSubject.hasSourceEquivalentTo( class BuilderTest { @Test - @DisplayName("generate builder for dag class") - fun generateBuilderForDagClass() { + @DisplayName("generate builder and task-reference twins for dag class") + fun generateBuilderAndRefForDagClass() { val compilation = compile( """ @@ -50,6 +50,7 @@ class BuilderTest { import org.apache.airflow.sdk.Builder; import org.apache.airflow.sdk.Client; import org.apache.airflow.sdk.Context; + import org.apache.airflow.sdk.Wiring; @Builder.Dag public class TestExample { @@ -58,13 +59,19 @@ class BuilderTest { @Builder.Task public int t2(Client client) { - return (Integer) client.getXCom("t0"); + return 7; } @Builder.Task public void t3(Context ctx, int value) { System.out.println(String.format("%s %s", ctx.ti, value)); } + + @Wiring + static void depends(TestExampleRef f) { + f.t1(); + f.t3(f.t2()); + } } """, ) @@ -80,21 +87,30 @@ class BuilderTest { import java.lang.Exception; import java.lang.Integer; import java.lang.Override; + import java.lang.String; + import java.util.List; import org.apache.airflow.sdk.Client; import org.apache.airflow.sdk.Context; import org.apache.airflow.sdk.DagDef; import org.apache.airflow.sdk.Task; - import org.apache.airflow.sdk.TaskDef; import org.apache.airflow.sdk.internal.ArgValues; + import org.apache.airflow.sdk.internal.Refs; public final class TestExampleBuilder { + public static final String DAG_ID = "TestExample"; + + public static DagDef dag() { + var dag = new DagDef(DAG_ID); + return dag; + } + 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)); + var dag = dag(); + TestExample.depends(new TestExampleRef(dag)); + Refs.requireRegistered(dag, List.of("t1", "t2", "t3")); return dag; } + public static final class T1 implements Task { @Override public void execute(Context context, Client client) throws Exception { @@ -117,6 +133,44 @@ class BuilderTest { } """, ) + assertThat(compilation) + .generatedSourceFile("org.apache.airflow.example.TestExampleRef") + .hasSourceEquivalentTo( + "org.apache.airflow.example.TestExampleRef", + """ + package org.apache.airflow.example; + + import java.lang.Integer; + import java.lang.Number; + import java.lang.Void; + import java.util.List; + import org.apache.airflow.sdk.DagDef; + import org.apache.airflow.sdk.In; + import org.apache.airflow.sdk.TaskDef; + import org.apache.airflow.sdk.TaskRef; + import org.apache.airflow.sdk.internal.Refs; + + public final class TestExampleRef { + private final DagDef dag; + + public TestExampleRef(DagDef dag) { + this.dag = dag; + } + + public TaskRef t1() { + return Refs.register(dag, new TaskDef("t1", TestExampleBuilder.T1.class), List.of()); + } + + public TaskRef t2() { + return Refs.register(dag, new TaskDef("t2", TestExampleBuilder.T2.class), List.of()); + } + + public TaskRef t3(In value) { + return Refs.register(dag, new TaskDef("t3", TestExampleBuilder.T3.class), List.of(value)); + } + } + """, + ) } @Test @@ -158,11 +212,19 @@ class BuilderTest { import org.apache.airflow.sdk.internal.ArgValues; public final class TestExampleBuilder { + public static final String DAG_ID = "TestExample"; + + public static DagDef dag() { + var dag = new DagDef(DAG_ID); + return dag; + } + public static DagDef build() { - var dag = new DagDef("TestExample"); + var dag = dag(); dag.addTask(new TaskDef("t", T.class)); return dag; } + public static final class T implements Task { @Override public void execute(Context context, Client client) throws Exception { @@ -219,11 +281,19 @@ class BuilderTest { import org.apache.airflow.sdk.internal.ArgValues; public final class TestExampleBuilder { + public static final String DAG_ID = "TestExample"; + + public static DagDef dag() { + var dag = new DagDef(DAG_ID); + return dag; + } + public static DagDef build() { - var dag = new DagDef("TestExample"); + var dag = dag(); dag.addTask(new TaskDef("t", T.class)); return dag; } + public static final class T implements Task { @Override public void execute(Context context, Client client) throws Exception { @@ -241,7 +311,193 @@ class BuilderTest { } @Test - @DisplayName("bind input-bundle fields by wire name") + @DisplayName("type twin inputs by declared parameter type") + fun generateRefTypesTwinInputs() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import java.util.List; + import org.apache.airflow.sdk.Builder; + import org.apache.airflow.sdk.Wiring; + @Builder.Dag + public class TestExample { + @Builder.Task + public String ps() { return "x"; } + + @Builder.Task + public void pv() {} + + @Builder.Task + public List pl() { return null; } + + @Builder.Task + public long pn() { return 1L; } + + @Builder.Task + public void t(String text, Object anything, List items, Integer boxed) {} + + @Wiring + static void depends(TestExampleRef f) { + f.t(f.ps(), f.pv(), f.pl(), f.pn()); + } + } + """, + ) + + assertThat(compilation).succeeded() + assertThat(compilation) + .generatedSourceFile("org.apache.airflow.example.TestExampleRef") + .hasSourceEquivalentTo( + "org.apache.airflow.example.TestExampleRef", + """ + package org.apache.airflow.example; + + import java.lang.Long; + import java.lang.Number; + import java.lang.String; + import java.lang.Void; + import java.util.List; + import org.apache.airflow.sdk.DagDef; + import org.apache.airflow.sdk.In; + import org.apache.airflow.sdk.TaskDef; + import org.apache.airflow.sdk.TaskRef; + import org.apache.airflow.sdk.internal.Refs; + + public final class TestExampleRef { + private final DagDef dag; + + public TestExampleRef(DagDef dag) { + this.dag = dag; + } + + public TaskRef ps() { + return Refs.register(dag, new TaskDef("ps", TestExampleBuilder.Ps.class), List.of()); + } + + public TaskRef pv() { + return Refs.register(dag, new TaskDef("pv", TestExampleBuilder.Pv.class), List.of()); + } + + public TaskRef> pl() { + return Refs.register(dag, new TaskDef("pl", TestExampleBuilder.Pl.class), List.of()); + } + + public TaskRef pn() { + return Refs.register(dag, new TaskDef("pn", TestExampleBuilder.Pn.class), List.of()); + } + + public TaskRef t(In text, In anything, + In> items, In boxed) { + return Refs.register(dag, new TaskDef("t", TestExampleBuilder.T.class), List.of(text, anything, items, boxed)); + } + } + """, + ) + } + + @Test + @DisplayName("lower explicit annotation attributes into config calls") + fun generateBuilderLowersConfigAttributes() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + import org.apache.airflow.sdk.Wiring; + @Builder.Dag(id = "cfg", schedule = "@daily", tags = {"a", "b"}, catchup = true, + startDate = "2026-01-01T00:00:00Z") + public class TestExample { + @Builder.Task(retries = 2, queue = "q", retryDelay = "PT5M", retryExponentialBackoff = 1.5) + public void t1() {} + + @Wiring + static void depends(TestExampleRef f) { + f.t1(); + } + } + """, + ) + + assertThat(compilation).succeeded() + assertThat(compilation) + .generatedSourceFile("org.apache.airflow.example.TestExampleBuilder") + .hasSourceEquivalentTo( + "org.apache.airflow.example.TestExampleBuilder", + """ + package org.apache.airflow.example; + + import java.lang.Exception; + import java.lang.Override; + import java.lang.String; + import java.time.OffsetDateTime; + import java.util.List; + import org.apache.airflow.sdk.Client; + import org.apache.airflow.sdk.Context; + import org.apache.airflow.sdk.DagDef; + import org.apache.airflow.sdk.Task; + import org.apache.airflow.sdk.internal.Refs; + + public final class TestExampleBuilder { + public static final String DAG_ID = "cfg"; + + public static DagDef dag() { + var dag = new DagDef(DAG_ID); + dag.config("schedule", "@daily"); + dag.config("tags", List.of("a", "b")); + dag.config("catchup", true); + dag.config("start_date", OffsetDateTime.parse("2026-01-01T00:00:00Z")); + return dag; + } + + public static DagDef build() { + var dag = dag(); + TestExample.depends(new TestExampleRef(dag)); + Refs.requireRegistered(dag, List.of("t1")); + return dag; + } + + public static final class T1 implements Task { + @Override + public void execute(Context context, Client client) throws Exception { + new TestExample().t1(); + } + } + } + """, + ) + assertThat(compilation) + .generatedSourceFile("org.apache.airflow.example.TestExampleRef") + .hasSourceEquivalentTo( + "org.apache.airflow.example.TestExampleRef", + """ + package org.apache.airflow.example; + + import java.lang.Void; + import java.time.Duration; + import java.util.List; + import org.apache.airflow.sdk.DagDef; + import org.apache.airflow.sdk.TaskDef; + import org.apache.airflow.sdk.TaskRef; + import org.apache.airflow.sdk.internal.Refs; + + public final class TestExampleRef { + private final DagDef dag; + + public TestExampleRef(DagDef dag) { + this.dag = dag; + } + + public TaskRef t1() { + return Refs.register(dag, new TaskDef("t1", TestExampleBuilder.T1.class).config("retries", 2).config("queue", "q").config("retry_delay", Duration.parse("PT5M")).config("retry_exponential_backoff", 1.5), List.of()); + } + } + """, + ) + } + + @Test + @DisplayName("bind input-bundle fields by wire name with a wholesale wiring fallback") fun generateBuilderBindsInputBundleFields() { val compilation = compile( @@ -251,7 +507,9 @@ class BuilderTest { import org.apache.airflow.sdk.ArgName; import org.apache.airflow.sdk.Builder; import org.apache.airflow.sdk.Client; + import org.apache.airflow.sdk.In; import org.apache.airflow.sdk.TaskInput; + import org.apache.airflow.sdk.Wiring; @Builder.Dag public class TestExample { public static class ScoreInput implements TaskInput { @@ -262,6 +520,11 @@ class BuilderTest { @Builder.Task public double score(Client client, ScoreInput input) { return input.threshold; } + + @Wiring + static void depends(TestExampleRef f) { + f.score(In.value(new ScoreInput())); + } } """, ) @@ -283,28 +546,70 @@ class BuilderTest { import org.apache.airflow.sdk.Context; import org.apache.airflow.sdk.DagDef; import org.apache.airflow.sdk.Task; - import org.apache.airflow.sdk.TaskDef; import org.apache.airflow.sdk.internal.ArgValues; + import org.apache.airflow.sdk.internal.Refs; public final class TestExampleBuilder { + public static final String DAG_ID = "TestExample"; + + public static DagDef dag() { + var dag = new DagDef(DAG_ID); + return dag; + } + public static DagDef build() { - var dag = new DagDef("TestExample"); - dag.addTask(new TaskDef("score", Score.class)); + var dag = dag(); + TestExample.depends(new TestExampleRef(dag)); + Refs.requireRegistered(dag, List.of("score")); return dag; } + public static final class Score implements Task { @Override public void execute(Context context, Client client) throws Exception { - TestExample.ScoreInput input = new TestExample.ScoreInput(); - input.region = ArgValues.optionalNamed(client, "region_code", String.class); - input.threshold = ArgValues.requiredNamed(client, "threshold", Double.class, "threshold"); - input.tags = (List) ArgValues.optionalNamed(client, "tags", List.class); + TestExample.ScoreInput input; + if (ArgValues.hasRuntimeBindings(client)) { + input = new TestExample.ScoreInput(); + input.region = ArgValues.optionalNamed(client, "region_code", String.class); + input.threshold = ArgValues.requiredNamed(client, "threshold", Double.class, "threshold"); + input.tags = (List) ArgValues.optionalNamed(client, "tags", List.class); + } else { + input = ArgValues.optionalInput(context, client, 0, TestExample.ScoreInput.class); + } client.setXCom(new TestExample().score(client, input)); } } } """, ) + assertThat(compilation) + .generatedSourceFile("org.apache.airflow.example.TestExampleRef") + .hasSourceEquivalentTo( + "org.apache.airflow.example.TestExampleRef", + """ + package org.apache.airflow.example; + + import java.lang.Double; + import java.util.List; + import org.apache.airflow.sdk.DagDef; + import org.apache.airflow.sdk.In; + import org.apache.airflow.sdk.TaskDef; + import org.apache.airflow.sdk.TaskRef; + import org.apache.airflow.sdk.internal.Refs; + + public final class TestExampleRef { + private final DagDef dag; + + public TestExampleRef(DagDef dag) { + this.dag = dag; + } + + public TaskRef score(In input) { + return Refs.register(dag, new TaskDef("score", TestExampleBuilder.Score.class), List.of(input)); + } + } + """, + ) } @Test @@ -430,8 +735,13 @@ class BuilderTest { "org.apache.airflow.example.TestExampleBuilder", """ package org.apache.airflow.example; + import java.lang.String; import org.apache.airflow.sdk.DagDef; - public final class TestExampleBuilder { public static DagDef build() { var dag = new DagDef("foo"); return dag; } } + public final class TestExampleBuilder { + public static final String DAG_ID = "foo"; + public static DagDef dag() { var dag = new DagDef(DAG_ID); return dag; } + public static DagDef build() { var dag = dag(); return dag; } + } """, ) } @@ -453,8 +763,13 @@ class BuilderTest { "org.apache.airflow.example.Foo", """ package org.apache.airflow.example; + import java.lang.String; import org.apache.airflow.sdk.DagDef; - public final class Foo { public static DagDef build() { var dag = new DagDef("TestExample"); return dag; } } + public final class Foo { + public static final String DAG_ID = "TestExample"; + public static DagDef dag() { var dag = new DagDef(DAG_ID); return dag; } + public static DagDef build() { var dag = dag(); return dag; } + } """, ) } @@ -467,8 +782,16 @@ class BuilderTest { """ package org.apache.airflow.example; import org.apache.airflow.sdk.Builder; + import org.apache.airflow.sdk.Wiring; @Builder.Dag - public class TestExample { @Builder.Task(id = "foo") public void t1() {} } + public class TestExample { + @Builder.Task(id = "foo") public void t1() {} + + @Wiring + static void depends(TestExampleRef f) { + f.t1(); + } + } """, ) @@ -480,15 +803,20 @@ class BuilderTest { package org.apache.airflow.example; import java.lang.Exception; import java.lang.Override; + import java.lang.String; + import java.util.List; import org.apache.airflow.sdk.Client; import org.apache.airflow.sdk.Context; import org.apache.airflow.sdk.DagDef; import org.apache.airflow.sdk.Task; - import org.apache.airflow.sdk.TaskDef; + import org.apache.airflow.sdk.internal.Refs; public final class TestExampleBuilder { + public static final String DAG_ID = "TestExample"; + public static DagDef dag() { var dag = new DagDef(DAG_ID); return dag; } public static DagDef build() { - var dag = new DagDef("TestExample"); - dag.addTask(new TaskDef("foo", T1.class)); + var dag = dag(); + TestExample.depends(new TestExampleRef(dag)); + Refs.requireRegistered(dag, List.of("foo")); return dag; } public static final class T1 implements Task { @@ -497,6 +825,130 @@ class BuilderTest { } """, ) + assertThat(compilation) + .generatedSourceFile("org.apache.airflow.example.TestExampleRef") + .hasSourceEquivalentTo( + "org.apache.airflow.example.TestExampleRef", + """ + package org.apache.airflow.example; + import java.lang.Void; + import java.util.List; + import org.apache.airflow.sdk.DagDef; + import org.apache.airflow.sdk.TaskDef; + import org.apache.airflow.sdk.TaskRef; + import org.apache.airflow.sdk.internal.Refs; + public final class TestExampleRef { + private final DagDef dag; + public TestExampleRef(DagDef dag) { this.dag = dag; } + public TaskRef t1() { + return Refs.register(dag, new TaskDef("foo", TestExampleBuilder.T1.class), List.of()); + } + } + """, + ) + } + + @Test + @DisplayName("reject wiring that feeds an incompatible upstream type") + fun rejectIncompatibleWiring() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + import org.apache.airflow.sdk.Wiring; + @Builder.Dag + public class TestExample { + @Builder.Task + public String ps() { return "x"; } + + @Builder.Task + public void t(int v) {} + + @Wiring + static void depends(TestExampleRef f) { + f.t(f.ps()); + } + } + """, + ) + assertThat(compilation).failed() + assertThat(compilation).hadErrorContaining("incompatible types") + } + + @Test + @DisplayName("reject more than one wiring method") + fun rejectMultipleWiringMethods() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + import org.apache.airflow.sdk.Wiring; + @Builder.Dag + public class TestExample { + @Builder.Task public void t1() {} + + @Wiring + static void one(TestExampleRef f) { f.t1(); } + + @Wiring + static void two(TestExampleRef f) {} + } + """, + ) + assertThat(compilation).failed() + assertThat(compilation).hadErrorContaining( + "Dag class TestExample declares more than one @Wiring method: one, two", + ) + } + + @Test + @DisplayName("reject a non-static wiring method") + fun rejectNonStaticWiring() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + import org.apache.airflow.sdk.Wiring; + @Builder.Dag + public class TestExample { + @Builder.Task public void t1() {} + + @Wiring + void depends(TestExampleRef f) { f.t1(); } + } + """, + ) + assertThat(compilation).failed() + assertThat(compilation).hadErrorContaining( + "@Wiring method 'depends' must be static and non-private", + ) + } + + @Test + @DisplayName("reject a wiring method with the wrong shape") + fun rejectWrongShapedWiring() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + import org.apache.airflow.sdk.Wiring; + @Builder.Dag + public class TestExample { + @Builder.Task public void t1() {} + + @Wiring + static void depends(TestExampleRef f, int extra) { f.t1(); } + } + """, + ) + assertThat(compilation).failed() + assertThat(compilation).hadErrorContaining( + "@Wiring method 'depends' must be void and take the generated TestExampleRef as its only parameter", + ) } @Test @@ -516,4 +968,46 @@ class BuilderTest { "Cannot create task from vararg function t1", ) } + + @Test + @DisplayName("reject duplicate task ids") + fun rejectDuplicateTaskIds() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + @Builder.Dag + public class TestExample { + @Builder.Task(id = "x") + public void t1() {} + + @Builder.Task(id = "x") + public void t2() {} + } + """, + ) + assertThat(compilation).failed() + assertThat(compilation).hadErrorContaining("Tasks in Dag have duplicate ID: x") + } + + @Test + @DisplayName("reject a duration attribute that is not ISO-8601") + fun rejectInvalidDurationAttribute() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + @Builder.Dag + public class TestExample { + @Builder.Task(retryDelay = "5 minutes") public void t() {} + } + """, + ) + assertThat(compilation).failed() + assertThat(compilation).hadErrorContaining( + "Annotation attribute 'retryDelay' is not valid ISO-8601: '5 minutes'", + ) + } } diff --git a/java-sdk/sdk/build.gradle.kts b/java-sdk/sdk/build.gradle.kts index 79dd9340d2147..f1852a78c1dc4 100644 --- a/java-sdk/sdk/build.gradle.kts +++ b/java-sdk/sdk/build.gradle.kts @@ -21,6 +21,7 @@ import java.io.File import java.net.URI import java.nio.file.Files import java.nio.file.StandardCopyOption +import java.time.Duration val airflowSupervisorSchemaVersion: String by project @@ -41,6 +42,8 @@ val pointersDir = layout.buildDirectory.dir("schema-pointers/main") val jsonSchemaPackage = "org.apache.airflow.sdk.execution.comm" val schemaModelsDir = layout.buildDirectory.dir("generate-resources/main/src/main/java") val discriminatorDir = layout.buildDirectory.dir("generated-resources/main/src/main/kotlin") +val dagSchemaInput = layout.projectDirectory.file("schema/dag-schema.json") +val dagDslDir = layout.buildDirectory.dir("generated-resources/dsl/src/main/kotlin") dependencies { compileOnly("com.github.spotbugs:spotbugs-annotations:4.9.8") @@ -210,6 +213,410 @@ abstract class SyncSupervisorSchemaTask : DefaultTask() { } } +// Keep the vendored Dag serialization schema in sync with the monorepo copy. +// The vendored file makes standalone (source-release) builds work; in-repo +// builds refresh it from airflow-core, and a prek hook guards against drift. +abstract class SyncDagSchemaTask : DefaultTask() { + @get:Internal + abstract val sourceFile: RegularFileProperty + + @get:Internal + abstract val targetFile: RegularFileProperty + + @TaskAction + fun sync() { + val src = sourceFile.get().asFile + if (!src.exists()) { + logger.lifecycle("Monorepo serialization schema not present; keeping vendored dag-schema.json.") + return + } + val dst = targetFile.get().asFile + if (dst.exists() && dst.readText() == src.readText()) { + logger.lifecycle("Vendored dag-schema.json is up-to-date.") + return + } + logger.lifecycle("Refreshing vendored dag-schema.json from ${src.path}") + src.copyTo(dst, overwrite = true) + } +} + +// Generate the Dag-authoring DSL surface from the Dag serialization schema: +// +// - org.apache.airflow.sdk.Builder and its nested Dag / Task annotations, +// whose configuration attributes mirror the scalar keys of the schema's +// "dag" and "operator" definitions (the annotation processor lowers +// explicitly-set attributes into DagDef.config / TaskDef.config calls), and +// - org.apache.airflow.sdk.internal.SchemaFields, the key -> type table that +// DagDef.config / TaskDef.config validate against at registration time. +// +// Field selection mirrors the Go SDK's TaskSpec generator: scalar properties +// only (string/integer/number/boolean plus timedelta/datetime refs), +// serializer-owned keys skipped ("_"-prefixed, schema-required, "has_on_" +// callbacks), and a documented exclusion list for Python-only concerns. An +// exclusion entry that stops matching an eligible key fails generation, so +// the list cannot go stale. +abstract class GenerateDagDslTask : DefaultTask() { + @get:InputFile + abstract val schemaFile: RegularFileProperty + + @get:OutputDirectory + abstract val targetDirectory: DirectoryProperty + + private data class DslField( + val key: String, + val attribute: String, + val fieldType: String, + val attrType: String, + val attrDefault: String, + val defaultJson: String?, + val doc: String, + ) + + private fun camelCase(key: String): String = + key + .split('_') + .filter { it.isNotEmpty() } + .mapIndexed { i, seg -> + if (i == 0) seg else seg.replaceFirstChar(Char::uppercase) + }.joinToString("") + + private fun quote(s: String): String = "\"" + s.replace("\\", "\\\\").replace("\"", "\\\"") + "\"" + + private fun resolveField( + key: String, + prop: com.fasterxml.jackson.databind.JsonNode, + typeOverride: String?, + ): DslField? { + val ref = prop.path("\$ref").asText("").substringAfterLast('/') + val schemaType = + when { + ref == "timedelta" -> "timedelta" + ref == "datetime" -> "datetime" + ref.isNotEmpty() -> return null + prop.path("type").isTextual -> prop.path("type").asText() + else -> return null + } + val default = prop.path("default") + val defaultJson = if (default.isMissingNode || default.isNull) null else default.toString() + val attribute = camelCase(key) + return when (schemaType) { + "string" -> + DslField( + key, + attribute, + "STRING", + "String", + quote(default.asText("")), + defaultJson, + "Schema key `$key`.", + ) + "boolean" -> + DslField( + key, + attribute, + "BOOLEAN", + "Boolean", + default.asBoolean(false).toString(), + defaultJson, + "Schema key `$key`.", + ) + "integer", "number" -> + if (typeOverride == "Double") { + DslField( + key, + attribute, + "NUMBER", + "Double", + if (defaultJson != null) default.asDouble().toString() else "-1.0", + defaultJson, + "Schema key `$key`.", + ) + } else { + DslField( + key, + attribute, + "INTEGER", + "Int", + if (defaultJson != null) default.asInt().toString() else "-1", + defaultJson, + "Schema key `$key`." + if (defaultJson == null) " Negative means unset." else "", + ) + } + "timedelta" -> + DslField( + key, + attribute, + "TIMEDELTA", + "String", + if (defaultJson != null) { + quote(Duration.ofSeconds(default.asLong()).toString()) + } else { + quote("") + }, + defaultJson, + "Schema key `$key`; an ISO-8601 duration such as `\"PT5M\"`. Empty means unset.", + ) + "datetime" -> + DslField( + key, + attribute, + "DATETIME", + "String", + quote(""), + defaultJson, + "Schema key `$key`; an ISO-8601 date-time such as `\"2026-01-01T00:00:00Z\"`. Empty means unset.", + ) + "array" -> + if (prop.path("items").path("type").asText("") == "string" || key == "tags") { + DslField( + key, + attribute, + "STRING_ARRAY", + "Array", + "[]", + defaultJson, + "Schema key `$key`.", + ) + } else { + null + } + else -> null + } + } + + @TaskAction + fun generate() { + // Python-only "operator" keys deliberately not exposed, mirroring the + // Go SDK's TaskSpec generator exclusion list. + val excludedTaskKeys = + setOf( + "doc", + "doc_json", + "doc_yaml", + "doc_rst", + "allow_nested_operators", + "multiple_outputs", + "start_from_trigger", + "is_setup", + "is_teardown", + "on_failure_fail_dagrun", + ) + // Dag-level keys exposed for configuration, mirroring the Go SDK's + // hand-curated DagSpec field list. "schedule" is virtual: the schema + // models it as the serializer-owned "timetable" object. + val dagAllowlist = + listOf( + "description", + "dag_display_name", + "doc_md", + "start_date", + "end_date", + "dagrun_timeout", + "tags", + "max_active_tasks", + "max_active_runs", + "max_consecutive_failed_dag_runs", + "catchup", + "fail_fast", + "render_template_as_native_obj", + "disable_bundle_versioning", + "is_paused_upon_creation", + ) + + val root = + com.fasterxml.jackson.databind + .ObjectMapper() + .readTree(schemaFile.get().asFile) + val dagProps = root.path("definitions").path("dag").path("properties") + val operator = root.path("definitions").path("operator") + val operatorRequired = operator.path("required").map { it.asText() }.toSet() + + val dagFields = + buildList { + add( + DslField( + "schedule", + "schedule", + "STRING", + "String", + quote(""), + null, + "`\"@once\"`, `\"@continuous\"`, a cron expression, or empty for no schedule.", + ), + ) + dagAllowlist.forEach { key -> + val prop = dagProps.path(key) + if (prop.isMissingNode) { + throw GradleException("Dag allowlist key '$key' is missing from the schema; update the allowlist") + } + add( + resolveField(key, prop, null) + ?: throw GradleException("Dag allowlist key '$key' is not a scalar the DSL can express"), + ) + } + } + + val excludedSeen = mutableSetOf() + val taskFields = + buildList { + operator.path("properties").fields().forEach { (key, prop) -> + val serializerOwned = + key.startsWith("_") || key in operatorRequired || key.startsWith("has_on_") + if (serializerOwned) return@forEach + if (key in excludedTaskKeys) { + excludedSeen += key + return@forEach + } + // retry_exponential_backoff is "number" with an integral + // default, but Python declares it float (a backoff + // multiplier), so the mechanical mapping would pick Int. + val override = if (key == "retry_exponential_backoff") "Double" else null + resolveField(key, prop, override)?.let { + if (it.fieldType != "STRING_ARRAY") add(it) + } + } + } + (excludedTaskKeys - excludedSeen).takeIf { it.isNotEmpty() }?.let { + throw GradleException("Excluded task keys match no eligible schema property; remove or fix: $it") + } + // "id"/"to" name the annotations' structural attributes, so a schema + // key camel-casing to either would silently shadow them. + (dagFields + taskFields).firstOrNull { it.attribute == "id" || it.attribute == "to" }?.let { + throw GradleException("Schema key '${it.key}' collides with a structural annotation attribute") + } + + val outDir = targetDirectory.get().asFile.also { it.deleteRecursively() } + + fun attrLines(fields: List): String = + fields.joinToString("\n") { f -> + " /** ${f.doc} */\n val ${f.attribute}: ${f.attrType} = ${f.attrDefault}," + } + + outDir.resolve("org/apache/airflow/sdk").apply { mkdirs() }.resolve("Builder.kt").writeText( + """ + |package org.apache.airflow.sdk + | + |// Generated from the Dag serialization schema (sdk/schema/dag-schema.json); do not edit by hand. + | + |/** + | * Container for the annotation-based Dag-authoring API. + | * + | * 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 + | * [Builder.Dag], plus a `Ref` twin class when the Dag class declares + | * a [Wiring] method. + | * + | * Example: + | * + | * ```java + | * @Builder.Dag(id = "my_pipeline", schedule = "@daily") + | * public class MyPipeline { + | * + | * @Builder.Task(id = "extract", retries = 2) + | * public long extract(Client client) { ... } + | * + | * @Builder.Task(id = "transform") + | * public long transform(Client client, long extracted) { ... } + | * + | * @Wiring + | * static void depends(MyPipelineRef f) { + | * f.transform(f.extract()); + | * } + | * } + | * ``` + | * + | * A task method's data parameters — everything other than the injected + | * [Client] and [Context] — receive, by position, the arguments the Python + | * `@task.stub` call site bound, falling back to the inputs the [Wiring] + | * method fed them. Keyword arguments bind by name instead through a single + | * [TaskInput] bundle parameter. + | * + | * The processor generates `MyPipelineBuilder.build()`, which returns a + | * fully wired-up [DagDef] ready to add to a [Bundle]. + | */ + |class Builder internal constructor() { + | /** + | * Annotation to automate a Dag-builder pattern. + | * + | * When applied on a class Foo, this generates a FooBuilder class with a + | * static build method to create the Dag structure automatically. + | * + | * Configuration attributes mirror the Dag serialization schema; only + | * attributes written explicitly at the use site are applied, so the + | * scheduler's own defaults win for everything left out. + | */ + | @Target(AnnotationTarget.CLASS) + | @MustBeDocumented + | annotation class Dag( + | /** Dag ID. Empty derives it from the annotated class's name. */ + | val id: String = "", + | /** Name of the generated builder class. Empty derives `Builder`. */ + | val to: String = "", + |${attrLines(dagFields)} + | ) + | + | /** + | * Annotation to automate task definition in a Dag-builder pattern. + | * + | * Configuration attributes mirror the Dag serialization schema; only + | * attributes written explicitly at the use site are applied. + | */ + | @Target(AnnotationTarget.FUNCTION) + | @MustBeDocumented + | annotation class Task( + | /** Task ID. Empty derives it from the annotated function's name. */ + | val id: String = "", + |${attrLines(taskFields)} + | ) + |} + | + """.trimMargin(), + ) + + fun tableLines(fields: List): String = + fields.joinToString("\n") { f -> + val defaultRepr = f.defaultJson?.let { quote(it) } ?: "null" + listOf( + " \"${f.key}\" to", + " Field(", + " \"${f.key}\",", + " \"${f.attribute}\",", + " FieldType.${f.fieldType},", + " $defaultRepr,", + " ),", + ).joinToString("\n") + } + + outDir.resolve("org/apache/airflow/sdk/internal").apply { mkdirs() }.resolve("SchemaFields.kt").writeText( + """ + |package org.apache.airflow.sdk.internal + | + |// Generated from the Dag serialization schema (sdk/schema/dag-schema.json); do not edit by hand. + | + |/** + | * Configuration keys accepted by `DagDef.config` and `TaskDef.config`, + | * keyed by Dag serialization schema property name. Public so that the + | * annotation processor can lower `@Builder.Dag` / `@Builder.Task` + | * attributes onto the same tables; not user-facing API. + | */ + |object SchemaFields { + | val DAG: Map = + | linkedMapOf( + |${tableLines(dagFields)} + | ) + | + | val TASK: Map = + | linkedMapOf( + |${tableLines(taskFields)} + | ) + |} + | + """.trimMargin(), + ) + } +} + val syncSupervisorSchema by tasks.registering(SyncSupervisorSchemaTask::class) { description = "Ensure the bundled Supervisor Schema is up-to-date with the Gradle property." schemaVersion = airflowSupervisorSchemaVersion @@ -232,6 +639,19 @@ tasks.register("generatePointers") { targetDirectory = pointersDir } +val syncDagSchema by tasks.registering(SyncDagSchemaTask::class) { + description = "Refresh the vendored Dag serialization schema from the monorepo copy when present." + sourceFile = layout.projectDirectory.file("../../airflow-core/src/airflow/serialization/schema.json") + targetFile = dagSchemaInput +} + +tasks.register("generateDagDsl") { + dependsOn(syncDagSchema) + description = "Generate the Builder.Dag/Builder.Task annotations and SchemaFields from the Dag serialization schema" + schemaFile = dagSchemaInput + targetDirectory = dagDslDir +} + val javadocJar by tasks.registering(Jar::class) { description = "Assembles Javadoc JAR from Dokka output" group = JavaBasePlugin.DOCUMENTATION_GROUP @@ -260,6 +680,7 @@ sourceSets { main { java.srcDir(tasks.named("generateJsonSchema2Pojo").map { schemaModelsDir }) kotlin.srcDir(tasks.named("generateDiscriminator").map { discriminatorDir }) + kotlin.srcDir(tasks.named("generateDagDsl").map { dagDslDir }) } } @@ -271,6 +692,11 @@ dokka { matchingRegex = """org\.apache\.airflow\.sdk\.execution.*""" suppress.set(true) } + // 'internal' is public only for the annotation processor's benefit. + perPackageOption { + matchingRegex = """org\.apache\.airflow\.sdk\.internal.*""" + suppress.set(true) + } } } @@ -288,15 +714,15 @@ tasks.named("compileKotlin") { } tasks.named("runKtlintCheckOverMainSourceSet") { - dependsOn("generateJsonSchema2Pojo", "generateDiscriminator") + dependsOn("generateJsonSchema2Pojo", "generateDiscriminator", "generateDagDsl") } tasks.matching { it.name.startsWith("dokkaGenerate") }.configureEach { - dependsOn("generateJsonSchema2Pojo", "generateDiscriminator") + dependsOn("generateJsonSchema2Pojo", "generateDiscriminator", "generateDagDsl") } tasks.withType { - dependsOn("generateJsonSchema2Pojo", "generateDiscriminator") + dependsOn("generateJsonSchema2Pojo", "generateDiscriminator", "generateDagDsl") manifest { attributes( "Airflow-Supervisor-Schema-Version" to airflowSupervisorSchemaVersion, diff --git a/java-sdk/sdk/schema/dag-schema.json b/java-sdk/sdk/schema/dag-schema.json new file mode 100644 index 0000000000000..b860ca5e1bf55 --- /dev/null +++ b/java-sdk/sdk/schema/dag-schema.json @@ -0,0 +1,496 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "$id": "https://airflow.apache.com/schemas/serialized-dags.json", + "definitions": { + "datetime": { + "description": "A date time, stored as fractional seconds since the epoch", + "type": "number" + }, + "timedelta": { + "type": "number", + "minimum": 0 + }, + "typed_timedelta": { + "type": "object", + "properties": { + "__type": { + "type": "string", + "const": "timedelta" + }, + "__var": { "$ref": "#/definitions/timedelta" } + }, + "required": [ + "__type", + "__var" + ], + "additionalProperties": false + }, + "typed_relativedelta": { + "type": "object", + "description": "A dateutil.relativedelta.relativedelta object", + "properties": { + "__type": { + "type": "string", + "const": "relativedelta" + }, + "__var": { + "type": "object", + "properties": { + "weekday": { + "type": "array", + "items": { "type": "integer" }, + "minItems": 1, + "maxItems": 2 + } + }, + "additionalProperties": { "type": "integer" } + } + } + }, + "timezone": { + "anyOf": [ + { "type": "string" }, + { "type": "integer" } + ] + }, + "asset_definition": { + "type": "object", + "properties": { + "uri": { "type": "string" }, + "name": { "type": "string" }, + "group": { "type": "string" }, + "extra": { + "anyOf": [ + {"type": "null"}, + { "$ref": "#/definitions/dict" } + ] + }, + "watchers": { + "type": "array", + "items": { "$ref": "#/definitions/trigger" } + } + }, + "required": [ "uri", "extra" ] + }, + "asset": { + "type": "object", + "properties": { + "uri": { "type": "string" }, + "extra": { + "anyOf": [ + {"type": "null"}, + { "$ref": "#/definitions/dict" } + ] + } + }, + "required": [ "uri", "extra" ] + }, + "typed_asset": { + "type": "object", + "properties": { + "__type": { + "type": "string", + "constant": "asset" + }, + "__var": { "$ref": "#/definitions/asset" } + }, + "required": [ + "__type", + "__var" + ], + "additionalProperties": false + }, + "typed_asset_cond": { + "type": "object", + "properties": { + "__type": { + "anyOf": [{ + "type": "string", + "constant": "asset_or" + }, + { + "type": "string", + "constant": "asset_and" + } + ] + }, + "__var": { + "type": "array", + "items": { + "anyOf": [ + {"$ref": "#/definitions/typed_asset"}, + { "$ref": "#/definitions/typed_asset_cond"} + ] + } + } + }, + "required": [ + "__type", + "__var" + ], + "additionalProperties": false + }, + "trigger": { + "type": "object", + "properties": { + "classpath": { "type": "string" }, + "kwargs": { "$ref": "#/definitions/dict" } + }, + "required": [ "classpath", "kwargs" ] + }, + "dict": { + "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}$" + }, + "extra_links": { + "type": "array", + "items": { + "type": "object", + "minProperties": 1, + "maxProperties": 1 + } + }, + "dag_dependencies": { + "type": "array", + "items": { + "type": "object" + } + }, + "dag": { + "type": "object", + "properties": { + "params": { "$ref": "#/definitions/params" }, + "dag_id": { "type": "string" }, + "tasks": { "$ref": "#/definitions/tasks" }, + "timezone": { "$ref": "#/definitions/timezone" }, + "owner_links": { "type": "object" }, + "timetable": { + "type": "object", + "properties": { + "type": { "type": "string" }, + "value": { "$ref": "#/definitions/dict" } + } + }, + "catchup": { "type": "boolean" }, + "allowed_run_types": { + "anyOf": [ + { "type": "array", "items": { "type": "string" } }, + { "type": "null" } + ] + }, + "fail_fast": { "type": "boolean", "default": false }, + "fileloc": { "type" : "string"}, + "relative_fileloc": { "type" : "string"}, + "bundle_name": { "anyOf": [{ "type": "null" }, { "type": "string" }] }, + "_processor_dags_folder": { + "anyOf": [ + { "type": "null" }, + { "type": "string" } + ] + }, + "dag_display_name": { "type" : "string"}, + "description": { "type" : "string"}, + "deadline": { + "anyOf": [ + { "$ref": "#/definitions/dict" }, + { + "type": "array", + "items": { "$ref": "#/definitions/dict" } + }, + { "type": "null" } + ] + }, + "_concurrency": { "type" : "number"}, + "max_active_tasks": { "type" : "number" }, + "max_active_runs": { "type" : "number" }, + "max_consecutive_failed_dag_runs": { "type" : "number" }, + "default_args": { "$ref": "#/definitions/dict" }, + "start_date": { "$ref": "#/definitions/datetime" }, + "end_date": { "$ref": "#/definitions/datetime" }, + "dagrun_timeout": { "$ref": "#/definitions/timedelta" }, + "doc_md": { "type" : "string"}, + "access_control": {"$ref": "#/definitions/dict" }, + "is_paused_upon_creation": { "type": "boolean" }, + "has_on_success_callback": { "type": "boolean", "default": false }, + "has_on_failure_callback": { "type": "boolean", "default": false }, + "render_template_as_native_obj": { "type": "boolean", "default": false }, + "tags": { "type": "array" }, + "task_group": {"anyOf": [ + { "type": "null" }, + { "$ref": "#/definitions/task_group" } + ]}, + "edge_info": { "$ref": "#/definitions/edge_info" }, + "dag_dependencies": { "$ref": "#/definitions/dag_dependencies" }, + "disable_bundle_versioning": {"type": "boolean" }, + "rerun_with_latest_version": {"type": ["boolean", "null"], "default": null} + }, + "required": [ + "dag_id", + "fileloc", + "tasks" + ], + "additionalProperties": false + }, + "tasks": { + "type": "array", + "additionalProperties": { "$ref": "#/definitions/operator" } + }, + "params": { + "type": "array", + "prefixItems": [ + { "type": "string" }, + { "$ref": "#/definitions/param" } + ], + "unevaluatedItems": false + }, + "param": { + "$comment": "A param for a dag / operator", + "type": "object", + "required": [ + "__class", + "default" + ], + "properties": { + "__class": { "type": "string" }, + "default": {}, + "description": {"anyOf": [{"type":"string"}, {"type":"null"}]}, + "schema": { "$ref": "#/definitions/dict" } + } + }, + "operator": { + "$comment": "A task/operator in a DAG", + "type": "object", + "required": [ + "task_type", + "_task_module", + "task_id", + "ui_color", + "ui_fgcolor", + "template_fields" + ], + "properties": { + "task_type": { "type": "string", "default": "BaseOperator"}, + "_task_module": { "type": "string" }, + "_operator_extra_links": { "$ref": "#/definitions/extra_links" }, + "task_id": { "type": "string" }, + "_task_display_name": { "type": "string" }, + "owner": { "type": "string", "default": "airflow" }, + "start_date": { "$ref": "#/definitions/datetime" }, + "end_date": { "$ref": "#/definitions/datetime" }, + "trigger_rule": { "type": "string", "default": "all_success" }, + "depends_on_past": { "type": "boolean", "default": false }, + "ignore_first_depends_on_past": { "type": "boolean", "default": false }, + "wait_for_past_depends_before_skipping": { "type": "boolean", "default": false }, + "wait_for_downstream": { "type": "boolean", "default": false }, + "retries": { "type": "number", "default": 0 }, + "queue": { "type": "string", "default": "default" }, + "pool": { "type": "string", "default": "default_pool" }, + "pool_slots": { "type": "number", "default": 1 }, + "execution_timeout": { "$ref": "#/definitions/timedelta" }, + "retry_delay": { "$ref": "#/definitions/timedelta", "default": 300.0 }, + "retry_exponential_backoff": { "type": "number", "default": 0 }, + "max_retry_delay": { "$ref": "#/definitions/timedelta" }, + "params": { "$ref": "#/definitions/params" }, + "priority_weight": { "type": "number", "default": 1 }, + "weight_rule": { "type": "string", "default": "downstream" }, + "executor": { "type": "string" }, + "executor_config": { "$ref": "#/definitions/dict" }, + "do_xcom_push": { "type": "boolean", "default": true }, + "email_on_failure": { "type": "boolean", "default": true }, + "email_on_retry": { "type": "boolean", "default": true }, + "ui_color": { "type": "string", "default": "#fff" }, + "ui_fgcolor": { "type": "string", "default": "#000" }, + "template_fields": { + "type": "array", + "items": { "type": "string" }, + "default": [] + }, + "template_ext": {"type": "array", "default": []}, + "template_fields_renderers": {"$ref": "#/definitions/dict", "default": {}}, + "downstream_task_ids": { + "type": "array", + "items": { "type": "string" }, + "default": [] + }, + "doc": { "type": "string" }, + "doc_md": { "type": "string" }, + "doc_json": { "type": "string" }, + "doc_yaml": { "type": "string" }, + "doc_rst": { "type": "string" }, + "_logger_name": { "type": "string" }, + "_needs_expansion": { "type": "boolean"}, + "_is_mapped": { "const": true, "$comment": "only present when True", "default": false }, + "_is_sensor": { "const": true, "$comment": "only present when True", "default": false }, + "partial_kwargs": { "type": "object" }, + "_disallow_kwargs_override": { "type": "boolean"}, + "_expand_input_attr": { "type": "string" }, + "map_index_template": { "type": "string" }, + "allow_nested_operators": { "type": "boolean", "default": true }, + "render_template_as_native_obj": { "anyOf": [{"type": "boolean"}, {"type": "null"}], "default": null }, + "inlets": {"type": "array", "default": []}, + "outlets": {"type": "array", "default": []}, + "has_on_execute_callback": {"type": "boolean", "default": false}, + "has_on_failure_callback": {"type": "boolean", "default": false}, + "has_on_skipped_callback": {"type": "boolean", "default": false}, + "has_on_success_callback": {"type": "boolean", "default": false}, + "has_on_retry_callback": {"type": "boolean", "default": false}, + "multiple_outputs": {"type": "boolean", "default": false}, + "start_from_trigger": {"type": "boolean", "default": false}, + "start_trigger_args": {"type": "object", "default": null}, + "is_setup": {"type": "boolean", "default": false}, + "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"}, + "_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"], + "partial_kwargs": ["expand_input", "_is_mapped"], + "_is_mapped": ["expand_input", "partial_kwargs"] + }, + "additionalProperties": true + }, + "task_group": { + "$comment": "A TaskGroup containing tasks", + "type": "object", + "required": [ + "_group_id", + "group_display_name", + "prefix_group_id", + "children", + "tooltip", + "ui_color", + "ui_fgcolor", + "upstream_group_ids", + "downstream_group_ids", + "upstream_task_ids", + "downstream_task_ids" + ], + "properties": { + "_group_id": {"anyOf": [{"type": "null"}, { "type": "string" }]}, + "group_display_name": {"type": "string" }, + "is_mapped": { "type": "boolean" }, + "prefix_group_id": { "type": "boolean" }, + "children": { "$ref": "#/definitions/dict" }, + "tooltip": { "type": "string" }, + "doc_md": { + "anyOf": [ + { "type": "string" }, + { "type": "null" } + ]}, + "ui_color": { "type": "string" }, + "ui_fgcolor": { "type": "string" }, + "upstream_group_ids": { + "type": "array", + "items": { "type": "string" } + }, + "downstream_group_ids": { + "type": "array", + "items": { "type": "string" } + }, + "upstream_task_ids": { + "type": "array", + "items": { "type": "string" } + }, + "downstream_task_ids": { + "type": "array", + "items": { "type": "string" } + } + }, + "additionalProperties": false + }, + "edge_info": { + "$comment": "Metadata about DAG edges", + "type": "object", + "additionalProperties": { + "type": "object", + "additionalProperties": { + "type": "object", + "properties": { + "label": { "type": "string" } + }, + "required": ["label"], + "additionalProperties": false + } + } + } + }, + + "type": "object", + "allOf": [ + { + "type": "object", + "properties": { + "__version": { + "type": "integer", + "exclusiveMinimum": 0 + }, + "dag": { "$ref": "#/definitions/dag" }, + "client_defaults": { + "type": "object", + "description": "SDK-specific default values that differ from schema defaults", + "properties": { + "tasks": { + "type": "object", + "description": "Task-level default overrides" + } + }, + "additionalProperties": false + } + }, + "additionalProperties": false, + "required": [ "__version", "dag" ] + } + ] +} 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 deleted file mode 100644 index 1e8ca63670838..0000000000000 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Builder.kt +++ /dev/null @@ -1,81 +0,0 @@ -/* - * 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 - -/** - * Container for the annotation-based Dag-authoring API. - * - * 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 [Builder.Dag]. - * - * Example: - * - * ```java - * @Builder.Dag(id = "my_pipeline") - * public class MyPipeline { - * - * @Builder.Task(id = "extract") - * public long extract(Client client) { ... } - * - * @Builder.Task(id = "transform") - * public long transform(Client client, long extracted) { ... } - * } - * ``` - * - * A task method's data parameters — everything other than the injected - * [Client] and [Context] — receive the arguments the Python `@task.stub` call - * site bound, by position. Keyword arguments bind by name instead through a - * single [TaskInput] bundle parameter. - * - * The processor generates `MyPipelineBuilder.build()`, which returns a - * fully wired-up [DagDef] ready to add to a [Bundle]. - */ -class Builder internal constructor() { - /** - * Annotation to automate a Dag-builder pattern. - * - * When applied on a class Foo, this generates a FooBuilder class with a - * static build method to create the Dag structure automatically. - * - * @param id Override the Dag ID. If empty or not provided, the annotated - * class's name is used by default. - * @param to Name of the Dag-builder class. If empty or not provided, use the - * annotated class name + "Builder". - */ - @Target(AnnotationTarget.CLASS) - @MustBeDocumented - annotation class Dag( - val id: String = "", - val to: String = "", - ) - - /** - * Annotation to automate task definition in a Dag-builder pattern. - * - * @param id Override the task ID. If empty or not provided, the annotated - * function's name is used by default. - */ - @Target(AnnotationTarget.FUNCTION) - @MustBeDocumented - annotation class Task( - val id: String = "", - ) -} 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 6cd549270a8a4..b44e957fc6394 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 @@ -26,12 +26,45 @@ package org.apache.airflow.sdk * [Server.serve] to start accepting task-execution requests. * * @property dags All registered Dags keyed by [DagDef.id]. - * @throws IllegalArgumentException if any two Dags share the same ID. + * @throws IllegalArgumentException if any two Dags share the same ID, if a + * task depends on an upstream that is not registered in its own Dag, or if + * the dependencies of a Dag contain a cycle. */ class Bundle( dags: Iterable, ) { internal val dags: Map = dags.associateByDagId() + + init { + for (dag in this.dags.values) { + for ((taskId, def) in dag.tasks) { + for (upstream in def.upstreams) { + require(dag.tasks[upstream.id] === upstream) { + "Task '$taskId' in Dag '${dag.id}' depends on task '${upstream.id}' " + + "that is not registered in the same Dag" + } + } + } + checkNoCycle(dag) + } + } +} + +// Ref-twin wiring cannot express a cycle, but TaskDef.dependsOn can. +private fun checkNoCycle(dag: DagDef) { + val visiting = mutableSetOf() + val done = mutableSetOf() + + fun visit(def: TaskDef) { + if (def.id in done) return + require(visiting.add(def.id)) { + "Task dependencies in Dag '${dag.id}' contain a cycle involving task '${def.id}'" + } + def.upstreams.forEach(::visit) + visiting -= def.id + done += def.id + } + dag.tasks.values.forEach(::visit) } private fun Iterable.associateByDagId(): Map { diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Context.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Context.kt index ece4d69b4f74b..c11b260845712 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Context.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Context.kt @@ -124,6 +124,9 @@ data class Context( @JvmField val dagRun: DagRun, @JvmField val ti: TaskInstance, ) { + /** Registration of the executing task; resolves wired data inputs. */ + internal var taskDef: TaskDef? = null + internal companion object { fun from(request: StartupDetails) = Context( diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/DagDef.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/DagDef.kt index a6bd50e949a1c..e7f6e260756c9 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/DagDef.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/DagDef.kt @@ -19,6 +19,8 @@ package org.apache.airflow.sdk +import org.apache.airflow.sdk.internal.SchemaFields +import org.apache.airflow.sdk.internal.checkConfigValue import kotlin.Throws /** @@ -31,9 +33,13 @@ import kotlin.Throws * class directly if you need to do low-level plumbing: * * ```java + * var extract = new TaskDef("extract", Extract.class).config("retries", 2); + * var load = new TaskDef("load", Load.class).dependsOn(extract); + * * var dag = new DagDef("java_etl") - * .addTask(new TaskDef("extract", Extract.class)) - * .addTask(new TaskDef("load", Load.class)); + * .config("schedule", "@daily") + * .addTask(extract) + * .addTask(load); * ``` * * @param id Dag identifier. Must contain only ASCII alphanumeric characters, @@ -45,13 +51,40 @@ class DagDef( val id: String, // TODO: charset check? ) { internal val tasks = linkedMapOf() + internal val dagConfig = linkedMapOf() + + /** + * Sets one Dag-level configuration value. + * + * Keys are the Dag serialization schema property names (for example + * `"schedule"`, `"description"`, `"tags"`, `"catchup"`); unknown keys and + * mismatched value types are rejected immediately, so mistakes surface at + * Dag-parse time. + * + * @param key Dag serialization schema property name. + * @param value Value matching the key's schema type. Durations take + * [java.time.Duration], date-times [java.time.OffsetDateTime] or + * [java.time.Instant], string arrays any `Iterable` of `String`. + * @return This Dag, for chaining. + * @throws IllegalArgumentException if the key is unknown or the value type + * does not match. + */ + fun config( + key: String, + value: Any?, + ): DagDef { + dagConfig[key] = checkConfigValue("Dag", SchemaFields.DAG, key, value) + return this + } /** * Registers a task with this 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. + * unique within a Dag. Upstream tasks referenced via [TaskDef.dependsOn] + * must be registered with the same Dag by the time it is added to a + * [Bundle]. * * @param task Task definition to register. * @return This Dag, for chaining. @@ -68,13 +101,35 @@ class DagDef( task.owner = this return this } + + /** + * Registers a task with this Dag, adding [upstreams] as its dependencies. + * + * Equivalent to `addTask(task.dependsOn(...))`; see [addTask] and + * [TaskDef.dependsOn]. + * + * @param task Task definition to register. + * @param upstreams Tasks that [task] depends on (its upstream tasks). + * @return This Dag, for chaining. + * @throws IllegalArgumentException if the task already belongs to a Dag or a + * task with the same ID is already registered. + */ + fun addTask( + task: TaskDef, + upstreams: List, + ): DagDef { + upstreams.forEach { task.dependsOn(it) } + return addTask(task) + } } /** - * One task definition: its ID and the class that implements it. + * One task definition: its ID, the class that implements it, its upstream + * dependencies, and its task-level configuration. * * ```java - * var extract = new TaskDef("extract", Extract.class); + * var extract = new TaskDef("extract", Extract.class).config("retries", 2); + * var load = new TaskDef("load", Load.class).dependsOn(extract); * ``` * * @param id Task identifier, unique within a [DagDef]. @@ -87,7 +142,48 @@ class TaskDef( val id: String, val definition: Class, ) { + internal val configValues = linkedMapOf() + internal val upstreams = linkedSetOf() + internal val inputs = mutableListOf>() internal var owner: DagDef? = null + + /** + * Sets one task-level configuration value. + * + * Keys are the Dag serialization schema property names (for example + * `"retries"`, `"queue"`, `"retry_delay"`); unknown keys and mismatched + * value types are rejected immediately, so mistakes surface at Dag-parse + * time. + * + * @param key Dag serialization schema property name, e.g. `"retries"`. + * @param value Value matching the key's schema type. Durations take + * [java.time.Duration], date-times [java.time.OffsetDateTime] or + * [java.time.Instant]. + * @return This task definition, for chaining. + * @throws IllegalArgumentException if the key is unknown or the value type + * does not match. + */ + fun config( + key: String, + value: Any?, + ): TaskDef { + configValues[key] = checkConfigValue("task", SchemaFields.TASK, key, value) + return this + } + + /** + * Declares that this task runs after [upstreams]. + * + * Referenced tasks must be registered with the same [DagDef] as this task by + * the time it is added to a [Bundle]. + * + * @param upstreams Tasks this task depends on. + * @return This task definition, for chaining. + */ + fun dependsOn(vararg upstreams: TaskDef): TaskDef { + this.upstreams += upstreams + return this + } } /** diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/In.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/In.kt new file mode 100644 index 0000000000000..9b0b9a232ebb8 --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/In.kt @@ -0,0 +1,64 @@ +/* + * 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 + +/** + * One data input of a task in a [Wiring] method: either the output of an + * upstream task (a [TaskRef] returned by another flow-twin call) or an + * inline literal created with [In.value]. + * + * ```java + * f.transform(f.extract()); // upstream output + * f.score(In.value(0.5)); // inline literal + * ``` + * + * @param T Declared type of the task parameter this input feeds. + */ +sealed class In { + companion object { + /** + * Wraps an inline literal as a task input. + * + * The value is delivered to the task parameter as-is; numeric values + * widen to the parameter's declared numeric type. + * + * @param value Literal to bind; may be null for nullable parameters. + */ + @JvmStatic + fun value(value: T?): In = LiteralIn(value) + } +} + +internal class LiteralIn( + internal val value: T?, +) : In() + +/** + * The output of a task registered by a flow-twin call in a [Wiring] method. + * + * Passing a handle to another twin call feeds this task's return value into + * that task's parameter and wires the dependency edge — the calls in the + * wiring method are the single way to declare dependencies. + * + * @param T Return type of the task this handle refers to. + */ +class TaskRef internal constructor( + internal val def: TaskDef, +) : In() diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Wiring.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Wiring.kt new file mode 100644 index 0000000000000..85f39b21bf88e --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Wiring.kt @@ -0,0 +1,49 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.airflow.sdk + +/** + * Marks the static method of a [Builder.Dag]-annotated class that wires its + * task graph, TaskFlow-style — the single place dependencies are declared. + * + * The method receives the generated twin class (`Ref`), whose + * methods mirror the [Builder.Task]-annotated methods with data parameters + * typed [In] and return values typed [TaskRef]. Calling a twin registers + * the task; passing one twin's return value into another feeds the upstream's + * output into the downstream's parameter and wires the dependency edge, all + * type-checked at compile time. The call graph is the task graph: + * + * ```java + * @Wiring + * static void depends(EtlPipelineRef f) { + * f.load(f.transform(f.extract())); + * } + * ``` + * + * Every [Builder.Task]-annotated method must be invoked exactly once; the + * generated `build()` fails at Dag-parse time otherwise. + * + * The wiring method is optional: a [Builder.Dag] class without one registers + * every task with no Java-side edges, which is the shape for stub-backed tasks + * whose graph is defined by a Python Dag file. + */ +@Target(AnnotationTarget.FUNCTION) +@MustBeDocumented +annotation class Wiring 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 1f1a01e674a58..5d54dada50528 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,11 +70,15 @@ internal object TaskRunner { request: StartupDetails, client: Client, ): Any { - val definition = - bundle.dags[request.ti.dagId]?.tasks[request.ti.taskId]?.definition + val taskDef = + bundle.dags[request.ti.dagId]?.tasks[request.ti.taskId] ?: return TaskResult.of(TaskState.State.REMOVED) return try { - definition.getDeclaredConstructor().newInstance().execute(Context.from(request), client) + val context = Context.from(request).also { it.taskDef = taskDef } + taskDef.definition + .getDeclaredConstructor() + .newInstance() + .execute(context, 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/main/kotlin/org/apache/airflow/sdk/internal/ArgValues.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/ArgValues.kt index cd56f38470668..e2d3ad4d9957f 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/ArgValues.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/ArgValues.kt @@ -25,18 +25,23 @@ import com.fasterxml.jackson.databind.ObjectMapper import com.fasterxml.jackson.databind.json.JsonMapper import org.apache.airflow.sdk.Client import org.apache.airflow.sdk.Context +import org.apache.airflow.sdk.In +import org.apache.airflow.sdk.LiteralIn import org.apache.airflow.sdk.MissingXComException +import org.apache.airflow.sdk.TaskRef import org.apache.airflow.sdk.execution.ArgBinding /** - * Resolves a task's data parameters from the arg bindings the supervisor - * delivered, and decodes their raw wire values into the declared parameter - * types. Public so that processor-generated task classes can call it; not - * user-facing API. + * Resolves a task's data parameters and decodes their raw wire values into the + * declared parameter types. Public so that processor-generated task classes + * can call it; not user-facing API. * - * The bindings come from the Python `@task.stub` call site, which is also the - * graph the scheduler ordered the run by. Flat data parameters resolve the - * binding at their position; input-bundle fields resolve bindings by name. + * Runtime arg bindings from the Python Dag file win over the Java-declared + * wiring: for a stub task the `@task.stub` call site is the graph the + * scheduler ordered the run by. Flat data parameters resolve the binding at + * their position; input-bundle fields resolve bindings by name. Only when the + * supervisor sent no bindings does resolution fall back to the inputs the + * `@Wiring` method recorded. */ object ArgValues { private val mapper: ObjectMapper = JsonMapper.builder().build().findAndRegisterModules() @@ -57,8 +62,8 @@ object ArgValues { type: Class, paramName: String, ): T { - val binding = bindingAt(context, client, position) - return decode(client.resolveBinding(binding), type) ?: throw missing(binding, paramName) + val resolved = resolveAt(context, client, position) + return decode(resolved.value, type) ?: throw resolved.missing(paramName) } /** @@ -74,7 +79,15 @@ object ArgValues { client: Client, position: Int, type: Class, - ): T? = decode(client.resolveBinding(bindingAt(context, client, position)), type) + ): T? = decode(resolveAt(context, client, position).value, type) + + /** + * Whether the supervisor delivered TaskFlow arg bindings for this run. + * Generated input-bundle code branches on this: bindings bind the bundle's + * fields by name, the wiring fallback decodes the bundle wholesale. + */ + @JvmStatic + fun hasRuntimeBindings(client: Client): Boolean = client.argBindings.isNotEmpty() /** * Resolves the runtime binding named [name] into [type] for an input-bundle @@ -117,19 +130,55 @@ object ArgValues { return decode(client.resolveBinding(binding), type) } - private fun bindingAt( + /** A resolved raw value plus the error to raise when it is null. */ + private class Resolved( + val value: Any?, + val missing: (String) -> MissingXComException, + ) + + private fun resolveAt( context: Context, client: Client, position: Int, - ): ArgBinding { + ): Resolved { val bindings = client.argBindings - check(position < bindings.size) { + if (bindings.isNotEmpty()) { + check(position < bindings.size) { + "Task '${context.ti.taskId}' declares a data parameter at position $position " + + "but the stub call bound only ${bindings.size} argument(s)" + } + val binding = bindings[position] + return Resolved(client.resolveBinding(binding)) { missing(binding, it) } + } + val input = inputAt(context, position) + return Resolved(resolveInput(input, client)) { missing(input, it) } + } + + private fun inputAt( + context: Context, + position: Int, + ): In<*> { + val inputs = + checkNotNull(context.taskDef?.inputs) { + "Task '${context.ti.taskId}' declares data parameters but has no wired inputs; " + + "register it through a @Wiring method" + } + check(position < inputs.size) { "Task '${context.ti.taskId}' declares a data parameter at position $position " + - "but the stub call bound only ${bindings.size} argument(s)" + "but only ${inputs.size} input(s) are wired" } - return bindings[position] + return inputs[position] } + private fun resolveInput( + input: In<*>, + client: Client, + ): Any? = + when (input) { + is TaskRef<*> -> client.getXCom(taskId = input.def.id) + is LiteralIn<*> -> input.value + } + private fun missing( binding: ArgBinding, target: String, @@ -145,6 +194,19 @@ object ArgValues { ) } + private fun missing( + input: In<*>, + target: String, + ): MissingXComException = + when (input) { + is TaskRef<*> -> MissingXComException(input.def.id, target) + is LiteralIn<*> -> + MissingXComException( + "'$target' has a primitive type but its wired literal input is null; " + + "declare a boxed type (e.g. Integer instead of int) to receive null.", + ) + } + internal fun decode( value: Any?, type: Class, diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/Fields.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/Fields.kt new file mode 100644 index 0000000000000..cc43db3186590 --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/Fields.kt @@ -0,0 +1,107 @@ +/* + * 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.internal + +import java.time.Duration +import java.time.Instant +import java.time.OffsetDateTime + +/** + * Value shape of one Dag serialization schema key. Public so that the + * annotation processor can lower `@Builder.Dag` / `@Builder.Task` attributes; + * not user-facing API. + */ +enum class FieldType { + STRING, + INTEGER, + NUMBER, + BOOLEAN, + STRING_ARRAY, + TIMEDELTA, + DATETIME, +} + +/** + * One configuration key from the Dag serialization schema. Public so that the + * annotation processor can lower `@Builder.Dag` / `@Builder.Task` attributes; + * not user-facing API. + * + * @property key Schema property name, e.g. `retry_delay`. + * @property attribute Annotation attribute name, e.g. `retryDelay`. + * @property type Accepted value shape. + * @property defaultJson Schema default as raw JSON, or `null` when the schema + * declares no default. + */ +class Field( + val key: String, + val attribute: String, + val type: FieldType, + val defaultJson: String?, +) + +/** + * Validates one `config(key, value)` call against a schema field table and + * returns the value to store. + * + * @throws IllegalArgumentException if the key is not a configurable schema + * key or the value does not match the key's type. + */ +internal fun checkConfigValue( + scope: String, + table: Map, + key: String, + value: Any?, +): Any { + val field = + requireNotNull(table[key]) { + "Unknown $scope config key: '$key'" + } + requireNotNull(value) { + "Value for $scope config key '$key' must not be null" + } + + fun mismatch(expected: String): Nothing = + throw IllegalArgumentException( + "Value for $scope config key '$key' must be $expected, got: ${value.javaClass.name}", + ) + return when (field.type) { + FieldType.STRING -> value as? String ?: mismatch("a String") + FieldType.BOOLEAN -> value as? Boolean ?: mismatch("a Boolean") + FieldType.NUMBER -> value as? Number ?: mismatch("a Number") + FieldType.INTEGER -> + when (value) { + is Byte, is Short, is Int, is Long -> value + else -> mismatch("an integral Number") + } + FieldType.TIMEDELTA -> value as? Duration ?: mismatch("a java.time.Duration") + FieldType.DATETIME -> + when (value) { + is OffsetDateTime -> value + is Instant -> value + else -> mismatch("a java.time.OffsetDateTime or java.time.Instant") + } + FieldType.STRING_ARRAY -> + when { + value is Iterable<*> && value.all { it is String } -> value.map { it as String } + value is Array<*> && value.all { it is String } -> value.map { it as String } + else -> mismatch("an Iterable of String") + } + } +} diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/Refs.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/Refs.kt new file mode 100644 index 0000000000000..97b809b3fb9d1 --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/Refs.kt @@ -0,0 +1,68 @@ +/* + * 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.internal + +import org.apache.airflow.sdk.DagDef +import org.apache.airflow.sdk.In +import org.apache.airflow.sdk.TaskDef +import org.apache.airflow.sdk.TaskRef + +/** + * Registration hooks called by processor-generated twin classes. Public + * so that generated code can call them; not user-facing API. + */ +object Refs { + /** + * Registers one task from a flow-twin call: records [inputs] as the task's + * data parameters in declaration order, wires a dependency edge for every + * upstream [TaskRef] among them, and adds the task to [dag]. + * + * @return The handle representing this task's output. + */ + @JvmStatic + fun register( + dag: DagDef, + def: TaskDef, + inputs: List>, + ): TaskRef { + inputs.filterIsInstance>().forEach { def.dependsOn(it.def) } + def.inputs += inputs + dag.addTask(def) + return TaskRef(def) + } + + /** + * Verifies that the user's `@Wiring` method registered every + * `@Builder.Task` method of the Dag class. + * + * @throws IllegalArgumentException naming the tasks the wiring missed. + */ + @JvmStatic + fun requireRegistered( + dag: DagDef, + taskIds: List, + ) { + val missing = taskIds.filterNot { it in dag.tasks } + require(missing.isEmpty()) { + "Wiring for Dag '${dag.id}' did not register task(s) ${missing.joinToString { "'$it'" }}: " + + "every @Builder.Task method must be invoked in the @Wiring method" + } + } +} 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 57050754e887e..42000ae496074 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 @@ -44,4 +44,79 @@ internal class BundleTest { Assertions.assertEquals("Dags in bundle have duplicate ID: dag", error.message) } + + @Test + @DisplayName("Should reject a task depending on an unregistered upstream") + fun shouldRejectUnregisteredUpstream() { + val missing = TaskDef("missing", NoOp::class.java) + val dag = DagDef("dag").addTask(TaskDef("t", NoOp::class.java).dependsOn(missing)) + + val error = + Assertions.assertThrows(IllegalArgumentException::class.java) { + Bundle(listOf(dag)) + } + + Assertions.assertEquals( + "Task 't' in Dag 'dag' depends on task 'missing' that is not registered in the same Dag", + error.message, + ) + } + + @Test + @DisplayName("Should reject a task depending on a task registered in another dag") + fun shouldRejectUpstreamFromAnotherDag() { + val foreign = TaskDef("u", NoOp::class.java) + val other = DagDef("other").addTask(foreign) + val dag = DagDef("dag").addTask(TaskDef("t", NoOp::class.java).dependsOn(foreign)) + + val error = + Assertions.assertThrows(IllegalArgumentException::class.java) { + Bundle(listOf(other, dag)) + } + + Assertions.assertEquals( + "Task 't' in Dag 'dag' depends on task 'u' that is not registered in the same Dag", + error.message, + ) + } + + @Test + @DisplayName("Should reject dependency cycles") + fun shouldRejectDependencyCycle() { + val a = TaskDef("a", NoOp::class.java) + val b = TaskDef("b", NoOp::class.java) + a.dependsOn(b) + b.dependsOn(a) + val dag = DagDef("dag").addTask(a).addTask(b) + + val error = + Assertions.assertThrows(IllegalArgumentException::class.java) { + Bundle(listOf(dag)) + } + + Assertions.assertEquals( + "Task dependencies in Dag 'dag' contain a cycle involving task 'a'", + error.message, + ) + } + + @Test + @DisplayName("Should accept a diamond-shaped dependency graph") + fun shouldAcceptDiamondGraph() { + val root = TaskDef("root", NoOp::class.java) + val left = TaskDef("left", NoOp::class.java).dependsOn(root) + val right = TaskDef("right", NoOp::class.java).dependsOn(root) + val join = TaskDef("join", NoOp::class.java).dependsOn(left, right) + val dag = DagDef("dag") + listOf(root, left, right, join).forEach(dag::addTask) + + Assertions.assertEquals(mapOf("dag" to dag), Bundle(listOf(dag)).dags) + } + + private class NoOp : Task { + override fun execute( + context: Context, + client: Client, + ) = Unit + } } diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ClientArgTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ClientArgTest.kt index e6b68629212b3..0a287330a94d5 100644 --- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ClientArgTest.kt +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ClientArgTest.kt @@ -27,6 +27,7 @@ import org.apache.airflow.sdk.execution.comm.TIRunContext import org.apache.airflow.sdk.execution.comm.VariableResult import org.apache.airflow.sdk.execution.comm.XComResult import org.apache.airflow.sdk.internal.ArgValues +import org.apache.airflow.sdk.internal.Refs import org.junit.jupiter.api.Assertions.assertEquals import org.junit.jupiter.api.Assertions.assertFalse import org.junit.jupiter.api.Assertions.assertNull @@ -91,12 +92,26 @@ private fun clientWith( return Client(startupDetails(argBindings), transport) to transport } +private class NoopClientArgTask : Task { + override fun execute( + context: Context, + client: Client, + ) = Unit +} + private fun taskContext(): Context = Context( dagRun = DagRun("d", "r", null, null, null, null, null, emptyMap()), ti = TaskInstance("d", "r", "t", null, 1), ) +/** A context whose task was Java-wired with the given inputs. */ +private fun contextWiredWith(inputs: List>): Context { + val def = TaskDef("t", NoopClientArgTask::class.java) + Refs.register(DagDef("d"), def, inputs) + return taskContext().also { it.taskDef = def } +} + internal class ClientArgTest { @Test @DisplayName("Should resolve a literal binding to its inline value, by position and by name") @@ -237,6 +252,28 @@ internal class ClientArgTest { assertEquals(listOf("upstream" to null), transport.pulls) } + @Test + @DisplayName("Should prefer the runtime binding at the position over the Java wiring") + fun shouldPreferRuntimeBinding() { + val context = contextWiredWith(listOf(In.value(9L))) + val (client, transport) = clientWith(listOf(mapOf("kind" to "literal", "name" to "value", "value" to 5L))) + + val resolved = ArgValues.requiredInput(context, client, 0, Integer::class.java, "value") + + assertEquals(5, resolved.toInt()) + assertEquals(emptyList>(), transport.pulls) + } + + @Test + @DisplayName("Should fall back to the Java wiring without runtime bindings") + fun shouldFallBackToWiring() { + val context = contextWiredWith(listOf(In.value(9L))) + val (client, _) = clientWith(null) + + assertFalse(ArgValues.hasRuntimeBindings(client)) + assertEquals(9, ArgValues.requiredInput(context, client, 0, Integer::class.java, "value").toInt()) + } + @Test @DisplayName("Should fail fast when the stub call bound fewer arguments than declared") fun shouldFailOnArityMismatch() { @@ -275,6 +312,7 @@ internal class ClientArgTest { xcoms = mapOf("upstream" to 0.5), ) + assertTrue(ArgValues.hasRuntimeBindings(client)) assertEquals("emea", ArgValues.optionalNamed(client, "region_code", String::class.java)) assertEquals(0.5, ArgValues.requiredNamed(client, "threshold", java.lang.Double::class.java, "threshold")) } 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 index f22c26432c30c..437ce873e48c6 100644 --- 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 @@ -22,6 +22,8 @@ package org.apache.airflow.sdk import org.junit.jupiter.api.Assertions import org.junit.jupiter.api.DisplayName import org.junit.jupiter.api.Test +import java.time.Duration +import java.time.OffsetDateTime internal class DagDefTest { private class NoOp : Task { @@ -83,4 +85,134 @@ internal class DagDefTest { Assertions.assertEquals("Task 'extract' already belongs to Dag 'dag'", error.message) } + + @Test + @DisplayName("Should store validated dag config values keyed by schema name") + fun shouldStoreDagConfigValues() { + val dag = + DagDef("dag") + .config("schedule", "@daily") + .config("description", "demo") + .config("catchup", true) + .config("max_active_runs", 3) + .config("dagrun_timeout", Duration.ofMinutes(5)) + .config("start_date", OffsetDateTime.parse("2026-01-01T00:00:00Z")) + .config("tags", listOf("a", "b")) + + Assertions.assertEquals( + mapOf( + "schedule" to "@daily", + "description" to "demo", + "catchup" to true, + "max_active_runs" to 3, + "dagrun_timeout" to Duration.ofMinutes(5), + "start_date" to OffsetDateTime.parse("2026-01-01T00:00:00Z"), + "tags" to listOf("a", "b"), + ), + dag.dagConfig, + ) + } + + @Test + @DisplayName("Should reject unknown dag config keys") + fun shouldRejectUnknownDagConfigKey() { + val error = + Assertions.assertThrows(IllegalArgumentException::class.java) { + DagDef("dag").config("scheduel", "@daily") + } + + Assertions.assertEquals("Unknown Dag config key: 'scheduel'", error.message) + } + + @Test + @DisplayName("Should reject dag config values of the wrong type") + fun shouldRejectMismatchedDagConfigValue() { + val error = + Assertions.assertThrows(IllegalArgumentException::class.java) { + DagDef("dag").config("catchup", "yes") + } + + Assertions.assertEquals( + "Value for Dag config key 'catchup' must be a Boolean, got: java.lang.String", + error.message, + ) + } + + @Test + @DisplayName("Should reject null dag config values") + fun shouldRejectNullDagConfigValue() { + val error = + Assertions.assertThrows(IllegalArgumentException::class.java) { + DagDef("dag").config("description", null) + } + + Assertions.assertEquals("Value for Dag config key 'description' must not be null", error.message) + } + + @Test + @DisplayName("Should reject non-integral values for integer dag config keys") + fun shouldRejectFractionalIntegerValue() { + val error = + Assertions.assertThrows(IllegalArgumentException::class.java) { + DagDef("dag").config("max_active_runs", 1.5) + } + + Assertions.assertEquals( + "Value for Dag config key 'max_active_runs' must be an integral Number, got: java.lang.Double", + error.message, + ) + } + + @Test + @DisplayName("Should store validated task config values on the task definition") + fun shouldStoreTaskConfigValues() { + val def = + TaskDef("extract", NoOp::class.java) + .config("retries", 2) + .config("queue", "q") + .config("retry_delay", Duration.ofMinutes(5)) + .config("retry_exponential_backoff", 1.5) + + Assertions.assertEquals( + mapOf( + "retries" to 2, + "queue" to "q", + "retry_delay" to Duration.ofMinutes(5), + "retry_exponential_backoff" to 1.5, + ), + def.configValues, + ) + } + + @Test + @DisplayName("Should reject unknown task config keys") + fun shouldRejectUnknownTaskConfigKey() { + val error = + Assertions.assertThrows(IllegalArgumentException::class.java) { + TaskDef("extract", NoOp::class.java).config("retrys", 1) + } + + Assertions.assertEquals("Unknown task config key: 'retrys'", error.message) + } + + @Test + @DisplayName("Should record upstream task definitions from dependsOn") + fun shouldRecordUpstreams() { + val extract = TaskDef("extract", NoOp::class.java) + val load = TaskDef("load", NoOp::class.java).dependsOn(extract) + DagDef("dag").addTask(extract).addTask(load) + + Assertions.assertEquals(emptySet(), extract.upstreams) + Assertions.assertEquals(setOf(extract), load.upstreams) + } + + @Test + @DisplayName("Should wire upstreams passed to the addTask overload") + fun shouldWireUpstreamsFromAddTaskOverload() { + val extract = TaskDef("extract", NoOp::class.java) + val load = TaskDef("load", NoOp::class.java) + DagDef("dag").addTask(extract).addTask(load, listOf(extract)) + + Assertions.assertEquals(setOf(extract), load.upstreams) + } } 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 5d303a3dd814a..78bb8d4c12e84 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 @@ -85,6 +85,19 @@ class TaskTest { Assertions.assertEquals(TaskState.State.FAILED, (result as TaskState).state) } + @Test + @DisplayName("Should thread the task definition into the execution context") + fun shouldThreadTaskDefIntoContext() { + val result = + runTask( + bundleWith("asserting", TaskDefAssertingTask::class.java), + startupDetails(taskId = "asserting"), + noOpClient(), + ) + + Assertions.assertInstanceOf(SucceedTask::class.java, result) + } + private fun bundleWith( taskId: String, taskClass: Class, @@ -171,4 +184,15 @@ class TaskTest { client: Client, ): Unit = throw NoClassDefFoundError("simulated") } + + class TaskDefAssertingTask : Task { + override fun execute( + context: Context, + client: Client, + ) { + check(context.taskDef?.id == context.ti.taskId) { + "expected the runner to thread the task definition into the context" + } + } + } } diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/ArgValuesTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/ArgValuesTest.kt new file mode 100644 index 0000000000000..d9e2a398168ff --- /dev/null +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/ArgValuesTest.kt @@ -0,0 +1,193 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +@file:Suppress("PLATFORM_CLASS_MAPPED_TO_KOTLIN") + +package org.apache.airflow.sdk.internal + +import org.apache.airflow.sdk.Client +import org.apache.airflow.sdk.Context +import org.apache.airflow.sdk.DagDef +import org.apache.airflow.sdk.DagRun +import org.apache.airflow.sdk.In +import org.apache.airflow.sdk.MissingXComException +import org.apache.airflow.sdk.Task +import org.apache.airflow.sdk.TaskDef +import org.apache.airflow.sdk.TaskInstance +import org.apache.airflow.sdk.TaskRef +import org.apache.airflow.sdk.execution.comm.ConnectionResult +import org.apache.airflow.sdk.execution.comm.StartupDetails +import org.apache.airflow.sdk.execution.comm.VariableResult +import org.apache.airflow.sdk.execution.comm.XComResult +import org.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.Assertions.assertNull +import org.junit.jupiter.api.Assertions.assertThrows +import org.junit.jupiter.api.DisplayName +import org.junit.jupiter.api.Test +import org.apache.airflow.sdk.execution.Client as Transport +import org.apache.airflow.sdk.execution.comm.TaskInstance as CommTaskInstance + +private class NoopArgTask : Task { + override fun execute( + context: Context, + client: Client, + ) = Unit +} + +/** Resolution of the inputs a `@Wiring` method recorded, without runtime bindings. */ +internal class ArgValuesTest { + private fun clientWith(xcomsByTask: Map): Client = + Client( + StartupDetails().also { + it.ti = + CommTaskInstance().also { ti -> + ti.taskId = "consumer" + ti.dagId = "d" + ti.runId = "r" + ti.tryNumber = 1 + } + }, + object : Transport { + override fun getConnection(id: String): ConnectionResult = throw NotImplementedError() + + override fun getVariable(key: String): VariableResult = throw NotImplementedError() + + override fun getXCom( + key: String, + dagId: String, + taskId: String, + runId: String, + mapIndex: Int?, + includePriorDates: Boolean, + ): XComResult = XComResult().also { it.value = xcomsByTask[taskId] } + + override fun setXCom( + key: String, + value: Any, + dagId: String, + taskId: String, + runId: String, + mapIndex: Int, + ): Unit = throw NotImplementedError() + }, + ) + + private fun contextFor(inputs: List>): Context { + val dag = DagDef("d") + inputs.filterIsInstance>().forEach { dag.addTask(it.def) } + val def = TaskDef("consumer", NoopArgTask::class.java) + Refs.register(dag, def, inputs) + return contextWithoutTaskDef().also { it.taskDef = def } + } + + private fun contextWithoutTaskDef(): Context = + Context( + dagRun = DagRun("d", "r", null, null, null, null, null, emptyMap()), + ti = TaskInstance("d", "r", "consumer", null, 1), + ) + + private fun handleFor(taskId: String): TaskRef = TaskRef(TaskDef(taskId, NoopArgTask::class.java)) + + @Test + @DisplayName("Should resolve a handle input from the upstream task's XCom") + fun shouldResolveHandleInputFromXCom() { + val context = contextFor(listOf(handleFor("producer"))) + + assertEquals( + 42L, + ArgValues.requiredInput(context, clientWith(mapOf("producer" to 42L)), 0, java.lang.Long::class.java, "value"), + ) + } + + @Test + @DisplayName("Should resolve a literal input without touching the client") + fun shouldResolveLiteralInput() { + val context = contextFor(listOf(In.value(7))) + + assertEquals(7L, ArgValues.requiredInput(context, clientWith(emptyMap()), 0, java.lang.Long::class.java, "value")) + } + + @Test + @DisplayName("Should throw MissingXComException when a required upstream pushed no value") + fun shouldThrowForMissingRequiredValue() { + val context = contextFor(listOf(handleFor("producer"))) + val client = clientWith(mapOf("producer" to null)) + + assertThrows(MissingXComException::class.java) { + ArgValues.requiredInput(context, client, 0, Integer::class.java, "value") + } + } + + @Test + @DisplayName("Should throw MissingXComException for a required null literal") + fun shouldThrowForRequiredNullLiteral() { + val context = contextFor(listOf(In.value(null))) + val client = clientWith(emptyMap()) + + val error = + assertThrows(MissingXComException::class.java) { + ArgValues.requiredInput(context, client, 0, Integer::class.java, "value") + } + + assertEquals( + "'value' has a primitive type but its wired literal input is null; " + + "declare a boxed type (e.g. Integer instead of int) to receive null.", + error.message, + ) + } + + @Test + @DisplayName("Should pass null through for optional inputs") + fun shouldPassNullThroughForOptionalInputs() { + val context = contextFor(listOf(handleFor("producer"))) + + assertNull(ArgValues.optionalInput(context, clientWith(mapOf("producer" to null)), 0, Integer::class.java)) + } + + @Test + @DisplayName("Should fail when the position has no wired input") + fun shouldFailOnUnwiredPosition() { + val context = contextFor(listOf()) + + val error = + assertThrows(IllegalStateException::class.java) { + ArgValues.optionalInput(context, clientWith(emptyMap()), 0, Integer::class.java) + } + + assertEquals( + "Task 'consumer' declares a data parameter at position 0 but only 0 input(s) are wired", + error.message, + ) + } + + @Test + @DisplayName("Should fail when the context carries no task definition") + fun shouldFailWithoutTaskDef() { + val error = + assertThrows(IllegalStateException::class.java) { + ArgValues.optionalInput(contextWithoutTaskDef(), clientWith(emptyMap()), 0, Integer::class.java) + } + + assertEquals( + "Task 'consumer' declares data parameters but has no wired inputs; " + + "register it through a @Wiring method", + error.message, + ) + } +} diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/RefsTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/RefsTest.kt new file mode 100644 index 0000000000000..c3831fb6b0c90 --- /dev/null +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/RefsTest.kt @@ -0,0 +1,88 @@ +/* + * 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.internal + +import org.apache.airflow.sdk.Client +import org.apache.airflow.sdk.Context +import org.apache.airflow.sdk.DagDef +import org.apache.airflow.sdk.In +import org.apache.airflow.sdk.LiteralIn +import org.apache.airflow.sdk.Task +import org.apache.airflow.sdk.TaskDef +import org.apache.airflow.sdk.TaskRef +import org.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.Assertions.assertThrows +import org.junit.jupiter.api.DisplayName +import org.junit.jupiter.api.Test + +private class NoopRefTask : Task { + override fun execute( + context: Context, + client: Client, + ) = Unit +} + +internal class RefsTest { + @Test + @DisplayName("Should register the task, record inputs, and wire handle edges") + fun shouldRegisterTaskWithInputsAndEdges() { + val dag = DagDef("d") + val producer = Refs.register(dag, TaskDef("p", NoopRefTask::class.java), listOf()) + Refs.register( + dag, + TaskDef("c", NoopRefTask::class.java), + listOf(producer, In.value(5)), + ) + + val consumerDef = dag.tasks.getValue("c") + assertEquals(setOf("p", "c"), dag.tasks.keys) + assertEquals(setOf(dag.tasks.getValue("p")), consumerDef.upstreams) + assertEquals(2, consumerDef.inputs.size) + assertEquals(dag.tasks.getValue("p"), (consumerDef.inputs[0] as TaskRef<*>).def) + assertEquals(5, (consumerDef.inputs[1] as LiteralIn<*>).value) + } + + @Test + @DisplayName("Should pass requireRegistered when every task was wired") + fun shouldPassRequireRegisteredWhenComplete() { + val dag = DagDef("d") + Refs.register(dag, TaskDef("t", NoopRefTask::class.java), listOf()) + + Refs.requireRegistered(dag, listOf("t")) + } + + @Test + @DisplayName("Should fail requireRegistered naming the tasks the wiring missed") + fun shouldFailRequireRegisteredNamingMissedTasks() { + val dag = DagDef("d") + Refs.register(dag, TaskDef("t", NoopRefTask::class.java), listOf()) + + val error = + assertThrows(IllegalArgumentException::class.java) { + Refs.requireRegistered(dag, listOf("t", "x", "y")) + } + + assertEquals( + "Wiring for Dag 'd' did not register task(s) 'x', 'y': " + + "every @Builder.Task method must be invoked in the @Wiring method", + error.message, + ) + } +} From 8de00c35cec2a795bae385f51ad2fb7ac957be66 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 5 Aug 2026 15:39:05 +0000 Subject: [PATCH 10/10] Java SDK: Serialize native Dags to DagSerialization v3 on parse requests A Java-authored Dag could not reach the scheduler on its own: the runtime answered task-execution requests only, so a Python stub file still had to exist purely to describe the Dag's structure. With dependency edges and schema-keyed configuration now recorded on the Java Dag model, the runtime has everything it needs to answer the coordinator's DagFileParseRequest the same way the Go SDK does, and the nativedag examples become real schedulable Dags with no Python counterpart. Native Java tasks deliberately emit no _arg_bindings: the execution API delivers bindings only for Python _StubOperator tasks, and a Java task always runs inside the JVM bundle that already holds its wired inputs, so the runtime resolves them locally. Cron schedules map to CronTriggerTimetable only, mirroring the Go SDK until the supervisor forwards the [scheduler] timetable flags over the coordinator protocol (see the TODO at the timetable serializer). --- .../kotlin/org/apache/airflow/sdk/Server.kt | 11 + .../org/apache/airflow/sdk/execution/Serde.kt | 302 ++++++++++++++++++ .../org/apache/airflow/sdk/ServerTest.kt | 62 ++++ .../apache/airflow/sdk/execution/SerdeTest.kt | 238 ++++++++++++++ 4 files changed, 613 insertions(+) create mode 100644 java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Serde.kt create mode 100644 java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/SerdeTest.kt diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt index 654baba944ba7..8dd7a840857db 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt @@ -33,8 +33,10 @@ import kotlinx.coroutines.runBlocking import org.apache.airflow.sdk.execution.CoordinatorComm import org.apache.airflow.sdk.execution.LogSender import org.apache.airflow.sdk.execution.Logger +import org.apache.airflow.sdk.execution.comm.DagFileParseRequest import org.apache.airflow.sdk.execution.comm.ErrorResponse import org.apache.airflow.sdk.execution.comm.StartupDetails +import org.apache.airflow.sdk.execution.parseDags import org.apache.airflow.sdk.execution.runTask import kotlin.text.substringAfterLast import kotlin.text.substringBeforeLast @@ -176,6 +178,7 @@ class Server( val frame = coordinator.readMessage() when (val body = frame.body) { is StartupDetails -> runTaskAndReport(bundle, body, coordinator) + is DagFileParseRequest -> parseDagsAndReport(bundle, body, coordinator) is ErrorResponse -> throw ApiError("[${body.error}] ${body.detail}") else -> throw ApiError("Unexpected initial frame (id=${frame.id})") } @@ -189,4 +192,12 @@ class Server( val result = runTask(bundle, startup, coordinator) coordinator.communicate(result) } + + private suspend fun parseDagsAndReport( + bundle: Bundle, + request: DagFileParseRequest, + coordinator: CoordinatorComm, + ) { + coordinator.communicate(parseDags(bundle, request)) + } } diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Serde.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Serde.kt new file mode 100644 index 0000000000000..1304fe43bcf3c --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Serde.kt @@ -0,0 +1,302 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.airflow.sdk.execution + +import com.fasterxml.jackson.databind.ObjectMapper +import org.apache.airflow.sdk.Bundle +import org.apache.airflow.sdk.DagDef +import org.apache.airflow.sdk.TaskDef +import org.apache.airflow.sdk.execution.comm.DagFileParseRequest +import org.apache.airflow.sdk.internal.Field +import org.apache.airflow.sdk.internal.SchemaFields +import java.nio.file.InvalidPathException +import java.nio.file.Paths +import java.time.Duration +import java.time.Instant +import java.time.OffsetDateTime + +// Serializes Dags to Airflow DagSerialization v3 JSON, mirroring the Go SDK's +// serde (go-sdk/pkg/execution/serde.go), which in turn matches Python's +// DagSerialization output. + +// Per-Dag defaults that Python resolves from [core] config when the Dag does +// not override them. The serializer always emits these fields (they have no +// JSON-schema default to omit against), so we fall back to the same values. +private const val DEFAULT_MAX_ACTIVE_TASKS_PER_DAG = 16 // [core] max_active_tasks_per_dag +private const val DEFAULT_MAX_ACTIVE_RUNS_PER_DAG = 16 // [core] max_active_runs_per_dag + +private val defaultsMapper = ObjectMapper() + +/** + * Processes a [DagFileParseRequest] by serialising every Dag registered on + * [bundle] to DagSerialization v3 and returning the result as a + * DagFileParsingResult body. + */ +internal fun parseDags( + bundle: Bundle, + request: DagFileParseRequest, +): Map { + val fileloc = request.file ?: "" + val relativeFileloc = computeRelativeFileloc(fileloc, request.bundlePath) + val serializedDags = + bundle.dags.values.map { dag -> + mapOf( + "data" to + mapOf( + "__version" to 3, + "dag" to serializeDag(dag, fileloc, relativeFileloc), + ), + ) + } + return linkedMapOf( + "type" to "DagFileParsingResult", + "fileloc" to fileloc, + "serialized_dags" to serializedDags, + ) +} + +/** + * Converts a [DagDef] to Airflow DagSerialization v3 format. Required fields are + * always present; config-driven fields follow the rules in [applyDagConfig] + * (some always emitted, some only when set). + */ +internal fun serializeDag( + dag: DagDef, + fileloc: String, + relativeFileloc: String, +): Map { + val downstream = linkedMapOf>() + dag.tasks.forEach { (taskId, def) -> + def.upstreams.forEach { upstream -> + downstream.getOrPut(upstream.id) { mutableListOf() } += taskId + } + } + + val result = + linkedMapOf( + "dag_id" to dag.id, + "fileloc" to fileloc, + "relative_fileloc" to relativeFileloc, + "timezone" to "UTC", + "timetable" to serializeTimetable(dag.dagConfig["schedule"] as String?), + "tasks" to dag.tasks.map { (taskId, def) -> serializeTask(taskId, def, downstream[taskId]) }, + "dag_dependencies" to emptyList(), + "task_group" to serializeTaskGroup(dag.tasks.keys), + "edge_info" to emptyMap(), + "params" to emptyList(), + "deadline" to null, + "allowed_run_types" to null, + ) + applyDagConfig(result, dag.dagConfig) + return result +} + +/** + * Converts one task to the Airflow serialization format. `downstream` is the + * inverted view of the Dag's upstream edges, sorted for stable JSON. + * + * Native Java tasks deliberately emit no `_arg_bindings`: the execution API + * delivers bindings only for Python `_StubOperator` tasks, and a Java task + * always executes inside the JVM bundle that already holds its wired inputs, + * so the runtime resolves them locally. + */ +private fun serializeTask( + taskId: String, + def: TaskDef, + downstream: List?, +): Map { + val data = + linkedMapOf( + "task_id" to taskId, + "task_type" to def.definition.simpleName, + "_task_module" to def.definition.packageName, + "language" to "java", + // Python's operator serializer always emits template_fields (its list + // value never matches the tuple default it is compared against), so it + // is unconditional here too. Java tasks have no template fields. + "template_fields" to emptyList(), + ) + // Emit only config entries that differ from their schema default, mirroring + // Python BaseSerialization's "omit hard-coded default" behavior. Operator + // fields are stored unwrapped, so the __type encoding is stripped. + def.configValues.forEach { (key, value) -> + if (!matchesSchemaDefault(SchemaFields.TASK[key], value)) { + data[key] = unwrapTypeEncoding(serializeValue(value)) + } + } + if (!downstream.isNullOrEmpty()) { + data["downstream_task_ids"] = downstream.sorted() + } + return mapOf( + "__type" to "operator", + "__var" to data, + ) +} + +/** + * Writes Dag-level config onto [data]. Fields with a JSON-schema default + * (description, dates, tags, fail_fast, ...) are omitted when unset. Fields + * with no schema default (catchup, disable_bundle_versioning, + * max_active_tasks, max_active_runs, max_consecutive_failed_dag_runs) are + * always emitted, because Python's serializer never omits them — it writes + * the resolved value, falling back to the matching `[core]` config default. + */ +private fun applyDagConfig( + data: MutableMap, + config: Map, +) { + listOf("description", "dag_display_name", "doc_md", "start_date", "end_date", "dagrun_timeout").forEach { key -> + config[key]?.let { data[key] = unwrapTypeEncoding(serializeValue(it)) } + } + (config["tags"] as? List<*>)?.let { tags -> + // Python stores tags in a set and serializes them sorted (for a stable + // dag_hash); mirror that regardless of registration order. + data["tags"] = tags.map { it.toString() }.sorted() + } + data["max_active_tasks"] = config["max_active_tasks"] ?: DEFAULT_MAX_ACTIVE_TASKS_PER_DAG + data["max_active_runs"] = config["max_active_runs"] ?: DEFAULT_MAX_ACTIVE_RUNS_PER_DAG + data["max_consecutive_failed_dag_runs"] = config["max_consecutive_failed_dag_runs"] ?: 0 + data["catchup"] = config["catchup"] ?: false + data["disable_bundle_versioning"] = config["disable_bundle_versioning"] ?: false + // fail_fast and render_template_as_native_obj have schema default false, so + // Python omits them when false; keep that behavior. + if (config["fail_fast"] == true) data["fail_fast"] = true + if (config["render_template_as_native_obj"] == true) data["render_template_as_native_obj"] = true + config["is_paused_upon_creation"]?.let { data["is_paused_upon_creation"] = it } +} + +// TODO: respect [scheduler] create_cron_data_intervals like Python's +// _create_timetable; the JVM bundle cannot read airflow.cfg, so the +// supervisor must send those flags over the coordinator protocol first. +// Mirrors the Go SDK's default-only behavior; tracked at +// https://github.com/apache/airflow/issues/67938 +private fun serializeTimetable(schedule: String?): Map = + when (schedule) { + null -> mapOf("__type" to "airflow.timetables.simple.NullTimetable", "__var" to emptyMap()) + "@once" -> mapOf("__type" to "airflow.timetables.simple.OnceTimetable", "__var" to emptyMap()) + "@continuous" -> + mapOf("__type" to "airflow.timetables.simple.ContinuousTimetable", "__var" to emptyMap()) + else -> + mapOf( + "__type" to "airflow.timetables.trigger.CronTriggerTimetable", + "__var" to + mapOf( + "expression" to schedule, + "timezone" to "UTC", + "interval" to 0.0, + "run_immediately" to false, + ), + ) + } + +/** Creates the flat root task group containing all task IDs. */ +private fun serializeTaskGroup(taskIds: Collection): Map = + mapOf( + "_group_id" to null, + "group_display_name" to "", + "prefix_group_id" to true, + "tooltip" to "", + "ui_color" to "CornflowerBlue", + "ui_fgcolor" to "#000", + "children" to taskIds.associateWith { listOf("operator", it) }, + "upstream_group_ids" to emptyList(), + "downstream_group_ids" to emptyList(), + "upstream_task_ids" to emptyList(), + "downstream_task_ids" to emptyList(), + ) + +/** + * Recursively serializes a value with Airflow's type/var encoding, matching + * Python's `BaseSerialization.serialize()` output: primitives pass through, + * date-times become `{"__type": "datetime", "__var": epoch_seconds}`, + * durations `{"__type": "timedelta", "__var": total_seconds}`, and maps + * `{"__type": "dict", "__var": {...}}`. + */ +internal fun serializeValue(value: Any?): Any? = + when (value) { + null -> null + is String, is Boolean, is Int, is Long, is Double -> value + is Byte, is Short -> (value as Number).toInt() + is Float -> value.toDouble() + is OffsetDateTime -> serializeValue(value.toInstant()) + is Instant -> + mapOf( + "__type" to "datetime", + "__var" to value.epochSecond + value.nano / 1e9, + ) + is Duration -> + mapOf( + "__type" to "timedelta", + "__var" to value.toNanos() / 1e9, + ) + is Map<*, *> -> + mapOf( + "__type" to "dict", + "__var" to value.entries.associate { (k, v) -> k.toString() to serializeValue(v) }, + ) + is List<*> -> value.map(::serializeValue) + is Array<*> -> value.map(::serializeValue) + else -> value + } + +/** + * Extracts the `__var` part from a type-encoded value: in Python's + * `serialize_to_json`, non-decorated fields are serialized then unwrapped. + */ +internal fun unwrapTypeEncoding(value: Any?): Any? { + val map = value as? Map<*, *> ?: return value + if ("__type" !in map) return value + return if ("__var" in map) map["__var"] else value +} + +/** Whether a config value equals the schema default and can be omitted. */ +private fun matchesSchemaDefault( + field: Field?, + value: Any, +): Boolean { + val defaultJson = field?.defaultJson ?: return false + val node = defaultsMapper.readTree(defaultJson) + return when (value) { + is String -> node.isTextual && node.asText() == value + is Boolean -> node.isBoolean && node.asBoolean() == value + is Number -> node.isNumber && node.asDouble() == value.toDouble() + is Duration -> node.isNumber && node.asDouble() == value.toNanos() / 1e9 + else -> false + } +} + +private fun computeRelativeFileloc( + fileloc: String, + bundlePath: String?, +): String { + if (fileloc.isEmpty()) return "" + if (bundlePath.isNullOrEmpty()) return "." + return try { + Paths + .get(bundlePath) + .relativize(Paths.get(fileloc)) + .toString() + .ifEmpty { "." } + } catch (e: InvalidPathException) { + "." + } catch (e: IllegalArgumentException) { + "." + } +} diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ServerTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ServerTest.kt index 6cae7a4c27b29..4f43e2416b8b8 100644 --- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ServerTest.kt +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ServerTest.kt @@ -27,6 +27,7 @@ import kotlinx.coroutines.runBlocking import org.apache.airflow.sdk.execution.CoordinatorComm import org.apache.airflow.sdk.execution.Frame import org.apache.airflow.sdk.execution.IncomingFrame +import org.apache.airflow.sdk.execution.RawFrame import org.apache.airflow.sdk.execution.comm.TaskState import org.junit.jupiter.api.Assertions import org.junit.jupiter.api.DisplayName @@ -37,6 +38,13 @@ import java.io.ByteArrayOutputStream import java.util.concurrent.ArrayBlockingQueue import java.util.concurrent.TimeUnit +private class ServerNoopTask : Task { + override fun execute( + context: Context, + client: Client, + ) = Unit +} + class ServerTest { private fun hexToBytes(hex: String): ByteArray = hex @@ -112,6 +120,60 @@ class ServerTest { comm.close() } + @Test + @DisplayName("Should serialize the bundle's dags when the initial frame is a parse request") + @Timeout(value = 30, unit = TimeUnit.SECONDS) + fun parsesDagsAndReportsResult() { + val toServer = ByteChannel(autoFlush = true) + val fromServer = ByteChannel(autoFlush = true) + val comm = CoordinatorComm(toServer, fromServer) + val server = Server(InetSocketAddress("localhost", 0), InetSocketAddress("localhost", 0)) + val bundle = Bundle(listOf(DagDef("parsed_dag").addTask(TaskDef("t", ServerNoopTask::class.java)))) + + val reported = ArrayBlockingQueue(1) + val supervisor = + Thread { + runBlocking { + toServer.writeFrame(parseRequestFrame(3, "/bundle/dags/java.py", "/bundle")) + val prefix = fromServer.readByteArray(4) + val payload = fromServer.readByteArray(Frame.parseLengthPrefix(prefix)) + val raw = Frame.decodeRaw(payload) + reported.put(raw) + toServer.writeFrame(ackFrame(raw.id)) + } + } + supervisor.start() + + runBlocking { server.dispatchTask(bundle, comm) } + supervisor.join() + + val body = reported.take().rawBody as Map<*, *> + Assertions.assertEquals("DagFileParsingResult", body["type"]) + Assertions.assertEquals("/bundle/dags/java.py", body["fileloc"]) + val dags = body["serialized_dags"] as List<*> + Assertions.assertEquals(1, dags.size) + val dag = ((dags[0] as Map<*, *>)["data"] as Map<*, *>)["dag"] as Map<*, *> + Assertions.assertEquals("parsed_dag", dag["dag_id"]) + comm.close() + } + + private fun parseRequestFrame( + id: Int, + file: String, + bundlePath: String, + ): ByteArray { + val out = ByteArrayOutputStream() + MessagePack.newDefaultPacker(out).use { packer -> + packer.packArrayHeader(2) + packer.packInt(id) + packer.packMapHeader(3) + packer.packString("type").packString("DagFileParseRequest") + packer.packString("file").packString(file) + packer.packString("bundle_path").packString(bundlePath) + } + return out.toByteArray() + } + private companion object { // [2, msg, null] with msg coming from // https://github.com/astronomer/airflow/blob/f39c8da8/task-sdk/tests/task_sdk/execution_time/test_comms.py#L73-L108 diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/SerdeTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/SerdeTest.kt new file mode 100644 index 0000000000000..c4e2c86736cc4 --- /dev/null +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/SerdeTest.kt @@ -0,0 +1,238 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.airflow.sdk.execution + +import org.apache.airflow.sdk.Bundle +import org.apache.airflow.sdk.Client +import org.apache.airflow.sdk.Context +import org.apache.airflow.sdk.DagDef +import org.apache.airflow.sdk.In +import org.apache.airflow.sdk.Task +import org.apache.airflow.sdk.TaskDef +import org.apache.airflow.sdk.execution.comm.DagFileParseRequest +import org.apache.airflow.sdk.internal.Refs +import org.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.Assertions.assertFalse +import org.junit.jupiter.api.Assertions.assertNull +import org.junit.jupiter.api.DisplayName +import org.junit.jupiter.api.Test +import java.time.Duration +import java.time.OffsetDateTime + +private class SerdeNoopTask : Task { + override fun execute( + context: Context, + client: Client, + ) = Unit +} + +@Suppress("UNCHECKED_CAST") +private fun taskData( + serialized: Map, + index: Int, +): Map { + val tasks = serialized["tasks"] as List> + assertEquals("operator", tasks[index]["__type"]) + return tasks[index]["__var"] as Map +} + +internal class SerdeTest { + @Test + @DisplayName("Should always emit required dag fields and config-backed defaults") + fun shouldEmitRequiredDagFields() { + val serialized = serializeDag(DagDef("d"), "/bundles/app/dags.jar", "app/dags.jar") + + assertEquals("d", serialized["dag_id"]) + assertEquals("/bundles/app/dags.jar", serialized["fileloc"]) + assertEquals("app/dags.jar", serialized["relative_fileloc"]) + assertEquals("UTC", serialized["timezone"]) + assertEquals( + mapOf("__type" to "airflow.timetables.simple.NullTimetable", "__var" to emptyMap()), + serialized["timetable"], + ) + assertEquals(emptyList(), serialized["tasks"]) + assertEquals(emptyList(), serialized["dag_dependencies"]) + assertEquals(emptyMap(), serialized["edge_info"]) + assertEquals(emptyList(), serialized["params"]) + assertNull(serialized["deadline"]) + assertNull(serialized["allowed_run_types"]) + assertEquals(16, serialized["max_active_tasks"]) + assertEquals(16, serialized["max_active_runs"]) + assertEquals(0, serialized["max_consecutive_failed_dag_runs"]) + assertEquals(false, serialized["catchup"]) + assertEquals(false, serialized["disable_bundle_versioning"]) + assertFalse("description" in serialized) + assertFalse("fail_fast" in serialized) + assertFalse("tags" in serialized) + } + + @Test + @DisplayName("Should map schedule strings to the matching timetable") + fun shouldMapScheduleToTimetable() { + val cron = serializeDag(DagDef("d").config("schedule", "@daily"), "", ".") + assertEquals( + mapOf( + "__type" to "airflow.timetables.trigger.CronTriggerTimetable", + "__var" to + mapOf( + "expression" to "@daily", + "timezone" to "UTC", + "interval" to 0.0, + "run_immediately" to false, + ), + ), + cron["timetable"], + ) + + val once = serializeDag(DagDef("d").config("schedule", "@once"), "", ".") + assertEquals( + mapOf("__type" to "airflow.timetables.simple.OnceTimetable", "__var" to emptyMap()), + once["timetable"], + ) + + val continuous = serializeDag(DagDef("d").config("schedule", "@continuous"), "", ".") + assertEquals( + mapOf("__type" to "airflow.timetables.simple.ContinuousTimetable", "__var" to emptyMap()), + continuous["timetable"], + ) + } + + @Test + @DisplayName("Should apply dag config values with Python's emit rules") + fun shouldApplyDagConfig() { + val dag = + DagDef("d") + .config("description", "demo") + .config("tags", listOf("b", "a")) + .config("catchup", true) + .config("fail_fast", true) + .config("max_active_runs", 3) + .config("dagrun_timeout", Duration.ofMinutes(5)) + .config("start_date", OffsetDateTime.parse("2026-01-01T00:00:00Z")) + + val serialized = serializeDag(dag, "", ".") + + assertEquals("demo", serialized["description"]) + assertEquals(listOf("a", "b"), serialized["tags"]) + assertEquals(true, serialized["catchup"]) + assertEquals(true, serialized["fail_fast"]) + assertEquals(3, serialized["max_active_runs"]) + assertEquals(300.0, serialized["dagrun_timeout"]) + assertEquals(1.7672256E9, serialized["start_date"]) + } + + @Test + @DisplayName("Should serialize tasks with identity fields, config, and sorted downstream ids") + fun shouldSerializeTasks() { + val extractDef = + TaskDef("extract", SerdeNoopTask::class.java) + .config("retries", 2) + .config("queue", "q") + .config("retry_delay", Duration.ofMinutes(10)) + val transformDef = + TaskDef("transform", SerdeNoopTask::class.java) + .dependsOn(extractDef) + // Explicitly at schema defaults: omitted from the serialized form. + .config("retries", 0) + .config("queue", "default") + .config("retry_delay", Duration.ofMinutes(5)) + val dag = DagDef("d").addTask(extractDef).addTask(transformDef) + + val serialized = serializeDag(dag, "", ".") + + val extract = taskData(serialized, 0) + assertEquals("extract", extract["task_id"]) + assertEquals("SerdeNoopTask", extract["task_type"]) + assertEquals("org.apache.airflow.sdk.execution", extract["_task_module"]) + assertEquals("java", extract["language"]) + assertEquals(emptyList(), extract["template_fields"]) + assertEquals(2, extract["retries"]) + assertEquals("q", extract["queue"]) + assertEquals(600.0, extract["retry_delay"]) + assertEquals(listOf("transform"), extract["downstream_task_ids"]) + + val transform = taskData(serialized, 1) + assertEquals("transform", transform["task_id"]) + assertFalse("retries" in transform) + assertFalse("queue" in transform) + assertFalse("retry_delay" in transform) + assertFalse("downstream_task_ids" in transform) + + assertEquals( + mapOf("extract" to listOf("operator", "extract"), "transform" to listOf("operator", "transform")), + (serialized["task_group"] as Map<*, *>)["children"], + ) + } + + @Test + @DisplayName("Should serialize wiring-registered dags with their data-flow edges") + fun shouldSerializeWiredDag() { + val dag = DagDef("d") + val extracted = Refs.register(dag, TaskDef("extract", SerdeNoopTask::class.java), listOf()) + Refs.register(dag, TaskDef("transform", SerdeNoopTask::class.java), listOf>(extracted)) + + val serialized = serializeDag(dag, "", ".") + + assertEquals(listOf("transform"), taskData(serialized, 0)["downstream_task_ids"]) + assertFalse("_arg_bindings" in taskData(serialized, 1)) + } + + @Test + @DisplayName("Should wrap parsed dags in a DagFileParsingResult body") + fun shouldBuildParsingResult() { + val bundle = Bundle(listOf(DagDef("d").addTask(TaskDef("t", SerdeNoopTask::class.java)))) + val request = + DagFileParseRequest().also { + it.file = "/bundles/app/dags.jar" + it.bundlePath = "/bundles" + } + + val result = parseDags(bundle, request) + + assertEquals("DagFileParsingResult", result["type"]) + assertEquals("/bundles/app/dags.jar", result["fileloc"]) + val dags = result["serialized_dags"] as List<*> + assertEquals(1, dags.size) + val data = (dags[0] as Map<*, *>)["data"] as Map<*, *> + assertEquals(3, data["__version"]) + val dag = data["dag"] as Map<*, *> + assertEquals("d", dag["dag_id"]) + assertEquals("app/dags.jar", dag["relative_fileloc"]) + } + + @Test + @DisplayName("Should encode temporals and nested maps with the type/var envelope") + fun shouldEncodeValuesWithTypeEnvelope() { + assertEquals( + mapOf("__type" to "timedelta", "__var" to 90.0), + serializeValue(Duration.ofSeconds(90)), + ) + assertEquals( + mapOf("__type" to "datetime", "__var" to 1.7672256E9), + serializeValue(OffsetDateTime.parse("2026-01-01T00:00:00Z")), + ) + assertEquals( + mapOf("__type" to "dict", "__var" to mapOf("k" to listOf(1, 2))), + serializeValue(mapOf("k" to listOf(1, 2))), + ) + assertEquals(42, unwrapTypeEncoding(mapOf("__type" to "timedelta", "__var" to 42))) + assertEquals(mapOf("plain" to 1), unwrapTypeEncoding(mapOf("plain" to 1))) + } +}