Skip to content

Commit ab4da0b

Browse files
authored
Java SDK: Declare a Dag's task graph with @Builder.Deps (#71189)
* Java SDK: Declare a Dag's task graph with @Builder.Deps An annotated Dag class could name its tasks but not say how they connect, so a Dag whose annotations carried everything else still needed a Python Dag file for its graph. Declaring the graph by calling the tasks lets javac check that each task is fed something its parameter accepts. Bundle.register(Class) now also takes the @Builder.Dag class itself and registers the Dag its generated builder builds, so there is still only one name to keep in sync. * Java SDK: Quote wired task ids as Java string literals * Java SDK: Validate the @Builder.Deps class at compile time * Java SDK: Reject task method names that clash with the wiring view * Java SDK: Type a wired numeric input by its parameter's own type * Java SDK: Register both the Dag and handlers of a class, unwrapping builder failures * Java SDK: Reject a raw null wiring argument and test the recorder's guards * Java SDK: Reword the native-Dag docs without em-dashes * Java SDK: Drop Bundle tests that repeat existing coverage * Java SDK: Require a @Builder.Deps class on every @Builder.Dag class * Java SDK: Back the k8s test's stub tasks with task handlers * Java SDK: Reject overloaded task and task-handler methods * Java SDK: Fix the Dag examples a required wiring class invalidates
1 parent 31b7748 commit ab4da0b

21 files changed

Lines changed: 1756 additions & 160 deletions

File tree

‎airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -331,6 +331,9 @@ Annotate a plain Java class and let the SDK generate the boilerplate at compile
331331
- Marks a method as a task of a Java-owned Dag. If ``id`` is omitted the method name is
332332
used. Further attributes (``retries``, ``queue``, ``retryDelay``, …) are Airflow's own
333333
task settings; only attributes written explicitly are applied.
334+
* - ``@Builder.Deps``
335+
- Marks the nested class that declares the task graph in Java, TaskFlow-style. Required for
336+
a Dag that Java owns end to end. See :ref:`java-sdk/native-dags`.
334337
* - ``TaskInput`` / ``@ArgName("...")``
335338
- Marks a class as a task's input, so keyword arguments bind by name instead of by position:
336339
each public field receives the argument whose name matches it, ignoring case and
@@ -591,6 +594,80 @@ Edges are checked when the Dag is registered with a ``Bundle``: an upstream that
591594
Dag, or to no Dag, and a cycle anywhere in the graph both fail there rather than at the first task
592595
run.
593596

