mirror of
https://github.com/apache/airflow.git
synced 2026-11-04 07:12:48 +00:00
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
This commit is contained in:
1 parent
2bda192870
commit
d4013cc2b5
27 files changed
+1274
-36
No files matched your search
+1
-1
@@ -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 † | ✗ | – | |
|
||||
|
||||
@@ -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:
|
||||
|
||||
+14
-1
@@ -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());
|
||||
|
||||
+20
-1
@@ -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;
|
||||
}
|
||||
}
|
||||
@@ -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(
|
||||
|
||||
@@ -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() {
|
||||
|
||||
@@ -56,7 +56,7 @@ if __name__ == "__main__":
|
||||
"--sdk",
|
||||
"java",
|
||||
"--supports",
|
||||
"literal_inputs",
|
||||
"literal_inputs,branch",
|
||||
"--",
|
||||
*serializer,
|
||||
]
|
||||
|
||||
@@ -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`.
|
||||
| *
|
||||
|
||||
@@ -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 † | ✗ | – | |
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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"
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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.
|
||||
*
|
||||
|
||||
@@ -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"
|
||||
}
|
||||
}
|
||||
@@ -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.
|
||||
*
|
||||
|
||||
@@ -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) }
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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() }
|
||||
@@ -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 =
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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())
|
||||
}
|
||||
}
|
||||
@@ -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>,
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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()
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@@ -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 },
|
||||
)
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user