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