597+
Wiring the graph with ``@Builder.Deps``
598+
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
599+
600+
For a Dag written with annotations, the graph is declared by a nested ``@Builder.Deps`` class. The
601+
annotation processor generates a ``<ClassName>Deps`` interface, the *wiring view*, with one method
602+
per ``@Builder.Task`` method: the injected ``Client`` and ``Context`` parameters are dropped, each
603+
data parameter becomes an ``Arg<T>``, and the return value becomes a ``TaskRef<T>``. Calling a view
604+
method registers its task, and passing the handle one returned into another feeds the upstream's
605+
output into the downstream's parameter *and* wires the data edge. The call graph is the task graph,
606+
and ``javac`` type-checks it.
607+
608+
Declare the wiring class as a ``static`` nested class of the Dag class that ``implements`` the
609+
generated view, with a no-argument ``depends()`` method:
610+
611+
.. code-block:: java
612+
613+
@Builder.Dag(
614+
id = "java_etl",
615+
schedule = "@daily",
616+
description = "Pure-Java Dag built with annotations",
617+
tags = {"example", "java-sdk"})
618+
public class EtlPipeline {
619+
620+
@Builder.Task(id = "extract", retries = 2)
621+
public long extract() {
622+
return 42L;
623+
}
624+
625+
@Builder.Task(id = "transform")
626+
public long transform(long extracted, double factor) {
627+
return (long) (extracted * factor);
628+
}
629+
630+
@Builder.Task(id = "load")
631+
public void load(long transformed) {
632+
// implement task logic
633+
}
634+
635+
@Builder.Task(id = "audit")
636+
public void audit() {
637+
// side effect only, no data in or out
638+
}
639+
640+
@Builder.Deps
641+
static class Wiring implements EtlPipelineDeps {
642+
void depends() {
643+
var rows = extract();
644+
load(transform(rows, lit(0.9)));
645+
rows.before(audit()); // ordering-only edge: audit waits for extract
646+
}
647+
}
648+
}
649+
650+
Every ``@Builder.Task`` method must be called in the wiring class; a task the wiring missed fails at
651+
Dag-parse time. ``lit(...)`` wires an inline constant where no upstream feeds a parameter. A bare
652+
``double`` cannot be an ``Arg``, so a constant is wrapped. A view method that takes no arguments
653+
returns the same handle every time, so it names one node wherever it appears; one that takes
654+
arguments is called once, and the wiring fails if it is called again with arguments, so hold its
655+
handle in a local and reuse that.
656+
657+
``before``, ``after`` and ``Flow.of`` work here exactly as they do on the interface surface; inside
658+
the wiring class ``Flow`` is inherited by simple name, so it needs no import and never collides with
659+
``java.util.concurrent.Flow``.
660+
661+
Every ``@Builder.Dag`` class declares a wiring class, because the graph is what the Dag owns. A
662+
class that supplies only task bodies, for a Dag a Python file declares, carries
663+
``@Builder.TaskHandler`` instead and contributes no Dag.
664+
665+
.. note::
666+
667+
Runtime argument bindings win over Java-declared wiring. When the supervisor delivers bindings
668+
for a run (see :ref:`java-sdk/arg-binding`), the binding at a parameter's position is what the
669+
task receives. Wired inputs are the fallback, which is what a native Java Dag always uses.
670+
594671
Configuration attributes
595672
~~~~~~~~~~~~~~~~~~~~~~~~
596673

‎java-sdk/README.md‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -745,9 +745,9 @@ E2E_TEST_MODE=java_sdk uv run --project airflow-e2e-tests pytest \
745745
- The annotation processor (`BuilderProcessor.kt`) uses `kapt`. The `Builder`
746746
class holding the `@Builder.Dag` / `@Builder.Task` annotations is generated
747747
from the Dag serialization schema by `:sdk:generateDagDsl` (vendored at
748-
`sdk/schema/dag-schema.json`). The `Arg`/`TaskRef` and `Deps`/`Flow` graph
749-
types are hand-written next to the rest of the public surface in
750-
`sdk/src/main/kotlin/org/apache/airflow/sdk/`. When adding annotation
748+
`sdk/schema/dag-schema.json`), `@Builder.Deps` included. The `Arg`/`TaskRef`
749+
and `Deps`/`Flow` graph types are hand-written next to the rest of the public
750+
surface in `sdk/src/main/kotlin/org/apache/airflow/sdk/`. When adding annotation
751751
behaviour, handle it in `BuilderProcessor.kt` and add a golden-output test in
752752
`processor/src/test/kotlin/`.
753753
- The Python coordinator subclasses `SubprocessCoordinator`. Do not reach into

