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 d4013cc2b5b52842f75ca6cf6360306600d69aa4 Author: Jason(Zhe-You) Liu <[email protected]> AuthorDate: Fri Oct 9 14:30:52 2026 +0800 Java SDK: Pick one of two tasks with a condition (#74438) * Java SDK: Pick one of two tasks with a condition A native Java Dag could only run every task it declared: it had no way to let a task's result decide what runs next, so a Dag that loads only when there are rows had to be split in two or fall back to a Python stub. `DagDef.If` and `@Builder.If` declare a task whose boolean names the side that runs; the other is skipped. The deciding task carries `_can_skip_downstream`, so Airflow skips a cleared side again rather than running it, and the SDK writes the `skipmixin_key` XCom before the skip, as Python's SkipMixin does. It skips only the side not taken, so a task after both sides needs a trigger rule such as `none_failed_min_one_success`. `If` is capitalized because `if` is a Java keyword and `@Builder.If` reads the same way. With no task id given, one is derived from the condition class's simple name, so `HasRows.class` becomes `hasRows`. The runtime reaches `SkipDownstreamTasks` through the transport client only, never through the task-facing `Client`, so a task body cannot skip tasks the Dag did not declare it may. * Java SDK: Let a wiring view name a condition's sides in more than one call * Java SDK: Reject a decider class registered as an ordinary task * Java SDK: Name a condition's sides Then and Else * Java SDK: Pin the order a condition skips and records what it skipped * Java SDK: Drop the unused switch in the condition test's decider --- .pre-commit-config.yaml | 2 +- java-sdk/README.md | 2 +- java-sdk/capabilities.yaml | 4 +- .../example/nativedag/AnnotationExample.java | 15 +- .../example/nativedag/InterfaceExample.java | 21 +- .../org/apache/airflow/sdk/BuilderProcessor.kt | 103 +++++++--- .../kotlin/org/apache/airflow/sdk/BuilderTest.kt | 198 ++++++++++++++++++ .../ci/prek/check_serialization_conformance.py | 2 +- java-sdk/sdk/build.gradle.kts | 33 +++ java-sdk/sdk/module.md | 2 +- .../main/kotlin/org/apache/airflow/sdk/Bundle.kt | 26 ++- .../kotlin/org/apache/airflow/sdk/Condition.kt | 224 +++++++++++++++++++++ .../main/kotlin/org/apache/airflow/sdk/DagDef.kt | 52 +++++ .../main/kotlin/org/apache/airflow/sdk/Decider.kt | 90 +++++++++ .../kotlin/org/apache/airflow/sdk/TaskGroupRef.kt | 32 +++ .../org/apache/airflow/sdk/execution/Client.kt | 14 ++ .../org/apache/airflow/sdk/execution/Serde.kt | 2 + .../org/apache/airflow/sdk/execution/Task.kt | 29 ++- .../kotlin/org/apache/airflow/sdk/internal/Ids.kt | 10 + .../org/apache/airflow/sdk/ArgTestSupport.kt | 2 + .../kotlin/org/apache/airflow/sdk/ClientTest.kt | 2 + .../kotlin/org/apache/airflow/sdk/ConditionTest.kt | 208 +++++++++++++++++++ .../airflow/sdk/conformance/SerializeJava.kt | 29 ++- .../org/apache/airflow/sdk/execution/CommTest.kt | 27 +++ .../org/apache/airflow/sdk/execution/TaskTest.kt | 150 ++++++++++++++ .../apache/airflow/sdk/internal/ArgValuesTest.kt | 2 + .../org/apache/airflow/sdk/internal/RefsTest.kt | 29 +++ 27 files changed, 1274 insertions(+), 36 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 8392915d926..37f87710dfc 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|DagDef|Deps|Arg|TaskGroupRef)\.kt$| + ^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/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 8a3284403b8..1e4e7478f39 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 † | ✗ | – | | +| capability: `branching` | SHOULD † | ✓ | 3.4 | DagDef.If / @Builder.If | | 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 75c6053fc89..33106f170fb 100644 --- a/java-sdk/capabilities.yaml +++ b/java-sdk/capabilities.yaml @@ -107,7 +107,9 @@ capabilities: supported: true since: "3.4" branching: - supported: false + supported: true + since: "3.4" + note: "DagDef.If / @Builder.If" 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 1c6ce70cc72..d9538fd391f 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 @@ -54,6 +54,18 @@ public class AnnotationExample { log.log(INFO, "Loaded {0}", transformed); } + @Builder.Task(id = "load_empty") + public void loadEmpty() { + log.log(INFO, "Nothing to load"); + } + + // A condition: its boolean picks one of the two loads, and the other is + // skipped. + @Builder.If(id = "has_rows") + public boolean hasRows(long transformed) { + return transformed > 0; + } + // A task group: everything it declares is prefixed with its id, so this is // the task "checks.audit". @Builder.TaskGroup(id = "checks") @@ -70,7 +82,8 @@ public class AnnotationExample { static class Wiring implements AnnotationExampleDeps { void depends() { var extracted = extract(); - load(transform(extracted, lit(1.5))); + var transformed = transform(extracted, lit(1.5)); + hasRows(transformed).Then(load(transformed)).Else(loadEmpty()); // 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 6047aeedd8e..cce10676dfd 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 @@ -55,6 +55,22 @@ public class InterfaceExample { } } + public static class LoadEmpty implements Task { + @Override + public void execute(Context context, Client client) { + log.log(INFO, "Nothing to load"); + } + } + + // A condition: its boolean picks one of the two loads, and the other is + // skipped. + public static class HasRows implements ConditionTask { + @Override + public boolean decide(Context context, Client client) { + return ((Number) client.getXCom("transform")).longValue() > 0; + } + } + public static DagDef build() { var dag = new DagDef("java_native_interface_example") @@ -69,8 +85,11 @@ public class InterfaceExample { .config("doc_md", "Extracts a value and pushes it as an XCom."); var transform = dag.task("transform", Transform.class); var load = dag.task("load", Load.class); + var loadEmpty = dag.task("load_empty", LoadEmpty.class); - transform.after(extract).before(load); + transform.after(extract); + // With no task id given, the condition takes one from its class: "hasRows". + dag.If(HasRows.class).after(transform).Then(load).Else(loadEmpty); 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 145ee4714c7..1848fc007f7 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 @@ -96,6 +96,7 @@ import org.apache.airflow.sdk.internal.builderName as generatedBuilderName @SupportedAnnotationTypes( "org.apache.airflow.sdk.Builder.Dag", "org.apache.airflow.sdk.Builder.Task", + "org.apache.airflow.sdk.Builder.If", "org.apache.airflow.sdk.Builder.TaskGroup", "org.apache.airflow.sdk.Builder.TaskHandler", "org.apache.airflow.sdk.Builder.Deps", @@ -356,29 +357,42 @@ 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(ParameterizedTypeName.get(TASK_HANDLE_TYPE, TypeName.get(decl.method.returnType).boxIfPossible())) + .returns( + if (condition) { + CONDITION_REF_TYPE + } else { + 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 } if (decl.dataParams.isEmpty()) { - method.addStatement($$"return $T.node($L, $L)", REFS_TYPE, group, def) + method.addStatement($$"return $L", wrap(CodeBlock.of($$"$T.node($L, $L)", REFS_TYPE, group, def))) } else { // The parameter names ride along so the serialized Dag can name each // argument, as the binding spec ADR-0007 defines requires. method.addStatement( - $$"return $T.call($L, $L, $T.of($L), $L)", - REFS_TYPE, - group, - def, - LIST_TYPE, - CodeBlock.join(decl.dataParams.map { CodeBlock.of($$"$S", it.name) }, ", "), - decl.dataParams.joinToString { it.name }, + $$"return $L", + wrap( + CodeBlock.of( + $$"$T.call($L, $L, $T.of($L), $L)", + REFS_TYPE, + group, + def, + LIST_TYPE, + CodeBlock.join(decl.dataParams.map { CodeBlock.of($$"$S", it.name) }, ", "), + decl.dataParams.joinToString { it.name }, + ), + ), ) } return method.build() @@ -396,7 +410,7 @@ class BuilderProcessor : AbstractProcessor() { CodeBlock .builder() .add($$"new $T($S, $L.class)", TASK_DEF_TYPE, decl.id, classRef) - explicitConfig(decl.method, TASK_ANNOTATION, TASK_STRUCTURAL_ATTRIBUTES, SchemaFields.TASK).forEach { (key, value) -> + explicitConfig(decl.method, decl.kind.annotation, TASK_STRUCTURAL_ATTRIBUTES, SchemaFields.TASK).forEach { (key, value) -> taskDef.add($$".config($S, $L)", key, value) } return taskDef.build() @@ -421,16 +435,24 @@ class BuilderProcessor : AbstractProcessor() { val tasks = mutableListOf<TaskDeclaration>() for (inner in el.enclosedElements) { if (inner !is ExecutableElement) continue - val ann = inner.getAnnotation(Builder.Task::class.java) ?: 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 if (inner.isVarArgs) throw IllegalArgumentException("Cannot create task from vararg function ${inner.simpleName}") - val localId = ann.id.ifBlank { inner.simpleName.toString() } + checkDeciderReturn(kind, inner) + val localId = (condition?.id ?: task!!.id).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 " + - "task id with @Builder.Task(id = \"$localId\")" + "task id with ${kind.spelling}(id = \"$localId\")" } tasks += - TaskDeclaration(inner, (path + localId).joinToString("."), collectDataParams(inner), el, classPath) + TaskDeclaration(inner, (path + localId).joinToString("."), collectDataParams(inner), el, classPath, kind) } val groups = mutableListOf<GroupDeclaration>() @@ -571,7 +593,7 @@ class BuilderProcessor : AbstractProcessor() { val name = decl.method.simpleName.toString() require(name !in reserved) { "Task method '$name' clashes with a member of the wiring view; rename the method and keep " + - "the task id with @Builder.Task(id = \"${decl.id.substringAfterLast('.')}\")" + "the task id with ${decl.kind.spelling}(id = \"${decl.id.substringAfterLast('.')}\")" } } val accessors = mutableMapOf<String, GroupDeclaration>() @@ -682,13 +704,31 @@ class BuilderProcessor : AbstractProcessor() { } } + /** + * Checks what a decider hands back, which is what the SDK reads its decision + * from. A plain task may return anything, including nothing. + */ + private fun checkDeciderReturn( + 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" + } + } + private fun buildTask(decl: TaskDeclaration): TypeSpec { + val condition = decl.kind == TaskKind.CONDITION val executeSpec = MethodSpec - .methodBuilder("execute") + .methodBuilder(if (condition) "decide" else "execute") .addAnnotation(Override::class.java) .addModifiers(Modifier.PUBLIC) - .returns(TypeName.VOID) + .returns(if (condition) TypeName.BOOLEAN else TypeName.VOID) .addParameter(CONTEXT_TYPE, "context") .addParameter(CLIENT_TYPE, "client") .addException(Exception::class.java) @@ -734,10 +774,11 @@ class BuilderProcessor : AbstractProcessor() { } } - if (inner.returnType.kind == TypeKind.VOID) { - $$"new $T().$L($L)" - } else { - $$"client.setXCom(new $T().$L($L))" + when { + // The SDK pushes a condition's result itself, after it has skipped the side not taken. + condition -> $$"return new $T().$L($L)" + inner.returnType.kind == TypeKind.VOID -> $$"new $T().$L($L)" + else -> $$"client.setXCom(new $T().$L($L))" }.also { executeSpec.addStatement( it, @@ -749,7 +790,7 @@ class BuilderProcessor : AbstractProcessor() { return TypeSpec .classBuilder(decl.className) - .addSuperinterface(Task::class.java) + .addSuperinterface(if (condition) CONDITION_TASK_TYPE else ClassName.get(Task::class.java)) .addModifiers(Modifier.PUBLIC, Modifier.FINAL, Modifier.STATIC) .addMethod(executeSpec.build()) .build() @@ -893,6 +934,7 @@ private class TaskDeclaration( val dataParams: List<DataParam>, val owner: TypeElement, val classPath: List<String> = emptyList(), + val kind: TaskKind = TaskKind.TASK, ) { val className: String = (classPath + method.simpleName.toString().replaceFirstChar(Char::uppercase)).joinToString("_") @@ -924,12 +966,27 @@ 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 CONDITION_TASK_TYPE = ClassName.get(ConditionTask::class.java) +private val CONDITION_REF_TYPE = ClassName.get(ConditionRef::class.java) +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") private val LIST_TYPE = ClassName.get(List::class.java) private const val DAG_ANNOTATION = "org.apache.airflow.sdk.Builder.Dag" -private const val TASK_ANNOTATION = "org.apache.airflow.sdk.Builder.Task" + +/** + * What a declared task is, and the annotation that declares it. A decider + * differs from a plain task in what it implements and in what its wiring-view + * method hands back, but carries the same configuration attributes. + */ +private enum class TaskKind( + val annotation: String, + val spelling: String, +) { + TASK("org.apache.airflow.sdk.Builder.Task", "@Builder.Task"), + CONDITION("org.apache.airflow.sdk.Builder.If", "@Builder.If"), +} private val RESERVED_VIEW_NAMES = setOf( 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 a30dbacf008..70802c405aa 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 @@ -2161,6 +2161,204 @@ class BuilderTest { ) } + @Test + @DisplayName("generate a condition task and a wiring view that names each side") + fun generateConditionTask() { + val compilation = + compile( + """ + package org.apache.airflow.example; + + import org.apache.airflow.sdk.Builder; + + @Builder.Dag(id = "etl") + public class TestExample { + @Builder.If(id = "has_rows") + public boolean hasRows() { + return true; + } + + @Builder.Task + public void load() {} + + @Builder.Deps + static class Wiring implements TestExampleDeps { + void depends() { + hasRows().Then(load()); + } + } + } + """, + ) + + 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.ConditionTask; + import org.apache.airflow.sdk.Context; + import org.apache.airflow.sdk.DagDef; + 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("has_rows", "load"), List.of(), new TestExample.Wiring()::depends); + } + + public static final class HasRows implements ConditionTask { + @Override + public boolean decide(Context context, Client client) throws Exception { + return new TestExample().hasRows(); + } + } + + public static final class Load implements Task { + @Override + public void execute(Context context, Client client) throws Exception { + new TestExample().load(); + } + } + } + """, + ) + 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.ConditionRef; + 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 ConditionRef hasRows() { + return ConditionRef.of(Refs.node("", new TaskDef("has_rows", TestExampleBuilder.HasRows.class))); + } + + default TaskRef<Void> load() { + return Refs.node("", new TaskDef("load", TestExampleBuilder.Load.class)); + } + } + """, + ) + } + + @Test + @DisplayName("let a wiring view name a condition's sides in more than one statement") + fun conditionNamedAcrossStatements() { + val compilation = + compile( + """ + package org.apache.airflow.example; + + import org.apache.airflow.sdk.Builder; + + @Builder.Dag(id = "etl") + public class TestExample { + @Builder.Task + public void extract() {} + + @Builder.If(id = "has_rows") + public boolean hasRows() { + return true; + } + + @Builder.Task + public void load() {} + + @Builder.Task + public void reportEmpty() {} + + @Builder.Deps + static class Wiring implements TestExampleDeps { + void depends() { + extract().before(hasRows()); + hasRows().Then(load()); + hasRows().Else(reportEmpty()); + } + } + } + """, + ) + + assertThat(compilation).succeeded() + } + + @Test + @DisplayName("reject a condition that does not return a boolean") + fun rejectNonBooleanCondition() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + @Builder.Dag + public class TestExample { + @Builder.If + public String hasRows() { + return "yes"; + } + + @Builder.Deps + static class Wiring implements TestExampleDeps { + void depends() {} + } + } + """, + ) + + assertThat(compilation).failed() + assertThat(compilation).hadErrorContaining( + "@Builder.If method 'hasRows' returns java.lang.String, but a condition returns boolean", + ) + } + + @Test + @DisplayName("reject a method that is both a task and a condition") + fun rejectTaskAndCondition() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + @Builder.Dag + public class TestExample { + @Builder.Task + @Builder.If + public boolean hasRows() { + return true; + } + + @Builder.Deps + static class Wiring implements TestExampleDeps { + void depends() {} + } + } + """, + ) + + assertThat(compilation).failed() + assertThat(compilation).hadErrorContaining( + "Method 'hasRows' carries both @Builder.Task and @Builder.If", + ) + } + @Test @DisplayName("map the Dag to the annotated class, not its generated builder") fun dagDeclaredByAnnotatedClass() { diff --git a/java-sdk/scripts/ci/prek/check_serialization_conformance.py b/java-sdk/scripts/ci/prek/check_serialization_conformance.py index 9869f4ccc5c..fe569c1ff7c 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", + "literal_inputs,branch", "--", *serializer, ] diff --git a/java-sdk/sdk/build.gradle.kts b/java-sdk/sdk/build.gradle.kts index 0159e1d2803..723f60085c2 100644 --- a/java-sdk/sdk/build.gradle.kts +++ b/java-sdk/sdk/build.gradle.kts @@ -581,6 +581,39 @@ abstract class GenerateDagDslTask : DefaultTask() { | ) | | /** + | * Marks a task method whose `boolean` picks one of two tasks; the + | * other is skipped. + | * + | * The method is an ordinary task method that returns `boolean`. Its + | * wiring-view method returns a [ConditionRef] rather than a + | * [TaskRef], so each side is named where the Dag is wired: + | * + | * ```java + | * @Builder.If(id = "has_rows") + | * public boolean hasRows(long rows) { return rows > 0; } + | * + | * @Builder.Deps + | * static class Wiring implements EtlDeps { + | * void depends() { hasRows(extract()).Then(load()).Else(reportEmpty()); } + | * } + | * ``` + | * + | * The condition skips only the side not taken, so a task that runs + | * after both sides needs a trigger rule that tolerates one 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 If( + | /** 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 3d11a42da5e..523dbaf54a4 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 † | ✗ | – | | +| capability: `branching` | SHOULD † | ✓ | 3.4 | DagDef.If / @Builder.If | | 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 bd055ab2c25..7937446c3f6 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 @@ -66,8 +66,8 @@ class Bundle( * @return This bundle, for chaining. * @throws IllegalArgumentException if another Dag shares its ID, task * handlers are already registered against it, a task depends on an - * upstream not registered in the same Dag, or the dependencies contain a - * cycle. + * 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. * @throws IllegalStateException if [Server.serve] has already been called. */ fun register(dag: DagDef): Bundle { @@ -83,11 +83,30 @@ class Bundle( "that is not registered in the same Dag" } } + val decides = DECIDER_TYPES.any { it.isAssignableFrom(def.definition) } + 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}'" + } else { + "Task '${def.id}' chooses which task runs, but '${def.definition.name}' decides nothing, in Dag '${dag.id}'" + } + } + def.decider?.let { decider -> + decider.describe(def)?.let { throw IllegalArgumentException("$it, in Dag '${dag.id}'") } + for (case in decider.cases) { + require(dag.tasks[case.id] === case) { + "Task '$taskId' in Dag '${dag.id}' can run task '${case.id}' " + + "that is not registered in the same Dag" + } + } + } } checkNoCycle(dag) require(dags.putIfAbsent(dag.id, dag) == null) { "Dags in bundle have duplicate ID: ${dag.id}" } + dag.registered = true return this } @@ -192,6 +211,9 @@ class Bundle( private fun checkOpen() = check(!served) { "Server.serve has already been called; register everything before serve" } } +/** Task types the SDK runs through a decider, so each needs the sides it chooses between. */ +private val DECIDER_TYPES = listOf(ConditionTask::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 // for readability since the scenario is unlikely; feel free to rewrite if it blows up for you. 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 new file mode 100644 index 00000000000..90150609c7b --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Condition.kt @@ -0,0 +1,224 @@ +/* + * 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 whose boolean decides which of two tasks runs; the other is skipped. + * + * Register one with [DagDef.If], then name the sides with [ConditionRef.Then] + * and [ConditionRef.Else]: + * + * ```java + * public class HasRows implements ConditionTask { + * @Override + * public boolean decide(Context context, Client client) { + * return ((Number) client.getXCom("extract")).longValue() > 0; + * } + * } + * ``` + * + * The SDK runs [decide] and pushes its result as this task's return value, so + * [execute] is never called. + * + * @see DagDef.If + */ +interface ConditionTask : Task { + /** + * Decides which side of the condition 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 True to run the [ConditionRef.Then] side, false the + * [ConditionRef.Else] side. + * @throws Exception on failure; the task instance is marked failed. + */ + @Throws(Exception::class) + fun decide( + context: Context, + client: Client, + ): Boolean + + /** Never called: the SDK runs a condition through [decide]. */ + override fun execute( + context: Context, + client: Client, + ): Unit = + throw IllegalStateException( + "Condition '${javaClass.name}' runs through decide(), so execute() is never called", + ) +} + +/** + * A condition registered with a Dag: name the task each outcome runs. + * + * ```java + * dag.If(HasRows.class).Then(load).Else(reportEmpty); + * ``` + * + * A named task runs after the condition, so naming it records that edge, as + * [Deps.Flow.before] would. The condition skips only the side not taken; what + * runs after both sides needs a trigger rule that tolerates one skipped + * upstream, such as `none_failed_min_one_success`. + * + * @see DagDef.If + */ +class ConditionRef private constructor( + private val ref: TaskRef<Boolean>, + private val sides: ConditionDef, +) : Deps.Flow { + companion object { + /** + * @suppress + * + * Marks a registered task as a condition. Public so a generated wiring + * view can call it; user code reaches a condition through [DagDef.If]. + * + * A wiring view gets the same task back each time it names it, so a second + * call reuses the condition rather than declaring it again. [DagDef.If] and + * [TaskGroupRef.If] never reach that path, since registering a task rejects + * a duplicate ID first. + */ + @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" } + return ConditionRef(ref, existing) + } + return ConditionRef(ref, ConditionDef().also { ref.def.decider = it }) + } + } + + /** Airflow task ID of the deciding task. */ + val id: String get() = ref.def.id + + /** + * Names the task that runs when the condition holds. + * + * @param task Task of the same Dag, which the condition skips when it does + * not hold. + * @return This condition, so [Else] can follow. + * @throws IllegalArgumentException if the task belongs to another Dag, or + * this side is already named. + */ + @Suppress("ktlint:standard:function-naming") + fun Then(task: TaskRef<*>): ConditionRef = name("Then", task) { sides.whenTrue = it } + + /** + * Names the task that runs when the condition does not hold. A condition + * needs no `Else`; without one, nothing is skipped when it holds. + * + * @param task Task of the same Dag, which the condition skips when it holds. + * @return This condition, for chaining. + * @throws IllegalArgumentException if the task belongs to another Dag, is + * the `Then` task, or this side is already named. + */ + @Suppress("ktlint:standard:function-naming") + fun Else(task: TaskRef<*>): ConditionRef = name("Else", task) { sides.whenFalse = it } + + /** + * 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 condition, for chaining. + * @throws IllegalArgumentException if the key is unknown or the value type + * does not match. + */ + fun config( + key: String, + value: Any?, + ): ConditionRef { + ref.config(key, value) + return this + } + + override fun nodes(): List<TaskDef> = ref.nodes() + + override fun before(vararg next: Deps.Flow): ConditionRef { + ref.before(*next) + return this + } + + override fun after(vararg previous: Deps.Flow): ConditionRef { + ref.after(*previous) + return this + } + + private fun name( + side: String, + task: TaskRef<*>, + assign: (TaskDef) -> Unit, + ): ConditionRef { + sides.check(ref.def, side, task.def) + assign(task.def) + // The side runs after the condition, which is what puts the condition in + // the serialized Dag as its upstream and lets registration see a cycle + // that runs through a condition. + ref.before(task) + return this + } +} + +/** The two sides of one condition, as [ConditionRef] named them. */ +internal class ConditionDef : DeciderDef { + var whenTrue: TaskDef? = null + var whenFalse: TaskDef? = null + + override val cases: List<TaskDef> get() = listOfNotNull(whenTrue, whenFalse) + + override fun describe(decider: TaskDef): String? = + if (whenTrue == null) { + "Condition '${decider.id}' names no task to run when it holds; call Then(...)" + } else { + null + } + + override fun decide( + decider: TaskDef, + instance: Task, + context: Context, + client: Client, + ): Decision { + val held = (instance as ConditionTask).decide(context, client) + return Decision(held, listOfNotNull((if (held) whenFalse else whenTrue)?.id)) + } + + fun check( + decider: TaskDef, + side: String, + task: TaskDef, + ) { + requireUnregistered(decider, "Condition") + val named = if (side == "Then") whenTrue else whenFalse + require(named == null) { + "Condition '${decider.id}' already runs '${named!!.id}' on its $side side; name each side once" + } + requireSameDag(decider, task, "Condition") + require(task !== whenTrue && task !== whenFalse) { + "Condition '${decider.id}' already runs '${task.id}' on its other side, so the condition " + + "would decide nothing" + } + } +} 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 a269dff8c2c..40ead4ae608 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 @@ -23,6 +23,7 @@ import org.apache.airflow.sdk.internal.DagSource import org.apache.airflow.sdk.internal.GROUP_ID import org.apache.airflow.sdk.internal.SchemaFields import org.apache.airflow.sdk.internal.checkConfigValue +import org.apache.airflow.sdk.internal.deriveTaskId import org.apache.airflow.sdk.internal.validateTaskInput import kotlin.Throws @@ -58,6 +59,13 @@ class DagDef( /** Edges with a task group at either end, in the order drawn. */ internal val groupEdges = linkedSetOf<Pair<Endpoint, Endpoint>>() + /** + * Whether a [Bundle] has taken this Dag. A decider checked when the Dag was + * registered cannot be changed afterwards, because nothing would check the + * change. + */ + internal var registered: Boolean = false + /** Outermost class that declared this Dag, or `null` if it could not be told. */ internal var declaringClass: Class<*>? = DagSource.capture() @@ -152,6 +160,47 @@ class DagDef( return this } + /** + * Declares a task whose boolean picks one of two tasks; the other is + * skipped. + * + * The task's ID is the class's simple name with its first character + * lowercased, so `HasRows.class` becomes `hasRows`. Use + * [If(id, definition)][If] to set it, and [ConditionRef.config] for the + * task's other settings: + * + * ```java + * dag.If(HasRows.class).config("retries", 2).Then(load).Else(reportEmpty); + * ``` + * + * @param definition Class that implements [ConditionTask]. Must have a + * public no-arg constructor. + * @return The condition, to name each side on. + * @throws IllegalArgumentException if a task with the same ID is already + * registered. + */ + @Suppress("ktlint:standard:function-naming") + fun If(definition: Class<out ConditionTask>): ConditionRef = If(deriveTaskId(definition), definition) + + /** + * Declares a task whose boolean picks one of two tasks, under the task ID + * [id]. + * + * @param id Task identifier, unique within this Dag. + * @param definition Class that implements [ConditionTask]. Must have a + * public no-arg constructor. + * @return The condition, to name each side on. + * @throws IllegalArgumentException if a task with the same ID is already + * registered. + * + * @see If + */ + @Suppress("ktlint:standard:function-naming") + fun If( + id: String, + definition: Class<out ConditionTask>, + ): ConditionRef = ConditionRef.of(task(id, definition)) + /** * Declares a task group of this Dag. * @@ -326,6 +375,9 @@ 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. */ + internal var decider: DeciderDef? = null + /** * Sets one task-level configuration value. * diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Decider.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Decider.kt new file mode 100644 index 00000000000..15f4234b20d --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Decider.kt @@ -0,0 +1,90 @@ +/* + * 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 + +/** + * What one deciding task can run, and what it therefore skips. + * + * A [TaskDef] carrying one is a decider: the runtime asks it which of its + * [cases] the task chose, skips the rest, and the serialized task carries + * `_can_skip_downstream` so Airflow re-skips a cleared one. + */ +internal interface DeciderDef { + /** Every task this decider can run, in the order they were named. */ + val cases: List<TaskDef> + + /** Why this decider is incomplete, or null when it is ready to register. */ + fun describe(decider: TaskDef): String? + + /** + * Runs the deciding task [instance] and reports what it chose. + * + * @throws Exception whatever the task body throws, which fails the task + * without pushing or skipping anything. + */ + fun decide( + decider: TaskDef, + instance: Task, + context: Context, + client: Client, + ): Decision +} + +/** + * What a deciding task chose: the value it pushes as its return value, and + * the task IDs of the cases it did not take. + */ +internal class Decision( + val value: Any, + val skipped: List<String>, +) + +/** Mirrors `SkipMixin.skip` in `task-sdk/src/airflow/sdk/bases/skipmixin.py`. */ +internal const val SKIPMIXIN_XCOM_KEY = "skipmixin_key" +internal const val SKIPMIXIN_SKIPPED = "skipped" + +/** + * Rejects a case named after the Dag was registered, which the checks that run + * at registration would never see. + */ +internal fun requireUnregistered( + decider: TaskDef, + what: String, +) { + val owner = decider.owner + require(owner?.registered != true) { + "$what '${decider.id}' of Dag '${owner?.id}' is already registered; name every case before " + + "the Dag is added to a Bundle" + } +} + +/** Rejects a case that belongs to another Dag, which no edge of this one can reach. */ +internal fun requireSameDag( + decider: TaskDef, + case: TaskDef, + what: String, +) { + val owner = decider.owner + val caseOwner = case.owner + require(owner == null || caseOwner == null || owner === caseOwner) { + "$what '${decider.id}' of Dag '${owner?.id}' cannot run task '${case.id}' of Dag " + + "'${caseOwner?.id}'; name a task of the same 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 dc6448a5eb0..e2d7d9613ef 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 @@ -19,6 +19,8 @@ package org.apache.airflow.sdk +import org.apache.airflow.sdk.internal.deriveTaskId + /** * A group of tasks in a Dag, shown in the Airflow UI as one node that expands: * Python's `TaskGroup`. @@ -74,6 +76,36 @@ class TaskGroupRef internal constructor( return TaskRef(def) } + /** + * Declares a condition in this group, as [DagDef.If] does for the Dag. + * + * @param definition Class that implements [ConditionTask]. Must have a + * public no-arg constructor. + * @return The condition, to name each side on. + * @throws IllegalArgumentException if the Dag already has a task or task + * group with the resulting ID. + */ + @Suppress("ktlint:standard:function-naming") + fun If(definition: Class<out ConditionTask>): ConditionRef = If(deriveTaskId(definition), definition) + + /** + * Declares a condition in this group under the task ID `<group ID>.<id>`. + * + * @param id Task ID within this group. + * @param definition Class that implements [ConditionTask]. Must have a + * public no-arg constructor. + * @return The condition, to name each side on. + * @throws IllegalArgumentException if the Dag already has a task or task + * group with the resulting ID. + * + * @see If + */ + @Suppress("ktlint:standard:function-naming") + fun If( + id: String, + definition: Class<out ConditionTask>, + ): ConditionRef = ConditionRef.of(task(id, definition)) + /** * Nests a task group inside this one. * diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Client.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Client.kt index a4bb15c44cf..f3174e77e2a 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Client.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Client.kt @@ -33,6 +33,7 @@ import org.apache.airflow.sdk.execution.comm.OKResponse import org.apache.airflow.sdk.execution.comm.PutVariable import org.apache.airflow.sdk.execution.comm.SetTaskStateStore import org.apache.airflow.sdk.execution.comm.SetXCom +import org.apache.airflow.sdk.execution.comm.SkipDownstreamTasks import org.apache.airflow.sdk.execution.comm.TaskStateStoreResult import org.apache.airflow.sdk.execution.comm.VariableResult import org.apache.airflow.sdk.execution.comm.XComResult @@ -102,6 +103,13 @@ interface Client { ) fun clearTaskStateStore(tiId: UUID) + + /** + * Marks [taskIds] of the running Dag run skipped, as a decider's untaken + * branches are. Airflow leaves a task instance that is already running, + * succeeded or failed alone. + */ + fun skipDownstreamTasks(taskIds: List<String>) } /** @@ -231,4 +239,10 @@ class CoordinatorClient( override fun clearTaskStateStore(tiId: UUID) { runBlocking { exec.communicate<OKResponse>(ClearTaskStateStore().also { it.tiId = tiId }) } } + + override fun skipDownstreamTasks(taskIds: List<String>) { + if (taskIds.isEmpty()) return + val message = SkipDownstreamTasks().also { it.tasks = ArrayList<Any>(taskIds) } + runBlocking { exec.communicate<Unit>(message) } + } } diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Serde.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Serde.kt index 89245cf6946..bdc719a3a68 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Serde.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Serde.kt @@ -143,6 +143,8 @@ private fun serializeTask( "is_stub" to true, ) argBindings(taskId, def)?.let { data["_arg_bindings"] = it } + // Lets NotPreviouslySkippedDep re-skip a cleared downstream, as Python's SkipMixin does. + if (def.decider != null) data["_can_skip_downstream"] = true // Emit only config entries that differ from their schema default, mirroring // Python BaseSerialization's "omit hard-coded default" behavior, which the Go // and TypeScript SDKs mirror too. Operator fields are stored unwrapped, so the diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Task.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Task.kt index 6490e221ee9..182dfe1cee6 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Task.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Task.kt @@ -23,6 +23,11 @@ import kotlinx.coroutines.CancellationException import org.apache.airflow.sdk.Bundle import org.apache.airflow.sdk.Client import org.apache.airflow.sdk.Context +import org.apache.airflow.sdk.DeciderDef +import org.apache.airflow.sdk.SKIPMIXIN_SKIPPED +import org.apache.airflow.sdk.SKIPMIXIN_XCOM_KEY +import org.apache.airflow.sdk.Task +import org.apache.airflow.sdk.TaskDef import org.apache.airflow.sdk.execution.comm.AssetProfile import org.apache.airflow.sdk.execution.comm.RetryTask import org.apache.airflow.sdk.execution.comm.StartupDetails @@ -66,6 +71,26 @@ internal object TaskResult { fun failure(shouldRetry: Boolean) = if (shouldRetry) retry() else of(TaskState.State.FAILED) } +/** + * Runs a deciding task: Airflow skips every case it did not choose, and its + * choice becomes the task's return value. + */ +private fun decide( + decider: DeciderDef, + taskDef: TaskDef, + instance: Task, + context: Context, + client: Client, +) { + val decision = decider.decide(taskDef, instance, context, client) + if (decision.skipped.isNotEmpty()) { + // Written before the skip so a cleared case is skipped again, as SkipMixin does. + client.setXCom(SKIPMIXIN_XCOM_KEY, mapOf(SKIPMIXIN_SKIPPED to decision.skipped)) + client.impl.skipDownstreamTasks(decision.skipped) + } + client.setXCom(value = decision.value) +} + internal object TaskRunner { val logger = Logger(TaskRunner::class) @@ -104,7 +129,9 @@ internal object TaskRunner { return TaskResult.failure(request.tiContext.shouldRetry) } return try { - instance.execute(Context.from(request, taskDef), client) + val context = Context.from(request, taskDef) + val decider = taskDef.decider + if (decider == null) instance.execute(context, client) else decide(decider, taskDef, instance, context, client) TaskResult.success() } catch (e: CancellationException) { throw e // Let coroutine cancellation propagate so the task coroutine unwinds. diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/Ids.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/Ids.kt index b64551e9b01..b2824be3c91 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/Ids.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/Ids.kt @@ -27,3 +27,13 @@ package org.apache.airflow.sdk.internal * against the same pattern the SDK enforces; not user-facing API. */ val GROUP_ID: Regex = Regex("[A-Za-z0-9_-]+") + +/** + * @suppress + * + * The task ID a task declared from a class alone carries: the class's simple + * name with its first character lowercased, so `HasRows.class` and a + * `hasRows` task method name agree. Public so the annotation processor can + * check an ID it derives the same way; not user-facing API. + */ +fun deriveTaskId(definition: Class<*>): String = definition.simpleName.replaceFirstChar { it.lowercaseChar() } diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ArgTestSupport.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ArgTestSupport.kt index 0929dfe4e4e..6c6af528765 100644 --- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ArgTestSupport.kt +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ArgTestSupport.kt @@ -90,6 +90,8 @@ internal class FakeXComTransport( ) = throw NotImplementedError() override fun clearTaskStateStore(tiId: UUID) = throw NotImplementedError() + + override fun skipDownstreamTasks(taskIds: List<String>) = throw NotImplementedError() } internal fun startupDetails(argBindings: List<Map<String, Any?>>?): StartupDetails = diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ClientTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ClientTest.kt index dce06ca0977..2ff51df4f36 100644 --- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ClientTest.kt +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ClientTest.kt @@ -104,6 +104,8 @@ private class FakeTransport( override fun clearTaskStateStore(tiId: UUID) { calls.add(StateStoreCall("clear", tiId)) } + + override fun skipDownstreamTasks(taskIds: List<String>) = Unit } class ClientTest { 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 new file mode 100644 index 00000000000..1cebf7ff8e1 --- /dev/null +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ConditionTest.kt @@ -0,0 +1,208 @@ +/* + * 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.apache.airflow.sdk.execution.serializeDag +import org.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.Assertions.assertFalse +import org.junit.jupiter.api.Assertions.assertThrows +import org.junit.jupiter.api.Assertions.assertTrue +import org.junit.jupiter.api.DisplayName +import org.junit.jupiter.api.Test + +/** A condition the tests register; it always holds. */ +class HasRows : ConditionTask { + override fun decide( + context: Context, + client: Client, + ): Boolean = true +} + +internal class ConditionTest { + private fun dagWithSides(): Triple<DagDef, ConditionRef, Pair<TaskRef<Unit>, TaskRef<Unit>>> { + val dag = DagDef("d") + val load = dag.task<Unit>("load", NoopTask::class.java) + val reportEmpty = dag.task<Unit>("report_empty", NoopTask::class.java) + return Triple(dag, dag.If(HasRows::class.java), load to reportEmpty) + } + + @Test + @DisplayName("Should take the condition's task ID from its class when none is given") + fun shouldDeriveTaskIdFromClass() { + val (dag, condition, sides) = dagWithSides() + condition.Then(sides.first) + + assertEquals("hasRows", condition.id) + assertTrue("hasRows" in dag.tasks) + } + + @Test + @DisplayName("Should run each named side after the condition") + fun shouldRunEachSideAfterTheCondition() { + val (dag, condition, sides) = dagWithSides() + val (load, reportEmpty) = sides + condition.Then(load).Else(reportEmpty) + + val decider = dag.tasks.getValue("hasRows") + assertEquals(setOf(decider), load.def.upstreams) + assertEquals(setOf(decider), reportEmpty.def.upstreams) + } + + @Test + @DisplayName("Should mark only a deciding task as able to skip what runs after it") + fun shouldMarkDeciderAsSkipping() { + val (dag, condition, sides) = dagWithSides() + condition.Then(sides.first) + + val tasks = serializeDag(dag, "", ".")["tasks"] as List<*> + + @Suppress("UNCHECKED_CAST") + val byId = tasks.associate { task -> ((task as Map<String, Any?>)["__var"] as Map<String, Any?>).let { it["task_id"] to it } } + assertEquals(true, (byId["hasRows"] as Map<*, *>)["_can_skip_downstream"]) + assertFalse("_can_skip_downstream" in (byId["load"] as Map<*, *>)) + } + + @Test + @DisplayName("Should reject naming the same side twice") + fun shouldRejectNamingASideTwice() { + val (_, condition, sides) = dagWithSides() + condition.Then(sides.first) + + val error = assertThrows(IllegalArgumentException::class.java) { condition.Then(sides.second) } + + assertEquals("Condition 'hasRows' already runs 'load' on its Then side; name each side once", error.message) + } + + @Test + @DisplayName("Should reject a condition whose sides are the same task") + fun shouldRejectIdenticalSides() { + val (_, condition, sides) = dagWithSides() + condition.Then(sides.first) + + val error = assertThrows(IllegalArgumentException::class.java) { condition.Else(sides.first) } + + assertEquals( + "Condition 'hasRows' already runs 'load' on its other side, so the condition would decide nothing", + error.message, + ) + } + + @Test + @DisplayName("Should reject a side declared in another Dag") + fun shouldRejectSideFromAnotherDag() { + val (_, condition, _) = dagWithSides() + val other = DagDef("other").task<Unit>("load", NoopTask::class.java) + + val error = assertThrows(IllegalArgumentException::class.java) { condition.Then(other) } + + assertEquals( + "Condition 'hasRows' of Dag 'd' cannot run task 'load' of Dag 'other'; name a task of the same Dag", + error.message, + ) + } + + @Test + @DisplayName("Should reject registering a condition that names no task to run") + fun shouldRejectConditionWithoutThen() { + val (dag, _, _) = dagWithSides() + + val error = assertThrows(IllegalArgumentException::class.java) { Bundle().register(dag) } + + assertEquals( + "Condition 'hasRows' names no task to run when it holds; call Then(...), in Dag 'd'", + error.message, + ) + } + + @Test + @DisplayName("Should reject registering a decider class as an ordinary task") + fun shouldRejectDeciderRegisteredAsTask() { + val dag = DagDef("d") + dag.task<Unit>("has_rows", HasRows::class.java) + + val error = assertThrows(IllegalArgumentException::class.java) { Bundle().register(dag) } + + 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'", + error.message, + ) + } + + @Test + @DisplayName("Should reject a task that chooses which task runs but decides nothing") + fun shouldRejectChooserThatDecidesNothing() { + val dag = DagDef("d") + val task = dag.task<Boolean>("gate", NoopTask::class.java) + ConditionRef.of(task) + + val error = assertThrows(IllegalArgumentException::class.java) { Bundle().register(dag) } + + assertEquals( + "Task 'gate' chooses which task runs, but '${NoopTask::class.java.name}' decides nothing, in Dag 'd'", + error.message, + ) + } + + @Test + @DisplayName("Should reuse the condition when a task already declared as one is declared again") + fun shouldReuseConditionDeclaredTwice() { + val (_, condition, sides) = dagWithSides() + val again = ConditionRef.of(TaskRef(condition.nodes().single())) + + condition.Then(sides.first) + again.Else(sides.second) + + val error = assertThrows(IllegalArgumentException::class.java) { again.Then(sides.second) } + assertEquals( + "Condition 'hasRows' already runs 'load' on its Then side; name each side once", + error.message, + ) + } + + @Test + @DisplayName("Should reject naming a side after the Dag was registered") + fun shouldRejectNamingASideAfterRegistration() { + val (dag, condition, sides) = dagWithSides() + condition.Then(sides.first) + Bundle().register(dag) + + val error = assertThrows(IllegalArgumentException::class.java) { condition.Else(sides.second) } + + assertEquals( + "Condition 'hasRows' of Dag 'd' is already registered; name every case before the Dag is " + + "added to a Bundle", + error.message, + ) + } + + @Test + @DisplayName("Should declare a condition inside a task group under the group's prefix") + fun shouldDeclareConditionInGroup() { + val dag = DagDef("d") + val checks = dag.taskGroup("checks") + val load = checks.task<Unit>("load", NoopTask::class.java) + val condition = checks.If("has_rows", HasRows::class.java) + condition.Then(load) + + assertEquals("checks.has_rows", condition.id) + assertEquals(listOf("checks.load", "checks.has_rows"), 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 7d9aa937010..4c2f7f84e4e 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 @@ -34,6 +34,8 @@ import com.fasterxml.jackson.dataformat.yaml.YAMLMapper import org.apache.airflow.sdk.Arg import org.apache.airflow.sdk.Bundle import org.apache.airflow.sdk.Client +import org.apache.airflow.sdk.ConditionRef +import org.apache.airflow.sdk.ConditionTask import org.apache.airflow.sdk.Context import org.apache.airflow.sdk.DagDef import org.apache.airflow.sdk.Deps @@ -55,6 +57,13 @@ class ConformanceTask : Task { ) = Unit } +class ConformanceCondition : ConditionTask { + override fun decide( + context: Context, + client: Client, + ) = true +} + 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") @@ -80,16 +89,21 @@ private fun buildDag(case: JsonNode): DagDef { groups[groupId] = if (parentId.isEmpty()) dag.taskGroup(localId) else groups.getValue(parentId).taskGroup(localId) } - val tasks = linkedMapOf<String, TaskRef<Unit>>() + 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>>() 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 ref = if (groupId.isEmpty()) { - dag.task<Unit>(localId, ConformanceTask::class.java) + dag.task<Any?>(localId, definition) } else { - groups.getValue(groupId).task(localId, ConformanceTask::class.java) + groups.getValue(groupId).task<Any?>(localId, definition) } + if (!branch.isMissingNode) conditions += 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 @@ -103,6 +117,11 @@ 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())) } + } + case.path("order_edges").forEach { edge -> val node = { id: String -> groups[id] as Deps.Flow? ?: tasks.getValue(id) } node(edge[0].asText()).before(node(edge[1].asText())) @@ -110,6 +129,10 @@ private fun buildDag(case: JsonNode): DagDef { return dag } +/** The handle of a task declared as a decider, whose type argument no caller reads. */ +@Suppress("UNCHECKED_CAST") +private fun asCondition(ref: TaskRef<*>): ConditionRef = ConditionRef.of(ref as TaskRef<Boolean>) + /** Reads a YAML value as the Java type the config key takes. */ private fun toValue( table: Map<String, Field>, diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/CommTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/CommTest.kt index 7f161ed9b5c..1d9b5cc4ace 100644 --- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/CommTest.kt +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/CommTest.kt @@ -508,6 +508,33 @@ class CommsTest { Assertions.assertEquals(tiId.toString(), clearBody["ti_id"]) } + @Test + @DisplayName("skipDownstreamTasks sends the task IDs to skip and accepts the supervisor's empty response") + @Timeout(value = 30, unit = TimeUnit.SECONDS) + fun skipDownstreamTasksSendsTaskIds() { + val (body, failure) = + roundTrip(::emptyResponseFrame) { it.impl.skipDownstreamTasks(listOf("report_empty", "audit")) } + + Assertions.assertNull(failure, "skipDownstreamTasks should return normally on an empty response, got $failure") + Assertions.assertEquals("SkipDownstreamTasks", body["type"]) + Assertions.assertEquals(listOf("report_empty", "audit"), body["tasks"]) + } + + @Test + @DisplayName("skipDownstreamTasks sends nothing when there is nothing to skip") + @Timeout(value = 30, unit = TimeUnit.SECONDS) + fun skipDownstreamTasksSendsNothingForEmptyList() { + val toClient = ByteChannel(autoFlush = true) + val fromClient = ByteChannel(autoFlush = true) + val comm = CoordinatorComm(toClient, fromClient) + + CoordinatorClient(comm).skipDownstreamTasks(emptyList()) + + runBlocking { fromClient.flushAndClose() } + Assertions.assertTrue(fromClient.isClosedForRead, "no frame should have been written") + comm.close() + } + @Test @DisplayName("Should fail a pending call when the coordinator socket closes") @Timeout(value = 30, unit = TimeUnit.SECONDS) 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 5ce508b708d..a1ca3f7553d 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 @@ -21,6 +21,7 @@ package org.apache.airflow.sdk.execution 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.Task @@ -182,6 +183,77 @@ class TaskTest { Assertions.assertInstanceOf(SucceedTask::class.java, result) } + @Test + @DisplayName("Should push a condition's result and skip the side it did not take") + fun shouldSkipTheSideNotTaken() { + TestCondition.decision = true + val transport = RecordingTransport() + + val result = + runTask(conditionBundle(withElse = true), startupDetails(taskId = "gate"), Client(startupDetails("gate"), transport)) + + Assertions.assertInstanceOf(SucceedTask::class.java, result) + Assertions.assertEquals(listOf("report_empty"), transport.skipped) + Assertions.assertEquals( + listOf("skipmixin_key" to mapOf("skipped" to listOf("report_empty")), "return_value" to true), + transport.xComs, + ) + Assertions.assertEquals( + listOf("xcom:skipmixin_key", "skip:[report_empty]", "xcom:return_value"), + transport.events, + ) + } + + @Test + @DisplayName("Should skip the then side when a condition does not hold") + fun shouldSkipThenSideWhenConditionIsFalse() { + TestCondition.decision = false + val transport = RecordingTransport() + + runTask(conditionBundle(withElse = true), startupDetails(taskId = "gate"), Client(startupDetails("gate"), transport)) + + Assertions.assertEquals(listOf("load"), transport.skipped) + Assertions.assertEquals(false, transport.xComs.last().second) + Assertions.assertEquals(listOf("xcom:skipmixin_key", "skip:[load]", "xcom:return_value"), transport.events) + } + + @Test + @DisplayName("Should skip nothing when a one-sided condition holds") + fun shouldSkipNothingWhenOneSidedConditionHolds() { + TestCondition.decision = true + val transport = RecordingTransport() + + runTask(conditionBundle(withElse = false), startupDetails(taskId = "gate"), Client(startupDetails("gate"), transport)) + + Assertions.assertEquals(emptyList<String>(), transport.skipped) + Assertions.assertEquals(listOf("return_value" to true), transport.xComs) + Assertions.assertEquals(listOf("xcom:return_value"), transport.events) + } + + @Test + @DisplayName("Should fail a condition that throws, without pushing or skipping anything") + fun shouldFailConditionThatThrows() { + TestCondition.decision = null + val transport = RecordingTransport() + + val result = + runTask(conditionBundle(withElse = true), startupDetails(taskId = "gate"), Client(startupDetails("gate"), 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 conditionBundle(withElse: Boolean): Bundle { + val dag = DagDef("test_dag") + val load = dag.task<Unit>("load", SuccessTask::class.java) + val reportEmpty = dag.task<Unit>("report_empty", SuccessTask::class.java) + val condition = dag.If("gate", TestCondition::class.java).Then(load) + if (withElse) condition.Else(reportEmpty) + return Bundle(listOf(dag)) + } + private fun bundleWith( taskId: String, taskClass: Class<out Task>, @@ -271,9 +343,87 @@ class TaskTest { ): Unit = throw UnsupportedOperationException("not used in test") override fun clearTaskStateStore(tiId: UUID): Unit = throw UnsupportedOperationException("not used in test") + + override fun skipDownstreamTasks(taskIds: List<String>): Unit = throw UnsupportedOperationException("not used in test") }, ) + /** Decides what [decision] holds; a null one throws, standing for a condition body that fails. */ + class TestCondition : ConditionTask { + override fun decide( + context: Context, + client: Client, + ): Boolean = decision ?: throw IllegalStateException("boom") + + companion object { + var decision: Boolean? = true + } + } + + /** 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>>() + val skipped = mutableListOf<String>() + val events = mutableListOf<String>() + + override fun setXCom( + key: String, + value: Any, + dagId: String, + taskId: String, + runId: String, + mapIndex: Int, + ) { + xComs += key to value + events += "xcom:$key" + } + + override fun skipDownstreamTasks(taskIds: List<String>) { + skipped += taskIds + events += "skip:$taskIds" + } + + override fun getConnection(id: String) = throw UnsupportedOperationException("not used in test") + + override fun getVariable(key: String) = throw UnsupportedOperationException("not used in test") + + override fun setVariable( + key: String, + value: String, + description: String?, + ): Unit = throw UnsupportedOperationException("not used in test") + + override fun deleteVariable(key: String): Unit = throw UnsupportedOperationException("not used in test") + + override fun getXCom( + key: String, + dagId: String, + taskId: String, + runId: String, + mapIndex: Int?, + includePriorDates: Boolean, + ) = throw UnsupportedOperationException("not used in test") + + override fun getTaskStateStore( + tiId: UUID, + key: String, + ) = throw UnsupportedOperationException("not used in test") + + override fun setTaskStateStore( + tiId: UUID, + key: String, + value: Any, + expiresAt: OffsetDateTime?, + ): Unit = throw UnsupportedOperationException("not used in test") + + override fun deleteTaskStateStore( + tiId: UUID, + key: String, + ): Unit = throw UnsupportedOperationException("not used in test") + + override fun clearTaskStateStore(tiId: UUID): Unit = throw UnsupportedOperationException("not used in test") + } + class SuccessTask : Task { override fun execute( context: Context, diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/ArgValuesTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/ArgValuesTest.kt index b9fd480e10c..c8ca101267c 100644 --- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/ArgValuesTest.kt +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/ArgValuesTest.kt @@ -128,6 +128,8 @@ internal class ArgValuesTest { ): Unit = throw NotImplementedError() override fun clearTaskStateStore(tiId: UUID): Unit = throw NotImplementedError() + + override fun skipDownstreamTasks(taskIds: List<String>): Unit = throw NotImplementedError() }, ) 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 3a06f7c472b..59e36c182bd 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 @@ -21,6 +21,8 @@ package org.apache.airflow.sdk.internal import org.apache.airflow.sdk.Arg import org.apache.airflow.sdk.Client +import org.apache.airflow.sdk.ConditionRef +import org.apache.airflow.sdk.ConditionTask import org.apache.airflow.sdk.Context import org.apache.airflow.sdk.DagDef import org.apache.airflow.sdk.Deps @@ -47,6 +49,13 @@ private class NoopRefTask : Task { ) = Unit } +private class NoopCondition : ConditionTask { + override fun decide( + context: Context, + client: Client, + ) = true +} + internal class RefsTest { @Test @DisplayName("Should register the task, record inputs, and wire handle edges") @@ -203,4 +212,24 @@ internal class RefsTest { assertEquals("Dag 'd' has no task group 'staging'", error.message) } + + @Test + @DisplayName("Should let a view name a condition's sides in separate statements") + fun shouldNameConditionSidesAcrossStatements() { + val dag = DagDef("d") + Refs.record(dag, listOf("has_rows", "load", "report_empty"), emptyList()) { + fun hasRows() = ConditionRef.of(Refs.node<Boolean>("", TaskDef("has_rows", NoopCondition::class.java))) + hasRows().Then(Refs.node<Unit>("", TaskDef("load", NoopRefTask::class.java))) + hasRows().Else(Refs.node<Unit>("", TaskDef("report_empty", NoopRefTask::class.java))) + } + + assertEquals( + listOf("load", "report_empty"), + dag.tasks + .getValue("has_rows") + .decider!! + .cases + .map { it.id }, + ) + } }
