Skip to content

Commit 542db5d

Browse files
committed
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. (cherry picked from commit af569b2)
1 parent 8125524 commit 542db5d

26 files changed

Lines changed: 1929 additions & 44 deletions

File tree

‎.pre-commit-config.yaml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -444,7 +444,7 @@ repos:
444444
^java-sdk/scripts/ci/prek/check_serialization_conformance\.py$|
445445
^java-sdk/sdk/schema/dag-schema\.json$|
446446
^java-sdk/sdk/build\.gradle\.kts$|
447-
^java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/(Branch|Bundle|Condition|DagDef|Decider|Deps|Arg|TaskGroupRef)\.kt$|
447+
^java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/(Branch|Bundle|Condition|DagDef|Decider|Deps|Arg|TaskGroupRef|TriggerDagRun)\.kt$|
448448
^java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/execution/Serde\.kt$|
449449
^java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/(Refs|Fields)\.kt$|
450450
^java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/conformance/.*$|

‎java-sdk/README.md‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -620,8 +620,8 @@ prek hook regenerate it.
620620
| state: `success` | MUST | ✓ | 3.3 | |
621621
| state: `failed` | MUST | ✓ | 3.3 | |
622622
| state: `up_for_retry` | MUST | ✓ | 3.3 | RetryTask |
623-
| state: `skipped` | SHOULD | ✗ | – | runtime does not emit TaskState skipped yet |
624-
| state: `deferred` | MAY | ✗ | – | runtime does not emit DeferTask yet |
623+
| state: `skipped` | SHOULD | ✓ | 3.4 | a TriggerDagRun task that skips an existing run |
624+
| state: `deferred` | MAY | ✓ | 3.4 | a TriggerDagRun task that waits with deferrable |
625625
| state: `up_for_reschedule` | MAY | ✗ | – | runtime does not emit RescheduleTask yet |
626626
| state: `awaiting_input` | MAY | ✗ | – | runtime does not emit AwaitInputTask yet |
627627
| state: `removed` | MAY | ✓ | 3.3 | |

‎java-sdk/capabilities.yaml‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -23,8 +23,8 @@ min_airflow_version: "3.3"
2323
# JAR manifest. The render hook fails if the two disagree.
2424
supervisor_schema_version: "2026-10-30"
2525

26-
# The runtime terminates a task with SucceedTask, RetryTask, or TaskState (failed/removed); it does
27-
# not yet emit skipped, DeferTask, RescheduleTask, or AwaitInputTask.
26+
# The runtime terminates a task with SucceedTask, RetryTask, DeferTask, or TaskState
27+
# (failed/skipped/removed); it does not yet emit RescheduleTask or AwaitInputTask.
2828
states:
2929
success:
3030
supported: true
@@ -37,11 +37,13 @@ states:
3737
since: "3.3"
3838
note: "RetryTask"
3939
skipped:
40-
supported: false
41-
note: "runtime does not emit TaskState skipped yet"
40+
supported: true
41+
since: "3.4"
42+
note: "a TriggerDagRun task that skips an existing run"
4243
deferred:
43-
supported: false
44-
note: "runtime does not emit DeferTask yet"
44+
supported: true
45+
since: "3.4"
46+
note: "a TriggerDagRun task that waits with deferrable"
4547
up_for_reschedule:
4648
supported: false
4749
note: "runtime does not emit RescheduleTask yet"

‎java-sdk/example/src/java/org/apache/airflow/example/nativedag/AnnotationExample.java‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,13 @@ public TaskId pickReport(long transformed) {
8686
: AnnotationExampleBuilder.TaskIds.REPORT_SHORT;
8787
}
8888

