This is an automated email from the ASF dual-hosted git repository.
jason810496 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new ebca0c8edcf Java SDK: Register tasks as first-class TaskDef objects
(#71057)
ebca0c8edcf is described below
commit ebca0c8edcf4d441381c3e66557363f88bef6b77
Author: Jason(Zhe-You) Liu <[email protected]>
AuthorDate: Fri Aug 21 15:00:41 2026 +0800
Java SDK: Register tasks as first-class TaskDef objects (#71057)
* 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.
The KubernetesExecutor lang-SDK example adopts the new names directly rather
than pinning to the old ones: CI now builds those artifacts from the branch
under test (#71527), so the example no longer has to compile against
whatever
the SDK on main happens to expose.
* Java SDK: Clarify task ownership validation
* Java SDK: Simplify manual task registration
Manual interface-based bundles are authored directly by users, so requiring
every call to wrap the task class in TaskDef adds noise to the common case
while the SDK can retain first-class task definitions internally.
---
.../language-sdks/java.rst | 12 +--
.../airflow/example/ExampleBundleBuilder.java | 2 +-
.../airflow/example/InterfaceExampleBuilder.java | 11 ++-
.../org/apache/airflow/sdk/BuilderProcessor.kt | 11 +--
.../kotlin/org/apache/airflow/sdk/BuilderTest.kt | 48 ++++++-----
.../apache/airflow/example/ScalaSparkExample.scala | 8 +-
.../main/kotlin/org/apache/airflow/sdk/Builder.kt | 4 +-
.../main/kotlin/org/apache/airflow/sdk/Bundle.kt | 20 ++---
.../org/apache/airflow/sdk/{Dag.kt => DagDef.kt} | 67 ++++++++++++---
.../org/apache/airflow/sdk/execution/Task.kt | 6 +-
.../kotlin/org/apache/airflow/sdk/BundleTest.kt | 4 +-
.../kotlin/org/apache/airflow/sdk/DagDefTest.kt | 97 ++++++++++++++++++++++
.../org/apache/airflow/sdk/execution/TaskTest.kt | 6 +-
.../airflow/k8sexample/K8sBundleBuilder.java | 2 +-
14 files changed, 221 insertions(+), 77 deletions(-)
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 c8b73b4626a..a3e27f03132 100644
--- a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst
+++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst
@@ -115,7 +115,7 @@ Java entry point
public class Main implements BundleBuilder {
@Override
- public Iterable<Dag> getDags() {
+ public Iterable<DagDef> getDags() {
return List.of(SalesPipelineBuilder.build()); // SalesPipelineBuilder
generated at compile time
}
@@ -203,7 +203,7 @@ Interface-based API
~~~~~~~~~~~~~~~~~~~
Implement the ``Task`` interface directly for full control over how tasks are
registered and how XComs are
-read.
+read. Each task is registered as a ``TaskDef`` on a ``DagDef``.
.. code-block:: java
@@ -224,10 +224,10 @@ Register tasks manually in a ``BundleBuilder``:
public class MyBundle implements BundleBuilder {
@Override
- public Iterable<Dag> getDags() {
- var dag = new Dag("my_dag");
- dag.addTask("fetch", FetchTask.class);
- dag.addTask("process", ProcessTask.class);
+ public Iterable<DagDef> getDags() {
+ var dag = new DagDef("my_dag")
+ .addTask("fetch", FetchTask.class)
+ .addTask("process", ProcessTask.class);
return List.of(dag);
}
}
diff --git
a/java-sdk/example/src/java/org/apache/airflow/example/ExampleBundleBuilder.java
b/java-sdk/example/src/java/org/apache/airflow/example/ExampleBundleBuilder.java
index f63d6c4d743..fa1a8607558 100644
---
a/java-sdk/example/src/java/org/apache/airflow/example/ExampleBundleBuilder.java
+++
b/java-sdk/example/src/java/org/apache/airflow/example/ExampleBundleBuilder.java
@@ -26,7 +26,7 @@ import org.jetbrains.annotations.NotNull;
public class ExampleBundleBuilder implements BundleBuilder {
@NotNull
@Override
- public Iterable<Dag> getDags() {
+ public Iterable<DagDef> getDags() {
return List.of(
InterfaceExampleBuilder.build(),
AnnotationExampleBuilder.build(),
diff --git
a/java-sdk/example/src/java/org/apache/airflow/example/InterfaceExampleBuilder.java
b/java-sdk/example/src/java/org/apache/airflow/example/InterfaceExampleBuilder.java
index 1c536c3cbf2..78e78eed269 100644
---
a/java-sdk/example/src/java/org/apache/airflow/example/InterfaceExampleBuilder.java
+++
b/java-sdk/example/src/java/org/apache/airflow/example/InterfaceExampleBuilder.java
@@ -71,11 +71,10 @@ public class InterfaceExampleBuilder {
}
}
- 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("extract", Extract.class)
+ .addTask("transform", Transform.class)
+ .addTask("load", Load.class);
}
}
diff --git
a/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt
b/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt
index 11202ffcdb4..56cbf1e76ad 100644
---
a/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt
+++
b/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt
@@ -54,8 +54,8 @@ import javax.tools.Diagnostic
* containing:
*
* - One inner class per [Builder.Task]-annotated method, implementing [Task].
- * - A static `build()` method that constructs the [Dag] and registers those
- * inner classes as tasks.
+ * - A static `build()` method that constructs the [DagDef] and registers those
+ * inner classes as [TaskDef]s.
*
* [Builder.XCom]-annotated parameters are resolved via `client.getXCom` in the
* generated `execute` body, with the result cast to the parameter's declared
@@ -102,8 +102,8 @@ class BuilderProcessor : AbstractProcessor() {
MethodSpec
.methodBuilder("build")
.addModifiers(Modifier.PUBLIC, Modifier.STATIC)
- .returns(ClassName.get(Dag::class.java))
- .addStatement($$"var dag = new $T($S)",
ClassName.get(Dag::class.java), ann.id.ifBlank { el.simpleName })
+ .returns(ClassName.get(DagDef::class.java))
+ .addStatement($$"var dag = new $T($S)",
ClassName.get(DagDef::class.java), ann.id.ifBlank { el.simpleName })
for (inner in el.enclosedElements) {
if (inner !is ExecutableElement) continue
@@ -116,7 +116,8 @@ class BuilderProcessor : AbstractProcessor() {
builderClass.addType(task.spec)
buildMethod.addStatement(
- $$"dag.addTask($S, $L.class)",
+ $$"dag.addTask(new $T($S, $L.class))",
+ ClassName.get(TaskDef::class.java),
ann.id.ifBlank { inner.simpleName },
innerName,
)
diff --git
a/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt
b/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt
index 3e28e1b009a..6a08979c561 100644
--- a/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt
+++ b/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt
@@ -83,16 +83,17 @@ class BuilderTest {
import java.util.Optional;
import org.apache.airflow.sdk.Client;
import org.apache.airflow.sdk.Context;
- import org.apache.airflow.sdk.Dag;
+ import org.apache.airflow.sdk.DagDef;
import org.apache.airflow.sdk.MissingXComException;
import org.apache.airflow.sdk.Task;
+ import org.apache.airflow.sdk.TaskDef;
public final class TestExampleBuilder {
- public static Dag build() {
- var dag = new Dag("TestExample");
- dag.addTask("t1", T1.class);
- dag.addTask("t2", T2.class);
- dag.addTask("t3", T3.class);
+ public static DagDef build() {
+ var dag = new DagDef("TestExample");
+ dag.addTask(new TaskDef("t1", T1.class));
+ dag.addTask(new TaskDef("t2", T2.class));
+ dag.addTask(new TaskDef("t3", T3.class));
return dag;
}
public static final class T1 implements Task {
@@ -157,14 +158,15 @@ class BuilderTest {
import java.util.Optional;
import org.apache.airflow.sdk.Client;
import org.apache.airflow.sdk.Context;
- import org.apache.airflow.sdk.Dag;
+ import org.apache.airflow.sdk.DagDef;
import org.apache.airflow.sdk.MissingXComException;
import org.apache.airflow.sdk.Task;
+ import org.apache.airflow.sdk.TaskDef;
public final class TestExampleBuilder {
- public static Dag build() {
- var dag = new Dag("TestExample");
- dag.addTask("t", T.class);
+ public static DagDef build() {
+ var dag = new DagDef("TestExample");
+ dag.addTask(new TaskDef("t", T.class));
return dag;
}
public static final class T implements Task {
@@ -220,14 +222,15 @@ class BuilderTest {
import java.util.Optional;
import org.apache.airflow.sdk.Client;
import org.apache.airflow.sdk.Context;
- import org.apache.airflow.sdk.Dag;
+ import org.apache.airflow.sdk.DagDef;
import org.apache.airflow.sdk.MissingXComException;
import org.apache.airflow.sdk.Task;
+ import org.apache.airflow.sdk.TaskDef;
public final class TestExampleBuilder {
- public static Dag build() {
- var dag = new Dag("TestExample");
- dag.addTask("t", T.class);
+ public static DagDef build() {
+ var dag = new DagDef("TestExample");
+ dag.addTask(new TaskDef("t", T.class));
return dag;
}
public static final class T implements Task {
@@ -261,8 +264,8 @@ class BuilderTest {
"org.apache.airflow.example.TestExampleBuilder",
"""
package org.apache.airflow.example;
- import org.apache.airflow.sdk.Dag;
- public final class TestExampleBuilder { public static Dag build() {
var dag = new Dag("foo"); return dag; } }
+ import org.apache.airflow.sdk.DagDef;
+ public final class TestExampleBuilder { public static DagDef build()
{ var dag = new DagDef("foo"); return dag; } }
""",
)
}
@@ -284,8 +287,8 @@ class BuilderTest {
"org.apache.airflow.example.Foo",
"""
package org.apache.airflow.example;
- import org.apache.airflow.sdk.Dag;
- public final class Foo { public static Dag build() { var dag = new
Dag("TestExample"); return dag; } }
+ import org.apache.airflow.sdk.DagDef;
+ public final class Foo { public static DagDef build() { var dag = new
DagDef("TestExample"); return dag; } }
""",
)
}
@@ -313,12 +316,13 @@ class BuilderTest {
import java.lang.Override;
import org.apache.airflow.sdk.Client;
import org.apache.airflow.sdk.Context;
- import org.apache.airflow.sdk.Dag;
+ import org.apache.airflow.sdk.DagDef;
import org.apache.airflow.sdk.Task;
+ import org.apache.airflow.sdk.TaskDef;
public final class TestExampleBuilder {
- public static Dag build() {
- var dag = new Dag("TestExample");
- dag.addTask("foo", T1.class);
+ public static DagDef build() {
+ var dag = new DagDef("TestExample");
+ dag.addTask(new TaskDef("foo", T1.class));
return dag;
}
public static final class T1 implements Task {
diff --git
a/java-sdk/scala_spark_example/src/main/scala/org/apache/airflow/example/ScalaSparkExample.scala
b/java-sdk/scala_spark_example/src/main/scala/org/apache/airflow/example/ScalaSparkExample.scala
index ae779bf0213..f7f32bad8e2 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}
import org.apache.logging.log4j.{LogManager, Logger}
import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.functions.sum
@@ -133,8 +133,8 @@ class SparkLoad extends Task {
}
object ScalaSparkExample {
- def build(): Dag =
- new Dag(SparkEtl.DagId)
+ def build(): DagDef =
+ new DagDef(SparkEtl.DagId)
.addTask(SparkEtl.ExtractTaskId, classOf[SparkExtract])
.addTask(SparkEtl.TransformTaskId, classOf[SparkTransform])
.addTask(SparkEtl.LoadTaskId, classOf[SparkLoad])
@@ -142,7 +142,7 @@ object ScalaSparkExample {
/** Bundle entry point served to Airflow's Java coordinator. */
object ScalaSparkBundleBuilder extends BundleBuilder {
- override def getDags(): java.lang.Iterable[Dag] =
java.util.List.of(ScalaSparkExample.build())
+ override def getDags(): java.lang.Iterable[DagDef] =
java.util.List.of(ScalaSparkExample.build())
def main(args: Array[String]): Unit =
Server.create(args).serve(new Bundle(getDags()))
diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Builder.kt
b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Builder.kt
index 3a5b84d2daf..9dbfbfbefc3 100644
--- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Builder.kt
+++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Builder.kt
@@ -24,7 +24,7 @@ package org.apache.airflow.sdk
*
* This class is not instantiated directly. Its nested annotations drive the
* `BuilderProcessor` annotation processor in the :processor project,
- * which generates a `*Builder` class for each class annotated with [Dag].
+ * which generates a `*Builder` class for each class annotated with
[Builder.Dag].
*
* Example:
*
@@ -41,7 +41,7 @@ package org.apache.airflow.sdk
* ```
*
* The processor generates `MyPipelineBuilder.build()`, which returns a
- * fully wired-up [Dag] ready to add to a [Bundle].
+ * fully wired-up [DagDef] ready to add to a [Bundle].
*/
class Builder internal constructor() {
/**
diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Bundle.kt
b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Bundle.kt
index 677ec48eb93..6cd549270a8 100644
--- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Bundle.kt
+++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Bundle.kt
@@ -20,22 +20,22 @@
package org.apache.airflow.sdk
/**
- * An immutable snapshot of all [Dag]s that this JVM process can execute.
+ * An immutable snapshot of all [DagDef]s that this JVM process can execute.
*
* Build a [Bundle] by implementing [BundleBuilder], then pass it to
* [Server.serve] to start accepting task-execution requests.
*
- * @property dags All registered Dags keyed by [Dag.id].
+ * @property dags All registered Dags keyed by [DagDef.id].
* @throws IllegalArgumentException if any two Dags share the same ID.
*/
class Bundle(
- dags: Iterable<Dag>,
+ dags: Iterable<DagDef>,
) {
- internal val dags: Map<String, Dag> = dags.associateByDagId()
+ internal val dags: Map<String, DagDef> = dags.associateByDagId()
}
-private fun Iterable<Dag>.associateByDagId(): Map<String, Dag> {
- val dagMap = linkedMapOf<String, Dag>()
+private fun Iterable<DagDef>.associateByDagId(): Map<String, DagDef> {
+ val dagMap = linkedMapOf<String, DagDef>()
for (dag in this) {
require(dagMap.putIfAbsent(dag.id, dag) == null) {
"Dags in bundle have duplicate ID: ${dag.id}"
@@ -45,14 +45,14 @@ private fun Iterable<Dag>.associateByDagId(): Map<String,
Dag> {
}
/**
- * 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<Dag> getDags() {
+ * public Iterable<DagDef> getDags() {
* return List.of(MyDagBuilder.build());
* }
*
@@ -64,14 +64,14 @@ private fun Iterable<Dag>.associateByDagId(): Map<String,
Dag> {
*/
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<Dag>
+ fun getDags(): Iterable<DagDef>
/**
* Constructs a [Bundle] from the Dags returned by [getDags].
diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Dag.kt
b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/DagDef.kt
similarity index 61%
rename from java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Dag.kt
rename to java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/DagDef.kt
index c9985803741..e3ba700a0fa 100644
--- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Dag.kt
+++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/DagDef.kt
@@ -24,46 +24,87 @@ import kotlin.Throws
/**
* A collection of tasks with directional dependencies.
*
- * Create a [Dag] directly and register tasks with [addTask].
+ * Create a [DagDef] directly and register [TaskDef]s with [addTask].
*
* The [Builder.Dag] annotation should generally be preferred in user code,
* where the annotation processor generates the wiring for you. Only use this
- * class directly if you need to do low-level plumbing.
+ * class directly if you need to do low-level plumbing:
+ *
+ * ```java
+ * var dag = new DagDef("java_etl")
+ * .addTask("extract", Extract.class)
+ * .addTask("load", Load.class);
+ * ```
*
* @param id Dag identifier. Must contain only ASCII alphanumeric characters,
* dashes, dots, or underscores; must be unique within a [Bundle].
*
* @see Builder.Dag
*/
-class Dag(
+class DagDef(
val id: String, // TODO: charset check?
) {
- internal var tasks = mutableMapOf<String, Class<out Task>>()
+ internal val tasks = linkedMapOf<String, TaskDef>()
/**
- * 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.
+ * Registers a task from its ID and implementation class.
*
* @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.
+ * @throws IllegalArgumentException if a task with the same ID is already
+ * registered.
*/
fun addTask(
id: String,
definition: Class<out Task>,
- ): Dag {
- require(tasks.putIfAbsent(id, definition) == null) {
- "Tasks in Dag have duplicate ID: $id"
+ ): DagDef = addTask(TaskDef(id, definition))
+
+ /**
+ * 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.
+ *
+ * @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 {
+ task.owner?.let { owner ->
+ throw IllegalArgumentException("Task '${task.id}' already belongs to Dag
'${owner.id}'")
}
+ require(tasks.putIfAbsent(task.id, task) == null) {
+ "Tasks in Dag have duplicate ID: ${task.id}"
+ }
+ task.owner = this
return this
}
}
+/**
+ * One task definition: its ID and the class that implements it.
+ *
+ * ```java
+ * var extract = new TaskDef("extract", Extract.class);
+ * ```
+ *
+ * @param id Task identifier, unique within a [DagDef].
+ * @param definition Class that implements [Task]. Must have a public no-arg
+ * constructor.
+ *
+ * @see Builder.Task
+ */
+class TaskDef(
+ val id: String,
+ val definition: Class<out Task>,
+) {
+ internal var owner: DagDef? = null
+}
+
/**
* A single unit of work executed by Airflow.
*
diff --git
a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Task.kt
b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Task.kt
index 3685644c0b7..bddc5256f55 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
@@ -71,9 +71,11 @@ internal object TaskRunner {
request: StartupDetails,
client: Client,
): Any {
- val task = bundle.dags[request.ti.dagId]?.tasks[request.ti.taskId] ?:
return TaskResult.of(TaskState.State.REMOVED)
+ val definition =
+ bundle.dags[request.ti.dagId]?.tasks[request.ti.taskId]?.definition
+ ?: return TaskResult.of(TaskState.State.REMOVED)
return try {
-
task.getDeclaredConstructor().newInstance().execute(Context.from(request),
client)
+
definition.getDeclaredConstructor().newInstance().execute(Context.from(request),
client)
TaskResult.success()
} catch (e: CancellationException) {
throw e // Let coroutine cancellation propagate so the task coroutine
unwinds.
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 0e4afb1894a..57050754e88 100644
--- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/BundleTest.kt
+++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/BundleTest.kt
@@ -27,7 +27,7 @@ internal class BundleTest {
@Test
@DisplayName("Should index dags by dagId")
fun shouldIndexDagsByDagId() {
- val dag = Dag("dag")
+ val dag = DagDef("dag")
val bundle = Bundle(listOf(dag))
@@ -39,7 +39,7 @@ internal class BundleTest {
fun shouldRejectDuplicateDagIds() {
val error =
Assertions.assertThrows(IllegalArgumentException::class.java) {
- Bundle(listOf(Dag("dag"), Dag("dag")))
+ Bundle(listOf(DagDef("dag"), DagDef("dag")))
}
Assertions.assertEquals("Dags in bundle have duplicate ID: dag",
error.message)
diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/DagDefTest.kt
b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/DagDefTest.kt
new file mode 100644
index 00000000000..a23f76bad82
--- /dev/null
+++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/DagDefTest.kt
@@ -0,0 +1,97 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.airflow.sdk
+
+import org.junit.jupiter.api.Assertions
+import org.junit.jupiter.api.DisplayName
+import org.junit.jupiter.api.Test
+
+internal class DagDefTest {
+ private class NoOp : Task {
+ override fun execute(
+ context: Context,
+ client: Client,
+ ) = Unit
+ }
+
+ @Test
+ @DisplayName("Should index tasks by taskId in registration order")
+ fun shouldIndexTasksByTaskId() {
+ val extract = TaskDef("extract", NoOp::class.java)
+ val load = TaskDef("load", NoOp::class.java)
+
+ val dag = DagDef("dag").addTask(extract).addTask(load)
+
+ Assertions.assertEquals(listOf("extract", "load"), dag.tasks.keys.toList())
+ Assertions.assertEquals(mapOf("extract" to extract, "load" to load),
dag.tasks)
+ }
+
+ @Test
+ @DisplayName("Should register a task from its id and implementation class")
+ fun shouldRegisterTaskFromIdAndImplementationClass() {
+ val dag = DagDef("dag").addTask("extract", NoOp::class.java)
+
+ val task = dag.tasks.getValue("extract")
+ Assertions.assertEquals("extract", task.id)
+ Assertions.assertEquals(NoOp::class.java, task.definition)
+ Assertions.assertEquals(dag, task.owner)
+ }
+
+ @Test
+ @DisplayName("Should reject duplicate task ids")
+ fun shouldRejectDuplicateTaskIds() {
+ val dag = DagDef("dag").addTask(TaskDef("extract", NoOp::class.java))
+
+ val error =
+ Assertions.assertThrows(IllegalArgumentException::class.java) {
+ dag.addTask(TaskDef("extract", NoOp::class.java))
+ }
+
+ Assertions.assertEquals("Tasks in Dag have duplicate ID: extract",
error.message)
+ }
+
+ @Test
+ @DisplayName("Should reject a task already registered with another dag")
+ fun shouldRejectTaskOwnedByAnotherDag() {
+ val extract = TaskDef("extract", NoOp::class.java)
+ DagDef("first").addTask(extract)
+
+ val error =
+ Assertions.assertThrows(IllegalArgumentException::class.java) {
+ DagDef("second").addTask(extract)
+ }
+
+ Assertions.assertEquals("Task 'extract' already belongs to Dag 'first'",
error.message)
+ }
+
+ @Test
+ @DisplayName("Should reject the same task registered twice with one dag")
+ fun shouldRejectTaskRegisteredTwice() {
+ val extract = TaskDef("extract", NoOp::class.java)
+ val dag = DagDef("dag").addTask(extract)
+
+ val error =
+ Assertions.assertThrows(IllegalArgumentException::class.java) {
+ dag.addTask(extract)
+ }
+
+ Assertions.assertEquals("Task 'extract' already belongs to Dag 'dag'",
error.message)
+ }
+}
diff --git
a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskTest.kt
b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskTest.kt
index 42a083b94aa..5d303a3dd81 100644
--- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskTest.kt
+++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskTest.kt
@@ -22,8 +22,9 @@ package org.apache.airflow.sdk.execution
import org.apache.airflow.sdk.Bundle
import org.apache.airflow.sdk.Client
import org.apache.airflow.sdk.Context
-import org.apache.airflow.sdk.Dag
+import org.apache.airflow.sdk.DagDef
import org.apache.airflow.sdk.Task
+import org.apache.airflow.sdk.TaskDef
import org.apache.airflow.sdk.execution.comm.BundleInfo
import org.apache.airflow.sdk.execution.comm.DagRun
import org.apache.airflow.sdk.execution.comm.RetryTask
@@ -88,8 +89,7 @@ class TaskTest {
taskId: String,
taskClass: Class<out Task>,
): Bundle {
- val dag = Dag("test_dag")
- dag.addTask(taskId, taskClass)
+ val dag = DagDef("test_dag").addTask(TaskDef(taskId, taskClass))
return Bundle(listOf(dag))
}
diff --git
a/kubernetes-tests/lang_sdk/java_example/src/java/org/apache/airflow/k8sexample/K8sBundleBuilder.java
b/kubernetes-tests/lang_sdk/java_example/src/java/org/apache/airflow/k8sexample/K8sBundleBuilder.java
index 8333567da83..c3b511b898c 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 @@ import org.jetbrains.annotations.NotNull;
public class K8sBundleBuilder implements BundleBuilder {
@NotNull
@Override
- public Iterable<Dag> getDags() {
+ public Iterable<DagDef> getDags() {
return List.of(CombinedExampleBuilder.build());
}
