You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by ja...@apache.org on 2026/10/09 06:31:12 UTC
(airflow) 03/04: Java SDK: Choose one of several tasks with a branch (#74439)
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
commit c3043c266c677a45531d4eae6a2e26ddd61d3096
Author: Jason(Zhe-You) Liu <68...@users.noreply.github.com>
AuthorDate: Fri Oct 9 14:30:52 2026 +0800
Java SDK: Choose one of several tasks with a branch (#74439)
* Java SDK: Choose one of several tasks with a branch
A condition picks between two tasks, but a Dag that routes work by size or
by region needs more than two, and splitting that into nested conditions
makes the graph unreadable.
`DagDef.Branch` and `@Builder.Branch` declare a task that chooses one of
the tasks listed with `option(...)`; every other one is skipped. The chosen
task's ID becomes the branch's return value, as it does for Python's branch
operators. A branch chooses exactly one case, so a task that runs after
several of them needs a trigger rule such as
`none_failed_min_one_success`.
Each authoring style names the case in the way that style can type-check. A
`BranchTask` returns the case's own class, so javac checks it; two cases
that share a class are rejected where the second is named, since the branch
could not tell them apart. A `@Builder.Branch` method has no such class to
name -- a task method's implementation is generated -- so it returns a
`TaskId` from the generated `<Dag>Builder.TaskIds`, which holds one
constant per task of the Dag. Either way nothing hard-codes a task id
string, and the generated holder is emitted only for a Dag that declares a
branch.
* Java SDK: Name the multi-way branch Switch, as the other SDKs do
* Java SDK: Name a switch case by its task class in both styles
* Java SDK: Say when a switch rejects a choice and when it fails
* Java SDK: Cover a generated switch that names a task it cannot choose
* Java SDK: Check a switch carries the chosen task's id
* Java SDK: Let a wiring view name a switch's cases in more than one call
---
.pre-commit-config.yaml | 2 +-
java-sdk/README.md | 2 +-
java-sdk/capabilities.yaml | 2 +-
.../example/nativedag/AnnotationExample.java | 20 ++
.../example/nativedag/InterfaceExample.java | 32 ++-
.../org/apache/airflow/sdk/BuilderProcessor.kt | 108 +++++++---
.../kotlin/org/apache/airflow/sdk/BuilderTest.kt | 212 ++++++++++++++++++-
.../ci/prek/check_serialization_conformance.py | 2 +-
java-sdk/sdk/build.gradle.kts | 35 +++
java-sdk/sdk/module.md | 2 +-
.../main/kotlin/org/apache/airflow/sdk/Bundle.kt | 7 +-
.../kotlin/org/apache/airflow/sdk/Condition.kt | 2 +-
.../main/kotlin/org/apache/airflow/sdk/DagDef.kt | 43 +++-
.../main/kotlin/org/apache/airflow/sdk/Switch.kt | 234 +++++++++++++++++++++
.../kotlin/org/apache/airflow/sdk/TaskGroupRef.kt | 30 +++
.../kotlin/org/apache/airflow/sdk/ConditionTest.kt | 2 +-
.../kotlin/org/apache/airflow/sdk/SwitchTest.kt | 207 ++++++++++++++++++
.../airflow/sdk/conformance/SerializeJava.kt | 140 +++++++++++-
.../org/apache/airflow/sdk/execution/TaskTest.kt | 98 +++++++++
.../org/apache/airflow/sdk/internal/RefsTest.kt | 36 ++++
20 files changed, 1165 insertions(+), 51 deletions(-)
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 37f87710dfc..836f020378a 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -461,7 +461,7 @@ repos:
^java-sdk/scripts/ci/prek/check_serialization_conformance\.py$|
^java-sdk/sdk/schema/dag-schema\.json$|
^java-sdk/sdk/build\.gradle\.kts$|
- ^java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/(Bundle|Condition|DagDef|Decider|Deps|Arg|TaskGroupRef)\.kt$|
+ ^java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/(Switch|Bundle|Condition|DagDef|Decider|Deps|Arg|TaskGroupRef)\.kt$|
^java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Serde\.kt$|
^java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/(Refs|Fields)\.kt$|
^java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/conformance/.*$|
diff --git a/java-sdk/README.md b/java-sdk/README.md
index 1e4e7478f39..9b9653a4974 100644
--- a/java-sdk/README.md
+++ b/java-sdk/README.md
@@ -643,7 +643,7 @@ prek hook regenerate it.
| capability: `task-args` | MUST † | ✓ | 3.4 | |
| capability: `dag-params` | MUST † | ✗ | – | |
| capability: `taskflow-dependencies` | MUST † | ✓ | 3.4 | |
-| capability: `branching` | SHOULD † | ✓ | 3.4 | DagDef.If / @Builder.If |
+| capability: `branching` | SHOULD † | ✓ | 3.4 | DagDef.If and DagDef.Switch, or @Builder.If and @Builder.Switch |
| capability: `dag-test` | SHOULD † | ✗ | – | |
| capability: `task-group` | MAY † | ✓ | 3.4 | |
| capability: `dynamic-task-mapping` | MAY † | ✗ | – | |
diff --git a/java-sdk/capabilities.yaml b/java-sdk/capabilities.yaml
index 33106f170fb..5a2ed1f5ae4 100644
--- a/java-sdk/capabilities.yaml
+++ b/java-sdk/capabilities.yaml
@@ -109,7 +109,7 @@ capabilities:
branching:
supported: true
since: "3.4"
- note: "DagDef.If / @Builder.If"
+ note: "DagDef.If and DagDef.Switch, or @Builder.If and @Builder.Switch"
dag-test:
supported: false
task-group:
diff --git a/java-sdk/example/src/java/org/apache/airflow/example/nativedag/AnnotationExample.java b/java-sdk/example/src/java/org/apache/airflow/example/nativedag/AnnotationExample.java
index d9538fd391f..16cdc555452 100644
--- a/java-sdk/example/src/java/org/apache/airflow/example/nativedag/AnnotationExample.java
+++ b/java-sdk/example/src/java/org/apache/airflow/example/nativedag/AnnotationExample.java
@@ -66,6 +66,25 @@ public class AnnotationExample {
return transformed > 0;
}
+ @Builder.Task(id = "report_long")
+ public void reportLong() {
+ log.log(INFO, "Long report");
+ }
+
+ @Builder.Task(id = "report_short")
+ public void reportShort() {
+ log.log(INFO, "Short report");
+ }
+
+ // A switch: it names the one case that runs by the class generated for it,
+ // and every other case is skipped.
+ @Builder.Switch(id = "pick_report")
+ public Class<? extends Task> pickReport(long transformed) {
+ return transformed > 100
+ ? AnnotationExampleBuilder.ReportLong.class
+ : AnnotationExampleBuilder.ReportShort.class;
+ }
+
// A task group: everything it declares is prefixed with its id, so this is
// the task "checks.audit".
@Builder.TaskGroup(id = "checks")
@@ -84,6 +103,7 @@ public class AnnotationExample {
var extracted = extract();
var transformed = transform(extracted, lit(1.5));
hasRows(transformed).Then(load(transformed)).Else(loadEmpty());
+ pickReport(transformed).Case(reportLong()).Case(reportShort());
// Ordering-only edge: the checks group runs after extract, with no data
// flowing.
extracted.before(checks());
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
index cce10676dfd..b8099931b7a 100644
--- 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
@@ -71,6 +71,31 @@ public class InterfaceExample {
}
}
+ public static class ReportLong implements Task {
+ @Override
+ public void execute(Context context, Client client) {
+ log.log(INFO, "Long report");
+ }
+ }
+
+ public static class ReportShort implements Task {
+ @Override
+ public void execute(Context context, Client client) {
+ log.log(INFO, "Short report");
+ }
+ }
+
+ // A switch: it names the one case that runs by its class, and every other
+ // case is skipped.
+ public static class PickReport implements SwitchTask {
+ @Override
+ public Class<? extends Task> choose(Context context, Client client) {
+ return ((Number) client.getXCom("transform")).longValue() > 100
+ ? ReportLong.class
+ : ReportShort.class;
+ }
+ }
+
public static DagDef build() {
var dag =
new DagDef("java_native_interface_example")
@@ -87,9 +112,14 @@ public class InterfaceExample {
var load = dag.task("load", Load.class);
var loadEmpty = dag.task("load_empty", LoadEmpty.class);
+ var reportLong = dag.task("report_long", ReportLong.class);
+ var reportShort = dag.task("report_short", ReportShort.class);
+
transform.after(extract);
- // With no task id given, the condition takes one from its class: "hasRows".
+ // With no task id given, a decider takes one from its class: "hasRows" and
+ // "pickReport".
dag.If(HasRows.class).after(transform).Then(load).Else(loadEmpty);
+ dag.Switch(PickReport.class).after(transform).Case(reportLong).Case(reportShort);
return dag;
}
}
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 1848fc007f7..97b810358a0 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
@@ -49,6 +49,7 @@ import javax.annotation.processing.RoundEnvironment
import javax.annotation.processing.SupportedAnnotationTypes
import javax.annotation.processing.SupportedSourceVersion
import javax.lang.model.SourceVersion
+import javax.lang.model.element.AnnotationMirror
import javax.lang.model.element.AnnotationValue
import javax.lang.model.element.Element
import javax.lang.model.element.ElementKind
@@ -56,8 +57,10 @@ import javax.lang.model.element.ExecutableElement
import javax.lang.model.element.Modifier
import javax.lang.model.element.TypeElement
import javax.lang.model.element.VariableElement
+import javax.lang.model.type.DeclaredType
import javax.lang.model.type.TypeKind
import javax.lang.model.type.TypeMirror
+import javax.lang.model.type.WildcardType
import javax.tools.Diagnostic
import java.lang.reflect.Modifier as ReflectModifier
import org.apache.airflow.sdk.internal.builderName as generatedBuilderName
@@ -97,6 +100,7 @@ import org.apache.airflow.sdk.internal.builderName as generatedBuilderName
"org.apache.airflow.sdk.Builder.Dag",
"org.apache.airflow.sdk.Builder.Task",
"org.apache.airflow.sdk.Builder.If",
+ "org.apache.airflow.sdk.Builder.Switch",
"org.apache.airflow.sdk.Builder.TaskGroup",
"org.apache.airflow.sdk.Builder.TaskHandler",
"org.apache.airflow.sdk.Builder.Deps",
@@ -357,24 +361,22 @@ class BuilderProcessor : AbstractProcessor() {
builderName: ClassName,
inGroup: Boolean,
): MethodSpec {
- val condition = decl.kind == TaskKind.CONDITION
val method =
MethodSpec
.methodBuilder(decl.method.simpleName.toString())
.addModifiers(Modifier.PUBLIC, Modifier.DEFAULT)
.returns(
- if (condition) {
- CONDITION_REF_TYPE
- } else {
- ParameterizedTypeName.get(TASK_HANDLE_TYPE, TypeName.get(decl.method.returnType).boxIfPossible())
- },
+ decl.kind.refType
+ ?: ParameterizedTypeName.get(TASK_HANDLE_TYPE, TypeName.get(decl.method.returnType).boxIfPossible()),
)
decl.dataParams.forEach { method.addParameter(inType(it.type), it.name) }
val def = taskDefCode(decl, CodeBlock.of($$"$T.$L", builderName, decl.className))
// The view knows the group it belongs to, so the recorder is told where
// the task goes instead of deriving it from the task's ID.
val group = if (inGroup) CodeBlock.of("groupId()") else CodeBlock.of($$"$S", "")
- val wrap = { call: CodeBlock -> if (condition) CodeBlock.of($$"$T.of($L)", CONDITION_REF_TYPE, call) else call }
+ val wrap = { call: CodeBlock ->
+ decl.kind.refType?.let { CodeBlock.of($$"$T.of($L)", it, call) } ?: call
+ }
if (decl.dataParams.isEmpty()) {
method.addStatement($$"return $L", wrap(CodeBlock.of($$"$T.node($L, $L)", REFS_TYPE, group, def)))
} else {
@@ -435,17 +437,19 @@ class BuilderProcessor : AbstractProcessor() {
val tasks = mutableListOf<TaskDeclaration>()
for (inner in el.enclosedElements) {
if (inner !is ExecutableElement) continue
- val task = inner.getAnnotation(Builder.Task::class.java)
- val condition = inner.getAnnotation(Builder.If::class.java)
- if (task == null && condition == null) continue
- require(task == null || condition == null) {
- "Method '${inner.simpleName}' carries both @Builder.Task and @Builder.If; a condition is a " +
- "task, so it carries @Builder.If alone"
- }
- val kind = if (condition != null) TaskKind.CONDITION else TaskKind.TASK
+ val declared =
+ TaskKind.entries.filter { kind -> inner.annotationMirrors.any { it.names(kind.annotation) } }
+ if (declared.isEmpty()) continue
+ val kind =
+ declared.singleOrNull()
+ ?: throw IllegalArgumentException(
+ "Method '${inner.simpleName}' carries ${declared.joinToString { it.spelling }}; a task is " +
+ "declared by one of them alone",
+ )
+ val annotated = declaredId(inner, kind)
if (inner.isVarArgs) throw IllegalArgumentException("Cannot create task from vararg function ${inner.simpleName}")
checkDeciderReturn(kind, inner)
- val localId = (condition?.id ?: task!!.id).ifBlank { inner.simpleName.toString() }
+ val localId = annotated.ifBlank { inner.simpleName.toString() }
require(tasks.none { it.method.simpleName.contentEquals(inner.simpleName) }) {
"Class ${el.simpleName} overloads task method '${inner.simpleName}'; a method's name is the " +
"name of its generated task class and of its wiring-view method, so rename one and keep its " +
@@ -712,24 +716,55 @@ class BuilderProcessor : AbstractProcessor() {
kind: TaskKind,
method: ExecutableElement,
) {
- if (kind != TaskKind.CONDITION) return
val returns = method.returnType
- val boolean = returns.kind == TypeKind.BOOLEAN || with(processingEnv) { isType(returns, BOXED_BOOLEAN_TYPE) }
- require(boolean) {
- "@Builder.If method '${method.simpleName}' returns $returns, but a condition returns boolean: " +
- "true runs the task named by Then, false the one named by Else"
+ when (kind) {
+ TaskKind.TASK -> return
+ TaskKind.CONDITION ->
+ require(returns.kind == TypeKind.BOOLEAN || with(processingEnv) { isType(returns, BOXED_BOOLEAN_TYPE) }) {
+ "@Builder.If method '${method.simpleName}' returns $returns, but a condition returns boolean: " +
+ "true runs the task named by Then, false the one named by Else"
+ }
+ TaskKind.SWITCH ->
+ require(returnsTaskClass(returns)) {
+ "@Builder.Switch method '${method.simpleName}' returns $returns, but a switch returns the " +
+ "class of the task it chose, as Class<? extends Task>"
+ }
}
}
+ /** Whether [type] is `Class<T>` or `Class<? extends T>` for a [Task] type `T`. */
+ private fun returnsTaskClass(type: TypeMirror): Boolean {
+ val declared = type as? DeclaredType ?: return false
+ if (!(declared.asElement() as TypeElement).qualifiedName.contentEquals("java.lang.Class")) return false
+ val argument = declared.typeArguments.singleOrNull() ?: return false
+ val bound = if (argument is WildcardType) argument.extendsBound else argument
+ return bound != null &&
+ with(processingEnv) { typeUtils.isAssignable(bound, elementUtils.getTypeElement(TASK_TYPE.canonicalName()).asType()) }
+ }
+
+ /** The `id` the declaring annotation sets, empty when it leaves it out. */
+ private fun declaredId(
+ method: ExecutableElement,
+ kind: TaskKind,
+ ): String {
+ val mirror = method.annotationMirrors.first { it.names(kind.annotation) }
+ val id = mirror.elementValues.entries.firstOrNull { it.key.simpleName.contentEquals("id") }
+ return id?.value?.value as String? ?: ""
+ }
+
private fun buildTask(decl: TaskDeclaration): TypeSpec {
- val condition = decl.kind == TaskKind.CONDITION
val executeSpec =
MethodSpec
- .methodBuilder(if (condition) "decide" else "execute")
+ .methodBuilder(decl.kind.bodyMethod)
.addAnnotation(Override::class.java)
.addModifiers(Modifier.PUBLIC)
- .returns(if (condition) TypeName.BOOLEAN else TypeName.VOID)
- .addParameter(CONTEXT_TYPE, "context")
+ .returns(
+ when (decl.kind) {
+ TaskKind.TASK -> TypeName.VOID
+ TaskKind.CONDITION -> TypeName.BOOLEAN
+ TaskKind.SWITCH -> TASK_CLASS_TYPE
+ },
+ ).addParameter(CONTEXT_TYPE, "context")
.addParameter(CLIENT_TYPE, "client")
.addException(Exception::class.java)
@@ -775,8 +810,8 @@ class BuilderProcessor : AbstractProcessor() {
}
when {
- // The SDK pushes a condition's result itself, after it has skipped the side not taken.
- condition -> $$"return new $T().$L($L)"
+ // The SDK pushes a decider's choice itself, after it has skipped what the choice rules out.
+ decl.kind != TaskKind.TASK -> $$"return new $T().$L($L)"
inner.returnType.kind == TypeKind.VOID -> $$"new $T().$L($L)"
else -> $$"client.setXCom(new $T().$L($L))"
}.also {
@@ -790,7 +825,7 @@ class BuilderProcessor : AbstractProcessor() {
return TypeSpec
.classBuilder(decl.className)
- .addSuperinterface(if (condition) CONDITION_TASK_TYPE else ClassName.get(Task::class.java))
+ .addSuperinterface(decl.kind.taskInterface)
.addModifiers(Modifier.PUBLIC, Modifier.FINAL, Modifier.STATIC)
.addMethod(executeSpec.build())
.build()
@@ -966,8 +1001,13 @@ private val DAG_SOURCE_TYPE = ClassName.get(DagSource::class.java)
private val REFS_TYPE = ClassName.get(Refs::class.java)
private val ARG_TYPE = ClassName.get(Arg::class.java)
private val TASK_HANDLE_TYPE = ClassName.get(TaskRef::class.java)
+private val TASK_TYPE = ClassName.get(Task::class.java)
private val CONDITION_TASK_TYPE = ClassName.get(ConditionTask::class.java)
private val CONDITION_REF_TYPE = ClassName.get(ConditionRef::class.java)
+private val SWITCH_TASK_TYPE = ClassName.get(SwitchTask::class.java)
+private val SWITCH_REF_TYPE = ClassName.get(SwitchRef::class.java)
+private val TASK_CLASS_TYPE =
+ ParameterizedTypeName.get(ClassName.get(Class::class.java), WildcardTypeName.subtypeOf(TASK_TYPE))
private val BOXED_BOOLEAN_TYPE = ClassName.get("java.lang", "Boolean")
private val DEPS_TYPE = ClassName.get(Deps::class.java)
private val GROUP_TYPE = DEPS_TYPE.nestedClass("TaskGroup")
@@ -983,11 +1023,19 @@ private const val DAG_ANNOTATION = "org.apache.airflow.sdk.Builder.Dag"
private enum class TaskKind(
val annotation: String,
val spelling: String,
+ /** Method the generated class implements, and the type its wiring view hands back. */
+ val bodyMethod: String,
+ val taskInterface: ClassName,
+ val refType: ClassName?,
) {
- TASK("org.apache.airflow.sdk.Builder.Task", "@Builder.Task"),
- CONDITION("org.apache.airflow.sdk.Builder.If", "@Builder.If"),
+ TASK("org.apache.airflow.sdk.Builder.Task", "@Builder.Task", "execute", TASK_TYPE, null),
+ CONDITION("org.apache.airflow.sdk.Builder.If", "@Builder.If", "decide", CONDITION_TASK_TYPE, CONDITION_REF_TYPE),
+ SWITCH("org.apache.airflow.sdk.Builder.Switch", "@Builder.Switch", "choose", SWITCH_TASK_TYPE, SWITCH_REF_TYPE),
}
+/** Whether this annotation is the one [name] qualifies. */
+private fun AnnotationMirror.names(name: String): Boolean = (annotationType.asElement() as TypeElement).qualifiedName.contentEquals(name)
+
private val RESERVED_VIEW_NAMES =
setOf(
"depends",
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 70802c405aa..2473ef9561a 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
@@ -2300,6 +2300,216 @@ class BuilderTest {
assertThat(compilation).succeeded()
}
+ @Test
+ @DisplayName("let a wiring view name a switch's cases in more than one statement")
+ fun switchNamedAcrossStatements() {
+ val compilation =
+ compile(
+ """
+ package org.apache.airflow.example;
+
+ import org.apache.airflow.sdk.Builder;
+ import org.apache.airflow.sdk.Task;
+
+ @Builder.Dag(id = "etl")
+ public class TestExample {
+ @Builder.Switch(id = "pick_path")
+ public Class<? extends Task> pickPath() {
+ return TestExampleBuilder.HandleLong.class;
+ }
+
+ @Builder.Task
+ public void handleLong() {}
+
+ @Builder.Task
+ public void handleShort() {}
+
+ @Builder.Deps
+ static class Wiring implements TestExampleDeps {
+ void depends() {
+ pickPath().Case(handleLong());
+ pickPath().Case(handleShort());
+ }
+ }
+ }
+ """,
+ )
+
+ assertThat(compilation).succeeded()
+ }
+
+ @Test
+ @DisplayName("generate a switch that names its case by the class generated for it")
+ fun generateSwitchTask() {
+ val compilation =
+ compile(
+ """
+ package org.apache.airflow.example;
+
+ import org.apache.airflow.sdk.Builder;
+ import org.apache.airflow.sdk.Task;
+
+ @Builder.Dag(id = "etl")
+ public class TestExample {
+ @Builder.Switch(id = "pick_path")
+ public Class<? extends Task> pickPath() {
+ return TestExampleBuilder.HandleLong.class;
+ }
+
+ @Builder.Task
+ public void handleLong() {}
+
+ @Builder.Task
+ public void handleShort() {}
+
+ @Builder.Deps
+ static class Wiring implements TestExampleDeps {
+ void depends() {
+ pickPath().Case(handleLong()).Case(handleShort());
+ }
+ }
+ }
+ """,
+ )
+
+ assertThat(compilation).succeeded()
+ assertThat(compilation)
+ .generatedSourceFile("org.apache.airflow.example.TestExampleBuilder")
+ .hasSourceEquivalentTo(
+ "org.apache.airflow.example.TestExampleBuilder",
+ """
+ package org.apache.airflow.example;
+
+ import java.lang.Class;
+ import java.lang.Exception;
+ import java.lang.Override;
+ import java.util.List;
+ import org.apache.airflow.sdk.Client;
+ import org.apache.airflow.sdk.Context;
+ import org.apache.airflow.sdk.DagDef;
+ import org.apache.airflow.sdk.SwitchTask;
+ import org.apache.airflow.sdk.Task;
+ import org.apache.airflow.sdk.internal.DagSource;
+ import org.apache.airflow.sdk.internal.Refs;
+
+ public final class TestExampleBuilder {
+ public static DagDef build() {
+ var dag = DagSource.declaredBy(new DagDef("etl"), TestExample.class);
+ return Refs.record(dag, List.of("pick_path", "handleLong", "handleShort"), List.of(), new TestExample.Wiring()::depends);
+ }
+
+ public static final class PickPath implements SwitchTask {
+ @Override
+ public Class<? extends Task> choose(Context context, Client client) throws Exception {
+ return new TestExample().pickPath();
+ }
+ }
+
+ public static final class HandleLong implements Task {
+ @Override
+ public void execute(Context context, Client client) throws Exception {
+ new TestExample().handleLong();
+ }
+ }
+
+ public static final class HandleShort implements Task {
+ @Override
+ public void execute(Context context, Client client) throws Exception {
+ new TestExample().handleShort();
+ }
+ }
+ }
+ """,
+ )
+ assertThat(compilation)
+ .generatedSourceFile("org.apache.airflow.example.TestExampleDeps")
+ .hasSourceEquivalentTo(
+ "org.apache.airflow.example.TestExampleDeps",
+ """
+ package org.apache.airflow.example;
+
+ import java.lang.Void;
+ import org.apache.airflow.sdk.Deps;
+ import org.apache.airflow.sdk.SwitchRef;
+ import org.apache.airflow.sdk.TaskDef;
+ import org.apache.airflow.sdk.TaskRef;
+ import org.apache.airflow.sdk.internal.Refs;
+
+ public interface TestExampleDeps extends Deps {
+ default SwitchRef pickPath() {
+ return SwitchRef.of(Refs.node("", new TaskDef("pick_path", TestExampleBuilder.PickPath.class)));
+ }
+
+ default TaskRef<Void> handleLong() {
+ return Refs.node("", new TaskDef("handleLong", TestExampleBuilder.HandleLong.class));
+ }
+
+ default TaskRef<Void> handleShort() {
+ return Refs.node("", new TaskDef("handleShort", TestExampleBuilder.HandleShort.class));
+ }
+ }
+ """,
+ )
+ }
+
+ @Test
+ @DisplayName("reject a switch that does not return the class of a task")
+ fun rejectNonTaskClassSwitch() {
+ val compilation =
+ compile(
+ """
+ package org.apache.airflow.example;
+ import org.apache.airflow.sdk.Builder;
+ @Builder.Dag
+ public class TestExample {
+ @Builder.Switch
+ public String pickPath() {
+ return "handleLong";
+ }
+
+ @Builder.Deps
+ static class Wiring implements TestExampleDeps {
+ void depends() {}
+ }
+ }
+ """,
+ )
+
+ assertThat(compilation).failed()
+ assertThat(compilation).hadErrorContaining(
+ "@Builder.Switch method 'pickPath' returns java.lang.String, but a switch returns the class of the task it chose",
+ )
+ }
+
+ @Test
+ @DisplayName("reject a switch that returns a class that is not a task")
+ fun rejectSwitchReturningNonTaskClass() {
+ val compilation =
+ compile(
+ """
+ package org.apache.airflow.example;
+ import org.apache.airflow.sdk.Builder;
+ @Builder.Dag
+ public class TestExample {
+ @Builder.Switch
+ public Class<String> pickPath() {
+ return String.class;
+ }
+
+ @Builder.Deps
+ static class Wiring implements TestExampleDeps {
+ void depends() {}
+ }
+ }
+ """,
+ )
+
+ assertThat(compilation).failed()
+ assertThat(compilation).hadErrorContaining(
+ "@Builder.Switch method 'pickPath' returns java.lang.Class<java.lang.String>, but a switch returns the class of the task it chose",
+ )
+ }
+
@Test
@DisplayName("reject a condition that does not return a boolean")
fun rejectNonBooleanCondition() {
@@ -2355,7 +2565,7 @@ class BuilderTest {
assertThat(compilation).failed()
assertThat(compilation).hadErrorContaining(
- "Method 'hasRows' carries both @Builder.Task and @Builder.If",
+ "Method 'hasRows' carries @Builder.Task, @Builder.If; a task is declared by one of them alone",
)
}
diff --git a/java-sdk/scripts/ci/prek/check_serialization_conformance.py b/java-sdk/scripts/ci/prek/check_serialization_conformance.py
index fe569c1ff7c..01a16e34219 100755
--- a/java-sdk/scripts/ci/prek/check_serialization_conformance.py
+++ b/java-sdk/scripts/ci/prek/check_serialization_conformance.py
@@ -56,7 +56,7 @@ if __name__ == "__main__":
"--sdk",
"java",
"--supports",
- "literal_inputs,branch",
+ "literal_inputs,branch,switch",
"--",
*serializer,
]
diff --git a/java-sdk/sdk/build.gradle.kts b/java-sdk/sdk/build.gradle.kts
index 723f60085c2..e30facde61e 100644
--- a/java-sdk/sdk/build.gradle.kts
+++ b/java-sdk/sdk/build.gradle.kts
@@ -614,6 +614,41 @@ abstract class GenerateDagDslTask : DefaultTask() {
| )
|
| /**
+ | * Marks a task method that chooses one of several tasks to run;
+ | * every other one is skipped.
+ | *
+ | * The method returns the class the Dag's builder generated for the task it
+ | * chose, as `Class<? extends Task>`. Its wiring-view method returns a
+ | * [SwitchRef], so the cases are listed where the Dag is wired:
+ | *
+ | * ```java
+ | * @Builder.Switch(id = "pick_path")
+ | * public Class<? extends Task> pickPath(long rows) {
+ | * return rows > 1000 ? EtlBuilder.HandleLong.class : EtlBuilder.HandleShort.class;
+ | * }
+ | *
+ | * @Builder.Deps
+ | * static class Wiring implements EtlDeps {
+ | * void depends() { pickPath(extract()).Case(handleLong()).Case(handleShort()); }
+ | * }
+ | * ```
+ | *
+ | * A switch chooses exactly one case, so a task that runs after
+ | * several of them needs a trigger rule that tolerates a skipped
+ | * upstream, such as `none_failed_min_one_success`.
+ | *
+ | * Configuration attributes are [Task]'s, and apply to the deciding
+ | * task.
+ | */
+ | @Target(AnnotationTarget.FUNCTION)
+ | @MustBeDocumented
+ | annotation class Switch(
+ | /** Task ID. Empty derives it from the annotated function's name. */
+ | val id: String = "",
+ |${attrLines(taskFields)}
+ | )
+ |
+ | /**
| * Marks a method as the Java body of a task the Python Dag file
| * declares with `@task.stub`.
| *
diff --git a/java-sdk/sdk/module.md b/java-sdk/sdk/module.md
index 523dbaf54a4..f03b936d7b9 100644
--- a/java-sdk/sdk/module.md
+++ b/java-sdk/sdk/module.md
@@ -60,7 +60,7 @@ meaning of each dimension is defined in the
| capability: `task-args` | MUST † | ✓ | 3.4 | |
| capability: `dag-params` | MUST † | ✗ | – | |
| capability: `taskflow-dependencies` | MUST † | ✓ | 3.4 | |
-| capability: `branching` | SHOULD † | ✓ | 3.4 | DagDef.If / @Builder.If |
+| capability: `branching` | SHOULD † | ✓ | 3.4 | DagDef.If and DagDef.Switch, or @Builder.If and @Builder.Switch |
| capability: `dag-test` | SHOULD † | ✗ | – | |
| capability: `task-group` | MAY † | ✓ | 3.4 | |
| capability: `dynamic-task-mapping` | MAY † | ✗ | – | |
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 7937446c3f6..5e817c05e93 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
@@ -67,7 +67,8 @@ class Bundle(
* @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, a task that decides which task
- * runs is not declared with `If`, or the dependencies contain a cycle.
+ * runs is not declared with `If` or `Switch`, or the dependencies contain
+ * a cycle.
* @throws IllegalStateException if [Server.serve] has already been called.
*/
fun register(dag: DagDef): Bundle {
@@ -87,7 +88,7 @@ class Bundle(
require(decides == (def.decider != null)) {
if (decides) {
"Task '${def.id}' runs '${def.definition.name}', which decides which task runs, but nothing names " +
- "what it chooses; declare it with If(...), in Dag '${dag.id}'"
+ "what it chooses; declare it with If(...) or Switch(...), in Dag '${dag.id}'"
} else {
"Task '${def.id}' chooses which task runs, but '${def.definition.name}' decides nothing, in Dag '${dag.id}'"
}
@@ -212,7 +213,7 @@ class Bundle(
}
/** Task types the SDK runs through a decider, so each needs the sides it chooses between. */
-private val DECIDER_TYPES = listOf(ConditionTask::class.java)
+private val DECIDER_TYPES = listOf(ConditionTask::class.java, SwitchTask::class.java)
// 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
diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Condition.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Condition.kt
index 90150609c7b..ea521f92028 100644
--- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Condition.kt
+++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Condition.kt
@@ -103,7 +103,7 @@ class ConditionRef private constructor(
@JvmStatic
fun of(ref: TaskRef<Boolean>): ConditionRef {
ref.def.decider?.let { existing ->
- require(existing is ConditionDef) { "Task '${ref.def.id}' already decides what to skip as a branch" }
+ require(existing is ConditionDef) { "Task '${ref.def.id}' already decides what to skip as a switch" }
return ConditionRef(ref, existing)
}
return ConditionRef(ref, ConditionDef().also { ref.def.decider = it })
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 40ead4ae608..9f79be81a6c 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
@@ -201,6 +201,47 @@ class DagDef(
definition: Class<out ConditionTask>,
): ConditionRef = ConditionRef.of(task(id, definition))
+ /**
+ * Declares a task that chooses one of several tasks to run; every other one
+ * is skipped.
+ *
+ * The task's ID is the class's simple name with its first character
+ * lowercased, so `PickPath.class` becomes `pickPath`. Use
+ * [Switch(id, definition)][Switch] to set it, and [SwitchRef.config] for the
+ * task's other settings:
+ *
+ * ```java
+ * dag.Switch(PickPath.class).Case(handleLong).Case(handleShort);
+ * ```
+ *
+ * @param definition Class that implements [SwitchTask]. Must have a public
+ * no-arg constructor.
+ * @return The switch, to list its cases on.
+ * @throws IllegalArgumentException if a task with the same ID is already
+ * registered.
+ */
+ @Suppress("ktlint:standard:function-naming")
+ fun Switch(definition: Class<out SwitchTask>): SwitchRef = Switch(deriveTaskId(definition), definition)
+
+ /**
+ * Declares a task that chooses one of several tasks to run, under the task
+ * ID [id].
+ *
+ * @param id Task identifier, unique within this Dag.
+ * @param definition Class that implements [SwitchTask]. Must have a public
+ * no-arg constructor.
+ * @return The switch, to list its cases on.
+ * @throws IllegalArgumentException if a task with the same ID is already
+ * registered.
+ *
+ * @see Switch
+ */
+ @Suppress("ktlint:standard:function-naming")
+ fun Switch(
+ id: String,
+ definition: Class<out SwitchTask>,
+ ): SwitchRef = SwitchRef.of(task<Any?>(id, definition))
+
/**
* Declares a task group of this Dag.
*
@@ -375,7 +416,7 @@ class TaskDef(
internal val upstreams = linkedSetOf<TaskDef>()
internal var owner: DagDef? = null
- /** What this task decides to run, for a condition or a branch; null otherwise. */
+ /** What this task decides to run, for a condition or a switch; null otherwise. */
internal var decider: DeciderDef? = null
/**
diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Switch.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Switch.kt
new file mode 100644
index 00000000000..a668db89609
--- /dev/null
+++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Switch.kt
@@ -0,0 +1,234 @@
+/*
+ * 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 kotlin.Throws
+
+/**
+ * A task that chooses one of several tasks to run; every other one is skipped.
+ *
+ * Register one with [DagDef.Switch], then list what it can choose with
+ * [SwitchRef.Case]. The choice is the class of the task to run:
+ *
+ * ```java
+ * public class PickPath implements SwitchTask {
+ * @Override
+ * public Class<? extends Task> choose(Context context, Client client) {
+ * return ((Number) client.getXCom("extract")).longValue() > 1000
+ * ? HandleLong.class
+ * : HandleShort.class;
+ * }
+ * }
+ * ```
+ *
+ * The SDK runs [choose] and pushes the chosen task's ID as this task's return
+ * value, so [execute] is never called.
+ *
+ * Only a [Task] class compiles as a choice. javac cannot tell whether it is a
+ * case of this switch, so a Task class that is not one compiles and fails when
+ * the task runs: the task instance is marked failed, and nothing is pushed or
+ * skipped.
+ *
+ * Because a switch names a case by its class, two tasks that run the same
+ * class cannot both be cases of it; the second [SwitchRef.Case] call throws.
+ *
+ * @see DagDef.Switch
+ */
+interface SwitchTask : Task {
+ /**
+ * Chooses the one case that runs.
+ *
+ * Any exception thrown marks the task instance as failed, and nothing is
+ * then pushed or skipped.
+ *
+ * @param context Runtime context for the current execution workload.
+ * @param client Client for Airflow API calls scoped to this execution.
+ * @return The class of one of this switch's cases.
+ * @throws Exception on failure; the task instance is marked failed.
+ */
+ @Throws(Exception::class)
+ fun choose(
+ context: Context,
+ client: Client,
+ ): Class<out Task>
+
+ /** Never called: the SDK runs a switch through [choose]. */
+ override fun execute(
+ context: Context,
+ client: Client,
+ ): Unit =
+ throw IllegalStateException(
+ "Switch '${javaClass.name}' runs through choose(), so execute() is never called",
+ )
+}
+
+/**
+ * A switch registered with a Dag: list the tasks it can choose between.
+ *
+ * ```java
+ * dag.Switch(PickPath.class).Case(handleLong).Case(handleShort);
+ * ```
+ *
+ * A case runs after the switch, so naming it records that edge, as
+ * [Deps.Flow.before] would. The switch chooses exactly one case and skips
+ * every other, so a task that runs after several cases needs a trigger rule
+ * that tolerates a skipped upstream, such as `none_failed_min_one_success`.
+ *
+ * @see DagDef.Switch
+ */
+class SwitchRef private constructor(
+ private val ref: TaskRef<*>,
+ private val cases: SwitchDef,
+) : Deps.Flow {
+ companion object {
+ /**
+ * @suppress
+ *
+ * Marks a registered task as a switch. Public so a generated wiring view
+ * can call it; user code reaches a switch through [DagDef.Switch].
+ *
+ * A wiring view gets the same task back each time it names it, so a second
+ * call reuses the switch rather than declaring it again. [DagDef.Switch] and
+ * [TaskGroupRef.Switch] never reach that path, since registering a task
+ * rejects a duplicate ID first.
+ */
+ @JvmStatic
+ fun of(ref: TaskRef<*>): SwitchRef {
+ ref.def.decider?.let { existing ->
+ require(existing is SwitchDef) { "Task '${ref.def.id}' already decides what to skip as a condition" }
+ return SwitchRef(ref, existing)
+ }
+ val definition = ref.def.definition
+ require(SwitchTask::class.java.isAssignableFrom(definition)) {
+ "Task '${ref.def.id}' runs '${definition.name}', which chooses nothing; a switch runs a " +
+ "SwitchTask"
+ }
+ return SwitchRef(ref, SwitchDef().also { ref.def.decider = it })
+ }
+ }
+
+ /** Airflow task ID of the deciding task. */
+ val id: String get() = ref.def.id
+
+ /**
+ * Adds a task this switch can choose. Every other case is skipped when the
+ * switch chooses this one.
+ *
+ * @param task Task of the same Dag.
+ * @return This switch, for chaining.
+ * @throws IllegalArgumentException if the task belongs to another Dag, is
+ * already a case of this switch, or runs the same class as one.
+ */
+ @Suppress("ktlint:standard:function-naming")
+ fun Case(task: TaskRef<*>): SwitchRef {
+ cases.add(ref.def, task.def)
+ // The case runs after the switch, which is what puts the switch in the
+ // serialized Dag as its upstream and lets registration see a cycle that
+ // runs through a switch.
+ ref.before(task)
+ return this
+ }
+
+ /**
+ * Sets one task-level configuration value on the deciding task.
+ *
+ * @param key Airflow task setting name.
+ * @param value Value matching the key's schema type.
+ * @return This switch, for chaining.
+ * @throws IllegalArgumentException if the key is unknown or the value type
+ * does not match.
+ */
+ fun config(
+ key: String,
+ value: Any?,
+ ): SwitchRef {
+ ref.config(key, value)
+ return this
+ }
+
+ override fun nodes(): List<TaskDef> = ref.nodes()
+
+ override fun before(vararg next: Deps.Flow): SwitchRef {
+ ref.before(*next)
+ return this
+ }
+
+ override fun after(vararg previous: Deps.Flow): SwitchRef {
+ ref.after(*previous)
+ return this
+ }
+}
+
+/** The cases of one switch, as [SwitchRef.Case] listed them. */
+internal class SwitchDef : DeciderDef {
+ private val options = mutableListOf<TaskDef>()
+
+ override val cases: List<TaskDef> get() = options
+
+ fun add(
+ decider: TaskDef,
+ case: TaskDef,
+ ) {
+ requireUnregistered(decider, "Switch")
+ requireSameDag(decider, case, "Switch")
+ require(options.none { it === case }) {
+ "Switch '${decider.id}' already chooses between '${case.id}' and others; name each case once"
+ }
+ // A switch names its case by class, so two cases sharing one would be
+ // indistinguishable.
+ options.firstOrNull { it.definition == case.definition }?.let { first ->
+ throw IllegalArgumentException(
+ "Switch '${decider.id}' cannot choose between '${first.id}' and '${case.id}': both run " +
+ "'${case.definition.name}', and a switch names its case by class",
+ )
+ }
+ options += case
+ }
+
+ override fun describe(decider: TaskDef): String? =
+ if (options.isEmpty()) "Switch '${decider.id}' has no task to choose between; call Case(...)" else null
+
+ override fun decide(
+ decider: TaskDef,
+ instance: Task,
+ context: Context,
+ client: Client,
+ ): Decision {
+ val chosen =
+ when (instance) {
+ is SwitchTask -> byClass(decider, instance.choose(context, client))
+ else -> throw IllegalStateException(
+ "Task '${decider.id}' is a switch, but '${instance.javaClass.name}' is not a SwitchTask",
+ )
+ }
+ return Decision(chosen.id, options.filterNot { it === chosen }.map { it.id })
+ }
+
+ private fun byClass(
+ decider: TaskDef,
+ definition: Class<out Task>?,
+ ): TaskDef =
+ options.firstOrNull { it.definition == definition }
+ ?: throw IllegalArgumentException(
+ "Switch '${decider.id}' chose ${definition?.name ?: "nothing"}, which is not one of its cases: ${describeCases()}",
+ )
+
+ private fun describeCases(): String = options.joinToString { "'${it.id}'" }
+}
diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/TaskGroupRef.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/TaskGroupRef.kt
index e2d7d9613ef..60d43f84fcd 100644
--- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/TaskGroupRef.kt
+++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/TaskGroupRef.kt
@@ -106,6 +106,36 @@ class TaskGroupRef internal constructor(
definition: Class<out ConditionTask>,
): ConditionRef = ConditionRef.of(task(id, definition))
+ /**
+ * Declares a switch in this group, as [DagDef.Switch] does for the Dag.
+ *
+ * @param definition Class that implements [SwitchTask]. Must have a public
+ * no-arg constructor.
+ * @return The switch, to list its cases on.
+ * @throws IllegalArgumentException if the Dag already has a task or task
+ * group with the resulting ID.
+ */
+ @Suppress("ktlint:standard:function-naming")
+ fun Switch(definition: Class<out SwitchTask>): SwitchRef = Switch(deriveTaskId(definition), definition)
+
+ /**
+ * Declares a switch in this group under the task ID `<group ID>.<id>`.
+ *
+ * @param id Task ID within this group.
+ * @param definition Class that implements [SwitchTask]. Must have a public
+ * no-arg constructor.
+ * @return The switch, to list its cases on.
+ * @throws IllegalArgumentException if the Dag already has a task or task
+ * group with the resulting ID.
+ *
+ * @see Switch
+ */
+ @Suppress("ktlint:standard:function-naming")
+ fun Switch(
+ id: String,
+ definition: Class<out SwitchTask>,
+ ): SwitchRef = SwitchRef.of(task<Any?>(id, definition))
+
/**
* Nests a task group inside this one.
*
diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ConditionTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ConditionTest.kt
index 1cebf7ff8e1..0fb13877837 100644
--- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ConditionTest.kt
+++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ConditionTest.kt
@@ -141,7 +141,7 @@ internal class ConditionTest {
assertEquals(
"Task 'has_rows' runs '${HasRows::class.java.name}', which decides which task runs, but nothing names " +
- "what it chooses; declare it with If(...), in Dag 'd'",
+ "what it chooses; declare it with If(...) or Switch(...), in Dag 'd'",
error.message,
)
}
diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/SwitchTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/SwitchTest.kt
new file mode 100644
index 00000000000..ef72edf3a4c
--- /dev/null
+++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/SwitchTest.kt
@@ -0,0 +1,207 @@
+/*
+ * 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.assertEquals
+import org.junit.jupiter.api.Assertions.assertThrows
+import org.junit.jupiter.api.DisplayName
+import org.junit.jupiter.api.Test
+
+/** A case of the switches under test; a second class so a switch can tell two cases apart. */
+internal class HandleLong : Task {
+ override fun execute(
+ context: Context,
+ client: Client,
+ ) = Unit
+}
+
+/** A switch that always chooses [HandleLong]. */
+class PickPath : SwitchTask {
+ override fun choose(
+ context: Context,
+ client: Client,
+ ): Class<out Task> = HandleLong::class.java
+}
+
+internal class SwitchTest {
+ private fun dagWithCases(): Triple<DagDef, SwitchRef, Pair<TaskRef<Unit>, TaskRef<Unit>>> {
+ val dag = DagDef("d")
+ val handleLong = dag.task<Unit>("handle_long", HandleLong::class.java)
+ val handleShort = dag.task<Unit>("handle_short", NoopTask::class.java)
+ return Triple(dag, dag.Switch(PickPath::class.java), handleLong to handleShort)
+ }
+
+ @Test
+ @DisplayName("Should take the switch's task ID from its class when none is given")
+ fun shouldDeriveTaskIdFromClass() {
+ val (dag, switch, cases) = dagWithCases()
+ switch.Case(cases.first)
+
+ assertEquals("pickPath", switch.id)
+ assertEquals(setOf(dag.tasks.getValue("pickPath")), cases.first.def.upstreams)
+ }
+
+ @Test
+ @DisplayName("Should reject naming the same case twice")
+ fun shouldRejectDuplicateCase() {
+ val (_, switch, cases) = dagWithCases()
+ switch.Case(cases.first)
+
+ val error = assertThrows(IllegalArgumentException::class.java) { switch.Case(cases.first) }
+
+ assertEquals(
+ "Switch 'pickPath' already chooses between 'handle_long' and others; name each case once",
+ error.message,
+ )
+ }
+
+ @Test
+ @DisplayName("Should reject two cases a class-named switch could not tell apart")
+ fun shouldRejectCasesSharingAClass() {
+ val dag = DagDef("d")
+ val first = dag.task<Unit>("first", HandleLong::class.java)
+ val second = dag.task<Unit>("second", HandleLong::class.java)
+ val switch = dag.Switch(PickPath::class.java).Case(first)
+
+ val error = assertThrows(IllegalArgumentException::class.java) { switch.Case(second) }
+
+ assertEquals(
+ "Switch 'pickPath' cannot choose between 'first' and 'second': both run " +
+ "'org.apache.airflow.sdk.HandleLong', and a switch names its case by class",
+ error.message,
+ )
+ }
+
+ @Test
+ @DisplayName("Should reject a case declared in another Dag")
+ fun shouldRejectCaseFromAnotherDag() {
+ val (_, switch, _) = dagWithCases()
+ val other = DagDef("other").task<Unit>("handle_long", HandleLong::class.java)
+
+ val error = assertThrows(IllegalArgumentException::class.java) { switch.Case(other) }
+
+ assertEquals(
+ "Switch 'pickPath' of Dag 'd' cannot run task 'handle_long' of Dag 'other'; " +
+ "name a task of the same Dag",
+ error.message,
+ )
+ }
+
+ @Test
+ @DisplayName("Should reject registering a switch with nothing to choose between")
+ fun shouldRejectSwitchWithoutCases() {
+ val (dag, _, _) = dagWithCases()
+
+ val error = assertThrows(IllegalArgumentException::class.java) { Bundle().register(dag) }
+
+ assertEquals(
+ "Switch 'pickPath' has no task to choose between; call Case(...), in Dag 'd'",
+ error.message,
+ )
+ }
+
+ @Test
+ @DisplayName("Should reject a switch over a task that chooses nothing")
+ fun shouldRejectSwitchOverAPlainTask() {
+ val dag = DagDef("d")
+
+ val error =
+ assertThrows(IllegalArgumentException::class.java) {
+ SwitchRef.of(dag.task<Unit>("plain", NoopTask::class.java))
+ }
+
+ assertEquals(
+ "Task 'plain' runs 'org.apache.airflow.sdk.NoopTask', which chooses nothing; " +
+ "a switch runs a SwitchTask",
+ error.message,
+ )
+ }
+
+ @Test
+ @DisplayName("Should reject naming a case after the Dag was registered")
+ fun shouldRejectCaseAfterRegistration() {
+ val (dag, switch, cases) = dagWithCases()
+ switch.Case(cases.first)
+ Bundle().register(dag)
+
+ val error = assertThrows(IllegalArgumentException::class.java) { switch.Case(cases.second) }
+
+ assertEquals(
+ "Switch 'pickPath' of Dag 'd' is already registered; name every case before the Dag is " +
+ "added to a Bundle",
+ error.message,
+ )
+ }
+
+ @Test
+ @DisplayName("Should declare a switch inside a task group under the group's prefix")
+ fun shouldDeclareSwitchInGroup() {
+ val dag = DagDef("d")
+ val reports = dag.taskGroup("reports")
+ val long = reports.task<Unit>("long", HandleLong::class.java)
+ val switch = reports.Switch("pick", PickPath::class.java).Case(long)
+
+ assertEquals("reports.pick", switch.id)
+ assertEquals(listOf("reports.long", "reports.pick"), dag.tasks.keys.toList())
+ }
+
+ @Test
+ @DisplayName("Should record the chosen case's task ID and skip the other case")
+ fun shouldRecordChosenTaskIdAndSkipTheOther() {
+ val (dag, switch, cases) = dagWithCases()
+ switch.Case(cases.first).Case(cases.second)
+ val pick = dag.tasks.getValue("pickPath")
+
+ val decision = pick.decider!!.decide(pick, PickPath(), taskContext(), clientWith(null).first)
+
+ assertEquals("handle_long", decision.value)
+ assertEquals(listOf("handle_short"), decision.skipped)
+ }
+
+ @Test
+ @DisplayName("Should reuse the switch when a task already declared as one is declared again")
+ fun shouldReuseSwitchDeclaredTwice() {
+ val (dag, switch, cases) = dagWithCases()
+ val again = SwitchRef.of(TaskRef<Unit>(dag.tasks.getValue("pickPath")))
+
+ switch.Case(cases.first)
+ again.Case(cases.second)
+
+ assertEquals(
+ listOf("handle_long", "handle_short"),
+ dag.tasks
+ .getValue("pickPath")
+ .decider!!
+ .cases
+ .map { it.id },
+ )
+ }
+
+ @Test
+ @DisplayName("Should reject declaring a condition as a switch")
+ fun shouldRejectSwitchOverACondition() {
+ val dag = DagDef("d")
+ val condition = dag.If(HasRows::class.java)
+
+ val error = assertThrows(IllegalArgumentException::class.java) { SwitchRef.of(TaskRef<Unit>(condition.nodes().single())) }
+
+ assertEquals("Task 'hasRows' already decides what to skip as a condition", error.message)
+ }
+}
diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/conformance/SerializeJava.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/conformance/SerializeJava.kt
index 4c2f7f84e4e..e327afd661e 100644
--- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/conformance/SerializeJava.kt
+++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/conformance/SerializeJava.kt
@@ -39,6 +39,8 @@ import org.apache.airflow.sdk.ConditionTask
import org.apache.airflow.sdk.Context
import org.apache.airflow.sdk.DagDef
import org.apache.airflow.sdk.Deps
+import org.apache.airflow.sdk.SwitchRef
+import org.apache.airflow.sdk.SwitchTask
import org.apache.airflow.sdk.Task
import org.apache.airflow.sdk.TaskGroupRef
import org.apache.airflow.sdk.TaskRef
@@ -50,13 +52,108 @@ import java.io.File
import java.time.Duration
import java.time.OffsetDateTime
-class ConformanceTask : Task {
+// A switch names its cases by class, so each task of a Dag runs a class of its own. The pool has
+// room for the largest Dag in test_dags.yaml and a few more; the serialized task type is not compared.
+class ConformanceTask1 : Task {
override fun execute(
context: Context,
client: Client,
) = Unit
}
+class ConformanceTask2 : Task {
+ override fun execute(
+ context: Context,
+ client: Client,
+ ) = Unit
+}
+
+class ConformanceTask3 : Task {
+ override fun execute(
+ context: Context,
+ client: Client,
+ ) = Unit
+}
+
+class ConformanceTask4 : Task {
+ override fun execute(
+ context: Context,
+ client: Client,
+ ) = Unit
+}
+
+class ConformanceTask5 : Task {
+ override fun execute(
+ context: Context,
+ client: Client,
+ ) = Unit
+}
+
+class ConformanceTask6 : Task {
+ override fun execute(
+ context: Context,
+ client: Client,
+ ) = Unit
+}
+
+class ConformanceTask7 : Task {
+ override fun execute(
+ context: Context,
+ client: Client,
+ ) = Unit
+}
+
+class ConformanceTask8 : Task {
+ override fun execute(
+ context: Context,
+ client: Client,
+ ) = Unit
+}
+
+class ConformanceTask9 : Task {
+ override fun execute(
+ context: Context,
+ client: Client,
+ ) = Unit
+}
+
+class ConformanceTask10 : Task {
+ override fun execute(
+ context: Context,
+ client: Client,
+ ) = Unit
+}
+
+class ConformanceTask11 : Task {
+ override fun execute(
+ context: Context,
+ client: Client,
+ ) = Unit
+}
+
+class ConformanceTask12 : Task {
+ override fun execute(
+ context: Context,
+ client: Client,
+ ) = Unit
+}
+
+private val taskPool: List<Class<out Task>> =
+ listOf(
+ ConformanceTask1::class.java,
+ ConformanceTask2::class.java,
+ ConformanceTask3::class.java,
+ ConformanceTask4::class.java,
+ ConformanceTask5::class.java,
+ ConformanceTask6::class.java,
+ ConformanceTask7::class.java,
+ ConformanceTask8::class.java,
+ ConformanceTask9::class.java,
+ ConformanceTask10::class.java,
+ ConformanceTask11::class.java,
+ ConformanceTask12::class.java,
+ )
+
class ConformanceCondition : ConditionTask {
override fun decide(
context: Context,
@@ -64,6 +161,14 @@ class ConformanceCondition : ConditionTask {
) = true
}
+// Never run, only serialized, so any task class will do.
+class ConformanceSwitch : SwitchTask {
+ override fun choose(
+ context: Context,
+ client: Client,
+ ): Class<out Task> = ConformanceTask1::class.java
+}
+
fun main(args: Array<String>) {
require(args.size == 2) { "usage: SerializeJava <test_dags.yaml> <output.json>" }
val cases = YAMLMapper().readTree(File(args[0])).path("dags")
@@ -90,20 +195,33 @@ private fun buildDag(case: JsonNode): DagDef {
}
val tasks = linkedMapOf<String, TaskRef<*>>()
- // A condition names tasks that may be declared after it, so the sides are wired once every task exists.
- val conditions = mutableListOf<Pair<ConditionRef, JsonNode>>()
+ val unusedTaskClasses = taskPool.iterator()
+ // A decider names tasks that may be declared after it, so its cases are wired once every task exists.
+ val deciders = mutableListOf<Pair<Deps.Flow, JsonNode>>()
case.path("tasks").forEach { task ->
val groupId = task.path("group").asText("")
val localId = task.path("task_id").asText()
val branch = task.path("branch")
- val definition = if (branch.isMissingNode) ConformanceTask::class.java else ConformanceCondition::class.java
+ val definition =
+ when {
+ branch.isMissingNode -> {
+ require(unusedTaskClasses.hasNext()) {
+ "Dag '${dag.id}' has more tasks than the ${taskPool.size} classes in taskPool; add a ConformanceTask class"
+ }
+ unusedTaskClasses.next()
+ }
+ branch.has("cases") -> ConformanceSwitch::class.java
+ else -> ConformanceCondition::class.java
+ }
val ref =
if (groupId.isEmpty()) {
dag.task<Any?>(localId, definition)
} else {
groups.getValue(groupId).task<Any?>(localId, definition)
}
- if (!branch.isMissingNode) conditions += asCondition(ref) to branch
+ if (!branch.isMissingNode) {
+ deciders += (if (branch.has("cases")) SwitchRef.of(ref) else asCondition(ref)) to branch
+ }
task.path("spec").fields().forEach { (key, value) -> ref.config(key, toValue(SchemaFields.TASK, key, value)) }
// A task's `upstream` handles and its `literals` are its call arguments, in that order, so the
// Dag carries the binding spec a stub call would. Names are positional, as the Go SDK names
@@ -117,9 +235,15 @@ private fun buildDag(case: JsonNode): DagDef {
tasks[ref.def.id] = ref
}
- conditions.forEach { (condition, branch) ->
- condition.Then(tasks.getValue(branch.path("then").asText()))
- branch.path("else").takeIf { !it.isMissingNode }?.let { condition.Else(tasks.getValue(it.asText())) }
+ deciders.forEach { (decider, branch) ->
+ when (decider) {
+ is SwitchRef -> branch.path("cases").forEach { decider.Case(tasks.getValue(it.asText())) }
+ is ConditionRef -> {
+ decider.Then(tasks.getValue(branch.path("then").asText()))
+ branch.path("else").takeIf { !it.isMissingNode }?.let { decider.Else(tasks.getValue(it.asText())) }
+ }
+ else -> throw IllegalStateException("Unknown decider: $decider")
+ }
}
case.path("order_edges").forEach { edge ->
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 a1ca3f7553d..db1478661c1 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
@@ -24,6 +24,7 @@ import org.apache.airflow.sdk.Client
import org.apache.airflow.sdk.ConditionTask
import org.apache.airflow.sdk.Context
import org.apache.airflow.sdk.DagDef
+import org.apache.airflow.sdk.SwitchTask
import org.apache.airflow.sdk.Task
import org.apache.airflow.sdk.TaskDef
import org.apache.airflow.sdk.execution.comm.BundleInfo
@@ -245,6 +246,62 @@ class TaskTest {
Assertions.assertEquals(emptyList<Pair<String, Any>>(), transport.xComs)
}
+ @Test
+ @DisplayName("Should push the case a switch chose and skip every other one")
+ fun shouldSkipTheCasesNotChosen() {
+ TestSwitch.choice = HandleLongTask::class.java
+ val transport = RecordingTransport()
+
+ val result = runTask(switchBundle(), startupDetails(taskId = "pick"), Client(startupDetails("pick"), transport))
+
+ Assertions.assertInstanceOf(SucceedTask::class.java, result)
+ Assertions.assertEquals(listOf("handle_short"), transport.skipped)
+ Assertions.assertEquals(
+ listOf("skipmixin_key" to mapOf("skipped" to listOf("handle_short")), "return_value" to "handle_long"),
+ transport.xComs,
+ )
+ }
+
+ @Test
+ @DisplayName("Should fail a switch that chose a task it cannot run")
+ fun shouldFailSwitchThatChoseAnUnknownCase() {
+ TestSwitch.choice = SuccessTask::class.java
+ val transport = RecordingTransport()
+
+ val result = runTask(switchBundle(), startupDetails(taskId = "pick"), Client(startupDetails("pick"), transport))
+
+ Assertions.assertInstanceOf(TaskState::class.java, result)
+ Assertions.assertEquals(TaskState.State.FAILED, (result as TaskState).state)
+ Assertions.assertEquals(emptyList<String>(), transport.skipped)
+ Assertions.assertEquals(emptyList<Pair<String, Any>>(), transport.xComs)
+ }
+
+ @Test
+ @DisplayName("Should fail a generated switch that chose a task of the Dag that is not one of its cases")
+ fun shouldFailGeneratedSwitchThatChoseANonCase() {
+ val dag = DagDef("test_dag")
+ val long = dag.task<Unit>("handle_long", HandleLongTask::class.java)
+ val short = dag.task<Unit>("handle_short", HandleShortTask::class.java)
+ dag.task<Unit>("unlisted", UnlistedTask::class.java)
+ dag.Switch("pick", TestGeneratedSwitch::class.java).Case(long).Case(short)
+ val transport = RecordingTransport()
+
+ val result = runTask(Bundle(listOf(dag)), startupDetails(taskId = "pick"), Client(startupDetails("pick"), transport))
+
+ Assertions.assertInstanceOf(TaskState::class.java, result)
+ Assertions.assertEquals(TaskState.State.FAILED, (result as TaskState).state)
+ Assertions.assertEquals(emptyList<String>(), transport.skipped)
+ Assertions.assertEquals(emptyList<Pair<String, Any>>(), transport.xComs)
+ }
+
+ private fun switchBundle(): Bundle {
+ val dag = DagDef("test_dag")
+ val long = dag.task<Unit>("handle_long", HandleLongTask::class.java)
+ val short = dag.task<Unit>("handle_short", HandleShortTask::class.java)
+ dag.Switch("pick", TestSwitch::class.java).Case(long).Case(short)
+ return Bundle(listOf(dag))
+ }
+
private fun conditionBundle(withElse: Boolean): Bundle {
val dag = DagDef("test_dag")
val load = dag.task<Unit>("load", SuccessTask::class.java)
@@ -360,6 +417,47 @@ class TaskTest {
}
}
+ class HandleLongTask : Task {
+ override fun execute(
+ context: Context,
+ client: Client,
+ ) = Unit
+ }
+
+ class HandleShortTask : Task {
+ override fun execute(
+ context: Context,
+ client: Client,
+ ) = Unit
+ }
+
+ /** Chooses whatever [choice] holds, so a test can drive it to a case or past one. */
+ class TestSwitch : SwitchTask {
+ override fun choose(
+ context: Context,
+ client: Client,
+ ): Class<out Task> = choice
+
+ companion object {
+ var choice: Class<out Task> = HandleLongTask::class.java
+ }
+ }
+
+ class UnlistedTask : Task {
+ override fun execute(
+ context: Context,
+ client: Client,
+ ) = Unit
+ }
+
+ /** The shape the annotation processor generates: it returns the class of a task, here one that is no case. */
+ class TestGeneratedSwitch : SwitchTask {
+ override fun choose(
+ context: Context,
+ client: Client,
+ ): Class<out Task> = UnlistedTask::class.java
+ }
+
/** Records what a deciding task pushed and asked to skip, and the order of the two. */
private class RecordingTransport : org.apache.airflow.sdk.execution.Client {
val xComs = mutableListOf<Pair<String, Any>>()
diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/RefsTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/RefsTest.kt
index 59e36c182bd..c764585e409 100644
--- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/RefsTest.kt
+++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/RefsTest.kt
@@ -27,6 +27,8 @@ import org.apache.airflow.sdk.Context
import org.apache.airflow.sdk.DagDef
import org.apache.airflow.sdk.Deps
import org.apache.airflow.sdk.LiteralArg
+import org.apache.airflow.sdk.SwitchRef
+import org.apache.airflow.sdk.SwitchTask
import org.apache.airflow.sdk.Task
import org.apache.airflow.sdk.TaskDef
import org.apache.airflow.sdk.TaskRef
@@ -56,6 +58,20 @@ private class NoopCondition : ConditionTask {
) = true
}
+private class OtherNoopRefTask : Task {
+ override fun execute(
+ context: Context,
+ client: Client,
+ ) = Unit
+}
+
+private class NoopSwitch : SwitchTask {
+ override fun choose(
+ context: Context,
+ client: Client,
+ ): Class<out Task> = NoopRefTask::class.java
+}
+
internal class RefsTest {
@Test
@DisplayName("Should register the task, record inputs, and wire handle edges")
@@ -232,4 +248,24 @@ internal class RefsTest {
.map { it.id },
)
}
+
+ @Test
+ @DisplayName("Should let a view name a switch's cases in separate statements")
+ fun shouldNameSwitchCasesAcrossStatements() {
+ val dag = DagDef("d")
+ Refs.record(dag, listOf("pick", "long", "short"), emptyList()) {
+ fun pick() = SwitchRef.of(Refs.node<Unit>("", TaskDef("pick", NoopSwitch::class.java)))
+ pick().Case(Refs.node<Unit>("", TaskDef("long", NoopRefTask::class.java)))
+ pick().Case(Refs.node<Unit>("", TaskDef("short", OtherNoopRefTask::class.java)))
+ }
+
+ assertEquals(
+ listOf("long", "short"),
+ dag.tasks
+ .getValue("pick")
+ .decider!!
+ .cases
+ .map { it.id },
+ )
+ }
}