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.

Author: jason810496

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

@@ -348,7 +348,8 @@ Interface-based API
 ~~~~~~~~~~~~~~~~~~~
 
 Implement the ``Task`` interface directly for full control over how tasks are registered and how XComs are
-read.  Each task is registered as a ``TaskDef`` on a ``DagDef``.
+read.  Each task is registered as a ``TaskDef`` on a ``DagDef``; ``TaskDef`` also carries
+``dependsOn(...)`` for declaring edges between task definitions.
 
 The runner creates a fresh instance of the task class through reflection for every task-instance run,
 which puts four constraints on the class:
@@ -525,6 +526,41 @@ An ``InputTask`` whose type argument is not a concrete ``TaskInput`` fails when
 rather than mid-run.  Plain ``Task`` remains the right interface for a task the Dag file
 calls with no arguments.
 
+.. _java-sdk/native-dags:
+
+Native Java Dags
+----------------
+
+A Dag can also be authored entirely in Java, with no Python stub file: the ``DagDef`` and
+``TaskDef`` objects hold the tasks, and Java declares the graph.
+
+Building the Dag in Java
+~~~~~~~~~~~~~~~~~~~~~~~~
+
+``dag.task(...)`` registers a task as it creates it and hands back a handle, so there is no second
+``addTask`` call to forget.  ``before`` and ``after`` draw every edge on this surface — Python's
+``a >> b`` and ``b << a`` — and the task body moves the data itself, by reading the upstream's XCom
+through ``Client``:
+
+.. code-block:: java
+
+    var dag = new DagDef("java_etl");
+
+    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);
+
+Both are variadic, so ``a.before(b, c)`` fans out and ``d.after(b, c)`` fans in, and both return
+their own receiver, so a chain reads from one task outwards.  ``Deps.Flow.of(a, b).before(c, d)``
+draws every edge between two sets in one call.  ``TaskDef.dependsOn(...)`` declares the same edges
+from the definitions rather than the handles, for code that builds ``TaskDef`` objects up front.
+
+Edges are checked when the Dag is registered with a ``Bundle``: an upstream that belongs to another
+Dag, or to no Dag, and a cycle anywhere in the graph both fail there rather than at the first task
+run.
+
 .. _java-sdk/logging:
 
 Logging

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

@@ -26,7 +26,8 @@ public static Bundle build() {
     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) {

java-sdk/example/src/java/org/apache/airflow/example/nativedag/InterfaceExample.java

@@ -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;
+  }
+}

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

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

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

@@ -28,7 +28,9 @@ import org.apache.airflow.sdk.internal.registrarName
  * [Server.serve] to start accepting task-execution requests.
  *
  * @property dags All registered Dags 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>,
@@ -45,10 +47,24 @@ 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.
+   * @throws IllegalArgumentException if another Dag shares its ID, a task
+   *    depends on an upstream not registered in the same Dag, or the
+   *    dependencies contain a cycle.
    */
   fun register(dag: DagDef): Bundle {
+    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}"
     }
@@ -85,8 +101,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.
@@ -103,6 +119,23 @@ class Bundle(
   }
 }
 
+// TaskDef.dependsOn can express a cycle, so reject one at registration time.
+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.
  *

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

@@ -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. Upstream tasks referenced via [TaskDef.dependsOn]
+   * 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.
@@ -84,13 +113,35 @@ class DagDef(
     task.owner = this
     return this
   }
+
+  /**
+   * Registers a task with this Dag, adding [upstreams] as its dependencies.
+   *
+   * Equivalent to `addTask(task.dependsOn(...))`; see [addTask] and
+   * [TaskDef.dependsOn].
+   *
+   * @param task Task definition to register.
+   * @param upstreams Tasks that [task] depends on (its upstream tasks).
+   * @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,
+    upstreams: List<TaskDef>,
+  ): DagDef {
+    upstreams.forEach { task.dependsOn(it) }
+    return addTask(task)
+  }
 }
 
 /**
- * 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.
  *
  * ```java
  * var extract = new TaskDef("extract", Extract.class);
+ * var load = new TaskDef("load", Load.class).dependsOn(extract);
  * ```
  *
  * @param id Task identifier, unique within a [DagDef].
@@ -110,7 +161,22 @@ class TaskDef(
     validateTaskInput(definition)
   }
 
+  internal val upstreams = linkedSetOf<TaskDef>()
   internal var owner: DagDef? = null
+
+  /**
+   * Declares that this task runs after [upstreams].
+   *
+   * Referenced tasks must be registered with the same [DagDef] as this task by
+   * the time it is added to a [Bundle].
+   *
+   * @param upstreams Tasks this task depends on.
+   * @return This task definition, for chaining.
+   */
+  fun dependsOn(vararg upstreams: TaskDef): TaskDef {
+    this.upstreams += upstreams
+    return this
+  }
 }
 
 /**