You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by ja...@apache.org on 2026/10/09 06:31:09 UTC
(airflow) branch main updated (6e5d9ddc185 -> c5f504366d5)
This is an automated email from the ASF dual-hosted git repository.
jason810496 pushed a change to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
from 6e5d9ddc185 Update the TypeScript SDK docs for Dags declared in TypeScript (#73875)
new 2bda1928703 Java SDK: Pack each native Dag's source file into the bundle JAR (#74096)
new d4013cc2b5b Java SDK: Pick one of two tasks with a condition (#74438)
new c3043c266c6 Java SDK: Choose one of several tasks with a branch (#74439)
new c5f504366d5 Java SDK: Trigger another Dag's run from a native Java Dag (#74440)
The 4 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
Summary of changes:
.pre-commit-config.yaml | 2 +-
.../0015-per-dag-source-in-bundle-artifact.md | 2 +
.../language-sdks/java.rst | 4 +
java-sdk/README.md | 10 +-
java-sdk/capabilities.yaml | 18 +-
.../airflow/example/ExampleBundleBuilder.java | 3 +-
.../example/nativedag/AnnotationExample.java | 44 +-
.../example/nativedag/InterfaceExample.java | 57 +-
.../airflow/example/nativedag/TargetExample.java | 49 ++
java-sdk/plugin/build.gradle.kts | 6 +
.../apache/airflow/sdk/plugin/AirflowSdkPlugin.kt | 55 +-
.../org/apache/airflow/sdk/plugin/ClassFiles.kt | 85 +++
.../apache/airflow/sdk/plugin/PackDagSources.kt | 256 ++++++++
.../airflow/sdk/plugin/AirflowSdkPluginTest.kt | 365 +++++++++++
.../apache/airflow/sdk/plugin/ClassFilesTest.kt | 80 +++
.../airflow/sdk/plugin/PackDagSourcesTest.kt | 89 +++
.../org/apache/airflow/sdk/plugin/TestSupport.kt} | 38 +-
.../org/apache/airflow/sdk/BuilderProcessor.kt | 209 ++++++-
.../kotlin/org/apache/airflow/sdk/BuilderTest.kt | 692 ++++++++++++++++++++-
.../ci/prek/check_serialization_conformance.py | 2 +-
java-sdk/sdk/build.gradle.kts | 68 ++
java-sdk/sdk/module.md | 6 +-
.../main/kotlin/org/apache/airflow/sdk/Bundle.kt | 27 +-
.../kotlin/org/apache/airflow/sdk/Condition.kt | 224 +++++++
.../main/kotlin/org/apache/airflow/sdk/DagDef.kt | 140 +++++
.../main/kotlin/org/apache/airflow/sdk/Decider.kt | 90 +++
.../main/kotlin/org/apache/airflow/sdk/Server.kt | 109 +++-
.../main/kotlin/org/apache/airflow/sdk/Switch.kt | 240 +++++++
.../kotlin/org/apache/airflow/sdk/TaskGroupRef.kt | 81 +++
.../kotlin/org/apache/airflow/sdk/TriggerDagRun.kt | 189 ++++++
.../org/apache/airflow/sdk/execution/Client.kt | 88 +++
.../org/apache/airflow/sdk/execution/Logger.kt | 5 +
.../org/apache/airflow/sdk/execution/Serde.kt | 129 +++-
.../org/apache/airflow/sdk/execution/Task.kt | 42 +-
.../apache/airflow/sdk/execution/TriggerRunner.kt | 324 ++++++++++
.../org/apache/airflow/sdk/internal/DagSource.kt | 66 ++
.../org/apache/airflow/sdk/internal/Fields.kt | 33 +
.../kotlin/org/apache/airflow/sdk/internal/Ids.kt | 10 +
.../apache/airflow/example/DagSourceFixtures.kt} | 49 +-
.../org/apache/airflow/sdk/ArgTestSupport.kt | 19 +
.../kotlin/org/apache/airflow/sdk/ClientTest.kt | 19 +
.../kotlin/org/apache/airflow/sdk/ConditionTest.kt | 208 +++++++
.../kotlin/org/apache/airflow/sdk/DagSourceTest.kt | 132 ++++
.../kotlin/org/apache/airflow/sdk/ServerTest.kt | 6 +-
.../kotlin/org/apache/airflow/sdk/SwitchTest.kt | 224 +++++++
.../org/apache/airflow/sdk/TriggerDagRunTest.kt | 275 ++++++++
.../airflow/sdk/conformance/SerializeJava.kt | 191 +++++-
.../org/apache/airflow/sdk/execution/CommTest.kt | 171 +++++
.../org/apache/airflow/sdk/execution/TaskTest.kt | 319 ++++++++++
.../airflow/sdk/execution/TriggerRunnerTest.kt | 473 ++++++++++++++
.../apache/airflow/sdk/internal/ArgValuesTest.kt | 19 +
.../org/apache/airflow/sdk/internal/RefsTest.kt | 65 ++
scripts/ci/lang_sdk_serialization/test_dags.yaml | 5 +-
.../airflow/sdk/coordinators/java/_dag_importer.py | 4 +-
.../coordinators/java/test_dag_importer.py | 19 +
55 files changed, 5967 insertions(+), 168 deletions(-)
create mode 100644 java-sdk/example/src/java/org/apache/airflow/example/nativedag/TargetExample.java
create mode 100644 java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/ClassFiles.kt
create mode 100644 java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/PackDagSources.kt
create mode 100644 java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPluginTest.kt
create mode 100644 java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/ClassFilesTest.kt
create mode 100644 java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/PackDagSourcesTest.kt
copy java-sdk/{example/src/java/org/apache/airflow/example/ExampleBundleBuilder.java => plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/TestSupport.kt} (50%)
create mode 100644 java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Condition.kt
create mode 100644 java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Decider.kt
create mode 100644 java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Switch.kt
create mode 100644 java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/TriggerDagRun.kt
create mode 100644 java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/TriggerRunner.kt
create mode 100644 java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/DagSource.kt
copy java-sdk/sdk/src/{testFixtures/kotlin/org/apache/airflow/sdk/execution/LogCapture.kt => test/kotlin/org/apache/airflow/example/DagSourceFixtures.kt} (57%)
create mode 100644 java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ConditionTest.kt
create mode 100644 java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/DagSourceTest.kt
create mode 100644 java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/SwitchTest.kt
create mode 100644 java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/TriggerDagRunTest.kt
create mode 100644 java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TriggerRunnerTest.kt
(airflow) 03/04: Java SDK: Choose one of several tasks with a branch (#74439)
Posted by ja...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
jason810496 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
commit c3043c266c677a45531d4eae6a2e26ddd61d3096
Author: Jason(Zhe-You) Liu <68...@users.noreply.github.com>
AuthorDate: Fri Oct 9 14:30:52 2026 +0800
Java SDK: Choose one of several tasks with a branch (#74439)
* Java SDK: Choose one of several tasks with a branch
A condition picks between two tasks, but a Dag that routes work by size or
by region needs more than two, and splitting that into nested conditions
makes the graph unreadable.
`DagDef.Branch` and `@Builder.Branch` declare a task that chooses one of
the tasks listed with `option(...)`; every other one is skipped. The chosen
task's ID becomes the branch's return value, as it does for Python's branch
operators. A branch chooses exactly one case, so a task that runs after
several of them needs a trigger rule such as
`none_failed_min_one_success`.
Each authoring style names the case in the way that style can type-check. A
`BranchTask` returns the case's own class, so javac checks it; two cases
that share a class are rejected where the second is named, since the branch
could not tell them apart. A `@Builder.Branch` method has no such class to
name -- a task method's implementation is generated -- so it returns a
`TaskId` from the generated `<Dag>Builder.TaskIds`, which holds one
constant per task of the Dag. Either way nothing hard-codes a task id
string, and the generated holder is emitted only for a Dag that declares a
branch.
* Java SDK: Name the multi-way branch Switch, as the other SDKs do
* Java SDK: Name a switch case by its task class in both styles
* Java SDK: Say when a switch rejects a choice and when it fails
* Java SDK: Cover a generated switch that names a task it cannot choose
* Java SDK: Check a switch carries the chosen task's id
* Java SDK: Let a wiring view name a switch's cases in more than one call
---
.pre-commit-config.yaml | 2 +-
java-sdk/README.md | 2 +-
java-sdk/capabilities.yaml | 2 +-
.../example/nativedag/AnnotationExample.java | 20 ++
.../example/nativedag/InterfaceExample.java | 32 ++-
.../org/apache/airflow/sdk/BuilderProcessor.kt | 108 +++++++---
.../kotlin/org/apache/airflow/sdk/BuilderTest.kt | 212 ++++++++++++++++++-
.../ci/prek/check_serialization_conformance.py | 2 +-
java-sdk/sdk/build.gradle.kts | 35 +++
java-sdk/sdk/module.md | 2 +-
.../main/kotlin/org/apache/airflow/sdk/Bundle.kt | 7 +-
.../kotlin/org/apache/airflow/sdk/Condition.kt | 2 +-
.../main/kotlin/org/apache/airflow/sdk/DagDef.kt | 43 +++-
.../main/kotlin/org/apache/airflow/sdk/Switch.kt | 234 +++++++++++++++++++++
.../kotlin/org/apache/airflow/sdk/TaskGroupRef.kt | 30 +++
.../kotlin/org/apache/airflow/sdk/ConditionTest.kt | 2 +-
.../kotlin/org/apache/airflow/sdk/SwitchTest.kt | 207 ++++++++++++++++++
.../airflow/sdk/conformance/SerializeJava.kt | 140 +++++++++++-
.../org/apache/airflow/sdk/execution/TaskTest.kt | 98 +++++++++
.../org/apache/airflow/sdk/internal/RefsTest.kt | 36 ++++
20 files changed, 1165 insertions(+), 51 deletions(-)
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 37f87710dfc..836f020378a 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -461,7 +461,7 @@ repos:
^java-sdk/scripts/ci/prek/check_serialization_conformance\.py$|
^java-sdk/sdk/schema/dag-schema\.json$|
^java-sdk/sdk/build\.gradle\.kts$|
- ^java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/(Bundle|Condition|DagDef|Decider|Deps|Arg|TaskGroupRef)\.kt$|
+ ^java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/(Switch|Bundle|Condition|DagDef|Decider|Deps|Arg|TaskGroupRef)\.kt$|
^java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Serde\.kt$|
^java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/(Refs|Fields)\.kt$|
^java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/conformance/.*$|
diff --git a/java-sdk/README.md b/java-sdk/README.md
index 1e4e7478f39..9b9653a4974 100644
--- a/java-sdk/README.md
+++ b/java-sdk/README.md
@@ -643,7 +643,7 @@ prek hook regenerate it.
| capability: `task-args` | MUST † | ✓ | 3.4 | |
| capability: `dag-params` | MUST † | ✗ | – | |
| capability: `taskflow-dependencies` | MUST † | ✓ | 3.4 | |
-| capability: `branching` | SHOULD † | ✓ | 3.4 | DagDef.If / @Builder.If |
+| capability: `branching` | SHOULD † | ✓ | 3.4 | DagDef.If and DagDef.Switch, or @Builder.If and @Builder.Switch |
| capability: `dag-test` | SHOULD † | ✗ | – | |
| capability: `task-group` | MAY † | ✓ | 3.4 | |
| capability: `dynamic-task-mapping` | MAY † | ✗ | – | |
diff --git a/java-sdk/capabilities.yaml b/java-sdk/capabilities.yaml
index 33106f170fb..5a2ed1f5ae4 100644
--- a/java-sdk/capabilities.yaml
+++ b/java-sdk/capabilities.yaml
@@ -109,7 +109,7 @@ capabilities:
branching:
supported: true
since: "3.4"
- note: "DagDef.If / @Builder.If"
+ note: "DagDef.If and DagDef.Switch, or @Builder.If and @Builder.Switch"
dag-test:
supported: false
task-group:
diff --git a/java-sdk/example/src/java/org/apache/airflow/example/nativedag/AnnotationExample.java b/java-sdk/example/src/java/org/apache/airflow/example/nativedag/AnnotationExample.java
index d9538fd391f..16cdc555452 100644
--- a/java-sdk/example/src/java/org/apache/airflow/example/nativedag/AnnotationExample.java
+++ b/java-sdk/example/src/java/org/apache/airflow/example/nativedag/AnnotationExample.java
@@ -66,6 +66,25 @@ public class AnnotationExample {
return transformed > 0;
}
+ @Builder.Task(id = "report_long")
+ public void reportLong() {
+ log.log(INFO, "Long report");
+ }
+
+ @Builder.Task(id = "report_short")
+ public void reportShort() {
+ log.log(INFO, "Short report");
+ }
+
+ // A switch: it names the one case that runs by the class generated for it,
+ // and every other case is skipped.
+ @Builder.Switch(id = "pick_report")
+ public Class<? extends Task> pickReport(long transformed) {
+ return transformed > 100
+ ? AnnotationExampleBuilder.ReportLong.class
+ : AnnotationExampleBuilder.ReportShort.class;
+ }
+
// A task group: everything it declares is prefixed with its id, so this is
// the task "checks.audit".
@Builder.TaskGroup(id = "checks")
@@ -84,6 +103,7 @@ public class AnnotationExample {
var extracted = extract();
var transformed = transform(extracted, lit(1.5));
hasRows(transformed).Then(load(transformed)).Else(loadEmpty());
+ pickReport(transformed).Case(reportLong()).Case(reportShort());
// Ordering-only edge: the checks group runs after extract, with no data
// flowing.
extracted.before(checks());
diff --git a/java-sdk/example/src/java/org/apache/airflow/example/nativedag/InterfaceExample.java b/java-sdk/example/src/java/org/apache/airflow/example/nativedag/InterfaceExample.java
index cce10676dfd..b8099931b7a 100644
--- a/java-sdk/example/src/java/org/apache/airflow/example/nativedag/InterfaceExample.java
+++ b/java-sdk/example/src/java/org/apache/airflow/example/nativedag/InterfaceExample.java
@@ -71,6 +71,31 @@ public class InterfaceExample {
}
}
+ public static class ReportLong implements Task {
+ @Override
+ public void execute(Context context, Client client) {
+ log.log(INFO, "Long report");
+ }
+ }
+
+ public static class ReportShort implements Task {
+ @Override
+ public void execute(Context context, Client client) {
+ log.log(INFO, "Short report");
+ }
+ }
+
+ // A switch: it names the one case that runs by its class, and every other
+ // case is skipped.
+ public static class PickReport implements SwitchTask {
+ @Override
+ public Class<? extends Task> choose(Context context, Client client) {
+ return ((Number) client.getXCom("transform")).longValue() > 100
+ ? ReportLong.class
+ : ReportShort.class;
+ }
+ }
+
public static DagDef build() {
var dag =
new DagDef("java_native_interface_example")
@@ -87,9 +112,14 @@ public class InterfaceExample {
var load = dag.task("load", Load.class);
var loadEmpty = dag.task("load_empty", LoadEmpty.class);
+ var reportLong = dag.task("report_long", ReportLong.class);
+ var reportShort = dag.task("report_short", ReportShort.class);
+
transform.after(extract);
- // With no task id given, the condition takes one from its class: "hasRows".
+ // With no task id given, a decider takes one from its class: "hasRows" and
+ // "pickReport".
dag.If(HasRows.class).after(transform).Then(load).Else(loadEmpty);
+ dag.Switch(PickReport.class).after(transform).Case(reportLong).Case(reportShort);
return dag;
}
}
diff --git a/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt b/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt
index 1848fc007f7..97b810358a0 100644
--- a/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt
+++ b/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt
@@ -49,6 +49,7 @@ import javax.annotation.processing.RoundEnvironment
import javax.annotation.processing.SupportedAnnotationTypes
import javax.annotation.processing.SupportedSourceVersion
import javax.lang.model.SourceVersion
+import javax.lang.model.element.AnnotationMirror
import javax.lang.model.element.AnnotationValue
import javax.lang.model.element.Element
import javax.lang.model.element.ElementKind
@@ -56,8 +57,10 @@ import javax.lang.model.element.ExecutableElement
import javax.lang.model.element.Modifier
import javax.lang.model.element.TypeElement
import javax.lang.model.element.VariableElement
+import javax.lang.model.type.DeclaredType
import javax.lang.model.type.TypeKind
import javax.lang.model.type.TypeMirror
+import javax.lang.model.type.WildcardType
import javax.tools.Diagnostic
import java.lang.reflect.Modifier as ReflectModifier
import org.apache.airflow.sdk.internal.builderName as generatedBuilderName
@@ -97,6 +100,7 @@ import org.apache.airflow.sdk.internal.builderName as generatedBuilderName
"org.apache.airflow.sdk.Builder.Dag",
"org.apache.airflow.sdk.Builder.Task",
"org.apache.airflow.sdk.Builder.If",
+ "org.apache.airflow.sdk.Builder.Switch",
"org.apache.airflow.sdk.Builder.TaskGroup",
"org.apache.airflow.sdk.Builder.TaskHandler",
"org.apache.airflow.sdk.Builder.Deps",
@@ -357,24 +361,22 @@ class BuilderProcessor : AbstractProcessor() {
builderName: ClassName,
inGroup: Boolean,
): MethodSpec {
- val condition = decl.kind == TaskKind.CONDITION
val method =
MethodSpec
.methodBuilder(decl.method.simpleName.toString())
.addModifiers(Modifier.PUBLIC, Modifier.DEFAULT)
.returns(
- if (condition) {
- CONDITION_REF_TYPE
- } else {
- ParameterizedTypeName.get(TASK_HANDLE_TYPE, TypeName.get(decl.method.returnType).boxIfPossible())
- },
+ decl.kind.refType
+ ?: ParameterizedTypeName.get(TASK_HANDLE_TYPE, TypeName.get(decl.method.returnType).boxIfPossible()),
)
decl.dataParams.forEach { method.addParameter(inType(it.type), it.name) }
val def = taskDefCode(decl, CodeBlock.of($$"$T.$L", builderName, decl.className))
// The view knows the group it belongs to, so the recorder is told where
// the task goes instead of deriving it from the task's ID.
val group = if (inGroup) CodeBlock.of("groupId()") else CodeBlock.of($$"$S", "")
- val wrap = { call: CodeBlock -> if (condition) CodeBlock.of($$"$T.of($L)", CONDITION_REF_TYPE, call) else call }
+ val wrap = { call: CodeBlock ->
+ decl.kind.refType?.let { CodeBlock.of($$"$T.of($L)", it, call) } ?: call
+ }
if (decl.dataParams.isEmpty()) {
method.addStatement($$"return $L", wrap(CodeBlock.of($$"$T.node($L, $L)", REFS_TYPE, group, def)))
} else {
@@ -435,17 +437,19 @@ class BuilderProcessor : AbstractProcessor() {
val tasks = mutableListOf<TaskDeclaration>()
for (inner in el.enclosedElements) {
if (inner !is ExecutableElement) continue
- val task = inner.getAnnotation(Builder.Task::class.java)
- val condition = inner.getAnnotation(Builder.If::class.java)
- if (task == null && condition == null) continue
- require(task == null || condition == null) {
- "Method '${inner.simpleName}' carries both @Builder.Task and @Builder.If; a condition is a " +
- "task, so it carries @Builder.If alone"
- }
- val kind = if (condition != null) TaskKind.CONDITION else TaskKind.TASK
+ val declared =
+ TaskKind.entries.filter { kind -> inner.annotationMirrors.any { it.names(kind.annotation) } }
+ if (declared.isEmpty()) continue
+ val kind =
+ declared.singleOrNull()
+ ?: throw IllegalArgumentException(
+ "Method '${inner.simpleName}' carries ${declared.joinToString { it.spelling }}; a task is " +
+ "declared by one of them alone",
+ )
+ val annotated = declaredId(inner, kind)
if (inner.isVarArgs) throw IllegalArgumentException("Cannot create task from vararg function ${inner.simpleName}")
checkDeciderReturn(kind, inner)
- val localId = (condition?.id ?: task!!.id).ifBlank { inner.simpleName.toString() }
+ val localId = annotated.ifBlank { inner.simpleName.toString() }
require(tasks.none { it.method.simpleName.contentEquals(inner.simpleName) }) {
"Class ${el.simpleName} overloads task method '${inner.simpleName}'; a method's name is the " +
"name of its generated task class and of its wiring-view method, so rename one and keep its " +
@@ -712,24 +716,55 @@ class BuilderProcessor : AbstractProcessor() {
kind: TaskKind,
method: ExecutableElement,
) {
- if (kind != TaskKind.CONDITION) return
val returns = method.returnType
- val boolean = returns.kind == TypeKind.BOOLEAN || with(processingEnv) { isType(returns, BOXED_BOOLEAN_TYPE) }
- require(boolean) {
- "@Builder.If method '${method.simpleName}' returns $returns, but a condition returns boolean: " +
- "true runs the task named by Then, false the one named by Else"
+ when (kind) {
+ TaskKind.TASK -> return
+ TaskKind.CONDITION ->
+ require(returns.kind == TypeKind.BOOLEAN || with(processingEnv) { isType(returns, BOXED_BOOLEAN_TYPE) }) {
+ "@Builder.If method '${method.simpleName}' returns $returns, but a condition returns boolean: " +
+ "true runs the task named by Then, false the one named by Else"
+ }
+ TaskKind.SWITCH ->
+ require(returnsTaskClass(returns)) {
+ "@Builder.Switch method '${method.simpleName}' returns $returns, but a switch returns the " +
+ "class of the task it chose, as Class<? extends Task>"
+ }
}
}
+ /** Whether [type] is `Class<T>` or `Class<? extends T>` for a [Task] type `T`. */
+ private fun returnsTaskClass(type: TypeMirror): Boolean {
+ val declared = type as? DeclaredType ?: return false
+ if (!(declared.asElement() as TypeElement).qualifiedName.contentEquals("java.lang.Class")) return false
+ val argument = declared.typeArguments.singleOrNull() ?: return false
+ val bound = if (argument is WildcardType) argument.extendsBound else argument
+ return bound != null &&
+ with(processingEnv) { typeUtils.isAssignable(bound, elementUtils.getTypeElement(TASK_TYPE.canonicalName()).asType()) }
+ }
+
+ /** The `id` the declaring annotation sets, empty when it leaves it out. */
+ private fun declaredId(
+ method: ExecutableElement,
+ kind: TaskKind,
+ ): String {
+ val mirror = method.annotationMirrors.first { it.names(kind.annotation) }
+ val id = mirror.elementValues.entries.firstOrNull { it.key.simpleName.contentEquals("id") }
+ return id?.value?.value as String? ?: ""
+ }
+
private fun buildTask(decl: TaskDeclaration): TypeSpec {
- val condition = decl.kind == TaskKind.CONDITION
val executeSpec =
MethodSpec
- .methodBuilder(if (condition) "decide" else "execute")
+ .methodBuilder(decl.kind.bodyMethod)
.addAnnotation(Override::class.java)
.addModifiers(Modifier.PUBLIC)
- .returns(if (condition) TypeName.BOOLEAN else TypeName.VOID)
- .addParameter(CONTEXT_TYPE, "context")
+ .returns(
+ when (decl.kind) {
+ TaskKind.TASK -> TypeName.VOID
+ TaskKind.CONDITION -> TypeName.BOOLEAN
+ TaskKind.SWITCH -> TASK_CLASS_TYPE
+ },
+ ).addParameter(CONTEXT_TYPE, "context")
.addParameter(CLIENT_TYPE, "client")
.addException(Exception::class.java)
@@ -775,8 +810,8 @@ class BuilderProcessor : AbstractProcessor() {
}
when {
- // The SDK pushes a condition's result itself, after it has skipped the side not taken.
- condition -> $$"return new $T().$L($L)"
+ // The SDK pushes a decider's choice itself, after it has skipped what the choice rules out.
+ decl.kind != TaskKind.TASK -> $$"return new $T().$L($L)"
inner.returnType.kind == TypeKind.VOID -> $$"new $T().$L($L)"
else -> $$"client.setXCom(new $T().$L($L))"
}.also {
@@ -790,7 +825,7 @@ class BuilderProcessor : AbstractProcessor() {
return TypeSpec
.classBuilder(decl.className)
- .addSuperinterface(if (condition) CONDITION_TASK_TYPE else ClassName.get(Task::class.java))
+ .addSuperinterface(decl.kind.taskInterface)
.addModifiers(Modifier.PUBLIC, Modifier.FINAL, Modifier.STATIC)
.addMethod(executeSpec.build())
.build()
@@ -966,8 +1001,13 @@ private val DAG_SOURCE_TYPE = ClassName.get(DagSource::class.java)
private val REFS_TYPE = ClassName.get(Refs::class.java)
private val ARG_TYPE = ClassName.get(Arg::class.java)
private val TASK_HANDLE_TYPE = ClassName.get(TaskRef::class.java)
+private val TASK_TYPE = ClassName.get(Task::class.java)
private val CONDITION_TASK_TYPE = ClassName.get(ConditionTask::class.java)
private val CONDITION_REF_TYPE = ClassName.get(ConditionRef::class.java)
+private val SWITCH_TASK_TYPE = ClassName.get(SwitchTask::class.java)
+private val SWITCH_REF_TYPE = ClassName.get(SwitchRef::class.java)
+private val TASK_CLASS_TYPE =
+ ParameterizedTypeName.get(ClassName.get(Class::class.java), WildcardTypeName.subtypeOf(TASK_TYPE))
private val BOXED_BOOLEAN_TYPE = ClassName.get("java.lang", "Boolean")
private val DEPS_TYPE = ClassName.get(Deps::class.java)
private val GROUP_TYPE = DEPS_TYPE.nestedClass("TaskGroup")
@@ -983,11 +1023,19 @@ private const val DAG_ANNOTATION = "org.apache.airflow.sdk.Builder.Dag"
private enum class TaskKind(
val annotation: String,
val spelling: String,
+ /** Method the generated class implements, and the type its wiring view hands back. */
+ val bodyMethod: String,
+ val taskInterface: ClassName,
+ val refType: ClassName?,
) {
- TASK("org.apache.airflow.sdk.Builder.Task", "@Builder.Task"),
- CONDITION("org.apache.airflow.sdk.Builder.If", "@Builder.If"),
+ TASK("org.apache.airflow.sdk.Builder.Task", "@Builder.Task", "execute", TASK_TYPE, null),
+ CONDITION("org.apache.airflow.sdk.Builder.If", "@Builder.If", "decide", CONDITION_TASK_TYPE, CONDITION_REF_TYPE),
+ SWITCH("org.apache.airflow.sdk.Builder.Switch", "@Builder.Switch", "choose", SWITCH_TASK_TYPE, SWITCH_REF_TYPE),
}
+/** Whether this annotation is the one [name] qualifies. */
+private fun AnnotationMirror.names(name: String): Boolean = (annotationType.asElement() as TypeElement).qualifiedName.contentEquals(name)
+
private val RESERVED_VIEW_NAMES =
setOf(
"depends",
diff --git a/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt b/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt
index 70802c405aa..2473ef9561a 100644
--- a/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt
+++ b/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt
@@ -2300,6 +2300,216 @@ class BuilderTest {
assertThat(compilation).succeeded()
}
+ @Test
+ @DisplayName("let a wiring view name a switch's cases in more than one statement")
+ fun switchNamedAcrossStatements() {
+ val compilation =
+ compile(
+ """
+ package org.apache.airflow.example;
+
+ import org.apache.airflow.sdk.Builder;
+ import org.apache.airflow.sdk.Task;
+
+ @Builder.Dag(id = "etl")
+ public class TestExample {
+ @Builder.Switch(id = "pick_path")
+ public Class<? extends Task> pickPath() {
+ return TestExampleBuilder.HandleLong.class;
+ }
+
+ @Builder.Task
+ public void handleLong() {}
+
+ @Builder.Task
+ public void handleShort() {}
+
+ @Builder.Deps
+ static class Wiring implements TestExampleDeps {
+ void depends() {
+ pickPath().Case(handleLong());
+ pickPath().Case(handleShort());
+ }
+ }
+ }
+ """,
+ )
+
+ assertThat(compilation).succeeded()
+ }
+
+ @Test
+ @DisplayName("generate a switch that names its case by the class generated for it")
+ fun generateSwitchTask() {
+ val compilation =
+ compile(
+ """
+ package org.apache.airflow.example;
+
+ import org.apache.airflow.sdk.Builder;
+ import org.apache.airflow.sdk.Task;
+
+ @Builder.Dag(id = "etl")
+ public class TestExample {
+ @Builder.Switch(id = "pick_path")
+ public Class<? extends Task> pickPath() {
+ return TestExampleBuilder.HandleLong.class;
+ }
+
+ @Builder.Task
+ public void handleLong() {}
+
+ @Builder.Task
+ public void handleShort() {}
+
+ @Builder.Deps
+ static class Wiring implements TestExampleDeps {
+ void depends() {
+ pickPath().Case(handleLong()).Case(handleShort());
+ }
+ }
+ }
+ """,
+ )
+
+ assertThat(compilation).succeeded()
+ assertThat(compilation)
+ .generatedSourceFile("org.apache.airflow.example.TestExampleBuilder")
+ .hasSourceEquivalentTo(
+ "org.apache.airflow.example.TestExampleBuilder",
+ """
+ package org.apache.airflow.example;
+
+ import java.lang.Class;
+ import java.lang.Exception;
+ import java.lang.Override;
+ import java.util.List;
+ import org.apache.airflow.sdk.Client;
+ import org.apache.airflow.sdk.Context;
+ import org.apache.airflow.sdk.DagDef;
+ import org.apache.airflow.sdk.SwitchTask;
+ import org.apache.airflow.sdk.Task;
+ import org.apache.airflow.sdk.internal.DagSource;
+ import org.apache.airflow.sdk.internal.Refs;
+
+ public final class TestExampleBuilder {
+ public static DagDef build() {
+ var dag = DagSource.declaredBy(new DagDef("etl"), TestExample.class);
+ return Refs.record(dag, List.of("pick_path", "handleLong", "handleShort"), List.of(), new TestExample.Wiring()::depends);
+ }
+
+ public static final class PickPath implements SwitchTask {
+ @Override
+ public Class<? extends Task> choose(Context context, Client client) throws Exception {
+ return new TestExample().pickPath();
+ }
+ }
+
+ public static final class HandleLong implements Task {
+ @Override
+ public void execute(Context context, Client client) throws Exception {
+ new TestExample().handleLong();
+ }
+ }
+
+ public static final class HandleShort implements Task {
+ @Override
+ public void execute(Context context, Client client) throws Exception {
+ new TestExample().handleShort();
+ }
+ }
+ }
+ """,
+ )
+ assertThat(compilation)
+ .generatedSourceFile("org.apache.airflow.example.TestExampleDeps")
+ .hasSourceEquivalentTo(
+ "org.apache.airflow.example.TestExampleDeps",
+ """
+ package org.apache.airflow.example;
+
+ import java.lang.Void;
+ import org.apache.airflow.sdk.Deps;
+ import org.apache.airflow.sdk.SwitchRef;
+ import org.apache.airflow.sdk.TaskDef;
+ import org.apache.airflow.sdk.TaskRef;
+ import org.apache.airflow.sdk.internal.Refs;
+
+ public interface TestExampleDeps extends Deps {
+ default SwitchRef pickPath() {
+ return SwitchRef.of(Refs.node("", new TaskDef("pick_path", TestExampleBuilder.PickPath.class)));
+ }
+
+ default TaskRef<Void> handleLong() {
+ return Refs.node("", new TaskDef("handleLong", TestExampleBuilder.HandleLong.class));
+ }
+
+ default TaskRef<Void> handleShort() {
+ return Refs.node("", new TaskDef("handleShort", TestExampleBuilder.HandleShort.class));
+ }
+ }
+ """,
+ )
+ }
+
+ @Test
+ @DisplayName("reject a switch that does not return the class of a task")
+ fun rejectNonTaskClassSwitch() {
+ val compilation =
+ compile(
+ """
+ package org.apache.airflow.example;
+ import org.apache.airflow.sdk.Builder;
+ @Builder.Dag
+ public class TestExample {
+ @Builder.Switch
+ public String pickPath() {
+ return "handleLong";
+ }
+
+ @Builder.Deps
+ static class Wiring implements TestExampleDeps {
+ void depends() {}
+ }
+ }
+ """,
+ )
+
+ assertThat(compilation).failed()
+ assertThat(compilation).hadErrorContaining(
+ "@Builder.Switch method 'pickPath' returns java.lang.String, but a switch returns the class of the task it chose",
+ )
+ }
+
+ @Test
+ @DisplayName("reject a switch that returns a class that is not a task")
+ fun rejectSwitchReturningNonTaskClass() {
+ val compilation =
+ compile(
+ """
+ package org.apache.airflow.example;
+ import org.apache.airflow.sdk.Builder;
+ @Builder.Dag
+ public class TestExample {
+ @Builder.Switch
+ public Class<String> pickPath() {
+ return String.class;
+ }
+
+ @Builder.Deps
+ static class Wiring implements TestExampleDeps {
+ void depends() {}
+ }
+ }
+ """,
+ )
+
+ assertThat(compilation).failed()
+ assertThat(compilation).hadErrorContaining(
+ "@Builder.Switch method 'pickPath' returns java.lang.Class<java.lang.String>, but a switch returns the class of the task it chose",
+ )
+ }
+
@Test
@DisplayName("reject a condition that does not return a boolean")
fun rejectNonBooleanCondition() {
@@ -2355,7 +2565,7 @@ class BuilderTest {
assertThat(compilation).failed()
assertThat(compilation).hadErrorContaining(
- "Method 'hasRows' carries both @Builder.Task and @Builder.If",
+ "Method 'hasRows' carries @Builder.Task, @Builder.If; a task is declared by one of them alone",
)
}
diff --git a/java-sdk/scripts/ci/prek/check_serialization_conformance.py b/java-sdk/scripts/ci/prek/check_serialization_conformance.py
index fe569c1ff7c..01a16e34219 100755
--- a/java-sdk/scripts/ci/prek/check_serialization_conformance.py
+++ b/java-sdk/scripts/ci/prek/check_serialization_conformance.py
@@ -56,7 +56,7 @@ if __name__ == "__main__":
"--sdk",
"java",
"--supports",
- "literal_inputs,branch",
+ "literal_inputs,branch,switch",
"--",
*serializer,
]
diff --git a/java-sdk/sdk/build.gradle.kts b/java-sdk/sdk/build.gradle.kts
index 723f60085c2..e30facde61e 100644
--- a/java-sdk/sdk/build.gradle.kts
+++ b/java-sdk/sdk/build.gradle.kts
@@ -614,6 +614,41 @@ abstract class GenerateDagDslTask : DefaultTask() {
| )
|
| /**
+ | * Marks a task method that chooses one of several tasks to run;
+ | * every other one is skipped.
+ | *
+ | * The method returns the class the Dag's builder generated for the task it
+ | * chose, as `Class<? extends Task>`. Its wiring-view method returns a
+ | * [SwitchRef], so the cases are listed where the Dag is wired:
+ | *
+ | * ```java
+ | * @Builder.Switch(id = "pick_path")
+ | * public Class<? extends Task> pickPath(long rows) {
+ | * return rows > 1000 ? EtlBuilder.HandleLong.class : EtlBuilder.HandleShort.class;
+ | * }
+ | *
+ | * @Builder.Deps
+ | * static class Wiring implements EtlDeps {
+ | * void depends() { pickPath(extract()).Case(handleLong()).Case(handleShort()); }
+ | * }
+ | * ```
+ | *
+ | * A switch chooses exactly one case, so a task that runs after
+ | * several of them needs a trigger rule that tolerates a skipped
+ | * upstream, such as `none_failed_min_one_success`.
+ | *
+ | * Configuration attributes are [Task]'s, and apply to the deciding
+ | * task.
+ | */
+ | @Target(AnnotationTarget.FUNCTION)
+ | @MustBeDocumented
+ | annotation class Switch(
+ | /** Task ID. Empty derives it from the annotated function's name. */
+ | val id: String = "",
+ |${attrLines(taskFields)}
+ | )
+ |
+ | /**
| * Marks a method as the Java body of a task the Python Dag file
| * declares with `@task.stub`.
| *
diff --git a/java-sdk/sdk/module.md b/java-sdk/sdk/module.md
index 523dbaf54a4..f03b936d7b9 100644
--- a/java-sdk/sdk/module.md
+++ b/java-sdk/sdk/module.md
@@ -60,7 +60,7 @@ meaning of each dimension is defined in the
| capability: `task-args` | MUST † | ✓ | 3.4 | |
| capability: `dag-params` | MUST † | ✗ | – | |
| capability: `taskflow-dependencies` | MUST † | ✓ | 3.4 | |
-| capability: `branching` | SHOULD † | ✓ | 3.4 | DagDef.If / @Builder.If |
+| capability: `branching` | SHOULD † | ✓ | 3.4 | DagDef.If and DagDef.Switch, or @Builder.If and @Builder.Switch |
| capability: `dag-test` | SHOULD † | ✗ | – | |
| capability: `task-group` | MAY † | ✓ | 3.4 | |
| capability: `dynamic-task-mapping` | MAY † | ✗ | – | |
diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Bundle.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Bundle.kt
index 7937446c3f6..5e817c05e93 100644
--- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Bundle.kt
+++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Bundle.kt
@@ -67,7 +67,8 @@ class Bundle(
* @throws IllegalArgumentException if another Dag shares its ID, task
* handlers are already registered against it, a task depends on an
* upstream not registered in the same Dag, a task that decides which task
- * runs is not declared with `If`, or the dependencies contain a cycle.
+ * runs is not declared with `If` or `Switch`, or the dependencies contain
+ * a cycle.
* @throws IllegalStateException if [Server.serve] has already been called.
*/
fun register(dag: DagDef): Bundle {
@@ -87,7 +88,7 @@ class Bundle(
require(decides == (def.decider != null)) {
if (decides) {
"Task '${def.id}' runs '${def.definition.name}', which decides which task runs, but nothing names " +
- "what it chooses; declare it with If(...), in Dag '${dag.id}'"
+ "what it chooses; declare it with If(...) or Switch(...), in Dag '${dag.id}'"
} else {
"Task '${def.id}' chooses which task runs, but '${def.definition.name}' decides nothing, in Dag '${dag.id}'"
}
@@ -212,7 +213,7 @@ class Bundle(
}
/** Task types the SDK runs through a decider, so each needs the sides it chooses between. */
-private val DECIDER_TYPES = listOf(ConditionTask::class.java)
+private val DECIDER_TYPES = listOf(ConditionTask::class.java, SwitchTask::class.java)
// Reject cycles produced by before and after at registration time. This is (non-tailrec-eligible)
// recursive and could blow up with deep dependency chains. I kept the recursive implementation
diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Condition.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Condition.kt
index 90150609c7b..ea521f92028 100644
--- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Condition.kt
+++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Condition.kt
@@ -103,7 +103,7 @@ class ConditionRef private constructor(
@JvmStatic
fun of(ref: TaskRef<Boolean>): ConditionRef {
ref.def.decider?.let { existing ->
- require(existing is ConditionDef) { "Task '${ref.def.id}' already decides what to skip as a branch" }
+ require(existing is ConditionDef) { "Task '${ref.def.id}' already decides what to skip as a switch" }
return ConditionRef(ref, existing)
}
return ConditionRef(ref, ConditionDef().also { ref.def.decider = it })
diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/DagDef.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/DagDef.kt
index 40ead4ae608..9f79be81a6c 100644
--- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/DagDef.kt
+++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/DagDef.kt
@@ -201,6 +201,47 @@ class DagDef(
definition: Class<out ConditionTask>,
): ConditionRef = ConditionRef.of(task(id, definition))
+ /**
+ * Declares a task that chooses one of several tasks to run; every other one
+ * is skipped.
+ *
+ * The task's ID is the class's simple name with its first character
+ * lowercased, so `PickPath.class` becomes `pickPath`. Use
+ * [Switch(id, definition)][Switch] to set it, and [SwitchRef.config] for the
+ * task's other settings:
+ *
+ * ```java
+ * dag.Switch(PickPath.class).Case(handleLong).Case(handleShort);
+ * ```
+ *
+ * @param definition Class that implements [SwitchTask]. Must have a public
+ * no-arg constructor.
+ * @return The switch, to list its cases on.
+ * @throws IllegalArgumentException if a task with the same ID is already
+ * registered.
+ */
+ @Suppress("ktlint:standard:function-naming")
+ fun Switch(definition: Class<out SwitchTask>): SwitchRef = Switch(deriveTaskId(definition), definition)
+
+ /**
+ * Declares a task that chooses one of several tasks to run, under the task
+ * ID [id].
+ *
+ * @param id Task identifier, unique within this Dag.
+ * @param definition Class that implements [SwitchTask]. Must have a public
+ * no-arg constructor.
+ * @return The switch, to list its cases on.
+ * @throws IllegalArgumentException if a task with the same ID is already
+ * registered.
+ *
+ * @see Switch
+ */
+ @Suppress("ktlint:standard:function-naming")
+ fun Switch(
+ id: String,
+ definition: Class<out SwitchTask>,
+ ): SwitchRef = SwitchRef.of(task<Any?>(id, definition))
+
/**
* Declares a task group of this Dag.
*
@@ -375,7 +416,7 @@ class TaskDef(
internal val upstreams = linkedSetOf<TaskDef>()
internal var owner: DagDef? = null
- /** What this task decides to run, for a condition or a branch; null otherwise. */
+ /** What this task decides to run, for a condition or a switch; null otherwise. */
internal var decider: DeciderDef? = null
/**
diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Switch.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Switch.kt
new file mode 100644
index 00000000000..a668db89609
--- /dev/null
+++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Switch.kt
@@ -0,0 +1,234 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.airflow.sdk
+
+import kotlin.Throws
+
+/**
+ * A task that chooses one of several tasks to run; every other one is skipped.
+ *
+ * Register one with [DagDef.Switch], then list what it can choose with
+ * [SwitchRef.Case]. The choice is the class of the task to run:
+ *
+ * ```java
+ * public class PickPath implements SwitchTask {
+ * @Override
+ * public Class<? extends Task> choose(Context context, Client client) {
+ * return ((Number) client.getXCom("extract")).longValue() > 1000
+ * ? HandleLong.class
+ * : HandleShort.class;
+ * }
+ * }
+ * ```
+ *
+ * The SDK runs [choose] and pushes the chosen task's ID as this task's return
+ * value, so [execute] is never called.
+ *
+ * Only a [Task] class compiles as a choice. javac cannot tell whether it is a
+ * case of this switch, so a Task class that is not one compiles and fails when
+ * the task runs: the task instance is marked failed, and nothing is pushed or
+ * skipped.
+ *
+ * Because a switch names a case by its class, two tasks that run the same
+ * class cannot both be cases of it; the second [SwitchRef.Case] call throws.
+ *
+ * @see DagDef.Switch
+ */
+interface SwitchTask : Task {
+ /**
+ * Chooses the one case that runs.
+ *
+ * Any exception thrown marks the task instance as failed, and nothing is
+ * then pushed or skipped.
+ *
+ * @param context Runtime context for the current execution workload.
+ * @param client Client for Airflow API calls scoped to this execution.
+ * @return The class of one of this switch's cases.
+ * @throws Exception on failure; the task instance is marked failed.
+ */
+ @Throws(Exception::class)
+ fun choose(
+ context: Context,
+ client: Client,
+ ): Class<out Task>
+
+ /** Never called: the SDK runs a switch through [choose]. */
+ override fun execute(
+ context: Context,
+ client: Client,
+ ): Unit =
+ throw IllegalStateException(
+ "Switch '${javaClass.name}' runs through choose(), so execute() is never called",
+ )
+}
+
+/**
+ * A switch registered with a Dag: list the tasks it can choose between.
+ *
+ * ```java
+ * dag.Switch(PickPath.class).Case(handleLong).Case(handleShort);
+ * ```
+ *
+ * A case runs after the switch, so naming it records that edge, as
+ * [Deps.Flow.before] would. The switch chooses exactly one case and skips
+ * every other, so a task that runs after several cases needs a trigger rule
+ * that tolerates a skipped upstream, such as `none_failed_min_one_success`.
+ *
+ * @see DagDef.Switch
+ */
+class SwitchRef private constructor(
+ private val ref: TaskRef<*>,
+ private val cases: SwitchDef,
+) : Deps.Flow {
+ companion object {
+ /**
+ * @suppress
+ *
+ * Marks a registered task as a switch. Public so a generated wiring view
+ * can call it; user code reaches a switch through [DagDef.Switch].
+ *
+ * A wiring view gets the same task back each time it names it, so a second
+ * call reuses the switch rather than declaring it again. [DagDef.Switch] and
+ * [TaskGroupRef.Switch] never reach that path, since registering a task
+ * rejects a duplicate ID first.
+ */
+ @JvmStatic
+ fun of(ref: TaskRef<*>): SwitchRef {
+ ref.def.decider?.let { existing ->
+ require(existing is SwitchDef) { "Task '${ref.def.id}' already decides what to skip as a condition" }
+ return SwitchRef(ref, existing)
+ }
+ val definition = ref.def.definition
+ require(SwitchTask::class.java.isAssignableFrom(definition)) {
+ "Task '${ref.def.id}' runs '${definition.name}', which chooses nothing; a switch runs a " +
+ "SwitchTask"
+ }
+ return SwitchRef(ref, SwitchDef().also { ref.def.decider = it })
+ }
+ }
+
+ /** Airflow task ID of the deciding task. */
+ val id: String get() = ref.def.id
+
+ /**
+ * Adds a task this switch can choose. Every other case is skipped when the
+ * switch chooses this one.
+ *
+ * @param task Task of the same Dag.
+ * @return This switch, for chaining.
+ * @throws IllegalArgumentException if the task belongs to another Dag, is
+ * already a case of this switch, or runs the same class as one.
+ */
+ @Suppress("ktlint:standard:function-naming")
+ fun Case(task: TaskRef<*>): SwitchRef {
+ cases.add(ref.def, task.def)
+ // The case runs after the switch, which is what puts the switch in the
+ // serialized Dag as its upstream and lets registration see a cycle that
+ // runs through a switch.
+ ref.before(task)
+ return this
+ }
+
+ /**
+ * Sets one task-level configuration value on the deciding task.
+ *
+ * @param key Airflow task setting name.
+ * @param value Value matching the key's schema type.
+ * @return This switch, for chaining.
+ * @throws IllegalArgumentException if the key is unknown or the value type
+ * does not match.
+ */
+ fun config(
+ key: String,
+ value: Any?,
+ ): SwitchRef {
+ ref.config(key, value)
+ return this
+ }
+
+ override fun nodes(): List<TaskDef> = ref.nodes()
+
+ override fun before(vararg next: Deps.Flow): SwitchRef {
+ ref.before(*next)
+ return this
+ }
+
+ override fun after(vararg previous: Deps.Flow): SwitchRef {
+ ref.after(*previous)
+ return this
+ }
+}
+
+/** The cases of one switch, as [SwitchRef.Case] listed them. */
+internal class SwitchDef : DeciderDef {
+ private val options = mutableListOf<TaskDef>()
+
+ override val cases: List<TaskDef> get() = options
+
+ fun add(
+ decider: TaskDef,
+ case: TaskDef,
+ ) {
+ requireUnregistered(decider, "Switch")
+ requireSameDag(decider, case, "Switch")
+ require(options.none { it === case }) {
+ "Switch '${decider.id}' already chooses between '${case.id}' and others; name each case once"
+ }
+ // A switch names its case by class, so two cases sharing one would be
+ // indistinguishable.
+ options.firstOrNull { it.definition == case.definition }?.let { first ->
+ throw IllegalArgumentException(
+ "Switch '${decider.id}' cannot choose between '${first.id}' and '${case.id}': both run " +
+ "'${case.definition.name}', and a switch names its case by class",
+ )
+ }
+ options += case
+ }
+
+ override fun describe(decider: TaskDef): String? =
+ if (options.isEmpty()) "Switch '${decider.id}' has no task to choose between; call Case(...)" else null
+
+ override fun decide(
+ decider: TaskDef,
+ instance: Task,
+ context: Context,
+ client: Client,
+ ): Decision {
+ val chosen =
+ when (instance) {
+ is SwitchTask -> byClass(decider, instance.choose(context, client))
+ else -> throw IllegalStateException(
+ "Task '${decider.id}' is a switch, but '${instance.javaClass.name}' is not a SwitchTask",
+ )
+ }
+ return Decision(chosen.id, options.filterNot { it === chosen }.map { it.id })
+ }
+
+ private fun byClass(
+ decider: TaskDef,
+ definition: Class<out Task>?,
+ ): TaskDef =
+ options.firstOrNull { it.definition == definition }
+ ?: throw IllegalArgumentException(
+ "Switch '${decider.id}' chose ${definition?.name ?: "nothing"}, which is not one of its cases: ${describeCases()}",
+ )
+
+ private fun describeCases(): String = options.joinToString { "'${it.id}'" }
+}
diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/TaskGroupRef.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/TaskGroupRef.kt
index e2d7d9613ef..60d43f84fcd 100644
--- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/TaskGroupRef.kt
+++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/TaskGroupRef.kt
@@ -106,6 +106,36 @@ class TaskGroupRef internal constructor(
definition: Class<out ConditionTask>,
): ConditionRef = ConditionRef.of(task(id, definition))
+ /**
+ * Declares a switch in this group, as [DagDef.Switch] does for the Dag.
+ *
+ * @param definition Class that implements [SwitchTask]. Must have a public
+ * no-arg constructor.
+ * @return The switch, to list its cases on.
+ * @throws IllegalArgumentException if the Dag already has a task or task
+ * group with the resulting ID.
+ */
+ @Suppress("ktlint:standard:function-naming")
+ fun Switch(definition: Class<out SwitchTask>): SwitchRef = Switch(deriveTaskId(definition), definition)
+
+ /**
+ * Declares a switch in this group under the task ID `<group ID>.<id>`.
+ *
+ * @param id Task ID within this group.
+ * @param definition Class that implements [SwitchTask]. Must have a public
+ * no-arg constructor.
+ * @return The switch, to list its cases on.
+ * @throws IllegalArgumentException if the Dag already has a task or task
+ * group with the resulting ID.
+ *
+ * @see Switch
+ */
+ @Suppress("ktlint:standard:function-naming")
+ fun Switch(
+ id: String,
+ definition: Class<out SwitchTask>,
+ ): SwitchRef = SwitchRef.of(task<Any?>(id, definition))
+
/**
* Nests a task group inside this one.
*
diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ConditionTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ConditionTest.kt
index 1cebf7ff8e1..0fb13877837 100644
--- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ConditionTest.kt
+++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ConditionTest.kt
@@ -141,7 +141,7 @@ internal class ConditionTest {
assertEquals(
"Task 'has_rows' runs '${HasRows::class.java.name}', which decides which task runs, but nothing names " +
- "what it chooses; declare it with If(...), in Dag 'd'",
+ "what it chooses; declare it with If(...) or Switch(...), in Dag 'd'",
error.message,
)
}
diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/SwitchTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/SwitchTest.kt
new file mode 100644
index 00000000000..ef72edf3a4c
--- /dev/null
+++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/SwitchTest.kt
@@ -0,0 +1,207 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.airflow.sdk
+
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Assertions.assertThrows
+import org.junit.jupiter.api.DisplayName
+import org.junit.jupiter.api.Test
+
+/** A case of the switches under test; a second class so a switch can tell two cases apart. */
+internal class HandleLong : Task {
+ override fun execute(
+ context: Context,
+ client: Client,
+ ) = Unit
+}
+
+/** A switch that always chooses [HandleLong]. */
+class PickPath : SwitchTask {
+ override fun choose(
+ context: Context,
+ client: Client,
+ ): Class<out Task> = HandleLong::class.java
+}
+
+internal class SwitchTest {
+ private fun dagWithCases(): Triple<DagDef, SwitchRef, Pair<TaskRef<Unit>, TaskRef<Unit>>> {
+ val dag = DagDef("d")
+ val handleLong = dag.task<Unit>("handle_long", HandleLong::class.java)
+ val handleShort = dag.task<Unit>("handle_short", NoopTask::class.java)
+ return Triple(dag, dag.Switch(PickPath::class.java), handleLong to handleShort)
+ }
+
+ @Test
+ @DisplayName("Should take the switch's task ID from its class when none is given")
+ fun shouldDeriveTaskIdFromClass() {
+ val (dag, switch, cases) = dagWithCases()
+ switch.Case(cases.first)
+
+ assertEquals("pickPath", switch.id)
+ assertEquals(setOf(dag.tasks.getValue("pickPath")), cases.first.def.upstreams)
+ }
+
+ @Test
+ @DisplayName("Should reject naming the same case twice")
+ fun shouldRejectDuplicateCase() {
+ val (_, switch, cases) = dagWithCases()
+ switch.Case(cases.first)
+
+ val error = assertThrows(IllegalArgumentException::class.java) { switch.Case(cases.first) }
+
+ assertEquals(
+ "Switch 'pickPath' already chooses between 'handle_long' and others; name each case once",
+ error.message,
+ )
+ }
+
+ @Test
+ @DisplayName("Should reject two cases a class-named switch could not tell apart")
+ fun shouldRejectCasesSharingAClass() {
+ val dag = DagDef("d")
+ val first = dag.task<Unit>("first", HandleLong::class.java)
+ val second = dag.task<Unit>("second", HandleLong::class.java)
+ val switch = dag.Switch(PickPath::class.java).Case(first)
+
+ val error = assertThrows(IllegalArgumentException::class.java) { switch.Case(second) }
+
+ assertEquals(
+ "Switch 'pickPath' cannot choose between 'first' and 'second': both run " +
+ "'org.apache.airflow.sdk.HandleLong', and a switch names its case by class",
+ error.message,
+ )
+ }
+
+ @Test
+ @DisplayName("Should reject a case declared in another Dag")
+ fun shouldRejectCaseFromAnotherDag() {
+ val (_, switch, _) = dagWithCases()
+ val other = DagDef("other").task<Unit>("handle_long", HandleLong::class.java)
+
+ val error = assertThrows(IllegalArgumentException::class.java) { switch.Case(other) }
+
+ assertEquals(
+ "Switch 'pickPath' of Dag 'd' cannot run task 'handle_long' of Dag 'other'; " +
+ "name a task of the same Dag",
+ error.message,
+ )
+ }
+
+ @Test
+ @DisplayName("Should reject registering a switch with nothing to choose between")
+ fun shouldRejectSwitchWithoutCases() {
+ val (dag, _, _) = dagWithCases()
+
+ val error = assertThrows(IllegalArgumentException::class.java) { Bundle().register(dag) }
+
+ assertEquals(
+ "Switch 'pickPath' has no task to choose between; call Case(...), in Dag 'd'",
+ error.message,
+ )
+ }
+
+ @Test
+ @DisplayName("Should reject a switch over a task that chooses nothing")
+ fun shouldRejectSwitchOverAPlainTask() {
+ val dag = DagDef("d")
+
+ val error =
+ assertThrows(IllegalArgumentException::class.java) {
+ SwitchRef.of(dag.task<Unit>("plain", NoopTask::class.java))
+ }
+
+ assertEquals(
+ "Task 'plain' runs 'org.apache.airflow.sdk.NoopTask', which chooses nothing; " +
+ "a switch runs a SwitchTask",
+ error.message,
+ )
+ }
+
+ @Test
+ @DisplayName("Should reject naming a case after the Dag was registered")
+ fun shouldRejectCaseAfterRegistration() {
+ val (dag, switch, cases) = dagWithCases()
+ switch.Case(cases.first)
+ Bundle().register(dag)
+
+ val error = assertThrows(IllegalArgumentException::class.java) { switch.Case(cases.second) }
+
+ assertEquals(
+ "Switch 'pickPath' of Dag 'd' is already registered; name every case before the Dag is " +
+ "added to a Bundle",
+ error.message,
+ )
+ }
+
+ @Test
+ @DisplayName("Should declare a switch inside a task group under the group's prefix")
+ fun shouldDeclareSwitchInGroup() {
+ val dag = DagDef("d")
+ val reports = dag.taskGroup("reports")
+ val long = reports.task<Unit>("long", HandleLong::class.java)
+ val switch = reports.Switch("pick", PickPath::class.java).Case(long)
+
+ assertEquals("reports.pick", switch.id)
+ assertEquals(listOf("reports.long", "reports.pick"), dag.tasks.keys.toList())
+ }
+
+ @Test
+ @DisplayName("Should record the chosen case's task ID and skip the other case")
+ fun shouldRecordChosenTaskIdAndSkipTheOther() {
+ val (dag, switch, cases) = dagWithCases()
+ switch.Case(cases.first).Case(cases.second)
+ val pick = dag.tasks.getValue("pickPath")
+
+ val decision = pick.decider!!.decide(pick, PickPath(), taskContext(), clientWith(null).first)
+
+ assertEquals("handle_long", decision.value)
+ assertEquals(listOf("handle_short"), decision.skipped)
+ }
+
+ @Test
+ @DisplayName("Should reuse the switch when a task already declared as one is declared again")
+ fun shouldReuseSwitchDeclaredTwice() {
+ val (dag, switch, cases) = dagWithCases()
+ val again = SwitchRef.of(TaskRef<Unit>(dag.tasks.getValue("pickPath")))
+
+ switch.Case(cases.first)
+ again.Case(cases.second)
+
+ assertEquals(
+ listOf("handle_long", "handle_short"),
+ dag.tasks
+ .getValue("pickPath")
+ .decider!!
+ .cases
+ .map { it.id },
+ )
+ }
+
+ @Test
+ @DisplayName("Should reject declaring a condition as a switch")
+ fun shouldRejectSwitchOverACondition() {
+ val dag = DagDef("d")
+ val condition = dag.If(HasRows::class.java)
+
+ val error = assertThrows(IllegalArgumentException::class.java) { SwitchRef.of(TaskRef<Unit>(condition.nodes().single())) }
+
+ assertEquals("Task 'hasRows' already decides what to skip as a condition", error.message)
+ }
+}
diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/conformance/SerializeJava.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/conformance/SerializeJava.kt
index 4c2f7f84e4e..e327afd661e 100644
--- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/conformance/SerializeJava.kt
+++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/conformance/SerializeJava.kt
@@ -39,6 +39,8 @@ import org.apache.airflow.sdk.ConditionTask
import org.apache.airflow.sdk.Context
import org.apache.airflow.sdk.DagDef
import org.apache.airflow.sdk.Deps
+import org.apache.airflow.sdk.SwitchRef
+import org.apache.airflow.sdk.SwitchTask
import org.apache.airflow.sdk.Task
import org.apache.airflow.sdk.TaskGroupRef
import org.apache.airflow.sdk.TaskRef
@@ -50,13 +52,108 @@ import java.io.File
import java.time.Duration
import java.time.OffsetDateTime
-class ConformanceTask : Task {
+// A switch names its cases by class, so each task of a Dag runs a class of its own. The pool has
+// room for the largest Dag in test_dags.yaml and a few more; the serialized task type is not compared.
+class ConformanceTask1 : Task {
override fun execute(
context: Context,
client: Client,
) = Unit
}
+class ConformanceTask2 : Task {
+ override fun execute(
+ context: Context,
+ client: Client,
+ ) = Unit
+}
+
+class ConformanceTask3 : Task {
+ override fun execute(
+ context: Context,
+ client: Client,
+ ) = Unit
+}
+
+class ConformanceTask4 : Task {
+ override fun execute(
+ context: Context,
+ client: Client,
+ ) = Unit
+}
+
+class ConformanceTask5 : Task {
+ override fun execute(
+ context: Context,
+ client: Client,
+ ) = Unit
+}
+
+class ConformanceTask6 : Task {
+ override fun execute(
+ context: Context,
+ client: Client,
+ ) = Unit
+}
+
+class ConformanceTask7 : Task {
+ override fun execute(
+ context: Context,
+ client: Client,
+ ) = Unit
+}
+
+class ConformanceTask8 : Task {
+ override fun execute(
+ context: Context,
+ client: Client,
+ ) = Unit
+}
+
+class ConformanceTask9 : Task {
+ override fun execute(
+ context: Context,
+ client: Client,
+ ) = Unit
+}
+
+class ConformanceTask10 : Task {
+ override fun execute(
+ context: Context,
+ client: Client,
+ ) = Unit
+}
+
+class ConformanceTask11 : Task {
+ override fun execute(
+ context: Context,
+ client: Client,
+ ) = Unit
+}
+
+class ConformanceTask12 : Task {
+ override fun execute(
+ context: Context,
+ client: Client,
+ ) = Unit
+}
+
+private val taskPool: List<Class<out Task>> =
+ listOf(
+ ConformanceTask1::class.java,
+ ConformanceTask2::class.java,
+ ConformanceTask3::class.java,
+ ConformanceTask4::class.java,
+ ConformanceTask5::class.java,
+ ConformanceTask6::class.java,
+ ConformanceTask7::class.java,
+ ConformanceTask8::class.java,
+ ConformanceTask9::class.java,
+ ConformanceTask10::class.java,
+ ConformanceTask11::class.java,
+ ConformanceTask12::class.java,
+ )
+
class ConformanceCondition : ConditionTask {
override fun decide(
context: Context,
@@ -64,6 +161,14 @@ class ConformanceCondition : ConditionTask {
) = true
}
+// Never run, only serialized, so any task class will do.
+class ConformanceSwitch : SwitchTask {
+ override fun choose(
+ context: Context,
+ client: Client,
+ ): Class<out Task> = ConformanceTask1::class.java
+}
+
fun main(args: Array<String>) {
require(args.size == 2) { "usage: SerializeJava <test_dags.yaml> <output.json>" }
val cases = YAMLMapper().readTree(File(args[0])).path("dags")
@@ -90,20 +195,33 @@ private fun buildDag(case: JsonNode): DagDef {
}
val tasks = linkedMapOf<String, TaskRef<*>>()
- // A condition names tasks that may be declared after it, so the sides are wired once every task exists.
- val conditions = mutableListOf<Pair<ConditionRef, JsonNode>>()
+ val unusedTaskClasses = taskPool.iterator()
+ // A decider names tasks that may be declared after it, so its cases are wired once every task exists.
+ val deciders = mutableListOf<Pair<Deps.Flow, JsonNode>>()
case.path("tasks").forEach { task ->
val groupId = task.path("group").asText("")
val localId = task.path("task_id").asText()
val branch = task.path("branch")
- val definition = if (branch.isMissingNode) ConformanceTask::class.java else ConformanceCondition::class.java
+ val definition =
+ when {
+ branch.isMissingNode -> {
+ require(unusedTaskClasses.hasNext()) {
+ "Dag '${dag.id}' has more tasks than the ${taskPool.size} classes in taskPool; add a ConformanceTask class"
+ }
+ unusedTaskClasses.next()
+ }
+ branch.has("cases") -> ConformanceSwitch::class.java
+ else -> ConformanceCondition::class.java
+ }
val ref =
if (groupId.isEmpty()) {
dag.task<Any?>(localId, definition)
} else {
groups.getValue(groupId).task<Any?>(localId, definition)
}
- if (!branch.isMissingNode) conditions += asCondition(ref) to branch
+ if (!branch.isMissingNode) {
+ deciders += (if (branch.has("cases")) SwitchRef.of(ref) else asCondition(ref)) to branch
+ }
task.path("spec").fields().forEach { (key, value) -> ref.config(key, toValue(SchemaFields.TASK, key, value)) }
// A task's `upstream` handles and its `literals` are its call arguments, in that order, so the
// Dag carries the binding spec a stub call would. Names are positional, as the Go SDK names
@@ -117,9 +235,15 @@ private fun buildDag(case: JsonNode): DagDef {
tasks[ref.def.id] = ref
}
- conditions.forEach { (condition, branch) ->
- condition.Then(tasks.getValue(branch.path("then").asText()))
- branch.path("else").takeIf { !it.isMissingNode }?.let { condition.Else(tasks.getValue(it.asText())) }
+ deciders.forEach { (decider, branch) ->
+ when (decider) {
+ is SwitchRef -> branch.path("cases").forEach { decider.Case(tasks.getValue(it.asText())) }
+ is ConditionRef -> {
+ decider.Then(tasks.getValue(branch.path("then").asText()))
+ branch.path("else").takeIf { !it.isMissingNode }?.let { decider.Else(tasks.getValue(it.asText())) }
+ }
+ else -> throw IllegalStateException("Unknown decider: $decider")
+ }
}
case.path("order_edges").forEach { edge ->
diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskTest.kt
index a1ca3f7553d..db1478661c1 100644
--- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskTest.kt
+++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskTest.kt
@@ -24,6 +24,7 @@ import org.apache.airflow.sdk.Client
import org.apache.airflow.sdk.ConditionTask
import org.apache.airflow.sdk.Context
import org.apache.airflow.sdk.DagDef
+import org.apache.airflow.sdk.SwitchTask
import org.apache.airflow.sdk.Task
import org.apache.airflow.sdk.TaskDef
import org.apache.airflow.sdk.execution.comm.BundleInfo
@@ -245,6 +246,62 @@ class TaskTest {
Assertions.assertEquals(emptyList<Pair<String, Any>>(), transport.xComs)
}
+ @Test
+ @DisplayName("Should push the case a switch chose and skip every other one")
+ fun shouldSkipTheCasesNotChosen() {
+ TestSwitch.choice = HandleLongTask::class.java
+ val transport = RecordingTransport()
+
+ val result = runTask(switchBundle(), startupDetails(taskId = "pick"), Client(startupDetails("pick"), transport))
+
+ Assertions.assertInstanceOf(SucceedTask::class.java, result)
+ Assertions.assertEquals(listOf("handle_short"), transport.skipped)
+ Assertions.assertEquals(
+ listOf("skipmixin_key" to mapOf("skipped" to listOf("handle_short")), "return_value" to "handle_long"),
+ transport.xComs,
+ )
+ }
+
+ @Test
+ @DisplayName("Should fail a switch that chose a task it cannot run")
+ fun shouldFailSwitchThatChoseAnUnknownCase() {
+ TestSwitch.choice = SuccessTask::class.java
+ val transport = RecordingTransport()
+
+ val result = runTask(switchBundle(), startupDetails(taskId = "pick"), Client(startupDetails("pick"), transport))
+
+ Assertions.assertInstanceOf(TaskState::class.java, result)
+ Assertions.assertEquals(TaskState.State.FAILED, (result as TaskState).state)
+ Assertions.assertEquals(emptyList<String>(), transport.skipped)
+ Assertions.assertEquals(emptyList<Pair<String, Any>>(), transport.xComs)
+ }
+
+ @Test
+ @DisplayName("Should fail a generated switch that chose a task of the Dag that is not one of its cases")
+ fun shouldFailGeneratedSwitchThatChoseANonCase() {
+ val dag = DagDef("test_dag")
+ val long = dag.task<Unit>("handle_long", HandleLongTask::class.java)
+ val short = dag.task<Unit>("handle_short", HandleShortTask::class.java)
+ dag.task<Unit>("unlisted", UnlistedTask::class.java)
+ dag.Switch("pick", TestGeneratedSwitch::class.java).Case(long).Case(short)
+ val transport = RecordingTransport()
+
+ val result = runTask(Bundle(listOf(dag)), startupDetails(taskId = "pick"), Client(startupDetails("pick"), transport))
+
+ Assertions.assertInstanceOf(TaskState::class.java, result)
+ Assertions.assertEquals(TaskState.State.FAILED, (result as TaskState).state)
+ Assertions.assertEquals(emptyList<String>(), transport.skipped)
+ Assertions.assertEquals(emptyList<Pair<String, Any>>(), transport.xComs)
+ }
+
+ private fun switchBundle(): Bundle {
+ val dag = DagDef("test_dag")
+ val long = dag.task<Unit>("handle_long", HandleLongTask::class.java)
+ val short = dag.task<Unit>("handle_short", HandleShortTask::class.java)
+ dag.Switch("pick", TestSwitch::class.java).Case(long).Case(short)
+ return Bundle(listOf(dag))
+ }
+
private fun conditionBundle(withElse: Boolean): Bundle {
val dag = DagDef("test_dag")
val load = dag.task<Unit>("load", SuccessTask::class.java)
@@ -360,6 +417,47 @@ class TaskTest {
}
}
+ class HandleLongTask : Task {
+ override fun execute(
+ context: Context,
+ client: Client,
+ ) = Unit
+ }
+
+ class HandleShortTask : Task {
+ override fun execute(
+ context: Context,
+ client: Client,
+ ) = Unit
+ }
+
+ /** Chooses whatever [choice] holds, so a test can drive it to a case or past one. */
+ class TestSwitch : SwitchTask {
+ override fun choose(
+ context: Context,
+ client: Client,
+ ): Class<out Task> = choice
+
+ companion object {
+ var choice: Class<out Task> = HandleLongTask::class.java
+ }
+ }
+
+ class UnlistedTask : Task {
+ override fun execute(
+ context: Context,
+ client: Client,
+ ) = Unit
+ }
+
+ /** The shape the annotation processor generates: it returns the class of a task, here one that is no case. */
+ class TestGeneratedSwitch : SwitchTask {
+ override fun choose(
+ context: Context,
+ client: Client,
+ ): Class<out Task> = UnlistedTask::class.java
+ }
+
/** Records what a deciding task pushed and asked to skip, and the order of the two. */
private class RecordingTransport : org.apache.airflow.sdk.execution.Client {
val xComs = mutableListOf<Pair<String, Any>>()
diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/RefsTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/RefsTest.kt
index 59e36c182bd..c764585e409 100644
--- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/RefsTest.kt
+++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/RefsTest.kt
@@ -27,6 +27,8 @@ import org.apache.airflow.sdk.Context
import org.apache.airflow.sdk.DagDef
import org.apache.airflow.sdk.Deps
import org.apache.airflow.sdk.LiteralArg
+import org.apache.airflow.sdk.SwitchRef
+import org.apache.airflow.sdk.SwitchTask
import org.apache.airflow.sdk.Task
import org.apache.airflow.sdk.TaskDef
import org.apache.airflow.sdk.TaskRef
@@ -56,6 +58,20 @@ private class NoopCondition : ConditionTask {
) = true
}
+private class OtherNoopRefTask : Task {
+ override fun execute(
+ context: Context,
+ client: Client,
+ ) = Unit
+}
+
+private class NoopSwitch : SwitchTask {
+ override fun choose(
+ context: Context,
+ client: Client,
+ ): Class<out Task> = NoopRefTask::class.java
+}
+
internal class RefsTest {
@Test
@DisplayName("Should register the task, record inputs, and wire handle edges")
@@ -232,4 +248,24 @@ internal class RefsTest {
.map { it.id },
)
}
+
+ @Test
+ @DisplayName("Should let a view name a switch's cases in separate statements")
+ fun shouldNameSwitchCasesAcrossStatements() {
+ val dag = DagDef("d")
+ Refs.record(dag, listOf("pick", "long", "short"), emptyList()) {
+ fun pick() = SwitchRef.of(Refs.node<Unit>("", TaskDef("pick", NoopSwitch::class.java)))
+ pick().Case(Refs.node<Unit>("", TaskDef("long", NoopRefTask::class.java)))
+ pick().Case(Refs.node<Unit>("", TaskDef("short", OtherNoopRefTask::class.java)))
+ }
+
+ assertEquals(
+ listOf("long", "short"),
+ dag.tasks
+ .getValue("pick")
+ .decider!!
+ .cases
+ .map { it.id },
+ )
+ }
}
(airflow) 01/04: Java SDK: Pack each native Dag's source file into the bundle JAR (#74096)
Posted by ja...@apache.org.
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 2bda192870390af1b9a912589e08fea439f7e3c7
Author: Jason(Zhe-You) Liu <68...@users.noreply.github.com>
AuthorDate: Fri Oct 9 14:30:52 2026 +0800
Java SDK: Pack each native Dag's source file into the bundle JAR (#74096)
* Java SDK: Record the class that declares each Dag and describe it on request
DagDef remembers the outermost class that built it, and a generated
@Builder.Dag builder names the annotated class instead of itself.
Server accepts --describe-sources <file> to write each Java-declared Dag's
declaring class and return without connecting, for the Gradle plugin.
* Java SDK: Pack each native Dag's source file into the bundle JAR
The packDagSources task runs mainClass in describe mode, finds each Dag's
source file through the class file's SourceFile attribute, and packs each
unique file with a sources.json mapping under META-INF/airflow/. The JAR
manifest names it with Airflow-Java-SDK-Sources.
* Java SDK: Bound the describe run and keep its payload to the bundle JAR
Every `Jar` task gained a dependency on `packDagSources`, so a sources or
javadoc JAR carried the Dag sources and claimed the
`Airflow-Java-SDK-Sources` attribute that the Python side picks a bundle
JAR out of a directory by. Only `jar` and `shadowJar` carry them now.
The describe run had no time bound, so a `main` that does not route
through `Server.create(args)` and starts serving instead held the build
open with no way out. It now runs under a timeout and falls back to
packing the entrypoint's source alone, which is what every other
describe failure already does.
Say in `serve` and `serveAsync` that a `--describe-sources` server
returns without connecting, and record in `sources.json` that its paths
are relative to the source directory, which is also where they sit under
`META-INF/airflow/sources/`.
* Java SDK: Show each Dag's own source and find it outside the package layout
The Java importer now passes the Dag ID to the source lookup, so the Code
view reads dag_source_paths instead of always showing the entrypoint.
The source locator falls back to a unique match on the SourceFile name
when the package does not match the directory, using an index built only
when the package path misses. An unresolvable entrypoint or Dag class is
now logged as a warning.
The describe task is no longer up to date after a failed run, so a
transient failure does not stick until a source changes. Its hint now
points at Server.create(args), which is what answers --describe-sources.
Add a shadowJar test, correct the comment on which JARs carry the sources
payload, and document that a Dag maps to the class that constructed it.
* Java SDK: Use a describe run's output even when the run itself failed
* Java SDK: Serve and describe sources from their own Server classes
* Java SDK: Pack the sources into the one JAR the bundle is built from
* Java SDK: Attach the plugin's documentation to the plugin class
* Java SDK: Drop the duplicated classes dir from the describe classpath
* Java SDK: Let a bundle set how long the describe run may take
---
.../0015-per-dag-source-in-bundle-artifact.md | 2 +
.../language-sdks/java.rst | 4 +
java-sdk/README.md | 4 +
java-sdk/plugin/build.gradle.kts | 6 +
.../apache/airflow/sdk/plugin/AirflowSdkPlugin.kt | 55 +++-
.../org/apache/airflow/sdk/plugin/ClassFiles.kt | 85 +++++
.../apache/airflow/sdk/plugin/PackDagSources.kt | 256 +++++++++++++++
.../airflow/sdk/plugin/AirflowSdkPluginTest.kt | 365 +++++++++++++++++++++
.../apache/airflow/sdk/plugin/ClassFilesTest.kt | 80 +++++
.../airflow/sdk/plugin/PackDagSourcesTest.kt | 89 +++++
.../org/apache/airflow/sdk/plugin/TestSupport.kt | 45 +++
.../org/apache/airflow/sdk/BuilderProcessor.kt | 10 +-
.../kotlin/org/apache/airflow/sdk/BuilderTest.kt | 81 ++++-
.../main/kotlin/org/apache/airflow/sdk/DagDef.kt | 4 +
.../main/kotlin/org/apache/airflow/sdk/Server.kt | 109 ++++--
.../org/apache/airflow/sdk/internal/DagSource.kt | 66 ++++
.../apache/airflow/example/DagSourceFixtures.kt | 51 +++
.../kotlin/org/apache/airflow/sdk/DagSourceTest.kt | 132 ++++++++
.../kotlin/org/apache/airflow/sdk/ServerTest.kt | 6 +-
.../airflow/sdk/coordinators/java/_dag_importer.py | 4 +-
.../coordinators/java/test_dag_importer.py | 19 ++
21 files changed, 1418 insertions(+), 55 deletions(-)
diff --git a/airflow-core/adr/lang-sdk/0015-per-dag-source-in-bundle-artifact.md b/airflow-core/adr/lang-sdk/0015-per-dag-source-in-bundle-artifact.md
index 602066d958f..902a4d43b9c 100644
--- a/airflow-core/adr/lang-sdk/0015-per-dag-source-in-bundle-artifact.md
+++ b/airflow-core/adr/lang-sdk/0015-per-dag-source-in-bundle-artifact.md
@@ -94,6 +94,8 @@ TypeScript is the reference ([#73723](https://github.com/apache/airflow/pull/737
packers adopt the same shape — per-Dag source resolution, de-duplication by path, and an
always-embedded entrypoint fallback — so a bundle reader treats every language the same way and the
Code tab behaves identically regardless of which SDK produced the artifact.
+A JVM SDK has no module body, so it resolves a Dag by the class that declared it, and a Dag built by a
+factory maps to the factory's class.
Go has no module evaluation, so the Go SDK records the file that calls `airflow.Dag`. A Dag built in
a factory function maps to the factory's file, not to the file that calls the factory.
diff --git a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst
index 12b65d85580..793eda35412 100644
--- a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst
+++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst
@@ -1144,6 +1144,10 @@ into the bundle JAR, so the Airflow UI can show the source of a native Java Dag
:ref:`java-sdk/native-dag-parsing`). To find those classes, the plugin runs ``mainClass`` once at build
time. If that run fails, the build logs a warning and packs only the ``mainClass`` source.
+A Dag is recorded against the class that constructed it. A Dag built by a factory therefore maps to the
+factory's class, and one built by a factory in a dependency JAR has no source in the project and falls
+back to the ``mainClass`` source.
+
.. note::
You only need the ``annotationProcessor`` entry if you use the annotation-based API. It is not needed for
diff --git a/java-sdk/README.md b/java-sdk/README.md
index 6b4ba3522fd..8a3284403b8 100644
--- a/java-sdk/README.md
+++ b/java-sdk/README.md
@@ -91,6 +91,10 @@ Now `cd example` into the example project, and
../gradlew bundle
```
+ The bundle JAR also carries the source file of each Dag declared in Java
+ (and of the entrypoint) under `META-INF/airflow/`, so Airflow can show a
+ Dag's source.
+
* Put the [DAG with stub tasks](./example/src/resources/dags) to somewhere Airflow can find.
* Ensure the `java` command is available in the same environments the Airflow
diff --git a/java-sdk/plugin/build.gradle.kts b/java-sdk/plugin/build.gradle.kts
index edf930bb6fc..7efcdc36e61 100644
--- a/java-sdk/plugin/build.gradle.kts
+++ b/java-sdk/plugin/build.gradle.kts
@@ -25,6 +25,12 @@ plugins {
dependencies {
implementation("com.gradleup.shadow:shadow-gradle-plugin:9.1.0") // Last supporting Java 11.
+
+ testImplementation(kotlin("test"))
+}
+
+tasks.withType<Test> {
+ useJUnitPlatform()
}
gradlePlugin {
diff --git a/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt
index 3756ef688e1..5e3a41b5ce0 100644
--- a/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt
+++ b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt
@@ -22,13 +22,16 @@ package org.apache.airflow.sdk.plugin
import com.github.jengelman.gradle.plugins.shadow.tasks.ShadowJar
import org.gradle.api.Plugin
import org.gradle.api.Project
+import org.gradle.api.plugins.JavaPluginExtension
import org.gradle.api.provider.Property
import org.gradle.api.tasks.Copy
import org.gradle.api.tasks.Input
import org.gradle.api.tasks.SourceSetContainer
import org.gradle.api.tasks.bundling.Jar
+import org.gradle.jvm.toolchain.JavaToolchainService
import java.lang.reflect.Modifier
import java.net.URLClassLoader
+import java.time.Duration
import java.util.jar.JarFile
import kotlin.jvm.java
@@ -39,6 +42,7 @@ import kotlin.jvm.java
* airflowBundle {
* mainClass = "com.example.ExampleBundleBuilder"
* // fatJar = false // opt out of shadow JAR creation
+ * // describeTimeout = Duration.ofSeconds(60)
* }
* ```
*/
@@ -61,6 +65,10 @@ abstract class AirflowBundleExtension {
*/
@get:Input
abstract val fatJar: Property<Boolean>
+
+ /** How long the describe run may take before the build gives up and packs only the entrypoint's source. */
+ @get:Input
+ abstract val describeTimeout: Property<Duration>
}
/**
@@ -93,6 +101,13 @@ abstract class AirflowBundleExtension {
* identify which version of the Supervisor Schema it should use to communicate
* with the built JAR.
*
+ * When `mainClass` is set, the bundle JAR also carries the source file of each
+ * Dag declared in Java and of the entrypoint, so Airflow can show a Dag's source.
+ * They are collected by the `packDagSources` task, which runs `mainClass` once in
+ * a describe mode, and are listed in `META-INF/airflow/sources.json`, named by the
+ * `Airflow-Java-SDK-Sources` manifest attribute. If that run fails, a warning is
+ * logged and only the entrypoint's source is packed.
+ *
* If `fatJar` is explicitly set to `false`, the `bundle` task builds a bare JAR
* containing only the Dag bundle, and collect all dependency JARs into the target
* directory instead. In this mode, `Airflow-Supervisor-Schema-Version` lives in
@@ -105,8 +120,31 @@ class AirflowSdkPlugin : Plugin<Project> {
val ext = project.extensions.create("airflowBundle", AirflowBundleExtension::class.java)
ext.fatJar.convention(true)
+ ext.describeTimeout.convention(Duration.ofSeconds(30))
project.afterEvaluate {
+ val main = project.extensions.getByType(SourceSetContainer::class.java).getByName("main")
+
+ val packTask =
+ project.tasks.register("packDagSources", PackDagSources::class.java) { task ->
+ task.group = "build"
+ task.description = "Collects the source file of each Dag and the entrypoint to pack into the bundle JAR."
+ task.dependsOn(project.tasks.named("classes"))
+ task.onlyIf { ext.mainClass.isPresent }
+ task.mainClass.set(ext.mainClass)
+ task.describeTimeout.set(ext.describeTimeout)
+ task.classesDirs.from(main.output.classesDirs)
+ task.runtimeClasspath.from(main.runtimeClasspath)
+ task.sourceDirs.from(main.allSource.srcDirs)
+ task.launcher.convention(
+ project.extensions
+ .getByType(JavaToolchainService::class.java)
+ .launcherFor(project.extensions.getByType(JavaPluginExtension::class.java).toolchain),
+ )
+ task.describeFile.set(project.layout.buildDirectory.file("airflow/describe-sources.json"))
+ task.sourcesDir.set(project.layout.buildDirectory.dir("airflow/sources"))
+ }
+
project.tasks.withType(Jar::class.java).configureEach { task ->
task.doFirst {
ext.mainClass.orNull?.let { className ->
@@ -117,11 +155,7 @@ class AirflowSdkPlugin : Plugin<Project> {
val classFiles =
project.objects.fileCollection().from(
- project.extensions
- .getByType(SourceSetContainer::class.java)
- .getByName("main")
- .output
- .classesDirs,
+ main.output.classesDirs,
project.configurations.getByName("runtimeClasspath"),
)
@@ -162,8 +196,17 @@ class AirflowSdkPlugin : Plugin<Project> {
}
}
+ fun packSourcesInto(jarTask: String) {
+ if (!ext.mainClass.isPresent) return
+ project.tasks.named(jarTask, Jar::class.java).configure { task ->
+ task.from(packTask.flatMap { it.sourcesDir })
+ task.manifest.attributes(mapOf(SOURCES_MANIFEST_ATTRIBUTE to SOURCES_JSON_PATH))
+ }
+ }
+
if (ext.fatJar.get()) {
project.plugins.apply("com.gradleup.shadow")
+ packSourcesInto("shadowJar")
val schemaVersionProvider =
project.providers.provider {
@@ -205,6 +248,8 @@ class AirflowSdkPlugin : Plugin<Project> {
task.into(project.layout.buildDirectory.dir("bundle"))
}
} else {
+ packSourcesInto("jar")
+
// bundle copies the thin JAR and all runtime dependency JARs into
// build/bundle/, mirroring what installDist puts in lib/.
project.tasks.register("bundle", Copy::class.java) { task ->
diff --git a/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/ClassFiles.kt b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/ClassFiles.kt
new file mode 100644
index 00000000000..ef693559501
--- /dev/null
+++ b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/ClassFiles.kt
@@ -0,0 +1,85 @@
+/*
+ * 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.plugin
+
+import java.io.ByteArrayInputStream
+import java.io.DataInputStream
+import java.io.IOException
+
+private const val CLASS_MAGIC = 0xCAFEBABE.toInt()
+
+/**
+ * Reads the `SourceFile` attribute of a class file, which javac, kotlinc and scalac all
+ * write, or returns `null` if the class has none or the bytes are not a class file.
+ */
+internal fun readSourceFile(classFile: ByteArray): String? =
+ try {
+ DataInputStream(ByteArrayInputStream(classFile)).use(::readSourceFile)
+ } catch (_: IOException) {
+ null
+ }
+
+private fun readSourceFile(input: DataInputStream): String? {
+ if (input.readInt() != CLASS_MAGIC) return null
+ input.skipBytes(4) // minor and major version
+
+ val utf8 = arrayOfNulls<String>(input.readUnsignedShort())
+ var index = 1
+ while (index < utf8.size) {
+ when (val tag = input.readUnsignedByte()) {
+ 1 -> utf8[index] = input.readUTF()
+ 3, 4, 9, 10, 11, 12, 17, 18 -> input.skipBytes(4)
+ 5, 6 -> {
+ input.skipBytes(8)
+ index++ // long and double take two slots
+ }
+ 7, 8, 16, 19, 20 -> input.skipBytes(2)
+ 15 -> input.skipBytes(3)
+ else -> throw IOException("Unknown constant pool tag $tag")
+ }
+ index++
+ }
+
+ input.skipBytes(6) // access flags, this class, super class
+ input.skipBytes(2 * input.readUnsignedShort()) // interfaces
+ repeat(2) { skipMembers(input) } // fields, then methods
+
+ repeat(input.readUnsignedShort()) {
+ val name = utf8.getOrNull(input.readUnsignedShort())
+ val length = input.readInt()
+ if (name == "SourceFile" && length == 2) return utf8.getOrNull(input.readUnsignedShort())
+ input.skipBytes(length)
+ }
+ return null
+}
+
+private fun skipMembers(input: DataInputStream) {
+ repeat(input.readUnsignedShort()) {
+ input.skipBytes(6) // access flags, name, descriptor
+ skipAttributes(input)
+ }
+}
+
+private fun skipAttributes(input: DataInputStream) {
+ repeat(input.readUnsignedShort()) {
+ input.skipBytes(2)
+ input.skipBytes(input.readInt())
+ }
+}
diff --git a/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/PackDagSources.kt b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/PackDagSources.kt
new file mode 100644
index 00000000000..63bafb05de9
--- /dev/null
+++ b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/PackDagSources.kt
@@ -0,0 +1,256 @@
+/*
+ * 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.plugin
+
+import groovy.json.JsonOutput
+import groovy.json.JsonSlurper
+import org.gradle.api.DefaultTask
+import org.gradle.api.file.ConfigurableFileCollection
+import org.gradle.api.file.DirectoryProperty
+import org.gradle.api.file.RegularFileProperty
+import org.gradle.api.provider.Property
+import org.gradle.api.tasks.CacheableTask
+import org.gradle.api.tasks.Classpath
+import org.gradle.api.tasks.Input
+import org.gradle.api.tasks.InputFiles
+import org.gradle.api.tasks.Nested
+import org.gradle.api.tasks.Optional
+import org.gradle.api.tasks.OutputDirectory
+import org.gradle.api.tasks.OutputFile
+import org.gradle.api.tasks.PathSensitive
+import org.gradle.api.tasks.PathSensitivity
+import org.gradle.api.tasks.TaskAction
+import org.gradle.jvm.toolchain.JavaLauncher
+import java.io.File
+import java.time.Duration
+import java.util.concurrent.TimeUnit
+
+internal const val SOURCES_MANIFEST_ATTRIBUTE = "Airflow-Java-SDK-Sources"
+internal const val SOURCES_JSON_PATH = "META-INF/airflow/sources.json"
+internal const val SOURCES_DIR_PATH = "META-INF/airflow/sources"
+
+private const val DRAIN_TIMEOUT_MILLIS = 2_000L
+
+/**
+ * Finds the source file of a class by its `SourceFile` attribute: the class
+ * file's package path plus that name, looked up in each source directory, or
+ * else the one source file of that name when the layout does not match the
+ * package.
+ */
+internal class SourceLocator(
+ private val classesDirs: Iterable<File>,
+ private val sourceDirs: Iterable<File>,
+) {
+ /** Every source file under the source dirs by file name, walked once and only if a lookup misses. */
+ private val pathsByName: Map<String, List<String>> by lazy {
+ sourceDirs
+ .flatMap { dir -> dir.walkTopDown().filter(File::isFile).map { it.relativeTo(dir).invariantSeparatorsPath } }
+ .groupBy { it.substringAfterLast('/') }
+ }
+
+ /** Path of the source file relative to its source directory, or `null` if it cannot be found. */
+ fun locate(className: String): String? {
+ val classFile =
+ classesDirs
+ .map { File(it, className.replace('.', '/') + ".class") }
+ .firstOrNull { it.isFile } ?: return null
+ val sourceFile =
+ readSourceFile(classFile.readBytes())?.takeIf { it.isNotEmpty() && '/' !in it && '\\' !in it } ?: return null
+ val packagePath = className.substringBeforeLast('.', "").replace('.', '/')
+ val byPackage = listOfNotNull(packagePath.ifEmpty { null }, sourceFile).joinToString("/")
+ if (sourceDirs.any { File(it, byPackage).isFile }) return byPackage
+ // Neither javac nor Kotlin requires the package to match the directory, so fall back to the
+ // file name when exactly one source tree holds it. An ambiguous name is left unresolved.
+ return pathsByName[sourceFile]?.distinct()?.singleOrNull()
+ }
+
+ fun file(relativePath: String): File = sourceDirs.map { File(it, relativePath) }.first { it.isFile }
+}
+
+/**
+ * The `META-INF/airflow/sources.json` body: `entrypoint_path` and each Dag ID
+ * in `dag_source_paths`.
+ *
+ * Every path is relative to the source directory the file was found in, which
+ * is also where it sits under `META-INF/airflow/sources/` in the JAR, so a
+ * reader resolves one by prefixing that directory.
+ */
+internal fun sourcesJson(
+ entrypointPath: String?,
+ dagSourcePaths: Map<String, String>,
+): String =
+ JsonOutput.prettyPrint(
+ JsonOutput.toJson(
+ linkedMapOf<String, Any>().apply {
+ entrypointPath?.let { put("entrypoint_path", it) }
+ put("dag_source_paths", dagSourcePaths)
+ },
+ ),
+ )
+
+/**
+ * Runs the bundle's `mainClass` in describe mode to learn which class declared
+ * each Dag, then collects those classes' source files, and the entry point's,
+ * for the bundle JAR to carry.
+ */
+@CacheableTask
+abstract class PackDagSources : DefaultTask() {
+ @get:Input
+ @get:Optional
+ abstract val mainClass: Property<String>
+
+ /**
+ * How long the describe run gets before the build gives up on it and packs
+ * the entrypoint's source alone. A `main` that does not reach
+ * `Server.create(args)` never answers `--describe-sources`, and one that
+ * starts serving instead would otherwise hold the build open.
+ */
+ @get:Input
+ abstract val describeTimeout: Property<Duration>
+
+ @get:Classpath
+ abstract val classesDirs: ConfigurableFileCollection
+
+ @get:Classpath
+ abstract val runtimeClasspath: ConfigurableFileCollection
+
+ @get:InputFiles
+ @get:PathSensitive(PathSensitivity.RELATIVE)
+ abstract val sourceDirs: ConfigurableFileCollection
+
+ @get:Nested
+ @get:Optional
+ abstract val launcher: Property<JavaLauncher>
+
+ @get:OutputFile
+ abstract val describeFile: RegularFileProperty
+
+ @get:OutputDirectory
+ abstract val sourcesDir: DirectoryProperty
+
+ private var describeFailed = false
+
+ init {
+ outputs.doNotCacheIf("the Dag describe run failed") { describeFailed }
+ // A failed run leaves describeFile missing, which Gradle would otherwise read as unchanged.
+ outputs.upToDateWhen { describeFile.get().asFile.isFile }
+ }
+
+ @TaskAction
+ fun pack() {
+ val describe = describeFile.get().asFile
+ val root = sourcesDir.get().asFile
+ describe.delete()
+ root.deleteRecursively()
+ describe.parentFile.mkdirs()
+
+ val main = mainClass.get()
+ val declaringClasses = describeDags(main, describe)
+ val locator = SourceLocator(classesDirs.files, sourceDirs.files)
+
+ val entrypoint = locator.locate(main)
+ if (entrypoint == null) {
+ logger.warn("No source file found for entrypoint class {}; the Code view will show no source for this JAR", main)
+ }
+ val dagPaths = linkedMapOf<String, String>()
+ declaringClasses.forEach { (dagId, className) ->
+ val path = locator.locate(className)
+ if (path == null) {
+ logger.warn("No source file found for class {} of Dag '{}'; it falls back to the entrypoint", className, dagId)
+ } else {
+ dagPaths[dagId] = path
+ }
+ }
+
+ (listOfNotNull(entrypoint) + dagPaths.values).distinct().forEach { path ->
+ val target = File(root, "$SOURCES_DIR_PATH/$path")
+ target.parentFile.mkdirs()
+ locator.file(path).copyTo(target)
+ }
+ File(root, SOURCES_JSON_PATH).apply {
+ parentFile.mkdirs()
+ writeText(sourcesJson(entrypoint, dagPaths) + "\n")
+ }
+ }
+
+ private fun describeDags(
+ main: String,
+ describe: File,
+ ): Map<String, String> {
+ var output = ""
+ val failure =
+ try {
+ val java =
+ launcher.orNull
+ ?.executablePath
+ ?.asFile
+ ?.absolutePath ?: "java"
+ val classpath = (classesDirs + runtimeClasspath).joinToString(File.pathSeparator) { it.absolutePath }
+ val process =
+ ProcessBuilder(java, "-cp", classpath, main, "--describe-sources", describe.absolutePath)
+ .redirectErrorStream(true)
+ .start()
+ // Drained on its own thread: a main that writes more than the pipe
+ // holds would otherwise block before the timeout could fire.
+ val drain = Thread { output = process.inputStream.bufferedReader().readText() }
+ drain.isDaemon = true
+ drain.start()
+ val timeout = describeTimeout.get()
+ val finished = process.waitFor(timeout.toMillis(), TimeUnit.MILLISECONDS)
+ if (!finished) process.destroyForcibly()
+ drain.join(DRAIN_TIMEOUT_MILLIS)
+ when {
+ !finished -> "it did not finish within ${timeout.seconds} seconds"
+ process.exitValue() != 0 -> "exit code ${process.exitValue()}"
+ else -> null
+ }
+ } catch (e: Exception) {
+ e.message ?: e.javaClass.simpleName
+ }
+ val log = output.trim()
+ val parsed =
+ describe
+ .takeIf { it.isFile }
+ ?.let {
+ @Suppress("UNCHECKED_CAST")
+ (runCatching { JsonSlurper().parse(it) }.getOrNull() as? Map<String, Any?>)
+ }
+ if (parsed != null) {
+ if (failure != null) {
+ logger.warn(
+ "'{}' wrote each Dag's source but {}; its sources are used anyway.{}",
+ main,
+ failure,
+ if (log.isEmpty()) "" else "\n$log",
+ )
+ }
+ return parsed.mapNotNull { (id, cls) -> (cls as? String)?.let { id to it } }.toMap(linkedMapOf())
+ }
+ describeFailed = true
+ val why = failure ?: "it wrote no valid --describe-sources file; does main pass its args to Server.create?"
+ logger.warn(
+ "Could not read each Dag's source from '{}' ({}); only its entrypoint source is packed.{}",
+ main,
+ why,
+ if (log.isEmpty()) "" else "\n$log",
+ )
+ return emptyMap()
+ }
+}
diff --git a/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPluginTest.kt b/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPluginTest.kt
new file mode 100644
index 00000000000..3297ae940db
--- /dev/null
+++ b/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPluginTest.kt
@@ -0,0 +1,365 @@
+/*
+ * 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.plugin
+
+import groovy.json.JsonSlurper
+import org.gradle.testkit.runner.BuildResult
+import org.gradle.testkit.runner.GradleRunner
+import org.gradle.testkit.runner.TaskOutcome
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Assertions.assertFalse
+import org.junit.jupiter.api.Assertions.assertNotEquals
+import org.junit.jupiter.api.Assertions.assertNull
+import org.junit.jupiter.api.Assertions.assertTrue
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.io.TempDir
+import java.io.File
+import java.util.jar.Attributes
+import java.util.jar.JarFile
+import java.util.jar.JarOutputStream
+import java.util.jar.Manifest
+
+class AirflowSdkPluginTest {
+ private fun main(describeBody: String) =
+ """
+ package com.example;
+
+ public class Main {
+ public static void main(String[] args) throws Exception {
+ if (args.length == 2 && args[0].equals("--describe-sources")) {
+ $describeBody
+ }
+ }
+ }
+ """.trimIndent()
+
+ private val describeDags =
+ """
+ java.nio.file.Files.write(java.nio.file.Paths.get(args[1]), (
+ "{\"orders\":\"com.example.Main\"," +
+ "\"reports\":\"com.example.dags.Reports\"," +
+ "\"reports_backfill\":\"com.example.dags.Reports\"," +
+ "\"ghost\":\"com.example.Ghost\"}").getBytes());
+ """.trimIndent()
+
+ private fun project(
+ dir: File,
+ mainBody: String = describeDags,
+ mainClass: String? = "com.example.Main",
+ ): File {
+ dir.write("settings.gradle", "rootProject.name = 'bundle-test'\n")
+ dir.write(
+ "build.gradle",
+ """
+ plugins { id 'org.apache.airflow.sdk' }
+ airflowBundle {
+ ${mainClass?.let { "mainClass = '$it'" } ?: ""}
+ fatJar = false
+ }
+ sourceSets { main { java.srcDir 'src/extra' } }
+ """.trimIndent(),
+ )
+ dir.write("src/main/java/com/example/Main.java", main(mainBody))
+ dir.write("src/extra/com/example/dags/Reports.java", "package com.example.dags; public class Reports {}")
+ return dir
+ }
+
+ private fun gradle(
+ dir: File,
+ vararg tasks: String,
+ ): BuildResult =
+ GradleRunner
+ .create()
+ .withPluginClasspath()
+ .withProjectDir(dir)
+ .withArguments(*tasks, "--stacktrace")
+ .build()
+
+ private fun entries(jar: File): List<String> =
+ JarFile(jar).use {
+ it
+ .entries()
+ .asSequence()
+ .map { e -> e.name }
+ .toList()
+ }
+
+ @Suppress("UNCHECKED_CAST")
+ private fun sourcesJson(jar: File): Map<String, Any> =
+ JarFile(jar).use {
+ JsonSlurper().parse(it.getInputStream(it.getJarEntry("META-INF/airflow/sources.json"))) as Map<String, Any>
+ }
+
+ @Test
+ fun packsEachDagSourceOnceAndTheEntrypoint(
+ @TempDir dir: File,
+ ) {
+ project(dir)
+ val result = gradle(dir, "jar")
+ assertEquals(TaskOutcome.SUCCESS, result.task(":packDagSources")!!.outcome)
+
+ val jar = File(dir, "build/libs/bundle-test.jar")
+ assertEquals(
+ listOf(
+ "META-INF/airflow/sources.json",
+ "META-INF/airflow/sources/com/example/Main.java",
+ "META-INF/airflow/sources/com/example/dags/Reports.java",
+ ),
+ entries(jar).filter { it.startsWith("META-INF/airflow/") && !it.endsWith("/") }.sorted(),
+ )
+ assertEquals(
+ mapOf(
+ "entrypoint_path" to "com/example/Main.java",
+ "dag_source_paths" to
+ mapOf(
+ "orders" to "com/example/Main.java",
+ "reports" to "com/example/dags/Reports.java",
+ "reports_backfill" to "com/example/dags/Reports.java",
+ ),
+ ),
+ sourcesJson(jar),
+ )
+ JarFile(jar).use {
+ assertEquals("META-INF/airflow/sources.json", it.manifest.mainAttributes.getValue("Airflow-Java-SDK-Sources"))
+ assertEquals("com.example.Main", it.manifest.mainAttributes.getValue("Main-Class"))
+ }
+ }
+
+ @Test
+ fun isUpToDateOnSecondRunAndRerunsWhenASourceChanges(
+ @TempDir dir: File,
+ ) {
+ project(dir)
+ gradle(dir, "jar")
+ assertEquals(TaskOutcome.UP_TO_DATE, gradle(dir, "jar").task(":packDagSources")!!.outcome)
+
+ dir.write("src/extra/com/example/dags/Reports.java", "package com.example.dags; public class Reports { int v; }")
+ assertEquals(TaskOutcome.SUCCESS, gradle(dir, "jar").task(":packDagSources")!!.outcome)
+ val packed = File(dir, "build/airflow/sources/META-INF/airflow/sources/com/example/dags/Reports.java")
+ assertTrue(packed.readText().contains("int v"))
+ }
+
+ @Test
+ fun worksWithConfigurationCache(
+ @TempDir dir: File,
+ ) {
+ project(dir)
+ gradle(dir, "jar", "--configuration-cache")
+ val again = gradle(dir, "jar", "--configuration-cache")
+ assertTrue(again.output.contains("Reusing configuration cache"))
+ }
+
+ @Test
+ fun fallsBackToEntrypointOnlyWhenTheDescribeRunFails(
+ @TempDir dir: File,
+ ) {
+ project(dir, mainBody = """throw new IllegalStateException("boom");""")
+ val result = gradle(dir, "jar")
+
+ assertEquals(TaskOutcome.SUCCESS, result.task(":jar")!!.outcome)
+ assertTrue(result.output.contains("only its entrypoint source is packed"), result.output)
+ assertTrue(result.output.contains("boom"), result.output)
+ assertEquals(
+ mapOf("entrypoint_path" to "com/example/Main.java", "dag_source_paths" to emptyMap<String, String>()),
+ sourcesJson(File(dir, "build/libs/bundle-test.jar")),
+ )
+ }
+
+ @Test
+ fun usesTheDescribeFileEvenWhenTheRunExitsNonZero(
+ @TempDir dir: File,
+ ) {
+ project(dir, mainBody = describeDags + "\nSystem.exit(3);")
+ val result = gradle(dir, "jar")
+
+ assertTrue(result.output.contains("its sources are used anyway"), result.output)
+ assertEquals(
+ mapOf(
+ "orders" to "com/example/Main.java",
+ "reports" to "com/example/dags/Reports.java",
+ "reports_backfill" to "com/example/dags/Reports.java",
+ ),
+ sourcesJson(File(dir, "build/libs/bundle-test.jar"))["dag_source_paths"],
+ )
+ }
+
+ @Test
+ fun givesUpOnADescribeRunThatOutlastsTheConfiguredTimeout(
+ @TempDir dir: File,
+ ) {
+ project(dir, mainBody = "Thread.sleep(60000);")
+ dir.write(
+ "build.gradle",
+ File(dir, "build.gradle").readText() + "\nairflowBundle { describeTimeout = java.time.Duration.ofSeconds(2) }\n",
+ )
+ val result = gradle(dir, "jar")
+
+ assertTrue(result.output.contains("it did not finish within 2 seconds"), result.output)
+ }
+
+ @Test
+ fun fallsBackToEntrypointOnlyWhenMainWritesNothing(
+ @TempDir dir: File,
+ ) {
+ project(dir, mainBody = "")
+ val result = gradle(dir, "jar")
+
+ assertTrue(result.output.contains("does main pass its args to Server.create?"), result.output)
+ assertEquals(
+ mapOf("entrypoint_path" to "com/example/Main.java", "dag_source_paths" to emptyMap<String, String>()),
+ sourcesJson(File(dir, "build/libs/bundle-test.jar")),
+ )
+ }
+
+ @Test
+ fun findsASourceWhosePackageDoesNotMatchItsDirectory(
+ @TempDir dir: File,
+ ) {
+ project(dir)
+ File(dir, "src/extra/com/example/dags/Reports.java").delete()
+ dir.write("src/main/java/dags/Reports.java", "package com.example.dags; public class Reports {}")
+ gradle(dir, "jar")
+
+ assertEquals(
+ mapOf(
+ "entrypoint_path" to "com/example/Main.java",
+ "dag_source_paths" to
+ mapOf(
+ "orders" to "com/example/Main.java",
+ "reports" to "dags/Reports.java",
+ "reports_backfill" to "dags/Reports.java",
+ ),
+ ),
+ sourcesJson(File(dir, "build/libs/bundle-test.jar")),
+ )
+ }
+
+ @Test
+ fun warnsWhenTheEntrypointSourceCannotBeFound(
+ @TempDir dir: File,
+ ) {
+ project(dir)
+ dir.write(
+ "build.gradle",
+ File(dir, "build.gradle").readText() +
+ "\ncompileJava.doLast { delete 'src/main/java/com/example/Main.java' }\n",
+ )
+ val result = gradle(dir, "jar")
+
+ assertTrue(result.output.contains("No source file found for entrypoint class com.example.Main"), result.output)
+ assertNull(sourcesJson(File(dir, "build/libs/bundle-test.jar"))["entrypoint_path"])
+ }
+
+ @Test
+ fun rerunsAfterAFailedDescribeRun(
+ @TempDir dir: File,
+ ) {
+ project(dir, mainBody = "")
+ gradle(dir, "jar")
+ val again = gradle(dir, "jar")
+
+ assertNotEquals(TaskOutcome.UP_TO_DATE, again.task(":packDagSources")!!.outcome)
+ assertTrue(again.output.contains("only its entrypoint source is packed"), again.output)
+ }
+
+ @Test
+ fun shadowJarCarriesTheSourcesPayloadByDefault(
+ @TempDir dir: File,
+ ) {
+ dir.write("settings.gradle", "rootProject.name = 'bundle-test'\n")
+ dir.write(
+ "build.gradle",
+ """
+ plugins { id 'org.apache.airflow.sdk' }
+ repositories { maven { url = uri('repo') } }
+ dependencies { implementation 'org.apache.airflow:airflow-sdk:1.0' }
+ airflowBundle { mainClass = 'com.example.Main' }
+ """.trimIndent(),
+ )
+ dir.write("src/main/java/com/example/Main.java", main(describeDags))
+ dir.write("src/main/java/com/example/dags/Reports.java", "package com.example.dags; public class Reports {}")
+ stubSdk(File(dir, "repo/org/apache/airflow/airflow-sdk/1.0"))
+
+ gradle(dir, "shadowJar")
+
+ val jar = File(dir, "build/libs/bundle-test-all.jar")
+ assertTrue(entries(jar).contains("META-INF/airflow/sources.json"))
+ JarFile(jar).use {
+ assertEquals("META-INF/airflow/sources.json", it.manifest.mainAttributes.getValue("Airflow-Java-SDK-Sources"))
+ assertEquals("1", it.manifest.mainAttributes.getValue("Airflow-Supervisor-Schema-Version"))
+ }
+ }
+
+ private fun stubSdk(dir: File) {
+ dir.mkdirs()
+ File(dir, "airflow-sdk-1.0.pom").writeText(
+ """
+ <project>
+ <modelVersion>4.0.0</modelVersion>
+ <groupId>org.apache.airflow</groupId>
+ <artifactId>airflow-sdk</artifactId>
+ <version>1.0</version>
+ </project>
+ """.trimIndent(),
+ )
+ val manifest =
+ Manifest().apply {
+ mainAttributes[Attributes.Name.MANIFEST_VERSION] = "1.0"
+ mainAttributes.putValue("Airflow-Supervisor-Schema-Version", "1")
+ }
+ JarOutputStream(File(dir, "airflow-sdk-1.0.jar").outputStream(), manifest).close()
+ }
+
+ @Test
+ fun leavesTheSourcesPayloadOutOfANonBundleJar(
+ @TempDir dir: File,
+ ) {
+ project(dir)
+ dir.write(
+ "build.gradle",
+ File(dir, "build.gradle").readText() + "\njava { withSourcesJar() }\n",
+ )
+
+ gradle(dir, "jar", "sourcesJar")
+
+ val sources = File(dir, "build/libs/bundle-test-sources.jar")
+ assertFalse(entries(sources).any { it.startsWith("META-INF/airflow") })
+ JarFile(sources).use { assertNull(it.manifest.mainAttributes.getValue("Airflow-Java-SDK-Sources")) }
+ JarFile(File(dir, "build/libs/bundle-test.jar")).use {
+ assertEquals(
+ "META-INF/airflow/sources.json",
+ it.manifest.mainAttributes.getValue("Airflow-Java-SDK-Sources"),
+ )
+ }
+ }
+
+ @Test
+ fun skipsEverythingWithoutMainClass(
+ @TempDir dir: File,
+ ) {
+ project(dir, mainClass = null)
+ val result = gradle(dir, "jar", "packDagSources")
+
+ assertEquals(TaskOutcome.SKIPPED, result.task(":packDagSources")!!.outcome)
+ val jar = File(dir, "build/libs/bundle-test.jar")
+ assertFalse(entries(jar).any { it.startsWith("META-INF/airflow") })
+ JarFile(jar).use { assertNull(it.manifest.mainAttributes.getValue("Airflow-Java-SDK-Sources")) }
+ }
+}
diff --git a/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/ClassFilesTest.kt b/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/ClassFilesTest.kt
new file mode 100644
index 00000000000..693a8a21c12
--- /dev/null
+++ b/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/ClassFilesTest.kt
@@ -0,0 +1,80 @@
+/*
+ * 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.plugin
+
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Assertions.assertNull
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.io.TempDir
+import java.io.File
+
+class ClassFilesTest {
+ private class Nested
+
+ private fun bytesOf(cls: Class<*>): ByteArray =
+ cls.getResourceAsStream("/" + cls.name.replace('.', '/') + ".class")!!.use { it.readBytes() }
+
+ @Test
+ fun readsSourceFileOfKotlinClass() {
+ assertEquals("ClassFilesTest.kt", readSourceFile(bytesOf(ClassFilesTest::class.java)))
+ }
+
+ @Test
+ fun readsSourceFileOfNestedClassFromItsOutermostFile() {
+ assertEquals("ClassFilesTest.kt", readSourceFile(bytesOf(Nested::class.java)))
+ }
+
+ @Test
+ fun readsSourceFileOfJdkClassesWithLongAndDoubleConstants() {
+ // These carry every constant pool kind, including long, double and invokedynamic entries.
+ assertEquals("String.java", readSourceFile(bytesOf(String::class.java)))
+ assertEquals("Long.java", readSourceFile(bytesOf(Long::class.javaObjectType)))
+ assertEquals("Double.java", readSourceFile(bytesOf(Double::class.javaObjectType)))
+ assertEquals("HashMap.java", readSourceFile(bytesOf(java.util.HashMap::class.java)))
+ }
+
+ @Test
+ fun readsSourceFileOfJavaClassWithSeveralTopLevelTypes(
+ @TempDir dir: File,
+ ) {
+ dir.write("p/Main.java", "package p; public class Main { static final long L = 1L; } class Other {}")
+ compileJava(dir, File(dir, "out"), "p/Main.java")
+
+ assertEquals("Main.java", readSourceFile(File(dir, "out/p/Main.class").readBytes()))
+ assertEquals("Main.java", readSourceFile(File(dir, "out/p/Other.class").readBytes()))
+ }
+
+ @Test
+ fun returnsNullWithoutSourceFileAttribute(
+ @TempDir dir: File,
+ ) {
+ dir.write("p/Main.java", "package p; public class Main {}")
+ compileJava(dir, File(dir, "out"), "p/Main.java", options = listOf("-g:none"))
+
+ assertNull(readSourceFile(File(dir, "out/p/Main.class").readBytes()))
+ }
+
+ @Test
+ fun returnsNullForBytesThatAreNotAClassFile() {
+ assertNull(readSourceFile(ByteArray(0)))
+ assertNull(readSourceFile("not a class file".toByteArray()))
+ assertNull(readSourceFile(bytesOf(ClassFilesTest::class.java).copyOf(40)))
+ }
+}
diff --git a/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/PackDagSourcesTest.kt b/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/PackDagSourcesTest.kt
new file mode 100644
index 00000000000..bca13df2396
--- /dev/null
+++ b/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/PackDagSourcesTest.kt
@@ -0,0 +1,89 @@
+/*
+ * 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.plugin
+
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Assertions.assertNull
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.io.TempDir
+import java.io.File
+
+class PackDagSourcesTest {
+ private fun project(dir: File): Pair<File, File> {
+ val src = File(dir, "src")
+ val extra = File(dir, "extra")
+ src.write("com/example/Main.java", "package com.example; public class Main {}")
+ src.write("com/example/Same.java", "package com.example; public class Same { class Inner {} }")
+ extra.write("com/example/dags/Reports.java", "package com.example.dags; public class Reports {} class Helper {}")
+ src.write("com/example/NoSource.java", "package com.example; public class NoSource {}")
+ val classes = File(dir, "classes")
+ compileJava(
+ src,
+ classes,
+ "com/example/Main.java",
+ "com/example/Same.java",
+ "com/example/NoSource.java",
+ )
+ compileJava(extra, classes, "com/example/dags/Reports.java")
+ File(src, "com/example/NoSource.java").delete()
+ return src to classes
+ }
+
+ @Test
+ fun locatesSourceByPackagePathAndSourceFileAttribute(
+ @TempDir dir: File,
+ ) {
+ val (src, classes) = project(dir)
+ val locator = SourceLocator(listOf(classes), listOf(src, File(dir, "extra")))
+
+ assertEquals("com/example/Main.java", locator.locate("com.example.Main"))
+ assertEquals("com/example/Same.java", locator.locate("com.example.Same\$Inner"))
+ assertEquals("com/example/dags/Reports.java", locator.locate("com.example.dags.Reports"))
+ assertEquals("com/example/dags/Reports.java", locator.locate("com.example.dags.Helper"))
+ assertEquals(File(dir, "extra/com/example/dags/Reports.java"), locator.file("com/example/dags/Reports.java"))
+ }
+
+ @Test
+ fun locatesNothingForUnknownClassOrMissingSource(
+ @TempDir dir: File,
+ ) {
+ val (src, classes) = project(dir)
+ val locator = SourceLocator(listOf(classes), listOf(src))
+
+ assertNull(locator.locate("com.example.Missing"))
+ assertNull(locator.locate("com.example.NoSource"))
+ assertNull(locator.locate("com.example.dags.Reports"))
+ }
+
+ @Test
+ fun rendersSourcesJson() {
+ assertEquals(
+ """{"entrypoint_path":"com/example/Main.java","dag_source_paths":{"orders":"com/example/Main.java"}}""",
+ compact(sourcesJson("com/example/Main.java", mapOf("orders" to "com/example/Main.java"))),
+ )
+ }
+
+ @Test
+ fun omitsEntrypointPathWhenUnresolved() {
+ assertEquals("""{"dag_source_paths":{}}""", compact(sourcesJson(null, emptyMap())))
+ }
+
+ private fun compact(json: String) = (groovy.json.JsonSlurper().parseText(json) as Map<*, *>).let(groovy.json.JsonOutput::toJson)
+}
diff --git a/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/TestSupport.kt b/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/TestSupport.kt
new file mode 100644
index 00000000000..d8acc86fcfc
--- /dev/null
+++ b/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/TestSupport.kt
@@ -0,0 +1,45 @@
+/*
+ * 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.plugin
+
+import java.io.File
+import javax.tools.ToolProvider
+
+/** Compiles [sources] (paths under [srcDir]) into [classesDir] with the JDK's own compiler. */
+internal fun compileJava(
+ srcDir: File,
+ classesDir: File,
+ vararg sources: String,
+ options: List<String> = emptyList(),
+) {
+ classesDir.mkdirs()
+ val compiler = checkNotNull(ToolProvider.getSystemJavaCompiler()) { "Tests need a JDK" }
+ val args = options + listOf("-d", classesDir.path) + sources.map { File(srcDir, it).path }
+ check(compiler.run(null, null, null, *args.toTypedArray()) == 0) { "javac failed for ${sources.toList()}" }
+}
+
+internal fun File.write(
+ relativePath: String,
+ text: String,
+): File =
+ File(this, relativePath).apply {
+ parentFile.mkdirs()
+ writeText(text)
+ }
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 cc0d9c594f8..145ee4714c7 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
@@ -30,6 +30,7 @@ import com.squareup.javapoet.TypeName
import com.squareup.javapoet.TypeSpec
import com.squareup.javapoet.WildcardTypeName
import org.apache.airflow.sdk.internal.ArgValues
+import org.apache.airflow.sdk.internal.DagSource
import org.apache.airflow.sdk.internal.Field
import org.apache.airflow.sdk.internal.FieldType
import org.apache.airflow.sdk.internal.GROUP_ID
@@ -249,7 +250,13 @@ class BuilderProcessor : AbstractProcessor() {
.methodBuilder("build")
.addModifiers(Modifier.PUBLIC, Modifier.STATIC)
.returns(DAG_DEF_TYPE)
- .addStatement($$"var dag = new $T($S)", DAG_DEF_TYPE, ann.id.ifBlank { el.simpleName })
+ .addStatement(
+ $$"var dag = $T.declaredBy(new $T($S), $T.class)",
+ DAG_SOURCE_TYPE,
+ DAG_DEF_TYPE,
+ ann.id.ifBlank { el.simpleName },
+ ClassName.get(el),
+ )
explicitConfig(el, DAG_ANNOTATION, DAG_STRUCTURAL_ATTRIBUTES, SchemaFields.DAG).forEach { (key, value) ->
buildMethod.addStatement($$"dag.config($S, $L)", key, value)
}
@@ -913,6 +920,7 @@ private val TASK_INPUT_TYPE = ClassName.get(TaskInput::class.java)
private val TASK_ARGS_TYPE = ClassName.get(TaskArgs::class.java)
private val TYPE_REF_TYPE = ClassName.get(TypeRef::class.java)
private val ARG_VALUES_TYPE = ClassName.get(ArgValues::class.java)
+private val DAG_SOURCE_TYPE = ClassName.get(DagSource::class.java)
private val REFS_TYPE = ClassName.get(Refs::class.java)
private val ARG_TYPE = ClassName.get(Arg::class.java)
private val TASK_HANDLE_TYPE = ClassName.get(TaskRef::class.java)
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 c3759fa1998..a30dbacf008 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
@@ -23,8 +23,11 @@ import com.google.testing.compile.CompilationSubject.assertThat
import com.google.testing.compile.Compiler
import com.google.testing.compile.JavaFileObjectSubject
import com.google.testing.compile.JavaFileObjects
+import org.junit.jupiter.api.Assertions
import org.junit.jupiter.api.DisplayName
import org.junit.jupiter.api.Test
+import java.nio.file.Files
+import javax.tools.JavaFileObject
private fun compile(source: String) =
Compiler.javac().withProcessors(BuilderProcessor()).compile(
@@ -93,12 +96,13 @@ class BuilderTest {
import org.apache.airflow.sdk.Context;
import org.apache.airflow.sdk.DagDef;
import org.apache.airflow.sdk.Task;
+ import org.apache.airflow.sdk.internal.DagSource;
import org.apache.airflow.sdk.internal.Refs;
import org.apache.airflow.sdk.internal.TaskArgs;
public final class TestExampleBuilder {
public static DagDef build() {
- var dag = new DagDef("TestExample");
+ var dag = DagSource.declaredBy(new DagDef("TestExample"), TestExample.class);
return Refs.record(dag, List.of("t1", "t2", "t3"), List.of(), new TestExample.Wiring()::depends);
}
@@ -208,12 +212,13 @@ class BuilderTest {
import org.apache.airflow.sdk.Context;
import org.apache.airflow.sdk.DagDef;
import org.apache.airflow.sdk.Task;
+ import org.apache.airflow.sdk.internal.DagSource;
import org.apache.airflow.sdk.internal.Refs;
import org.apache.airflow.sdk.internal.TaskArgs;
public final class TestExampleBuilder {
public static DagDef build() {
- var dag = new DagDef("TestExample");
+ var dag = DagSource.declaredBy(new DagDef("TestExample"), TestExample.class);
return Refs.record(dag, List.of("t"), List.of(), new TestExample.Wiring()::depends);
}
@@ -275,13 +280,14 @@ class BuilderTest {
import org.apache.airflow.sdk.Context;
import org.apache.airflow.sdk.DagDef;
import org.apache.airflow.sdk.Task;
+ import org.apache.airflow.sdk.internal.DagSource;
import org.apache.airflow.sdk.internal.Refs;
import org.apache.airflow.sdk.internal.TaskArgs;
import org.apache.airflow.sdk.internal.TypeRef;
public final class TestExampleBuilder {
public static DagDef build() {
- var dag = new DagDef("TestExample");
+ var dag = DagSource.declaredBy(new DagDef("TestExample"), TestExample.class);
return Refs.record(dag, List.of("t"), List.of(), new TestExample.Wiring()::depends);
}
@@ -429,11 +435,12 @@ class BuilderTest {
import org.apache.airflow.sdk.Context;
import org.apache.airflow.sdk.DagDef;
import org.apache.airflow.sdk.Task;
+ import org.apache.airflow.sdk.internal.DagSource;
import org.apache.airflow.sdk.internal.Refs;
public final class TestExampleBuilder {
public static DagDef build() {
- var dag = new DagDef("cfg");
+ var dag = DagSource.declaredBy(new DagDef("cfg"), TestExample.class);
dag.config("schedule", "@daily");
dag.config("tags", List.of("a", "b"));
dag.config("catchup", true);
@@ -518,12 +525,13 @@ class BuilderTest {
import org.apache.airflow.sdk.Context;
import org.apache.airflow.sdk.DagDef;
import org.apache.airflow.sdk.Task;
+ import org.apache.airflow.sdk.internal.DagSource;
import org.apache.airflow.sdk.internal.Refs;
import org.apache.airflow.sdk.internal.TaskArgs;
public final class TestExampleBuilder {
public static DagDef build() {
- var dag = new DagDef("TestExample");
+ var dag = DagSource.declaredBy(new DagDef("TestExample"), TestExample.class);
return Refs.record(dag, List.of("t"), List.of(), new TestExample.Wiring()::depends);
}
public static final class T implements Task {
@@ -588,12 +596,13 @@ class BuilderTest {
import org.apache.airflow.sdk.DagDef;
import org.apache.airflow.sdk.Task;
import org.apache.airflow.sdk.internal.ArgValues;
+ import org.apache.airflow.sdk.internal.DagSource;
import org.apache.airflow.sdk.internal.Refs;
import org.apache.airflow.sdk.internal.TaskArgs;
public final class TestExampleBuilder {
public static DagDef build() {
- var dag = new DagDef("TestExample");
+ var dag = DagSource.declaredBy(new DagDef("TestExample"), TestExample.class);
return Refs.record(dag, List.of("flat", "named"), List.of(), new TestExample.Wiring()::depends);
}
@@ -811,10 +820,11 @@ class BuilderTest {
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 = new DagDef("foo");
+ var dag = DagSource.declaredBy(new DagDef("foo"), TestExample.class);
return Refs.record(dag, List.of(), List.of(), new TestExample.Wiring()::depends);
}
}
@@ -847,10 +857,11 @@ class BuilderTest {
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 Foo {
public static DagDef build() {
- var dag = new DagDef("TestExample");
+ var dag = DagSource.declaredBy(new DagDef("TestExample"), TestExample.class);
return Refs.record(dag, List.of(), List.of(), new TestExample.Wiring()::depends);
}
}
@@ -894,11 +905,12 @@ class BuilderTest {
import org.apache.airflow.sdk.Context;
import org.apache.airflow.sdk.DagDef;
import org.apache.airflow.sdk.Task;
+ import org.apache.airflow.sdk.internal.DagSource;
import org.apache.airflow.sdk.internal.Refs;
public final class TestExampleBuilder {
public static DagDef build() {
- var dag = new DagDef("TestExample");
+ var dag = DagSource.declaredBy(new DagDef("TestExample"), TestExample.class);
return Refs.record(dag, List.of("foo"), List.of(), new TestExample.Wiring()::depends);
}
@@ -1965,11 +1977,12 @@ class BuilderTest {
import org.apache.airflow.sdk.DagDef;
import org.apache.airflow.sdk.Task;
import org.apache.airflow.sdk.internal.ArgValues;
+ import org.apache.airflow.sdk.internal.DagSource;
import org.apache.airflow.sdk.internal.Refs;
public final class TestExampleBuilder {
public static DagDef build() {
- var dag = new DagDef("TestExample");
+ var dag = DagSource.declaredBy(new DagDef("TestExample"), TestExample.class);
return Refs.record(dag, List.of("score"), List.of(), new TestExample.Wiring()::depends);
}
@@ -2147,4 +2160,52 @@ class BuilderTest {
"Nested class 'Inner' holding @Builder.TaskHandler methods must be static",
)
}
+
+ @Test
+ @DisplayName("map the Dag to the annotated class, not its generated builder")
+ fun dagDeclaredByAnnotatedClass() {
+ val compilation =
+ compile(
+ """
+ package org.apache.airflow.example;
+ import org.apache.airflow.sdk.Builder;
+ @Builder.Dag(id = "orders")
+ public class TestExample {
+ @Builder.Task
+ public void t1() {}
+
+ @Builder.Deps
+ static class Wiring implements TestExampleDeps {
+ void depends() {
+ t1();
+ }
+ }
+ }
+ """,
+ )
+ assertThat(compilation).succeeded()
+
+ 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)
+ }
+ }
+ val bundle = Bundle().register(loader.loadClass("org.apache.airflow.example.TestExample"))
+ val target = Files.createTempFile("sources", ".json").toFile()
+ Server.create(arrayOf("--describe-sources", target.path)).serve(bundle)
+
+ Assertions.assertEquals("""{"orders":"org.apache.airflow.example.TestExample"}""", target.readText())
+ }
}
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 ff6f740af18..a269dff8c2c 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
@@ -19,6 +19,7 @@
package org.apache.airflow.sdk
+import org.apache.airflow.sdk.internal.DagSource
import org.apache.airflow.sdk.internal.GROUP_ID
import org.apache.airflow.sdk.internal.SchemaFields
import org.apache.airflow.sdk.internal.checkConfigValue
@@ -57,6 +58,9 @@ class DagDef(
/** Edges with a task group at either end, in the order drawn. */
internal val groupEdges = linkedSetOf<Pair<Endpoint, Endpoint>>()
+ /** Outermost class that declared this Dag, or `null` if it could not be told. */
+ internal var declaringClass: Class<*>? = DagSource.capture()
+
/**
* Sets one Dag-level configuration value.
*
diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt
index 47b44662b99..5847f5b7326 100644
--- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt
+++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt
@@ -19,7 +19,10 @@
package org.apache.airflow.sdk
+import com.fasterxml.jackson.databind.ObjectMapper
import com.xenomachina.argparser.ArgParser
+import com.xenomachina.argparser.MissingValueException
+import com.xenomachina.argparser.default
import io.ktor.network.selector.SelectorManager
import io.ktor.network.sockets.InetSocketAddress
import io.ktor.network.sockets.aSocket
@@ -38,6 +41,7 @@ import org.apache.airflow.sdk.execution.comm.ErrorResponse
import org.apache.airflow.sdk.execution.comm.StartupDetails
import org.apache.airflow.sdk.execution.parseDags
import org.apache.airflow.sdk.execution.runTask
+import java.io.File
import kotlin.text.substringAfterLast
import kotlin.text.substringBeforeLast
@@ -50,12 +54,16 @@ private class Args(
address.substringAfterLast(':').toInt(),
)
- val comm by parser.storing("--comm", help = "Address (host:port) to communicate with parent") {
- parseAddress(this)
- }
- val logs by parser.storing("--logs", help = "Address (host:port) to send Airflow logs to") {
- parseAddress(this)
- }
+ val comm by parser
+ .storing("--comm", help = "Address (host:port) to communicate with parent") { parseAddress(this) }
+ .default(null)
+ val logs by parser
+ .storing("--logs", help = "Address (host:port) to send Airflow logs to") { parseAddress(this) }
+ .default(null)
+ val describeSources by parser
+ .storing("--describe-sources", help = "File to write each Dag's declaring class to, instead of serving") {
+ File(this)
+ }.default(null)
}
/**
@@ -83,33 +91,7 @@ class ApiError(
* The process exits when the coordinator closes the connection (normally after
* one task-instance execution).
*/
-class Server(
- private val comm: InetSocketAddress,
- private val logs: InetSocketAddress,
-) {
- companion object {
- /**
- * Parses coordinator addresses from command-line arguments and returns a
- * ready-to-use [Server].
- *
- * The arguments are supplied automatically by Airflow and are not intended
- * to be constructed by hand:
- *
- * * `--comm host:port` address for task-execution messages.
- * * `--logs host:port` address for log forwarding.
- *
- * @param args Command-line arguments as received by `main`.
- * @return A configured [Server] ready to call [serve].
- */
- @JvmStatic
- fun create(args: Array<String>): Server {
- val args = ArgParser(args).parseInto(::Args)
- return Server(args.comm, args.logs)
- }
- }
-
- private val logger = Logger(Server::class)
-
+sealed class Server {
/**
* Blocking entry point: connects to the coordinator and serves task-execution
* requests from the given [bundle].
@@ -119,6 +101,9 @@ class Server(
* The call returns when the coordinator closes the connection (normally after
* one task-instance execution).
*
+ * A [Server] created from `--describe-sources` writes that file and returns
+ * without connecting anywhere.
+ *
* @param bundle Bundle containing all Dags this process can execute.
*
* @see [serveAsync]
@@ -136,6 +121,9 @@ class Server(
* coordinator closes the connection (normally after one task-instance
* execution). The coroutine returns once both channels have been closed.
*
+ * A [Server] created from `--describe-sources` writes that file and returns
+ * without connecting anywhere.
+ *
* Use this variant when calling from an existing coroutine scope; use the
* blocking [serve] from a plain `main` method.
*
@@ -143,7 +131,59 @@ class Server(
*
* @see [serve]
*/
- suspend fun serveAsync(bundle: Bundle) =
+ abstract suspend fun serveAsync(bundle: Bundle)
+
+ companion object {
+ /**
+ * Parses coordinator addresses from command-line arguments and returns a
+ * ready-to-use [Server].
+ *
+ * The arguments are supplied automatically by Airflow and are not intended
+ * to be constructed by hand:
+ *
+ * * `--comm host:port` address for task-execution messages.
+ * * `--logs host:port` address for log forwarding.
+ * * `--describe-sources file` makes [serve] write a JSON object mapping each
+ * Java-declared Dag ID to the binary name of its declaring class into `file`
+ * and return without connecting anywhere; `--comm` and `--logs` are not
+ * required then. It is for the Gradle plugin, not for users.
+ *
+ * @param args Command-line arguments as received by `main`.
+ * @return A configured [Server] ready to call [serve].
+ */
+ @JvmStatic
+ fun create(args: Array<String>): Server {
+ val args = ArgParser(args).parseInto(::Args)
+ args.describeSources?.let { return SourceDescriber(it) }
+ return CoordinatorServer(
+ args.comm ?: throw MissingValueException("--comm"),
+ args.logs ?: throw MissingValueException("--logs"),
+ )
+ }
+ }
+}
+
+/** Writes each Java-declared Dag's declaring class to [target] instead of serving. */
+internal class SourceDescriber(
+ private val target: File,
+) : Server() {
+ override suspend fun serveAsync(bundle: Bundle) {
+ bundle.finalizeRegistration()
+ val sources = linkedMapOf<String, String>()
+ bundle.dags.values.forEach { dag -> dag.declaringClass?.let { sources[dag.id] = it.name } }
+ target.absoluteFile.parentFile?.mkdirs()
+ ObjectMapper().writeValue(target, sources)
+ }
+}
+
+/** Serves task-execution requests over the coordinator's [comm] and [logs] sockets. */
+class CoordinatorServer(
+ private val comm: InetSocketAddress,
+ private val logs: InetSocketAddress,
+) : Server() {
+ private val logger = Logger(CoordinatorServer::class)
+
+ override suspend fun serveAsync(bundle: Bundle) {
coroutineScope {
bundle.finalizeRegistration()
val deferral = CompletableDeferred<Unit>()
@@ -175,6 +215,7 @@ class Server(
}
}
}
+ }
internal suspend fun dispatchTask(
bundle: Bundle,
diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/DagSource.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/DagSource.kt
new file mode 100644
index 00000000000..809dd818a3d
--- /dev/null
+++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/DagSource.kt
@@ -0,0 +1,66 @@
+/*
+ * 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.internal
+
+import org.apache.airflow.sdk.DagDef
+
+/**
+ * @suppress
+ *
+ * Tracks which class declared a [DagDef], so the bundle can ship that class's
+ * source file. Public so that processor-generated builders can call [declaredBy];
+ * not user-facing API.
+ */
+object DagSource {
+ private const val SDK_PACKAGE = "org.apache.airflow.sdk."
+ private val IGNORED_PREFIXES = listOf(SDK_PACKAGE, "java.", "javax.", "jdk.", "sun.", "kotlin.", "kotlinx.")
+
+ private val walker = StackWalker.getInstance(StackWalker.Option.RETAIN_CLASS_REFERENCE)
+
+ /**
+ * Names [declaring] as the class that declared [dag], replacing what was
+ * captured at construction. A generated builder calls this so the Dag points
+ * at the annotated class rather than the builder.
+ */
+ @JvmStatic
+ fun declaredBy(
+ dag: DagDef,
+ declaring: Class<*>,
+ ): DagDef {
+ dag.declaringClass = outermost(declaring)
+ return dag
+ }
+
+ /** The outermost class of the first caller outside the SDK and the standard libraries. */
+ internal fun capture(): Class<*>? =
+ walker.walk { frames ->
+ frames
+ .map { it.declaringClass }
+ .filter { c -> IGNORED_PREFIXES.none { c.name.startsWith(it) } }
+ .findFirst()
+ .map { outermost(it) }
+ .orElse(null)
+ }
+
+ internal fun outermost(cls: Class<*>): Class<*> {
+ var current = cls
+ while (true) current = current.enclosingClass ?: return current
+ }
+}
diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/example/DagSourceFixtures.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/example/DagSourceFixtures.kt
new file mode 100644
index 00000000000..bdcf697060b
--- /dev/null
+++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/example/DagSourceFixtures.kt
@@ -0,0 +1,51 @@
+/*
+ * 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.example
+
+import org.apache.airflow.sdk.DagDef
+
+// Lives outside org.apache.airflow.sdk, as user code does, so the SDK does not skip its frames.
+class DagSourceFixtures {
+ fun plain() = DagDef("plain")
+
+ fun fromLambda(): DagDef {
+ lateinit var dag: DagDef
+ Runnable { dag = DagDef("lambda") }.run()
+ return dag
+ }
+
+ fun fromAnonymous(): DagDef {
+ lateinit var dag: DagDef
+ object : Runnable {
+ override fun run() {
+ dag = DagDef("anonymous")
+ }
+ }.run()
+ return dag
+ }
+
+ class Nested {
+ fun make() = DagDef("nested")
+
+ class Deeper {
+ fun make() = DagDef("deeper")
+ }
+ }
+}
diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/DagSourceTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/DagSourceTest.kt
new file mode 100644
index 00000000000..39a3b740aca
--- /dev/null
+++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/DagSourceTest.kt
@@ -0,0 +1,132 @@
+/*
+ * 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 com.fasterxml.jackson.databind.ObjectMapper
+import com.xenomachina.argparser.SystemExitException
+import org.apache.airflow.example.DagSourceFixtures
+import org.apache.airflow.sdk.internal.DagSource
+import org.junit.jupiter.api.Assertions
+import org.junit.jupiter.api.DisplayName
+import org.junit.jupiter.api.Test
+import org.junit.jupiter.api.io.TempDir
+import java.io.File
+
+internal class DagSourceTest {
+ private class NoOp : Task {
+ override fun execute(
+ context: Context,
+ client: Client,
+ ) = Unit
+ }
+
+ @Test
+ @DisplayName("Should record the class that constructed the Dag")
+ fun recordsConstructingClass() {
+ Assertions.assertEquals(DagSourceFixtures::class.java, DagSourceFixtures().plain().declaringClass)
+ }
+
+ @Test
+ @DisplayName("Should record the outermost class for nested, lambda and anonymous callers")
+ fun recordsOutermostClass() {
+ val fixtures = DagSourceFixtures()
+ Assertions.assertEquals(DagSourceFixtures::class.java, DagSourceFixtures.Nested().make().declaringClass)
+ Assertions.assertEquals(
+ DagSourceFixtures::class.java,
+ DagSourceFixtures.Nested
+ .Deeper()
+ .make()
+ .declaringClass,
+ )
+ Assertions.assertEquals(DagSourceFixtures::class.java, fixtures.fromLambda().declaringClass)
+ Assertions.assertEquals(DagSourceFixtures::class.java, fixtures.fromAnonymous().declaringClass)
+ }
+
+ @Test
+ @DisplayName("Should let a generated builder name the annotated class instead of itself")
+ fun declaredByOverridesCapturedClass() {
+ val dag = DagSource.declaredBy(DagDef("dag"), DagSourceFixtures.Nested::class.java)
+ Assertions.assertEquals(DagSourceFixtures::class.java, dag.declaringClass)
+ }
+
+ private fun describe(
+ dir: File,
+ bundle: Bundle,
+ ): Map<*, *> {
+ val target = File(dir, "out/sources.json")
+ Server.create(arrayOf("--describe-sources", target.path)).serve(bundle)
+ return ObjectMapper().readValue(target, Map::class.java)
+ }
+
+ @Test
+ @DisplayName("Should write each Java-declared Dag's declaring class and skip task-handler Dags")
+ fun describeSourcesWritesDeclaringClasses(
+ @TempDir dir: File,
+ ) {
+ val fixtures = DagSourceFixtures()
+ val bundle =
+ Bundle(listOf(fixtures.plain(), DagSourceFixtures.Nested().make()))
+ .register("python_owned", "t", NoOp::class.java)
+
+ Assertions.assertEquals(
+ mapOf("plain" to DagSourceFixtures::class.java.name, "nested" to DagSourceFixtures::class.java.name),
+ describe(dir, bundle),
+ )
+ }
+
+ @Test
+ @DisplayName("Should omit a Dag whose declaring class is unknown")
+ fun describeSourcesOmitsUnknownDeclaringClass(
+ @TempDir dir: File,
+ ) {
+ val unknown = DagDef("unknown").also { it.declaringClass = null }
+ val bundle = Bundle(listOf(DagSourceFixtures().plain(), unknown))
+
+ Assertions.assertEquals(mapOf("plain" to DagSourceFixtures::class.java.name), describe(dir, bundle))
+ }
+
+ @Test
+ @DisplayName("Should write an empty object when no Dag is declared in Java")
+ fun describeSourcesWithoutJavaDags(
+ @TempDir dir: File,
+ ) {
+ Assertions.assertEquals(emptyMap<Any, Any>(), describe(dir, Bundle()))
+ }
+
+ @Test
+ @DisplayName("Should use the binary name of the declaring class")
+ fun describeSourcesUsesBinaryName(
+ @TempDir dir: File,
+ ) {
+ val dag = DagDef("dag").also { it.declaringClass = DagSourceFixtures.Nested::class.java }
+ Assertions.assertEquals(
+ mapOf("dag" to "org.apache.airflow.example.DagSourceFixtures\$Nested"),
+ describe(dir, Bundle(listOf(dag))),
+ )
+ }
+
+ @Test
+ @DisplayName("Should still require --comm and --logs without --describe-sources")
+ fun createStillRequiresAddresses() {
+ Assertions.assertThrows(SystemExitException::class.java) { Server.create(arrayOf("--comm", "localhost:1")) }
+ Assertions.assertThrows(SystemExitException::class.java) { Server.create(arrayOf("--logs", "localhost:1")) }
+ Assertions.assertThrows(SystemExitException::class.java) { Server.create(emptyArray()) }
+ }
+}
diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ServerTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ServerTest.kt
index d6e93fe2114..7c7ab91b2ab 100644
--- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ServerTest.kt
+++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ServerTest.kt
@@ -76,7 +76,7 @@ class ServerTest {
val toServer = ByteChannel(autoFlush = true)
val fromServer = ByteChannel(autoFlush = true)
val comm = CoordinatorComm(toServer, fromServer)
- val server = Server(InetSocketAddress("localhost", 0), InetSocketAddress("localhost", 0))
+ val server = CoordinatorServer(InetSocketAddress("localhost", 0), InetSocketAddress("localhost", 0))
val reported = ArrayBlockingQueue<IncomingFrame>(1)
val supervisor =
@@ -110,7 +110,7 @@ class ServerTest {
val toServer = ByteChannel(autoFlush = true)
val fromServer = ByteChannel(autoFlush = true)
val comm = CoordinatorComm(toServer, fromServer)
- val server = Server(InetSocketAddress("localhost", 0), InetSocketAddress("localhost", 0))
+ val server = CoordinatorServer(InetSocketAddress("localhost", 0), InetSocketAddress("localhost", 0))
Assertions.assertThrows(ApiError::class.java) {
runBlocking {
@@ -128,7 +128,7 @@ class ServerTest {
val toServer = ByteChannel(autoFlush = true)
val fromServer = ByteChannel(autoFlush = true)
val comm = CoordinatorComm(toServer, fromServer)
- val server = Server(InetSocketAddress("localhost", 0), InetSocketAddress("localhost", 0))
+ val server = CoordinatorServer(InetSocketAddress("localhost", 0), InetSocketAddress("localhost", 0))
val bundle = Bundle(listOf(DagDef("parsed_dag").addTask(TaskDef("t", ServerNoopTask::class.java))))
val reported = ArrayBlockingQueue<RawFrame>(1)
diff --git a/task-sdk/src/airflow/sdk/coordinators/java/_dag_importer.py b/task-sdk/src/airflow/sdk/coordinators/java/_dag_importer.py
index 6a20f3d6018..b1e9283d26e 100644
--- a/task-sdk/src/airflow/sdk/coordinators/java/_dag_importer.py
+++ b/task-sdk/src/airflow/sdk/coordinators/java/_dag_importer.py
@@ -74,9 +74,9 @@ class JavaDagImporter(CoordinatorDagImporter):
return not wanted or main_class == wanted
def get_source_code(self, definition: DagDefinition, dag_id: str | None = None) -> DagSourceCode:
- """Return the entrypoint source the JAR embeds, or a notice when it embeds none."""
+ """Return the source the JAR embeds for *dag_id*, else the entrypoint's, else a notice."""
with definition.as_file() as path, zipfile.ZipFile(path) as zf:
- info = _find_source_entry(zf)
+ info = _find_source_entry(zf, dag_id)
if info is None:
return DagSourceCode(source_code=_NO_SOURCE, language="java")
if info.file_size > _MAX_SOURCE_BYTES:
diff --git a/task-sdk/tests/task_sdk/coordinators/java/test_dag_importer.py b/task-sdk/tests/task_sdk/coordinators/java/test_dag_importer.py
index 45ecda314ee..e3952299fbf 100644
--- a/task-sdk/tests/task_sdk/coordinators/java/test_dag_importer.py
+++ b/task-sdk/tests/task_sdk/coordinators/java/test_dag_importer.py
@@ -227,6 +227,25 @@ class TestGetSourceCode:
with zipfile.ZipFile(jar) as zf:
assert _find_source_entry(zf, dag_id).filename == f"META-INF/airflow/sources/{expected}"
+ @pytest.mark.parametrize(
+ ("dag_id", "expected"),
+ [
+ pytest.param("reports", "class Reports", id="mapped"),
+ pytest.param("orders", "class Dags", id="unmapped"),
+ pytest.param(None, "class Dags", id="no-dag-id"),
+ ],
+ )
+ def test_get_source_code_reads_the_source_of_the_dag(self, tmp_path, dag_id, expected):
+ jar = self._jar(
+ tmp_path,
+ index={"entrypoint_path": ENTRYPOINT, "dag_source_paths": {"reports": REPORTS}},
+ files={ENTRYPOINT: "class Dags", REPORTS: "class Reports"},
+ )
+
+ result = _importer().get_source_code(_definition(jar), dag_id)
+
+ assert result.source_code == expected
+
@pytest.mark.parametrize(
("attributes", "index", "files"),
[
(airflow) 02/04: Java SDK: Pick one of two tasks with a condition (#74438)
Posted by ja...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
jason810496 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
commit d4013cc2b5b52842f75ca6cf6360306600d69aa4
Author: Jason(Zhe-You) Liu <68...@users.noreply.github.com>
AuthorDate: Fri Oct 9 14:30:52 2026 +0800
Java SDK: Pick one of two tasks with a condition (#74438)
* Java SDK: Pick one of two tasks with a condition
A native Java Dag could only run every task it declared: it had no way to
let a task's result decide what runs next, so a Dag that loads only when
there are rows had to be split in two or fall back to a Python stub.
`DagDef.If` and `@Builder.If` declare a task whose boolean names the side
that runs; the other is skipped. The deciding task carries
`_can_skip_downstream`, so Airflow skips a cleared side again rather than
running it, and the SDK writes the `skipmixin_key` XCom before the skip, as
Python's SkipMixin does. It skips only the side not taken, so a task after
both sides needs a trigger rule such as `none_failed_min_one_success`.
`If` is capitalized because `if` is a Java keyword and `@Builder.If` reads
the same way. With no task id given, one is derived from the condition
class's simple name, so `HasRows.class` becomes `hasRows`.
The runtime reaches `SkipDownstreamTasks` through the transport client
only, never through the task-facing `Client`, so a task body cannot skip
tasks the Dag did not declare it may.
* Java SDK: Let a wiring view name a condition's sides in more than one call
* Java SDK: Reject a decider class registered as an ordinary task
* Java SDK: Name a condition's sides Then and Else
* Java SDK: Pin the order a condition skips and records what it skipped
* Java SDK: Drop the unused switch in the condition test's decider
---
.pre-commit-config.yaml | 2 +-
java-sdk/README.md | 2 +-
java-sdk/capabilities.yaml | 4 +-
.../example/nativedag/AnnotationExample.java | 15 +-
.../example/nativedag/InterfaceExample.java | 21 +-
.../org/apache/airflow/sdk/BuilderProcessor.kt | 103 +++++++---
.../kotlin/org/apache/airflow/sdk/BuilderTest.kt | 198 ++++++++++++++++++
.../ci/prek/check_serialization_conformance.py | 2 +-
java-sdk/sdk/build.gradle.kts | 33 +++
java-sdk/sdk/module.md | 2 +-
.../main/kotlin/org/apache/airflow/sdk/Bundle.kt | 26 ++-
.../kotlin/org/apache/airflow/sdk/Condition.kt | 224 +++++++++++++++++++++
.../main/kotlin/org/apache/airflow/sdk/DagDef.kt | 52 +++++
.../main/kotlin/org/apache/airflow/sdk/Decider.kt | 90 +++++++++
.../kotlin/org/apache/airflow/sdk/TaskGroupRef.kt | 32 +++
.../org/apache/airflow/sdk/execution/Client.kt | 14 ++
.../org/apache/airflow/sdk/execution/Serde.kt | 2 +
.../org/apache/airflow/sdk/execution/Task.kt | 29 ++-
.../kotlin/org/apache/airflow/sdk/internal/Ids.kt | 10 +
.../org/apache/airflow/sdk/ArgTestSupport.kt | 2 +
.../kotlin/org/apache/airflow/sdk/ClientTest.kt | 2 +
.../kotlin/org/apache/airflow/sdk/ConditionTest.kt | 208 +++++++++++++++++++
.../airflow/sdk/conformance/SerializeJava.kt | 29 ++-
.../org/apache/airflow/sdk/execution/CommTest.kt | 27 +++
.../org/apache/airflow/sdk/execution/TaskTest.kt | 150 ++++++++++++++
.../apache/airflow/sdk/internal/ArgValuesTest.kt | 2 +
.../org/apache/airflow/sdk/internal/RefsTest.kt | 29 +++
27 files changed, 1274 insertions(+), 36 deletions(-)
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 8392915d926..37f87710dfc 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -461,7 +461,7 @@ repos:
^java-sdk/scripts/ci/prek/check_serialization_conformance\.py$|
^java-sdk/sdk/schema/dag-schema\.json$|
^java-sdk/sdk/build\.gradle\.kts$|
- ^java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/(Bundle|DagDef|Deps|Arg|TaskGroupRef)\.kt$|
+ ^java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/(Bundle|Condition|DagDef|Decider|Deps|Arg|TaskGroupRef)\.kt$|
^java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Serde\.kt$|
^java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/(Refs|Fields)\.kt$|
^java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/conformance/.*$|
diff --git a/java-sdk/README.md b/java-sdk/README.md
index 8a3284403b8..1e4e7478f39 100644
--- a/java-sdk/README.md
+++ b/java-sdk/README.md
@@ -643,7 +643,7 @@ prek hook regenerate it.
| capability: `task-args` | MUST † | ✓ | 3.4 | |
| capability: `dag-params` | MUST † | ✗ | – | |
| capability: `taskflow-dependencies` | MUST † | ✓ | 3.4 | |
-| capability: `branching` | SHOULD † | ✗ | – | |
+| capability: `branching` | SHOULD † | ✓ | 3.4 | DagDef.If / @Builder.If |
| capability: `dag-test` | SHOULD † | ✗ | – | |
| capability: `task-group` | MAY † | ✓ | 3.4 | |
| capability: `dynamic-task-mapping` | MAY † | ✗ | – | |
diff --git a/java-sdk/capabilities.yaml b/java-sdk/capabilities.yaml
index 75c6053fc89..33106f170fb 100644
--- a/java-sdk/capabilities.yaml
+++ b/java-sdk/capabilities.yaml
@@ -107,7 +107,9 @@ capabilities:
supported: true
since: "3.4"
branching:
- supported: false
+ supported: true
+ since: "3.4"
+ note: "DagDef.If / @Builder.If"
dag-test:
supported: false
task-group:
diff --git a/java-sdk/example/src/java/org/apache/airflow/example/nativedag/AnnotationExample.java b/java-sdk/example/src/java/org/apache/airflow/example/nativedag/AnnotationExample.java
index 1c6ce70cc72..d9538fd391f 100644
--- a/java-sdk/example/src/java/org/apache/airflow/example/nativedag/AnnotationExample.java
+++ b/java-sdk/example/src/java/org/apache/airflow/example/nativedag/AnnotationExample.java
@@ -54,6 +54,18 @@ public class AnnotationExample {
log.log(INFO, "Loaded {0}", transformed);
}
+ @Builder.Task(id = "load_empty")
+ public void loadEmpty() {
+ log.log(INFO, "Nothing to load");
+ }
+
+ // A condition: its boolean picks one of the two loads, and the other is
+ // skipped.
+ @Builder.If(id = "has_rows")
+ public boolean hasRows(long transformed) {
+ return transformed > 0;
+ }
+
// A task group: everything it declares is prefixed with its id, so this is
// the task "checks.audit".
@Builder.TaskGroup(id = "checks")
@@ -70,7 +82,8 @@ public class AnnotationExample {
static class Wiring implements AnnotationExampleDeps {
void depends() {
var extracted = extract();
- load(transform(extracted, lit(1.5)));
+ var transformed = transform(extracted, lit(1.5));
+ hasRows(transformed).Then(load(transformed)).Else(loadEmpty());
// Ordering-only edge: the checks group runs after extract, with no data
// flowing.
extracted.before(checks());
diff --git a/java-sdk/example/src/java/org/apache/airflow/example/nativedag/InterfaceExample.java b/java-sdk/example/src/java/org/apache/airflow/example/nativedag/InterfaceExample.java
index 6047aeedd8e..cce10676dfd 100644
--- a/java-sdk/example/src/java/org/apache/airflow/example/nativedag/InterfaceExample.java
+++ b/java-sdk/example/src/java/org/apache/airflow/example/nativedag/InterfaceExample.java
@@ -55,6 +55,22 @@ public class InterfaceExample {
}
}
+ public static class LoadEmpty implements Task {
+ @Override
+ public void execute(Context context, Client client) {
+ log.log(INFO, "Nothing to load");
+ }
+ }
+
+ // A condition: its boolean picks one of the two loads, and the other is
+ // skipped.
+ public static class HasRows implements ConditionTask {
+ @Override
+ public boolean decide(Context context, Client client) {
+ return ((Number) client.getXCom("transform")).longValue() > 0;
+ }
+ }
+
public static DagDef build() {
var dag =
new DagDef("java_native_interface_example")
@@ -69,8 +85,11 @@ public class InterfaceExample {
.config("doc_md", "Extracts a value and pushes it as an XCom.");
var transform = dag.task("transform", Transform.class);
var load = dag.task("load", Load.class);
+ var loadEmpty = dag.task("load_empty", LoadEmpty.class);
- transform.after(extract).before(load);
+ transform.after(extract);
+ // With no task id given, the condition takes one from its class: "hasRows".
+ dag.If(HasRows.class).after(transform).Then(load).Else(loadEmpty);
return dag;
}
}
diff --git a/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt b/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt
index 145ee4714c7..1848fc007f7 100644
--- a/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt
+++ b/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt
@@ -96,6 +96,7 @@ import org.apache.airflow.sdk.internal.builderName as generatedBuilderName
@SupportedAnnotationTypes(
"org.apache.airflow.sdk.Builder.Dag",
"org.apache.airflow.sdk.Builder.Task",
+ "org.apache.airflow.sdk.Builder.If",
"org.apache.airflow.sdk.Builder.TaskGroup",
"org.apache.airflow.sdk.Builder.TaskHandler",
"org.apache.airflow.sdk.Builder.Deps",
@@ -356,29 +357,42 @@ class BuilderProcessor : AbstractProcessor() {
builderName: ClassName,
inGroup: Boolean,
): MethodSpec {
+ val condition = decl.kind == TaskKind.CONDITION
val method =
MethodSpec
.methodBuilder(decl.method.simpleName.toString())
.addModifiers(Modifier.PUBLIC, Modifier.DEFAULT)
- .returns(ParameterizedTypeName.get(TASK_HANDLE_TYPE, TypeName.get(decl.method.returnType).boxIfPossible()))
+ .returns(
+ if (condition) {
+ CONDITION_REF_TYPE
+ } else {
+ ParameterizedTypeName.get(TASK_HANDLE_TYPE, TypeName.get(decl.method.returnType).boxIfPossible())
+ },
+ )
decl.dataParams.forEach { method.addParameter(inType(it.type), it.name) }
val def = taskDefCode(decl, CodeBlock.of($$"$T.$L", builderName, decl.className))
// The view knows the group it belongs to, so the recorder is told where
// the task goes instead of deriving it from the task's ID.
val group = if (inGroup) CodeBlock.of("groupId()") else CodeBlock.of($$"$S", "")
+ val wrap = { call: CodeBlock -> if (condition) CodeBlock.of($$"$T.of($L)", CONDITION_REF_TYPE, call) else call }
if (decl.dataParams.isEmpty()) {
- method.addStatement($$"return $T.node($L, $L)", REFS_TYPE, group, def)
+ method.addStatement($$"return $L", wrap(CodeBlock.of($$"$T.node($L, $L)", REFS_TYPE, group, def)))
} else {
// The parameter names ride along so the serialized Dag can name each
// argument, as the binding spec ADR-0007 defines requires.
method.addStatement(
- $$"return $T.call($L, $L, $T.of($L), $L)",
- REFS_TYPE,
- group,
- def,
- LIST_TYPE,
- CodeBlock.join(decl.dataParams.map { CodeBlock.of($$"$S", it.name) }, ", "),
- decl.dataParams.joinToString { it.name },
+ $$"return $L",
+ wrap(
+ CodeBlock.of(
+ $$"$T.call($L, $L, $T.of($L), $L)",
+ REFS_TYPE,
+ group,
+ def,
+ LIST_TYPE,
+ CodeBlock.join(decl.dataParams.map { CodeBlock.of($$"$S", it.name) }, ", "),
+ decl.dataParams.joinToString { it.name },
+ ),
+ ),
)
}
return method.build()
@@ -396,7 +410,7 @@ class BuilderProcessor : AbstractProcessor() {
CodeBlock
.builder()
.add($$"new $T($S, $L.class)", TASK_DEF_TYPE, decl.id, classRef)
- explicitConfig(decl.method, TASK_ANNOTATION, TASK_STRUCTURAL_ATTRIBUTES, SchemaFields.TASK).forEach { (key, value) ->
+ explicitConfig(decl.method, decl.kind.annotation, TASK_STRUCTURAL_ATTRIBUTES, SchemaFields.TASK).forEach { (key, value) ->
taskDef.add($$".config($S, $L)", key, value)
}
return taskDef.build()
@@ -421,16 +435,24 @@ class BuilderProcessor : AbstractProcessor() {
val tasks = mutableListOf<TaskDeclaration>()
for (inner in el.enclosedElements) {
if (inner !is ExecutableElement) continue
- val ann = inner.getAnnotation(Builder.Task::class.java) ?: continue
+ val task = inner.getAnnotation(Builder.Task::class.java)
+ val condition = inner.getAnnotation(Builder.If::class.java)
+ if (task == null && condition == null) continue
+ require(task == null || condition == null) {
+ "Method '${inner.simpleName}' carries both @Builder.Task and @Builder.If; a condition is a " +
+ "task, so it carries @Builder.If alone"
+ }
+ val kind = if (condition != null) TaskKind.CONDITION else TaskKind.TASK
if (inner.isVarArgs) throw IllegalArgumentException("Cannot create task from vararg function ${inner.simpleName}")
- val localId = ann.id.ifBlank { inner.simpleName.toString() }
+ checkDeciderReturn(kind, inner)
+ val localId = (condition?.id ?: task!!.id).ifBlank { inner.simpleName.toString() }
require(tasks.none { it.method.simpleName.contentEquals(inner.simpleName) }) {
"Class ${el.simpleName} overloads task method '${inner.simpleName}'; a method's name is the " +
"name of its generated task class and of its wiring-view method, so rename one and keep its " +
- "task id with @Builder.Task(id = \"$localId\")"
+ "task id with ${kind.spelling}(id = \"$localId\")"
}
tasks +=
- TaskDeclaration(inner, (path + localId).joinToString("."), collectDataParams(inner), el, classPath)
+ TaskDeclaration(inner, (path + localId).joinToString("."), collectDataParams(inner), el, classPath, kind)
}
val groups = mutableListOf<GroupDeclaration>()
@@ -571,7 +593,7 @@ class BuilderProcessor : AbstractProcessor() {
val name = decl.method.simpleName.toString()
require(name !in reserved) {
"Task method '$name' clashes with a member of the wiring view; rename the method and keep " +
- "the task id with @Builder.Task(id = \"${decl.id.substringAfterLast('.')}\")"
+ "the task id with ${decl.kind.spelling}(id = \"${decl.id.substringAfterLast('.')}\")"
}
}
val accessors = mutableMapOf<String, GroupDeclaration>()
@@ -682,13 +704,31 @@ class BuilderProcessor : AbstractProcessor() {
}
}
+ /**
+ * Checks what a decider hands back, which is what the SDK reads its decision
+ * from. A plain task may return anything, including nothing.
+ */
+ private fun checkDeciderReturn(
+ kind: TaskKind,
+ method: ExecutableElement,
+ ) {
+ if (kind != TaskKind.CONDITION) return
+ val returns = method.returnType
+ val boolean = returns.kind == TypeKind.BOOLEAN || with(processingEnv) { isType(returns, BOXED_BOOLEAN_TYPE) }
+ require(boolean) {
+ "@Builder.If method '${method.simpleName}' returns $returns, but a condition returns boolean: " +
+ "true runs the task named by Then, false the one named by Else"
+ }
+ }
+
private fun buildTask(decl: TaskDeclaration): TypeSpec {
+ val condition = decl.kind == TaskKind.CONDITION
val executeSpec =
MethodSpec
- .methodBuilder("execute")
+ .methodBuilder(if (condition) "decide" else "execute")
.addAnnotation(Override::class.java)
.addModifiers(Modifier.PUBLIC)
- .returns(TypeName.VOID)
+ .returns(if (condition) TypeName.BOOLEAN else TypeName.VOID)
.addParameter(CONTEXT_TYPE, "context")
.addParameter(CLIENT_TYPE, "client")
.addException(Exception::class.java)
@@ -734,10 +774,11 @@ class BuilderProcessor : AbstractProcessor() {
}
}
- if (inner.returnType.kind == TypeKind.VOID) {
- $$"new $T().$L($L)"
- } else {
- $$"client.setXCom(new $T().$L($L))"
+ when {
+ // The SDK pushes a condition's result itself, after it has skipped the side not taken.
+ condition -> $$"return new $T().$L($L)"
+ inner.returnType.kind == TypeKind.VOID -> $$"new $T().$L($L)"
+ else -> $$"client.setXCom(new $T().$L($L))"
}.also {
executeSpec.addStatement(
it,
@@ -749,7 +790,7 @@ class BuilderProcessor : AbstractProcessor() {
return TypeSpec
.classBuilder(decl.className)
- .addSuperinterface(Task::class.java)
+ .addSuperinterface(if (condition) CONDITION_TASK_TYPE else ClassName.get(Task::class.java))
.addModifiers(Modifier.PUBLIC, Modifier.FINAL, Modifier.STATIC)
.addMethod(executeSpec.build())
.build()
@@ -893,6 +934,7 @@ private class TaskDeclaration(
val dataParams: List<DataParam>,
val owner: TypeElement,
val classPath: List<String> = emptyList(),
+ val kind: TaskKind = TaskKind.TASK,
) {
val className: String =
(classPath + method.simpleName.toString().replaceFirstChar(Char::uppercase)).joinToString("_")
@@ -924,12 +966,27 @@ private val DAG_SOURCE_TYPE = ClassName.get(DagSource::class.java)
private val REFS_TYPE = ClassName.get(Refs::class.java)
private val ARG_TYPE = ClassName.get(Arg::class.java)
private val TASK_HANDLE_TYPE = ClassName.get(TaskRef::class.java)
+private val CONDITION_TASK_TYPE = ClassName.get(ConditionTask::class.java)
+private val CONDITION_REF_TYPE = ClassName.get(ConditionRef::class.java)
+private val BOXED_BOOLEAN_TYPE = ClassName.get("java.lang", "Boolean")
private val DEPS_TYPE = ClassName.get(Deps::class.java)
private val GROUP_TYPE = DEPS_TYPE.nestedClass("TaskGroup")
private val LIST_TYPE = ClassName.get(List::class.java)
private const val DAG_ANNOTATION = "org.apache.airflow.sdk.Builder.Dag"
-private const val TASK_ANNOTATION = "org.apache.airflow.sdk.Builder.Task"
+
+/**
+ * What a declared task is, and the annotation that declares it. A decider
+ * differs from a plain task in what it implements and in what its wiring-view
+ * method hands back, but carries the same configuration attributes.
+ */
+private enum class TaskKind(
+ val annotation: String,
+ val spelling: String,
+) {
+ TASK("org.apache.airflow.sdk.Builder.Task", "@Builder.Task"),
+ CONDITION("org.apache.airflow.sdk.Builder.If", "@Builder.If"),
+}
private val RESERVED_VIEW_NAMES =
setOf(
diff --git a/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt b/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt
index a30dbacf008..70802c405aa 100644
--- a/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt
+++ b/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt
@@ -2161,6 +2161,204 @@ class BuilderTest {
)
}
+ @Test
+ @DisplayName("generate a condition task and a wiring view that names each side")
+ fun generateConditionTask() {
+ val compilation =
+ compile(
+ """
+ package org.apache.airflow.example;
+
+ import org.apache.airflow.sdk.Builder;
+
+ @Builder.Dag(id = "etl")
+ public class TestExample {
+ @Builder.If(id = "has_rows")
+ public boolean hasRows() {
+ return true;
+ }
+
+ @Builder.Task
+ public void load() {}
+
+ @Builder.Deps
+ static class Wiring implements TestExampleDeps {
+ void depends() {
+ hasRows().Then(load());
+ }
+ }
+ }
+ """,
+ )
+
+ assertThat(compilation).succeeded()
+ assertThat(compilation)
+ .generatedSourceFile("org.apache.airflow.example.TestExampleBuilder")
+ .hasSourceEquivalentTo(
+ "org.apache.airflow.example.TestExampleBuilder",
+ """
+ package org.apache.airflow.example;
+
+ import java.lang.Exception;
+ import java.lang.Override;
+ import java.util.List;
+ import org.apache.airflow.sdk.Client;
+ import org.apache.airflow.sdk.ConditionTask;
+ import org.apache.airflow.sdk.Context;
+ import org.apache.airflow.sdk.DagDef;
+ import org.apache.airflow.sdk.Task;
+ import org.apache.airflow.sdk.internal.DagSource;
+ import org.apache.airflow.sdk.internal.Refs;
+
+ public final class TestExampleBuilder {
+ public static DagDef build() {
+ var dag = DagSource.declaredBy(new DagDef("etl"), TestExample.class);
+ return Refs.record(dag, List.of("has_rows", "load"), List.of(), new TestExample.Wiring()::depends);
+ }
+
+ public static final class HasRows implements ConditionTask {
+ @Override
+ public boolean decide(Context context, Client client) throws Exception {
+ return new TestExample().hasRows();
+ }
+ }
+
+ public static final class Load implements Task {
+ @Override
+ public void execute(Context context, Client client) throws Exception {
+ new TestExample().load();
+ }
+ }
+ }
+ """,
+ )
+ assertThat(compilation)
+ .generatedSourceFile("org.apache.airflow.example.TestExampleDeps")
+ .hasSourceEquivalentTo(
+ "org.apache.airflow.example.TestExampleDeps",
+ """
+ package org.apache.airflow.example;
+
+ import java.lang.Void;
+ import org.apache.airflow.sdk.ConditionRef;
+ import org.apache.airflow.sdk.Deps;
+ import org.apache.airflow.sdk.TaskDef;
+ import org.apache.airflow.sdk.TaskRef;
+ import org.apache.airflow.sdk.internal.Refs;
+
+ public interface TestExampleDeps extends Deps {
+ default ConditionRef hasRows() {
+ return ConditionRef.of(Refs.node("", new TaskDef("has_rows", TestExampleBuilder.HasRows.class)));
+ }
+
+ default TaskRef<Void> load() {
+ return Refs.node("", new TaskDef("load", TestExampleBuilder.Load.class));
+ }
+ }
+ """,
+ )
+ }
+
+ @Test
+ @DisplayName("let a wiring view name a condition's sides in more than one statement")
+ fun conditionNamedAcrossStatements() {
+ val compilation =
+ compile(
+ """
+ package org.apache.airflow.example;
+
+ import org.apache.airflow.sdk.Builder;
+
+ @Builder.Dag(id = "etl")
+ public class TestExample {
+ @Builder.Task
+ public void extract() {}
+
+ @Builder.If(id = "has_rows")
+ public boolean hasRows() {
+ return true;
+ }
+
+ @Builder.Task
+ public void load() {}
+
+ @Builder.Task
+ public void reportEmpty() {}
+
+ @Builder.Deps
+ static class Wiring implements TestExampleDeps {
+ void depends() {
+ extract().before(hasRows());
+ hasRows().Then(load());
+ hasRows().Else(reportEmpty());
+ }
+ }
+ }
+ """,
+ )
+
+ assertThat(compilation).succeeded()
+ }
+
+ @Test
+ @DisplayName("reject a condition that does not return a boolean")
+ fun rejectNonBooleanCondition() {
+ val compilation =
+ compile(
+ """
+ package org.apache.airflow.example;
+ import org.apache.airflow.sdk.Builder;
+ @Builder.Dag
+ public class TestExample {
+ @Builder.If
+ public String hasRows() {
+ return "yes";
+ }
+
+ @Builder.Deps
+ static class Wiring implements TestExampleDeps {
+ void depends() {}
+ }
+ }
+ """,
+ )
+
+ assertThat(compilation).failed()
+ assertThat(compilation).hadErrorContaining(
+ "@Builder.If method 'hasRows' returns java.lang.String, but a condition returns boolean",
+ )
+ }
+
+ @Test
+ @DisplayName("reject a method that is both a task and a condition")
+ fun rejectTaskAndCondition() {
+ val compilation =
+ compile(
+ """
+ package org.apache.airflow.example;
+ import org.apache.airflow.sdk.Builder;
+ @Builder.Dag
+ public class TestExample {
+ @Builder.Task
+ @Builder.If
+ public boolean hasRows() {
+ return true;
+ }
+
+ @Builder.Deps
+ static class Wiring implements TestExampleDeps {
+ void depends() {}
+ }
+ }
+ """,
+ )
+
+ assertThat(compilation).failed()
+ assertThat(compilation).hadErrorContaining(
+ "Method 'hasRows' carries both @Builder.Task and @Builder.If",
+ )
+ }
+
@Test
@DisplayName("map the Dag to the annotated class, not its generated builder")
fun dagDeclaredByAnnotatedClass() {
diff --git a/java-sdk/scripts/ci/prek/check_serialization_conformance.py b/java-sdk/scripts/ci/prek/check_serialization_conformance.py
index 9869f4ccc5c..fe569c1ff7c 100755
--- a/java-sdk/scripts/ci/prek/check_serialization_conformance.py
+++ b/java-sdk/scripts/ci/prek/check_serialization_conformance.py
@@ -56,7 +56,7 @@ if __name__ == "__main__":
"--sdk",
"java",
"--supports",
- "literal_inputs",
+ "literal_inputs,branch",
"--",
*serializer,
]
diff --git a/java-sdk/sdk/build.gradle.kts b/java-sdk/sdk/build.gradle.kts
index 0159e1d2803..723f60085c2 100644
--- a/java-sdk/sdk/build.gradle.kts
+++ b/java-sdk/sdk/build.gradle.kts
@@ -581,6 +581,39 @@ abstract class GenerateDagDslTask : DefaultTask() {
| )
|
| /**
+ | * Marks a task method whose `boolean` picks one of two tasks; the
+ | * other is skipped.
+ | *
+ | * The method is an ordinary task method that returns `boolean`. Its
+ | * wiring-view method returns a [ConditionRef] rather than a
+ | * [TaskRef], so each side is named where the Dag is wired:
+ | *
+ | * ```java
+ | * @Builder.If(id = "has_rows")
+ | * public boolean hasRows(long rows) { return rows > 0; }
+ | *
+ | * @Builder.Deps
+ | * static class Wiring implements EtlDeps {
+ | * void depends() { hasRows(extract()).Then(load()).Else(reportEmpty()); }
+ | * }
+ | * ```
+ | *
+ | * The condition skips only the side not taken, so a task that runs
+ | * after both sides needs a trigger rule that tolerates one skipped
+ | * upstream, such as `none_failed_min_one_success`.
+ | *
+ | * Configuration attributes are [Task]'s, and apply to the deciding
+ | * task.
+ | */
+ | @Target(AnnotationTarget.FUNCTION)
+ | @MustBeDocumented
+ | annotation class If(
+ | /** Task ID. Empty derives it from the annotated function's name. */
+ | val id: String = "",
+ |${attrLines(taskFields)}
+ | )
+ |
+ | /**
| * Marks a method as the Java body of a task the Python Dag file
| * declares with `@task.stub`.
| *
diff --git a/java-sdk/sdk/module.md b/java-sdk/sdk/module.md
index 3d11a42da5e..523dbaf54a4 100644
--- a/java-sdk/sdk/module.md
+++ b/java-sdk/sdk/module.md
@@ -60,7 +60,7 @@ meaning of each dimension is defined in the
| capability: `task-args` | MUST † | ✓ | 3.4 | |
| capability: `dag-params` | MUST † | ✗ | – | |
| capability: `taskflow-dependencies` | MUST † | ✓ | 3.4 | |
-| capability: `branching` | SHOULD † | ✗ | – | |
+| capability: `branching` | SHOULD † | ✓ | 3.4 | DagDef.If / @Builder.If |
| capability: `dag-test` | SHOULD † | ✗ | – | |
| capability: `task-group` | MAY † | ✓ | 3.4 | |
| capability: `dynamic-task-mapping` | MAY † | ✗ | – | |
diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Bundle.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Bundle.kt
index bd055ab2c25..7937446c3f6 100644
--- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Bundle.kt
+++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Bundle.kt
@@ -66,8 +66,8 @@ class Bundle(
* @return This bundle, for chaining.
* @throws IllegalArgumentException if another Dag shares its ID, task
* handlers are already registered against it, a task depends on an
- * upstream not registered in the same Dag, or the dependencies contain a
- * cycle.
+ * upstream not registered in the same Dag, a task that decides which task
+ * runs is not declared with `If`, or the dependencies contain a cycle.
* @throws IllegalStateException if [Server.serve] has already been called.
*/
fun register(dag: DagDef): Bundle {
@@ -83,11 +83,30 @@ class Bundle(
"that is not registered in the same Dag"
}
}
+ val decides = DECIDER_TYPES.any { it.isAssignableFrom(def.definition) }
+ require(decides == (def.decider != null)) {
+ if (decides) {
+ "Task '${def.id}' runs '${def.definition.name}', which decides which task runs, but nothing names " +
+ "what it chooses; declare it with If(...), in Dag '${dag.id}'"
+ } else {
+ "Task '${def.id}' chooses which task runs, but '${def.definition.name}' decides nothing, in Dag '${dag.id}'"
+ }
+ }
+ def.decider?.let { decider ->
+ decider.describe(def)?.let { throw IllegalArgumentException("$it, in Dag '${dag.id}'") }
+ for (case in decider.cases) {
+ require(dag.tasks[case.id] === case) {
+ "Task '$taskId' in Dag '${dag.id}' can run task '${case.id}' " +
+ "that is not registered in the same Dag"
+ }
+ }
+ }
}
checkNoCycle(dag)
require(dags.putIfAbsent(dag.id, dag) == null) {
"Dags in bundle have duplicate ID: ${dag.id}"
}
+ dag.registered = true
return this
}
@@ -192,6 +211,9 @@ class Bundle(
private fun checkOpen() = check(!served) { "Server.serve has already been called; register everything before serve" }
}
+/** Task types the SDK runs through a decider, so each needs the sides it chooses between. */
+private val DECIDER_TYPES = listOf(ConditionTask::class.java)
+
// Reject cycles produced by before and after at registration time. This is (non-tailrec-eligible)
// recursive and could blow up with deep dependency chains. I kept the recursive implementation
// for readability since the scenario is unlikely; feel free to rewrite if it blows up for you.
diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Condition.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Condition.kt
new file mode 100644
index 00000000000..90150609c7b
--- /dev/null
+++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Condition.kt
@@ -0,0 +1,224 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.airflow.sdk
+
+import kotlin.Throws
+
+/**
+ * A task whose boolean decides which of two tasks runs; the other is skipped.
+ *
+ * Register one with [DagDef.If], then name the sides with [ConditionRef.Then]
+ * and [ConditionRef.Else]:
+ *
+ * ```java
+ * public class HasRows implements ConditionTask {
+ * @Override
+ * public boolean decide(Context context, Client client) {
+ * return ((Number) client.getXCom("extract")).longValue() > 0;
+ * }
+ * }
+ * ```
+ *
+ * The SDK runs [decide] and pushes its result as this task's return value, so
+ * [execute] is never called.
+ *
+ * @see DagDef.If
+ */
+interface ConditionTask : Task {
+ /**
+ * Decides which side of the condition runs.
+ *
+ * Any exception thrown marks the task instance as failed, and nothing is
+ * then pushed or skipped.
+ *
+ * @param context Runtime context for the current execution workload.
+ * @param client Client for Airflow API calls scoped to this execution.
+ * @return True to run the [ConditionRef.Then] side, false the
+ * [ConditionRef.Else] side.
+ * @throws Exception on failure; the task instance is marked failed.
+ */
+ @Throws(Exception::class)
+ fun decide(
+ context: Context,
+ client: Client,
+ ): Boolean
+
+ /** Never called: the SDK runs a condition through [decide]. */
+ override fun execute(
+ context: Context,
+ client: Client,
+ ): Unit =
+ throw IllegalStateException(
+ "Condition '${javaClass.name}' runs through decide(), so execute() is never called",
+ )
+}
+
+/**
+ * A condition registered with a Dag: name the task each outcome runs.
+ *
+ * ```java
+ * dag.If(HasRows.class).Then(load).Else(reportEmpty);
+ * ```
+ *
+ * A named task runs after the condition, so naming it records that edge, as
+ * [Deps.Flow.before] would. The condition skips only the side not taken; what
+ * runs after both sides needs a trigger rule that tolerates one skipped
+ * upstream, such as `none_failed_min_one_success`.
+ *
+ * @see DagDef.If
+ */
+class ConditionRef private constructor(
+ private val ref: TaskRef<Boolean>,
+ private val sides: ConditionDef,
+) : Deps.Flow {
+ companion object {
+ /**
+ * @suppress
+ *
+ * Marks a registered task as a condition. Public so a generated wiring
+ * view can call it; user code reaches a condition through [DagDef.If].
+ *
+ * A wiring view gets the same task back each time it names it, so a second
+ * call reuses the condition rather than declaring it again. [DagDef.If] and
+ * [TaskGroupRef.If] never reach that path, since registering a task rejects
+ * a duplicate ID first.
+ */
+ @JvmStatic
+ fun of(ref: TaskRef<Boolean>): ConditionRef {
+ ref.def.decider?.let { existing ->
+ require(existing is ConditionDef) { "Task '${ref.def.id}' already decides what to skip as a branch" }
+ return ConditionRef(ref, existing)
+ }
+ return ConditionRef(ref, ConditionDef().also { ref.def.decider = it })
+ }
+ }
+
+ /** Airflow task ID of the deciding task. */
+ val id: String get() = ref.def.id
+
+ /**
+ * Names the task that runs when the condition holds.
+ *
+ * @param task Task of the same Dag, which the condition skips when it does
+ * not hold.
+ * @return This condition, so [Else] can follow.
+ * @throws IllegalArgumentException if the task belongs to another Dag, or
+ * this side is already named.
+ */
+ @Suppress("ktlint:standard:function-naming")
+ fun Then(task: TaskRef<*>): ConditionRef = name("Then", task) { sides.whenTrue = it }
+
+ /**
+ * Names the task that runs when the condition does not hold. A condition
+ * needs no `Else`; without one, nothing is skipped when it holds.
+ *
+ * @param task Task of the same Dag, which the condition skips when it holds.
+ * @return This condition, for chaining.
+ * @throws IllegalArgumentException if the task belongs to another Dag, is
+ * the `Then` task, or this side is already named.
+ */
+ @Suppress("ktlint:standard:function-naming")
+ fun Else(task: TaskRef<*>): ConditionRef = name("Else", task) { sides.whenFalse = it }
+
+ /**
+ * Sets one task-level configuration value on the deciding task.
+ *
+ * @param key Airflow task setting name.
+ * @param value Value matching the key's schema type.
+ * @return This condition, for chaining.
+ * @throws IllegalArgumentException if the key is unknown or the value type
+ * does not match.
+ */
+ fun config(
+ key: String,
+ value: Any?,
+ ): ConditionRef {
+ ref.config(key, value)
+ return this
+ }
+
+ override fun nodes(): List<TaskDef> = ref.nodes()
+
+ override fun before(vararg next: Deps.Flow): ConditionRef {
+ ref.before(*next)
+ return this
+ }
+
+ override fun after(vararg previous: Deps.Flow): ConditionRef {
+ ref.after(*previous)
+ return this
+ }
+
+ private fun name(
+ side: String,
+ task: TaskRef<*>,
+ assign: (TaskDef) -> Unit,
+ ): ConditionRef {
+ sides.check(ref.def, side, task.def)
+ assign(task.def)
+ // The side runs after the condition, which is what puts the condition in
+ // the serialized Dag as its upstream and lets registration see a cycle
+ // that runs through a condition.
+ ref.before(task)
+ return this
+ }
+}
+
+/** The two sides of one condition, as [ConditionRef] named them. */
+internal class ConditionDef : DeciderDef {
+ var whenTrue: TaskDef? = null
+ var whenFalse: TaskDef? = null
+
+ override val cases: List<TaskDef> get() = listOfNotNull(whenTrue, whenFalse)
+
+ override fun describe(decider: TaskDef): String? =
+ if (whenTrue == null) {
+ "Condition '${decider.id}' names no task to run when it holds; call Then(...)"
+ } else {
+ null
+ }
+
+ override fun decide(
+ decider: TaskDef,
+ instance: Task,
+ context: Context,
+ client: Client,
+ ): Decision {
+ val held = (instance as ConditionTask).decide(context, client)
+ return Decision(held, listOfNotNull((if (held) whenFalse else whenTrue)?.id))
+ }
+
+ fun check(
+ decider: TaskDef,
+ side: String,
+ task: TaskDef,
+ ) {
+ requireUnregistered(decider, "Condition")
+ val named = if (side == "Then") whenTrue else whenFalse
+ require(named == null) {
+ "Condition '${decider.id}' already runs '${named!!.id}' on its $side side; name each side once"
+ }
+ requireSameDag(decider, task, "Condition")
+ require(task !== whenTrue && task !== whenFalse) {
+ "Condition '${decider.id}' already runs '${task.id}' on its other side, so the condition " +
+ "would decide nothing"
+ }
+ }
+}
diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/DagDef.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/DagDef.kt
index a269dff8c2c..40ead4ae608 100644
--- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/DagDef.kt
+++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/DagDef.kt
@@ -23,6 +23,7 @@ import org.apache.airflow.sdk.internal.DagSource
import org.apache.airflow.sdk.internal.GROUP_ID
import org.apache.airflow.sdk.internal.SchemaFields
import org.apache.airflow.sdk.internal.checkConfigValue
+import org.apache.airflow.sdk.internal.deriveTaskId
import org.apache.airflow.sdk.internal.validateTaskInput
import kotlin.Throws
@@ -58,6 +59,13 @@ class DagDef(
/** Edges with a task group at either end, in the order drawn. */
internal val groupEdges = linkedSetOf<Pair<Endpoint, Endpoint>>()
+ /**
+ * Whether a [Bundle] has taken this Dag. A decider checked when the Dag was
+ * registered cannot be changed afterwards, because nothing would check the
+ * change.
+ */
+ internal var registered: Boolean = false
+
/** Outermost class that declared this Dag, or `null` if it could not be told. */
internal var declaringClass: Class<*>? = DagSource.capture()
@@ -152,6 +160,47 @@ class DagDef(
return this
}
+ /**
+ * Declares a task whose boolean picks one of two tasks; the other is
+ * skipped.
+ *
+ * The task's ID is the class's simple name with its first character
+ * lowercased, so `HasRows.class` becomes `hasRows`. Use
+ * [If(id, definition)][If] to set it, and [ConditionRef.config] for the
+ * task's other settings:
+ *
+ * ```java
+ * dag.If(HasRows.class).config("retries", 2).Then(load).Else(reportEmpty);
+ * ```
+ *
+ * @param definition Class that implements [ConditionTask]. Must have a
+ * public no-arg constructor.
+ * @return The condition, to name each side on.
+ * @throws IllegalArgumentException if a task with the same ID is already
+ * registered.
+ */
+ @Suppress("ktlint:standard:function-naming")
+ fun If(definition: Class<out ConditionTask>): ConditionRef = If(deriveTaskId(definition), definition)
+
+ /**
+ * Declares a task whose boolean picks one of two tasks, under the task ID
+ * [id].
+ *
+ * @param id Task identifier, unique within this Dag.
+ * @param definition Class that implements [ConditionTask]. Must have a
+ * public no-arg constructor.
+ * @return The condition, to name each side on.
+ * @throws IllegalArgumentException if a task with the same ID is already
+ * registered.
+ *
+ * @see If
+ */
+ @Suppress("ktlint:standard:function-naming")
+ fun If(
+ id: String,
+ definition: Class<out ConditionTask>,
+ ): ConditionRef = ConditionRef.of(task(id, definition))
+
/**
* Declares a task group of this Dag.
*
@@ -326,6 +375,9 @@ class TaskDef(
internal val upstreams = linkedSetOf<TaskDef>()
internal var owner: DagDef? = null
+ /** What this task decides to run, for a condition or a branch; null otherwise. */
+ internal var decider: DeciderDef? = null
+
/**
* Sets one task-level configuration value.
*
diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Decider.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Decider.kt
new file mode 100644
index 00000000000..15f4234b20d
--- /dev/null
+++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Decider.kt
@@ -0,0 +1,90 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.airflow.sdk
+
+/**
+ * What one deciding task can run, and what it therefore skips.
+ *
+ * A [TaskDef] carrying one is a decider: the runtime asks it which of its
+ * [cases] the task chose, skips the rest, and the serialized task carries
+ * `_can_skip_downstream` so Airflow re-skips a cleared one.
+ */
+internal interface DeciderDef {
+ /** Every task this decider can run, in the order they were named. */
+ val cases: List<TaskDef>
+
+ /** Why this decider is incomplete, or null when it is ready to register. */
+ fun describe(decider: TaskDef): String?
+
+ /**
+ * Runs the deciding task [instance] and reports what it chose.
+ *
+ * @throws Exception whatever the task body throws, which fails the task
+ * without pushing or skipping anything.
+ */
+ fun decide(
+ decider: TaskDef,
+ instance: Task,
+ context: Context,
+ client: Client,
+ ): Decision
+}
+
+/**
+ * What a deciding task chose: the value it pushes as its return value, and
+ * the task IDs of the cases it did not take.
+ */
+internal class Decision(
+ val value: Any,
+ val skipped: List<String>,
+)
+
+/** Mirrors `SkipMixin.skip` in `task-sdk/src/airflow/sdk/bases/skipmixin.py`. */
+internal const val SKIPMIXIN_XCOM_KEY = "skipmixin_key"
+internal const val SKIPMIXIN_SKIPPED = "skipped"
+
+/**
+ * Rejects a case named after the Dag was registered, which the checks that run
+ * at registration would never see.
+ */
+internal fun requireUnregistered(
+ decider: TaskDef,
+ what: String,
+) {
+ val owner = decider.owner
+ require(owner?.registered != true) {
+ "$what '${decider.id}' of Dag '${owner?.id}' is already registered; name every case before " +
+ "the Dag is added to a Bundle"
+ }
+}
+
+/** Rejects a case that belongs to another Dag, which no edge of this one can reach. */
+internal fun requireSameDag(
+ decider: TaskDef,
+ case: TaskDef,
+ what: String,
+) {
+ val owner = decider.owner
+ val caseOwner = case.owner
+ require(owner == null || caseOwner == null || owner === caseOwner) {
+ "$what '${decider.id}' of Dag '${owner?.id}' cannot run task '${case.id}' of Dag " +
+ "'${caseOwner?.id}'; name a task of the same Dag"
+ }
+}
diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/TaskGroupRef.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/TaskGroupRef.kt
index dc6448a5eb0..e2d7d9613ef 100644
--- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/TaskGroupRef.kt
+++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/TaskGroupRef.kt
@@ -19,6 +19,8 @@
package org.apache.airflow.sdk
+import org.apache.airflow.sdk.internal.deriveTaskId
+
/**
* A group of tasks in a Dag, shown in the Airflow UI as one node that expands:
* Python's `TaskGroup`.
@@ -74,6 +76,36 @@ class TaskGroupRef internal constructor(
return TaskRef(def)
}
+ /**
+ * Declares a condition in this group, as [DagDef.If] does for the Dag.
+ *
+ * @param definition Class that implements [ConditionTask]. Must have a
+ * public no-arg constructor.
+ * @return The condition, to name each side on.
+ * @throws IllegalArgumentException if the Dag already has a task or task
+ * group with the resulting ID.
+ */
+ @Suppress("ktlint:standard:function-naming")
+ fun If(definition: Class<out ConditionTask>): ConditionRef = If(deriveTaskId(definition), definition)
+
+ /**
+ * Declares a condition in this group under the task ID `<group ID>.<id>`.
+ *
+ * @param id Task ID within this group.
+ * @param definition Class that implements [ConditionTask]. Must have a
+ * public no-arg constructor.
+ * @return The condition, to name each side on.
+ * @throws IllegalArgumentException if the Dag already has a task or task
+ * group with the resulting ID.
+ *
+ * @see If
+ */
+ @Suppress("ktlint:standard:function-naming")
+ fun If(
+ id: String,
+ definition: Class<out ConditionTask>,
+ ): ConditionRef = ConditionRef.of(task(id, definition))
+
/**
* Nests a task group inside this one.
*
diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Client.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Client.kt
index a4bb15c44cf..f3174e77e2a 100644
--- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Client.kt
+++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Client.kt
@@ -33,6 +33,7 @@ import org.apache.airflow.sdk.execution.comm.OKResponse
import org.apache.airflow.sdk.execution.comm.PutVariable
import org.apache.airflow.sdk.execution.comm.SetTaskStateStore
import org.apache.airflow.sdk.execution.comm.SetXCom
+import org.apache.airflow.sdk.execution.comm.SkipDownstreamTasks
import org.apache.airflow.sdk.execution.comm.TaskStateStoreResult
import org.apache.airflow.sdk.execution.comm.VariableResult
import org.apache.airflow.sdk.execution.comm.XComResult
@@ -102,6 +103,13 @@ interface Client {
)
fun clearTaskStateStore(tiId: UUID)
+
+ /**
+ * Marks [taskIds] of the running Dag run skipped, as a decider's untaken
+ * branches are. Airflow leaves a task instance that is already running,
+ * succeeded or failed alone.
+ */
+ fun skipDownstreamTasks(taskIds: List<String>)
}
/**
@@ -231,4 +239,10 @@ class CoordinatorClient(
override fun clearTaskStateStore(tiId: UUID) {
runBlocking { exec.communicate<OKResponse>(ClearTaskStateStore().also { it.tiId = tiId }) }
}
+
+ override fun skipDownstreamTasks(taskIds: List<String>) {
+ if (taskIds.isEmpty()) return
+ val message = SkipDownstreamTasks().also { it.tasks = ArrayList<Any>(taskIds) }
+ runBlocking { exec.communicate<Unit>(message) }
+ }
}
diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Serde.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Serde.kt
index 89245cf6946..bdc719a3a68 100644
--- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Serde.kt
+++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Serde.kt
@@ -143,6 +143,8 @@ private fun serializeTask(
"is_stub" to true,
)
argBindings(taskId, def)?.let { data["_arg_bindings"] = it }
+ // Lets NotPreviouslySkippedDep re-skip a cleared downstream, as Python's SkipMixin does.
+ if (def.decider != null) data["_can_skip_downstream"] = true
// Emit only config entries that differ from their schema default, mirroring
// Python BaseSerialization's "omit hard-coded default" behavior, which the Go
// and TypeScript SDKs mirror too. Operator fields are stored unwrapped, so the
diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Task.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Task.kt
index 6490e221ee9..182dfe1cee6 100644
--- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Task.kt
+++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Task.kt
@@ -23,6 +23,11 @@ import kotlinx.coroutines.CancellationException
import org.apache.airflow.sdk.Bundle
import org.apache.airflow.sdk.Client
import org.apache.airflow.sdk.Context
+import org.apache.airflow.sdk.DeciderDef
+import org.apache.airflow.sdk.SKIPMIXIN_SKIPPED
+import org.apache.airflow.sdk.SKIPMIXIN_XCOM_KEY
+import org.apache.airflow.sdk.Task
+import org.apache.airflow.sdk.TaskDef
import org.apache.airflow.sdk.execution.comm.AssetProfile
import org.apache.airflow.sdk.execution.comm.RetryTask
import org.apache.airflow.sdk.execution.comm.StartupDetails
@@ -66,6 +71,26 @@ internal object TaskResult {
fun failure(shouldRetry: Boolean) = if (shouldRetry) retry() else of(TaskState.State.FAILED)
}
+/**
+ * Runs a deciding task: Airflow skips every case it did not choose, and its
+ * choice becomes the task's return value.
+ */
+private fun decide(
+ decider: DeciderDef,
+ taskDef: TaskDef,
+ instance: Task,
+ context: Context,
+ client: Client,
+) {
+ val decision = decider.decide(taskDef, instance, context, client)
+ if (decision.skipped.isNotEmpty()) {
+ // Written before the skip so a cleared case is skipped again, as SkipMixin does.
+ client.setXCom(SKIPMIXIN_XCOM_KEY, mapOf(SKIPMIXIN_SKIPPED to decision.skipped))
+ client.impl.skipDownstreamTasks(decision.skipped)
+ }
+ client.setXCom(value = decision.value)
+}
+
internal object TaskRunner {
val logger = Logger(TaskRunner::class)
@@ -104,7 +129,9 @@ internal object TaskRunner {
return TaskResult.failure(request.tiContext.shouldRetry)
}
return try {
- instance.execute(Context.from(request, taskDef), client)
+ val context = Context.from(request, taskDef)
+ val decider = taskDef.decider
+ if (decider == null) instance.execute(context, client) else decide(decider, taskDef, instance, context, client)
TaskResult.success()
} catch (e: CancellationException) {
throw e // Let coroutine cancellation propagate so the task coroutine unwinds.
diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/Ids.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/Ids.kt
index b64551e9b01..b2824be3c91 100644
--- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/Ids.kt
+++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/Ids.kt
@@ -27,3 +27,13 @@ package org.apache.airflow.sdk.internal
* against the same pattern the SDK enforces; not user-facing API.
*/
val GROUP_ID: Regex = Regex("[A-Za-z0-9_-]+")
+
+/**
+ * @suppress
+ *
+ * The task ID a task declared from a class alone carries: the class's simple
+ * name with its first character lowercased, so `HasRows.class` and a
+ * `hasRows` task method name agree. Public so the annotation processor can
+ * check an ID it derives the same way; not user-facing API.
+ */
+fun deriveTaskId(definition: Class<*>): String = definition.simpleName.replaceFirstChar { it.lowercaseChar() }
diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ArgTestSupport.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ArgTestSupport.kt
index 0929dfe4e4e..6c6af528765 100644
--- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ArgTestSupport.kt
+++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ArgTestSupport.kt
@@ -90,6 +90,8 @@ internal class FakeXComTransport(
) = throw NotImplementedError()
override fun clearTaskStateStore(tiId: UUID) = throw NotImplementedError()
+
+ override fun skipDownstreamTasks(taskIds: List<String>) = throw NotImplementedError()
}
internal fun startupDetails(argBindings: List<Map<String, Any?>>?): StartupDetails =
diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ClientTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ClientTest.kt
index dce06ca0977..2ff51df4f36 100644
--- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ClientTest.kt
+++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ClientTest.kt
@@ -104,6 +104,8 @@ private class FakeTransport(
override fun clearTaskStateStore(tiId: UUID) {
calls.add(StateStoreCall("clear", tiId))
}
+
+ override fun skipDownstreamTasks(taskIds: List<String>) = Unit
}
class ClientTest {
diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ConditionTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ConditionTest.kt
new file mode 100644
index 00000000000..1cebf7ff8e1
--- /dev/null
+++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ConditionTest.kt
@@ -0,0 +1,208 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.airflow.sdk
+
+import org.apache.airflow.sdk.execution.serializeDag
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Assertions.assertFalse
+import org.junit.jupiter.api.Assertions.assertThrows
+import org.junit.jupiter.api.Assertions.assertTrue
+import org.junit.jupiter.api.DisplayName
+import org.junit.jupiter.api.Test
+
+/** A condition the tests register; it always holds. */
+class HasRows : ConditionTask {
+ override fun decide(
+ context: Context,
+ client: Client,
+ ): Boolean = true
+}
+
+internal class ConditionTest {
+ private fun dagWithSides(): Triple<DagDef, ConditionRef, Pair<TaskRef<Unit>, TaskRef<Unit>>> {
+ val dag = DagDef("d")
+ val load = dag.task<Unit>("load", NoopTask::class.java)
+ val reportEmpty = dag.task<Unit>("report_empty", NoopTask::class.java)
+ return Triple(dag, dag.If(HasRows::class.java), load to reportEmpty)
+ }
+
+ @Test
+ @DisplayName("Should take the condition's task ID from its class when none is given")
+ fun shouldDeriveTaskIdFromClass() {
+ val (dag, condition, sides) = dagWithSides()
+ condition.Then(sides.first)
+
+ assertEquals("hasRows", condition.id)
+ assertTrue("hasRows" in dag.tasks)
+ }
+
+ @Test
+ @DisplayName("Should run each named side after the condition")
+ fun shouldRunEachSideAfterTheCondition() {
+ val (dag, condition, sides) = dagWithSides()
+ val (load, reportEmpty) = sides
+ condition.Then(load).Else(reportEmpty)
+
+ val decider = dag.tasks.getValue("hasRows")
+ assertEquals(setOf(decider), load.def.upstreams)
+ assertEquals(setOf(decider), reportEmpty.def.upstreams)
+ }
+
+ @Test
+ @DisplayName("Should mark only a deciding task as able to skip what runs after it")
+ fun shouldMarkDeciderAsSkipping() {
+ val (dag, condition, sides) = dagWithSides()
+ condition.Then(sides.first)
+
+ val tasks = serializeDag(dag, "", ".")["tasks"] as List<*>
+
+ @Suppress("UNCHECKED_CAST")
+ val byId = tasks.associate { task -> ((task as Map<String, Any?>)["__var"] as Map<String, Any?>).let { it["task_id"] to it } }
+ assertEquals(true, (byId["hasRows"] as Map<*, *>)["_can_skip_downstream"])
+ assertFalse("_can_skip_downstream" in (byId["load"] as Map<*, *>))
+ }
+
+ @Test
+ @DisplayName("Should reject naming the same side twice")
+ fun shouldRejectNamingASideTwice() {
+ val (_, condition, sides) = dagWithSides()
+ condition.Then(sides.first)
+
+ val error = assertThrows(IllegalArgumentException::class.java) { condition.Then(sides.second) }
+
+ assertEquals("Condition 'hasRows' already runs 'load' on its Then side; name each side once", error.message)
+ }
+
+ @Test
+ @DisplayName("Should reject a condition whose sides are the same task")
+ fun shouldRejectIdenticalSides() {
+ val (_, condition, sides) = dagWithSides()
+ condition.Then(sides.first)
+
+ val error = assertThrows(IllegalArgumentException::class.java) { condition.Else(sides.first) }
+
+ assertEquals(
+ "Condition 'hasRows' already runs 'load' on its other side, so the condition would decide nothing",
+ error.message,
+ )
+ }
+
+ @Test
+ @DisplayName("Should reject a side declared in another Dag")
+ fun shouldRejectSideFromAnotherDag() {
+ val (_, condition, _) = dagWithSides()
+ val other = DagDef("other").task<Unit>("load", NoopTask::class.java)
+
+ val error = assertThrows(IllegalArgumentException::class.java) { condition.Then(other) }
+
+ assertEquals(
+ "Condition 'hasRows' of Dag 'd' cannot run task 'load' of Dag 'other'; name a task of the same Dag",
+ error.message,
+ )
+ }
+
+ @Test
+ @DisplayName("Should reject registering a condition that names no task to run")
+ fun shouldRejectConditionWithoutThen() {
+ val (dag, _, _) = dagWithSides()
+
+ val error = assertThrows(IllegalArgumentException::class.java) { Bundle().register(dag) }
+
+ assertEquals(
+ "Condition 'hasRows' names no task to run when it holds; call Then(...), in Dag 'd'",
+ error.message,
+ )
+ }
+
+ @Test
+ @DisplayName("Should reject registering a decider class as an ordinary task")
+ fun shouldRejectDeciderRegisteredAsTask() {
+ val dag = DagDef("d")
+ dag.task<Unit>("has_rows", HasRows::class.java)
+
+ val error = assertThrows(IllegalArgumentException::class.java) { Bundle().register(dag) }
+
+ assertEquals(
+ "Task 'has_rows' runs '${HasRows::class.java.name}', which decides which task runs, but nothing names " +
+ "what it chooses; declare it with If(...), in Dag 'd'",
+ error.message,
+ )
+ }
+
+ @Test
+ @DisplayName("Should reject a task that chooses which task runs but decides nothing")
+ fun shouldRejectChooserThatDecidesNothing() {
+ val dag = DagDef("d")
+ val task = dag.task<Boolean>("gate", NoopTask::class.java)
+ ConditionRef.of(task)
+
+ val error = assertThrows(IllegalArgumentException::class.java) { Bundle().register(dag) }
+
+ assertEquals(
+ "Task 'gate' chooses which task runs, but '${NoopTask::class.java.name}' decides nothing, in Dag 'd'",
+ error.message,
+ )
+ }
+
+ @Test
+ @DisplayName("Should reuse the condition when a task already declared as one is declared again")
+ fun shouldReuseConditionDeclaredTwice() {
+ val (_, condition, sides) = dagWithSides()
+ val again = ConditionRef.of(TaskRef(condition.nodes().single()))
+
+ condition.Then(sides.first)
+ again.Else(sides.second)
+
+ val error = assertThrows(IllegalArgumentException::class.java) { again.Then(sides.second) }
+ assertEquals(
+ "Condition 'hasRows' already runs 'load' on its Then side; name each side once",
+ error.message,
+ )
+ }
+
+ @Test
+ @DisplayName("Should reject naming a side after the Dag was registered")
+ fun shouldRejectNamingASideAfterRegistration() {
+ val (dag, condition, sides) = dagWithSides()
+ condition.Then(sides.first)
+ Bundle().register(dag)
+
+ val error = assertThrows(IllegalArgumentException::class.java) { condition.Else(sides.second) }
+
+ assertEquals(
+ "Condition 'hasRows' of Dag 'd' is already registered; name every case before the Dag is " +
+ "added to a Bundle",
+ error.message,
+ )
+ }
+
+ @Test
+ @DisplayName("Should declare a condition inside a task group under the group's prefix")
+ fun shouldDeclareConditionInGroup() {
+ val dag = DagDef("d")
+ val checks = dag.taskGroup("checks")
+ val load = checks.task<Unit>("load", NoopTask::class.java)
+ val condition = checks.If("has_rows", HasRows::class.java)
+ condition.Then(load)
+
+ assertEquals("checks.has_rows", condition.id)
+ assertEquals(listOf("checks.load", "checks.has_rows"), dag.tasks.keys.toList())
+ }
+}
diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/conformance/SerializeJava.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/conformance/SerializeJava.kt
index 7d9aa937010..4c2f7f84e4e 100644
--- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/conformance/SerializeJava.kt
+++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/conformance/SerializeJava.kt
@@ -34,6 +34,8 @@ import com.fasterxml.jackson.dataformat.yaml.YAMLMapper
import org.apache.airflow.sdk.Arg
import org.apache.airflow.sdk.Bundle
import org.apache.airflow.sdk.Client
+import org.apache.airflow.sdk.ConditionRef
+import org.apache.airflow.sdk.ConditionTask
import org.apache.airflow.sdk.Context
import org.apache.airflow.sdk.DagDef
import org.apache.airflow.sdk.Deps
@@ -55,6 +57,13 @@ class ConformanceTask : Task {
) = Unit
}
+class ConformanceCondition : ConditionTask {
+ override fun decide(
+ context: Context,
+ client: Client,
+ ) = true
+}
+
fun main(args: Array<String>) {
require(args.size == 2) { "usage: SerializeJava <test_dags.yaml> <output.json>" }
val cases = YAMLMapper().readTree(File(args[0])).path("dags")
@@ -80,16 +89,21 @@ private fun buildDag(case: JsonNode): DagDef {
groups[groupId] = if (parentId.isEmpty()) dag.taskGroup(localId) else groups.getValue(parentId).taskGroup(localId)
}
- val tasks = linkedMapOf<String, TaskRef<Unit>>()
+ val tasks = linkedMapOf<String, TaskRef<*>>()
+ // A condition names tasks that may be declared after it, so the sides are wired once every task exists.
+ val conditions = mutableListOf<Pair<ConditionRef, JsonNode>>()
case.path("tasks").forEach { task ->
val groupId = task.path("group").asText("")
val localId = task.path("task_id").asText()
+ val branch = task.path("branch")
+ val definition = if (branch.isMissingNode) ConformanceTask::class.java else ConformanceCondition::class.java
val ref =
if (groupId.isEmpty()) {
- dag.task<Unit>(localId, ConformanceTask::class.java)
+ dag.task<Any?>(localId, definition)
} else {
- groups.getValue(groupId).task(localId, ConformanceTask::class.java)
+ groups.getValue(groupId).task<Any?>(localId, definition)
}
+ if (!branch.isMissingNode) conditions += asCondition(ref) to branch
task.path("spec").fields().forEach { (key, value) -> ref.config(key, toValue(SchemaFields.TASK, key, value)) }
// A task's `upstream` handles and its `literals` are its call arguments, in that order, so the
// Dag carries the binding spec a stub call would. Names are positional, as the Go SDK names
@@ -103,6 +117,11 @@ private fun buildDag(case: JsonNode): DagDef {
tasks[ref.def.id] = ref
}
+ conditions.forEach { (condition, branch) ->
+ condition.Then(tasks.getValue(branch.path("then").asText()))
+ branch.path("else").takeIf { !it.isMissingNode }?.let { condition.Else(tasks.getValue(it.asText())) }
+ }
+
case.path("order_edges").forEach { edge ->
val node = { id: String -> groups[id] as Deps.Flow? ?: tasks.getValue(id) }
node(edge[0].asText()).before(node(edge[1].asText()))
@@ -110,6 +129,10 @@ private fun buildDag(case: JsonNode): DagDef {
return dag
}
+/** The handle of a task declared as a decider, whose type argument no caller reads. */
+@Suppress("UNCHECKED_CAST")
+private fun asCondition(ref: TaskRef<*>): ConditionRef = ConditionRef.of(ref as TaskRef<Boolean>)
+
/** Reads a YAML value as the Java type the config key takes. */
private fun toValue(
table: Map<String, Field>,
diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/CommTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/CommTest.kt
index 7f161ed9b5c..1d9b5cc4ace 100644
--- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/CommTest.kt
+++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/CommTest.kt
@@ -508,6 +508,33 @@ class CommsTest {
Assertions.assertEquals(tiId.toString(), clearBody["ti_id"])
}
+ @Test
+ @DisplayName("skipDownstreamTasks sends the task IDs to skip and accepts the supervisor's empty response")
+ @Timeout(value = 30, unit = TimeUnit.SECONDS)
+ fun skipDownstreamTasksSendsTaskIds() {
+ val (body, failure) =
+ roundTrip(::emptyResponseFrame) { it.impl.skipDownstreamTasks(listOf("report_empty", "audit")) }
+
+ Assertions.assertNull(failure, "skipDownstreamTasks should return normally on an empty response, got $failure")
+ Assertions.assertEquals("SkipDownstreamTasks", body["type"])
+ Assertions.assertEquals(listOf("report_empty", "audit"), body["tasks"])
+ }
+
+ @Test
+ @DisplayName("skipDownstreamTasks sends nothing when there is nothing to skip")
+ @Timeout(value = 30, unit = TimeUnit.SECONDS)
+ fun skipDownstreamTasksSendsNothingForEmptyList() {
+ val toClient = ByteChannel(autoFlush = true)
+ val fromClient = ByteChannel(autoFlush = true)
+ val comm = CoordinatorComm(toClient, fromClient)
+
+ CoordinatorClient(comm).skipDownstreamTasks(emptyList())
+
+ runBlocking { fromClient.flushAndClose() }
+ Assertions.assertTrue(fromClient.isClosedForRead, "no frame should have been written")
+ comm.close()
+ }
+
@Test
@DisplayName("Should fail a pending call when the coordinator socket closes")
@Timeout(value = 30, unit = TimeUnit.SECONDS)
diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskTest.kt
index 5ce508b708d..a1ca3f7553d 100644
--- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskTest.kt
+++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/TaskTest.kt
@@ -21,6 +21,7 @@ package org.apache.airflow.sdk.execution
import org.apache.airflow.sdk.Bundle
import org.apache.airflow.sdk.Client
+import org.apache.airflow.sdk.ConditionTask
import org.apache.airflow.sdk.Context
import org.apache.airflow.sdk.DagDef
import org.apache.airflow.sdk.Task
@@ -182,6 +183,77 @@ class TaskTest {
Assertions.assertInstanceOf(SucceedTask::class.java, result)
}
+ @Test
+ @DisplayName("Should push a condition's result and skip the side it did not take")
+ fun shouldSkipTheSideNotTaken() {
+ TestCondition.decision = true
+ val transport = RecordingTransport()
+
+ val result =
+ runTask(conditionBundle(withElse = true), startupDetails(taskId = "gate"), Client(startupDetails("gate"), transport))
+
+ Assertions.assertInstanceOf(SucceedTask::class.java, result)
+ Assertions.assertEquals(listOf("report_empty"), transport.skipped)
+ Assertions.assertEquals(
+ listOf("skipmixin_key" to mapOf("skipped" to listOf("report_empty")), "return_value" to true),
+ transport.xComs,
+ )
+ Assertions.assertEquals(
+ listOf("xcom:skipmixin_key", "skip:[report_empty]", "xcom:return_value"),
+ transport.events,
+ )
+ }
+
+ @Test
+ @DisplayName("Should skip the then side when a condition does not hold")
+ fun shouldSkipThenSideWhenConditionIsFalse() {
+ TestCondition.decision = false
+ val transport = RecordingTransport()
+
+ runTask(conditionBundle(withElse = true), startupDetails(taskId = "gate"), Client(startupDetails("gate"), transport))
+
+ Assertions.assertEquals(listOf("load"), transport.skipped)
+ Assertions.assertEquals(false, transport.xComs.last().second)
+ Assertions.assertEquals(listOf("xcom:skipmixin_key", "skip:[load]", "xcom:return_value"), transport.events)
+ }
+
+ @Test
+ @DisplayName("Should skip nothing when a one-sided condition holds")
+ fun shouldSkipNothingWhenOneSidedConditionHolds() {
+ TestCondition.decision = true
+ val transport = RecordingTransport()
+
+ runTask(conditionBundle(withElse = false), startupDetails(taskId = "gate"), Client(startupDetails("gate"), transport))
+
+ Assertions.assertEquals(emptyList<String>(), transport.skipped)
+ Assertions.assertEquals(listOf("return_value" to true), transport.xComs)
+ Assertions.assertEquals(listOf("xcom:return_value"), transport.events)
+ }
+
+ @Test
+ @DisplayName("Should fail a condition that throws, without pushing or skipping anything")
+ fun shouldFailConditionThatThrows() {
+ TestCondition.decision = null
+ val transport = RecordingTransport()
+
+ val result =
+ runTask(conditionBundle(withElse = true), startupDetails(taskId = "gate"), Client(startupDetails("gate"), transport))
+
+ Assertions.assertInstanceOf(TaskState::class.java, result)
+ Assertions.assertEquals(TaskState.State.FAILED, (result as TaskState).state)
+ Assertions.assertEquals(emptyList<String>(), transport.skipped)
+ Assertions.assertEquals(emptyList<Pair<String, Any>>(), transport.xComs)
+ }
+
+ private fun conditionBundle(withElse: Boolean): Bundle {
+ val dag = DagDef("test_dag")
+ val load = dag.task<Unit>("load", SuccessTask::class.java)
+ val reportEmpty = dag.task<Unit>("report_empty", SuccessTask::class.java)
+ val condition = dag.If("gate", TestCondition::class.java).Then(load)
+ if (withElse) condition.Else(reportEmpty)
+ return Bundle(listOf(dag))
+ }
+
private fun bundleWith(
taskId: String,
taskClass: Class<out Task>,
@@ -271,9 +343,87 @@ class TaskTest {
): Unit = throw UnsupportedOperationException("not used in test")
override fun clearTaskStateStore(tiId: UUID): Unit = throw UnsupportedOperationException("not used in test")
+
+ override fun skipDownstreamTasks(taskIds: List<String>): Unit = throw UnsupportedOperationException("not used in test")
},
)
+ /** Decides what [decision] holds; a null one throws, standing for a condition body that fails. */
+ class TestCondition : ConditionTask {
+ override fun decide(
+ context: Context,
+ client: Client,
+ ): Boolean = decision ?: throw IllegalStateException("boom")
+
+ companion object {
+ var decision: Boolean? = true
+ }
+ }
+
+ /** Records what a deciding task pushed and asked to skip, and the order of the two. */
+ private class RecordingTransport : org.apache.airflow.sdk.execution.Client {
+ val xComs = mutableListOf<Pair<String, Any>>()
+ val skipped = mutableListOf<String>()
+ val events = mutableListOf<String>()
+
+ override fun setXCom(
+ key: String,
+ value: Any,
+ dagId: String,
+ taskId: String,
+ runId: String,
+ mapIndex: Int,
+ ) {
+ xComs += key to value
+ events += "xcom:$key"
+ }
+
+ override fun skipDownstreamTasks(taskIds: List<String>) {
+ skipped += taskIds
+ events += "skip:$taskIds"
+ }
+
+ override fun getConnection(id: String) = throw UnsupportedOperationException("not used in test")
+
+ override fun getVariable(key: String) = throw UnsupportedOperationException("not used in test")
+
+ override fun setVariable(
+ key: String,
+ value: String,
+ description: String?,
+ ): Unit = throw UnsupportedOperationException("not used in test")
+
+ override fun deleteVariable(key: String): Unit = throw UnsupportedOperationException("not used in test")
+
+ override fun getXCom(
+ key: String,
+ dagId: String,
+ taskId: String,
+ runId: String,
+ mapIndex: Int?,
+ includePriorDates: Boolean,
+ ) = throw UnsupportedOperationException("not used in test")
+
+ override fun getTaskStateStore(
+ tiId: UUID,
+ key: String,
+ ) = throw UnsupportedOperationException("not used in test")
+
+ override fun setTaskStateStore(
+ tiId: UUID,
+ key: String,
+ value: Any,
+ expiresAt: OffsetDateTime?,
+ ): Unit = throw UnsupportedOperationException("not used in test")
+
+ override fun deleteTaskStateStore(
+ tiId: UUID,
+ key: String,
+ ): Unit = throw UnsupportedOperationException("not used in test")
+
+ override fun clearTaskStateStore(tiId: UUID): Unit = throw UnsupportedOperationException("not used in test")
+ }
+
class SuccessTask : Task {
override fun execute(
context: Context,
diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/ArgValuesTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/ArgValuesTest.kt
index b9fd480e10c..c8ca101267c 100644
--- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/ArgValuesTest.kt
+++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/ArgValuesTest.kt
@@ -128,6 +128,8 @@ internal class ArgValuesTest {
): Unit = throw NotImplementedError()
override fun clearTaskStateStore(tiId: UUID): Unit = throw NotImplementedError()
+
+ override fun skipDownstreamTasks(taskIds: List<String>): Unit = throw NotImplementedError()
},
)
diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/RefsTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/RefsTest.kt
index 3a06f7c472b..59e36c182bd 100644
--- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/RefsTest.kt
+++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/RefsTest.kt
@@ -21,6 +21,8 @@ package org.apache.airflow.sdk.internal
import org.apache.airflow.sdk.Arg
import org.apache.airflow.sdk.Client
+import org.apache.airflow.sdk.ConditionRef
+import org.apache.airflow.sdk.ConditionTask
import org.apache.airflow.sdk.Context
import org.apache.airflow.sdk.DagDef
import org.apache.airflow.sdk.Deps
@@ -47,6 +49,13 @@ private class NoopRefTask : Task {
) = Unit
}
+private class NoopCondition : ConditionTask {
+ override fun decide(
+ context: Context,
+ client: Client,
+ ) = true
+}
+
internal class RefsTest {
@Test
@DisplayName("Should register the task, record inputs, and wire handle edges")
@@ -203,4 +212,24 @@ internal class RefsTest {
assertEquals("Dag 'd' has no task group 'staging'", error.message)
}
+
+ @Test
+ @DisplayName("Should let a view name a condition's sides in separate statements")
+ fun shouldNameConditionSidesAcrossStatements() {
+ val dag = DagDef("d")
+ Refs.record(dag, listOf("has_rows", "load", "report_empty"), emptyList()) {
+ fun hasRows() = ConditionRef.of(Refs.node<Boolean>("", TaskDef("has_rows", NoopCondition::class.java)))
+ hasRows().Then(Refs.node<Unit>("", TaskDef("load", NoopRefTask::class.java)))
+ hasRows().Else(Refs.node<Unit>("", TaskDef("report_empty", NoopRefTask::class.java)))
+ }
+
+ assertEquals(
+ listOf("load", "report_empty"),
+ dag.tasks
+ .getValue("has_rows")
+ .decider!!
+ .cases
+ .map { it.id },
+ )
+ }
}
(airflow) 04/04: Java SDK: Trigger another Dag's run from a native Java Dag (#74440)
Posted by ja...@apache.org.
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 <68...@users.noreply.github.com>
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