From ebca0c8edcf4d441381c3e66557363f88bef6b77 Mon Sep 17 00:00:00 2001 From: "Jason(Zhe-You) Liu" <68415893+jason810496@users.noreply.github.com> Date: Fri, 21 Aug 2026 15:00:41 +0800 Subject: [PATCH] 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>`, 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 +- .../example/InterfaceExampleBuilder.java | 11 +-- .../apache/airflow/sdk/BuilderProcessor.kt | 11 ++- .../org/apache/airflow/sdk/BuilderTest.kt | 48 ++++----- .../airflow/example/ScalaSparkExample.scala | 8 +- .../kotlin/org/apache/airflow/sdk/Builder.kt | 4 +- .../kotlin/org/apache/airflow/sdk/Bundle.kt | 20 ++-- .../apache/airflow/sdk/{Dag.kt => DagDef.kt} | 67 ++++++++++--- .../org/apache/airflow/sdk/execution/Task.kt | 6 +- .../org/apache/airflow/sdk/BundleTest.kt | 4 +- .../org/apache/airflow/sdk/DagDefTest.kt | 97 +++++++++++++++++++ .../apache/airflow/sdk/execution/TaskTest.kt | 6 +- .../airflow/k8sexample/K8sBundleBuilder.java | 2 +- 14 files changed, 221 insertions(+), 77 deletions(-) rename java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/{Dag.kt => DagDef.kt} (61%) create mode 100644 java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/DagDefTest.kt 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 getDags() { + public Iterable 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 getDags() { - var dag = new Dag("my_dag"); - dag.addTask("fetch", FetchTask.class); - dag.addTask("process", ProcessTask.class); + public Iterable 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 getDags() { + public Iterable 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, + dags: Iterable, ) { - internal val dags: Map = dags.associateByDagId() + internal val dags: Map = dags.associateByDagId() } -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 +45,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 +64,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/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>() + internal val tasks = linkedMapOf() /** - * 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, - ): 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, +) { + 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, ): 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 getDags() { + public Iterable getDags() { return List.of(CombinedExampleBuilder.build()); }