Skip to content

Commit d62bdb0

Browse files
jason810496jason810496
committed
Java SDK: Register tasks as first-class TaskDef objects
`dag.addTask("extract", Extract.class)` stored tasks as a plain `Map<String, Class<out Task>>`, which leaves nowhere to hang anything else a task needs: dependency edges, task-level configuration, and argument wiring all have to attach to a per-task object, and a map of classes cannot carry them. Introducing that object now keeps those follow-ups additive instead of forcing another break of the registration API later. The annotation surface keeps `Builder.Dag` / `Builder.Task`, and the interface users implement keeps the `Task` name, so the definition objects are `DagDef` and `TaskDef` -- a pairing that stays unambiguous next to `Task` at a use site. The SDK is pre-1.0, so the old string-keyed overload is removed outright rather than deprecated.
1 parent b9ff1d7 commit d62bdb0

14 files changed

Lines changed: 202 additions & 84 deletions

File tree

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

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -115,7 +115,7 @@ Java entry point
115115
116116
public class Main implements BundleBuilder {
117117
@Override
118-
public Iterable<Dag> getDags() {
118+
public Iterable<DagDef> getDags() {
119119
return List.of(SalesPipelineBuilder.build()); // SalesPipelineBuilder generated at compile time
120120
}
121121
@@ -203,7 +203,7 @@ Interface-based API
203203
~~~~~~~~~~~~~~~~~~~
204204

205205
Implement the ``Task`` interface directly for full control over how tasks are registered and how XComs are
206-
read.
206+
read. Each task is registered as a ``TaskDef`` on a ``DagDef``.
207207

208208
.. code-block:: java
209209
@@ -224,10 +224,10 @@ Register tasks manually in a ``BundleBuilder``:
224224
225225
public class MyBundle implements BundleBuilder {
226226
@Override
227-
public Iterable<Dag> getDags() {
228-
var dag = new Dag("my_dag");
229-
dag.addTask("fetch", FetchTask.class);
230-
dag.addTask("process", ProcessTask.class);
227+
public Iterable<DagDef> getDags() {
228+
var dag = new DagDef("my_dag")
229+
.addTask(new TaskDef("fetch", FetchTask.class))
230+
.addTask(new TaskDef("process", ProcessTask.class));
231231
return List.of(dag);
232232
}
233233
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@
2626
public class ExampleBundleBuilder implements BundleBuilder {
2727
@NotNull
2828
@Override
29-
public Iterable<Dag> getDags() {
29+
public Iterable<DagDef> getDags() {
3030
return List.of(
3131
InterfaceExampleBuilder.build(),
3232
AnnotationExampleBuilder.build(),

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

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -71,11 +71,10 @@ public void execute(@NotNull Context context, Client client) {
7171
}
7272
}
7373

74-
public static Dag build() {
75-
var dag = new Dag("java_interface_example");
76-
dag.addTask("extract", Extract.class);
77-
dag.addTask("transform", Transform.class);
78-
dag.addTask("load", Load.class);
79-
return dag;
74+
public static DagDef build() {
75+
return new DagDef("java_interface_example")
76+
.addTask(new TaskDef("extract", Extract.class))
77+
.addTask(new TaskDef("transform", Transform.class))
78+
.addTask(new TaskDef("load", Load.class));
8079
}
8180
}

‎java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -54,8 +54,8 @@ import javax.tools.Diagnostic
5454
* containing:
5555
*
5656
* - One inner class per [Builder.Task]-annotated method, implementing [Task].
57-
* - A static `build()` method that constructs the [Dag] and registers those
58-
* inner classes as tasks.
57+
* - A static `build()` method that constructs the [DagDef] and registers those
58+
* inner classes as [TaskDef]s.
5959
*
6060
* [Builder.XCom]-annotated parameters are resolved via `client.getXCom` in the
6161
* generated `execute` body, with the result cast to the parameter's declared
@@ -102,8 +102,8 @@ class BuilderProcessor : AbstractProcessor() {
102102
MethodSpec
103103
.methodBuilder("build")
104104
.addModifiers(Modifier.PUBLIC, Modifier.STATIC)
105-
.returns(ClassName.get(Dag::class.java))
106-
.addStatement($$"var dag = new $T($S)", ClassName.get(Dag::class.java), ann.id.ifBlank { el.simpleName })
105+
.returns(ClassName.get(DagDef::class.java))
106+
.addStatement($$"var dag = new $T($S)", ClassName.get(DagDef::class.java), ann.id.ifBlank { el.simpleName })
107107

108108
for (inner in el.enclosedElements) {
109109
if (inner !is ExecutableElement) continue
@@ -116,7 +116,8 @@ class BuilderProcessor : AbstractProcessor() {
116116
builderClass.addType(task.spec)
117117

118118
buildMethod.addStatement(
119-
$$"dag.addTask($S, $L.class)",
119+
$$"dag.addTask(new $T($S, $L.class))",
120+
ClassName.get(TaskDef::class.java),
120121
ann.id.ifBlank { inner.simpleName },
121122
innerName,
122123
)

‎java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt‎

Lines changed: 26 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -83,16 +83,17 @@ class BuilderTest {
8383
import java.util.Optional;
8484
import org.apache.airflow.sdk.Client;
8585
import org.apache.airflow.sdk.Context;
86-
import org.apache.airflow.sdk.Dag;
86+
import org.apache.airflow.sdk.DagDef;
8787
import org.apache.airflow.sdk.MissingXComException;
8888
import org.apache.airflow.sdk.Task;
89+
import org.apache.airflow.sdk.TaskDef;
8990
9091
public final class TestExampleBuilder {
91-
public static Dag build() {
92-
var dag = new Dag("TestExample");
93-
dag.addTask("t1", T1.class);
94-
dag.addTask("t2", T2.class);
95-
dag.addTask("t3", T3.class);
92+
public static DagDef build() {
93+
var dag = new DagDef("TestExample");
94+
dag.addTask(new TaskDef("t1", T1.class));
95+
dag.addTask(new TaskDef("t2", T2.class));
96+
dag.addTask(new TaskDef("t3", T3.class));
9697
return dag;
9798
}
9899
public static final class T1 implements Task {
@@ -157,14 +158,15 @@ class BuilderTest {
157158
import java.util.Optional;
158159
import org.apache.airflow.sdk.Client;
159160
import org.apache.airflow.sdk.Context;
160-
import org.apache.airflow.sdk.Dag;
161+
import org.apache.airflow.sdk.DagDef;
161162
import org.apache.airflow.sdk.MissingXComException;
162163
import org.apache.airflow.sdk.Task;
164+
import org.apache.airflow.sdk.TaskDef;
163165
164166
public final class TestExampleBuilder {
165-
public static Dag build() {
166-
var dag = new Dag("TestExample");
167-
dag.addTask("t", T.class);
167+
public static DagDef build() {
168+
var dag = new DagDef("TestExample");
169+
dag.addTask(new TaskDef("t", T.class));
168170
return dag;
169171
}
170172
public static final class T implements Task {
@@ -220,14 +222,15 @@ class BuilderTest {
220222
import java.util.Optional;
221223
import org.apache.airflow.sdk.Client;
222224
import org.apache.airflow.sdk.Context;
223-
import org.apache.airflow.sdk.Dag;
225+
import org.apache.airflow.sdk.DagDef;
224226
import org.apache.airflow.sdk.MissingXComException;
225227
import org.apache.airflow.sdk.Task;
228+
import org.apache.airflow.sdk.TaskDef;
226229
227230
public final class TestExampleBuilder {
228-
public static Dag build() {
229-
var dag = new Dag("TestExample");
230-
dag.addTask("t", T.class);
231+
public static DagDef build() {
232+
var dag = new DagDef("TestExample");
233+
dag.addTask(new TaskDef("t", T.class));
231234
return dag;
232235
}
233236
public static final class T implements Task {
@@ -261,8 +264,8 @@ class BuilderTest {
261264
"org.apache.airflow.example.TestExampleBuilder",
262265
"""
263266
package org.apache.airflow.example;
264-
import org.apache.airflow.sdk.Dag;
265-
public final class TestExampleBuilder { public static Dag build() { var dag = new Dag("foo"); return dag; } }
267+
import org.apache.airflow.sdk.DagDef;
268+
public final class TestExampleBuilder { public static DagDef build() { var dag = new DagDef("foo"); return dag; } }
266269
""",
267270
)
268271
}
@@ -284,8 +287,8 @@ class BuilderTest {
284287
"org.apache.airflow.example.Foo",
285288
"""
286289
package org.apache.airflow.example;
287-
import org.apache.airflow.sdk.Dag;
288-
public final class Foo { public static Dag build() { var dag = new Dag("TestExample"); return dag; } }
290+
import org.apache.airflow.sdk.DagDef;
291+
public final class Foo { public static DagDef build() { var dag = new DagDef("TestExample"); return dag; } }
289292
""",
290293
)
291294
}
@@ -313,12 +316,13 @@ class BuilderTest {
313316
import java.lang.Override;
314317
import org.apache.airflow.sdk.Client;
315318
import org.apache.airflow.sdk.Context;
316-
import org.apache.airflow.sdk.Dag;
319+
import org.apache.airflow.sdk.DagDef;
317320
import org.apache.airflow.sdk.Task;
321+
import org.apache.airflow.sdk.TaskDef;
318322
public final class TestExampleBuilder {
319-
public static Dag build() {
320-
var dag = new Dag("TestExample");
321-
dag.addTask("foo", T1.class);
323+
public static DagDef build() {
324+
var dag = new DagDef("TestExample");
325+
dag.addTask(new TaskDef("foo", T1.class));
322326
return dag;
323327
}
324328
public static final class T1 implements Task {

‎java-sdk/scala_spark_example/src/main/scala/org/apache/airflow/example/ScalaSparkExample.scala‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919

2020
package org.apache.airflow.example
2121

22-
import org.apache.airflow.sdk.{Bundle, BundleBuilder, Client, Context, Dag, Server, Task}
22+
import org.apache.airflow.sdk.{Bundle, BundleBuilder, Client, Context, DagDef, Server, Task, TaskDef}
2323
import org.apache.logging.log4j.{LogManager, Logger}
2424
import org.apache.spark.sql.SparkSession
2525
import org.apache.spark.sql.functions.sum
@@ -133,16 +133,16 @@ class SparkLoad extends Task {
133133
}
134134

135135
object ScalaSparkExample {
136-
def build(): Dag =
137-
new Dag(SparkEtl.DagId)
138-
.addTask(SparkEtl.ExtractTaskId, classOf[SparkExtract])
139-
.addTask(SparkEtl.TransformTaskId, classOf[SparkTransform])
140-
.addTask(SparkEtl.LoadTaskId, classOf[SparkLoad])
136+
def build(): DagDef =
137+
new DagDef(SparkEtl.DagId)
138+
.addTask(new TaskDef(SparkEtl.ExtractTaskId, classOf[SparkExtract]))
139+
.addTask(new TaskDef(SparkEtl.TransformTaskId, classOf[SparkTransform]))
140+
.addTask(new TaskDef(SparkEtl.LoadTaskId, classOf[SparkLoad]))
141141
}
142142

143143
/** Bundle entry point served to Airflow's Java coordinator. */
144144
object ScalaSparkBundleBuilder extends BundleBuilder {
145-
override def getDags(): java.lang.Iterable[Dag] = java.util.List.of(ScalaSparkExample.build())
145+
override def getDags(): java.lang.Iterable[DagDef] = java.util.List.of(ScalaSparkExample.build())
146146

147147
def main(args: Array[String]): Unit =
148148
Server.create(args).serve(new Bundle(getDags()))

‎java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Builder.kt‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@ package org.apache.airflow.sdk
2424
*
2525
* This class is not instantiated directly. Its nested annotations drive the
2626
* `BuilderProcessor` annotation processor in the :processor project,
27-
* which generates a `*Builder` class for each class annotated with [Dag].
27+
* which generates a `*Builder` class for each class annotated with [Builder.Dag].
2828
*
2929
* Example:
3030
*
@@ -41,7 +41,7 @@ package org.apache.airflow.sdk
4141
* ```
4242
*
4343
* The processor generates `MyPipelineBuilder.build()`, which returns a
44-
* fully wired-up [Dag] ready to add to a [Bundle].
44+
* fully wired-up [DagDef] ready to add to a [Bundle].
4545
*/
4646
class Builder internal constructor() {
4747
/**

‎java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Bundle.kt‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -20,22 +20,22 @@
2020
package org.apache.airflow.sdk
2121

2222
/**
23-
* An immutable snapshot of all [Dag]s that this JVM process can execute.
23+
* An immutable snapshot of all [DagDef]s that this JVM process can execute.
2424
*
2525
* Build a [Bundle] by implementing [BundleBuilder], then pass it to
2626
* [Server.serve] to start accepting task-execution requests.
2727
*
28-
* @property dags All registered Dags keyed by [Dag.id].
28+
* @property dags All registered Dags keyed by [DagDef.id].
2929
* @throws IllegalArgumentException if any two Dags share the same ID.
3030
*/
3131
class Bundle(
32-
dags: Iterable<Dag>,
32+
dags: Iterable<DagDef>,
3333
) {
34-
internal val dags: Map<String, Dag> = dags.associateByDagId()
34+
internal val dags: Map<String, DagDef> = dags.associateByDagId()
3535
}
3636

37-
private fun Iterable<Dag>.associateByDagId(): Map<String, Dag> {
38-
val dagMap = linkedMapOf<String, Dag>()
37+
private fun Iterable<DagDef>.associateByDagId(): Map<String, DagDef> {
38+
val dagMap = linkedMapOf<String, DagDef>()
3939
for (dag in this) {
4040
require(dagMap.putIfAbsent(dag.id, dag) == null) {
4141
"Dags in bundle have duplicate ID: ${dag.id}"
@@ -45,14 +45,14 @@ private fun Iterable<Dag>.associateByDagId(): Map<String, Dag> {
4545
}
4646

4747
/**
48-
* Entry point for declaring the [Dag]s that this bundle contains.
48+
* Entry point for declaring the [DagDef]s that this bundle contains.
4949
*
5050
* Implement this interface to create a Dag bundle to be served by [Server].
5151
*
5252
* ```java
5353
* public class MyBundleBuilder implements BundleBuilder {
5454
* @Override
55-
* public Iterable<Dag> getDags() {
55+
* public Iterable<DagDef> getDags() {
5656
* return List.of(MyDagBuilder.build());
5757
* }
5858
*
@@ -64,14 +64,14 @@ private fun Iterable<Dag>.associateByDagId(): Map<String, Dag> {
6464
*/
6565
interface BundleBuilder {
6666
/**
67-
* Returns all [Dag]s that belong to this bundle.
67+
* Returns all [DagDef]s that belong to this bundle.
6868
*
6969
* Called once during [build]; Dag IDs must be unique across the returned
7070
* collection.
7171
*
7272
* @throws IllegalArgumentException if any two Dags share the same ID.
7373
*/
74-
fun getDags(): Iterable<Dag>
74+
fun getDags(): Iterable<DagDef>
7575

7676
/**
7777
* Constructs a [Bundle] from the Dags returned by [getDags].

java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Dag.kt renamed to java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/DagDef.kt

Lines changed: 43 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -24,46 +24,72 @@ import kotlin.Throws
2424
/**
2525
* A collection of tasks with directional dependencies.
2626
*
27-
* Create a [Dag] directly and register tasks with [addTask].
27+
* Create a [DagDef] directly and register [TaskDef]s with [addTask].
2828
*
2929
* The [Builder.Dag] annotation should generally be preferred in user code,
3030
* where the annotation processor generates the wiring for you. Only use this
31-
* class directly if you need to do low-level plumbing.
31+
* class directly if you need to do low-level plumbing:
32+
*
33+
* ```java
34+
* var dag = new DagDef("java_etl")
35+
* .addTask(new TaskDef("extract", Extract.class))
36+
* .addTask(new TaskDef("load", Load.class));
37+
* ```
3238
*
3339
* @param id Dag identifier. Must contain only ASCII alphanumeric characters,
3440
* dashes, dots, or underscores; must be unique within a [Bundle].
3541
*
3642
* @see Builder.Dag
3743
*/
38-
class Dag(
44+
class DagDef(
3945
val id: String, // TODO: charset check?
4046
) {
41-
internal var tasks = mutableMapOf<String, Class<out Task>>()
47+
internal val tasks = linkedMapOf<String, TaskDef>()
4248

4349
/**
4450
* Registers a task with this Dag.
4551
*
46-
* The class must have a public no-argument constructor and implement [Task].
47-
* Task IDs must be unique within a Dag.
52+
* A [TaskDef] belongs to at most one [DagDef]; registering the same instance
53+
* with a second Dag, or twice with the same one, fails. Task IDs must be
54+
* unique within a Dag.
4855
*
49-
* @param id Task identifier, unique within this Dag.
50-
* @param definition Class that implements [Task]. Must have a public no-arg
51-
* constructor.
56+
* @param task Task definition to register.
5257
* @return This Dag, for chaining.
53-
* @throws IllegalArgumentException if a task already exists in the Dag with
54-
* the same ID.
58+
* @throws IllegalArgumentException if the task already belongs to a Dag or a
59+
* task with the same ID is already registered.
5560
*/
56-
fun addTask(
57-
id: String,
58-
definition: Class<out Task>,
59-
): Dag {
60-
require(tasks.putIfAbsent(id, definition) == null) {
61-
"Tasks in Dag have duplicate ID: $id"
61+
fun addTask(task: TaskDef): DagDef {
62+
require(task.owner == null) {
63+
"Task '${task.id}' already belongs to Dag '${task.owner?.id}'"
64+
}
65+
require(tasks.putIfAbsent(task.id, task) == null) {
66+
"Tasks in Dag have duplicate ID: ${task.id}"
6267
}
68+
task.owner = this
6369
return this
6470
}
6571
}
6672

73+
/**
74+
* One task definition: its ID and the class that implements it.
75+
*
76+
* ```java
77+
* var extract = new TaskDef("extract", Extract.class);
78+
* ```
79+
*
80+
* @param id Task identifier, unique within a [DagDef].
81+
* @param definition Class that implements [Task]. Must have a public no-arg
82+
* constructor.
83+
*
84+
* @see Builder.Task
85+
*/
86+
class TaskDef(
87+
val id: String,
88+
val definition: Class<out Task>,
89+
) {
90+
internal var owner: DagDef? = null
91+
}
92+
6793
/**
6894
* A single unit of work executed by Airflow.
6995
*

0 commit comments

Comments
 (0)