‎java-sdk/adr/0002-native-dag-interface.md‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -65,7 +65,8 @@ public class EtlPipeline { // extends nothing of ours; your own base class stays
6565
public void audit(Client client) { /* side effect only, no data in or out */ }
6666

6767
@Builder.Task(id = "notify")
68-
public void notify(Client client) { /* side effect only, no data in or out */ }
68+
// "notify" is Object.notify, so the method takes another name and keeps the task id.
69+
public void alert(Client client) { /* side effect only, no data in or out */ }
6970

7071
@Builder.Deps
7172
static class Wiring implements EtlPipelineDeps {
@@ -77,7 +78,7 @@ public class EtlPipeline { // extends nothing of ours; your own base class stays
7778

7879
// non-TaskFlow (ordering-only) edges: sequence with no data flowing
7980
rows.then(audit()); // extract >> audit
80-
Flow.of(loaded, audit()).then(notify()); // [load, audit] >> notify
81+
Flow.of(loaded, audit()).then(alert()); // [load, audit] >> notify
8182
}
8283
}
8384

@@ -101,7 +102,7 @@ interface EtlPipelineDeps extends Deps { // Deps: the shared base that nests Flo
101102
default TaskRef<Long> transform(Arg<Long> extracted, Arg<Double> threshold) { return Flow.call("transform", extracted, threshold); }
102103
default TaskRef<Void> load(Arg<Long> transformed) { return Flow.call("load", transformed); }
103104
default TaskRef<Void> audit() { return Flow.node("audit"); }
104-
default TaskRef<Void> notify() { return Flow.node("notify"); }
105+
default TaskRef<Void> alert() { return Flow.node("notify"); }
105106
}
106107
```
107108

‎java-sdk/example/src/java/org/apache/airflow/example/ExampleBundleBuilder.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,12 +21,15 @@
2121

2222
import org.apache.airflow.sdk.*;
2323

24+
// One bundle serves every surface: Dags built in Java, and the handler classes
25+
// whose Dags the Python file owns.
2426
public class ExampleBundleBuilder {
2527
public static Bundle build() {
2628
return new Bundle()
2729
.register(InterfaceExampleBuilder.build())
2830
.register(AnnotationExample.class)
2931
.register(XComCastingExample.class)
32+
.register(org.apache.airflow.example.nativedag.AnnotationExample.class)
3033
.register(org.apache.airflow.example.nativedag.InterfaceExample.build());
3134
}
3235

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,73 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
// "native" is a Java keyword, so the native-Dag examples live in "nativedag".
21+
package org.apache.airflow.example.nativedag;
22+
23+
import static java.lang.System.Logger.Level.INFO;
24+
25+
import org.apache.airflow.sdk.*;
26+
27+
// A Dag defined entirely in Java, annotation-style. The @Builder.Dag and
28+
// @Builder.Task attributes carry the configuration and the @Builder.Deps class
29+
// declares the graph -- the single place dependencies are defined.
30+
@Builder.Dag(
31+
id = "java_native_annotation_example",
32+
description = "Pure-Java Dag authored with annotations",
33+
schedule = "@daily",
34+
startDate = "2026-01-01T00:00:00Z",
35+
catchup = false,
36+
tags = {"example", "java-sdk"})
37+
public class AnnotationExample {
38+
private static final System.Logger log = System.getLogger(AnnotationExample.class.getName());
39+
40+
@Builder.Task(id = "extract", retries = 2)
41+
public long extract() {
42+
log.log(INFO, "Extracting a value");
43+
return 42L;
44+
}
45+
46+
@Builder.Task(id = "transform")
47+
public long transform(long extracted, double factor) {
48+
log.log(INFO, "Transforming {0} by {1}", extracted, factor);
49+
return (long) (extracted * factor);
50+
}
51+
52+
@Builder.Task(id = "load")
53+
public void load(long transformed) {
54+
log.log(INFO, "Loaded {0}", transformed);
55+
}
56+
57+
@Builder.Task(id = "audit")
58+
public void audit() {
59+
log.log(INFO, "Audited the run");
60+
}
61+
62+
// Implements the generated wiring view, so javac type-checks the graph:
63+
// extract() yields a TaskRef<Long>, which is an Arg<Long> for transform.
64+
@Builder.Deps
65+
static class Wiring implements AnnotationExampleDeps {
66+
void depends() {
67+
var extracted = extract();
68+
load(transform(extracted, lit(1.5)));
69+
// Ordering-only edge: audit runs after extract, with no data flowing.
70+
extracted.before(audit());
71+
}
72+
}
73+
}

0 commit comments

Comments
 (0)