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 77a1953cd41 Java SDK: Draw a Java-authored Dag's edges with before and
after (#73594)
77a1953cd41 is described below
commit 77a1953cd4193a6f09e08c25b3681d052d9a3fe5
Author: Jason(Zhe-You) Liu <[email protected]>
AuthorDate: Sat Sep 26 20:39:34 2026 +0800
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 <[email protected]>
---
.../language-sdks/java.rst | 34 ++++++++
.../airflow/example/ExampleBundleBuilder.java | 3 +-
.../example/nativedag/InterfaceExample.java | 68 ++++++++++++++++
.../src/main/kotlin/org/apache/airflow/sdk/Arg.kt | 53 ++++++++++++
.../main/kotlin/org/apache/airflow/sdk/Bundle.kt | 47 +++++++++--
.../main/kotlin/org/apache/airflow/sdk/DagDef.kt | 52 ++++++++++--
.../src/main/kotlin/org/apache/airflow/sdk/Deps.kt | 93 ++++++++++++++++++++++
.../kotlin/org/apache/airflow/sdk/BundleTest.kt | 79 ++++++++++++++++++
.../kotlin/org/apache/airflow/sdk/DagDefTest.kt | 75 +++++++++++++++++
9 files changed, 492 insertions(+), 12 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 c3d385cc5e0..1561ecfffcf 100644
--- a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst
+++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst
@@ -532,6 +532,40 @@ 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. ``Flow.of(a,
b).before(c, d)``, from
+``org.apache.airflow.sdk.Deps.Flow``, draws every edge between two sets in one
call.
+
+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
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 c8e0ac75f4e..d75f5762ba1 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,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) {
diff --git
a/java-sdk/example/src/java/org/apache/airflow/example/nativedag/InterfaceExample.java
b/java-sdk/example/src/java/org/apache/airflow/example/nativedag/InterfaceExample.java
new file mode 100644
index 00000000000..58c33218254
--- /dev/null
+++
b/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;
+ }
+}
diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Arg.kt
b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Arg.kt
new file mode 100644
index 00000000000..1834a548e6a
--- /dev/null
+++ b/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
+ }
+}
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 447148b2152..e1dff618894 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
@@ -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.
*
diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/DagDef.kt
b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/DagDef.kt
index 135d651f5a9..90f77d86dee 100644
--- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/DagDef.kt
+++ b/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. 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
+ }
}
/**
diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Deps.kt
b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Deps.kt
new file mode 100644
index 00000000000..7860a69a0c8
--- /dev/null
+++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Deps.kt
@@ -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
+}
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 207fe8b2442..29c5fc8a99b 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
@@ -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() {
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
index a23f76bad82..5ad4125e985 100644
--- 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
@@ -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 }
}
