This is an automated email from the ASF dual-hosted git repository. jason810496 pushed a commit to branch main in repository https://gitbox.apache.org/repos/asf/airflow.git
commit c5f504366d5835d2028611ee871fe9b4fab4af9f Author: Jason(Zhe-You) Liu <[email protected]> AuthorDate: Fri Oct 9 14:30:53 2026 +0800 Java SDK: Trigger another Dag's run from a native Java Dag (#74440) * Java SDK: Trigger another Dag's run from a native Java Dag A native Java Dag could not start another Dag's run. Airflow's TriggerDagRunOperator is Python, and a Java bundle has no Python Dag file to declare one in, so a Java Dag that fans out to another Dag needed a Python counterpart for that one task. A `TriggerDagRun` passed to `dag.task` declares a task that has no Java body to run; the Java runtime runs it as `TriggerDagRunOperator` does. Its settings are named as Python names the operator's parameters: ```java dag.task("trigger_downstream", new TriggerDagRun("downstream_etl") .config("wait_for_completion", true) .config("conf", Map.of("rows", 2))); ``` An annotated Dag declares one from an ordinary `@Builder.Task` method that returns a `TriggerDagRun`. The serialized Dag has to carry what is triggered, so that method runs when the Dag is built rather than when the task runs, and therefore takes no parameters. The task serializes as Python's own `TriggerDagRunOperator`: its template fields, its UI colour and extra link, and a `dag_dependencies` entry, so the Airflow UI draws the dependency between the two Dags. Python writes only an operator's template fields and keeps the rest in the Dag file; a Java Dag has no such file, so every setting the author made is written too. The runtime records the triggered run's link and ID as XComs, skips or fails on a run that already exists, and waits for completion either by polling or, with `deferrable`, by deferring to the Python triggerer's DagStateTrigger and resuming from the event it fires. That makes this the first Java runtime to emit DeferTask and a skipped TaskState. `scripts/ci/lang_sdk_serialization/test_dags.yaml` now gives the shared trigger Dag a real logical date instead of a Jinja template, because the SDK takes one as a date-time. The Go SDK is making the same change in #74417; whichever lands second drops that one line, and until then the Go serialization conformance check reports that single difference. * Java SDK: Point both example triggers at a Dag that triggers nothing * Java SDK: Say why a switch cannot choose a task that triggers a Dag run * Java SDK: Snapshot a TriggerDagRun when its task is declared * Java SDK: Check a TriggerDagRun setting the way every other config key is checked * Java SDK: Describe what a trigger task pushes and how it names a run * Java SDK: Correct which runtimes already run trigger tasks * Java SDK: Format a trigger task's date-times in one place * Java SDK: Reject a trigger method that throws a checked exception * Java SDK: Give the trigger fixture a logical date every SDK can write * Java SDK: Run a trigger task through runTask in tests * Java SDK: Check a trigger task polls the run it started --- .pre-commit-config.yaml | 2 +- java-sdk/README.md | 4 +- java-sdk/capabilities.yaml | 14 +- .../airflow/example/ExampleBundleBuilder.java | 3 +- .../example/nativedag/AnnotationExample.java | 11 +- .../example/nativedag/InterfaceExample.java | 6 + .../airflow/example/nativedag/TargetExample.java | 49 +++ .../org/apache/airflow/sdk/BuilderProcessor.kt | 58 ++- .../kotlin/org/apache/airflow/sdk/BuilderTest.kt | 203 +++++++++ .../ci/prek/check_serialization_conformance.py | 2 +- java-sdk/sdk/module.md | 4 +- .../main/kotlin/org/apache/airflow/sdk/DagDef.kt | 43 ++ .../main/kotlin/org/apache/airflow/sdk/Switch.kt | 8 +- .../kotlin/org/apache/airflow/sdk/TaskGroupRef.kt | 19 + .../kotlin/org/apache/airflow/sdk/TriggerDagRun.kt | 189 ++++++++ .../org/apache/airflow/sdk/execution/Client.kt | 74 ++++ .../org/apache/airflow/sdk/execution/Logger.kt | 5 + .../org/apache/airflow/sdk/execution/Serde.kt | 127 +++++- .../org/apache/airflow/sdk/execution/Task.kt | 13 + .../apache/airflow/sdk/execution/TriggerRunner.kt | 324 ++++++++++++++ .../org/apache/airflow/sdk/internal/Fields.kt | 33 ++ .../org/apache/airflow/sdk/ArgTestSupport.kt | 17 + .../kotlin/org/apache/airflow/sdk/ClientTest.kt | 17 + .../kotlin/org/apache/airflow/sdk/SwitchTest.kt | 17 + .../org/apache/airflow/sdk/TriggerDagRunTest.kt | 275 ++++++++++++ .../airflow/sdk/conformance/SerializeJava.kt | 40 +- .../org/apache/airflow/sdk/execution/CommTest.kt | 144 +++++++ .../org/apache/airflow/sdk/execution/TaskTest.kt | 71 ++++ .../airflow/sdk/execution/TriggerRunnerTest.kt | 473 +++++++++++++++++++++ .../apache/airflow/sdk/internal/ArgValuesTest.kt | 17 + scripts/ci/lang_sdk_serialization/test_dags.yaml | 5 +- 31 files changed, 2220 insertions(+), 47 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 836f020378a..3bbe42de850 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -461,7 +461,7 @@ repos: ^java-sdk/scripts/ci/prek/check_serialization_conformance\.py$| ^java-sdk/sdk/schema/dag-schema\.json$| ^java-sdk/sdk/build\.gradle\.kts$| - ^java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/(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 9b9653a4974..fda035f80dc 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 5a2ed1f5ae4..84ee13a9dda 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 76175db8a6d..30da2aa8e11 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 class ExampleBundleBuilder { .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 16cdc555452..0bb55ceacd0 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 AnnotationExample { : 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 @@ public class AnnotationExample { 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 b8099931b7a..fcf095ee72e 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 class InterfaceExample { 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 00000000000..cf8d4514295 --- /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 97b810358a0..e101c0e8572 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<String, TaskDeclaration>() - 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 2473ef9561a..5d8c89a058b 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<Void> 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 01a16e34219..677ffcdb217 100755 --- a/java-sdk/scripts/ci/prek/check_serialization_conformance.py +++ b/java-sdk/scripts/ci/prek/check_serialization_conformance.py @@ -56,7 +56,7 @@ if __name__ == "__main__": "--sdk", "java", "--supports", - "literal_inputs,branch,switch", + "literal_inputs,branch,switch,trigger_dag_run", "--", *serializer, ] diff --git a/java-sdk/sdk/module.md b/java-sdk/sdk/module.md index f03b936d7b9..2cc58f52b4d 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 9f79be81a6c..37015ca9197 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<Void> { + 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<out Task>, ) : 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 a668db89609..5bbd93b4da1 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 60d43f84fcd..363f13072cd 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 `<group ID>.<id>`. + * @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<Void> { + 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 00000000000..c78d1a9fd5f --- /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<String, Field> = + 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<String, Any>() + + /** + * 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<String, Any?>()) { (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 f3174e77e2a..317dac7ebb0 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<String>) + + /** + * 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<String, Any?>?, + 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<Any>(taskIds) } runBlocking { exec.communicate<Unit>(message) } } + + override fun triggerDagRun( + dagId: String, + runId: String, + logicalDate: OffsetDateTime?, + runAfter: OffsetDateTime?, + conf: Map<String, Any?>?, + 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<OKResponse>(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<DagRunStateResult>(message) }.state.toString() + } + + override fun isDagPaused(dagId: String): Boolean = + runBlocking { exec.communicate<DagResult>(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 228d59e7240..04757386aa4 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<String, Any> = emptyMap(), ) = log(Level.DEBUG, message, arguments) + fun info( + message: String, + arguments: Map<String, Any> = emptyMap(), + ) = log(Level.INFO, message, arguments) + fun warning( message: String, arguments: Map<String, Any> = 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 bdc719a3a68..e5ad4d6eacf 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<Any?>(), + "dag_dependencies" to serializeDagDependencies(dag), "task_group" to serializeTaskGroups(dag, expansion), "edge_info" to emptyMap<String, Any?>(), "params" to emptyList<Any?>(), @@ -127,22 +128,24 @@ private fun serializeTask( def: TaskDef, downstream: List<String>?, ): Map<String, Any?> { - val data = - linkedMapOf<String, Any?>( - "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<Any?>(), - // 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<String, Any?>("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<Any?>() + // 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<String, Any?>, + 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<Map<String, Any?>> = + 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 182dfe1cee6..54cdb2b3911 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 00000000000..3af0ce30a59 --- /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<String>("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<Duration>("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<String, Any?>() + } + + /** `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<Any?, Any?>() + 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<String, Any?>? { + 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<String> = + trigger.setting<List<String>>("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<String> = trigger.setting<List<String>>("failed_states") ?: listOf("failed") + + private fun deferrable( + trigger: TriggerDagRun, + env: (String) -> String?, + ): Boolean = + trigger.setting<Boolean>("deferrable") + ?: booleanEnv(env, "AIRFLOW__OPERATORS__DEFAULT_DEFERRABLE", false) +} + +@Suppress("UNCHECKED_CAST") +private fun <T> 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 cc43db31865..5b813b5c733 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 6c6af528765..95edc43bbcf 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<String>) = throw NotImplementedError() + + override fun triggerDagRun( + dagId: String, + runId: String, + logicalDate: OffsetDateTime?, + runAfter: OffsetDateTime?, + conf: Map<String, Any?>?, + 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<Map<String, Any?>>?): StartupDetails = diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ClientTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ClientTest.kt index 2ff51df4f36..239e2bbf954 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<String>) = Unit + + override fun triggerDagRun( + dagId: String, + runId: String, + logicalDate: OffsetDateTime?, + runAfter: OffsetDateTime?, + conf: Map<String, Any?>?, + 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 ef72edf3a4c..83bd8305d9f 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 00000000000..56c0bb410b7 --- /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<String, Any?> { + val dag = DagDef("d") + dag.task("trigger", trigger) + val tasks = serializeDag(dag, "", ".")["tasks"] as List<Map<String, Any?>> + return tasks.single()["__var"] as Map<String, Any?> + } + + @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<Map<String, Any?>> + + 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<String, Any>("n" to 1) + val list = mutableListOf<Any>("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 e327afd661e..f7d11aa4284 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<Any?>(localId, definition) - } else { - groups.getValue(groupId).task<Any?>(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<Any?>(localId, definition) + else -> groups.getValue(groupId).task<Any?>(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<Boolean>) @@ -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 1d9b5cc4ace..20fba213948 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 <T> transportRoundTrip( + response: (Int) -> ByteArray, + call: (CoordinatorClient) -> T, + ): Triple<Map<*, *>, T?, Throwable?> { + val toClient = ByteChannel(autoFlush = true) + val fromClient = ByteChannel(autoFlush = true) + val comm = CoordinatorComm(toClient, fromClient) + val client = CoordinatorClient(comm) + + val requests = ConcurrentLinkedQueue<RawFrame>() + 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 db1478661c1..a18907a4e25 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<Pair<String, Any>>(), 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<Unit>("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<String>): Unit = throw UnsupportedOperationException("not used in test") + + override fun triggerDagRun( + dagId: String, + runId: String, + logicalDate: OffsetDateTime?, + runAfter: OffsetDateTime?, + conf: Map<String, Any?>?, + 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<Pair<String, Any>>() val skipped = mutableListOf<String>() val events = mutableListOf<String>() + val triggered = mutableListOf<String>() + 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<String, Any?>?, + 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 00000000000..ad0b629fa2a --- /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<String, Any?>?, + val resetDagRun: Boolean, + val note: String?, +) + +private class FakeTrigger( + private val alreadyExists: Boolean = false, + private val paused: Boolean = false, + private val states: List<String> = emptyList(), +) : org.apache.airflow.sdk.execution.Client { + val triggered = mutableListOf<TriggerCall>() + val xComs = mutableListOf<Pair<String, Any>>() + var polls = 0 + + override fun triggerDagRun( + dagId: String, + runId: String, + logicalDate: OffsetDateTime?, + runAfter: OffsetDateTime?, + conf: Map<String, Any?>?, + 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<String>): 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<String, String> = emptyMap(), + slept: MutableList<Duration> = 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<Duration>() + + 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 c8ca101267c..cafbd5e40d2 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<String>): Unit = throw NotImplementedError() + + override fun triggerDagRun( + dagId: String, + runId: String, + logicalDate: OffsetDateTime?, + runAfter: OffsetDateTime?, + conf: Map<String, Any?>?, + 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 ae5a756426d..174accfdd4f 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
