Java SDK: Draw a Java-authored Dag's edges with before and after (#73594)

* Java SDK: Draw a Java-authored Dag's edges with before and after A Dag authored in Java had no way to say what runs after what. Without a Python Dag file to hold the graph, every task was a root, so a Dag that Java owns end to end could not be expressed at all. The verbs are before and after, matching the TypeScript SDK, and leaving then free for the conditional control flow that reads as a then-branch. They are the only way to draw an edge, so one graph has one spelling however it was built. * Add notes on delibrate implementation choices --------- Co-authored-by: Tzu-ping Chung <tp@astronomer.io>
2026-11-04 07:12:48 +00:00 · 2026-09-26 20:39:34 +08:00
1 parent 7e5c580cd6
commit 77a1953cd4
@@ -26,7 +26,8 @@ public class ExampleBundleBuilder {
    return new Bundle()
        .register(InterfaceExampleBuilder.build())
        .register(AnnotationExample.class)
-        .register(XComCastingExample.class);
+        .register(XComCastingExample.class)
+        .register(org.apache.airflow.example.nativedag.InterfaceExample.build());
  }

  public static void main(String[] args) {
@@ -0,0 +1,68 @@
+/*
+ * 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.
+ */
+
+// "native" is a Java keyword, so the native-Dag examples live in "nativedag".
+package org.apache.airflow.example.nativedag;
+
+import static java.lang.System.Logger.Level.INFO;
+
+import org.apache.airflow.sdk.*;
+
+// A Dag defined entirely in Java, interface-style: no Python stub file
+// describes it. dag.task registers a task as it creates it and hands back the
+// handle, and `before`/`after` wire the graph -- Java's spelling of `>>` and `<<`.
+public class InterfaceExample {
+  private static final System.Logger log = System.getLogger(InterfaceExample.class.getName());
+
+  public static class Extract implements Task {
+    @Override
+    public void execute(Context context, Client client) {
+      log.log(INFO, "Extracting a value");
+      client.setXCom(42L);
+    }
+  }
+
+  public static class Transform implements Task {
+    @Override
+    public void execute(Context context, Client client) {
+      var extracted = ((Number) client.getXCom("extract")).longValue();
+      log.log(INFO, "Transforming {0}", extracted);
+      client.setXCom(extracted * 2);
+    }
+  }
+
+  public static class Load implements Task {
+    @Override
+    public void execute(Context context, Client client) {
+      var transformed = client.getXCom("transform");
+      log.log(INFO, "Loaded {0}", transformed);
+    }
+  }
+
+  public static DagDef build() {
+    var dag = new DagDef("java_native_interface_example");
+
+    var extract = dag.task("extract", Extract.class);
+    var transform = dag.task("transform", Transform.class);
+    var load = dag.task("load", Load.class);
+
+    transform.after(extract).before(load);
+    return dag;
+  }
+}
@@ -0,0 +1,53 @@
+/*
+ * 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
+
+/**
+ * A value a task can be given, of which [TaskRef] — the output of an upstream
+ * task — is the only form so far.
+ *
+ * @param T Type of the value.
+ */
+sealed class Arg<T>
+
+/**
+ * The output of a registered task, and the task's place in the flow.
+ *
+ * [Deps.Flow.before] and [Deps.Flow.after] wire an ordering-only edge from
+ * this task, where nothing flows but the sequence.
+ *
+ * @param T Return type of the task this handle refers to.
+ */
+class TaskRef<T> internal constructor(
+  internal val def: TaskDef,
+) : Arg<T>(),
+  Deps.Flow {
+  override fun nodes(): List<TaskDef> = listOf(def)
+
+  override fun before(vararg next: Deps.Flow): TaskRef<T> {
+    super<Deps.Flow>.before(*next)
+    return this
+  }
+
+  override fun after(vararg previous: Deps.Flow): TaskRef<T> {
+    super<Deps.Flow>.after(*previous)
+    return this
+  }
+}
@@ -29,7 +29,9 @@ import org.apache.airflow.sdk.internal.registrarName
 * the running task.
 *
 * @property dags Dags declared in Java, keyed by [DagDef.id].
- * @throws IllegalArgumentException if any two Dags share the same ID.
+ * @throws IllegalArgumentException if any two Dags share the same ID, if a
+ *    task depends on an upstream that is not registered in its own Dag, or if
+ *    the dependencies of a Dag contain a cycle.
 */
 class Bundle(
  dags: Iterable<DagDef>,
@@ -40,6 +42,8 @@ class Bundle(
  /** Dags the Python file owns, holding the task handlers registered for them. */
  internal val taskHandlers = linkedMapOf<String, DagDef>()

+  // This only guards the serve boundary, not registers racing each other. This is fine since
+  // we only encourage one sync register() chain; single-threaded by contract.
  @Volatile
  private var served = false

@@ -53,9 +57,14 @@ class Bundle(
  /**
   * Registers a Dag.
   *
+   * The Dag is checked as it is registered, so a bad edge fails here rather
+   * than at the first task run.
+   *
   * @return This bundle, for chaining.
-   * @throws IllegalArgumentException if another Dag shares its ID, or task
-   *    handlers are already registered against it.
+   * @throws IllegalArgumentException if another Dag shares its ID, task
+   *    handlers are already registered against it, a task depends on an
+   *    upstream not registered in the same Dag, or the dependencies contain a
+   *    cycle.
   * @throws IllegalStateException if [Server.serve] has already been called.
   */
  fun register(dag: DagDef): Bundle {
@@ -64,6 +73,15 @@ class Bundle(
      "Dag '${dag.id}' already has registered task handlers; a Dag declared in Java owns its " +
        "own tasks, so one Dag ID cannot have both"
    }
+    for ((taskId, def) in dag.tasks) {
+      for (upstream in def.upstreams) {
+        require(dag.tasks[upstream.id] === upstream) {
+          "Task '$taskId' in Dag '${dag.id}' depends on task '${upstream.id}' " +
+            "that is not registered in the same Dag"
+        }
+      }
+    }
+    checkNoCycle(dag)
    require(dags.putIfAbsent(dag.id, dag) == null) {
      "Dags in bundle have duplicate ID: ${dag.id}"
    }
@@ -101,8 +119,8 @@ class Bundle(
   * Registers one task implementation against a Dag the Python file owns, for
   * a task with no annotation to read the ids from.
   *
-   * The Dag is created on first use: a stub-backed Dag exists only so the
-   * runtime can find the task, and its graph lives in the Python Dag file.
+   * The Dag is created on first use and holds only the tasks registered
+   * here; its graph lives in the Python Dag file.
   *
   * @param dagId Dag ID as declared in the Python Dag file.
   * @param taskId Task ID as declared by the `@task.stub` function.
@@ -144,6 +162,25 @@ class Bundle(
  private fun checkOpen() = check(!served) { "Server.serve has already been called; register everything before serve" }
 }

+// Reject cycles produced by before and after at registration time. This is (non-tailrec-eligible)
+// recursive and could blow up with deep dependency chains. I kept the recursive implementation
+// for readability since the scenario is unlikely; feel free to rewrite if it blows up for you.
+private fun checkNoCycle(dag: DagDef) {
+  val visiting = mutableSetOf<String>()
+  val done = mutableSetOf<String>()
+
+  fun visit(def: TaskDef) {
+    if (def.id in done) return
+    require(visiting.add(def.id)) {
+      "Task dependencies in Dag '${dag.id}' contain a cycle involving task '${def.id}'"
+    }
+    def.upstreams.forEach(::visit)
+    visiting -= def.id
+    done += def.id
+  }
+  dag.tasks.values.forEach(::visit)
+}
+
 /**
 * Entry point for declaring the [DagDef]s that this bundle contains.
 *
@@ -32,9 +32,9 @@ import kotlin.Throws
 * 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);
+ * var dag = new DagDef("java_etl");
+ * var extract = dag.task("extract", Extract.class);
+ * extract.before(dag.task("load", Load.class));
 * ```
 *
 * @param id Dag identifier. Must contain only ASCII alphanumeric characters,
@@ -62,12 +62,41 @@ class DagDef(
    definition: Class<out Task>,
  ): DagDef = addTask(TaskDef(id, definition))

+  /**
+   * Creates a task, registers it, and hands back its handle — so there is no
+   * second `addTask` call to forget, and the handle is ready to wire edges
+   * with [Deps.Flow.before].
+   *
+   * ```java
+   * var extract = dag.task("extract", Extract.class);
+   * var load = dag.task("load", Load.class);
+   * extract.before(load);
+   * ```
+   *
+   * @param id Task identifier, unique within this Dag.
+   * @param definition Class that implements [Task]. Must have a public no-arg
+   *    constructor.
+   * @return The handle representing this task.
+   * @throws IllegalArgumentException if a task with the same ID is already
+   *    registered.
+   */
+  fun <T> task(
+    id: String,
+    definition: Class<out Task>,
+  ): TaskRef<T> {
+    val def = TaskDef(id, definition)
+    addTask(def)
+    return TaskRef(def)
+  }
+
  /**
   * 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.
+   * unique within a Dag. Tasks named as upstreams by [Deps.Flow.before] or
+   * [Deps.Flow.after] must be registered with the same Dag by the time it is
+   * added to a [Bundle].
   *
   * @param task Task definition to register.
   * @return This Dag, for chaining.
@@ -87,10 +116,14 @@ class DagDef(
 }

 /**
- * One task definition: its ID and the class that implements it.
+ * One task definition: its ID, the class that implements it, and its upstream
+ * dependencies.
+ *
+ * Edges are drawn on the handles that [DagDef.task] returns, not here:
 *
 * ```java
- * var extract = new TaskDef("extract", Extract.class);
+ * var dag = new DagDef("java_etl");
+ * dag.addTask(new TaskDef("extract", Extract.class));
 * ```
 *
 * @param id Task identifier, unique within a [DagDef].
@@ -110,7 +143,14 @@ class TaskDef(
    validateTaskInput(definition)
  }

+  internal val upstreams = linkedSetOf<TaskDef>()
  internal var owner: DagDef? = null
+
+  /** Records that this task runs after [upstreams], backing [Deps.Flow.before] and [Deps.Flow.after]. */
+  internal fun dependsOn(vararg upstreams: TaskDef): TaskDef {
+    this.upstreams += upstreams
+    return this
+  }
 }

 /**
@@ -0,0 +1,93 @@
+/*
+ * 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
+
+/** Vocabulary for declaring a Dag's task graph in Java. */
+interface Deps {
+  /**
+   * A point in the task graph: one task, or a set of them.
+   *
+   * [Flow] is Java's spelling of Python's `>>` and `<<`, for a dependency
+   * where nothing flows but the ordering. An edge that carries a value is
+   * declared by passing the upstream's handle instead.
+   */
+  interface Flow {
+    /** The tasks at this point in the flow. */
+    fun nodes(): List<TaskDef>
+
+    /**
+     * Runs the tasks here before each of [next], carrying no value — Java's
+     * spelling of Python's `>>`.
+     *
+     * ```java
+     * loaded.before(cleaned, notified); // load >> [cleanup, notify]
+     * ```
+     *
+     * Variadic, so one call fans out, and it returns its own receiver: a
+     * fan-out has no single next task to hand back. Declaring an edge that
+     * already exists changes nothing.
+     *
+     * @param next Tasks that run after the ones here.
+     * @return This point in the flow.
+     */
+    fun before(vararg next: Flow): Flow {
+      val upstreams = nodes()
+      next.flatMap { it.nodes() }.forEach { downstream -> upstreams.forEach { downstream.dependsOn(it) } }
+      return this
+    }
+
+    /**
+     * Runs the tasks here after each of [previous], carrying no value —
+     * Python's `<<`.
+     *
+     * ```java
+     * cleaned.after(loaded, transformed); // [load, transform] >> cleanup
+     * ```
+     *
+     * @param previous Tasks that run before the ones here.
+     * @return This point in the flow.
+     */
+    fun after(vararg previous: Flow): Flow {
+      val downstreams = nodes()
+      previous.flatMap { it.nodes() }.forEach { upstream -> downstreams.forEach { it.dependsOn(upstream) } }
+      return this
+    }
+
+    companion object {
+      /**
+       * Treats several tasks as one point in the flow, so a single call draws
+       * every edge between two sets:
+       *
+       * ```java
+       * Flow.of(a, b).before(c, d); // [a, b] >> [c, d]
+       * ```
+       */
+      @JvmStatic
+      fun of(vararg flows: Flow): Flow = FlowSet(flows.flatMap { it.nodes() })
+    }
+  }
+}
+
+/** Several tasks as one point in the flow, which no single [TaskRef] can represent. */
+internal class FlowSet(
+  private val nodes: List<TaskDef>,
+) : Deps.Flow {
+  override fun nodes(): List<TaskDef> = nodes
+}
@@ -55,6 +55,85 @@ internal class BundleTest {
    Assertions.assertEquals("Dags in bundle have duplicate ID: dag", error.message)
  }

+  @Test
+  @DisplayName("Should reject a task depending on an unregistered upstream")
+  fun shouldRejectUnregisteredUpstream() {
+    val missing = TaskDef("missing", NoOp::class.java)
+    val dag = DagDef("dag").addTask(TaskDef("t", NoOp::class.java).dependsOn(missing))
+
+    val error =
+      Assertions.assertThrows(IllegalArgumentException::class.java) {
+        Bundle(listOf(dag))
+      }
+
+    Assertions.assertEquals(
+      "Task 't' in Dag 'dag' depends on task 'missing' that is not registered in the same Dag",
+      error.message,
+    )
+  }
+
+  @Test
+  @DisplayName("Should reject a task depending on a task registered in another dag")
+  fun shouldRejectUpstreamFromAnotherDag() {
+    val foreign = TaskDef("u", NoOp::class.java)
+    val other = DagDef("other").addTask(foreign)
+    val dag = DagDef("dag").addTask(TaskDef("t", NoOp::class.java).dependsOn(foreign))
+
+    val error =
+      Assertions.assertThrows(IllegalArgumentException::class.java) {
+        Bundle(listOf(other, dag))
+      }
+
+    Assertions.assertEquals(
+      "Task 't' in Dag 'dag' depends on task 'u' that is not registered in the same Dag",
+      error.message,
+    )
+  }
+
+  @Test
+  @DisplayName("Should reject dependency cycles")
+  fun shouldRejectDependencyCycle() {
+    val a = TaskDef("a", NoOp::class.java)
+    val b = TaskDef("b", NoOp::class.java)
+    a.dependsOn(b)
+    b.dependsOn(a)
+    val dag = DagDef("dag").addTask(a).addTask(b)
+
+    val error =
+      Assertions.assertThrows(IllegalArgumentException::class.java) {
+        Bundle(listOf(dag))
+      }
+
+    Assertions.assertEquals(
+      "Task dependencies in Dag 'dag' contain a cycle involving task 'a'",
+      error.message,
+    )
+  }
+
+  @Test
+  @DisplayName("Should accept a diamond-shaped dependency graph")
+  fun shouldAcceptDiamondGraph() {
+    val root = TaskDef("root", NoOp::class.java)
+    val left = TaskDef("left", NoOp::class.java).dependsOn(root)
+    val right = TaskDef("right", NoOp::class.java).dependsOn(root)
+    val join = TaskDef("join", NoOp::class.java).dependsOn(left, right)
+    val dag = DagDef("dag")
+    listOf(root, left, right, join).forEach(dag::addTask)
+
+    Assertions.assertEquals(mapOf("dag" to dag), Bundle(listOf(dag)).dags)
+  }
+
+  @Test
+  @DisplayName("Should leave the bundle unchanged when a Dag fails validation")
+  fun shouldNotRegisterInvalidDag() {
+    val bundle = Bundle()
+    val dag = DagDef("dag").addTask(TaskDef("t", NoOp::class.java).dependsOn(TaskDef("missing", NoOp::class.java)))
+
+    Assertions.assertThrows(IllegalArgumentException::class.java) { bundle.register(dag) }
+
+    Assertions.assertEquals(emptySet<String>(), bundle.dags.keys)
+  }
+
  @Test
  @DisplayName("Should find the registrar generated for a nested handler class")
  fun shouldFindRegistrarOfNestedHandlerClass() {
@@ -19,6 +19,7 @@

 package org.apache.airflow.sdk

+import org.apache.airflow.sdk.Deps.Flow
 import org.junit.jupiter.api.Assertions
 import org.junit.jupiter.api.DisplayName
 import org.junit.jupiter.api.Test
@@ -94,4 +95,78 @@ internal class DagDefTest {

    Assertions.assertEquals("Task 'extract' already belongs to Dag 'dag'", error.message)
  }
+
+  @Test
+  @DisplayName("Should register a task and hand back its handle")
+  fun shouldRegisterTaskFromFactory() {
+    val dag = DagDef("dag")
+
+    val extract = dag.task<Long>("extract", NoOp::class.java)
+
+    Assertions.assertEquals(listOf("extract"), dag.tasks.keys.toList())
+    Assertions.assertEquals(listOf(dag.tasks.getValue("extract")), extract.nodes())
+  }
+
+  @Test
+  @DisplayName("Should fan out ordering-only edges with before")
+  fun shouldFanOutOrderingEdgesWithBefore() {
+    val dag = DagDef("dag")
+    val extract = dag.task<Long>("extract", NoOp::class.java)
+    val left = dag.task<Unit>("left", NoOp::class.java)
+    val right = dag.task<Unit>("right", NoOp::class.java)
+
+    extract.before(left, right)
+
+    Assertions.assertEquals(setOf("extract"), upstreamsOf(dag, "left"))
+    Assertions.assertEquals(setOf("extract"), upstreamsOf(dag, "right"))
+  }
+
+  @Test
+  @DisplayName("Should fan in ordering-only edges with after")
+  fun shouldFanInOrderingEdgesWithAfter() {
+    val dag = DagDef("dag")
+    val left = dag.task<Unit>("left", NoOp::class.java)
+    val right = dag.task<Unit>("right", NoOp::class.java)
+    val join = dag.task<Unit>("join", NoOp::class.java)
+
+    join.after(left, right)
+
+    Assertions.assertEquals(setOf("left", "right"), upstreamsOf(dag, "join"))
+  }
+
+  @Test
+  @DisplayName("Should return the receiver so before and after chain on one task")
+  fun shouldChainBeforeAndAfterOnOneTask() {
+    val dag = DagDef("dag")
+    val extract = dag.task<Long>("extract", NoOp::class.java)
+    val transform = dag.task<Long>("transform", NoOp::class.java)
+    val load = dag.task<Unit>("load", NoOp::class.java)
+
+    transform.after(extract).before(load)
+
+    Assertions.assertEquals(setOf("extract"), upstreamsOf(dag, "transform"))
+    Assertions.assertEquals(setOf("transform"), upstreamsOf(dag, "load"))
+  }
+
+  @Test
+  @DisplayName("Should draw every edge between two sets with Flow.of")
+  fun shouldDrawEdgesBetweenSetsWithFlowOf() {
+    val dag = DagDef("dag")
+    val left = dag.task<Unit>("left", NoOp::class.java)
+    val right = dag.task<Unit>("right", NoOp::class.java)
+    val join = dag.task<Unit>("join", NoOp::class.java)
+
+    Flow.of(left, right).before(join)
+
+    Assertions.assertEquals(setOf("left", "right"), upstreamsOf(dag, "join"))
+  }
+
+  private fun upstreamsOf(
+    dag: DagDef,
+    taskId: String,
+  ): Set<String> =
+    dag.tasks
+      .getValue(taskId)
+      .upstreams
+      .mapTo(mutableSetOf()) { it.id }
 }