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 c8b73b4626a28..280925dd240ba 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 ~~~~~~~~~~~~~~~~ @@ -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 } @@ -160,16 +160,29 @@ 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. - * - ``@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`). + - 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`` + 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 +197,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; @@ -203,7 +213,9 @@ 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``; 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 @@ -224,16 +236,193 @@ 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 fetch = new TaskDef("fetch", FetchTask.class).config("retries", 2); + var process = new TaskDef("process", ProcessTask.class).dependsOn(fetch); + var dag = new DagDef("my_dag") + .config("schedule", "@daily") + .addTask(fetch) + .addTask(process); return List.of(dag); } } 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/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 @@ -428,7 +617,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/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/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/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/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/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"), [ 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/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/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/ExampleBundleBuilder.java b/java-sdk/example/src/java/org/apache/airflow/example/ExampleBundleBuilder.java index f63d6c4d74337..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 @@ -26,11 +26,13 @@ public class ExampleBundleBuilder implements BundleBuilder { @NotNull @Override - public Iterable getDags() { + 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/InterfaceExampleBuilder.java b/java-sdk/example/src/java/org/apache/airflow/example/InterfaceExampleBuilder.java index 1c536c3cbf2c7..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,17 +67,19 @@ 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"); } } - 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/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/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/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 11202ffcdb455..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,20 +23,34 @@ 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 java.util.Optional +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 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 @@ -50,16 +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 [Dag] and registers those - * inner classes as tasks. + * - 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. * - * [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` 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) @@ -72,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, @@ -90,81 +122,311 @@ 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(ClassName.get(Dag::class.java)) - .addStatement($$"var dag = new $T($S)", ClassName.get(Dag::class.java), ann.id.ifBlank { el.simpleName }) + .returns(DAG_DEF_TYPE) + .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 ann = 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 + } - val task = buildTask(innerName, inner, el) - builderClass.addType(task.spec) + /** + * 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($S, $L.class)", - ann.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, - ): 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 inner = decl.method + val dataByName = decl.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)) + + 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 { + // 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() + } } + if (inner.returnType.kind == TypeKind.VOID) { $$"new $T().$L($L)" } else { @@ -178,80 +440,181 @@ 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(decl.className) + .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 [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. + */ +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 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, 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 3e28e1b009afd..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, @Builder.XCom(task = "t2") int value) { + 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()); + } } """, ) @@ -78,23 +85,32 @@ 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 java.lang.String; + import java.util.List; import org.apache.airflow.sdk.Client; import org.apache.airflow.sdk.Context; - import org.apache.airflow.sdk.Dag; - import org.apache.airflow.sdk.MissingXComException; + import org.apache.airflow.sdk.DagDef; import org.apache.airflow.sdk.Task; + import org.apache.airflow.sdk.internal.ArgValues; + import org.apache.airflow.sdk.internal.Refs; 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 final String DAG_ID = "TestExample"; + + public static DagDef dag() { + var dag = new DagDef(DAG_ID); return dag; } + + public static DagDef build() { + 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 { @@ -110,35 +126,67 @@ 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); } } } """, ) + 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 - @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) {} } """, ) @@ -152,33 +200,38 @@ 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.Dag; - import org.apache.airflow.sdk.MissingXComException; + 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 Dag build() { - var dag = new Dag("TestExample"); - dag.addTask("t", T.class); + 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(); + 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 { - 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); } } } @@ -187,20 +240,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) {} } """, ) @@ -214,34 +266,456 @@ 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.Dag; - import org.apache.airflow.sdk.MissingXComException; + 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 Dag build() { - var dag = new Dag("TestExample"); - dag.addTask("t", T.class); + 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(); + 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 { - 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); + } + } + } + """, + ) + } + + @Test + @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( + """ + 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.In; + import org.apache.airflow.sdk.TaskInput; + import org.apache.airflow.sdk.Wiring; + @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; } + + @Wiring + static void depends(TestExampleRef f) { + f.score(In.value(new ScoreInput())); + } + } + """, + ) + + 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.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 = 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; + 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 + @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 @@ -261,8 +735,13 @@ 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 java.lang.String; + import org.apache.airflow.sdk.DagDef; + 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; } + } """, ) } @@ -284,8 +763,13 @@ 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 java.lang.String; + import org.apache.airflow.sdk.DagDef; + 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; } + } """, ) } @@ -298,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(); + } + } """, ) @@ -311,14 +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.Dag; + 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 Dag build() { - var dag = new Dag("TestExample"); - dag.addTask("foo", T1.class); + 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(); + TestExample.depends(new TestExampleRef(dag)); + Refs.requireRegistered(dag, List.of("foo")); return dag; } public static final class T1 implements Task { @@ -327,23 +825,129 @@ 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("generate builder for dag class with invalid task parameter") - fun generateBuilderForDagClassWithInvalidTaskParameter() { + @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(id = "foo") public void t1(String client) {} } + 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( - "Unsupported task parameter 'client' with type: java.lang.String", + "@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", ) } @@ -364,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/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/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/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 deleted file mode 100644 index 3a5b84d2daf84..0000000000000 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Builder.kt +++ /dev/null @@ -1,88 +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 [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, @Builder.XCom(task = "extract") long extracted) { ... } - * } - * ``` - * - * The processor generates `MyPipelineBuilder.build()`, which returns a - * fully wired-up [Dag] 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 = "", - ) - - /** - * 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/Bundle.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Bundle.kt index 677ec48eb9311..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 @@ -20,22 +20,55 @@ 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]. - * @throws IllegalArgumentException if any two Dags share the same ID. + * @property dags All registered Dags keyed by [DagDef.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, + dags: Iterable, ) { - internal val dags: Map = dags.associateByDagId() + 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 { - 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 +78,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 +97,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/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/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/Dag.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Dag.kt deleted file mode 100644 index c998580374169..0000000000000 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Dag.kt +++ /dev/null @@ -1,98 +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 - -import kotlin.Throws - -/** - * A collection of tasks with directional dependencies. - * - * Create a [Dag] directly and register tasks 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. - * - * @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( - val id: String, // TODO: charset check? -) { - internal var tasks = mutableMapOf>() - - /** - * 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. - * - * @param id Task identifier, unique within this Dag. - * @param definition Class that implements [Task]. Must have a public no-arg - * constructor. - * @return This Dag, for chaining. - * @throws IllegalArgumentException if a task already exists in the Dag with - * the same ID. - */ - fun addTask( - id: String, - definition: Class, - ): Dag { - require(tasks.putIfAbsent(id, definition) == null) { - "Tasks in Dag have duplicate ID: $id" - } - return this - } -} - -/** - * A single unit of work executed by Airflow. - * - * Prefer using the [Builder.Task] annotation with [Builder.Dag] to have the - * annotation processor generate an implementation for you. Only use this - * interface if you need to do low-level plumbing. - * - * Implement this interface to define task logic. Airflow instantiates the class - * via its no-argument constructor, then calls [execute] once per task-instance - * run. - * - * @see Builder.Dag - * @see Builder.Task - */ -interface Task { - /** - * Executes this task. - * - * Any exception thrown marks the task instance as failed. Use [client] to - * read connections, variables, pull XComs, or to push an XCom for downstream - * tasks. - * - * @param context Runtime context for the current execution workload. - * @param client Client for Airflow API calls scoped to this exxecution. - * @throws Exception on failure; the task instance is marked failed. - */ - @Throws(Exception::class) - fun execute( - context: Context, - client: Client, - ) -} 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 new file mode 100644 index 0000000000000..e7f6e260756c9 --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/DagDef.kt @@ -0,0 +1,220 @@ +/* + * 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.apache.airflow.sdk.internal.SchemaFields +import org.apache.airflow.sdk.internal.checkConfigValue +import kotlin.Throws + +/** + * A collection of tasks with directional dependencies. + * + * 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: + * + * ```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") + * .config("schedule", "@daily") + * .addTask(extract) + * .addTask(load); + * ``` + * + * @param id Dag identifier. Must contain only ASCII alphanumeric characters, + * dashes, dots, or underscores; must be unique within a [Bundle]. + * + * @see Builder.Dag + */ +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. 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. + * @throws IllegalArgumentException if the task already belongs to a Dag or a + * task with the same ID is already registered. + */ + 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 + } + + /** + * 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, the class that implements it, its upstream + * dependencies, and its task-level configuration. + * + * ```java + * 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]. + * @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 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 + } +} + +/** + * A single unit of work executed by Airflow. + * + * Prefer using the [Builder.Task] annotation with [Builder.Dag] to have the + * annotation processor generate an implementation for you. Only use this + * interface if you need to do low-level plumbing. + * + * Implement this interface to define task logic. Airflow instantiates the class + * via its no-argument constructor, then calls [execute] once per task-instance + * run. + * + * @see Builder.Dag + * @see Builder.Task + */ +interface Task { + /** + * Executes this task. + * + * Any exception thrown marks the task instance as failed. Use [client] to + * read connections, variables, pull XComs, or to push an XCom for downstream + * tasks. + * + * @param context Runtime context for the current execution workload. + * @param client Client for Airflow API calls scoped to this exxecution. + * @throws Exception on failure; the task instance is marked failed. + */ + @Throws(Exception::class) + fun execute( + context: Context, + client: Client, + ) +} 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/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/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/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/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/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/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..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,9 +70,15 @@ 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 taskDef = + bundle.dags[request.ti.dagId]?.tasks[request.ti.taskId] + ?: return TaskResult.of(TaskState.State.REMOVED) return try { - task.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 new file mode 100644 index 0000000000000..e2d3ad4d9957f --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/ArgValues.kt @@ -0,0 +1,237 @@ +/* + * 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.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 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. + * + * 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() + + /** + * 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 resolved = resolveAt(context, client, position) + return decode(resolved.value, type) ?: throw resolved.missing(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(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 + * 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) + } + + /** 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, + ): Resolved { + val bindings = client.argBindings + 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 only ${inputs.size} input(s) are wired" + } + 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, + 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.", + ) + } + + 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, + ): 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/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 0e4afb1894a7f..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 @@ -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,9 +39,84 @@ 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) } + + @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 new file mode 100644 index 0000000000000..0a287330a94d5 --- /dev/null +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ClientArgTest.kt @@ -0,0 +1,335 @@ +/* + * 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.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.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 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") + 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 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() { + 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), + ) + + 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")) + } + + @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, + ) + } +} 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..437ce873e48c6 --- /dev/null +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/DagDefTest.kt @@ -0,0 +1,218 @@ +/* + * 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 +import java.time.Duration +import java.time.OffsetDateTime + +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) + } + + @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/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))) + } +} 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..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 @@ -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 @@ -84,12 +85,24 @@ 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, ): Bundle { - val dag = Dag("test_dag") - dag.addTask(taskId, taskClass) + val dag = DagDef("test_dag").addTask(TaskDef(taskId, taskClass)) return Bundle(listOf(dag)) } @@ -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, + ) + } +} 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()); } 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/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..d40f6bd4c7587 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,248 @@ 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: + 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): + 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 +324,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 +339,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 +376,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..95b029aeac7b2 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,378 @@ 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)) + + @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)() + 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 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 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; 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" },