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.

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(),

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);
   }
 }

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

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 {

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

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() {
   /**

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].

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.
  *

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.

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)

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)
+  }
+}

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