From 71e897830715744f175b000ef0f8c3ebdcfff679 Mon Sep 17 00:00:00 2001 From: "Jason(Zhe-You) Liu" <68415893+jason810496@users.noreply.github.com> Date: Thu, 22 Oct 2026 20:16:08 +0800 Subject: [PATCH] [v3-4-test] Java SDK: Let a native Java Dag set one queue for all its tasks (#74538) * Java SDK: Let a native Java Dag set one queue for all its tasks Every task of a native Java Dag runs on the Java coordinator, so each one needs a queue that queue_to_coordinator routes there. Airflow's Dag has no queue setting, so authors had to repeat the same queue on every task. The Go and TypeScript SDKs already take a Dag-level queue that each task inherits unless it sets its own; this gives Java the same. * Reject a blank Dag-level queue and fix the native Dag queue docs A blank queue matched no coordinator, so every task stayed queued. The parsing docs still said there was no Dag-level queue and routed a queue the examples do not use. * Reject a blank Dag-level queue in the annotation processor * Carry the blank rule on the schema field, for Dag and task queues alike * Set the queue once per Dag in the e2e native bundle * Put the blank check on Field and tidy the comments --------- (cherry picked from commit fc9c4b0b0036a481da1ae773e76f6885a470a20f) Co-authored-by: Jason(Zhe-You) Liu <68415893+jason810496@users.noreply.github.com> --- .../language-sdks/java.rst | 15 +++-- .../airflow/e2e/NativeBundleBuilder.java | 36 +++++------ .../airflow/e2e/nativedag/AnnotationDag.java | 25 ++++---- .../airflow/e2e/nativedag/DynamicDags.java | 5 +- .../airflow/e2e/nativedag/TargetDag.java | 5 +- .../test_java_sdk_native_dag.py | 17 ++--- .../example/nativedag/AnnotationExample.java | 1 + .../example/nativedag/InterfaceExample.java | 1 + .../example/nativedag/TargetExample.java | 1 + .../apache/airflow/sdk/AirflowProcessor.kt | 6 +- .../airflow/sdk/AirflowProcessorTest.kt | 53 ++++++++++++++- java-sdk/sdk/build.gradle.kts | 28 +++++++- .../kotlin/org/apache/airflow/sdk/DagDef.kt | 17 +++-- .../org/apache/airflow/sdk/TriggerDagRun.kt | 1 + .../org/apache/airflow/sdk/execution/Serde.kt | 10 ++- .../org/apache/airflow/sdk/internal/Fields.kt | 22 ++++++- .../org/apache/airflow/sdk/DagDefTest.kt | 28 ++++++++ .../apache/airflow/sdk/execution/SerdeTest.kt | 64 +++++++++++++++++++ 18 files changed, 271 insertions(+), 64 deletions(-) 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 f1ba8c35874d5..a41abe2a5b06a 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst @@ -352,7 +352,8 @@ Annotate a plain Java class and let the SDK generate the boilerplate at compile * - ``@Airflow.Dag(id = "...")`` - Marks a class as a Dag that Java itself owns. Attributes (``schedule``, ``description``, ``tags``, ``catchup``, …) are Airflow's own Dag settings; only attributes written - explicitly are applied. See :ref:`java-sdk/native-dags`. + explicitly are applied. ``queue`` is the queue each task runs on unless the task sets its + own. See :ref:`java-sdk/native-dags`. * - ``@Airflow.Task(id = "...")`` - Marks a method as a task of a Java-owned Dag. If ``id`` is omitted the method name is used. Further attributes (``retries``, ``queue``, ``retryDelay``, …) are Airflow's own @@ -642,6 +643,11 @@ Native Java Dags A Dag can also be authored entirely in Java: the annotations (or the ``DagDef`` / ``TaskDef`` objects) carry the configuration, and Java declares the graph. +Every task of the Dag runs on the Java coordinator, so it needs a queue that +:ref:`queue_to_coordinator ` sends there. Set ``queue`` once on the Dag, with +``@Airflow.Dag(queue = "java")`` or ``dag.config("queue", "java")``, and each task inherits it, including a +``TriggerDagRun`` task. A task's own ``queue`` wins over the Dag's. A ``queue`` that is set cannot be blank. + Building the Dag in Java ~~~~~~~~~~~~~~~~~~~~~~~~ @@ -922,7 +928,7 @@ The Dag processor runs the JAR's main class to list its Dags, so it needs a Java "kwargs": {"java_executable": "/usr/lib/jvm/java-17-openjdk/bin/java"} } } - queue_to_coordinator = {"java-native": "java-native"} + queue_to_coordinator = {"java": "java-native"} Once a ``JavaCoordinator`` is configured, the Dag processor parses the executable JARs of every Dag bundle, so it needs this ``[sdk]`` configuration and a JRE. With one ``JavaCoordinator``, it parses them all. @@ -953,9 +959,8 @@ several Java coordinators, an unmapped bundle's JARs fail to parse too. ``Main-Class`` are parsed. List JARs that should not be parsed in ``.airflowignore``. * Two JARs in the bundle that set the same ``Main-Class`` fail to parse, because the JVM would load the classes of only one of them. Keep one in the bundle. -* Set ``queue`` on every task, with ``@Airflow.Task(queue = "java-native")`` or - ``TaskDef.config("queue", "java-native")``, so it runs on the coordinator's queue. There is no - Dag-level queue yet. +* Set ``queue`` on the Dag, as :ref:`java-sdk/native-dags` shows, so its tasks run on a queue that + ``queue_to_coordinator`` routes to a ``JavaCoordinator``, such as ``java`` above. * A task runs the JAR of its Dag on the ``JavaCoordinator`` that its queue routes to, which need not be the one that parsed the JAR. For example, a queue can route to a coordinator that uses another JDK. A task whose queue routes to another kind of coordinator fails without retries. A task whose JAR is missing, diff --git a/airflow-e2e-tests/java-native-bundle/src/java/org/apache/airflow/e2e/NativeBundleBuilder.java b/airflow-e2e-tests/java-native-bundle/src/java/org/apache/airflow/e2e/NativeBundleBuilder.java index 0c52b3fe426cb..b20f324c4132e 100644 --- a/airflow-e2e-tests/java-native-bundle/src/java/org/apache/airflow/e2e/NativeBundleBuilder.java +++ b/airflow-e2e-tests/java-native-bundle/src/java/org/apache/airflow/e2e/NativeBundleBuilder.java @@ -31,9 +31,9 @@ * Bundle for the native Dag E2E tests: Dags declared entirely in Java, parsed by the Dag processor * from this JAR. * - *

This is the bundle's main class, so it is the Dag source the Airflow UI shows. Every task sets - * {@code queue}, which {@code queue_to_coordinator} routes to the {@code java-jdk} coordinator. - * The {@code java-native} coordinator only parses this JAR. + *

This is the bundle's main class, so it is the Dag source the Airflow UI shows. Each Dag sets + * {@code queue} once and every task inherits it, which {@code queue_to_coordinator} routes to the + * {@code java-jdk} coordinator. The {@code java-native} coordinator only parses this JAR. */ public class NativeBundleBuilder { public static final String QUEUE = "java-native"; @@ -110,34 +110,30 @@ public static DagDef buildDag() { new DagDef("java_native_e2e") .config("description", "Native Java Dag of the Airflow E2E tests") .config("catchup", false) - .config("tags", List.of("java-sdk", "native", "e2e")); + .config("tags", List.of("java-sdk", "native", "e2e")) + .config("queue", QUEUE); - var extract = dag.task("extract", Extract.class).config("queue", QUEUE); - var transform = dag.task("transform", Transform.class).config("queue", QUEUE); - var load = dag.task("load", Load.class).config("queue", QUEUE); + var extract = dag.task("extract", Extract.class); + var transform = dag.task("transform", Transform.class); + var load = dag.task("load", Load.class); transform.after(extract).before(load); - var reportMany = dag.task("report_many", ReportMany.class).config("queue", QUEUE); - var reportFew = dag.task("report_few", ReportFew.class).config("queue", QUEUE); - dag.If("has_rows", HasRows.class).after(transform).config("queue", QUEUE).Then(reportMany).Else(reportFew); + var reportMany = dag.task("report_many", ReportMany.class); + var reportFew = dag.task("report_few", ReportFew.class); + dag.If("has_rows", HasRows.class).after(transform).Then(reportMany).Else(reportFew); - var reportLong = dag.task("report_long", ReportLong.class).config("queue", QUEUE); - var reportShort = dag.task("report_short", ReportShort.class).config("queue", QUEUE); - dag.Switch("pick_report", PickReport.class) - .after(transform) - .config("queue", QUEUE) - .Case(reportLong) - .Case(reportShort); + var reportLong = dag.task("report_long", ReportLong.class); + var reportShort = dag.task("report_short", ReportShort.class); + dag.Switch("pick_report", PickReport.class).after(transform).Case(reportLong).Case(reportShort); // 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_e2e")).config("queue", QUEUE); + var trigger = dag.task("trigger_downstream", new TriggerDagRun("java_native_target_e2e")); load.before(trigger); // Ordering-only edge: the checks group runs after extract, with no data flowing. var checks = dag.taskGroup("checks"); - checks.task("audit", Audit.class).config("queue", QUEUE); + checks.task("audit", Audit.class); extract.before(checks); return dag; diff --git a/airflow-e2e-tests/java-native-bundle/src/java/org/apache/airflow/e2e/nativedag/AnnotationDag.java b/airflow-e2e-tests/java-native-bundle/src/java/org/apache/airflow/e2e/nativedag/AnnotationDag.java index 78531c2ccb5cd..3e245c3074f8d 100644 --- a/airflow-e2e-tests/java-native-bundle/src/java/org/apache/airflow/e2e/nativedag/AnnotationDag.java +++ b/airflow-e2e-tests/java-native-bundle/src/java/org/apache/airflow/e2e/nativedag/AnnotationDag.java @@ -24,58 +24,59 @@ import org.apache.airflow.sdk.*; -/** A native Java Dag declared with annotations; its tasks use the interface Dag's queue. */ +/** A native Java Dag declared with annotations; it sets the same queue as the interface Dag. */ @Airflow.Dag( id = "java_native_annotation_e2e", description = "Native Java Dag of the Airflow E2E tests, declared with annotations", catchup = false, + queue = QUEUE, tags = {"java-sdk", "native", "e2e"}) public class AnnotationDag { - @Airflow.Task(id = "extract", queue = QUEUE) + @Airflow.Task(id = "extract") public long extract() { return 42L; } - @Airflow.Task(id = "transform", queue = QUEUE) + @Airflow.Task(id = "transform") public long transform(long extracted, double factor) { return (long) (extracted * factor); } /** Fails the run unless the value reached it, so a successful run proves the XCom flow. */ - @Airflow.Task(id = "load", queue = QUEUE) + @Airflow.Task(id = "load") public void load(long transformed) { if (transformed != 63L) { throw new IllegalStateException("load expected 63 from transform, got " + transformed); } } - @Airflow.Task(id = "report_many", queue = QUEUE) + @Airflow.Task(id = "report_many") public void reportMany() {} - @Airflow.Task(id = "report_few", queue = QUEUE) + @Airflow.Task(id = "report_few") public void reportFew() {} /** Picks one of {@link #reportMany} and {@link #reportFew}; the other is skipped. */ - @Airflow.If(id = "has_rows", queue = QUEUE) + @Airflow.If(id = "has_rows") public boolean hasRows(long transformed) { return transformed > 0; } - @Airflow.Task(id = "report_long", queue = QUEUE) + @Airflow.Task(id = "report_long") public void reportLong() {} - @Airflow.Task(id = "report_short", queue = QUEUE) + @Airflow.Task(id = "report_short") public void reportShort() {} /** Names the one of {@link #reportLong} and {@link #reportShort} that runs; the other is skipped. */ - @Airflow.Switch(id = "pick_report", queue = QUEUE) + @Airflow.Switch(id = "pick_report") public Class pickReport(long transformed) { return transformed > 100 ? AnnotationDagBuilder.ReportLong.class : AnnotationDagBuilder.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. - @Airflow.Task(id = "trigger_downstream", queue = QUEUE) + @Airflow.Task(id = "trigger_downstream") public TriggerDagRun triggerDownstream() { return new TriggerDagRun("java_native_target_e2e"); } @@ -84,7 +85,7 @@ public TriggerDagRun triggerDownstream() { // the task "checks.audit". @Airflow.TaskGroup(id = "checks") static class Checks { - @Airflow.Task(id = "audit", queue = QUEUE) + @Airflow.Task(id = "audit") public void audit() {} } diff --git a/airflow-e2e-tests/java-native-bundle/src/java/org/apache/airflow/e2e/nativedag/DynamicDags.java b/airflow-e2e-tests/java-native-bundle/src/java/org/apache/airflow/e2e/nativedag/DynamicDags.java index 88f4cf7268f76..4014c5337d7fb 100644 --- a/airflow-e2e-tests/java-native-bundle/src/java/org/apache/airflow/e2e/nativedag/DynamicDags.java +++ b/airflow-e2e-tests/java-native-bundle/src/java/org/apache/airflow/e2e/nativedag/DynamicDags.java @@ -65,8 +65,9 @@ public static List load() throws IOException, URISyntaxException { var dag = new DagDef(dagId) .config("catchup", false) - .config("tags", List.of("java-sdk", "java_dynamic")); - dag.task("greet", Greet.class).config("queue", QUEUE); + .config("tags", List.of("java-sdk", "java_dynamic")) + .config("queue", QUEUE); + dag.task("greet", Greet.class); dags.add(dag); } return dags; diff --git a/airflow-e2e-tests/java-native-bundle/src/java/org/apache/airflow/e2e/nativedag/TargetDag.java b/airflow-e2e-tests/java-native-bundle/src/java/org/apache/airflow/e2e/nativedag/TargetDag.java index 2e1be90d106aa..d5ad4e1adbebd 100644 --- a/airflow-e2e-tests/java-native-bundle/src/java/org/apache/airflow/e2e/nativedag/TargetDag.java +++ b/airflow-e2e-tests/java-native-bundle/src/java/org/apache/airflow/e2e/nativedag/TargetDag.java @@ -39,8 +39,9 @@ public static DagDef build() { new DagDef("java_native_target_e2e") .config("description", "Native Java Dag the e2e's other native Dags trigger") .config("catchup", false) - .config("tags", List.of("java-sdk", "native", "e2e")); - dag.task("receive", Receive.class).config("queue", QUEUE); + .config("tags", List.of("java-sdk", "native", "e2e")) + .config("queue", QUEUE); + dag.task("receive", Receive.class); return dag; } } diff --git a/airflow-e2e-tests/tests/airflow_e2e_tests/java_sdk_tests/test_java_sdk_native_dag.py b/airflow-e2e-tests/tests/airflow_e2e_tests/java_sdk_tests/test_java_sdk_native_dag.py index 41219f4017018..54000e47326fe 100644 --- a/airflow-e2e-tests/tests/airflow_e2e_tests/java_sdk_tests/test_java_sdk_native_dag.py +++ b/airflow-e2e-tests/tests/airflow_e2e_tests/java_sdk_tests/test_java_sdk_native_dag.py @@ -24,11 +24,12 @@ No Python file declares these Dags. The ``airflow-e2e-tests/java-native-bundle`` JAR sits in the Dag bundle. There are four ``JavaCoordinator``s, so ``[sdk] dag_bundle_to_coordinator`` picks the -``java-native`` one to parse it: the Dag processor runs the JAR through it. Each task sets -``queue="java-native"``, which routes to ``java-jdk``. That coordinator runs the same JAR from the Dag's -own bundle, whatever its ``task_handler_bundle_name`` says. ``java_native_e2e`` is declared with the interface -API, ``java_native_annotation_e2e`` with annotations. Both also exercise a task group, an ``If``, a -``Switch`` and a ``TriggerDagRun`` of ``java_native_target_e2e``, the third Dag the JAR declares. +``java-native`` one to parse it: the Dag processor runs the JAR through it. Each Dag sets +``queue="java-native"`` once and every task inherits it, which routes to ``java-jdk``. That +coordinator runs the same JAR from the Dag's own bundle, whatever its ``task_handler_bundle_name`` +says. ``java_native_e2e`` is declared with the interface API, ``java_native_annotation_e2e`` with +annotations. Both also exercise a task group, an ``If``, a ``Switch`` and a ``TriggerDagRun`` of +``java_native_target_e2e``, the third Dag the JAR declares. """ from __future__ import annotations @@ -43,7 +44,7 @@ # The Dag processor starts a JVM to parse the JAR, and each task starts another. _JAVA_TASK_TIMEOUT = 600 -# The queue the Dag's tasks set, which ``queue_to_coordinator`` routes to ``java-jdk``. +# The queue the Dag sets for all its tasks, which ``queue_to_coordinator`` routes to ``java-jdk``. _QUEUE = "java-native" # The Dag ``trigger_downstream`` starts a run of; not triggered through the client itself. @@ -149,8 +150,8 @@ def test_the_graph_is_the_one_java_declared(parsed_dags: AirflowClient, native_d @_by_dag_id -def test_every_task_sets_the_native_queue(parsed_dags: AirflowClient, native_dag: _NativeDag): - """The Java DSL has no Dag-level queue, so each task sets its own.""" +def test_every_task_takes_the_dag_queue(parsed_dags: AirflowClient, native_dag: _NativeDag): + """The Dag sets ``queue`` once, and every task runs on it.""" tasks = parsed_dags.get_tasks(native_dag.dag_id).get("tasks", []) assert {task["task_id"]: task["queue"] for task in tasks} == dict.fromkeys(native_dag.downstream, _QUEUE) 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 6e751dd724a7a..e0cfb6458332f 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 @@ -31,6 +31,7 @@ id = "java_native_annotation_example", description = "Pure-Java Dag authored with annotations", schedule = "@daily", + queue = "java", startDate = "2026-01-01T00:00:00Z", catchup = false, tags = {"example", "java-sdk"}) 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 47a8f6786af12..a90378ad35ff0 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 @@ -101,6 +101,7 @@ public static DagDef build() { new DagDef("java_native_interface_example") .config("description", "Pure-Java Dag authored with the interface API") .config("schedule", "@daily") + .config("queue", "java") .config("catchup", false) .config("tags", List.of("example", "java-sdk")); 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 index cf8d451429545..f428650dc51b8 100644 --- 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 @@ -41,6 +41,7 @@ public static DagDef build() { var dag = new DagDef("java_native_target_example") .config("description", "Pure-Java Dag that the other native examples trigger") + .config("queue", "java") .config("catchup", false) .config("tags", List.of("example", "java-sdk")); dag.task("receive", Receive.class); diff --git a/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/AirflowProcessor.kt b/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/AirflowProcessor.kt index 36d2f1cfdd69c..53cfab63946df 100644 --- a/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/AirflowProcessor.kt +++ b/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/AirflowProcessor.kt @@ -712,7 +712,11 @@ class AirflowProcessor : AbstractProcessor() { value: AnnotationValue, ): CodeBlock = when (field.type) { - FieldType.STRING -> CodeBlock.of($$"$S", value.value) + FieldType.STRING -> { + val text = value.value as String + field.checkString(text) { "Annotation attribute '${field.attribute}'" } + CodeBlock.of($$"$S", text) + } FieldType.BOOLEAN, FieldType.INTEGER, FieldType.NUMBER -> CodeBlock.of($$"$L", value.value) FieldType.STRING_ARRAY -> { @Suppress("UNCHECKED_CAST") diff --git a/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/AirflowProcessorTest.kt b/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/AirflowProcessorTest.kt index 11be0f957f1b3..e41f6b82bbd44 100644 --- a/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/AirflowProcessorTest.kt +++ b/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/AirflowProcessorTest.kt @@ -403,7 +403,7 @@ class AirflowProcessorTest { """ package org.apache.airflow.example; import org.apache.airflow.sdk.Airflow; - @Airflow.Dag(id = "cfg", schedule = "@daily", tags = {"a", "b"}, catchup = true, + @Airflow.Dag(id = "cfg", schedule = "@daily", queue = "java", tags = {"a", "b"}, catchup = true, startDate = "2026-01-01T00:00:00Z") public class TestExample { @Airflow.Task(retries = 2, queue = "q", retryDelay = "PT5M", retryExponentialBackoff = 1.5) @@ -442,6 +442,7 @@ class AirflowProcessorTest { public static DagDef build() { var dag = DagSource.declaredBy(new DagDef("cfg"), TestExample.class); dag.config("schedule", "@daily"); + dag.config("queue", "java"); dag.config("tags", List.of("a", "b")); dag.config("catchup", true); dag.config("start_date", OffsetDateTime.parse("2026-01-01T00:00:00Z")); @@ -2074,6 +2075,56 @@ class AirflowProcessorTest { ) } + @Test + @DisplayName("reject a blank Dag-level queue") + fun rejectBlankDagQueue() { + for (queue in listOf("", " ")) { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Airflow; + @Airflow.Dag(queue = "$queue") + public class TestExample { + @Airflow.Task public void t() {} + + @Airflow.Deps + static class Wiring implements TestExampleDeps { + void depends() { t(); } + } + } + """, + ) + assertThat(compilation).failed() + assertThat(compilation).hadErrorContaining("Annotation attribute 'queue' must not be blank") + } + } + + @Test + @DisplayName("reject a blank task-level queue") + fun rejectBlankTaskQueue() { + for (queue in listOf("", " ")) { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Airflow; + @Airflow.Dag + public class TestExample { + @Airflow.Task(queue = "$queue") public void t() {} + + @Airflow.Deps + static class Wiring implements TestExampleDeps { + void depends() { t(); } + } + } + """, + ) + assertThat(compilation).failed() + assertThat(compilation).hadErrorContaining("Annotation attribute 'queue' must not be blank") + } + } + @Test @DisplayName("escape quotes and backslashes in a string-array attribute") fun generateBuilderEscapesStringArrayValues() { diff --git a/java-sdk/sdk/build.gradle.kts b/java-sdk/sdk/build.gradle.kts index 2ba31187168e6..825f481b950ff 100644 --- a/java-sdk/sdk/build.gradle.kts +++ b/java-sdk/sdk/build.gradle.kts @@ -284,8 +284,15 @@ abstract class GenerateDagDslTask : DefaultTask() { val attrDefault: String, val defaultJson: String?, val doc: String, + val notBlank: Boolean = false, ) + // Schema keys whose value cannot be blank, in every table that holds them. + private val notBlankKeys = setOf("queue") + + private fun constrained(fields: List): List = + fields.map { if (it.key in notBlankKeys) it.copy(notBlank = true) else it } + private fun camelCase(key: String): String = key .split('_') @@ -461,6 +468,22 @@ abstract class GenerateDagDslTask : DefaultTask() { "`\"@once\"`, `\"@continuous\"`, a cron expression, or empty for no schedule.", ), ) + // The schema has no Dag-level queue; a Python Dag sets one through default_args instead. + // Every task of a Java Dag runs on the Java coordinator, so one queue on the Dag routes all of them there. + if (!dagProps.path("queue").isMissingNode) { + throw GradleException("The schema now has a Dag-level 'queue'; resolve it from the schema instead") + } + add( + DslField( + "queue", + "queue", + "STRING", + "String", + quote(""), + null, + "Queue each task of the Dag runs on, unless the task sets its own `queue`. It cannot be blank.", + ), + ) dagAllowlist.forEach { key -> val prop = dagProps.path(key) if (prop.isMissingNode) { @@ -471,7 +494,7 @@ abstract class GenerateDagDslTask : DefaultTask() { ?: throw GradleException("Dag allowlist key '$key' is not a scalar the DSL can express"), ) } - } + }.let(::constrained) val excludedSeen = mutableSetOf() val taskFields = @@ -490,7 +513,7 @@ abstract class GenerateDagDslTask : DefaultTask() { val override = if (key == "retry_exponential_backoff") "Double" else null resolveField(key, prop, override)?.let { add(it) } } - } + }.let(::constrained) (excludedTaskKeys - excludedSeen).takeIf { it.isNotEmpty() }?.let { throw GradleException("Excluded task keys match no eligible schema property; remove or fix: $it") } @@ -756,6 +779,7 @@ abstract class GenerateDagDslTask : DefaultTask() { " \"${f.attribute}\",", " FieldType.${f.fieldType},", " $defaultRepr,", + " ${f.notBlank},", " ),", ).joinToString("\n") } 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 913817d090920..5964b4552b177 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 @@ -87,13 +87,16 @@ class DagDef( * mismatched value types are rejected on the call, so mistakes surface where * the Dag is defined. * + * `"queue"` is the one key Airflow's Dag has no setting for. + * It is the queue every task of the Dag runs on, unless the task sets its own `"queue"`, and it cannot be blank. + * * @param key Airflow Dag setting name. * @param value Value matching the key's schema type. Durations take * [java.time.Duration], date-times [java.time.OffsetDateTime] or * [java.time.Instant], string arrays any `Iterable` of `String`. * @return This Dag, for chaining. - * @throws IllegalArgumentException if the key is unknown or the value type - * does not match. + * @throws IllegalArgumentException if the key is unknown, the value type + * does not match, or `"queue"` is blank. */ fun config( key: String, @@ -542,17 +545,17 @@ class TaskDef( * Sets one task-level configuration value. * * Keys are Airflow's own task setting names (for example `"retries"`, - * `"queue"`, `"retry_delay"`); unknown keys and mismatched - * value types are rejected on the call, so mistakes surface where the task is - * defined. + * `"queue"`, `"retry_delay"`); unknown keys, mismatched + * value types and a blank `"queue"` are rejected on the call, so mistakes + * surface where the task is defined. * * @param key Airflow task setting name, e.g. `"retries"`. * @param value Value matching the key's schema type. Durations take * [java.time.Duration], date-times [java.time.OffsetDateTime] or * [java.time.Instant]. * @return This task definition, for chaining. - * @throws IllegalArgumentException if the key is unknown or the value type - * does not match. + * @throws IllegalArgumentException if the key is unknown, the value type + * does not match, or `"queue"` is blank. */ fun config( key: String, 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 index 15a6a5b3d0da2..c1cff8e9a8e2f 100644 --- 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 @@ -66,6 +66,7 @@ private val TRIGGER_FIELDS: Map = * 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. + * Like a Java task, it runs on the Dag's `queue` unless the task sets its own. * * @param dagId `trigger_dag_id`: the Dag to trigger. * @throws IllegalArgumentException if [dagId] is empty. 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 0482fc96ad808..c01592b8eea6e 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 @@ -102,6 +102,7 @@ internal fun serializeDag( } } + val dagQueue = dag.dagConfig["queue"] as String? val result = linkedMapOf( "dag_id" to dag.id, @@ -109,7 +110,7 @@ internal fun serializeDag( "relative_fileloc" to relativeFileloc, "timezone" to dagTimezone(dag.dagConfig), "timetable" to serializeTimetable(dag.id, dag.dagConfig), - "tasks" to dag.tasks.map { (taskId, def) -> serializeTask(taskId, def, downstream[taskId]) }, + "tasks" to dag.tasks.map { (taskId, def) -> serializeTask(taskId, def, downstream[taskId], dagQueue) }, "dag_dependencies" to serializeDagDependencies(dag), "task_group" to serializeTaskGroups(dag, expansion), "edge_info" to serializeEdgeInfo(dag, expansion), @@ -124,11 +125,13 @@ internal fun serializeDag( /** * Converts one task to the Airflow serialization format. `downstream` is the * inverted view of the Dag's upstream edges, sorted for stable JSON. + * `dagQueue` is the Dag's queue, which the task takes unless it sets its own. */ private fun serializeTask( taskId: String, def: TaskDef, downstream: List?, + dagQueue: String?, ): Map { val data = linkedMapOf("task_id" to taskId) val trigger = def.trigger @@ -150,13 +153,16 @@ private fun serializeTask( } // Lets NotPreviouslySkippedDep re-skip a cleared downstream, as Python's SkipMixin does. if (def.decider != null) data["_can_skip_downstream"] = true + // A trigger task takes the Dag's queue too, because the Java runtime runs it. + val config = + if (dagQueue == null || "queue" in def.configValues) def.configValues else def.configValues + ("queue" to dagQueue) // 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 // __type encoding is stripped. If core grows a task-level fill_config_defaults, // every SDK has to keep explicitly set values instead, or an explicit retries=0 // reads as unset and picks up the configured default. - def.configValues.forEach { (key, value) -> + config.forEach { (key, value) -> if (key !in OMITTED_TASK_KEYS && !matchesSchemaDefault(SchemaFields.TASK[key], value)) { data[key] = unwrapTypeEncoding(serializeValue(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 0afea428eedb0..bd03688517109 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 @@ -57,13 +57,28 @@ private val DAG_RUN_STATE_NAMES = listOf("queued", "running", "success", "failed * @property type Accepted value shape. * @property defaultJson Schema default as raw JSON, or `null` when the schema * declares no default. + * @property notBlank Whether a value for this key must not be blank. */ class Field( val key: String, val attribute: String, val type: FieldType, val defaultJson: String?, -) + val notBlank: Boolean = false, +) { + /** + * Enforces the constraints this key carries on a string value. [describe] + * names the value, so each caller keeps its own wording. + * + * @throws IllegalArgumentException if [text] breaks a constraint. + */ + fun checkString( + text: String, + describe: () -> String, + ) { + require(!notBlank || text.isNotBlank()) { "${describe()} must not be blank" } + } +} /** * Validates one `config(key, value)` call against a schema field table and @@ -91,7 +106,10 @@ internal fun checkConfigValue( "Value for $scope config key '$key' must be $expected, got: ${value.javaClass.name}", ) return when (field.type) { - FieldType.STRING -> value as? String ?: mismatch("a String") + FieldType.STRING -> + (value as? String ?: mismatch("a String")).also { text -> + field.checkString(text) { "Value for $scope config key '$key'" } + } FieldType.BOOLEAN -> value as? Boolean ?: mismatch("a Boolean") FieldType.NUMBER -> value as? Number ?: mismatch("a Number") FieldType.INTEGER -> diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/DagDefTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/DagDefTest.kt index 70bb24552e09d..23f35e6668e3e 100644 --- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/DagDefTest.kt +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/DagDefTest.kt @@ -175,6 +175,7 @@ internal class DagDefTest { .config("dagrun_timeout", Duration.ofMinutes(5)) .config("start_date", OffsetDateTime.parse("2026-01-01T00:00:00Z")) .config("tags", listOf("a", "b")) + .config("queue", "java") Assertions.assertEquals( mapOf( @@ -185,6 +186,7 @@ internal class DagDefTest { "dagrun_timeout" to Duration.ofMinutes(5), "start_date" to OffsetDateTime.parse("2026-01-01T00:00:00Z"), "tags" to listOf("a", "b"), + "queue" to "java", ), dag.dagConfig, ) @@ -226,6 +228,32 @@ internal class DagDefTest { Assertions.assertEquals("Value for Dag config key 'description' must not be null", error.message) } + @Test + @DisplayName("Should reject a blank dag queue") + fun shouldRejectBlankDagQueue() { + for (queue in listOf("", " ")) { + val error = + Assertions.assertThrows(IllegalArgumentException::class.java) { + DagDef("dag").config("queue", queue) + } + + Assertions.assertEquals("Value for Dag config key 'queue' must not be blank", error.message) + } + } + + @Test + @DisplayName("Should reject a blank task queue") + fun shouldRejectBlankTaskQueue() { + for (queue in listOf("", " ")) { + val error = + Assertions.assertThrows(IllegalArgumentException::class.java) { + DagDef("dag").task("extract", NoOp::class.java).config("queue", queue) + } + + Assertions.assertEquals("Value for task config key 'queue' must not be blank", error.message) + } + } + @Test @DisplayName("Should reject non-integral values for integer dag config keys") fun shouldRejectFractionalIntegerValue() { diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/SerdeTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/SerdeTest.kt index e025d9a5b154f..d8f71edb92aa7 100644 --- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/SerdeTest.kt +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/execution/SerdeTest.kt @@ -22,10 +22,13 @@ package org.apache.airflow.sdk.execution import org.apache.airflow.sdk.Arg 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.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.DagFileParseRequest import org.apache.airflow.sdk.internal.Refs import org.junit.jupiter.api.Assertions.assertEquals @@ -44,6 +47,20 @@ private class SerdeNoopTask : Task { ) = Unit } +private class SerdeCondition : ConditionTask { + override fun decide( + context: Context, + client: Client, + ) = true +} + +private class SerdeSwitch : SwitchTask { + override fun choose( + context: Context, + client: Client, + ) = SerdeNoopTask::class.java +} + @Suppress("UNCHECKED_CAST") private fun taskData( serialized: Map, @@ -54,6 +71,12 @@ private fun taskData( return tasks[index]["__var"] as Map } +@Suppress("UNCHECKED_CAST") +private fun queuesByTaskId(serialized: Map): Map = + (serialized["tasks"] as List>) + .map { it["__var"] as Map } + .associate { it["task_id"] as String to it["queue"] } + internal class SerdeTest { @Test @DisplayName("Should emit required dag fields and leave unset config-backed fields out") @@ -186,6 +209,47 @@ internal class SerdeTest { ) } + @Test + @DisplayName("Should give every task the Dag's queue unless the task sets its own") + fun shouldGiveEachTaskTheDagQueue() { + val dag = DagDef("d").config("queue", "java") + val extract = dag.task("extract", SerdeNoopTask::class.java) + dag.task("heavy", SerdeNoopTask::class.java).config("queue", "java_large") + dag.task("on_default", SerdeNoopTask::class.java).config("queue", "default") + dag.If("has_rows", SerdeCondition::class.java).after(extract) + dag.Switch("pick", SerdeSwitch::class.java).after(extract) + dag.task("trigger", TriggerDagRun("reports")) + dag.task("trigger_on_python", TriggerDagRun("reports")).config("queue", "python") + + val serialized = serializeDag(dag, "", ".") + + assertEquals( + mapOf( + "extract" to "java", + "heavy" to "java_large", + "on_default" to null, + "has_rows" to "java", + "pick" to "java", + "trigger" to "java", + "trigger_on_python" to "python", + ), + queuesByTaskId(serialized), + ) + assertFalse("queue" in serialized) + } + + @Test + @DisplayName("Should leave the queue out when the Dag's queue is the schema default or unset") + fun shouldLeaveOutDefaultOrUnsetDagQueue() { + val onDefault = DagDef("d").config("queue", "default") + onDefault.task("t", SerdeNoopTask::class.java) + val unset = DagDef("d") + unset.task("t", SerdeNoopTask::class.java) + + assertFalse("queue" in taskData(serializeDag(onDefault, "", "."), 0)) + assertFalse("queue" in taskData(serializeDag(unset, "", "."), 0)) + } + @Test @DisplayName("Should serialize nested task groups with their own edges") fun shouldSerializeTaskGroups() {