diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 5c04a604289b1..9d7b2b197c9ee 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -444,7 +444,7 @@ repos: ^java-sdk/scripts/ci/prek/check_serialization_conformance\.py$| ^java-sdk/sdk/schema/dag-schema\.json$| ^java-sdk/sdk/build\.gradle\.kts$| - ^java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/(Switch|Bundle|Condition|DagDef|Decider|Deps|Arg|TaskGroupRef)\.kt$| + ^java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/(Switch|Bundle|Condition|DagDef|Decider|Deps|Arg|TaskGroupRef|TriggerDagRun)\.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 9b9653a497474..fda035f80dcc6 100644 --- a/java-sdk/README.md +++ b/java-sdk/README.md @@ -620,8 +620,8 @@ prek hook regenerate it. | state: `success` | MUST | ✓ | 3.3 | | | state: `failed` | MUST | ✓ | 3.3 | | | state: `up_for_retry` | MUST | ✓ | 3.3 | RetryTask | -| state: `skipped` | SHOULD | ✗ | – | runtime does not emit TaskState skipped yet | -| state: `deferred` | MAY | ✗ | – | runtime does not emit DeferTask yet | +| state: `skipped` | SHOULD | ✓ | 3.4 | a TriggerDagRun task that skips an existing run | +| state: `deferred` | MAY | ✓ | 3.4 | a TriggerDagRun task that waits with deferrable | | state: `up_for_reschedule` | MAY | ✗ | – | runtime does not emit RescheduleTask yet | | state: `awaiting_input` | MAY | ✗ | – | runtime does not emit AwaitInputTask yet | | state: `removed` | MAY | ✓ | 3.3 | | diff --git a/java-sdk/capabilities.yaml b/java-sdk/capabilities.yaml index 5a2ed1f5ae465..84ee13a9dda3b 100644 --- a/java-sdk/capabilities.yaml +++ b/java-sdk/capabilities.yaml @@ -23,8 +23,8 @@ min_airflow_version: "3.3" # JAR manifest. The render hook fails if the two disagree. supervisor_schema_version: "2026-10-30" -# The runtime terminates a task with SucceedTask, RetryTask, or TaskState (failed/removed); it does -# not yet emit skipped, DeferTask, RescheduleTask, or AwaitInputTask. +# The runtime terminates a task with SucceedTask, RetryTask, DeferTask, or TaskState +# (failed/skipped/removed); it does not yet emit RescheduleTask or AwaitInputTask. states: success: supported: true @@ -37,11 +37,13 @@ states: since: "3.3" note: "RetryTask" skipped: - supported: false - note: "runtime does not emit TaskState skipped yet" + supported: true + since: "3.4" + note: "a TriggerDagRun task that skips an existing run" deferred: - supported: false - note: "runtime does not emit DeferTask yet" + supported: true + since: "3.4" + note: "a TriggerDagRun task that waits with deferrable" up_for_reschedule: supported: false note: "runtime does not emit RescheduleTask yet" diff --git a/java-sdk/example/src/java/org/apache/airflow/example/ExampleBundleBuilder.java b/java-sdk/example/src/java/org/apache/airflow/example/ExampleBundleBuilder.java index 76175db8a6d5c..30da2aa8e116f 100644 --- a/java-sdk/example/src/java/org/apache/airflow/example/ExampleBundleBuilder.java +++ b/java-sdk/example/src/java/org/apache/airflow/example/ExampleBundleBuilder.java @@ -30,7 +30,8 @@ public static Bundle build() { .register(AnnotationExample.class) .register(XComCastingExample.class) .register(org.apache.airflow.example.nativedag.AnnotationExample.class) - .register(org.apache.airflow.example.nativedag.InterfaceExample.build()); + .register(org.apache.airflow.example.nativedag.InterfaceExample.build()) + .register(org.apache.airflow.example.nativedag.TargetExample.build()); } public static void main(String[] args) { 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 16cdc555452a9..0bb55ceacd0f5 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 @@ -85,6 +85,13 @@ public Class pickReport(long transformed) { : AnnotationExampleBuilder.ReportShort.class; } + // A task that starts a run of another Dag. The method runs when this Dag is + // built, not when the task runs, so it takes no arguments. + @Builder.Task(id = "trigger_downstream") + public TriggerDagRun triggerDownstream() { + return new TriggerDagRun("java_native_target_example").config("note", "from the annotation example"); + } + // A task group: everything it declares is prefixed with its id, so this is // the task "checks.audit". @Builder.TaskGroup(id = "checks") @@ -102,8 +109,10 @@ static class Wiring implements AnnotationExampleDeps { void depends() { var extracted = extract(); var transformed = transform(extracted, lit(1.5)); - hasRows(transformed).Then(load(transformed)).Else(loadEmpty()); + var loaded = load(transformed); + hasRows(transformed).Then(loaded).Else(loadEmpty()); pickReport(transformed).Case(reportLong()).Case(reportShort()); + loaded.before(triggerDownstream()); // 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 b8099931b7aae..fcf095ee72ec7 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 @@ -114,12 +114,18 @@ public static DagDef build() { var reportLong = dag.task("report_long", ReportLong.class); var reportShort = dag.task("report_short", ReportShort.class); + // A task that starts a run of another Dag; it runs no Java code. + var trigger = + dag.task( + "trigger_downstream", + new TriggerDagRun("java_native_target_example").config("wait_for_completion", false)); transform.after(extract); // With no task id given, a decider takes one from its class: "hasRows" and // "pickReport". dag.If(HasRows.class).after(transform).Then(load).Else(loadEmpty); dag.Switch(PickReport.class).after(transform).Case(reportLong).Case(reportShort); + load.before(trigger); return dag; } } diff --git a/java-sdk/example/src/java/org/apache/airflow/example/nativedag/TargetExample.java b/java-sdk/example/src/java/org/apache/airflow/example/nativedag/TargetExample.java new file mode 100644 index 0000000000000..cf8d451429545 --- /dev/null +++ b/java-sdk/example/src/java/org/apache/airflow/example/nativedag/TargetExample.java @@ -0,0 +1,49 @@ +/* + * 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. + */ + +// "native" is a Java keyword, so the native-Dag examples live in "nativedag". +package org.apache.airflow.example.nativedag; + +import static java.lang.System.Logger.Level.INFO; + +import java.util.List; +import org.apache.airflow.sdk.*; + +// The Dag the other two native examples trigger. It triggers nothing itself, +// so the examples cannot trigger each other in a loop. +public class TargetExample { + private static final System.Logger log = System.getLogger(TargetExample.class.getName()); + + public static class Receive implements Task { + @Override + public void execute(Context context, Client client) { + log.log(INFO, "Triggered by another Dag"); + } + } + + public static DagDef build() { + var dag = + new DagDef("java_native_target_example") + .config("description", "Pure-Java Dag that the other native examples trigger") + .config("catchup", false) + .config("tags", List.of("example", "java-sdk")); + dag.task("receive", Receive.class); + 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 97b810358a0d2..e101c0e8572ae 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 @@ -278,7 +278,7 @@ class BuilderProcessor : AbstractProcessor() { ) builderClass.addMethod(buildMethod.build()) - declarations.forEach { builderClass.addType(buildTask(it)) } + declarations.filterNot { it.kind == TaskKind.TRIGGER }.forEach { builderClass.addType(buildTask(it)) } return builderClass.build() } @@ -366,8 +366,11 @@ class BuilderProcessor : AbstractProcessor() { .methodBuilder(decl.method.simpleName.toString()) .addModifiers(Modifier.PUBLIC, Modifier.DEFAULT) .returns( - decl.kind.refType - ?: ParameterizedTypeName.get(TASK_HANDLE_TYPE, TypeName.get(decl.method.returnType).boxIfPossible()), + when { + decl.kind.refType != null -> decl.kind.refType + decl.kind == TaskKind.TRIGGER -> ParameterizedTypeName.get(TASK_HANDLE_TYPE, VOID_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)) @@ -411,7 +414,13 @@ class BuilderProcessor : AbstractProcessor() { val taskDef = CodeBlock .builder() - .add($$"new $T($S, $L.class)", TASK_DEF_TYPE, decl.id, classRef) + .apply { + if (decl.kind == TaskKind.TRIGGER) { + add($$"new $T($S, new $T().$L())", TASK_DEF_TYPE, decl.id, ClassName.get(decl.owner), decl.method.simpleName) + } else { + add($$"new $T($S, $L.class)", TASK_DEF_TYPE, decl.id, classRef) + } + } explicitConfig(decl.method, decl.kind.annotation, TASK_STRUCTURAL_ATTRIBUTES, SchemaFields.TASK).forEach { (key, value) -> taskDef.add($$".config($S, $L)", key, value) } @@ -438,17 +447,33 @@ class BuilderProcessor : AbstractProcessor() { for (inner in el.enclosedElements) { if (inner !is ExecutableElement) continue val declared = - TaskKind.entries.filter { kind -> inner.annotationMirrors.any { it.names(kind.annotation) } } + TaskKind.declaring.filter { kind -> inner.annotationMirrors.any { it.names(kind.annotation) } } if (declared.isEmpty()) continue - val kind = + val declaredKind = declared.singleOrNull() ?: throw IllegalArgumentException( "Method '${inner.simpleName}' carries ${declared.joinToString { it.spelling }}; a task is " + "declared by one of them alone", ) + // A task method that hands back a TriggerDagRun declares what to trigger + // rather than a body to run, so it is read when the Dag is built. + val triggers = + declaredKind == TaskKind.TASK && with(processingEnv) { isType(inner.returnType, TRIGGER_TYPE) } + val kind = if (triggers) TaskKind.TRIGGER else declaredKind val annotated = declaredId(inner, kind) if (inner.isVarArgs) throw IllegalArgumentException("Cannot create task from vararg function ${inner.simpleName}") checkDeciderReturn(kind, inner) + if (kind == TaskKind.TRIGGER) { + require(inner.parameters.isEmpty()) { + "@Builder.Task method '${inner.simpleName}' returns a TriggerDagRun, so it runs when the Dag " + + "is built rather than when the task runs; it takes no parameters" + } + val checked = inner.thrownTypes.filterNot { isUnchecked(it) } + require(checked.isEmpty()) { + "@Builder.Task method '${inner.simpleName}' returns a TriggerDagRun, so it runs when the Dag " + + "is built, where a checked exception cannot be thrown; it must not throw: ${checked.joinToString()}" + } + } val localId = annotated.ifBlank { inner.simpleName.toString() } require(tasks.none { it.method.simpleName.contentEquals(inner.simpleName) }) { "Class ${el.simpleName} overloads task method '${inner.simpleName}'; a method's name is the " + @@ -572,7 +597,8 @@ class BuilderProcessor : AbstractProcessor() { } } val byClassName = mutableMapOf() - declarations.forEach { decl -> + // A task that triggers a Dag run has no generated class to collide with. + declarations.filterNot { it.kind == TaskKind.TRIGGER }.forEach { decl -> byClassName.put(decl.className, decl)?.let { first -> throw IllegalArgumentException( "Task methods '${first.id}' and '${decl.id}' both generate the task class " + @@ -694,6 +720,8 @@ class BuilderProcessor : AbstractProcessor() { parseTemporal(field, text) { OffsetDateTime.parse(text) } CodeBlock.of($$"$T.parse($S)", ClassName.get(OffsetDateTime::class.java), text) } + FieldType.DAG_RUN_STATES, FieldType.JSON_OBJECT -> + error("Field '${field.key}' is not an annotation attribute") } private fun parseTemporal( @@ -718,7 +746,7 @@ class BuilderProcessor : AbstractProcessor() { ) { val returns = method.returnType when (kind) { - TaskKind.TASK -> return + TaskKind.TASK, TaskKind.TRIGGER -> return TaskKind.CONDITION -> require(returns.kind == TypeKind.BOOLEAN || with(processingEnv) { isType(returns, BOXED_BOOLEAN_TYPE) }) { "@Builder.If method '${method.simpleName}' returns $returns, but a condition returns boolean: " + @@ -760,7 +788,7 @@ class BuilderProcessor : AbstractProcessor() { .addModifiers(Modifier.PUBLIC) .returns( when (decl.kind) { - TaskKind.TASK -> TypeName.VOID + TaskKind.TASK, TaskKind.TRIGGER -> TypeName.VOID TaskKind.CONDITION -> TypeName.BOOLEAN TaskKind.SWITCH -> TASK_CLASS_TYPE }, @@ -1008,6 +1036,8 @@ private val SWITCH_TASK_TYPE = ClassName.get(SwitchTask::class.java) private val SWITCH_REF_TYPE = ClassName.get(SwitchRef::class.java) private val TASK_CLASS_TYPE = ParameterizedTypeName.get(ClassName.get(Class::class.java), WildcardTypeName.subtypeOf(TASK_TYPE)) +private val TRIGGER_TYPE = ClassName.get(TriggerDagRun::class.java) +private val VOID_TYPE = ClassName.get("java.lang", "Void") 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") @@ -1029,8 +1059,18 @@ private enum class TaskKind( val refType: ClassName?, ) { TASK("org.apache.airflow.sdk.Builder.Task", "@Builder.Task", "execute", TASK_TYPE, null), + TRIGGER("org.apache.airflow.sdk.Builder.Task", "@Builder.Task", "execute", TASK_TYPE, null), CONDITION("org.apache.airflow.sdk.Builder.If", "@Builder.If", "decide", CONDITION_TASK_TYPE, CONDITION_REF_TYPE), SWITCH("org.apache.airflow.sdk.Builder.Switch", "@Builder.Switch", "choose", SWITCH_TASK_TYPE, SWITCH_REF_TYPE), + ; + + companion object { + /** + * The kinds an annotation declares. TRIGGER is left out: it shares + * [TASK]'s annotation and is told apart by the method's return type. + */ + val declaring = listOf(TASK, CONDITION, SWITCH) + } } /** Whether this annotation is the one [name] qualifies. */ diff --git a/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt b/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt index 2473ef9561aa6..5d8c89a058b06 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 @@ -2452,6 +2452,209 @@ class BuilderTest { ) } + @Test + @DisplayName("read a task that triggers a Dag run when the Dag is built") + fun generateTriggerDagRunTask() { + val compilation = + compile( + """ + package org.apache.airflow.example; + + import org.apache.airflow.sdk.Builder; + import org.apache.airflow.sdk.TriggerDagRun; + + @Builder.Dag(id = "etl") + public class TestExample { + @Builder.Task(id = "trigger_downstream", retries = 2) + public TriggerDagRun triggerDownstream() { + return new TriggerDagRun("downstream_etl").config("wait_for_completion", true); + } + + @Builder.Deps + static class Wiring implements TestExampleDeps { + void depends() { + triggerDownstream(); + } + } + } + """, + ) + + assertThat(compilation).succeeded() + assertThat(compilation) + .generatedSourceFile("org.apache.airflow.example.TestExampleBuilder") + .hasSourceEquivalentTo( + "org.apache.airflow.example.TestExampleBuilder", + """ + package org.apache.airflow.example; + + import java.util.List; + import org.apache.airflow.sdk.DagDef; + 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("trigger_downstream"), List.of(), new TestExample.Wiring()::depends); + } + } + """, + ) + assertThat(compilation) + .generatedSourceFile("org.apache.airflow.example.TestExampleDeps") + .hasSourceEquivalentTo( + "org.apache.airflow.example.TestExampleDeps", + """ + package org.apache.airflow.example; + + import java.lang.Void; + import org.apache.airflow.sdk.Deps; + import org.apache.airflow.sdk.TaskDef; + import org.apache.airflow.sdk.TaskRef; + import org.apache.airflow.sdk.internal.Refs; + + public interface TestExampleDeps extends Deps { + default TaskRef triggerDownstream() { + return Refs.node("", new TaskDef("trigger_downstream", new TestExample().triggerDownstream()).config("retries", 2)); + } + } + """, + ) + } + + @Test + @DisplayName("build a Dag that reuses one task handle as a condition side and as an upstream") + fun buildADagThatReusesATaskHandle() { + val compilation = + compile( + """ + package org.apache.airflow.example; + + import org.apache.airflow.sdk.Builder; + import org.apache.airflow.sdk.TriggerDagRun; + + @Builder.Dag(id = "etl") + public class TestExample { + @Builder.Task + public long extract() { + return 42L; + } + + @Builder.Task + public void load(long rows) {} + + @Builder.Task + public void loadEmpty() {} + + @Builder.If(id = "has_rows") + public boolean hasRows(long rows) { + return rows > 0; + } + + @Builder.Task(id = "trigger_downstream") + public TriggerDagRun triggerDownstream() { + return new TriggerDagRun("downstream_etl"); + } + + @Builder.Deps + static class Wiring implements TestExampleDeps { + void depends() { + var extracted = extract(); + var loaded = load(extracted); + hasRows(extracted).Then(loaded).Else(loadEmpty()); + loaded.before(triggerDownstream()); + } + } + } + """, + ) + + assertThat(compilation).succeeded() + // The wiring only runs when the Dag is built, so load it and build it. + val classes = + compilation + .generatedFiles() + .filter { it.kind == JavaFileObject.Kind.CLASS } + .associate { + it.name + .removePrefix("/CLASS_OUTPUT/") + .removeSuffix(".class") + .replace('/', '.') to it.openInputStream().readBytes() + } + val loader = + object : ClassLoader(javaClass.classLoader) { + override fun findClass(name: String): Class<*> { + val bytes = classes[name] ?: throw ClassNotFoundException(name) + return defineClass(name, bytes, 0, bytes.size) + } + } + + Bundle().register(loader.loadClass("org.apache.airflow.example.TestExample")) + } + + @Test + @DisplayName("reject a task that triggers a Dag run and takes parameters") + fun rejectTriggerDagRunWithParameters() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + import org.apache.airflow.sdk.TriggerDagRun; + @Builder.Dag + public class TestExample { + @Builder.Task + public TriggerDagRun triggerDownstream(long rows) { + return new TriggerDagRun("downstream_etl"); + } + + @Builder.Deps + static class Wiring implements TestExampleDeps { + void depends() {} + } + } + """, + ) + + assertThat(compilation).failed() + assertThat(compilation).hadErrorContaining( + "@Builder.Task method 'triggerDownstream' returns a TriggerDagRun, so it runs when the Dag is " + + "built rather than when the task runs; it takes no parameters", + ) + } + + @Test + @DisplayName("reject a task that triggers a Dag run and throws a checked exception") + fun rejectTriggerDagRunThrowingCheckedException() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + import org.apache.airflow.sdk.TriggerDagRun; + @Builder.Dag + public class TestExample { + @Builder.Task + public TriggerDagRun triggerDownstream() throws java.io.IOException { + return new TriggerDagRun("downstream_etl"); + } + + @Builder.Deps + static class Wiring implements TestExampleDeps { + void depends() {} + } + } + """, + ) + + assertThat(compilation).failed() + assertThat(compilation).hadErrorContaining( + "@Builder.Task method 'triggerDownstream' returns a TriggerDagRun, so it runs when the Dag is " + + "built, where a checked exception cannot be thrown; it must not throw: java.io.IOException", + ) + } + @Test @DisplayName("reject a switch that does not return the class of a task") fun rejectNonTaskClassSwitch() { diff --git a/java-sdk/scripts/ci/prek/check_serialization_conformance.py b/java-sdk/scripts/ci/prek/check_serialization_conformance.py index 01a16e3421987..677ffcdb2177d 100755 --- a/java-sdk/scripts/ci/prek/check_serialization_conformance.py +++ b/java-sdk/scripts/ci/prek/check_serialization_conformance.py @@ -56,7 +56,7 @@ "--sdk", "java", "--supports", - "literal_inputs,branch,switch", + "literal_inputs,branch,switch,trigger_dag_run", "--", *serializer, ] diff --git a/java-sdk/sdk/module.md b/java-sdk/sdk/module.md index f03b936d7b9b4..2cc58f52b4d1c 100644 --- a/java-sdk/sdk/module.md +++ b/java-sdk/sdk/module.md @@ -37,8 +37,8 @@ meaning of each dimension is defined in the | state: `success` | MUST | ✓ | 3.3 | | | state: `failed` | MUST | ✓ | 3.3 | | | state: `up_for_retry` | MUST | ✓ | 3.3 | RetryTask | -| state: `skipped` | SHOULD | ✗ | – | runtime does not emit TaskState skipped yet | -| state: `deferred` | MAY | ✗ | – | runtime does not emit DeferTask yet | +| state: `skipped` | SHOULD | ✓ | 3.4 | a TriggerDagRun task that skips an existing run | +| state: `deferred` | MAY | ✓ | 3.4 | a TriggerDagRun task that waits with deferrable | | state: `up_for_reschedule` | MAY | ✗ | – | runtime does not emit RescheduleTask yet | | state: `awaiting_input` | MAY | ✗ | – | runtime does not emit AwaitInputTask yet | | state: `removed` | MAY | ✓ | 3.3 | | 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 9f79be81a6c4a..37015ca91978f 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 @@ -160,6 +160,31 @@ class DagDef( return this } + /** + * Declares a task that starts a run of another Dag. + * + * ```java + * dag.task("trigger_downstream", new TriggerDagRun("downstream_etl")); + * ``` + * + * The task runs no Java code and takes no arguments, so it has no ID to + * derive and names one here. + * + * @param id Task identifier, unique within this Dag. + * @param trigger What to trigger, and how. + * @return The handle representing this task. + * @throws IllegalArgumentException if a task with the same ID is already + * registered. + */ + fun task( + id: String, + trigger: TriggerDagRun, + ): TaskRef { + val def = TaskDef(id, trigger) + addTask(def) + return TaskRef(def) + } + /** * Declares a task whose boolean picks one of two tasks; the other is * skipped. @@ -404,6 +429,21 @@ class TaskDef( val id: String, val definition: Class, ) : Endpoint { + /** + * Declares a task that starts a run of another Dag instead of running Java + * code. + * + * [DagDef.task] is the spelling user code uses; this is what a generated + * wiring view calls. + * + * @param id Task identifier, unique within a [DagDef]. + * @param trigger What to trigger, and how; settings made on it later do not + * reach this task. + */ + constructor(id: String, trigger: TriggerDagRun) : this(id, TriggerDagRunPlaceholder::class.java) { + this.trigger = trigger.snapshot() + } + init { validateTaskInput(definition) } @@ -419,6 +459,9 @@ class TaskDef( /** What this task decides to run, for a condition or a switch; null otherwise. */ internal var decider: DeciderDef? = null + /** The Dag run this task starts, for a task declared from a [TriggerDagRun]; null otherwise. */ + internal var trigger: TriggerDagRun? = null + /** * Sets one task-level configuration value. * diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Switch.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Switch.kt index a668db89609a0..5bbd93b4da13c 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Switch.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Switch.kt @@ -134,7 +134,8 @@ class SwitchRef private constructor( * @param task Task of the same Dag. * @return This switch, for chaining. * @throws IllegalArgumentException if the task belongs to another Dag, is - * already a case of this switch, or runs the same class as one. + * already a case of this switch, triggers a Dag run, or runs the same class + * as one. */ @Suppress("ktlint:standard:function-naming") fun Case(task: TaskRef<*>): SwitchRef { @@ -191,6 +192,11 @@ internal class SwitchDef : DeciderDef { require(options.none { it === case }) { "Switch '${decider.id}' already chooses between '${case.id}' and others; name each case once" } + require(case.trigger == null) { + "Switch '${decider.id}' cannot choose '${case.id}': a switch names its case by its class, and a " + + "task that triggers a Dag run has none. Put the trigger downstream of an ordinary task and " + + "let the switch choose that task, since Airflow skips the downstream of a skipped task" + } // A switch names its case by class, so two cases sharing one would be // indistinguishable. options.firstOrNull { it.definition == case.definition }?.let { first -> 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 60d43f84fcd89..363f13072cd13 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 @@ -76,6 +76,25 @@ class TaskGroupRef internal constructor( return TaskRef(def) } + /** + * Declares a task in this group that starts a run of another Dag, as + * [DagDef.task] does for the Dag. + * + * @param id Task ID within this group; the task's ID is `.`. + * @param trigger What to trigger, and how. + * @return The handle representing this task. + * @throws IllegalArgumentException if the Dag already has a task or task + * group with the resulting ID. + */ + fun task( + id: String, + trigger: TriggerDagRun, + ): TaskRef { + val def = TaskDef(qualify(id), trigger) + adopt(def) + return TaskRef(def) + } + /** * Declares a condition in this group, as [DagDef.If] does for the Dag. * diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/TriggerDagRun.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/TriggerDagRun.kt new file mode 100644 index 0000000000000..c78d1a9fd5fc2 --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/TriggerDagRun.kt @@ -0,0 +1,189 @@ +/* + * 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.internal.Field +import org.apache.airflow.sdk.internal.FieldType +import org.apache.airflow.sdk.internal.checkConfigValue +import java.time.Duration +import java.time.Instant +import java.time.OffsetDateTime + +/** + * Settings named as `TriggerDagRunOperator` names its parameters. The Dag ID + * is the constructor argument, so `trigger_dag_id` is not among them. + */ +private val TRIGGER_FIELDS: Map = + listOf( + Field("trigger_run_id", "triggerRunId", FieldType.STRING, null), + Field("conf", "conf", FieldType.JSON_OBJECT, null), + Field("logical_date", "logicalDate", FieldType.DATETIME, null), + Field("run_after", "runAfter", FieldType.DATETIME, null), + Field("reset_dag_run", "resetDagRun", FieldType.BOOLEAN, null), + Field("wait_for_completion", "waitForCompletion", FieldType.BOOLEAN, null), + Field("poke_interval", "pokeInterval", FieldType.TIMEDELTA, null), + Field("allowed_states", "allowedStates", FieldType.DAG_RUN_STATES, null), + Field("failed_states", "failedStates", FieldType.DAG_RUN_STATES, null), + Field("skip_when_already_exists", "skipWhenAlreadyExists", FieldType.BOOLEAN, null), + Field("fail_when_dag_is_paused", "failWhenDagIsPaused", FieldType.BOOLEAN, null), + Field("note", "note", FieldType.STRING, null), + Field("deferrable", "deferrable", FieldType.BOOLEAN, null), + ).associateBy { it.key } + +/** + * A task that starts a run of another Dag, in place of a task class. + * + * Pass one to [DagDef.task] and the Java runtime runs it as Airflow's + * `TriggerDagRunOperator` does, so the Dag needs no Python worker for it: + * + * ```java + * dag.task("trigger_downstream", + * new TriggerDagRun("downstream_etl") + * .config("wait_for_completion", true) + * .config("conf", Map.of("rows", 2))); + * ``` + * + * In an annotated Dag, a `@Builder.Task` method returns one instead; the + * method runs when the Dag is built, not when the task runs. + * + * The task runs no Java code and takes no arguments. It pushes the triggered + * run's ID, and the link the "Triggered DAG" extra link reads. It renders no + * templates, so a value such as `"{{ ds }}"` reaches the new run unchanged. + * + * @param dagId `trigger_dag_id`: the Dag to trigger. + * @throws IllegalArgumentException if [dagId] is empty. + */ +class TriggerDagRun( + val dagId: String, +) { + init { + require(dagId.isNotEmpty()) { "TriggerDagRun needs the ID of the Dag to trigger" } + } + + internal val settings = linkedMapOf() + + /** + * Sets one `TriggerDagRunOperator` setting, named as Python names the + * parameter. + * + * | Key | Value | + * | --- | --- | + * | `trigger_run_id` | `String`; generated when unset, and for a run with no logical date from `run_after` plus a random eight-character suffix | + * | `conf` | `Map` with string keys and JSON values | + * | `logical_date` | [OffsetDateTime] or [Instant]; the trigger time when neither this nor `run_after` is set, and none when only `run_after` is | + * | `run_after` | [OffsetDateTime] or [Instant] | + * | `reset_dag_run` | `Boolean`: clear a run with the same ID instead of failing | + * | `wait_for_completion` | `Boolean`: hold this task open until the run finishes | + * | `poke_interval` | [Duration], whole seconds; 60 seconds when unset | + * | `allowed_states` | states that count as success; `success` when unset | + * | `failed_states` | states that count as failure; `failed` when unset | + * | `skip_when_already_exists` | `Boolean`: skip rather than fail when the run exists | + * | `fail_when_dag_is_paused` | `Boolean`: fail rather than trigger a paused Dag | + * | `note` | `String` recorded against the triggered run | + * | `deferrable` | `Boolean`: while waiting, defer instead of holding the worker | + * + * A `deferrable` left unset follows Airflow's `[operators] default_deferrable`. + * + * @param key One of the keys above. + * @param value Value of the shape that key takes. + * @return This task, for chaining. + * @throws IllegalArgumentException if the key is unknown or the value does + * not match it. + */ + fun config( + key: String, + value: Any?, + ): TriggerDagRun { + require(key != "trigger_dag_id") { + "The Dag to trigger is the TriggerDagRun argument, not the config key 'trigger_dag_id'" + } + val checked = checkConfigValue("TriggerDagRun", TRIGGER_FIELDS, key, value) + settings[key] = + when (key) { + "poke_interval" -> { + val duration = checked as Duration + // Python's poke_interval is a number of seconds, and the serialized Dag carries it as one. + require(!duration.isNegative && duration.nano == 0) { + "Value for TriggerDagRun config key '$key' is $duration; it must be a whole, " + + "non-negative number of seconds" + } + duration + } + "conf" -> jsonValue(key, checked)!! + // Python generates a run ID for an empty one, which the runtime does not. + "trigger_run_id", "note" -> { + require((checked as String).isNotEmpty()) { + "Value for TriggerDagRun config key '$key' must not be empty" + } + checked + } + else -> checked + } + return this + } + + /** A copy that later calls on this builder cannot change. */ + internal fun snapshot(): TriggerDagRun = TriggerDagRun(dagId).also { it.settings.putAll(settings) } +} + +/** [value] as plain JSON, as a fresh structure, rejecting anything that has no JSON form. */ +private fun jsonValue( + key: String, + value: Any?, +): Any? = + when (value) { + null, is String, is Boolean, is Int, is Long, is Short, is Byte -> value + is Double -> + value.also { + require(it.isFinite()) { + "Value for TriggerDagRun config key '$key' must hold only finite numbers, got: $it" + } + } + is Float -> jsonValue(key, value.toDouble()) + is Collection<*> -> value.map { jsonValue(key, it) } + is Array<*> -> value.map { jsonValue(key, it) } + is Map<*, *> -> + value.entries.associateTo(linkedMapOf()) { (entryKey, entry) -> + require(entryKey is String) { + "Value for TriggerDagRun config key '$key' must have String keys, got: ${entryKey?.javaClass?.name}" + } + entryKey to jsonValue(key, entry) + } + else -> + throw IllegalArgumentException( + "Value for TriggerDagRun config key '$key' must hold only JSON values (null, String, " + + "Boolean, Number, List, Map), got: ${value.javaClass.name}", + ) + } + +/** + * Stands in for the class of a task that triggers a Dag run. The runtime runs + * such a task itself and never instantiates this, so it exists only so that + * every [TaskDef] names a class. + */ +internal class TriggerDagRunPlaceholder : Task { + override fun execute( + context: Context, + client: Client, + ): Unit = + throw IllegalStateException( + "A task that triggers a Dag run is run by the SDK, so its placeholder is never executed", + ) +} 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 f3174e77e2ade..317dac7ebb0a6 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 @@ -22,10 +22,14 @@ package org.apache.airflow.sdk.execution import kotlinx.coroutines.runBlocking import org.apache.airflow.sdk.execution.comm.ClearTaskStateStore import org.apache.airflow.sdk.execution.comm.ConnectionResult +import org.apache.airflow.sdk.execution.comm.DagResult +import org.apache.airflow.sdk.execution.comm.DagRunStateResult import org.apache.airflow.sdk.execution.comm.DeleteTaskStateStore import org.apache.airflow.sdk.execution.comm.DeleteVariable import org.apache.airflow.sdk.execution.comm.ErrorResponse import org.apache.airflow.sdk.execution.comm.GetConnection +import org.apache.airflow.sdk.execution.comm.GetDag +import org.apache.airflow.sdk.execution.comm.GetDagRunState import org.apache.airflow.sdk.execution.comm.GetTaskStateStore import org.apache.airflow.sdk.execution.comm.GetVariable import org.apache.airflow.sdk.execution.comm.GetXCom @@ -35,6 +39,7 @@ 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.TriggerDagRun import org.apache.airflow.sdk.execution.comm.VariableResult import org.apache.airflow.sdk.execution.comm.XComResult import java.time.OffsetDateTime @@ -110,6 +115,32 @@ interface Client { * succeeded or failed alone. */ fun skipDownstreamTasks(taskIds: List) + + /** + * Starts a run of another Dag, as a task declared from a + * [org.apache.airflow.sdk.TriggerDagRun] does. + * + * @return True when a run with that ID already exists, which the task + * decides what to do about, rather than an error. + */ + fun triggerDagRun( + dagId: String, + runId: String, + logicalDate: OffsetDateTime?, + runAfter: OffsetDateTime?, + conf: Map?, + resetDagRun: Boolean, + note: String?, + ): Boolean + + /** State of one Dag run, which a task that waits polls. */ + fun getDagRunState( + dagId: String, + runId: String, + ): String + + /** Whether the Dag is paused, which a task may refuse to trigger. */ + fun isDagPaused(dagId: String): Boolean } /** @@ -245,4 +276,47 @@ class CoordinatorClient( val message = SkipDownstreamTasks().also { it.tasks = ArrayList(taskIds) } runBlocking { exec.communicate(message) } } + + override fun triggerDagRun( + dagId: String, + runId: String, + logicalDate: OffsetDateTime?, + runAfter: OffsetDateTime?, + conf: Map?, + resetDagRun: Boolean, + note: String?, + ): Boolean { + val message = + TriggerDagRun().also { + it.dagId = dagId + it.runId = runId + it.logicalDate = logicalDate + it.runAfter = runAfter + it.conf = conf + it.resetDagRun = resetDagRun + it.note = note + } + // The supervisor answers a trigger with an OKResponse, and reports a run that already exists as + // the DAGRUN_ALREADY_EXISTS error, which the task decides what to do about. + val response = + runBlocking { + exec.communicateOrNullIf(message, ErrorResponse.ErrorType.DAGRUN_ALREADY_EXISTS) + } + return response == null + } + + override fun getDagRunState( + dagId: String, + runId: String, + ): String { + val message = + GetDagRunState().also { + it.dagId = dagId + it.runId = runId + } + return runBlocking { exec.communicate(message) }.state.toString() + } + + override fun isDagPaused(dagId: String): Boolean = + runBlocking { exec.communicate(GetDag().also { it.dagId = dagId }) }.isPaused == true } diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Logger.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Logger.kt index 228d59e724099..04757386aa46b 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Logger.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Logger.kt @@ -198,6 +198,11 @@ internal class Logger( arguments: Map = emptyMap(), ) = log(Level.DEBUG, message, arguments) + fun info( + message: String, + arguments: Map = emptyMap(), + ) = log(Level.INFO, message, arguments) + fun warning( message: String, arguments: Map = emptyMap(), 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 bdc719a3a685e..e5ad4d6eacf6a 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 @@ -28,6 +28,7 @@ import org.apache.airflow.sdk.LiteralArg import org.apache.airflow.sdk.TaskDef import org.apache.airflow.sdk.TaskGroupRef import org.apache.airflow.sdk.TaskRef +import org.apache.airflow.sdk.TriggerDagRun import org.apache.airflow.sdk.execution.comm.DagFileParseRequest import org.apache.airflow.sdk.internal.Field import org.apache.airflow.sdk.internal.SchemaFields @@ -107,7 +108,7 @@ internal fun serializeDag( "timezone" to dagTimezone(dag.dagConfig), "timetable" to serializeTimetable(dag.id, dag.dagConfig), "tasks" to dag.tasks.map { (taskId, def) -> serializeTask(taskId, def, downstream[taskId]) }, - "dag_dependencies" to emptyList(), + "dag_dependencies" to serializeDagDependencies(dag), "task_group" to serializeTaskGroups(dag, expansion), "edge_info" to emptyMap(), "params" to emptyList(), @@ -127,22 +128,24 @@ private fun serializeTask( def: TaskDef, downstream: List?, ): Map { - val data = - linkedMapOf( - "task_id" to taskId, - "task_type" to def.definition.simpleName, - "_task_module" to def.definition.packageName, - "language" to "java", - // Python's operator serializer always emits template_fields (its list - // value never matches the tuple default it is compared against), so it - // is unconditional here too. Java tasks have no template fields. - "template_fields" to emptyList(), - // What marks a task whose arguments Airflow resolves per instance for a - // runtime outside Python, as `@task.stub` does on the Python side. - // `get_arg_bindings` reads nothing without it. - "is_stub" to true, - ) - argBindings(taskId, def)?.let { data["_arg_bindings"] = it } + val data = linkedMapOf("task_id" to taskId) + val trigger = def.trigger + if (trigger != null) { + writeTriggerDagRun(data, trigger) + } else { + data["task_type"] = def.definition.simpleName + data["_task_module"] = def.definition.packageName + data["language"] = "java" + // Python's operator serializer always emits template_fields (its list + // value never matches the tuple default it is compared against), so it + // is unconditional here too. Java tasks have no template fields. + data["template_fields"] = emptyList() + // What marks a task whose arguments Airflow resolves per instance for a + // runtime outside Python, as `@task.stub` does on the Python side. + // `get_arg_bindings` reads nothing without it. + data["is_stub"] = 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 @@ -165,6 +168,96 @@ private fun serializeTask( ) } +/** What Python's serializer writes for a `TriggerDagRunOperator`. */ +private const val TRIGGER_TASK_TYPE = "TriggerDagRunOperator" +private const val TRIGGER_TASK_MODULE = "airflow.providers.standard.operators.trigger_dagrun" +private const val TRIGGER_UI_COLOR = "#ffefeb" + +private val TRIGGER_TEMPLATE_FIELDS = + listOf( + "trigger_dag_id", + "trigger_run_id", + "logical_date", + "conf", + "wait_for_completion", + "skip_when_already_exists", + ) + +/** + * Writes a task declared from a [TriggerDagRun] as Python writes a + * `TriggerDagRunOperator`, so the Airflow UI and the Dag dependency graph + * treat it as one. + * + * Python writes an operator's template fields and leaves its other parameters + * out, because a Python Dag file holds them. A Java Dag has no such file, so + * every setting the author made is written under the name Python's parameter + * carries. The runtime runs the task from the Dag in its own bundle and reads + * none of this back. + */ +private fun writeTriggerDagRun( + data: MutableMap, + trigger: TriggerDagRun, +) { + data["task_type"] = TRIGGER_TASK_TYPE + data["_task_module"] = TRIGGER_TASK_MODULE + data["ui_color"] = TRIGGER_UI_COLOR + data["template_fields"] = TRIGGER_TEMPLATE_FIELDS + data["template_fields_renderers"] = mapOf("conf" to "py") + data["_operator_extra_links"] = mapOf("Triggered DAG" to "_link_TriggerDagRunLink") + data["trigger_dag_id"] = trigger.dagId + + val settings = trigger.settings + // Python's logical_date defaults to NOTSET, which lets the operator pick the + // trigger time, and writes that sentinel as its name. A template field holds + // a datetime as str(datetime), not as a {"__type": "datetime"} object. + data["logical_date"] = asOffsetDateTime(settings["logical_date"])?.let { pythonIsoformat(it, ' ') } ?: "NOTSET" + // Python writes these two template fields whatever they hold. + data["wait_for_completion"] = settings["wait_for_completion"] ?: false + data["skip_when_already_exists"] = settings["skip_when_already_exists"] ?: false + settings["trigger_run_id"]?.let { data["trigger_run_id"] = it } + settings["conf"]?.let { data["conf"] = it } + // A field that is not a template field is read back with + // BaseSerialization.deserialize, which returns a datetime only for the + // type-encoded form. + settings["run_after"]?.let { data["run_after"] = serializeValue(it) } + settings["reset_dag_run"]?.takeIf { it == true }?.let { data["reset_dag_run"] = true } + settings["poke_interval"]?.let { data["poke_interval"] = (it as Duration).seconds.toInt() } + settings["allowed_states"]?.let { data["allowed_states"] = it } + settings["failed_states"]?.let { data["failed_states"] = it } + settings["fail_when_dag_is_paused"]?.takeIf { it == true }?.let { data["fail_when_dag_is_paused"] = true } + settings["note"]?.let { data["note"] = it } + settings["deferrable"]?.let { data["deferrable"] = it } +} + +/** + * The Dags this one triggers, which is what the UI draws the dependency graph + * between Dags from, sorted as Python sorts them. + */ +private fun serializeDagDependencies(dag: DagDef): List> = + dag.tasks + .filterValues { it.trigger != null } + .map { (taskId, def) -> + mapOf( + "source" to dag.id, + "target" to def.trigger!!.dagId, + // Python labels the dependency with the task's display name, which a Java task cannot set + // yet, and falls back to the task ID as this does. + "label" to taskId, + "dependency_type" to "trigger", + "dependency_id" to taskId, + ) + } + // Python sorts DagDependency as the ordered dataclass it is, field by field. + .sortedWith( + compareBy( + { it["source"] as String }, + { it["target"] as String }, + { it["label"] as String }, + { it["dependency_type"] as String }, + { it["dependency_id"] as String }, + ), + ) + /** * The task's arguments as the binding spec Airflow records, one entry per * argument in the order the Dag's call passed them, as 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 182dfe1cee6e3..54cdb2b3911d9 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 @@ -102,6 +102,19 @@ internal object TaskRunner { val taskDef = bundle.taskDef(request.ti.dagId, request.ti.taskId) ?: return TaskResult.of(TaskState.State.REMOVED) + taskDef.trigger?.let { trigger -> + return try { + TriggerRunner.run(trigger, request, client) + } catch (e: CancellationException) { + throw e + } catch (e: Throwable) { + logger.error( + "Error triggering Dag run", + mapOf("ti" to request.ti, "error" to e, "trace" to e.stackTraceToString()), + ) + TaskResult.failure(request.tiContext.shouldRetry) + } + } val definition = taskDef.definition val instance = try { diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/TriggerRunner.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/TriggerRunner.kt new file mode 100644 index 0000000000000..3af0ce30a59a2 --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/TriggerRunner.kt @@ -0,0 +1,324 @@ +/* + * 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. + */ + +// Mirrors TriggerDagRunOperator.execute and execute_complete in the standard +// provider, as the TypeScript SDK runtime does. + +package org.apache.airflow.sdk.execution + +import org.apache.airflow.sdk.Client +import org.apache.airflow.sdk.TriggerDagRun +import org.apache.airflow.sdk.execution.comm.DeferTask +import org.apache.airflow.sdk.execution.comm.StartupDetails +import org.apache.airflow.sdk.execution.comm.TaskState +import java.security.SecureRandom +import java.time.Duration +import java.time.Instant +import java.time.OffsetDateTime +import java.time.ZoneOffset +import java.time.format.DateTimeFormatter +import kotlin.random.asKotlinRandom + +/** `DagStateTrigger`, the Python triggerer runs the wait of a deferred task with. */ +private const val DAG_STATE_TRIGGER = "airflow.providers.standard.triggers.external_task.DagStateTrigger" + +/** `TriggerDagRunLink().xcom_key`: the XCom the "Triggered DAG" extra link reads. */ +private const val LINK_XCOM_KEY = "_link_TriggerDagRunLink" +private const val RUN_ID_XCOM_KEY = "trigger_run_id" + +/** `TRIGGER_FAIL_REPR`: the `next_method` a failed or timed-out trigger resumes with. */ +private const val TRIGGER_FAIL = "__fail__" +private const val EXECUTE_COMPLETE = "execute_complete" + +private val DEFAULT_POKE_INTERVAL = Duration.ofSeconds(60) + +internal object TriggerRunner { + private val logger = Logger(TriggerRunner::class) + + /** + * Runs a task declared from a [TriggerDagRun], and reports the state it + * leaves the task instance in. + * + * A task resumed by the triggerer reads its outcome from the event the + * trigger fired with, rather than triggering a second run. + */ + fun run( + trigger: TriggerDagRun, + request: StartupDetails, + client: Client, + env: (String) -> String? = System::getenv, + sleep: (Duration) -> Unit = { Thread.sleep(it.toMillis()) }, + now: () -> OffsetDateTime = { OffsetDateTime.now(ZoneOffset.UTC) }, + randomSuffix: () -> String = ::randomRunIdSuffix, + ): Any { + val nextMethod = request.tiContext?.nextMethod as String? + if (nextMethod != null) return resume(trigger, request, nextMethod) + + val runAfter = asOffsetDateTime(trigger.settings["run_after"]) + // `DagRun.generate_run_id`: a run the task gave only a run_after has no logical date, and is + // named after run_after plus a random suffix. Otherwise the logical date names it, and is the + // trigger time when the task set none. + val configuredLogicalDate = asOffsetDateTime(trigger.settings["logical_date"]) + val logicalDate = + when { + configuredLogicalDate != null -> configuredLogicalDate + runAfter == null -> now() + else -> null + } + val namesRun = runAfter ?: logicalDate!! + val runId = + trigger.setting("trigger_run_id") + ?: ("manual__${pythonIsoformat(namesRun)}" + if (logicalDate == null) "_${randomSuffix()}" else "") + + if (trigger.flag("fail_when_dag_is_paused") && client.impl.isDagPaused(trigger.dagId)) { + return failed(request, "Dag ${trigger.dagId} is paused") + } + + logger.info("Triggering Dag Run.", mapOf("trigger_dag_id" to trigger.dagId)) + client.setXCom(LINK_XCOM_KEY, dagRunUrl(env, trigger.dagId, runId)) + val alreadyExists = + client.impl.triggerDagRun( + dagId = trigger.dagId, + runId = runId, + logicalDate = logicalDate, + runAfter = runAfter, + conf = trigger.setting("conf"), + resetDagRun = trigger.flag("reset_dag_run"), + note = trigger.setting("note"), + ) + if (alreadyExists) { + if (trigger.flag("skip_when_already_exists")) { + logger.info( + "Dag Run already exists, skipping task as skip_when_already_exists is set to True.", + mapOf("dag_id" to trigger.dagId), + ) + return TaskResult.of(TaskState.State.SKIPPED) + } + logger.error("Dag Run already exists, marking task as failed.", mapOf("dag_id" to trigger.dagId)) + return TaskResult.of(TaskState.State.FAILED) + } + logger.info("Dag Run triggered successfully.", mapOf("trigger_dag_id" to trigger.dagId)) + client.setXCom(RUN_ID_XCOM_KEY, runId) + + if (!trigger.flag("wait_for_completion")) { + if (deferrable(trigger, env)) { + logger.info( + "Ignoring deferrable=True because wait_for_completion=False. " + + "Task will complete immediately without waiting for the triggered DAG run.", + mapOf("trigger_dag_id" to trigger.dagId), + ) + } + return TaskResult.success() + } + + val pokeInterval = trigger.setting("poke_interval") ?: DEFAULT_POKE_INTERVAL + if (deferrable(trigger, env)) { + logger.info("Pausing task as DEFERRED.", mapOf("trigger_dag_id" to trigger.dagId, "run_id" to runId)) + return defer(trigger, request, runId, pokeInterval, env) + } + + while (true) { + logger.info( + "Waiting for dag run to complete execution in allowed state.", + mapOf("dag_id" to trigger.dagId, "run_id" to runId, "allowed_state" to allowedStates(trigger)), + ) + sleep(pokeInterval) + val state = client.impl.getDagRunState(trigger.dagId, runId) + if (state in failedStates(trigger)) { + logger.error("DagRun finished with failed state.", mapOf("dag_id" to trigger.dagId, "state" to state)) + return failed(request, "${trigger.dagId} failed with failed state $state") + } + if (state in allowedStates(trigger)) { + logger.info("DagRun finished with allowed state.", mapOf("dag_id" to trigger.dagId, "state" to state)) + return TaskResult.success() + } + logger.debug( + "DagRun not yet in allowed or failed state.", + mapOf("dag_id" to trigger.dagId, "state" to state), + ) + } + } + + /** `DagStateTrigger.serialize()`, key for key, so the Python triggerer runs the wait. */ + private fun defer( + trigger: TriggerDagRun, + request: StartupDetails, + runId: String, + pokeInterval: Duration, + env: (String) -> String?, + ): DeferTask = + DeferTask().also { + it.state = "deferred" + it.classpath = DAG_STATE_TRIGGER + it.triggerKwargs = + mapOf( + "dag_id" to trigger.dagId, + "states" to allowedStates(trigger) + failedStates(trigger), + "poll_interval" to pokeInterval.seconds.toInt(), + "run_ids" to listOf(runId), + "execution_dates" to null, + ) + it.triggerTimeout = null + // `_defer_task` hands the trigger the task's queue only when triggerer queues are enabled. + it.queue = if (booleanEnv(env, "AIRFLOW__TRIGGERER__QUEUES_ENABLED", false)) request.ti.queue else null + it.nextMethod = EXECUTE_COMPLETE + it.nextKwargs = emptyMap() + } + + /** `BaseOperator.resume_execution` for this task: `__fail__` or `execute_complete`. */ + private fun resume( + trigger: TriggerDagRun, + request: StartupDetails, + nextMethod: String, + ): Any { + val kwargs = request.tiContext?.nextKwargs as? Map<*, *> ?: emptyMap() + if (nextMethod == TRIGGER_FAIL) { + (kwargs["traceback"] as? List<*>)?.let { logger.error("Trigger failed:\n${it.joinToString("\n")}") } + return failed(request, (kwargs["error"] ?: "Unknown").toString()) + } + if (nextMethod != EXECUTE_COMPLETE) { + return failed(request, "Task cannot resume with next_method \"$nextMethod\"") + } + val event = + decodeEvent(kwargs["event"]) + ?: return failed(request, "Task resumed with an event it cannot read: ${kwargs["event"]}") + val runIds = + event["run_ids"] as? List<*> ?: return failed( + request, + "Task resumed with an event it cannot read: ${kwargs["event"]}", + ) + val failedRunIds = + runIds.filter { runId -> + val state = event[runId.toString()].toString() + if (state in failedStates(trigger)) return@filter true + if (state in allowedStates(trigger)) { + logger.info( + "Triggered Dag run finished with allowed state.", + mapOf("dag_id" to trigger.dagId, "state" to state, "run_id" to runId.toString()), + ) + } + false + } + if (failedRunIds.isNotEmpty()) { + return failed( + request, + "${trigger.dagId} failed with failed states ${failedStates(trigger)} for run_ids $failedRunIds", + ) + } + return TaskResult.success() + } + + /** + * The payload `DagStateTrigger` fired with: `(classpath, data)`. The + * triggerer stores it with serde, which encodes a tuple as + * `{__classname__, __data__}`. + */ + private fun decodeEvent(event: Any?): Map? { + val pair = + if (event is Map<*, *> && event["__classname__"] == "builtins.tuple") event["__data__"] else event + val items = pair as? List<*> ?: return null + if (items.size != 2) return null + val data = items[1] as? Map<*, *> ?: return null + return data.entries.associate { (key, value) -> key.toString() to value } + } + + private fun failed( + request: StartupDetails, + message: String, + ): Any { + logger.error(message) + return TaskResult.failure(request.tiContext?.shouldRetry == true) + } + + private fun allowedStates(trigger: TriggerDagRun): List = + trigger.setting>("allowed_states")?.takeIf { it.isNotEmpty() } ?: listOf("success") + + // An empty failed_states that was set means no state fails the task, so it is kept as written. + private fun failedStates(trigger: TriggerDagRun): List = trigger.setting>("failed_states") ?: listOf("failed") + + private fun deferrable( + trigger: TriggerDagRun, + env: (String) -> String?, + ): Boolean = + trigger.setting("deferrable") + ?: booleanEnv(env, "AIRFLOW__OPERATORS__DEFAULT_DEFERRABLE", false) +} + +@Suppress("UNCHECKED_CAST") +private fun TriggerDagRun.setting(key: String): T? = settings[key] as T? + +private fun TriggerDagRun.flag(key: String): Boolean = settings[key] == true + +/** + * An Airflow boolean option from the environment. It takes what Go's `strconv.ParseBool` takes + * rather than every spelling `configparser` accepts, so `yes` and `on` are rejected rather than + * read as true. + */ +internal fun booleanEnv( + env: (String) -> String?, + name: String, + fallback: Boolean, +): Boolean { + val raw = env(name) ?: return fallback + return when (raw.trim().lowercase()) { + "t", "true", "1" -> true + "f", "false", "0" -> false + else -> throw IllegalArgumentException("$name is \"$raw\", which is not a boolean; use true or false") + } +} + +/** `build_airflow_dagrun_url`, on `[api] base_url` from the environment. */ +internal fun dagRunUrl( + env: (String) -> String?, + dagId: String, + runId: String, +): String { + val base = env("AIRFLOW__API__BASE_URL")?.takeIf { it.isNotEmpty() } ?: "/" + return "${base.trimEnd('/')}/dags/$dagId/runs/$runId" +} + +/** The eight characters `get_random_string` adds to the ID of a run with no logical date. */ +private const val RUN_ID_SUFFIX_CHARS = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789" + +private val RUN_ID_RANDOM = SecureRandom().asKotlinRandom() + +internal fun randomRunIdSuffix(): String = (1..8).map { RUN_ID_SUFFIX_CHARS.random(RUN_ID_RANDOM) }.joinToString("") + +/** + * `datetime.isoformat(sep)` of a UTC instant, which is how Python spells a run ID's date. `str(datetime)` is the + * same with a space as [separator]. + */ +internal fun pythonIsoformat( + moment: OffsetDateTime, + separator: Char = 'T', +): String { + val utc = moment.withOffsetSameInstant(ZoneOffset.UTC) + val micros = utc.nano / 1000 + val seconds = utc.format(DateTimeFormatter.ofPattern("uuuu-MM-dd'$separator'HH:mm:ss")) + return if (micros == 0) "$seconds+00:00" else "$seconds.%06d+00:00".format(micros) +} + +/** Reads a setting that may be stored as either temporal type. */ +internal fun asOffsetDateTime(value: Any?): OffsetDateTime? = + when (value) { + null -> null + is OffsetDateTime -> value + is Instant -> value.atOffset(ZoneOffset.UTC) + else -> throw IllegalArgumentException("Not a date-time: $value") + } diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/Fields.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/Fields.kt index cc43db3186590..5b813b5c7333f 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/Fields.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/Fields.kt @@ -36,8 +36,17 @@ enum class FieldType { STRING_ARRAY, TIMEDELTA, DATETIME, + + /** A list of Dag run states, as `allowed_states` takes. */ + DAG_RUN_STATES, + + /** A map with string keys, as `conf` takes. */ + JSON_OBJECT, } +/** A Dag run state, as a `DAG_RUN_STATES` field names one. */ +private val DAG_RUN_STATE_NAMES = listOf("queued", "running", "success", "failed") + /** * One configuration key from the Dag serialization schema. Public so that the * annotation processor can lower `@Builder.Dag` / `@Builder.Task` attributes; @@ -103,5 +112,29 @@ internal fun checkConfigValue( value is Array<*> && value.all { it is String } -> value.map { it as String } else -> mismatch("an Iterable of String") } + FieldType.DAG_RUN_STATES -> { + val states = + when (value) { + is Iterable<*> -> value.toList() + is Array<*> -> value.toList() + else -> mismatch("an Iterable of Dag run states") + } + states.map { state -> + require(state is String && state in DAG_RUN_STATE_NAMES) { + "Value for $scope config key '$key' holds $state, which is not a Dag run state; use one " + + "of ${DAG_RUN_STATE_NAMES.joinToString()}" + } + state as String + } + } + FieldType.JSON_OBJECT -> { + val map = value as? Map<*, *> ?: mismatch("a Map") + map.entries.associate { (mapKey, entry) -> + require(mapKey is String) { + "Value for $scope config key '$key' must have String keys, got: ${mapKey?.javaClass?.name}" + } + mapKey to entry + } + } } } 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 6c6af5287650e..95edc43bbcfc1 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 @@ -92,6 +92,23 @@ internal class FakeXComTransport( override fun clearTaskStateStore(tiId: UUID) = throw NotImplementedError() override fun skipDownstreamTasks(taskIds: List) = throw NotImplementedError() + + override fun triggerDagRun( + dagId: String, + runId: String, + logicalDate: OffsetDateTime?, + runAfter: OffsetDateTime?, + conf: Map?, + resetDagRun: Boolean, + note: String?, + ): Boolean = throw NotImplementedError() + + override fun getDagRunState( + dagId: String, + runId: String, + ): String = throw NotImplementedError() + + override fun isDagPaused(dagId: String): Boolean = throw NotImplementedError() } internal fun startupDetails(argBindings: List>?): 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 2ff51df4f36be..239e2bbf95496 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 @@ -106,6 +106,23 @@ private class FakeTransport( } override fun skipDownstreamTasks(taskIds: List) = Unit + + override fun triggerDagRun( + dagId: String, + runId: String, + logicalDate: OffsetDateTime?, + runAfter: OffsetDateTime?, + conf: Map?, + resetDagRun: Boolean, + note: String?, + ): Boolean = throw NotImplementedError() + + override fun getDagRunState( + dagId: String, + runId: String, + ): String = throw NotImplementedError() + + override fun isDagPaused(dagId: String): Boolean = throw NotImplementedError() } class ClientTest { diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/SwitchTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/SwitchTest.kt index ef72edf3a4c0b..83bd8305d9f86 100644 --- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/SwitchTest.kt +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/SwitchTest.kt @@ -89,6 +89,23 @@ internal class SwitchTest { ) } + @Test + @DisplayName("Should reject a task that triggers a Dag run as a case") + fun shouldRejectTriggerTaskAsCase() { + val dag = DagDef("d") + val trigger = dag.task("trigger_x", TriggerDagRun("x")) + val switch = dag.Switch(PickPath::class.java) + + val error = assertThrows(IllegalArgumentException::class.java) { switch.Case(trigger) } + + assertEquals( + "Switch 'pickPath' cannot choose 'trigger_x': a switch names its case by its class, and a " + + "task that triggers a Dag run has none. Put the trigger downstream of an ordinary task and " + + "let the switch choose that task, since Airflow skips the downstream of a skipped task", + error.message, + ) + } + @Test @DisplayName("Should reject a case declared in another Dag") fun shouldRejectCaseFromAnotherDag() { diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/TriggerDagRunTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/TriggerDagRunTest.kt new file mode 100644 index 0000000000000..56c0bb410b7a1 --- /dev/null +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/TriggerDagRunTest.kt @@ -0,0 +1,275 @@ +/* + * 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.DisplayName +import org.junit.jupiter.api.Test +import java.time.Duration +import java.time.OffsetDateTime + +internal class TriggerDagRunTest { + @Suppress("UNCHECKED_CAST") + private fun serializeTrigger(trigger: TriggerDagRun): Map { + val dag = DagDef("d") + dag.task("trigger", trigger) + val tasks = serializeDag(dag, "", ".")["tasks"] as List> + return tasks.single()["__var"] as Map + } + + @Test + @DisplayName("Should reject a trigger that names no Dag") + fun shouldRejectEmptyDagId() { + val error = assertThrows(IllegalArgumentException::class.java) { TriggerDagRun("") } + + assertEquals("TriggerDagRun needs the ID of the Dag to trigger", error.message) + } + + @Test + @DisplayName("Should point at the constructor for the Dag being triggered") + fun shouldRejectTriggerDagIdAsConfig() { + val error = + assertThrows(IllegalArgumentException::class.java) { + TriggerDagRun("downstream").config("trigger_dag_id", "other") + } + + assertEquals( + "The Dag to trigger is the TriggerDagRun argument, not the config key 'trigger_dag_id'", + error.message, + ) + } + + @Test + @DisplayName("Should reject a poke interval that is not a whole number of seconds") + fun shouldRejectFractionalPokeInterval() { + val error = + assertThrows(IllegalArgumentException::class.java) { + TriggerDagRun("downstream").config("poke_interval", Duration.ofMillis(1500)) + } + + assertEquals( + "Value for TriggerDagRun config key 'poke_interval' is PT1.5S; it must be a whole, " + + "non-negative number of seconds", + error.message, + ) + } + + @Test + @DisplayName("Should reject a state no Dag run can be in") + fun shouldRejectUnknownState() { + val error = + assertThrows(IllegalArgumentException::class.java) { + TriggerDagRun("downstream").config("allowed_states", listOf("success", "done")) + } + + assertEquals( + "Value for TriggerDagRun config key 'allowed_states' holds done, which is not a Dag run " + + "state; use one of queued, running, success, failed", + error.message, + ) + } + + @Test + @DisplayName("Should reject a conf value that has no JSON form when the setting is made") + fun shouldRejectConfWithoutJsonForm() { + val error = + assertThrows(IllegalArgumentException::class.java) { + TriggerDagRun("downstream").config("conf", mapOf("at" to Duration.ofSeconds(1))) + } + + assertEquals( + "Value for TriggerDagRun config key 'conf' must hold only JSON values (null, String, " + + "Boolean, Number, List, Map), got: java.time.Duration", + error.message, + ) + } + + @Test + @DisplayName("Should reject a conf holding a number JSON has no form for") + fun shouldRejectNonFiniteConf() { + val error = + assertThrows(IllegalArgumentException::class.java) { + TriggerDagRun("downstream").config("conf", mapOf("ratio" to listOf(Double.NaN))) + } + + assertEquals( + "Value for TriggerDagRun config key 'conf' must hold only finite numbers, got: NaN", + error.message, + ) + } + + @Test + @DisplayName("Should reject an empty run ID or note") + fun shouldRejectEmptyRunIdAndNote() { + listOf("trigger_run_id", "note").forEach { key -> + val error = + assertThrows(IllegalArgumentException::class.java) { TriggerDagRun("downstream").config(key, "") } + + assertEquals("Value for TriggerDagRun config key '$key' must not be empty", error.message) + } + } + + @Test + @DisplayName("Should check a value against its key's shape as other config does") + fun shouldReportMismatchInSharedShape() { + val error = + assertThrows(IllegalArgumentException::class.java) { + TriggerDagRun("downstream").config("wait_for_completion", "yes") + } + + assertEquals( + "Value for TriggerDagRun config key 'wait_for_completion' must be a Boolean, got: java.lang.String", + error.message, + ) + } + + @Test + @DisplayName("Should reject an unknown setting") + fun shouldRejectUnknownKey() { + val error = + assertThrows(IllegalArgumentException::class.java) { + TriggerDagRun("downstream").config("waitForCompletion", true) + } + + assertEquals("Unknown TriggerDagRun config key: 'waitForCompletion'", error.message) + } + + @Test + @DisplayName("Should write a trigger task as Python writes a TriggerDagRunOperator") + fun shouldSerializeDefaults() { + val data = serializeTrigger(TriggerDagRun("downstream")) + + assertEquals("TriggerDagRunOperator", data["task_type"]) + assertEquals("airflow.providers.standard.operators.trigger_dagrun", data["_task_module"]) + assertEquals("#ffefeb", data["ui_color"]) + assertEquals(mapOf("Triggered DAG" to "_link_TriggerDagRunLink"), data["_operator_extra_links"]) + assertEquals(mapOf("conf" to "py"), data["template_fields_renderers"]) + assertEquals("downstream", data["trigger_dag_id"]) + assertEquals("NOTSET", data["logical_date"]) + assertEquals(false, data["wait_for_completion"]) + assertEquals(false, data["skip_when_already_exists"]) + // A Java task's markers belong to a task the Java runtime runs a body for. + assertFalse("language" in data) + assertFalse("is_stub" in data) + assertFalse("reset_dag_run" in data) + } + + @Test + @DisplayName("Should write a logical date as Python writes a template field, and run_after type-encoded") + fun shouldSerializeTemporalSettings() { + val data = + serializeTrigger( + TriggerDagRun("downstream") + .config("logical_date", OffsetDateTime.parse("2026-09-30T01:02:03Z")) + .config("run_after", OffsetDateTime.parse("2026-09-30T00:00:00Z")) + .config("poke_interval", Duration.ofSeconds(30)), + ) + + assertEquals("2026-09-30 01:02:03+00:00", data["logical_date"]) + assertEquals(mapOf("__type" to "datetime", "__var" to 1790726400.0), data["run_after"]) + assertEquals(30, data["poke_interval"]) + } + + @Test + @DisplayName("Should record the Dags a Dag triggers, sorted as Python sorts them") + fun shouldRecordDagDependencies() { + val dag = DagDef("d") + dag.task("second", TriggerDagRun("beta")) + dag.task("first", TriggerDagRun("alpha")) + + val dependencies = serializeDag(dag, "", ".")["dag_dependencies"] as List<*> + + assertEquals( + listOf( + mapOf( + "source" to "d", + "target" to "alpha", + "label" to "first", + "dependency_type" to "trigger", + "dependency_id" to "first", + ), + mapOf( + "source" to "d", + "target" to "beta", + "label" to "second", + "dependency_type" to "trigger", + "dependency_id" to "second", + ), + ), + dependencies, + ) + } + + @Test + @DisplayName("Should carry a trigger task's own Airflow settings") + fun shouldCarryTaskConfig() { + val dag = DagDef("d") + dag.task("trigger", TriggerDagRun("downstream")).config("retries", 2) + + @Suppress("UNCHECKED_CAST") + val tasks = serializeDag(dag, "", ".")["tasks"] as List> + + assertEquals(2, (tasks.single()["__var"] as Map<*, *>)["retries"]) + } + + @Test + @DisplayName("Should declare a trigger task inside a task group under the group's prefix") + fun shouldDeclareTriggerInGroup() { + val dag = DagDef("d") + val ref = dag.taskGroup("downstream").task("trigger", TriggerDagRun("other")) + + assertEquals("downstream.trigger", ref.def.id) + assertEquals(listOf("downstream.trigger"), dag.tasks.keys.toList()) + } + + @Test + @DisplayName("Should not let a setting made after registration reach the registered task") + fun shouldSnapshotTriggerOnRegistration() { + val dag = DagDef("d") + val nested = mutableMapOf("n" to 1) + val list = mutableListOf("a") + val trigger = TriggerDagRun("other").config("conf", mapOf("nested" to nested, "list" to list)) + val ref = dag.task("trigger", trigger) + + trigger.config("note", "late") + nested["n"] = 2 + list += "b" + + val registered = ref.def.trigger!!.settings + assertFalse("note" in registered) + assertEquals(mapOf("nested" to mapOf("n" to 1), "list" to listOf("a")), registered["conf"]) + } + + @Test + @DisplayName("Should keep two tasks registered from one builder apart") + fun shouldKeepTasksFromOneBuilderApart() { + val dag = DagDef("d") + val trigger = TriggerDagRun("other").config("note", "first") + val first = dag.task("first", trigger) + trigger.config("note", "second") + val second = dag.task("second", trigger) + + assertEquals("first", first.def.trigger!!.settings["note"]) + assertEquals("second", second.def.trigger!!.settings["note"]) + } +} 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 e327afd661e4b..f7d11aa42846a 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 @@ -44,6 +44,7 @@ import org.apache.airflow.sdk.SwitchTask import org.apache.airflow.sdk.Task import org.apache.airflow.sdk.TaskGroupRef import org.apache.airflow.sdk.TaskRef +import org.apache.airflow.sdk.TriggerDagRun import org.apache.airflow.sdk.execution.serializeDag import org.apache.airflow.sdk.internal.Field import org.apache.airflow.sdk.internal.FieldType @@ -202,6 +203,7 @@ private fun buildDag(case: JsonNode): DagDef { val groupId = task.path("group").asText("") val localId = task.path("task_id").asText() val branch = task.path("branch") + val trigger = task.path("trigger_dag_run") val definition = when { branch.isMissingNode -> { @@ -214,10 +216,13 @@ private fun buildDag(case: JsonNode): DagDef { else -> ConformanceCondition::class.java } val ref = - if (groupId.isEmpty()) { - dag.task(localId, definition) - } else { - groups.getValue(groupId).task(localId, definition) + when { + !trigger.isMissingNode -> + triggerDagRun(trigger).let { + if (groupId.isEmpty()) dag.task(localId, it) else groups.getValue(groupId).task(localId, it) + } + groupId.isEmpty() -> dag.task(localId, definition) + else -> groups.getValue(groupId).task(localId, definition) } if (!branch.isMissingNode) { deciders += (if (branch.has("cases")) SwitchRef.of(ref) else asCondition(ref)) to branch @@ -253,6 +258,30 @@ private fun buildDag(case: JsonNode): DagDef { return dag } +/** Reads the template fields of a TriggerDagRunOperator from a task of test_dags.yaml. */ +private fun triggerDagRun(node: JsonNode): TriggerDagRun { + val trigger = TriggerDagRun(node.path("trigger_dag_id").asText()) + node.fields().forEach { (key, value) -> + when (key) { + "trigger_dag_id" -> Unit + "logical_date", "run_after" -> trigger.config(key, triggerDateTime(value.asText())) + "conf" -> trigger.config(key, toJsonValue(value) as Map<*, *>) + "wait_for_completion", "skip_when_already_exists", "reset_dag_run", "fail_when_dag_is_paused", "deferrable" -> + trigger.config(key, value.asBoolean()) + "poke_interval" -> trigger.config(key, Duration.ofSeconds(value.asLong())) + "allowed_states", "failed_states" -> trigger.config(key, value.map { it.asText() }) + else -> trigger.config(key, value.asText()) + } + } + return trigger +} + +/** + * A trigger task's `logical_date` or `run_after` from the fixture. The fixture spells one the way + * Python spells a datetime, with a space, so the separator is normalised before parsing. + */ +private fun triggerDateTime(text: String): OffsetDateTime = OffsetDateTime.parse(text.replaceFirst(" ", "T")) + /** 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) @@ -271,7 +300,8 @@ private fun toValue( // `!datetime` is an ISO 8601 timestamp, and `!timedelta` a number of seconds. FieldType.DATETIME -> OffsetDateTime.parse(node.asText()) FieldType.TIMEDELTA -> Duration.ofNanos((node.asText().toDouble() * 1e9).toLong()) - FieldType.STRING_ARRAY -> node.map { it.asText() } + FieldType.STRING_ARRAY, FieldType.DAG_RUN_STATES -> node.map { it.asText() } + FieldType.JSON_OBJECT -> toJsonValue(node)!! } /** Reads a YAML literal as the plain value `Serde` writes out. */ 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 1d9b5cc4ace72..20fba213948f5 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 @@ -253,6 +253,85 @@ class CommsTest { return (requests.single().rawBody as Map<*, *>) to failure } + /** + * Like [roundTrip], but on the transport client, which is where the calls a + * trigger or a decider makes live. + */ + private fun transportRoundTrip( + response: (Int) -> ByteArray, + call: (CoordinatorClient) -> T, + ): Triple, T?, Throwable?> { + val toClient = ByteChannel(autoFlush = true) + val fromClient = ByteChannel(autoFlush = true) + val comm = CoordinatorComm(toClient, fromClient) + val client = CoordinatorClient(comm) + + val requests = ConcurrentLinkedQueue() + val server = + Thread { + val request = readRequest(fromClient) + requests.add(request) + runBlocking { toClient.writeFrame(response(request.id)) } + } + server.start() + val outcome = runCatching { call(client) } + server.join() + comm.close() + + return Triple(requests.single().rawBody as Map<*, *>, outcome.getOrNull(), outcome.exceptionOrNull()) + } + + private fun dagRunExistsResponseFrame(id: Int): ByteArray { + val out = ByteArrayOutputStream() + MessagePack.newDefaultPacker(out).use { packer -> + packer.packArrayHeader(3) + packer.packInt(id) + packer.packNil() + packer.packMapHeader(3) + packer.packString("type") + packer.packString("ErrorResponse") + packer.packString("error") + packer.packString("DAGRUN_ALREADY_EXISTS") + packer.packString("detail") + packer.packMapHeader(1) + packer.packString("dag_id") + packer.packString("downstream") + } + return out.toByteArray() + } + + private fun dagRunStateResultFrame(id: Int): ByteArray { + val out = ByteArrayOutputStream() + MessagePack.newDefaultPacker(out).use { packer -> + packer.packArrayHeader(3) + packer.packInt(id) + packer.packMapHeader(2) + packer.packString("type") + packer.packString("DagRunStateResult") + packer.packString("state") + packer.packString("running") + packer.packNil() + } + return out.toByteArray() + } + + private fun dagResultFrame(id: Int): ByteArray { + val out = ByteArrayOutputStream() + MessagePack.newDefaultPacker(out).use { packer -> + packer.packArrayHeader(3) + packer.packInt(id) + packer.packMapHeader(3) + packer.packString("type") + packer.packString("DagResult") + packer.packString("dag_id") + packer.packString("downstream") + packer.packString("is_paused") + packer.packBoolean(true) + packer.packNil() + } + return out.toByteArray() + } + private suspend fun ByteChannel.writeFrame(payload: ByteArray) { writeByteArray(Frame.lengthPrefix(payload.size.toUInt())) writeByteArray(payload) @@ -380,6 +459,71 @@ class CommsTest { comm.close() } + @Test + @DisplayName("triggerDagRun takes the OKResponse the supervisor answers a trigger with") + @Timeout(value = 30, unit = TimeUnit.SECONDS) + fun triggerDagRunAcceptsOkResponse() { + val (body, alreadyExists, failure) = + transportRoundTrip(::okResponseFrame) { + it.triggerDagRun( + dagId = "downstream", + runId = "manual__2026-09-30T00:00:00+00:00", + logicalDate = OffsetDateTime.parse("2026-09-30T00:00:00Z"), + runAfter = null, + conf = mapOf("rows" to 2), + resetDagRun = false, + note = null, + ) + } + + Assertions.assertNull(failure, "triggerDagRun should return normally on an OKResponse, got $failure") + Assertions.assertEquals(false, alreadyExists) + Assertions.assertEquals("TriggerDagRun", body["type"]) + Assertions.assertEquals("downstream", body["dag_id"]) + Assertions.assertEquals("manual__2026-09-30T00:00:00+00:00", body["run_id"]) + // MessagePack packs a small integer as a byte, so the conf is compared by how it reads. + Assertions.assertEquals("{rows=2}", body["conf"].toString()) + } + + @Test + @DisplayName("triggerDagRun reports an existing run rather than failing") + @Timeout(value = 30, unit = TimeUnit.SECONDS) + fun triggerDagRunReportsAnExistingRun() { + val (_, alreadyExists, failure) = + transportRoundTrip(::dagRunExistsResponseFrame) { + it.triggerDagRun("downstream", "run", null, null, null, false, null) + } + + Assertions.assertNull(failure, "an existing run is a value, not an error, got $failure") + Assertions.assertEquals(true, alreadyExists) + } + + @Test + @DisplayName("getDagRunState reads the state out of a DagRunStateResult") + @Timeout(value = 30, unit = TimeUnit.SECONDS) + fun getDagRunStateReadsTheState() { + val (body, state, failure) = + transportRoundTrip(::dagRunStateResultFrame) { it.getDagRunState("downstream", "run") } + + Assertions.assertNull(failure, "getDagRunState should return normally, got $failure") + Assertions.assertEquals("running", state) + Assertions.assertEquals("GetDagRunState", body["type"]) + Assertions.assertEquals("downstream", body["dag_id"]) + Assertions.assertEquals("run", body["run_id"]) + } + + @Test + @DisplayName("isDagPaused reads is_paused out of a DagResult") + @Timeout(value = 30, unit = TimeUnit.SECONDS) + fun isDagPausedReadsTheFlag() { + val (body, paused, failure) = transportRoundTrip(::dagResultFrame) { it.isDagPaused("downstream") } + + Assertions.assertNull(failure, "isDagPaused should return normally, got $failure") + Assertions.assertEquals(true, paused) + Assertions.assertEquals("GetDag", body["type"]) + Assertions.assertEquals("downstream", body["dag_id"]) + } + @Test @DisplayName("setVariable keeps a null description on the wire so the supervisor accepts the request") @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 db1478661c186..a18907a4e25e8 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 @@ -27,6 +27,7 @@ import org.apache.airflow.sdk.DagDef import org.apache.airflow.sdk.SwitchTask import org.apache.airflow.sdk.Task import org.apache.airflow.sdk.TaskDef +import org.apache.airflow.sdk.TriggerDagRun import org.apache.airflow.sdk.execution.comm.BundleInfo import org.apache.airflow.sdk.execution.comm.DagRun import org.apache.airflow.sdk.execution.comm.RetryTask @@ -294,6 +295,36 @@ class TaskTest { Assertions.assertEquals(emptyList>(), transport.xComs) } + @Test + @DisplayName("Should run a trigger task itself and push the ID of the run it started") + fun shouldRunTriggerTask() { + val transport = RecordingTransport() + + val result = runTask(triggerBundle(), startupDetails(taskId = "trigger"), Client(startupDetails("trigger"), transport)) + + Assertions.assertInstanceOf(SucceedTask::class.java, result) + Assertions.assertEquals(listOf("downstream"), transport.triggered) + Assertions.assertTrue("trigger_run_id" in transport.xComs.map { it.first }) { "pushed: ${transport.xComs}" } + } + + @Test + @DisplayName("Should fail a trigger task whose client throws") + fun shouldFailTriggerTaskWhenClientThrows() { + val transport = RecordingTransport().also { it.triggerFailure = IllegalStateException("boom") } + + val result = runTask(triggerBundle(), startupDetails(taskId = "trigger"), Client(startupDetails("trigger"), transport)) + + Assertions.assertInstanceOf(TaskState::class.java, result) + Assertions.assertEquals(TaskState.State.FAILED, (result as TaskState).state) + Assertions.assertFalse("trigger_run_id" in transport.xComs.map { it.first }) + } + + private fun triggerBundle(): Bundle { + val dag = DagDef("test_dag") + dag.task("trigger", TriggerDagRun("downstream")) + return Bundle(listOf(dag)) + } + private fun switchBundle(): Bundle { val dag = DagDef("test_dag") val long = dag.task("handle_long", HandleLongTask::class.java) @@ -402,6 +433,23 @@ class TaskTest { override fun clearTaskStateStore(tiId: UUID): Unit = throw UnsupportedOperationException("not used in test") override fun skipDownstreamTasks(taskIds: List): Unit = throw UnsupportedOperationException("not used in test") + + override fun triggerDagRun( + dagId: String, + runId: String, + logicalDate: OffsetDateTime?, + runAfter: OffsetDateTime?, + conf: Map?, + resetDagRun: Boolean, + note: String?, + ): Boolean = throw UnsupportedOperationException("not used in test") + + override fun getDagRunState( + dagId: String, + runId: String, + ): String = throw UnsupportedOperationException("not used in test") + + override fun isDagPaused(dagId: String): Boolean = throw UnsupportedOperationException("not used in test") }, ) @@ -463,6 +511,8 @@ class TaskTest { val xComs = mutableListOf>() val skipped = mutableListOf() val events = mutableListOf() + val triggered = mutableListOf() + var triggerFailure: Throwable? = null override fun setXCom( key: String, @@ -520,6 +570,27 @@ class TaskTest { ): Unit = throw UnsupportedOperationException("not used in test") override fun clearTaskStateStore(tiId: UUID): Unit = throw UnsupportedOperationException("not used in test") + + override fun triggerDagRun( + dagId: String, + runId: String, + logicalDate: OffsetDateTime?, + runAfter: OffsetDateTime?, + conf: Map?, + resetDagRun: Boolean, + note: String?, + ): Boolean { + triggered += dagId + triggerFailure?.let { throw it } + return false + } + + override fun getDagRunState( + dagId: String, + runId: String, + ): String = throw UnsupportedOperationException("not used in test") + + override fun isDagPaused(dagId: String): Boolean = throw UnsupportedOperationException("not used in test") } class SuccessTask : Task { diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TriggerRunnerTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TriggerRunnerTest.kt new file mode 100644 index 0000000000000..ad0b629fa2a24 --- /dev/null +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TriggerRunnerTest.kt @@ -0,0 +1,473 @@ +/* + * 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.execution + +import org.apache.airflow.sdk.Client +import org.apache.airflow.sdk.TriggerDagRun +import org.apache.airflow.sdk.execution.comm.DagRun +import org.apache.airflow.sdk.execution.comm.DeferTask +import org.apache.airflow.sdk.execution.comm.RetryTask +import org.apache.airflow.sdk.execution.comm.StartupDetails +import org.apache.airflow.sdk.execution.comm.SucceedTask +import org.apache.airflow.sdk.execution.comm.TIRunContext +import org.apache.airflow.sdk.execution.comm.TaskInstance +import org.apache.airflow.sdk.execution.comm.TaskState +import org.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.Assertions.assertInstanceOf +import org.junit.jupiter.api.Assertions.assertNull +import org.junit.jupiter.api.Assertions.assertTrue +import org.junit.jupiter.api.DisplayName +import org.junit.jupiter.api.Test +import java.time.Duration +import java.time.OffsetDateTime +import java.util.UUID + +/** One call to [FakeTrigger.triggerDagRun], so a test can read what the task asked Airflow for. */ +private class TriggerCall( + val dagId: String, + val runId: String, + val logicalDate: OffsetDateTime?, + val runAfter: OffsetDateTime?, + val conf: Map?, + val resetDagRun: Boolean, + val note: String?, +) + +private class FakeTrigger( + private val alreadyExists: Boolean = false, + private val paused: Boolean = false, + private val states: List = emptyList(), +) : org.apache.airflow.sdk.execution.Client { + val triggered = mutableListOf() + val xComs = mutableListOf>() + var polls = 0 + + override fun triggerDagRun( + dagId: String, + runId: String, + logicalDate: OffsetDateTime?, + runAfter: OffsetDateTime?, + conf: Map?, + resetDagRun: Boolean, + note: String?, + ): Boolean { + triggered += TriggerCall(dagId, runId, logicalDate, runAfter, conf, resetDagRun, note) + return alreadyExists + } + + override fun getDagRunState( + dagId: String, + runId: String, + ): String { + val started = triggered.last() + assertEquals(started.dagId, dagId) + assertEquals(started.runId, runId) + return states[polls++.coerceAtMost(states.lastIndex)] + } + + override fun isDagPaused(dagId: String): Boolean = paused + + override fun setXCom( + key: String, + value: Any, + dagId: String, + taskId: String, + runId: String, + mapIndex: Int, + ) { + xComs += key to value + } + + 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") + + override fun skipDownstreamTasks(taskIds: List): Unit = throw UnsupportedOperationException("not used in test") +} + +/** The trigger time the tests run at, so a generated run ID is the same on every run. */ +private val NOW: OffsetDateTime = OffsetDateTime.parse("2026-10-08T05:00:00Z") + +internal class TriggerRunnerTest { + private fun startupDetails( + nextMethod: String? = null, + nextKwargs: Any? = null, + shouldRetry: Boolean = false, + ): StartupDetails = + StartupDetails().also { + it.ti = + TaskInstance().also { ti -> + ti.id = UUID.randomUUID() + ti.taskId = "trigger" + ti.dagId = "test_dag" + ti.runId = "manual__2026-03-31T00:00:00+00:00" + ti.tryNumber = 1 + ti.dagVersionId = UUID.randomUUID() + ti.queue = "workers" + } + it.tiContext = + TIRunContext().apply { + dagRun = DagRun().apply { dagId = "test_dag" } + this.shouldRetry = shouldRetry + this.nextMethod = nextMethod + this.nextKwargs = nextKwargs + } + } + + private fun run( + trigger: TriggerDagRun, + transport: FakeTrigger, + details: StartupDetails = startupDetails(), + env: Map = emptyMap(), + slept: MutableList = mutableListOf(), + ): Any = + TriggerRunner.run( + trigger, + details, + Client(details, transport), + env = env::get, + sleep = { slept += it }, + now = { NOW }, + randomSuffix = { "RaNd0m88" }, + ) + + @Test + @DisplayName("Should trigger the run, record its link and succeed without waiting") + fun shouldTriggerAndSucceed() { + val transport = FakeTrigger() + + val result = + run( + TriggerDagRun("downstream").config("conf", mapOf("rows" to 2)), + transport, + env = mapOf("AIRFLOW__API__BASE_URL" to "https://airflow.example.com/"), + ) + + assertInstanceOf(SucceedTask::class.java, result) + val call = transport.triggered.single() + assertEquals("downstream", call.dagId) + assertEquals(mapOf("rows" to 2), call.conf) + // With neither a logical date nor a run_after, the trigger time is both. + assertEquals("manual__2026-10-08T05:00:00+00:00", call.runId) + assertEquals(NOW, call.logicalDate) + assertNull(call.runAfter) + assertEquals( + listOf( + "_link_TriggerDagRunLink" to "https://airflow.example.com/dags/downstream/runs/${call.runId}", + "trigger_run_id" to call.runId, + ), + transport.xComs, + ) + } + + @Test + @DisplayName("Should trigger the run under the ID and settings the task carries") + fun shouldUseConfiguredRunId() { + val transport = FakeTrigger() + val logicalDate = OffsetDateTime.parse("2026-09-30T00:00:00Z") + + run( + TriggerDagRun("downstream") + .config("trigger_run_id", "fixed") + .config("logical_date", logicalDate) + .config("run_after", logicalDate) + .config("reset_dag_run", true) + .config("note", "from java"), + transport, + ) + + val call = transport.triggered.single() + assertEquals("fixed", call.runId) + assertEquals(logicalDate, call.logicalDate) + assertEquals(logicalDate, call.runAfter) + assertTrue(call.resetDagRun) + assertEquals("from java", call.note) + } + + @Test + @DisplayName("Should name the run after the logical date the task set") + fun shouldNameTheRunAfterTheLogicalDate() { + val transport = FakeTrigger() + val logicalDate = OffsetDateTime.parse("2026-09-30T00:00:00Z") + + run(TriggerDagRun("downstream").config("logical_date", logicalDate), transport) + + val call = transport.triggered.single() + assertEquals("manual__2026-09-30T00:00:00+00:00", call.runId) + assertEquals(logicalDate, call.logicalDate) + assertNull(call.runAfter) + } + + @Test + @DisplayName("Should leave a run with only a run_after no logical date, and add a random suffix") + fun shouldLeaveRunAfterOnlyRunWithoutALogicalDate() { + val transport = FakeTrigger() + val runAfter = OffsetDateTime.parse("2026-12-01T00:00:00Z") + + run(TriggerDagRun("downstream").config("run_after", runAfter), transport) + + val call = transport.triggered.single() + assertEquals("manual__2026-12-01T00:00:00+00:00_RaNd0m88", call.runId) + assertNull(call.logicalDate) + assertEquals(runAfter, call.runAfter) + } + + @Test + @DisplayName("Should skip rather than fail an existing run when the task says so") + fun shouldSkipWhenAlreadyExists() { + val transport = FakeTrigger(alreadyExists = true) + + val result = run(TriggerDagRun("downstream").config("skip_when_already_exists", true), transport) + + assertEquals(TaskState.State.SKIPPED, (result as TaskState).state) + } + + @Test + @DisplayName("Should fail an existing run the task does not skip") + fun shouldFailWhenAlreadyExists() { + val transport = FakeTrigger(alreadyExists = true) + + val result = run(TriggerDagRun("downstream"), transport) + + assertEquals(TaskState.State.FAILED, (result as TaskState).state) + assertEquals(listOf("_link_TriggerDagRunLink"), transport.xComs.map { it.first }) + } + + @Test + @DisplayName("Should refuse to trigger a paused Dag when the task says so") + fun shouldFailWhenDagIsPaused() { + val transport = FakeTrigger(paused = true) + + val result = run(TriggerDagRun("downstream").config("fail_when_dag_is_paused", true), transport) + + assertEquals(TaskState.State.FAILED, (result as TaskState).state) + assertTrue(transport.triggered.isEmpty()) + } + + @Test + @DisplayName("Should trigger a Dag that is not paused though the task would fail on a paused one") + fun shouldTriggerUnpausedDagWhenFailingOnPaused() { + val transport = FakeTrigger(paused = false) + + val result = run(TriggerDagRun("downstream").config("fail_when_dag_is_paused", true), transport) + + assertInstanceOf(SucceedTask::class.java, result) + assertEquals(listOf("downstream"), transport.triggered.map { it.dagId }) + } + + @Test + @DisplayName("Should wait until the triggered run reaches an allowed state") + fun shouldWaitForCompletion() { + val transport = FakeTrigger(states = listOf("running", "success")) + val slept = mutableListOf() + + val result = + run( + TriggerDagRun("downstream") + .config("wait_for_completion", true) + .config("poke_interval", Duration.ofSeconds(5)), + transport, + slept = slept, + ) + + assertInstanceOf(SucceedTask::class.java, result) + assertEquals(listOf(Duration.ofSeconds(5), Duration.ofSeconds(5)), slept) + } + + @Test + @DisplayName("Should retry when the triggered run reaches a failed state and the task has tries left") + fun shouldFailWhenRunFails() { + val transport = FakeTrigger(states = listOf("failed")) + + val result = + run( + TriggerDagRun("downstream").config("wait_for_completion", true).config("poke_interval", Duration.ZERO), + transport, + details = startupDetails(shouldRetry = true), + ) + + assertInstanceOf(RetryTask::class.java, result) + } + + @Test + @DisplayName("Should defer the wait to the triggerer when the task is deferrable") + fun shouldDefer() { + val transport = FakeTrigger() + + val result = + run( + TriggerDagRun("downstream") + .config("wait_for_completion", true) + .config("deferrable", true) + .config("poke_interval", Duration.ofSeconds(30)) + .config("failed_states", listOf("failed", "queued")), + transport, + env = mapOf("AIRFLOW__TRIGGERER__QUEUES_ENABLED" to "true"), + ) + + val defer = assertInstanceOf(DeferTask::class.java, result) + assertEquals("airflow.providers.standard.triggers.external_task.DagStateTrigger", defer.classpath) + assertEquals("execute_complete", defer.nextMethod) + assertEquals("workers", defer.queue) + assertEquals( + mapOf( + "dag_id" to "downstream", + "states" to listOf("success", "failed", "queued"), + "poll_interval" to 30, + "run_ids" to listOf(transport.triggered.single().runId), + "execution_dates" to null, + ), + defer.triggerKwargs, + ) + } + + @Test + @DisplayName("Should follow the default deferrable setting when the task sets none") + fun shouldDeferByDefaultWhenConfigured() { + val result = + run( + TriggerDagRun("downstream").config("wait_for_completion", true), + FakeTrigger(), + env = mapOf("AIRFLOW__OPERATORS__DEFAULT_DEFERRABLE" to "true"), + ) + + assertInstanceOf(DeferTask::class.java, result) + } + + @Test + @DisplayName("Should leave the deferred task's queue unset when triggerer queues are not enabled") + fun shouldNotQueueDeferralWithoutTriggererQueues() { + val result = + run( + TriggerDagRun("downstream").config("wait_for_completion", true).config("deferrable", true), + FakeTrigger(), + ) + + assertNull(assertInstanceOf(DeferTask::class.java, result).queue) + } + + @Test + @DisplayName("Should not defer a task that does not wait") + fun shouldNotDeferWithoutWaiting() { + val transport = FakeTrigger() + + val result = run(TriggerDagRun("downstream").config("deferrable", true), transport) + + assertInstanceOf(SucceedTask::class.java, result) + } + + @Test + @DisplayName("Should succeed when the triggerer reports an allowed state") + fun shouldResumeWithAllowedState() { + val transport = FakeTrigger() + val event = + mapOf( + "__classname__" to "builtins.tuple", + "__data__" to listOf("unused", mapOf("run_ids" to listOf("r1"), "r1" to "success")), + ) + + val result = + run( + TriggerDagRun("downstream").config("wait_for_completion", true), + transport, + details = startupDetails(nextMethod = "execute_complete", nextKwargs = mapOf("event" to event)), + ) + + assertInstanceOf(SucceedTask::class.java, result) + assertTrue(transport.triggered.isEmpty()) + } + + @Test + @DisplayName("Should fail when the triggerer reports a failed state") + fun shouldResumeWithFailedState() { + val transport = FakeTrigger() + val event = listOf("unused", mapOf("run_ids" to listOf("r1"), "r1" to "failed")) + + val result = + run( + TriggerDagRun("downstream").config("wait_for_completion", true), + transport, + details = startupDetails(nextMethod = "execute_complete", nextKwargs = mapOf("event" to event)), + ) + + assertEquals(TaskState.State.FAILED, (result as TaskState).state) + } + + @Test + @DisplayName("Should fail when the trigger itself failed") + fun shouldResumeFromTriggerFailure() { + val result = + run( + TriggerDagRun("downstream"), + FakeTrigger(), + details = + startupDetails( + nextMethod = "__fail__", + nextKwargs = mapOf("error" to "boom", "traceback" to listOf("line one", "line two")), + ), + ) + + assertEquals(TaskState.State.FAILED, (result as TaskState).state) + } + + @Test + @DisplayName("Should fail when resumed with a method it does not know") + fun shouldRejectUnknownResumeMethod() { + val result = run(TriggerDagRun("downstream"), FakeTrigger(), details = startupDetails(nextMethod = "other")) + + assertEquals(TaskState.State.FAILED, (result as TaskState).state) + } +} 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 c8ca101267c1b..cafbd5e40d281 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 @@ -130,6 +130,23 @@ internal class ArgValuesTest { override fun clearTaskStateStore(tiId: UUID): Unit = throw NotImplementedError() override fun skipDownstreamTasks(taskIds: List): Unit = throw NotImplementedError() + + override fun triggerDagRun( + dagId: String, + runId: String, + logicalDate: OffsetDateTime?, + runAfter: OffsetDateTime?, + conf: Map?, + resetDagRun: Boolean, + note: String?, + ): Boolean = throw NotImplementedError() + + override fun getDagRunState( + dagId: String, + runId: String, + ): String = throw NotImplementedError() + + override fun isDagPaused(dagId: String): Boolean = throw NotImplementedError() }, ) diff --git a/scripts/ci/lang_sdk_serialization/test_dags.yaml b/scripts/ci/lang_sdk_serialization/test_dags.yaml index ae5a756426d92..174accfdd4f33 100644 --- a/scripts/ci/lang_sdk_serialization/test_dags.yaml +++ b/scripts/ci/lang_sdk_serialization/test_dags.yaml @@ -295,6 +295,9 @@ dags: group: checks # Tasks that trigger other Dags, one with every template field set and one with the defaults. + # logical_date is a template field, so Python writes whatever the author gave it. Spelling it the + # way Python spells a datetime keeps an SDK that takes a date-time and one that passes the field + # through as a string both writing what Python writes. - dag_id: conformance_trigger_dag_run requires: [trigger_dag_run] spec: @@ -305,7 +308,7 @@ dags: trigger_dag_id: conformance_minimal trigger_run_id: "triggered_{{ run_id }}" conf: {rows: 2, ratio: 0.5, nested: {tags: [a, b]}} - logical_date: "{{ ds }}" + logical_date: "2026-09-30 00:00:00+00:00" wait_for_completion: true skip_when_already_exists: true - task_id: trigger_default