Skip to content

Commit 6227100

Browse files
committed
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. (cherry picked from commit de94bc5)
1 parent 715775b commit 6227100

16 files changed

Lines changed: 1090 additions & 116 deletions

File tree

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

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -314,6 +314,9 @@ Annotate a plain Java class and let the SDK generate the boilerplate at compile
314314
- Marks a method as a task of a Java-owned Dag. If ``id`` is omitted the method name is
315315
used. Further attributes (``retries``, ``queue``, ``retryDelay``, …) mirror the Dag
316316
serialization schema; only attributes written explicitly are applied.
317+
* - ``@Builder.Deps``
318+
- Marks the nested class that declares the task graph in Java, TaskFlow-style. Only needed
319+
for a Dag that has no Python stub file. See :ref:`java-sdk/native-dags`.
317320
* - ``TaskInput`` / ``@ArgName("...")``
318321
- Marks a class as a task's input, so keyword arguments bind by name instead of by position:
319322
each public field receives the argument whose name matches it, ignoring case and
@@ -574,6 +577,79 @@ Edges are checked when the Dag is registered with a ``Bundle``: an upstream that
574577
Dag, or to no Dag, and a cycle anywhere in the graph both fail there rather than at the first task
575578
run.
576579

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

‎java-sdk/README.md‎

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

‎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)