89+
// A task that starts a run of another Dag. The method runs when this Dag is
90+
// built, not when the task runs, so it takes no arguments.
91+
@Builder.Task(id = "trigger_downstream")
92+
public TriggerDagRun triggerDownstream() {
93+
return new TriggerDagRun("java_native_interface_example").config("note", "from the annotation example");
94+
}
95+
8996
// A task group: everything it declares is prefixed with its id, so this is
9097
// the task "checks.audit".
9198
@Builder.TaskGroup(id = "checks")
@@ -103,8 +110,10 @@ static class Wiring implements AnnotationExampleDeps {
103110
void depends() {
104111
var extracted = extract();
105112
var transformed = transform(extracted, lit(1.5));
106-
hasRows(transformed).then(load(transformed)).orElse(loadEmpty());
113+
var loaded = load(transformed);
114+
hasRows(transformed).then(loaded).orElse(loadEmpty());
107115
pickReport(transformed).option(reportLong()).option(reportShort());
116+
loaded.before(triggerDownstream());
108117
// Ordering-only edge: the checks group runs after extract, with no data
109118
// flowing.
110119
extracted.before(checks());

‎java-sdk/example/src/java/org/apache/airflow/example/nativedag/InterfaceExample.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -114,12 +114,18 @@ public static DagDef build() {
114114

115115
var reportLong = dag.task("report_long", ReportLong.class);
116116
var reportShort = dag.task("report_short", ReportShort.class);
117+
// A task that starts a run of another Dag; it runs no Java code.
118+
var trigger =
119+
dag.task(
120+
"trigger_downstream",
121+
new TriggerDagRun("java_native_annotation_example").config("wait_for_completion", false));
117122

118123
transform.after(extract);
119124
// With no task id given, a decider takes one from its class: "hasRows" and
120125
// "pickReport".
121126
dag.If(HasRows.class).after(transform).then(load).orElse(loadEmpty);
122127
dag.Branch(PickReport.class).after(transform).option(reportLong).option(reportShort);
128+
load.before(trigger);
123129
return dag;
124130
}
125131
}

‎java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt‎

Lines changed: 42 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -277,7 +277,7 @@ class BuilderProcessor : AbstractProcessor() {
277277
)
278278
builderClass.addMethod(buildMethod.build())
279279

280-
declarations.forEach { builderClass.addType(buildTask(it)) }
280+
declarations.filterNot { it.kind == TaskKind.TRIGGER }.forEach { builderClass.addType(buildTask(it)) }
281281
// Only a branch needs to name a task in code, so a Dag without one keeps
282282
// the generated builder to the classes that run its tasks.
283283
if (declarations.any { it.kind == TaskKind.BRANCH }) builderClass.addType(buildTaskIds(declarations))
@@ -394,8 +394,11 @@ class BuilderProcessor : AbstractProcessor() {
394394
.methodBuilder(decl.method.simpleName.toString())
395395
.addModifiers(Modifier.PUBLIC, Modifier.DEFAULT)
396396
.returns(
397-
decl.kind.refType
398-
?: ParameterizedTypeName.get(TASK_HANDLE_TYPE, TypeName.get(decl.method.returnType).boxIfPossible()),
397+
when {
398+
decl.kind.refType != null -> decl.kind.refType
399+
decl.kind == TaskKind.TRIGGER -> ParameterizedTypeName.get(TASK_HANDLE_TYPE, VOID_TYPE)
400+
else -> ParameterizedTypeName.get(TASK_HANDLE_TYPE, TypeName.get(decl.method.returnType).boxIfPossible())
401+
},
399402
)
400403
decl.dataParams.forEach { method.addParameter(inType(it.type), it.name) }
401404
val def = taskDefCode(decl, CodeBlock.of($$"$T.$L", builderName, decl.className))
@@ -439,7 +442,13 @@ class BuilderProcessor : AbstractProcessor() {
439442
val taskDef =
440443
CodeBlock
441444
.builder()
442-
.add($$"new $T($S, $L.class)", TASK_DEF_TYPE, decl.id, classRef)
445+
.apply {
446+
if (decl.kind == TaskKind.TRIGGER) {
447+
add($$"new $T($S, new $T().$L())", TASK_DEF_TYPE, decl.id, ClassName.get(decl.owner), decl.method.simpleName)
448+
} else {
449+
add($$"new $T($S, $L.class)", TASK_DEF_TYPE, decl.id, classRef)
450+
}
451+
}
443452
explicitConfig(decl.method, decl.kind.annotation, TASK_STRUCTURAL_ATTRIBUTES, SchemaFields.TASK).forEach { (key, value) ->
444453
taskDef.add($$".config($S, $L)", key, value)
445454
}
@@ -466,17 +475,28 @@ class BuilderProcessor : AbstractProcessor() {
466475
for (inner in el.enclosedElements) {
467476
if (inner !is ExecutableElement) continue
468477
val declared =
469-
TaskKind.entries.filter { kind -> inner.annotationMirrors.any { it.names(kind.annotation) } }
478+
TaskKind.declaring.filter { kind -> inner.annotationMirrors.any { it.names(kind.annotation) } }
470479
if (declared.isEmpty()) continue
471-
val kind =
480+
val declaredKind =
472481
declared.singleOrNull()
473482
?: throw IllegalArgumentException(
474483
"Method '${inner.simpleName}' carries ${declared.joinToString { it.spelling }}; a task is " +
475484
"declared by one of them alone",
476485
)
486+
// A task method that hands back a TriggerDagRun declares what to trigger
487+
// rather than a body to run, so it is read when the Dag is built.
488+
val triggers =
489+
declaredKind == TaskKind.TASK && with(processingEnv) { isType(inner.returnType, TRIGGER_TYPE) }
490+
val kind = if (triggers) TaskKind.TRIGGER else declaredKind
477491
val annotated = declaredId(inner, kind)
478492
if (inner.isVarArgs) throw IllegalArgumentException("Cannot create task from vararg function ${inner.simpleName}")
479493
checkDeciderReturn(kind, inner)
494+
if (kind == TaskKind.TRIGGER) {
495+
require(inner.parameters.isEmpty()) {
496+
"@Builder.Task method '${inner.simpleName}' returns a TriggerDagRun, so it runs when the Dag " +
497+
"is built rather than when the task runs; it takes no parameters"
498+
}
499+
}
480500
val localId = annotated.ifBlank { inner.simpleName.toString() }
481501
require(tasks.none { it.method.simpleName.contentEquals(inner.simpleName) }) {
482502
"Class ${el.simpleName} overloads task method '${inner.simpleName}'; a method's name is the " +
@@ -623,7 +643,8 @@ class BuilderProcessor : AbstractProcessor() {
623643
}
624644
}
625645
val byClassName = mutableMapOf<String, TaskDeclaration>()
626-
declarations.forEach { decl ->
646+
// A task that triggers a Dag run has no generated class to collide with.
647+
declarations.filterNot { it.kind == TaskKind.TRIGGER }.forEach { decl ->
627648
byClassName.put(decl.className, decl)?.let { first ->
628649
throw IllegalArgumentException(
629650
"Task methods '${first.id}' and '${decl.id}' both generate the task class " +
@@ -769,7 +790,7 @@ class BuilderProcessor : AbstractProcessor() {
769790
) {
770791
val returns = method.returnType
771792
when (kind) {
772-
TaskKind.TASK -> return
793+
TaskKind.TASK, TaskKind.TRIGGER -> return
773794
TaskKind.CONDITION ->
774795
require(returns.kind == TypeKind.BOOLEAN || with(processingEnv) { isType(returns, BOXED_BOOLEAN_TYPE) }) {
775796
"@Builder.If method '${method.simpleName}' returns $returns, but a condition returns boolean: " +
@@ -801,7 +822,7 @@ class BuilderProcessor : AbstractProcessor() {
801822
.addModifiers(Modifier.PUBLIC)
802823
.returns(
803824
when (decl.kind) {
804-
TaskKind.TASK -> TypeName.VOID
825+
TaskKind.TASK, TaskKind.TRIGGER -> TypeName.VOID
805826
TaskKind.CONDITION -> TypeName.BOOLEAN
806827
TaskKind.BRANCH -> TASK_ID_TYPE
807828
},
@@ -1048,6 +1069,8 @@ private val CONDITION_REF_TYPE = ClassName.get(ConditionRef::class.java)
10481069
private val BRANCH_TASK_TYPE = ClassName.get(TaskIdBranchTask::class.java)
10491070
private val BRANCH_REF_TYPE = ClassName.get(BranchRef::class.java)
10501071
private val TASK_ID_TYPE = ClassName.get(TaskId::class.java)
1072+
private val TRIGGER_TYPE = ClassName.get(TriggerDagRun::class.java)
1073+
private val VOID_TYPE = ClassName.get("java.lang", "Void")
10511074

10521075
/** Name of the generated holder of a Dag's task ids, which no task class may take. */
10531076
private const val TASK_IDS = "TaskIds"
@@ -1072,8 +1095,18 @@ private enum class TaskKind(
10721095
val refType: ClassName?,
10731096
) {
10741097
TASK("org.apache.airflow.sdk.Builder.Task", "@Builder.Task", "execute", TASK_TYPE, null),
1098+
TRIGGER("org.apache.airflow.sdk.Builder.Task", "@Builder.Task", "execute", TASK_TYPE, null),
10751099
CONDITION("org.apache.airflow.sdk.Builder.If", "@Builder.If", "decide", CONDITION_TASK_TYPE, CONDITION_REF_TYPE),
10761100
BRANCH("org.apache.airflow.sdk.Builder.Branch", "@Builder.Branch", "choose", BRANCH_TASK_TYPE, BRANCH_REF_TYPE),
1101+
;
1102+
1103+
companion object {
1104+
/**
1105+
* The kinds an annotation declares. TRIGGER is left out: it shares
1106+
* [TASK]'s annotation and is told apart by the method's return type.
1107+
*/
1108+
val declaring = listOf(TASK, CONDITION, BRANCH)
1109+
}
10771110
}
10781111

10791112
/** Whether this annotation is the one [name] qualifies. */

‎java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt‎

Lines changed: 172 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2450,6 +2450,178 @@ class BuilderTest {
24502450
)
24512451
}
24522452

2453+
@Test
2454+
@DisplayName("read a task that triggers a Dag run when the Dag is built")
2455+
fun generateTriggerDagRunTask() {
2456+
val compilation =
2457+
compile(
2458+
"""
2459+
package org.apache.airflow.example;
2460+
2461+
import org.apache.airflow.sdk.Builder;
2462+
import org.apache.airflow.sdk.TriggerDagRun;
2463+
2464+
@Builder.Dag(id = "etl")
2465+
public class TestExample {
2466+
@Builder.Task(id = "trigger_downstream", retries = 2)
2467+
public TriggerDagRun triggerDownstream() {
2468+
return new TriggerDagRun("downstream_etl").config("wait_for_completion", true);
2469+
}
2470+
2471+
@Builder.Deps
2472+
static class Wiring implements TestExampleDeps {
2473+
void depends() {
2474+
triggerDownstream();
2475+
}
2476+
}
2477+
}
2478+
""",
2479+
)
2480+
2481+
assertThat(compilation).succeeded()
2482+
assertThat(compilation)
2483+
.generatedSourceFile("org.apache.airflow.example.TestExampleBuilder")
2484+
.hasSourceEquivalentTo(
2485+
"org.apache.airflow.example.TestExampleBuilder",
2486+
"""
2487+
package org.apache.airflow.example;
2488+
2489+
import java.util.List;
2490+
import org.apache.airflow.sdk.DagDef;
2491+
import org.apache.airflow.sdk.internal.DagSource;
2492+
import org.apache.airflow.sdk.internal.Refs;
2493+
2494+
public final class TestExampleBuilder {
2495+
public static DagDef build() {
2496+
var dag = DagSource.declaredBy(new DagDef("etl"), TestExample.class);
2497+
return Refs.record(dag, List.of("trigger_downstream"), List.of(), new TestExample.Wiring()::depends);
2498+
}
2499+
}
2500+
""",
2501+
)
2502+
assertThat(compilation)
2503+
.generatedSourceFile("org.apache.airflow.example.TestExampleDeps")
2504+
.hasSourceEquivalentTo(
2505+
"org.apache.airflow.example.TestExampleDeps",
2506+
"""
2507+
package org.apache.airflow.example;
2508+
2509+
import java.lang.Void;
2510+
import org.apache.airflow.sdk.Deps;
2511+
import org.apache.airflow.sdk.TaskDef;
2512+
import org.apache.airflow.sdk.TaskRef;
2513+
import org.apache.airflow.sdk.internal.Refs;
2514+
2515+
public interface TestExampleDeps extends Deps {
2516+
default TaskRef<Void> triggerDownstream() {
2517+
return Refs.node("", new TaskDef("trigger_downstream", new TestExample().triggerDownstream()).config("retries", 2));
2518+
}
2519+
}
2520+
""",
2521+
)
2522+
}
2523+
2524+
@Test
2525+
@DisplayName("build a Dag that reuses one task handle as a condition side and as an upstream")
2526+
fun buildADagThatReusesATaskHandle() {
2527+
val compilation =
2528+
compile(
2529+
"""
2530+
package org.apache.airflow.example;
2531+
2532+
import org.apache.airflow.sdk.Builder;
2533+
import org.apache.airflow.sdk.TriggerDagRun;
2534+
2535+
@Builder.Dag(id = "etl")
2536+
public class TestExample {
2537+
@Builder.Task
2538+
public long extract() {
2539+
return 42L;
2540+
}
2541+
2542+
@Builder.Task
2543+
public void load(long rows) {}
2544+
2545+
@Builder.Task
2546+
public void loadEmpty() {}
2547+
2548+
@Builder.If(id = "has_rows")
2549+
public boolean hasRows(long rows) {
2550+
return rows > 0;
2551+
}
2552+
2553+
@Builder.Task(id = "trigger_downstream")
2554+
public TriggerDagRun triggerDownstream() {
2555+
return new TriggerDagRun("downstream_etl");
2556+
}
2557+
2558+
@Builder.Deps
2559+
static class Wiring implements TestExampleDeps {
2560+
void depends() {
2561+
var extracted = extract();
2562+
var loaded = load(extracted);
2563+
hasRows(extracted).then(loaded).orElse(loadEmpty());
2564+
loaded.before(triggerDownstream());
2565+
}
2566+
}
2567+
}
2568+
""",
2569+
)
2570+
2571+
assertThat(compilation).succeeded()
2572+
// The wiring only runs when the Dag is built, so load it and build it.
2573+
val classes =
2574+
compilation
2575+
.generatedFiles()
2576+
.filter { it.kind == JavaFileObject.Kind.CLASS }
2577+
.associate {
2578+
it.name
2579+
.removePrefix("/CLASS_OUTPUT/")
2580+
.removeSuffix(".class")
2581+
.replace('/', '.') to it.openInputStream().readBytes()
2582+
}
2583+
val loader =
2584+
object : ClassLoader(javaClass.classLoader) {
2585+
override fun findClass(name: String): Class<*> {
2586+
val bytes = classes[name] ?: throw ClassNotFoundException(name)
2587+
return defineClass(name, bytes, 0, bytes.size)
2588+
}
2589+
}
2590+
2591+
Bundle().register(loader.loadClass("org.apache.airflow.example.TestExample"))
2592+
}
2593+
2594+
@Test
2595+
@DisplayName("reject a task that triggers a Dag run and takes parameters")
2596+
fun rejectTriggerDagRunWithParameters() {
2597+
val compilation =
2598+
compile(
2599+
"""
2600+
package org.apache.airflow.example;
2601+
import org.apache.airflow.sdk.Builder;
2602+
import org.apache.airflow.sdk.TriggerDagRun;
2603+
@Builder.Dag
2604+
public class TestExample {
2605+
@Builder.Task
2606+
public TriggerDagRun triggerDownstream(long rows) {
2607+
return new TriggerDagRun("downstream_etl");
2608+
}
2609+
2610+
@Builder.Deps
2611+
static class Wiring implements TestExampleDeps {
2612+
void depends() {}
2613+
}
2614+
}
2615+
""",
2616+
)
2617+
2618+
assertThat(compilation).failed()
2619+
assertThat(compilation).hadErrorContaining(
2620+
"@Builder.Task method 'triggerDownstream' returns a TriggerDagRun, so it runs when the Dag is " +
2621+
"built rather than when the task runs; it takes no parameters",
2622+
)
2623+
}
2624+
24532625
@Test
24542626
@DisplayName("reject a branch that does not return a task id")
24552627
fun rejectNonTaskIdBranch() {

0 commit comments

Comments
 (0)