From f3c2472e2446a8d4fe3e6b86a195a9b5834b6b39 Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Thu, 8 Oct 2026 04:44:40 +0000 Subject: [PATCH 1/7] 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 `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. --- .pre-commit-config.yaml | 2 +- java-sdk/README.md | 2 +- java-sdk/capabilities.yaml | 2 +- .../example/nativedag/AnnotationExample.java | 21 ++ .../example/nativedag/InterfaceExample.java | 32 +- .../apache/airflow/sdk/BuilderProcessor.kt | 162 ++++++++-- .../org/apache/airflow/sdk/BuilderTest.kt | 258 +++++++++++++++- .../prek/check_serialization_conformance.py | 2 +- java-sdk/sdk/build.gradle.kts | 37 +++ java-sdk/sdk/module.md | 2 +- .../kotlin/org/apache/airflow/sdk/Branch.kt | 291 ++++++++++++++++++ .../kotlin/org/apache/airflow/sdk/Bundle.kt | 7 +- .../kotlin/org/apache/airflow/sdk/DagDef.kt | 41 +++ .../org/apache/airflow/sdk/TaskGroupRef.kt | 30 ++ .../org/apache/airflow/sdk/BranchTest.kt | 175 +++++++++++ .../org/apache/airflow/sdk/ConditionTest.kt | 2 +- .../airflow/sdk/conformance/SerializeJava.kt | 39 ++- .../apache/airflow/sdk/execution/TaskTest.kt | 91 ++++++ 18 files changed, 1148 insertions(+), 48 deletions(-) create mode 100644 java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Branch.kt create mode 100644 java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/BranchTest.kt diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 192f1dcddeb1b..323828a34f094 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -444,7 +444,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/(Branch|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 1e4e7478f3942..58e3bcad92393 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.Branch, or @Builder.If and @Builder.Branch | | 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 33106f170fb3e..b151d2562baa7 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.Branch, or @Builder.If and @Builder.Branch" 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 d9538fd391f7d..d5f66ac64affc 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,26 @@ public boolean hasRows(long transformed) { 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 branch: it names the one case that runs, and every other case is + // skipped. TaskIds is generated beside the builder, so a case that is not a + // task of this Dag does not compile. + @Builder.Branch(id = "pick_report") + public TaskId pickReport(long transformed) { + return transformed > 100 + ? AnnotationExampleBuilder.TaskIds.REPORT_LONG + : AnnotationExampleBuilder.TaskIds.REPORT_SHORT; + } + // A task group: everything it declares is prefixed with its id, so this is // the task "checks.audit". @Builder.TaskGroup(id = "checks") @@ -84,6 +104,7 @@ void depends() { var extracted = extract(); var transformed = transform(extracted, lit(1.5)); hasRows(transformed).Then(load(transformed)).Else(loadEmpty()); + pickReport(transformed).option(reportLong()).option(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 cce10676dfd8d..b6190406c5fe5 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 boolean decide(Context context, Client client) { } } + 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 branch: it names the one case that runs by its class, and every other + // case is skipped. + public static class PickReport implements BranchTask { + @Override + public Class 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 static DagDef build() { 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.Branch(PickReport.class).after(transform).option(reportLong).option(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 1848fc007f78f..bce9de83dd0d6 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 @@ -23,6 +23,7 @@ package org.apache.airflow.sdk import com.squareup.javapoet.ClassName import com.squareup.javapoet.CodeBlock +import com.squareup.javapoet.FieldSpec import com.squareup.javapoet.JavaFile import com.squareup.javapoet.MethodSpec import com.squareup.javapoet.ParameterizedTypeName @@ -49,6 +50,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 @@ -97,6 +99,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.Branch", "org.apache.airflow.sdk.Builder.TaskGroup", "org.apache.airflow.sdk.Builder.TaskHandler", "org.apache.airflow.sdk.Builder.Deps", @@ -275,9 +278,38 @@ class BuilderProcessor : AbstractProcessor() { builderClass.addMethod(buildMethod.build()) declarations.forEach { builderClass.addType(buildTask(it)) } + // Only a branch needs to name a task in code, so a Dag without one keeps + // the generated builder to the classes that run its tasks. + if (declarations.any { it.kind == TaskKind.BRANCH }) builderClass.addType(buildTaskIds(declarations)) return builderClass.build() } + /** + * Generates the `TaskIds` holder a `@Builder.Branch` method names its case + * with: one constant per task of the Dag, so a choice that is not a task of + * this Dag does not compile. + */ + private fun buildTaskIds(declarations: List): TypeSpec { + val holder = + TypeSpec + .classBuilder(TASK_IDS) + .addModifiers(Modifier.PUBLIC, Modifier.STATIC, Modifier.FINAL) + .addJavadoc( + "The task ids of this Dag, for a {@code @Builder.Branch} method to name its case with.\n\n" + + "

Every task of the Dag has one, so naming a task that is not a case of the branch\n" + + "compiles and fails when the branch runs.\n", + ) + declarations.forEach { decl -> + holder.addField( + FieldSpec + .builder(TASK_ID_TYPE, constantName(decl.id), Modifier.PUBLIC, Modifier.STATIC, Modifier.FINAL) + .initializer($$"$T.of($S)", TASK_ID_TYPE, decl.id) + .build(), + ) + } + return holder.build() + } + /** * Generates the Dag's wiring view: one default method per task, with the * injected arguments stripped, each data argument lifted to [Arg] and the @@ -357,24 +389,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 +465,19 @@ class BuilderProcessor : AbstractProcessor() { val tasks = mutableListOf() 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 " + @@ -567,6 +599,29 @@ class BuilderProcessor : AbstractProcessor() { "Dag has both a task and a task group with ID '${group.fullId}'; rename one" } } + if (declarations.any { it.kind == TaskKind.BRANCH }) { + val byConstant = mutableMapOf() + declarations.forEach { decl -> + val constant = constantName(decl.id) + require(SourceVersion.isName(constant)) { + "Task '${decl.id}' becomes the constant '$constant' of the generated TaskIds, which is not " + + "a Java name; give the task an id a branch can name it by, with ${decl.kind.spelling}(id = \"...\")" + } + byConstant.put(constant, decl)?.let { first -> + throw IllegalArgumentException( + "Tasks '${first.id}' and '${decl.id}' both become the constant " + + "'$constant' of the generated TaskIds; rename one so a branch can tell them apart", + ) + } + } + declarations.firstOrNull { it.className == TASK_IDS }?.let { decl -> + throw IllegalArgumentException( + "Task method '${decl.method.simpleName}' generates the class '$TASK_IDS', which is the " + + "holder of this Dag's task ids; rename the method and keep its task id with " + + "${decl.kind.spelling}(id = \"${decl.id.substringAfterLast('.')}\")", + ) + } + } val byClassName = mutableMapOf() declarations.forEach { decl -> byClassName.put(decl.className, decl)?.let { first -> @@ -712,24 +767,45 @@ 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.BRANCH -> + require(with(processingEnv) { isType(returns, TASK_ID_TYPE) }) { + "@Builder.Branch method '${method.simpleName}' returns $returns, but a branch returns a " + + "TaskId: name the case it chose with a constant of the generated TaskIds" + } } } + /** 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.BRANCH -> TASK_ID_TYPE + }, + ).addParameter(CONTEXT_TYPE, "context") .addParameter(CLIENT_TYPE, "client") .addException(Exception::class.java) @@ -775,8 +851,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 +866,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 +1042,15 @@ 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 BRANCH_TASK_TYPE = ClassName.get(TaskIdBranchTask::class.java) +private val BRANCH_REF_TYPE = ClassName.get(BranchRef::class.java) +private val TASK_ID_TYPE = ClassName.get(TaskId::class.java) + +/** Name of the generated holder of a Dag's task ids, which no task class may take. */ +private const val TASK_IDS = "TaskIds" 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 +1066,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), + BRANCH("org.apache.airflow.sdk.Builder.Branch", "@Builder.Branch", "choose", BRANCH_TASK_TYPE, BRANCH_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", @@ -1018,6 +1109,17 @@ private val RESERVED_GROUP_VIEW_NAMES: Set = private val DAG_STRUCTURAL_ATTRIBUTES = setOf("id", "to") private val TASK_STRUCTURAL_ATTRIBUTES = setOf("id") +/** + * The `TaskIds` constant for a task id: its words in upper case, joined by + * underscores, so `handleLong` and `checks.audit` become `HANDLE_LONG` and + * `CHECKS_AUDIT`. + */ +private fun constantName(id: String): String = + id + .replace(Regex("([a-z0-9])([A-Z])"), "$1_$2") + .replace(Regex("[^A-Za-z0-9]+"), "_") + .uppercase() + private fun TypeName.boxIfPossible(): TypeName = if (this == TypeName.VOID || isPrimitive) box() else this private fun ProcessingEnvironment.isType( 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 70802c405aa2a..4193bc211b6ea 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,262 @@ class BuilderTest { assertThat(compilation).succeeded() } + @Test + @DisplayName("generate a branch that names its case with a generated task-id constant") + fun generateBranchTask() { + val compilation = + compile( + """ + package org.apache.airflow.example; + + import org.apache.airflow.sdk.Builder; + import org.apache.airflow.sdk.TaskId; + + @Builder.Dag(id = "etl") + public class TestExample { + @Builder.Branch(id = "pick_path") + public TaskId pickPath() { + return TestExampleBuilder.TaskIds.HANDLE_LONG; + } + + @Builder.Task + public void handleLong() {} + + @Builder.Task + public void handleShort() {} + + @Builder.Deps + static class Wiring implements TestExampleDeps { + void depends() { + pickPath().option(handleLong()).option(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.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.Task; + import org.apache.airflow.sdk.TaskId; + import org.apache.airflow.sdk.TaskIdBranchTask; + 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 TaskIdBranchTask { + @Override + public TaskId 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(); + } + } + + /** + * The task ids of this Dag, for a {@code @Builder.Branch} method to name its case with. + */ + public static final class TaskIds { + public static final TaskId PICK_PATH = TaskId.of("pick_path"); + + public static final TaskId HANDLE_LONG = TaskId.of("handleLong"); + + public static final TaskId HANDLE_SHORT = TaskId.of("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.BranchRef; + import org.apache.airflow.sdk.Deps; + 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 BranchRef pickPath() { + return BranchRef.of(Refs.node("", new TaskDef("pick_path", TestExampleBuilder.PickPath.class))); + } + + default TaskRef handleLong() { + return Refs.node("", new TaskDef("handleLong", TestExampleBuilder.HandleLong.class)); + } + + default TaskRef handleShort() { + return Refs.node("", new TaskDef("handleShort", TestExampleBuilder.HandleShort.class)); + } + } + """, + ) + } + + @Test + @DisplayName("reject a task method whose generated class would take the TaskIds name") + fun rejectTaskNamedTaskIds() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + import org.apache.airflow.sdk.TaskId; + @Builder.Dag + public class TestExample { + @Builder.Branch + public TaskId pick() { + return TestExampleBuilder.TaskIds.TASK_IDS; + } + + @Builder.Task + public void taskIds() {} + + @Builder.Deps + static class Wiring implements TestExampleDeps { + void depends() {} + } + } + """, + ) + + assertThat(compilation).failed() + assertThat(compilation).hadErrorContaining( + "Task method 'taskIds' generates the class 'TaskIds', which is the holder of this Dag's task ids", + ) + } + + @Test + @DisplayName("reject a task id whose generated constant is not a Java name") + fun rejectTaskIdWithoutAJavaName() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + import org.apache.airflow.sdk.TaskId; + @Builder.Dag + public class TestExample { + @Builder.Branch + public TaskId pick() { + return null; + } + + @Builder.Task(id = "2nd_pass") + public void second() {} + + @Builder.Deps + static class Wiring implements TestExampleDeps { + void depends() {} + } + } + """, + ) + + assertThat(compilation).failed() + assertThat(compilation).hadErrorContaining( + "Task '2nd_pass' becomes the constant '2ND_PASS' of the generated TaskIds, which is not a Java name", + ) + } + + @Test + @DisplayName("reject a branch that does not return a task id") + fun rejectNonTaskIdBranch() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + @Builder.Dag + public class TestExample { + @Builder.Branch + public String pickPath() { + return "handleLong"; + } + + @Builder.Deps + static class Wiring implements TestExampleDeps { + void depends() {} + } + } + """, + ) + + assertThat(compilation).failed() + assertThat(compilation).hadErrorContaining( + "@Builder.Branch method 'pickPath' returns java.lang.String, but a branch returns a TaskId", + ) + } + + @Test + @DisplayName("reject two tasks whose generated task-id constants would collide") + fun rejectCollidingTaskIdConstants() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + import org.apache.airflow.sdk.TaskId; + @Builder.Dag + public class TestExample { + @Builder.Branch + public TaskId pick() { + return TestExampleBuilder.TaskIds.HANDLE_LONG; + } + + @Builder.Task + public void handleLong() {} + + @Builder.Task(id = "handle_long") + public void other() {} + + @Builder.Deps + static class Wiring implements TestExampleDeps { + void depends() {} + } + } + """, + ) + + assertThat(compilation).failed() + assertThat(compilation).hadErrorContaining( + "Tasks 'handleLong' and 'handle_long' both become the constant 'HANDLE_LONG' of the generated TaskIds", + ) + } + @Test @DisplayName("reject a condition that does not return a boolean") fun rejectNonBooleanCondition() { @@ -2355,7 +2611,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 fe569c1ff7cd4..01a16e3421987 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 @@ "--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 723f60085c2c7..fe4044913e747 100644 --- a/java-sdk/sdk/build.gradle.kts +++ b/java-sdk/sdk/build.gradle.kts @@ -614,6 +614,43 @@ 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 [TaskId] of the task it chose, which the + | * generated `Builder.TaskIds` declares a constant for, so a + | * choice that is not a task of this Dag does not compile. Its + | * wiring-view method returns a [BranchRef], so the cases are listed + | * where the Dag is wired: + | * + | * ```java + | * @Builder.Branch(id = "pick_path") + | * public TaskId pickPath(long rows) { + | * return rows > 1000 ? EtlBuilder.TaskIds.HANDLE_LONG : EtlBuilder.TaskIds.HANDLE_SHORT; + | * } + | * + | * @Builder.Deps + | * static class Wiring implements EtlDeps { + | * void depends() { pickPath(extract()).option(handleLong()).option(handleShort()); } + | * } + | * ``` + | * + | * A branch 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 Branch( + | /** 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 523dbaf54a4b3..9a74418e3ab6a 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.Branch, or @Builder.If and @Builder.Branch | | 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/Branch.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Branch.kt new file mode 100644 index 0000000000000..ce444e448d17d --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Branch.kt @@ -0,0 +1,291 @@ +/* + * 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 + +/** + * The Airflow task ID of one task of a Dag. + * + * A `@Builder.Branch` method names the task it chose with one of these. The + * generated `Builder.TaskIds` holds a constant per task of the Dag, so + * the choice is checked where it is written rather than when the task runs. + * + * @property value The task ID. + */ +class TaskId private constructor( + val value: String, +) { + companion object { + /** + * @suppress + * + * The ID [value] names. Public so a generated `TaskIds` holder can declare + * its constants; user code names a task through those constants. + */ + @JvmStatic + fun of(value: String): TaskId = TaskId(value) + } + + override fun equals(other: Any?): Boolean = other is TaskId && other.value == value + + override fun hashCode(): Int = value.hashCode() + + override fun toString(): String = value +} + +/** + * A task that chooses one of several tasks to run; every other one is skipped. + * + * Register one with [DagDef.Branch], then list what it can choose with + * [BranchRef.option]. The choice is the task's own class, so javac checks it: + * + * ```java + * public class PickPath implements BranchTask { + * @Override + * public Class 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. + * + * Because a branch names a case by its class, no two of its cases may be + * registered from the same class; registering the Dag reports that. + * + * @see DagDef.Branch + */ +interface BranchTask : 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 branch's cases. + * @throws Exception on failure; the task instance is marked failed. + */ + @Throws(Exception::class) + fun choose( + context: Context, + client: Client, + ): Class + + /** Never called: the SDK runs a branch through [choose]. */ + override fun execute( + context: Context, + client: Client, + ): Unit = + throw IllegalStateException( + "Branch '${javaClass.name}' runs through choose(), so execute() is never called", + ) +} + +/** + * @suppress + * + * What a generated `@Builder.Branch` task implements: a branch whose method + * named the case with a [TaskId]. Public so generated code can implement it; + * a Dag written against [DagDef.Branch] implements [BranchTask] instead. + */ +interface TaskIdBranchTask : Task { + @Throws(Exception::class) + fun choose( + context: Context, + client: Client, + ): TaskId + + override fun execute( + context: Context, + client: Client, + ): Unit = + throw IllegalStateException( + "Branch '${javaClass.name}' runs through choose(), so execute() is never called", + ) +} + +/** + * A branch registered with a Dag: list the tasks it can choose between. + * + * ```java + * dag.Branch(PickPath.class).option(handleLong).option(handleShort); + * ``` + * + * A case runs after the branch, so naming it records that edge, as + * [Deps.Flow.before] would. The branch 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.Branch + */ +class BranchRef private constructor( + private val ref: TaskRef<*>, + private val cases: BranchDef, +) : Deps.Flow { + companion object { + /** + * @suppress + * + * Marks a registered task as a branch. Public so a generated wiring view + * can call it; user code reaches a branch through [DagDef.Branch]. + */ + @JvmStatic + fun of(ref: TaskRef<*>): BranchRef { + require(ref.def.decider == null) { + "Task '${ref.def.id}' already decides what to skip; declare it once" + } + val definition = ref.def.definition + require( + BranchTask::class.java.isAssignableFrom(definition) || + TaskIdBranchTask::class.java.isAssignableFrom(definition), + ) { + "Task '${ref.def.id}' runs '${definition.name}', which chooses nothing; a branch runs a " + + "BranchTask" + } + return BranchRef(ref, BranchDef().also { ref.def.decider = it }) + } + } + + /** Airflow task ID of the deciding task. */ + val id: String get() = ref.def.id + + /** + * Adds a task this branch can choose. Every other case is skipped when the + * branch chooses this one. + * + * @param task Task of the same Dag. + * @return This branch, for chaining. + * @throws IllegalArgumentException if the task belongs to another Dag or is + * already a case of this branch. + */ + fun option(task: TaskRef<*>): BranchRef { + cases.add(ref.def, task.def) + // The case runs after the branch, which is what puts the branch in the + // serialized Dag as its upstream and lets registration see a cycle that + // runs through a branch. + 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 branch, for chaining. + * @throws IllegalArgumentException if the key is unknown or the value type + * does not match. + */ + fun config( + key: String, + value: Any?, + ): BranchRef { + ref.config(key, value) + return this + } + + override fun nodes(): List = ref.nodes() + + override fun before(vararg next: Deps.Flow): BranchRef { + ref.before(*next) + return this + } + + override fun after(vararg previous: Deps.Flow): BranchRef { + ref.after(*previous) + return this + } +} + +/** The cases of one branch, as [BranchRef.option] listed them. */ +internal class BranchDef : DeciderDef { + private val options = mutableListOf() + + override val cases: List get() = options + + fun add( + decider: TaskDef, + case: TaskDef, + ) { + requireUnregistered(decider, "Branch") + requireSameDag(decider, case, "Branch") + require(options.none { it === case }) { + "Branch '${decider.id}' already chooses between '${case.id}' and others; name each case once" + } + // A BranchTask names its case by class, so two cases sharing one would be + // indistinguishable. Checked here rather than at registration so the error + // points at the second option(...) call. + if (BranchTask::class.java.isAssignableFrom(decider.definition)) { + options.firstOrNull { it.definition == case.definition }?.let { first -> + throw IllegalArgumentException( + "Branch '${decider.id}' cannot choose between '${first.id}' and '${case.id}': both run " + + "'${case.definition.name}', and a branch names its case by class", + ) + } + } + options += case + } + + override fun describe(decider: TaskDef): String? = + if (options.isEmpty()) "Branch '${decider.id}' has no task to choose between; call option(...)" else null + + override fun decide( + decider: TaskDef, + instance: Task, + context: Context, + client: Client, + ): Decision { + val chosen = + when (instance) { + is BranchTask -> byClass(decider, instance.choose(context, client)) + is TaskIdBranchTask -> byId(decider, instance.choose(context, client).value) + else -> throw IllegalStateException( + "Task '${decider.id}' is a branch, but '${instance.javaClass.name}' is not a BranchTask", + ) + } + return Decision(chosen.id, options.filterNot { it === chosen }.map { it.id }) + } + + private fun byClass( + decider: TaskDef, + definition: Class?, + ): TaskDef = + options.firstOrNull { it.definition == definition } + ?: throw IllegalArgumentException( + "Branch '${decider.id}' chose ${definition?.name ?: "nothing"}, which is not one of its cases: ${describeCases()}", + ) + + private fun byId( + decider: TaskDef, + taskId: String, + ): TaskDef = + options.firstOrNull { it.id == taskId } + ?: throw IllegalArgumentException( + "Branch '${decider.id}' chose '$taskId', 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/Bundle.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Bundle.kt index 7937446c3f648..d34927d28b6ac 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 `Branch`, 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 Branch(...), 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, BranchTask::class.java, TaskIdBranchTask::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/DagDef.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/DagDef.kt index 40ead4ae60856..ae5556ff4bd01 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, ): 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 + * [Branch(id, definition)][Branch] to set it, and [BranchRef.config] for the + * task's other settings: + * + * ```java + * dag.Branch(PickPath.class).option(handleLong).option(handleShort); + * ``` + * + * @param definition Class that implements [BranchTask]. Must have a public + * no-arg constructor. + * @return The branch, to list its cases on. + * @throws IllegalArgumentException if a task with the same ID is already + * registered. + */ + @Suppress("ktlint:standard:function-naming") + fun Branch(definition: Class): BranchRef = Branch(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 [BranchTask]. Must have a public + * no-arg constructor. + * @return The branch, to list its cases on. + * @throws IllegalArgumentException if a task with the same ID is already + * registered. + * + * @see Branch + */ + @Suppress("ktlint:standard:function-naming") + fun Branch( + id: String, + definition: Class, + ): BranchRef = BranchRef.of(task(id, definition)) + /** * Declares a task group of this Dag. * 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 e2d7d9613ef4d..3fb8c3e7d5063 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, ): ConditionRef = ConditionRef.of(task(id, definition)) + /** + * Declares a branch in this group, as [DagDef.Branch] does for the Dag. + * + * @param definition Class that implements [BranchTask]. Must have a public + * no-arg constructor. + * @return The branch, 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 Branch(definition: Class): BranchRef = Branch(deriveTaskId(definition), definition) + + /** + * Declares a branch in this group under the task ID `.`. + * + * @param id Task ID within this group. + * @param definition Class that implements [BranchTask]. Must have a public + * no-arg constructor. + * @return The branch, to list its cases on. + * @throws IllegalArgumentException if the Dag already has a task or task + * group with the resulting ID. + * + * @see Branch + */ + @Suppress("ktlint:standard:function-naming") + fun Branch( + id: String, + definition: Class, + ): BranchRef = BranchRef.of(task(id, definition)) + /** * Nests a task group inside this one. * diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/BranchTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/BranchTest.kt new file mode 100644 index 0000000000000..e37d7a38d3719 --- /dev/null +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/BranchTest.kt @@ -0,0 +1,175 @@ +/* + * 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 branches under test; a second class so a branch can tell two cases apart. */ +internal class HandleLong : Task { + override fun execute( + context: Context, + client: Client, + ) = Unit +} + +/** A branch that chooses whatever [choice] holds. */ +class PickPath : BranchTask { + override fun choose( + context: Context, + client: Client, + ): Class = choice + + companion object { + var choice: Class = HandleLong::class.java + } +} + +internal class BranchTest { + private fun dagWithCases(): Triple, TaskRef>> { + val dag = DagDef("d") + val handleLong = dag.task("handle_long", HandleLong::class.java) + val handleShort = dag.task("handle_short", NoopTask::class.java) + return Triple(dag, dag.Branch(PickPath::class.java), handleLong to handleShort) + } + + @Test + @DisplayName("Should take the branch's task ID from its class when none is given") + fun shouldDeriveTaskIdFromClass() { + val (dag, branch, cases) = dagWithCases() + branch.option(cases.first) + + assertEquals("pickPath", branch.id) + assertEquals(setOf(dag.tasks.getValue("pickPath")), cases.first.def.upstreams) + } + + @Test + @DisplayName("Should reject naming the same case twice") + fun shouldRejectDuplicateCase() { + val (_, branch, cases) = dagWithCases() + branch.option(cases.first) + + val error = assertThrows(IllegalArgumentException::class.java) { branch.option(cases.first) } + + assertEquals( + "Branch 'pickPath' already chooses between 'handle_long' and others; name each case once", + error.message, + ) + } + + @Test + @DisplayName("Should reject two cases a class-named branch could not tell apart") + fun shouldRejectCasesSharingAClass() { + val dag = DagDef("d") + val first = dag.task("first", HandleLong::class.java) + val second = dag.task("second", HandleLong::class.java) + val branch = dag.Branch(PickPath::class.java).option(first) + + val error = assertThrows(IllegalArgumentException::class.java) { branch.option(second) } + + assertEquals( + "Branch 'pickPath' cannot choose between 'first' and 'second': both run " + + "'org.apache.airflow.sdk.HandleLong', and a branch names its case by class", + error.message, + ) + } + + @Test + @DisplayName("Should reject a case declared in another Dag") + fun shouldRejectCaseFromAnotherDag() { + val (_, branch, _) = dagWithCases() + val other = DagDef("other").task("handle_long", HandleLong::class.java) + + val error = assertThrows(IllegalArgumentException::class.java) { branch.option(other) } + + assertEquals( + "Branch '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 branch with nothing to choose between") + fun shouldRejectBranchWithoutCases() { + val (dag, _, _) = dagWithCases() + + val error = assertThrows(IllegalArgumentException::class.java) { Bundle().register(dag) } + + assertEquals( + "Branch 'pickPath' has no task to choose between; call option(...), in Dag 'd'", + error.message, + ) + } + + @Test + @DisplayName("Should reject a branch over a task that chooses nothing") + fun shouldRejectBranchOverAPlainTask() { + val dag = DagDef("d") + + val error = + assertThrows(IllegalArgumentException::class.java) { + BranchRef.of(dag.task("plain", NoopTask::class.java)) + } + + assertEquals( + "Task 'plain' runs 'org.apache.airflow.sdk.NoopTask', which chooses nothing; " + + "a branch runs a BranchTask", + error.message, + ) + } + + @Test + @DisplayName("Should reject naming a case after the Dag was registered") + fun shouldRejectCaseAfterRegistration() { + val (dag, branch, cases) = dagWithCases() + branch.option(cases.first) + Bundle().register(dag) + + val error = assertThrows(IllegalArgumentException::class.java) { branch.option(cases.second) } + + assertEquals( + "Branch '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 branch inside a task group under the group's prefix") + fun shouldDeclareBranchInGroup() { + val dag = DagDef("d") + val reports = dag.taskGroup("reports") + val long = reports.task("long", HandleLong::class.java) + val branch = reports.Branch("pick", PickPath::class.java).option(long) + + assertEquals("reports.pick", branch.id) + assertEquals(listOf("reports.long", "reports.pick"), dag.tasks.keys.toList()) + } + + @Test + @DisplayName("Should carry the task ID of a case rather than its class") + fun taskIdNamesOneTask() { + assertEquals(TaskId.of("handle_long"), TaskId.of("handle_long")) + assertEquals("handle_long", TaskId.of("handle_long").toString()) + } +} 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 1cebf7ff8e1b1..eca33565c5b35 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 Branch(...), in Dag 'd'", 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 4c2f7f84e4eea..9e12006485a5b 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 @@ -32,6 +32,7 @@ import com.fasterxml.jackson.databind.JsonNode import com.fasterxml.jackson.databind.ObjectMapper import com.fasterxml.jackson.dataformat.yaml.YAMLMapper import org.apache.airflow.sdk.Arg +import org.apache.airflow.sdk.BranchRef import org.apache.airflow.sdk.Bundle import org.apache.airflow.sdk.Client import org.apache.airflow.sdk.ConditionRef @@ -41,6 +42,8 @@ import org.apache.airflow.sdk.DagDef import org.apache.airflow.sdk.Deps import org.apache.airflow.sdk.Task import org.apache.airflow.sdk.TaskGroupRef +import org.apache.airflow.sdk.TaskId +import org.apache.airflow.sdk.TaskIdBranchTask import org.apache.airflow.sdk.TaskRef import org.apache.airflow.sdk.execution.serializeDag import org.apache.airflow.sdk.internal.Field @@ -64,6 +67,15 @@ class ConformanceCondition : ConditionTask { ) = true } +// Names its case by task ID: every conformance task runs the same class, so a +// branch that named a case by class could not tell two of them apart. +class ConformanceBranch : TaskIdBranchTask { + override fun choose( + context: Context, + client: Client, + ) = TaskId.of("") +} + fun main(args: Array) { require(args.size == 2) { "usage: SerializeJava " } val cases = YAMLMapper().readTree(File(args[0])).path("dags") @@ -90,20 +102,27 @@ private fun buildDag(case: JsonNode): DagDef { } val tasks = linkedMapOf>() - // A condition names tasks that may be declared after it, so the sides are wired once every task exists. - val conditions = mutableListOf>() + // A decider names tasks that may be declared after it, so its cases are wired once every task exists. + val deciders = mutableListOf>() 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 -> ConformanceTask::class.java + branch.has("cases") -> ConformanceBranch::class.java + else -> ConformanceCondition::class.java + } val ref = if (groupId.isEmpty()) { dag.task(localId, definition) } else { groups.getValue(groupId).task(localId, definition) } - if (!branch.isMissingNode) conditions += asCondition(ref) to branch + if (!branch.isMissingNode) { + deciders += (if (branch.has("cases")) BranchRef.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 +136,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 BranchRef -> branch.path("cases").forEach { decider.option(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 a1ca3f7553db2..15956625e5246 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 @@ -19,6 +19,8 @@ package org.apache.airflow.sdk.execution +import org.apache.airflow.sdk.BranchRef +import org.apache.airflow.sdk.BranchTask import org.apache.airflow.sdk.Bundle import org.apache.airflow.sdk.Client import org.apache.airflow.sdk.ConditionTask @@ -26,6 +28,8 @@ import org.apache.airflow.sdk.Context import org.apache.airflow.sdk.DagDef import org.apache.airflow.sdk.Task import org.apache.airflow.sdk.TaskDef +import org.apache.airflow.sdk.TaskId +import org.apache.airflow.sdk.TaskIdBranchTask import org.apache.airflow.sdk.execution.comm.BundleInfo import org.apache.airflow.sdk.execution.comm.DagRun import org.apache.airflow.sdk.execution.comm.RetryTask @@ -245,6 +249,59 @@ class TaskTest { Assertions.assertEquals(emptyList>(), transport.xComs) } + @Test + @DisplayName("Should push the case a branch chose and skip every other one") + fun shouldSkipTheCasesNotChosen() { + TestBranch.choice = HandleLongTask::class.java + val transport = RecordingTransport() + + val result = runTask(branchBundle(), 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 branch that chose a task it cannot run") + fun shouldFailBranchThatChoseAnUnknownCase() { + TestBranch.choice = SuccessTask::class.java + val transport = RecordingTransport() + + val result = runTask(branchBundle(), 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(), transport.skipped) + Assertions.assertEquals(emptyList>(), transport.xComs) + } + + @Test + @DisplayName("Should let a generated branch name its case by task ID") + fun shouldChooseByTaskId() { + val dag = DagDef("test_dag") + val long = dag.task("handle_long", HandleLongTask::class.java) + val short = dag.task("handle_short", HandleShortTask::class.java) + BranchRef.of(dag.task("pick", TestTaskIdBranch::class.java)).option(long).option(short) + val transport = RecordingTransport() + + runTask(Bundle(listOf(dag)), startupDetails(taskId = "pick"), Client(startupDetails("pick"), transport)) + + Assertions.assertEquals(listOf("handle_long"), transport.skipped) + Assertions.assertEquals("handle_short", transport.xComs.last().second) + } + + private fun branchBundle(): Bundle { + val dag = DagDef("test_dag") + val long = dag.task("handle_long", HandleLongTask::class.java) + val short = dag.task("handle_short", HandleShortTask::class.java) + dag.Branch("pick", TestBranch::class.java).option(long).option(short) + return Bundle(listOf(dag)) + } + private fun conditionBundle(withElse: Boolean): Bundle { val dag = DagDef("test_dag") val load = dag.task("load", SuccessTask::class.java) @@ -360,6 +417,40 @@ 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 TestBranch : BranchTask { + override fun choose( + context: Context, + client: Client, + ): Class = choice + + companion object { + var choice: Class = HandleLongTask::class.java + } + } + + /** The shape the annotation processor generates: a branch that names its case by ID. */ + class TestTaskIdBranch : TaskIdBranchTask { + override fun choose( + context: Context, + client: Client, + ): TaskId = TaskId.of("handle_short") + } + /** 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>() From 2594c90237ee054bedff29a25bd07e4686c1eb6e Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Thu, 8 Oct 2026 11:23:14 +0000 Subject: [PATCH 2/7] Java SDK: Name the multi-way branch Switch, as the other SDKs do --- .pre-commit-config.yaml | 2 +- java-sdk/README.md | 2 +- java-sdk/capabilities.yaml | 2 +- .../example/nativedag/AnnotationExample.java | 6 +- .../example/nativedag/InterfaceExample.java | 6 +- .../apache/airflow/sdk/BuilderProcessor.kt | 32 ++--- .../org/apache/airflow/sdk/BuilderTest.kt | 34 +++--- java-sdk/sdk/build.gradle.kts | 10 +- java-sdk/sdk/module.md | 2 +- .../kotlin/org/apache/airflow/sdk/Bundle.kt | 6 +- .../org/apache/airflow/sdk/Condition.kt | 2 +- .../kotlin/org/apache/airflow/sdk/DagDef.kt | 24 ++-- .../airflow/sdk/{Branch.kt => Switch.kt} | 113 +++++++++--------- .../org/apache/airflow/sdk/TaskGroupRef.kt | 22 ++-- .../org/apache/airflow/sdk/ConditionTest.kt | 2 +- .../sdk/{BranchTest.kt => SwitchTest.kt} | 74 ++++++------ .../airflow/sdk/conformance/SerializeJava.kt | 12 +- .../apache/airflow/sdk/execution/TaskTest.kt | 34 +++--- 18 files changed, 193 insertions(+), 192 deletions(-) rename java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/{Branch.kt => Switch.kt} (66%) rename java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/{BranchTest.kt => SwitchTest.kt} (68%) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 323828a34f094..5c04a604289b1 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -444,7 +444,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/(Branch|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 58e3bcad92393..9b9653a497474 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 and DagDef.Branch, or @Builder.If and @Builder.Branch | +| 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 b151d2562baa7..5a2ed1f5ae465 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 and DagDef.Branch, or @Builder.If and @Builder.Branch" + 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 d5f66ac64affc..c1ceef805b681 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 @@ -76,10 +76,10 @@ public void reportShort() { log.log(INFO, "Short report"); } - // A branch: it names the one case that runs, and every other case is + // A switch: it names the one case that runs, and every other case is // skipped. TaskIds is generated beside the builder, so a case that is not a // task of this Dag does not compile. - @Builder.Branch(id = "pick_report") + @Builder.Switch(id = "pick_report") public TaskId pickReport(long transformed) { return transformed > 100 ? AnnotationExampleBuilder.TaskIds.REPORT_LONG @@ -104,7 +104,7 @@ void depends() { var extracted = extract(); var transformed = transform(extracted, lit(1.5)); hasRows(transformed).Then(load(transformed)).Else(loadEmpty()); - pickReport(transformed).option(reportLong()).option(reportShort()); + 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 b6190406c5fe5..b8099931b7aae 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 @@ -85,9 +85,9 @@ public void execute(Context context, Client client) { } } - // A branch: it names the one case that runs by its class, and every other + // A switch: it names the one case that runs by its class, and every other // case is skipped. - public static class PickReport implements BranchTask { + public static class PickReport implements SwitchTask { @Override public Class choose(Context context, Client client) { return ((Number) client.getXCom("transform")).longValue() > 100 @@ -119,7 +119,7 @@ public static DagDef build() { // 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.Branch(PickReport.class).after(transform).option(reportLong).option(reportShort); + 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 bce9de83dd0d6..a15cdb9e0bf56 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 @@ -99,7 +99,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.Branch", + "org.apache.airflow.sdk.Builder.Switch", "org.apache.airflow.sdk.Builder.TaskGroup", "org.apache.airflow.sdk.Builder.TaskHandler", "org.apache.airflow.sdk.Builder.Deps", @@ -278,14 +278,14 @@ class BuilderProcessor : AbstractProcessor() { builderClass.addMethod(buildMethod.build()) declarations.forEach { builderClass.addType(buildTask(it)) } - // Only a branch needs to name a task in code, so a Dag without one keeps + // Only a switch needs to name a task in code, so a Dag without one keeps // the generated builder to the classes that run its tasks. - if (declarations.any { it.kind == TaskKind.BRANCH }) builderClass.addType(buildTaskIds(declarations)) + if (declarations.any { it.kind == TaskKind.SWITCH }) builderClass.addType(buildTaskIds(declarations)) return builderClass.build() } /** - * Generates the `TaskIds` holder a `@Builder.Branch` method names its case + * Generates the `TaskIds` holder a `@Builder.Switch` method names its case * with: one constant per task of the Dag, so a choice that is not a task of * this Dag does not compile. */ @@ -295,9 +295,9 @@ class BuilderProcessor : AbstractProcessor() { .classBuilder(TASK_IDS) .addModifiers(Modifier.PUBLIC, Modifier.STATIC, Modifier.FINAL) .addJavadoc( - "The task ids of this Dag, for a {@code @Builder.Branch} method to name its case with.\n\n" + - "

Every task of the Dag has one, so naming a task that is not a case of the branch\n" + - "compiles and fails when the branch runs.\n", + "The task ids of this Dag, for a {@code @Builder.Switch} method to name its case with.\n\n" + + "

Every task of the Dag has one, so naming a task that is not a case of the switch\n" + + "compiles and fails when the switch runs.\n", ) declarations.forEach { decl -> holder.addField( @@ -599,18 +599,18 @@ class BuilderProcessor : AbstractProcessor() { "Dag has both a task and a task group with ID '${group.fullId}'; rename one" } } - if (declarations.any { it.kind == TaskKind.BRANCH }) { + if (declarations.any { it.kind == TaskKind.SWITCH }) { val byConstant = mutableMapOf() declarations.forEach { decl -> val constant = constantName(decl.id) require(SourceVersion.isName(constant)) { "Task '${decl.id}' becomes the constant '$constant' of the generated TaskIds, which is not " + - "a Java name; give the task an id a branch can name it by, with ${decl.kind.spelling}(id = \"...\")" + "a Java name; give the task an id a switch can name it by, with ${decl.kind.spelling}(id = \"...\")" } byConstant.put(constant, decl)?.let { first -> throw IllegalArgumentException( "Tasks '${first.id}' and '${decl.id}' both become the constant " + - "'$constant' of the generated TaskIds; rename one so a branch can tell them apart", + "'$constant' of the generated TaskIds; rename one so a switch can tell them apart", ) } } @@ -775,9 +775,9 @@ class BuilderProcessor : AbstractProcessor() { "@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.BRANCH -> + TaskKind.SWITCH -> require(with(processingEnv) { isType(returns, TASK_ID_TYPE) }) { - "@Builder.Branch method '${method.simpleName}' returns $returns, but a branch returns a " + + "@Builder.Switch method '${method.simpleName}' returns $returns, but a switch returns a " + "TaskId: name the case it chose with a constant of the generated TaskIds" } } @@ -803,7 +803,7 @@ class BuilderProcessor : AbstractProcessor() { when (decl.kind) { TaskKind.TASK -> TypeName.VOID TaskKind.CONDITION -> TypeName.BOOLEAN - TaskKind.BRANCH -> TASK_ID_TYPE + TaskKind.SWITCH -> TASK_ID_TYPE }, ).addParameter(CONTEXT_TYPE, "context") .addParameter(CLIENT_TYPE, "client") @@ -1045,8 +1045,8 @@ 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 BRANCH_TASK_TYPE = ClassName.get(TaskIdBranchTask::class.java) -private val BRANCH_REF_TYPE = ClassName.get(BranchRef::class.java) +private val SWITCH_TASK_TYPE = ClassName.get(TaskIdSwitchTask::class.java) +private val SWITCH_REF_TYPE = ClassName.get(SwitchRef::class.java) private val TASK_ID_TYPE = ClassName.get(TaskId::class.java) /** Name of the generated holder of a Dag's task ids, which no task class may take. */ @@ -1073,7 +1073,7 @@ private enum class TaskKind( ) { 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), - BRANCH("org.apache.airflow.sdk.Builder.Branch", "@Builder.Branch", "choose", BRANCH_TASK_TYPE, BRANCH_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. */ 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 4193bc211b6ea..72ee6149f7360 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 @@ -2301,8 +2301,8 @@ class BuilderTest { } @Test - @DisplayName("generate a branch that names its case with a generated task-id constant") - fun generateBranchTask() { + @DisplayName("generate a switch that names its case with a generated task-id constant") + fun generateSwitchTask() { val compilation = compile( """ @@ -2313,7 +2313,7 @@ class BuilderTest { @Builder.Dag(id = "etl") public class TestExample { - @Builder.Branch(id = "pick_path") + @Builder.Switch(id = "pick_path") public TaskId pickPath() { return TestExampleBuilder.TaskIds.HANDLE_LONG; } @@ -2327,7 +2327,7 @@ class BuilderTest { @Builder.Deps static class Wiring implements TestExampleDeps { void depends() { - pickPath().option(handleLong()).option(handleShort()); + pickPath().Case(handleLong()).Case(handleShort()); } } } @@ -2350,7 +2350,7 @@ class BuilderTest { import org.apache.airflow.sdk.DagDef; import org.apache.airflow.sdk.Task; import org.apache.airflow.sdk.TaskId; - import org.apache.airflow.sdk.TaskIdBranchTask; + import org.apache.airflow.sdk.TaskIdSwitchTask; import org.apache.airflow.sdk.internal.DagSource; import org.apache.airflow.sdk.internal.Refs; @@ -2360,7 +2360,7 @@ class BuilderTest { return Refs.record(dag, List.of("pick_path", "handleLong", "handleShort"), List.of(), new TestExample.Wiring()::depends); } - public static final class PickPath implements TaskIdBranchTask { + public static final class PickPath implements TaskIdSwitchTask { @Override public TaskId choose(Context context, Client client) throws Exception { return new TestExample().pickPath(); @@ -2382,7 +2382,7 @@ class BuilderTest { } /** - * The task ids of this Dag, for a {@code @Builder.Branch} method to name its case with. + * The task ids of this Dag, for a {@code @Builder.Switch} method to name its case with. */ public static final class TaskIds { public static final TaskId PICK_PATH = TaskId.of("pick_path"); @@ -2402,15 +2402,15 @@ class BuilderTest { package org.apache.airflow.example; import java.lang.Void; - import org.apache.airflow.sdk.BranchRef; 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 BranchRef pickPath() { - return BranchRef.of(Refs.node("", new TaskDef("pick_path", TestExampleBuilder.PickPath.class))); + default SwitchRef pickPath() { + return SwitchRef.of(Refs.node("", new TaskDef("pick_path", TestExampleBuilder.PickPath.class))); } default TaskRef handleLong() { @@ -2436,7 +2436,7 @@ class BuilderTest { import org.apache.airflow.sdk.TaskId; @Builder.Dag public class TestExample { - @Builder.Branch + @Builder.Switch public TaskId pick() { return TestExampleBuilder.TaskIds.TASK_IDS; } @@ -2469,7 +2469,7 @@ class BuilderTest { import org.apache.airflow.sdk.TaskId; @Builder.Dag public class TestExample { - @Builder.Branch + @Builder.Switch public TaskId pick() { return null; } @@ -2492,8 +2492,8 @@ class BuilderTest { } @Test - @DisplayName("reject a branch that does not return a task id") - fun rejectNonTaskIdBranch() { + @DisplayName("reject a switch that does not return a task id") + fun rejectNonTaskIdSwitch() { val compilation = compile( """ @@ -2501,7 +2501,7 @@ class BuilderTest { import org.apache.airflow.sdk.Builder; @Builder.Dag public class TestExample { - @Builder.Branch + @Builder.Switch public String pickPath() { return "handleLong"; } @@ -2516,7 +2516,7 @@ class BuilderTest { assertThat(compilation).failed() assertThat(compilation).hadErrorContaining( - "@Builder.Branch method 'pickPath' returns java.lang.String, but a branch returns a TaskId", + "@Builder.Switch method 'pickPath' returns java.lang.String, but a switch returns a TaskId", ) } @@ -2531,7 +2531,7 @@ class BuilderTest { import org.apache.airflow.sdk.TaskId; @Builder.Dag public class TestExample { - @Builder.Branch + @Builder.Switch public TaskId pick() { return TestExampleBuilder.TaskIds.HANDLE_LONG; } diff --git a/java-sdk/sdk/build.gradle.kts b/java-sdk/sdk/build.gradle.kts index fe4044913e747..3c8d3edae93c3 100644 --- a/java-sdk/sdk/build.gradle.kts +++ b/java-sdk/sdk/build.gradle.kts @@ -620,22 +620,22 @@ abstract class GenerateDagDslTask : DefaultTask() { | * The method returns the [TaskId] of the task it chose, which the | * generated `Builder.TaskIds` declares a constant for, so a | * choice that is not a task of this Dag does not compile. Its - | * wiring-view method returns a [BranchRef], so the cases are listed + | * wiring-view method returns a [SwitchRef], so the cases are listed | * where the Dag is wired: | * | * ```java - | * @Builder.Branch(id = "pick_path") + | * @Builder.Switch(id = "pick_path") | * public TaskId pickPath(long rows) { | * return rows > 1000 ? EtlBuilder.TaskIds.HANDLE_LONG : EtlBuilder.TaskIds.HANDLE_SHORT; | * } | * | * @Builder.Deps | * static class Wiring implements EtlDeps { - | * void depends() { pickPath(extract()).option(handleLong()).option(handleShort()); } + | * void depends() { pickPath(extract()).Case(handleLong()).Case(handleShort()); } | * } | * ``` | * - | * A branch chooses exactly one case, so a task that runs after + | * 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`. | * @@ -644,7 +644,7 @@ abstract class GenerateDagDslTask : DefaultTask() { | */ | @Target(AnnotationTarget.FUNCTION) | @MustBeDocumented - | annotation class Branch( + | annotation class Switch( | /** Task ID. Empty derives it from the annotated function's name. */ | val id: String = "", |${attrLines(taskFields)} diff --git a/java-sdk/sdk/module.md b/java-sdk/sdk/module.md index 9a74418e3ab6a..f03b936d7b9b4 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 and DagDef.Branch, or @Builder.If and @Builder.Branch | +| 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 d34927d28b6ac..813b9fc0eb17a 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,7 @@ 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 `Branch`, or the dependencies contain + * runs is not declared with `If` or `Switch`, or the dependencies contain * a cycle. * @throws IllegalStateException if [Server.serve] has already been called. */ @@ -88,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(...) or Branch(...), 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}'" } @@ -213,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, BranchTask::class.java, TaskIdBranchTask::class.java) +private val DECIDER_TYPES = listOf(ConditionTask::class.java, SwitchTask::class.java, TaskIdSwitchTask::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 90150609c7b64..ea521f9202836 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): 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 ae5556ff4bd01..9f79be81a6c4a 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 @@ -207,40 +207,40 @@ class DagDef( * * The task's ID is the class's simple name with its first character * lowercased, so `PickPath.class` becomes `pickPath`. Use - * [Branch(id, definition)][Branch] to set it, and [BranchRef.config] for the + * [Switch(id, definition)][Switch] to set it, and [SwitchRef.config] for the * task's other settings: * * ```java - * dag.Branch(PickPath.class).option(handleLong).option(handleShort); + * dag.Switch(PickPath.class).Case(handleLong).Case(handleShort); * ``` * - * @param definition Class that implements [BranchTask]. Must have a public + * @param definition Class that implements [SwitchTask]. Must have a public * no-arg constructor. - * @return The branch, to list its cases on. + * @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 Branch(definition: Class): BranchRef = Branch(deriveTaskId(definition), definition) + fun Switch(definition: Class): 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 [BranchTask]. Must have a public + * @param definition Class that implements [SwitchTask]. Must have a public * no-arg constructor. - * @return The branch, to list its cases on. + * @return The switch, to list its cases on. * @throws IllegalArgumentException if a task with the same ID is already * registered. * - * @see Branch + * @see Switch */ @Suppress("ktlint:standard:function-naming") - fun Branch( + fun Switch( id: String, - definition: Class, - ): BranchRef = BranchRef.of(task(id, definition)) + definition: Class, + ): SwitchRef = SwitchRef.of(task(id, definition)) /** * Declares a task group of this Dag. @@ -416,7 +416,7 @@ class TaskDef( internal val upstreams = linkedSetOf() 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/Branch.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Switch.kt similarity index 66% rename from java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Branch.kt rename to java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Switch.kt index ce444e448d17d..87832ce05539b 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Branch.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Switch.kt @@ -24,7 +24,7 @@ import kotlin.Throws /** * The Airflow task ID of one task of a Dag. * - * A `@Builder.Branch` method names the task it chose with one of these. The + * A `@Builder.Switch` method names the task it chose with one of these. The * generated `Builder.TaskIds` holds a constant per task of the Dag, so * the choice is checked where it is written rather than when the task runs. * @@ -54,11 +54,11 @@ class TaskId private constructor( /** * A task that chooses one of several tasks to run; every other one is skipped. * - * Register one with [DagDef.Branch], then list what it can choose with - * [BranchRef.option]. The choice is the task's own class, so javac checks it: + * Register one with [DagDef.Switch], then list what it can choose with + * [SwitchRef.Case]. The choice is the task's own class, so javac checks it: * * ```java - * public class PickPath implements BranchTask { + * public class PickPath implements SwitchTask { * @Override * public Class choose(Context context, Client client) { * return ((Number) client.getXCom("extract")).longValue() > 1000 @@ -71,12 +71,12 @@ class TaskId private constructor( * The SDK runs [choose] and pushes the chosen task's ID as this task's return * value, so [execute] is never called. * - * Because a branch names a case by its class, no two of its cases may be + * Because a switch names a case by its class, no two of its cases may be * registered from the same class; registering the Dag reports that. * - * @see DagDef.Branch + * @see DagDef.Switch */ -interface BranchTask : Task { +interface SwitchTask : Task { /** * Chooses the one case that runs. * @@ -85,7 +85,7 @@ interface BranchTask : Task { * * @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 branch's cases. + * @return The class of one of this switch's cases. * @throws Exception on failure; the task instance is marked failed. */ @Throws(Exception::class) @@ -94,24 +94,24 @@ interface BranchTask : Task { client: Client, ): Class - /** Never called: the SDK runs a branch through [choose]. */ + /** Never called: the SDK runs a switch through [choose]. */ override fun execute( context: Context, client: Client, ): Unit = throw IllegalStateException( - "Branch '${javaClass.name}' runs through choose(), so execute() is never called", + "Switch '${javaClass.name}' runs through choose(), so execute() is never called", ) } /** * @suppress * - * What a generated `@Builder.Branch` task implements: a branch whose method + * What a generated `@Builder.Switch` task implements: a switch whose method * named the case with a [TaskId]. Public so generated code can implement it; - * a Dag written against [DagDef.Branch] implements [BranchTask] instead. + * a Dag written against [DagDef.Switch] implements [SwitchTask] instead. */ -interface TaskIdBranchTask : Task { +interface TaskIdSwitchTask : Task { @Throws(Exception::class) fun choose( context: Context, @@ -123,49 +123,49 @@ interface TaskIdBranchTask : Task { client: Client, ): Unit = throw IllegalStateException( - "Branch '${javaClass.name}' runs through choose(), so execute() is never called", + "Switch '${javaClass.name}' runs through choose(), so execute() is never called", ) } /** - * A branch registered with a Dag: list the tasks it can choose between. + * A switch registered with a Dag: list the tasks it can choose between. * * ```java - * dag.Branch(PickPath.class).option(handleLong).option(handleShort); + * dag.Switch(PickPath.class).Case(handleLong).Case(handleShort); * ``` * - * A case runs after the branch, so naming it records that edge, as - * [Deps.Flow.before] would. The branch chooses exactly one case and skips + * 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.Branch + * @see DagDef.Switch */ -class BranchRef private constructor( +class SwitchRef private constructor( private val ref: TaskRef<*>, - private val cases: BranchDef, + private val cases: SwitchDef, ) : Deps.Flow { companion object { /** * @suppress * - * Marks a registered task as a branch. Public so a generated wiring view - * can call it; user code reaches a branch through [DagDef.Branch]. + * Marks a registered task as a switch. Public so a generated wiring view + * can call it; user code reaches a switch through [DagDef.Switch]. */ @JvmStatic - fun of(ref: TaskRef<*>): BranchRef { + fun of(ref: TaskRef<*>): SwitchRef { require(ref.def.decider == null) { "Task '${ref.def.id}' already decides what to skip; declare it once" } val definition = ref.def.definition require( - BranchTask::class.java.isAssignableFrom(definition) || - TaskIdBranchTask::class.java.isAssignableFrom(definition), + SwitchTask::class.java.isAssignableFrom(definition) || + TaskIdSwitchTask::class.java.isAssignableFrom(definition), ) { - "Task '${ref.def.id}' runs '${definition.name}', which chooses nothing; a branch runs a " + - "BranchTask" + "Task '${ref.def.id}' runs '${definition.name}', which chooses nothing; a switch runs a " + + "SwitchTask" } - return BranchRef(ref, BranchDef().also { ref.def.decider = it }) + return SwitchRef(ref, SwitchDef().also { ref.def.decider = it }) } } @@ -173,19 +173,20 @@ class BranchRef private constructor( val id: String get() = ref.def.id /** - * Adds a task this branch can choose. Every other case is skipped when the - * branch chooses this one. + * 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 branch, for chaining. + * @return This switch, for chaining. * @throws IllegalArgumentException if the task belongs to another Dag or is - * already a case of this branch. + * already a case of this switch. */ - fun option(task: TaskRef<*>): BranchRef { + @Suppress("ktlint:standard:function-naming") + fun Case(task: TaskRef<*>): SwitchRef { cases.add(ref.def, task.def) - // The case runs after the branch, which is what puts the branch in the + // 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 branch. + // runs through a switch. ref.before(task) return this } @@ -195,33 +196,33 @@ class BranchRef private constructor( * * @param key Airflow task setting name. * @param value Value matching the key's schema type. - * @return This branch, for chaining. + * @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?, - ): BranchRef { + ): SwitchRef { ref.config(key, value) return this } override fun nodes(): List = ref.nodes() - override fun before(vararg next: Deps.Flow): BranchRef { + override fun before(vararg next: Deps.Flow): SwitchRef { ref.before(*next) return this } - override fun after(vararg previous: Deps.Flow): BranchRef { + override fun after(vararg previous: Deps.Flow): SwitchRef { ref.after(*previous) return this } } -/** The cases of one branch, as [BranchRef.option] listed them. */ -internal class BranchDef : DeciderDef { +/** The cases of one switch, as [SwitchRef.Case] listed them. */ +internal class SwitchDef : DeciderDef { private val options = mutableListOf() override val cases: List get() = options @@ -230,19 +231,19 @@ internal class BranchDef : DeciderDef { decider: TaskDef, case: TaskDef, ) { - requireUnregistered(decider, "Branch") - requireSameDag(decider, case, "Branch") + requireUnregistered(decider, "Switch") + requireSameDag(decider, case, "Switch") require(options.none { it === case }) { - "Branch '${decider.id}' already chooses between '${case.id}' and others; name each case once" + "Switch '${decider.id}' already chooses between '${case.id}' and others; name each case once" } - // A BranchTask names its case by class, so two cases sharing one would be + // A SwitchTask names its case by class, so two cases sharing one would be // indistinguishable. Checked here rather than at registration so the error - // points at the second option(...) call. - if (BranchTask::class.java.isAssignableFrom(decider.definition)) { + // points at the second Case(...) call. + if (SwitchTask::class.java.isAssignableFrom(decider.definition)) { options.firstOrNull { it.definition == case.definition }?.let { first -> throw IllegalArgumentException( - "Branch '${decider.id}' cannot choose between '${first.id}' and '${case.id}': both run " + - "'${case.definition.name}', and a branch names its case by class", + "Switch '${decider.id}' cannot choose between '${first.id}' and '${case.id}': both run " + + "'${case.definition.name}', and a switch names its case by class", ) } } @@ -250,7 +251,7 @@ internal class BranchDef : DeciderDef { } override fun describe(decider: TaskDef): String? = - if (options.isEmpty()) "Branch '${decider.id}' has no task to choose between; call option(...)" else null + if (options.isEmpty()) "Switch '${decider.id}' has no task to choose between; call Case(...)" else null override fun decide( decider: TaskDef, @@ -260,10 +261,10 @@ internal class BranchDef : DeciderDef { ): Decision { val chosen = when (instance) { - is BranchTask -> byClass(decider, instance.choose(context, client)) - is TaskIdBranchTask -> byId(decider, instance.choose(context, client).value) + is SwitchTask -> byClass(decider, instance.choose(context, client)) + is TaskIdSwitchTask -> byId(decider, instance.choose(context, client).value) else -> throw IllegalStateException( - "Task '${decider.id}' is a branch, but '${instance.javaClass.name}' is not a BranchTask", + "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 }) @@ -275,7 +276,7 @@ internal class BranchDef : DeciderDef { ): TaskDef = options.firstOrNull { it.definition == definition } ?: throw IllegalArgumentException( - "Branch '${decider.id}' chose ${definition?.name ?: "nothing"}, which is not one of its cases: ${describeCases()}", + "Switch '${decider.id}' chose ${definition?.name ?: "nothing"}, which is not one of its cases: ${describeCases()}", ) private fun byId( @@ -284,7 +285,7 @@ internal class BranchDef : DeciderDef { ): TaskDef = options.firstOrNull { it.id == taskId } ?: throw IllegalArgumentException( - "Branch '${decider.id}' chose '$taskId', which is not one of its cases: ${describeCases()}", + "Switch '${decider.id}' chose '$taskId', 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 3fb8c3e7d5063..60d43f84fcd89 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 @@ -107,34 +107,34 @@ class TaskGroupRef internal constructor( ): ConditionRef = ConditionRef.of(task(id, definition)) /** - * Declares a branch in this group, as [DagDef.Branch] does for the Dag. + * Declares a switch in this group, as [DagDef.Switch] does for the Dag. * - * @param definition Class that implements [BranchTask]. Must have a public + * @param definition Class that implements [SwitchTask]. Must have a public * no-arg constructor. - * @return The branch, to list its cases on. + * @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 Branch(definition: Class): BranchRef = Branch(deriveTaskId(definition), definition) + fun Switch(definition: Class): SwitchRef = Switch(deriveTaskId(definition), definition) /** - * Declares a branch in this group under the task ID `.`. + * Declares a switch in this group under the task ID `.`. * * @param id Task ID within this group. - * @param definition Class that implements [BranchTask]. Must have a public + * @param definition Class that implements [SwitchTask]. Must have a public * no-arg constructor. - * @return The branch, to list its cases on. + * @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 Branch + * @see Switch */ @Suppress("ktlint:standard:function-naming") - fun Branch( + fun Switch( id: String, - definition: Class, - ): BranchRef = BranchRef.of(task(id, definition)) + definition: Class, + ): SwitchRef = SwitchRef.of(task(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 eca33565c5b35..0fb138778371f 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(...) or Branch(...), 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/BranchTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/SwitchTest.kt similarity index 68% rename from java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/BranchTest.kt rename to java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/SwitchTest.kt index e37d7a38d3719..dbbd97caaa7b1 100644 --- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/BranchTest.kt +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/SwitchTest.kt @@ -24,7 +24,7 @@ import org.junit.jupiter.api.Assertions.assertThrows import org.junit.jupiter.api.DisplayName import org.junit.jupiter.api.Test -/** A case of the branches under test; a second class so a branch can tell two cases apart. */ +/** 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, @@ -32,8 +32,8 @@ internal class HandleLong : Task { ) = Unit } -/** A branch that chooses whatever [choice] holds. */ -class PickPath : BranchTask { +/** A switch that chooses whatever [choice] holds. */ +class PickPath : SwitchTask { override fun choose( context: Context, client: Client, @@ -44,51 +44,51 @@ class PickPath : BranchTask { } } -internal class BranchTest { - private fun dagWithCases(): Triple, TaskRef>> { +internal class SwitchTest { + private fun dagWithCases(): Triple, TaskRef>> { val dag = DagDef("d") val handleLong = dag.task("handle_long", HandleLong::class.java) val handleShort = dag.task("handle_short", NoopTask::class.java) - return Triple(dag, dag.Branch(PickPath::class.java), handleLong to handleShort) + return Triple(dag, dag.Switch(PickPath::class.java), handleLong to handleShort) } @Test - @DisplayName("Should take the branch's task ID from its class when none is given") + @DisplayName("Should take the switch's task ID from its class when none is given") fun shouldDeriveTaskIdFromClass() { - val (dag, branch, cases) = dagWithCases() - branch.option(cases.first) + val (dag, switch, cases) = dagWithCases() + switch.Case(cases.first) - assertEquals("pickPath", branch.id) + 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 (_, branch, cases) = dagWithCases() - branch.option(cases.first) + val (_, switch, cases) = dagWithCases() + switch.Case(cases.first) - val error = assertThrows(IllegalArgumentException::class.java) { branch.option(cases.first) } + val error = assertThrows(IllegalArgumentException::class.java) { switch.Case(cases.first) } assertEquals( - "Branch 'pickPath' already chooses between 'handle_long' and others; name each case once", + "Switch 'pickPath' already chooses between 'handle_long' and others; name each case once", error.message, ) } @Test - @DisplayName("Should reject two cases a class-named branch could not tell apart") + @DisplayName("Should reject two cases a class-named switch could not tell apart") fun shouldRejectCasesSharingAClass() { val dag = DagDef("d") val first = dag.task("first", HandleLong::class.java) val second = dag.task("second", HandleLong::class.java) - val branch = dag.Branch(PickPath::class.java).option(first) + val switch = dag.Switch(PickPath::class.java).Case(first) - val error = assertThrows(IllegalArgumentException::class.java) { branch.option(second) } + val error = assertThrows(IllegalArgumentException::class.java) { switch.Case(second) } assertEquals( - "Branch 'pickPath' cannot choose between 'first' and 'second': both run " + - "'org.apache.airflow.sdk.HandleLong', and a branch names its case by class", + "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, ) } @@ -96,44 +96,44 @@ internal class BranchTest { @Test @DisplayName("Should reject a case declared in another Dag") fun shouldRejectCaseFromAnotherDag() { - val (_, branch, _) = dagWithCases() + val (_, switch, _) = dagWithCases() val other = DagDef("other").task("handle_long", HandleLong::class.java) - val error = assertThrows(IllegalArgumentException::class.java) { branch.option(other) } + val error = assertThrows(IllegalArgumentException::class.java) { switch.Case(other) } assertEquals( - "Branch 'pickPath' of Dag 'd' cannot run task 'handle_long' of Dag 'other'; " + + "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 branch with nothing to choose between") - fun shouldRejectBranchWithoutCases() { + @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( - "Branch 'pickPath' has no task to choose between; call option(...), in Dag 'd'", + "Switch 'pickPath' has no task to choose between; call Case(...), in Dag 'd'", error.message, ) } @Test - @DisplayName("Should reject a branch over a task that chooses nothing") - fun shouldRejectBranchOverAPlainTask() { + @DisplayName("Should reject a switch over a task that chooses nothing") + fun shouldRejectSwitchOverAPlainTask() { val dag = DagDef("d") val error = assertThrows(IllegalArgumentException::class.java) { - BranchRef.of(dag.task("plain", NoopTask::class.java)) + SwitchRef.of(dag.task("plain", NoopTask::class.java)) } assertEquals( "Task 'plain' runs 'org.apache.airflow.sdk.NoopTask', which chooses nothing; " + - "a branch runs a BranchTask", + "a switch runs a SwitchTask", error.message, ) } @@ -141,28 +141,28 @@ internal class BranchTest { @Test @DisplayName("Should reject naming a case after the Dag was registered") fun shouldRejectCaseAfterRegistration() { - val (dag, branch, cases) = dagWithCases() - branch.option(cases.first) + val (dag, switch, cases) = dagWithCases() + switch.Case(cases.first) Bundle().register(dag) - val error = assertThrows(IllegalArgumentException::class.java) { branch.option(cases.second) } + val error = assertThrows(IllegalArgumentException::class.java) { switch.Case(cases.second) } assertEquals( - "Branch 'pickPath' of Dag 'd' is already registered; name every case before the Dag is " + + "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 branch inside a task group under the group's prefix") - fun shouldDeclareBranchInGroup() { + @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("long", HandleLong::class.java) - val branch = reports.Branch("pick", PickPath::class.java).option(long) + val switch = reports.Switch("pick", PickPath::class.java).Case(long) - assertEquals("reports.pick", branch.id) + assertEquals("reports.pick", switch.id) assertEquals(listOf("reports.long", "reports.pick"), dag.tasks.keys.toList()) } 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 9e12006485a5b..91adeb027c843 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 @@ -32,7 +32,6 @@ import com.fasterxml.jackson.databind.JsonNode import com.fasterxml.jackson.databind.ObjectMapper import com.fasterxml.jackson.dataformat.yaml.YAMLMapper import org.apache.airflow.sdk.Arg -import org.apache.airflow.sdk.BranchRef import org.apache.airflow.sdk.Bundle import org.apache.airflow.sdk.Client import org.apache.airflow.sdk.ConditionRef @@ -40,10 +39,11 @@ 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.Task import org.apache.airflow.sdk.TaskGroupRef import org.apache.airflow.sdk.TaskId -import org.apache.airflow.sdk.TaskIdBranchTask +import org.apache.airflow.sdk.TaskIdSwitchTask import org.apache.airflow.sdk.TaskRef import org.apache.airflow.sdk.execution.serializeDag import org.apache.airflow.sdk.internal.Field @@ -69,7 +69,7 @@ class ConformanceCondition : ConditionTask { // Names its case by task ID: every conformance task runs the same class, so a // branch that named a case by class could not tell two of them apart. -class ConformanceBranch : TaskIdBranchTask { +class ConformanceSwitch : TaskIdSwitchTask { override fun choose( context: Context, client: Client, @@ -111,7 +111,7 @@ private fun buildDag(case: JsonNode): DagDef { val definition = when { branch.isMissingNode -> ConformanceTask::class.java - branch.has("cases") -> ConformanceBranch::class.java + branch.has("cases") -> ConformanceSwitch::class.java else -> ConformanceCondition::class.java } val ref = @@ -121,7 +121,7 @@ private fun buildDag(case: JsonNode): DagDef { groups.getValue(groupId).task(localId, definition) } if (!branch.isMissingNode) { - deciders += (if (branch.has("cases")) BranchRef.of(ref) else asCondition(ref)) to branch + 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 @@ -138,7 +138,7 @@ private fun buildDag(case: JsonNode): DagDef { deciders.forEach { (decider, branch) -> when (decider) { - is BranchRef -> branch.path("cases").forEach { decider.option(tasks.getValue(it.asText())) } + 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())) } 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 15956625e5246..bc8f9d721363d 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 @@ -19,17 +19,17 @@ package org.apache.airflow.sdk.execution -import org.apache.airflow.sdk.BranchRef -import org.apache.airflow.sdk.BranchTask import org.apache.airflow.sdk.Bundle 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.SwitchRef +import org.apache.airflow.sdk.SwitchTask import org.apache.airflow.sdk.Task import org.apache.airflow.sdk.TaskDef import org.apache.airflow.sdk.TaskId -import org.apache.airflow.sdk.TaskIdBranchTask +import org.apache.airflow.sdk.TaskIdSwitchTask import org.apache.airflow.sdk.execution.comm.BundleInfo import org.apache.airflow.sdk.execution.comm.DagRun import org.apache.airflow.sdk.execution.comm.RetryTask @@ -250,12 +250,12 @@ class TaskTest { } @Test - @DisplayName("Should push the case a branch chose and skip every other one") + @DisplayName("Should push the case a switch chose and skip every other one") fun shouldSkipTheCasesNotChosen() { - TestBranch.choice = HandleLongTask::class.java + TestSwitch.choice = HandleLongTask::class.java val transport = RecordingTransport() - val result = runTask(branchBundle(), startupDetails(taskId = "pick"), Client(startupDetails("pick"), transport)) + val result = runTask(switchBundle(), startupDetails(taskId = "pick"), Client(startupDetails("pick"), transport)) Assertions.assertInstanceOf(SucceedTask::class.java, result) Assertions.assertEquals(listOf("handle_short"), transport.skipped) @@ -266,12 +266,12 @@ class TaskTest { } @Test - @DisplayName("Should fail a branch that chose a task it cannot run") - fun shouldFailBranchThatChoseAnUnknownCase() { - TestBranch.choice = SuccessTask::class.java + @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(branchBundle(), startupDetails(taskId = "pick"), Client(startupDetails("pick"), transport)) + 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) @@ -280,12 +280,12 @@ class TaskTest { } @Test - @DisplayName("Should let a generated branch name its case by task ID") + @DisplayName("Should let a generated switch name its case by task ID") fun shouldChooseByTaskId() { val dag = DagDef("test_dag") val long = dag.task("handle_long", HandleLongTask::class.java) val short = dag.task("handle_short", HandleShortTask::class.java) - BranchRef.of(dag.task("pick", TestTaskIdBranch::class.java)).option(long).option(short) + SwitchRef.of(dag.task("pick", TestTaskIdSwitch::class.java)).Case(long).Case(short) val transport = RecordingTransport() runTask(Bundle(listOf(dag)), startupDetails(taskId = "pick"), Client(startupDetails("pick"), transport)) @@ -294,11 +294,11 @@ class TaskTest { Assertions.assertEquals("handle_short", transport.xComs.last().second) } - private fun branchBundle(): Bundle { + private fun switchBundle(): Bundle { val dag = DagDef("test_dag") val long = dag.task("handle_long", HandleLongTask::class.java) val short = dag.task("handle_short", HandleShortTask::class.java) - dag.Branch("pick", TestBranch::class.java).option(long).option(short) + dag.Switch("pick", TestSwitch::class.java).Case(long).Case(short) return Bundle(listOf(dag)) } @@ -432,7 +432,7 @@ class TaskTest { } /** Chooses whatever [choice] holds, so a test can drive it to a case or past one. */ - class TestBranch : BranchTask { + class TestSwitch : SwitchTask { override fun choose( context: Context, client: Client, @@ -443,8 +443,8 @@ class TaskTest { } } - /** The shape the annotation processor generates: a branch that names its case by ID. */ - class TestTaskIdBranch : TaskIdBranchTask { + /** The shape the annotation processor generates: a switch that names its case by ID. */ + class TestTaskIdSwitch : TaskIdSwitchTask { override fun choose( context: Context, client: Client, From 141e5cedb9cb4ea2de955a8dac40992a81007e1f Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Thu, 8 Oct 2026 11:25:30 +0000 Subject: [PATCH 3/7] Java SDK: Name a switch case by its task class in both styles --- .../example/nativedag/AnnotationExample.java | 11 +- .../apache/airflow/sdk/BuilderProcessor.kt | 92 +++----------- .../org/apache/airflow/sdk/BuilderTest.kt | 116 +++--------------- java-sdk/sdk/build.gradle.kts | 12 +- .../kotlin/org/apache/airflow/sdk/Bundle.kt | 2 +- .../kotlin/org/apache/airflow/sdk/Switch.kt | 89 ++------------ .../org/apache/airflow/sdk/SwitchTest.kt | 7 -- .../airflow/sdk/conformance/SerializeJava.kt | 115 +++++++++++++++-- .../apache/airflow/sdk/execution/TaskTest.kt | 26 ---- 9 files changed, 163 insertions(+), 307 deletions(-) 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 c1ceef805b681..16cdc555452a9 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 @@ -76,14 +76,13 @@ public void reportShort() { log.log(INFO, "Short report"); } - // A switch: it names the one case that runs, and every other case is - // skipped. TaskIds is generated beside the builder, so a case that is not a - // task of this Dag does not compile. + // 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 TaskId pickReport(long transformed) { + public Class pickReport(long transformed) { return transformed > 100 - ? AnnotationExampleBuilder.TaskIds.REPORT_LONG - : AnnotationExampleBuilder.TaskIds.REPORT_SHORT; + ? AnnotationExampleBuilder.ReportLong.class + : AnnotationExampleBuilder.ReportShort.class; } // A task group: everything it declares is prefixed with its id, so this is 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 a15cdb9e0bf56..97b810358a0d2 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 @@ -23,7 +23,6 @@ package org.apache.airflow.sdk import com.squareup.javapoet.ClassName import com.squareup.javapoet.CodeBlock -import com.squareup.javapoet.FieldSpec import com.squareup.javapoet.JavaFile import com.squareup.javapoet.MethodSpec import com.squareup.javapoet.ParameterizedTypeName @@ -58,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 @@ -278,38 +279,9 @@ class BuilderProcessor : AbstractProcessor() { builderClass.addMethod(buildMethod.build()) declarations.forEach { builderClass.addType(buildTask(it)) } - // Only a switch needs to name a task in code, so a Dag without one keeps - // the generated builder to the classes that run its tasks. - if (declarations.any { it.kind == TaskKind.SWITCH }) builderClass.addType(buildTaskIds(declarations)) return builderClass.build() } - /** - * Generates the `TaskIds` holder a `@Builder.Switch` method names its case - * with: one constant per task of the Dag, so a choice that is not a task of - * this Dag does not compile. - */ - private fun buildTaskIds(declarations: List): TypeSpec { - val holder = - TypeSpec - .classBuilder(TASK_IDS) - .addModifiers(Modifier.PUBLIC, Modifier.STATIC, Modifier.FINAL) - .addJavadoc( - "The task ids of this Dag, for a {@code @Builder.Switch} method to name its case with.\n\n" + - "

Every task of the Dag has one, so naming a task that is not a case of the switch\n" + - "compiles and fails when the switch runs.\n", - ) - declarations.forEach { decl -> - holder.addField( - FieldSpec - .builder(TASK_ID_TYPE, constantName(decl.id), Modifier.PUBLIC, Modifier.STATIC, Modifier.FINAL) - .initializer($$"$T.of($S)", TASK_ID_TYPE, decl.id) - .build(), - ) - } - return holder.build() - } - /** * Generates the Dag's wiring view: one default method per task, with the * injected arguments stripped, each data argument lifted to [Arg] and the @@ -599,29 +571,6 @@ class BuilderProcessor : AbstractProcessor() { "Dag has both a task and a task group with ID '${group.fullId}'; rename one" } } - if (declarations.any { it.kind == TaskKind.SWITCH }) { - val byConstant = mutableMapOf() - declarations.forEach { decl -> - val constant = constantName(decl.id) - require(SourceVersion.isName(constant)) { - "Task '${decl.id}' becomes the constant '$constant' of the generated TaskIds, which is not " + - "a Java name; give the task an id a switch can name it by, with ${decl.kind.spelling}(id = \"...\")" - } - byConstant.put(constant, decl)?.let { first -> - throw IllegalArgumentException( - "Tasks '${first.id}' and '${decl.id}' both become the constant " + - "'$constant' of the generated TaskIds; rename one so a switch can tell them apart", - ) - } - } - declarations.firstOrNull { it.className == TASK_IDS }?.let { decl -> - throw IllegalArgumentException( - "Task method '${decl.method.simpleName}' generates the class '$TASK_IDS', which is the " + - "holder of this Dag's task ids; rename the method and keep its task id with " + - "${decl.kind.spelling}(id = \"${decl.id.substringAfterLast('.')}\")", - ) - } - } val byClassName = mutableMapOf() declarations.forEach { decl -> byClassName.put(decl.className, decl)?.let { first -> @@ -776,13 +725,23 @@ class BuilderProcessor : AbstractProcessor() { "true runs the task named by Then, false the one named by Else" } TaskKind.SWITCH -> - require(with(processingEnv) { isType(returns, TASK_ID_TYPE) }) { - "@Builder.Switch method '${method.simpleName}' returns $returns, but a switch returns a " + - "TaskId: name the case it chose with a constant of the generated TaskIds" + require(returnsTaskClass(returns)) { + "@Builder.Switch method '${method.simpleName}' returns $returns, but a switch returns the " + + "class of the task it chose, as Class" } } } + /** Whether [type] is `Class` or `Class` 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, @@ -803,7 +762,7 @@ class BuilderProcessor : AbstractProcessor() { when (decl.kind) { TaskKind.TASK -> TypeName.VOID TaskKind.CONDITION -> TypeName.BOOLEAN - TaskKind.SWITCH -> TASK_ID_TYPE + TaskKind.SWITCH -> TASK_CLASS_TYPE }, ).addParameter(CONTEXT_TYPE, "context") .addParameter(CLIENT_TYPE, "client") @@ -1045,12 +1004,10 @@ 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(TaskIdSwitchTask::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_ID_TYPE = ClassName.get(TaskId::class.java) - -/** Name of the generated holder of a Dag's task ids, which no task class may take. */ -private const val TASK_IDS = "TaskIds" +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") @@ -1109,17 +1066,6 @@ private val RESERVED_GROUP_VIEW_NAMES: Set = private val DAG_STRUCTURAL_ATTRIBUTES = setOf("id", "to") private val TASK_STRUCTURAL_ATTRIBUTES = setOf("id") -/** - * The `TaskIds` constant for a task id: its words in upper case, joined by - * underscores, so `handleLong` and `checks.audit` become `HANDLE_LONG` and - * `CHECKS_AUDIT`. - */ -private fun constantName(id: String): String = - id - .replace(Regex("([a-z0-9])([A-Z])"), "$1_$2") - .replace(Regex("[^A-Za-z0-9]+"), "_") - .uppercase() - private fun TypeName.boxIfPossible(): TypeName = if (this == TypeName.VOID || isPrimitive) box() else this private fun ProcessingEnvironment.isType( 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 72ee6149f7360..e0c60a3fc8289 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 @@ -2301,7 +2301,7 @@ class BuilderTest { } @Test - @DisplayName("generate a switch that names its case with a generated task-id constant") + @DisplayName("generate a switch that names its case by the class generated for it") fun generateSwitchTask() { val compilation = compile( @@ -2309,13 +2309,13 @@ class BuilderTest { package org.apache.airflow.example; import org.apache.airflow.sdk.Builder; - import org.apache.airflow.sdk.TaskId; + import org.apache.airflow.sdk.Task; @Builder.Dag(id = "etl") public class TestExample { @Builder.Switch(id = "pick_path") - public TaskId pickPath() { - return TestExampleBuilder.TaskIds.HANDLE_LONG; + public Class pickPath() { + return TestExampleBuilder.HandleLong.class; } @Builder.Task @@ -2342,15 +2342,15 @@ class BuilderTest { """ 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.TaskId; - import org.apache.airflow.sdk.TaskIdSwitchTask; import org.apache.airflow.sdk.internal.DagSource; import org.apache.airflow.sdk.internal.Refs; @@ -2360,9 +2360,9 @@ class BuilderTest { return Refs.record(dag, List.of("pick_path", "handleLong", "handleShort"), List.of(), new TestExample.Wiring()::depends); } - public static final class PickPath implements TaskIdSwitchTask { + public static final class PickPath implements SwitchTask { @Override - public TaskId choose(Context context, Client client) throws Exception { + public Class choose(Context context, Client client) throws Exception { return new TestExample().pickPath(); } } @@ -2380,17 +2380,6 @@ class BuilderTest { new TestExample().handleShort(); } } - - /** - * The task ids of this Dag, for a {@code @Builder.Switch} method to name its case with. - */ - public static final class TaskIds { - public static final TaskId PICK_PATH = TaskId.of("pick_path"); - - public static final TaskId HANDLE_LONG = TaskId.of("handleLong"); - - public static final TaskId HANDLE_SHORT = TaskId.of("handleShort"); - } } """, ) @@ -2426,74 +2415,8 @@ class BuilderTest { } @Test - @DisplayName("reject a task method whose generated class would take the TaskIds name") - fun rejectTaskNamedTaskIds() { - val compilation = - compile( - """ - package org.apache.airflow.example; - import org.apache.airflow.sdk.Builder; - import org.apache.airflow.sdk.TaskId; - @Builder.Dag - public class TestExample { - @Builder.Switch - public TaskId pick() { - return TestExampleBuilder.TaskIds.TASK_IDS; - } - - @Builder.Task - public void taskIds() {} - - @Builder.Deps - static class Wiring implements TestExampleDeps { - void depends() {} - } - } - """, - ) - - assertThat(compilation).failed() - assertThat(compilation).hadErrorContaining( - "Task method 'taskIds' generates the class 'TaskIds', which is the holder of this Dag's task ids", - ) - } - - @Test - @DisplayName("reject a task id whose generated constant is not a Java name") - fun rejectTaskIdWithoutAJavaName() { - val compilation = - compile( - """ - package org.apache.airflow.example; - import org.apache.airflow.sdk.Builder; - import org.apache.airflow.sdk.TaskId; - @Builder.Dag - public class TestExample { - @Builder.Switch - public TaskId pick() { - return null; - } - - @Builder.Task(id = "2nd_pass") - public void second() {} - - @Builder.Deps - static class Wiring implements TestExampleDeps { - void depends() {} - } - } - """, - ) - - assertThat(compilation).failed() - assertThat(compilation).hadErrorContaining( - "Task '2nd_pass' becomes the constant '2ND_PASS' of the generated TaskIds, which is not a Java name", - ) - } - - @Test - @DisplayName("reject a switch that does not return a task id") - fun rejectNonTaskIdSwitch() { + @DisplayName("reject a switch that does not return the class of a task") + fun rejectNonTaskClassSwitch() { val compilation = compile( """ @@ -2516,32 +2439,25 @@ class BuilderTest { assertThat(compilation).failed() assertThat(compilation).hadErrorContaining( - "@Builder.Switch method 'pickPath' returns java.lang.String, but a switch returns a TaskId", + "@Builder.Switch method 'pickPath' returns java.lang.String, but a switch returns the class of the task it chose", ) } @Test - @DisplayName("reject two tasks whose generated task-id constants would collide") - fun rejectCollidingTaskIdConstants() { + @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; - import org.apache.airflow.sdk.TaskId; @Builder.Dag public class TestExample { @Builder.Switch - public TaskId pick() { - return TestExampleBuilder.TaskIds.HANDLE_LONG; + public Class pickPath() { + return String.class; } - @Builder.Task - public void handleLong() {} - - @Builder.Task(id = "handle_long") - public void other() {} - @Builder.Deps static class Wiring implements TestExampleDeps { void depends() {} @@ -2552,7 +2468,7 @@ class BuilderTest { assertThat(compilation).failed() assertThat(compilation).hadErrorContaining( - "Tasks 'handleLong' and 'handle_long' both become the constant 'HANDLE_LONG' of the generated TaskIds", + "@Builder.Switch method 'pickPath' returns java.lang.Class, but a switch returns the class of the task it chose", ) } diff --git a/java-sdk/sdk/build.gradle.kts b/java-sdk/sdk/build.gradle.kts index 3c8d3edae93c3..e30facde61e40 100644 --- a/java-sdk/sdk/build.gradle.kts +++ b/java-sdk/sdk/build.gradle.kts @@ -617,16 +617,14 @@ 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 [TaskId] of the task it chose, which the - | * generated `Builder.TaskIds` declares a constant for, so a - | * choice that is not a task of this Dag does not compile. Its - | * wiring-view method returns a [SwitchRef], so the cases are listed - | * where the Dag is wired: + | * The method returns the class the Dag's builder generated for the task it + | * chose, as `Class`. Its wiring-view method returns a + | * [SwitchRef], so the cases are listed where the Dag is wired: | * | * ```java | * @Builder.Switch(id = "pick_path") - | * public TaskId pickPath(long rows) { - | * return rows > 1000 ? EtlBuilder.TaskIds.HANDLE_LONG : EtlBuilder.TaskIds.HANDLE_SHORT; + | * public Class pickPath(long rows) { + | * return rows > 1000 ? EtlBuilder.HandleLong.class : EtlBuilder.HandleShort.class; | * } | * | * @Builder.Deps 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 813b9fc0eb17a..5e817c05e9316 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 @@ -213,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, SwitchTask::class.java, TaskIdSwitchTask::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/Switch.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Switch.kt index 87832ce05539b..e2cfa28d3d7a9 100644 --- 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 @@ -21,36 +21,6 @@ package org.apache.airflow.sdk import kotlin.Throws -/** - * The Airflow task ID of one task of a Dag. - * - * A `@Builder.Switch` method names the task it chose with one of these. The - * generated `Builder.TaskIds` holds a constant per task of the Dag, so - * the choice is checked where it is written rather than when the task runs. - * - * @property value The task ID. - */ -class TaskId private constructor( - val value: String, -) { - companion object { - /** - * @suppress - * - * The ID [value] names. Public so a generated `TaskIds` holder can declare - * its constants; user code names a task through those constants. - */ - @JvmStatic - fun of(value: String): TaskId = TaskId(value) - } - - override fun equals(other: Any?): Boolean = other is TaskId && other.value == value - - override fun hashCode(): Int = value.hashCode() - - override fun toString(): String = value -} - /** * A task that chooses one of several tasks to run; every other one is skipped. * @@ -71,8 +41,8 @@ class TaskId private constructor( * The SDK runs [choose] and pushes the chosen task's ID as this task's return * value, so [execute] is never called. * - * Because a switch names a case by its class, no two of its cases may be - * registered from the same class; registering the Dag reports that. + * Because a switch names a case by its class, two tasks that run the same + * class cannot both be cases of it; registering the Dag reports that. * * @see DagDef.Switch */ @@ -104,29 +74,6 @@ interface SwitchTask : Task { ) } -/** - * @suppress - * - * What a generated `@Builder.Switch` task implements: a switch whose method - * named the case with a [TaskId]. Public so generated code can implement it; - * a Dag written against [DagDef.Switch] implements [SwitchTask] instead. - */ -interface TaskIdSwitchTask : Task { - @Throws(Exception::class) - fun choose( - context: Context, - client: Client, - ): TaskId - - 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. * @@ -158,10 +105,7 @@ class SwitchRef private constructor( "Task '${ref.def.id}' already decides what to skip; declare it once" } val definition = ref.def.definition - require( - SwitchTask::class.java.isAssignableFrom(definition) || - TaskIdSwitchTask::class.java.isAssignableFrom(definition), - ) { + require(SwitchTask::class.java.isAssignableFrom(definition)) { "Task '${ref.def.id}' runs '${definition.name}', which chooses nothing; a switch runs a " + "SwitchTask" } @@ -236,16 +180,13 @@ internal class SwitchDef : DeciderDef { require(options.none { it === case }) { "Switch '${decider.id}' already chooses between '${case.id}' and others; name each case once" } - // A SwitchTask names its case by class, so two cases sharing one would be - // indistinguishable. Checked here rather than at registration so the error - // points at the second Case(...) call. - if (SwitchTask::class.java.isAssignableFrom(decider.definition)) { - 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", - ) - } + // 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 } @@ -262,7 +203,6 @@ internal class SwitchDef : DeciderDef { val chosen = when (instance) { is SwitchTask -> byClass(decider, instance.choose(context, client)) - is TaskIdSwitchTask -> byId(decider, instance.choose(context, client).value) else -> throw IllegalStateException( "Task '${decider.id}' is a switch, but '${instance.javaClass.name}' is not a SwitchTask", ) @@ -279,14 +219,5 @@ internal class SwitchDef : DeciderDef { "Switch '${decider.id}' chose ${definition?.name ?: "nothing"}, which is not one of its cases: ${describeCases()}", ) - private fun byId( - decider: TaskDef, - taskId: String, - ): TaskDef = - options.firstOrNull { it.id == taskId } - ?: throw IllegalArgumentException( - "Switch '${decider.id}' chose '$taskId', which is not one of its cases: ${describeCases()}", - ) - private fun describeCases(): String = options.joinToString { "'${it.id}'" } } 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 index dbbd97caaa7b1..c066f43c00577 100644 --- 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 @@ -165,11 +165,4 @@ internal class SwitchTest { assertEquals("reports.pick", switch.id) assertEquals(listOf("reports.long", "reports.pick"), dag.tasks.keys.toList()) } - - @Test - @DisplayName("Should carry the task ID of a case rather than its class") - fun taskIdNamesOneTask() { - assertEquals(TaskId.of("handle_long"), TaskId.of("handle_long")) - assertEquals("handle_long", TaskId.of("handle_long").toString()) - } } 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 91adeb027c843..e327afd661e4b 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 @@ -40,10 +40,9 @@ 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.TaskId -import org.apache.airflow.sdk.TaskIdSwitchTask import org.apache.airflow.sdk.TaskRef import org.apache.airflow.sdk.execution.serializeDag import org.apache.airflow.sdk.internal.Field @@ -53,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> = + 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, @@ -67,13 +161,12 @@ class ConformanceCondition : ConditionTask { ) = true } -// Names its case by task ID: every conformance task runs the same class, so a -// branch that named a case by class could not tell two of them apart. -class ConformanceSwitch : TaskIdSwitchTask { +// Never run, only serialized, so any task class will do. +class ConformanceSwitch : SwitchTask { override fun choose( context: Context, client: Client, - ) = TaskId.of("") + ): Class = ConformanceTask1::class.java } fun main(args: Array) { @@ -102,6 +195,7 @@ private fun buildDag(case: JsonNode): DagDef { } val tasks = linkedMapOf>() + 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>() case.path("tasks").forEach { task -> @@ -110,7 +204,12 @@ private fun buildDag(case: JsonNode): DagDef { val branch = task.path("branch") val definition = when { - branch.isMissingNode -> ConformanceTask::class.java + 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 } 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 bc8f9d721363d..2ea52dbea0cf3 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,12 +24,9 @@ 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.SwitchRef import org.apache.airflow.sdk.SwitchTask import org.apache.airflow.sdk.Task import org.apache.airflow.sdk.TaskDef -import org.apache.airflow.sdk.TaskId -import org.apache.airflow.sdk.TaskIdSwitchTask import org.apache.airflow.sdk.execution.comm.BundleInfo import org.apache.airflow.sdk.execution.comm.DagRun import org.apache.airflow.sdk.execution.comm.RetryTask @@ -279,21 +276,6 @@ class TaskTest { Assertions.assertEquals(emptyList>(), transport.xComs) } - @Test - @DisplayName("Should let a generated switch name its case by task ID") - fun shouldChooseByTaskId() { - val dag = DagDef("test_dag") - val long = dag.task("handle_long", HandleLongTask::class.java) - val short = dag.task("handle_short", HandleShortTask::class.java) - SwitchRef.of(dag.task("pick", TestTaskIdSwitch::class.java)).Case(long).Case(short) - val transport = RecordingTransport() - - runTask(Bundle(listOf(dag)), startupDetails(taskId = "pick"), Client(startupDetails("pick"), transport)) - - Assertions.assertEquals(listOf("handle_long"), transport.skipped) - Assertions.assertEquals("handle_short", transport.xComs.last().second) - } - private fun switchBundle(): Bundle { val dag = DagDef("test_dag") val long = dag.task("handle_long", HandleLongTask::class.java) @@ -443,14 +425,6 @@ class TaskTest { } } - /** The shape the annotation processor generates: a switch that names its case by ID. */ - class TestTaskIdSwitch : TaskIdSwitchTask { - override fun choose( - context: Context, - client: Client, - ): TaskId = TaskId.of("handle_short") - } - /** 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>() From 35fd6f24a512b4ef31b1a3e1706936be8d17d7bc Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Thu, 8 Oct 2026 11:25:59 +0000 Subject: [PATCH 4/7] Java SDK: Say when a switch rejects a choice and when it fails --- .../main/kotlin/org/apache/airflow/sdk/Switch.kt | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) 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 index e2cfa28d3d7a9..7657006be7ebd 100644 --- 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 @@ -25,7 +25,7 @@ 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 task's own class, so javac checks it: + * [SwitchRef.Case]. The choice is the class of the task to run: * * ```java * public class PickPath implements SwitchTask { @@ -41,8 +41,13 @@ import kotlin.Throws * 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; registering the Dag reports that. + * class cannot both be cases of it; the second [SwitchRef.Case] call throws. * * @see DagDef.Switch */ @@ -122,8 +127,8 @@ class SwitchRef private constructor( * * @param task Task of the same Dag. * @return This switch, for chaining. - * @throws IllegalArgumentException if the task belongs to another Dag or is - * already a case of this switch. + * @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 { From 08909559642c38df009f5a8480b8e6887c2300cc Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Thu, 8 Oct 2026 11:26:16 +0000 Subject: [PATCH 5/7] Java SDK: Cover a generated switch that names a task it cannot choose --- .../apache/airflow/sdk/execution/TaskTest.kt | 33 +++++++++++++++++++ 1 file changed, 33 insertions(+) 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 2ea52dbea0cf3..db1478661c186 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 @@ -276,6 +276,24 @@ class TaskTest { Assertions.assertEquals(emptyList>(), 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("handle_long", HandleLongTask::class.java) + val short = dag.task("handle_short", HandleShortTask::class.java) + dag.task("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(), transport.skipped) + Assertions.assertEquals(emptyList>(), transport.xComs) + } + private fun switchBundle(): Bundle { val dag = DagDef("test_dag") val long = dag.task("handle_long", HandleLongTask::class.java) @@ -425,6 +443,21 @@ class TaskTest { } } + 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 = 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>() From b10b0dc0595b65fc34e4353e33c0a053aa16dbc9 Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Thu, 8 Oct 2026 11:26:58 +0000 Subject: [PATCH 6/7] Java SDK: Check a switch carries the chosen task's id --- .../org/apache/airflow/sdk/SwitchTest.kt | 21 +++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) 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 index c066f43c00577..3b19fa6553d32 100644 --- 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 @@ -32,16 +32,12 @@ internal class HandleLong : Task { ) = Unit } -/** A switch that chooses whatever [choice] holds. */ +/** A switch that always chooses [HandleLong]. */ class PickPath : SwitchTask { override fun choose( context: Context, client: Client, - ): Class = choice - - companion object { - var choice: Class = HandleLong::class.java - } + ): Class = HandleLong::class.java } internal class SwitchTest { @@ -165,4 +161,17 @@ internal class SwitchTest { 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) + } } From e643f06dc68fcf68bcbea92c4d715b381037bd96 Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Thu, 8 Oct 2026 11:27:53 +0000 Subject: [PATCH 7/7] Java SDK: Let a wiring view name a switch's cases in more than one call --- .../org/apache/airflow/sdk/BuilderTest.kt | 38 +++++++++++++++++++ .../kotlin/org/apache/airflow/sdk/Switch.kt | 10 ++++- .../org/apache/airflow/sdk/SwitchTest.kt | 30 +++++++++++++++ .../apache/airflow/sdk/internal/RefsTest.kt | 36 ++++++++++++++++++ 4 files changed, 112 insertions(+), 2 deletions(-) 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 e0c60a3fc8289..2473ef9561aa6 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,44 @@ 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 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() { 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 index 7657006be7ebd..a668db89609a0 100644 --- 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 @@ -103,11 +103,17 @@ class SwitchRef private constructor( * * 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 { - require(ref.def.decider == null) { - "Task '${ref.def.id}' already decides what to skip; declare it once" + 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)) { 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 index 3b19fa6553d32..ef72edf3a4c0b 100644 --- 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 @@ -174,4 +174,34 @@ internal class SwitchTest { 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(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(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/internal/RefsTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/RefsTest.kt index 59e36c182bd77..c764585e409fe 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 = 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("", TaskDef("pick", NoopSwitch::class.java))) + pick().Case(Refs.node("", TaskDef("long", NoopRefTask::class.java))) + pick().Case(Refs.node("", TaskDef("short", OtherNoopRefTask::class.java))) + } + + assertEquals( + listOf("long", "short"), + dag.tasks + .getValue("pick") + .decider!! + .cases + .map { it.id }, + ) + } }