From 0ae164080c0d213dff48cd164a81b3a22dc854d0 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 23 Sep 2026 06:58:24 +0000 Subject: [PATCH 01/13] Java SDK: Declare a Dag's task graph with @Builder.Deps An annotated Dag class could name its tasks but not say how they connect, so a Dag whose annotations carried everything else still needed a Python Dag file for its graph. Declaring the graph by calling the tasks lets javac check that each task is fed something its parameter accepts. Bundle.register(Class) now also takes the @Builder.Dag class itself and registers the Dag its generated builder builds, so there is still only one name to keep in sync. --- .../language-sdks/java.rst | 76 ++++ java-sdk/README.md | 6 +- .../airflow/example/ExampleBundleBuilder.java | 3 + .../example/nativedag/AnnotationExample.java | 73 ++++ .../apache/airflow/sdk/BuilderProcessor.kt | 258 ++++++++++--- .../org/apache/airflow/sdk/BuilderTest.kt | 356 +++++++++++++++++- java-sdk/sdk/build.gradle.kts | 39 +- .../main/kotlin/org/apache/airflow/sdk/Arg.kt | 22 +- .../kotlin/org/apache/airflow/sdk/Bundle.kt | 49 ++- .../kotlin/org/apache/airflow/sdk/Deps.kt | 17 +- .../org/apache/airflow/sdk/internal/Refs.kt | 114 ++++++ .../org/apache/airflow/sdk/ArgTestSupport.kt | 5 +- .../org/apache/airflow/sdk/BundleTest.kt | 61 ++- .../org/apache/airflow/sdk/InputTaskTest.kt | 6 +- .../airflow/sdk/internal/ArgValuesTest.kt | 5 +- .../apache/airflow/sdk/internal/RefsTest.kt | 115 ++++++ 16 files changed, 1090 insertions(+), 115 deletions(-) create mode 100644 java-sdk/example/src/java/org/apache/airflow/example/nativedag/AnnotationExample.java create mode 100644 java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/Refs.kt create mode 100644 java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/RefsTest.kt 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 3f69772e4d92d..e5edd6b8b9a0c 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst @@ -314,6 +314,9 @@ Annotate a plain Java class and let the SDK generate the boilerplate at compile - 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 task settings; only attributes written explicitly are applied. + * - ``@Builder.Deps`` + - Marks the nested class that declares the task graph in Java, TaskFlow-style. Required for + a Dag that Java owns end to end. See :ref:`java-sdk/native-dags`. * - ``TaskInput`` / ``@ArgName("...")`` - Marks a class as a task's input, so keyword arguments bind by name instead of by position: each public field receives the argument whose name matches it, ignoring case and @@ -574,6 +577,79 @@ Edges are checked when the Dag is registered with a ``Bundle``: an upstream that Dag, or to no Dag, and a cycle anywhere in the graph both fail there rather than at the first task run. +Wiring the graph with ``@Builder.Deps`` +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +For a Dag written with annotations, the graph is declared by a nested ``@Builder.Deps`` class. The +annotation processor generates a ``Deps`` interface — the *wiring view* — with one method +per ``@Builder.Task`` method: the injected ``Client`` and ``Context`` parameters are dropped, each +data parameter becomes an ``Arg``, and the return value becomes a ``TaskRef``. Calling a view +method registers its task, and passing the handle one returned into another feeds the upstream's +output into the downstream's parameter *and* wires the data edge. The call graph is the task graph, +and ``javac`` type-checks it. + +Declare the wiring class as a ``static`` nested class of the Dag class that ``implements`` the +generated view, with a no-argument ``depends()`` method: + +.. code-block:: java + + @Builder.Dag( + id = "java_etl", + schedule = "@daily", + description = "Pure-Java Dag built with annotations", + tags = {"example", "java-sdk"}) + public class EtlPipeline { + + @Builder.Task(id = "extract", retries = 2) + public long extract() { + return 42L; + } + + @Builder.Task(id = "transform") + public long transform(long extracted, double factor) { + return (long) (extracted * factor); + } + + @Builder.Task(id = "load") + public void load(long transformed) { + // implement task logic + } + + @Builder.Task(id = "audit") + public void audit() { + // side effect only, no data in or out + } + + @Builder.Deps + static class Wiring implements EtlPipelineDeps { + void depends() { + var rows = extract(); + load(transform(rows, lit(0.9))); + rows.before(audit()); // ordering-only edge: audit waits for extract + } + } + } + +Every ``@Builder.Task`` method must be called in the wiring class; a task the wiring missed fails at +Dag-parse time. ``lit(...)`` wires an inline constant where no upstream feeds a parameter — a bare +``double`` cannot be an ``Arg``, so a constant is wrapped. A view method that takes no arguments +returns the same handle every time, so it names one node wherever it appears; one that takes +arguments is called once, and the wiring fails if it is called again with arguments, so hold its +handle in a local and reuse that. + +``before``, ``after`` and ``Flow.of`` work here exactly as they do on the interface surface; inside +the wiring class ``Flow`` is inherited by simple name, so it needs no import and never collides with +``java.util.concurrent.Flow``. + +The wiring class is what makes the Dag native. A Dag class without one registers every task with +no Java-side edges, which is the task-handler shape rather than a Dag Java owns. + +.. note:: + + Runtime argument bindings win over Java-declared wiring. When the supervisor delivers bindings + for a run (see :ref:`java-sdk/arg-binding`), the binding at a parameter's position is what the + task receives. Wired inputs are the fallback, which is what a native Java Dag always uses. + Configuration attributes ~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/java-sdk/README.md b/java-sdk/README.md index b8600ab5d44fa..7fd9b6871d766 100644 --- a/java-sdk/README.md +++ b/java-sdk/README.md @@ -729,9 +729,9 @@ E2E_TEST_MODE=java_sdk uv run --project airflow-e2e-tests pytest \ - The annotation processor (`BuilderProcessor.kt`) uses `kapt`. The `Builder` class holding the `@Builder.Dag` / `@Builder.Task` annotations is generated from the Dag serialization schema by `:sdk:generateDagDsl` (vendored at - `sdk/schema/dag-schema.json`). The `Arg`/`TaskRef` and `Deps`/`Flow` graph - types are hand-written next to the rest of the public surface in - `sdk/src/main/kotlin/org/apache/airflow/sdk/`. When adding annotation + `sdk/schema/dag-schema.json`), `@Builder.Deps` included. The `Arg`/`TaskRef` + and `Deps`/`Flow` graph types are hand-written next to the rest of the public + surface in `sdk/src/main/kotlin/org/apache/airflow/sdk/`. When adding annotation behaviour, handle it in `BuilderProcessor.kt` and add a golden-output test in `processor/src/test/kotlin/`. - The Python coordinator subclasses `SubprocessCoordinator`. Do not reach into 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 d75f5762ba12d..76175db8a6d5c 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 @@ -21,12 +21,15 @@ import org.apache.airflow.sdk.*; +// One bundle serves every surface: Dags built in Java, and the handler classes +// whose Dags the Python file owns. public class ExampleBundleBuilder { public static Bundle build() { return new Bundle() .register(InterfaceExampleBuilder.build()) .register(AnnotationExample.class) .register(XComCastingExample.class) + .register(org.apache.airflow.example.nativedag.AnnotationExample.class) .register(org.apache.airflow.example.nativedag.InterfaceExample.build()); } 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 new file mode 100644 index 0000000000000..c2a3bc89f75cd --- /dev/null +++ b/java-sdk/example/src/java/org/apache/airflow/example/nativedag/AnnotationExample.java @@ -0,0 +1,73 @@ +/* + * 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 org.apache.airflow.sdk.*; + +// A Dag defined entirely in Java, annotation-style. The @Builder.Dag and +// @Builder.Task attributes carry the configuration and the @Builder.Deps class +// declares the graph -- the single place dependencies are defined. +@Builder.Dag( + id = "java_native_annotation_example", + description = "Pure-Java Dag authored with annotations", + schedule = "@daily", + startDate = "2026-01-01T00:00:00Z", + catchup = false, + tags = {"example", "java-sdk"}) +public class AnnotationExample { + private static final System.Logger log = System.getLogger(AnnotationExample.class.getName()); + + @Builder.Task(id = "extract", retries = 2) + public long extract() { + log.log(INFO, "Extracting a value"); + return 42L; + } + + @Builder.Task(id = "transform") + public long transform(long extracted, double factor) { + log.log(INFO, "Transforming {0} by {1}", extracted, factor); + return (long) (extracted * factor); + } + + @Builder.Task(id = "load") + public void load(long transformed) { + log.log(INFO, "Loaded {0}", transformed); + } + + @Builder.Task(id = "audit") + public void audit() { + log.log(INFO, "Audited the run"); + } + + // Implements the generated wiring view, so javac type-checks the graph: + // extract() yields a TaskRef, which is an Arg for transform. + @Builder.Deps + static class Wiring implements AnnotationExampleDeps { + void depends() { + var extracted = extract(); + load(transform(extracted, lit(1.5))); + // Ordering-only edge: audit runs after extract, with no data flowing. + extracted.before(audit()); + } + } +} 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 21b4165728e3b..f97f72997d0a0 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 @@ -28,9 +28,11 @@ import com.squareup.javapoet.MethodSpec import com.squareup.javapoet.ParameterizedTypeName 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.Field import org.apache.airflow.sdk.internal.FieldType +import org.apache.airflow.sdk.internal.Refs import org.apache.airflow.sdk.internal.SchemaFields import org.apache.airflow.sdk.internal.TaskArgs import org.apache.airflow.sdk.internal.TypeRef @@ -65,19 +67,27 @@ import javax.tools.Diagnostic * `META-INF/services/javax.annotation.processing.Processor`; not intended to be * instantiated or referenced directly. * - * For each class annotated with [Builder.Dag], generates a `*Builder` class - * containing: + * For each class annotated with [Builder.Dag], generates: * - * - One inner class per [Builder.Task]-annotated method, implementing [Task]. - * - A static `build()` method that constructs the [DagDef], lowers every - * explicitly-written `@Builder.Dag` attribute into a `DagDef.config` call, - * and registers those inner classes as [TaskDef]s, each carrying its - * explicitly-written `@Builder.Task` attributes the same way. + * - A `*Builder` class containing one inner class per [Builder.Task]-annotated + * method (implementing [Task]), and a static `build()` that constructs the + * [DagDef], lowers every explicitly-written `@Builder.Dag` attribute into a + * `DagDef.config` call, then runs the class's [Builder.Deps] class and + * verifies it registered every task — or, when the class declares none, + * registers every task with no Java-side edges. + * - A `*Deps` wiring-view interface (only when a [Builder.Deps] class exists) + * whose methods mirror the task methods: injectable parameters ([Client], + * [Context]) are dropped, data parameters become [Arg]-typed inputs, and the + * return value becomes a [TaskRef]. Calling one registers the task with its + * explicitly-written `@Builder.Task` attributes lowered into `TaskDef.config` + * calls; passing one call's handle to another wires the dependency edge and + * feeds the upstream's return-value XCom into the downstream's parameter, + * type-checked by javac through the [Arg] / [TaskRef] generics. * - * In the generated `execute` body, a task's data parameters resolve against the - * arg bindings the supervisor delivered for the run: flat parameters through - * [TaskArgs], by their position among the data parameters, and [TaskInput] - * fields through [ArgValues], by argument name. Non-`void` return values are + * In the generated `execute` bodies, a task's data parameters resolve against + * the arg bindings the supervisor delivered for the run: flat parameters + * through [TaskArgs], by their position among the data parameters, and + * [TaskInput] fields through [ArgValues], by argument name. Non-`void` return values are * forwarded to `client.setXCom`. */ @SupportedAnnotationTypes( @@ -108,12 +118,18 @@ class BuilderProcessor : AbstractProcessor() { roundEnv.getElementsAnnotatedWith(Builder.Dag::class.java).filterIsInstance().forEach { el -> with(processingEnv) { runCatching { + val packageName = elementUtils.getPackageOf(el).qualifiedName.toString() + val declarations = collectTasks(el) + val deps = findDeps(el) + val builderName = ClassName.get(packageName, dagAnnotation(el).to.ifBlank { "${el.simpleName}Builder" }) + val depsName = ClassName.get(packageName, "${el.simpleName}Deps") JavaFile - .builder( - elementUtils.getPackageOf(el).qualifiedName.toString(), - buildDag(el), - ).build() + .builder(packageName, buildBuilder(el, declarations, deps, builderName)) + .build() .writeTo(filer) + if (deps != null) { + JavaFile.builder(packageName, buildDeps(el, declarations, builderName, depsName)).build().writeTo(filer) + } }.onFailure { e -> messager.printMessage( Diagnostic.Kind.ERROR, @@ -161,24 +177,31 @@ class BuilderProcessor : AbstractProcessor() { require(handler.dag.isNotBlank()) { "@Builder.TaskHandler on '${inner.simpleName}' must name the Dag the Python file declares" } - val innerName = inner.simpleName.toString().replaceFirstChar(Char::uppercase) - registrar.addType(buildTask(innerName, inner, el)) + val decl = TaskDeclaration(inner, handler.task.ifBlank { inner.simpleName.toString() }, collectDataParams(inner)) + registrar.addType(buildTask(decl, el)) registerInto.addStatement( $$"bundle.register($S, $S, $L.class)", handler.dag, - handler.task.ifBlank { inner.simpleName }, - innerName, + decl.id, + decl.className, ) } return registrar.addMethod(registerInto.build()).build() } - private fun buildDag(el: TypeElement): TypeSpec { - val ann = el.getAnnotation(Builder.Dag::class.java)!! + private fun dagAnnotation(el: TypeElement): Builder.Dag = el.getAnnotation(Builder.Dag::class.java)!! + + private fun buildBuilder( + el: TypeElement, + declarations: List, + deps: TypeElement?, + builderName: ClassName, + ): TypeSpec { + val ann = dagAnnotation(el) val builderClass = TypeSpec - .classBuilder(ann.to.ifBlank { "${el.simpleName}Builder" }) + .classBuilder(builderName) .addModifiers(Modifier.PUBLIC, Modifier.FINAL) val buildMethod = @@ -190,46 +213,160 @@ class BuilderProcessor : AbstractProcessor() { explicitConfig(el, DAG_ANNOTATION, DAG_STRUCTURAL_ATTRIBUTES, SchemaFields.DAG).forEach { (key, value) -> buildMethod.addStatement($$"dag.config($S, $L)", key, value) } - - for (inner in el.enclosedElements) { - if (inner !is ExecutableElement) continue - if (inner.isVarArgs) throw IllegalArgumentException("Cannot create task from vararg function ${inner.simpleName}") - - val taskAnn = inner.getAnnotation(Builder.Task::class.java) ?: continue - val innerName = inner.simpleName.toString().replaceFirstChar(Char::uppercase) - - builderClass.addType(buildTask(innerName, inner, el)) - + if (deps != null) { buildMethod.addStatement( - $$"dag.addTask($L)", - taskDefCode(inner, taskAnn.id.ifBlank { inner.simpleName.toString() }, innerName), + $$"return $T.record(dag, $T.of($L), new $T()::depends)", + REFS_TYPE, + ClassName.get(List::class.java), + declarations.joinToString { "\"${it.id}\"" }, + ClassName.get(deps), ) + } else { + // No wiring class: register every task with no Java-side edges, which + // is the task-handler shape rather than a Dag Java owns. + declarations.forEach { decl -> + buildMethod.addStatement($$"dag.addTask($L)", taskDefCode(decl, CodeBlock.of($$"$L", decl.className))) + } + buildMethod.addStatement("return dag") } - - buildMethod.addStatement("return dag") builderClass.addMethod(buildMethod.build()) + + declarations.forEach { builderClass.addType(buildTask(it, el)) } return builderClass.build() } /** - * Emits `new TaskDef(id, .class)` with the explicitly-written + * Generates the Dag's wiring view: one default method per task, with the + * injected arguments stripped, each data argument lifted to [Arg] and the + * return lifted to [TaskRef]. + * + * It is an interface so the `@Builder.Deps` class can *implement* it and + * keep its own `extends` free, and so the Dag class's real task methods -- + * which differ only in their injected arguments -- do not clash with it. + */ + private fun buildDeps( + el: TypeElement, + declarations: List, + builderName: ClassName, + depsName: ClassName, + ): TypeSpec { + val view = + TypeSpec + .interfaceBuilder(depsName) + .addModifiers(Modifier.PUBLIC) + .addSuperinterface(DEPS_TYPE) + .addJavadoc( + "Wiring view of {@link \$T}'s task methods, for declaring its task graph.\n\n" + + "

Calling one registers its task with the Dag being built; passing the handle it\n" + + "returned into another call feeds the upstream's output into that task's parameter\n" + + "and wires the data edge. {@code before} and {@code after} wire an ordering-only edge.\n", + ClassName.get(el), + ) + + for (decl in declarations) { + 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())) + decl.dataParams.forEach { method.addParameter(inType(it.type), it.name) } + val def = taskDefCode(decl, CodeBlock.of($$"$T.$L", builderName, decl.className)) + if (decl.dataParams.isEmpty()) { + method.addStatement($$"return $T.node($L)", REFS_TYPE, def) + } else { + method.addStatement( + $$"return $T.call($L, $L)", + REFS_TYPE, + def, + decl.dataParams.joinToString { it.name }, + ) + } + view.addMethod(method.build()) + } + return view.build() + } + + /** + * Emits `new TaskDef(id, .class)` with the explicitly-written * `@Builder.Task` attributes lowered into chained `.config` calls. */ private fun taskDefCode( - method: ExecutableElement, - id: String, - className: String, + decl: TaskDeclaration, + classRef: CodeBlock, ): CodeBlock { val taskDef = CodeBlock .builder() - .add($$"new $T($S, $L.class)", TASK_DEF_TYPE, id, className) - explicitConfig(method, TASK_ANNOTATION, TASK_STRUCTURAL_ATTRIBUTES, SchemaFields.TASK).forEach { (key, value) -> + .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) -> taskDef.add($$".config($S, $L)", key, value) } return taskDef.build() } + /** + * Maps a data parameter's declared type to its twin-input type. Numeric + * parameters accept any numeric upstream (`Arg`, widened + * at run time); `Object`, raw `Map`, and raw `List` parameters accept any + * upstream (`Arg`, decoded loosely at run time); everything else accepts + * covariant matches of the declared type (`Arg`). + */ + private fun inType(paramType: TypeMirror): TypeName { + val boxed = TypeName.get(paramType).boxIfPossible() + val argument = + when { + isNumeric(paramType) -> WildcardTypeName.subtypeOf(TypeName.get(Number::class.java)) + else -> WildcardTypeName.subtypeOf(boxed) + } + return ParameterizedTypeName.get(ARG_TYPE, argument) + } + + private fun isNumeric(t: TypeMirror): Boolean = t.kind in NUMERIC_KINDS || TypeName.get(t) in BOXED_NUMERICS + + private fun collectTasks(el: TypeElement): List { + val declarations = mutableListOf() + for (inner in el.enclosedElements) { + if (inner !is ExecutableElement) continue + val ann = inner.getAnnotation(Builder.Task::class.java) ?: continue + if (inner.isVarArgs) throw IllegalArgumentException("Cannot create task from vararg function ${inner.simpleName}") + val id = ann.id.ifBlank { inner.simpleName.toString() } + require(declarations.none { it.id == id }) { "Tasks in Dag have duplicate ID: $id" } + declarations += TaskDeclaration(inner, id, collectDataParams(inner)) + } + return declarations + } + + /** + * Finds and validates the class's `@Builder.Deps` wiring class. It is + * optional: without one, every task registers with no Java-side edges. + */ + private fun findDeps(el: TypeElement): TypeElement? { + val classes = + el.enclosedElements + .filterIsInstance() + .filter { it.getAnnotation(Builder.Deps::class.java) != null } + if (classes.isEmpty()) return null + val deps = + classes.singleOrNull() + ?: throw IllegalArgumentException( + "Dag class ${el.simpleName} declares more than one @Builder.Deps class: " + + classes.joinToString { it.simpleName.toString() }, + ) + require(Modifier.STATIC in deps.modifiers && Modifier.PRIVATE !in deps.modifiers) { + "@Builder.Deps class '${deps.simpleName}' must be static and non-private" + } + require(deps.enclosedElements.filterIsInstance().any { it.isNoArgDepends() }) { + "@Builder.Deps class '${deps.simpleName}' must declare a non-private, no-argument depends() method" + } + return deps + } + + private fun ExecutableElement.isNoArgDepends(): Boolean = + simpleName.contentEquals("depends") && + parameters.isEmpty() && + Modifier.PRIVATE !in modifiers && + Modifier.STATIC !in modifiers + /** * Lowers the explicitly-written configuration attributes of [element]'s * [annotationName] annotation into (schema key, value code) pairs. Only @@ -299,8 +436,7 @@ class BuilderProcessor : AbstractProcessor() { } private fun buildTask( - name: String, - inner: ExecutableElement, + decl: TaskDeclaration, parent: TypeElement, ): TypeSpec { val executeSpec = @@ -313,8 +449,8 @@ class BuilderProcessor : AbstractProcessor() { .addParameter(CLIENT_TYPE, "client") .addException(Exception::class.java) - val dataParams = collectDataParams(inner) - val dataByName = dataParams.associateBy { it.name } + val inner = decl.method + val dataByName = decl.dataParams.associateBy { it.name } val innerArgs = with(processingEnv) { inner.parameters.joinToString { param -> @@ -327,9 +463,9 @@ class BuilderProcessor : AbstractProcessor() { } } - val taken = dataParams.mapTo(mutableSetOf()) { it.local } + val taken = decl.dataParams.mapTo(mutableSetOf()) { it.local } val argsLocal = generateSequence("args") { "${it}_" }.first { it !in taken } - val flatParams = dataParams.filterNot { it.isTaskInput } + val flatParams = decl.dataParams.filterNot { it.isTaskInput } if (flatParams.isNotEmpty()) { executeSpec.addStatement( $$"$T $L = $T.of(context, client, $L)", @@ -339,7 +475,7 @@ class BuilderProcessor : AbstractProcessor() { flatParams.size, ) } - dataParams.forEach { param -> + decl.dataParams.forEach { param -> val paramType = TypeName.get(param.type) if (param.isTaskInput) { executeSpec.addStatement( @@ -368,7 +504,7 @@ class BuilderProcessor : AbstractProcessor() { } return TypeSpec - .classBuilder(name) + .classBuilder(decl.className) .addSuperinterface(Task::class.java) .addModifiers(Modifier.PUBLIC, Modifier.FINAL, Modifier.STATIC) .addMethod(executeSpec.build()) @@ -479,6 +615,15 @@ class BuilderProcessor : AbstractProcessor() { } } +/** One [Builder.Task]-annotated method with its resolved id and data parameters. */ +private class TaskDeclaration( + val method: ExecutableElement, + val id: String, + val dataParams: List, +) { + val className: String = method.simpleName.toString().replaceFirstChar(Char::uppercase) +} + /** * One data parameter of a task method, positioned among its peers, read into * [local] by the generated body. [isTaskInput] marks a [TaskInput] parameter, @@ -501,6 +646,10 @@ 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 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 DEPS_TYPE = ClassName.get(Deps::class.java) private const val DAG_ANNOTATION = "org.apache.airflow.sdk.Builder.Dag" private const val TASK_ANNOTATION = "org.apache.airflow.sdk.Builder.Task" @@ -508,6 +657,15 @@ private const val TASK_ANNOTATION = "org.apache.airflow.sdk.Builder.Task" private val DAG_STRUCTURAL_ATTRIBUTES = setOf("id", "to") private val TASK_STRUCTURAL_ATTRIBUTES = setOf("id") +private val NUMERIC_KINDS = + setOf(TypeKind.BYTE, TypeKind.SHORT, TypeKind.INT, TypeKind.LONG, TypeKind.FLOAT, TypeKind.DOUBLE) + +private val BOXED_NUMERICS: Set = + setOf(TypeName.BYTE, TypeName.SHORT, TypeName.INT, TypeName.LONG, TypeName.FLOAT, TypeName.DOUBLE) + .mapTo(mutableSetOf()) { it.box() } + +private fun TypeName.boxIfPossible(): TypeName = if (this == TypeName.VOID || isPrimitive) box() else this + private fun ProcessingEnvironment.isType( t: TypeMirror, c: ClassName, 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 6decad1a90605..d4c8a5cf4550d 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 @@ -40,8 +40,8 @@ private fun JavaFileObjectSubject.hasSourceEquivalentTo( class BuilderTest { @Test - @DisplayName("generate builder for dag class") - fun generateBuilderForDagClass() { + @DisplayName("generate builder and task-reference twins for dag class") + fun generateBuilderAndRefForDagClass() { val compilation = compile( """ @@ -65,6 +65,14 @@ class BuilderTest { public void t3(Context ctx, int value) { System.out.println(String.format("%s %s", ctx.ti, value)); } + + @Builder.Deps + static class Wiring implements TestExampleDeps { + void depends() { + t1(); + t3(t2()); + } + } } """, ) @@ -80,20 +88,18 @@ class BuilderTest { import java.lang.Exception; import java.lang.Integer; 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.Task; - import org.apache.airflow.sdk.TaskDef; + 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"); - dag.addTask(new TaskDef("t1", T1.class)); - dag.addTask(new TaskDef("t2", T2.class)); - dag.addTask(new TaskDef("t3", T3.class)); - return dag; + return Refs.record(dag, List.of("t1", "t2", "t3"), new TestExample.Wiring()::depends); } public static final class T1 implements Task { @@ -121,6 +127,44 @@ class BuilderTest { } """, ) + assertThat(compilation) + .generatedSourceFile("org.apache.airflow.example.TestExampleDeps") + .hasSourceEquivalentTo( + "org.apache.airflow.example.TestExampleDeps", + """ + package org.apache.airflow.example; + + import java.lang.Integer; + import java.lang.Number; + import java.lang.Void; + import org.apache.airflow.sdk.Arg; + 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; + + /** + * Wiring view of {@link TestExample}'s task methods, for declaring its task graph. + * + *

Calling one registers its task with the Dag being built; passing the handle it + * returned into another call feeds the upstream's output into that task's parameter + * and wires the data edge. {@code before} and {@code after} wire an ordering-only edge. + */ + public interface TestExampleDeps extends Deps { + default TaskRef t1() { + return Refs.node(new TaskDef("t1", TestExampleBuilder.T1.class)); + } + + default TaskRef t2() { + return Refs.node(new TaskDef("t2", TestExampleBuilder.T2.class)); + } + + default TaskRef t3(Arg value) { + return Refs.call(new TaskDef("t3", TestExampleBuilder.T3.class), value); + } + } + """, + ) } @Test @@ -249,6 +293,94 @@ class BuilderTest { ) } + @Test + @DisplayName("type twin inputs by declared parameter type") + fun generateRefTypesTwinInputs() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import java.util.List; + import org.apache.airflow.sdk.Builder; + @Builder.Dag + public class TestExample { + @Builder.Task + public String ps() { return "x"; } + + @Builder.Task + public void pv() {} + + @Builder.Task + public List pl() { return null; } + + @Builder.Task + public long pn() { return 1L; } + + @Builder.Task + public void t(String text, Object anything, List items, Integer boxed) {} + + @Builder.Deps + static class Wiring implements TestExampleDeps { + void depends() { + t(ps(), pv(), pl(), pn()); + } + } + } + """, + ) + + assertThat(compilation).succeeded() + assertThat(compilation) + .generatedSourceFile("org.apache.airflow.example.TestExampleDeps") + .hasSourceEquivalentTo( + "org.apache.airflow.example.TestExampleDeps", + """ + package org.apache.airflow.example; + + import java.lang.Long; + import java.lang.Number; + import java.lang.String; + import java.lang.Void; + import java.util.List; + import org.apache.airflow.sdk.Arg; + 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; + + /** + * Wiring view of {@link TestExample}'s task methods, for declaring its task graph. + * + *

Calling one registers its task with the Dag being built; passing the handle it + * returned into another call feeds the upstream's output into that task's parameter + * and wires the data edge. {@code before} and {@code after} wire an ordering-only edge. + */ + public interface TestExampleDeps extends Deps { + default TaskRef ps() { + return Refs.node(new TaskDef("ps", TestExampleBuilder.Ps.class)); + } + + default TaskRef pv() { + return Refs.node(new TaskDef("pv", TestExampleBuilder.Pv.class)); + } + + default TaskRef> pl() { + return Refs.node(new TaskDef("pl", TestExampleBuilder.Pl.class)); + } + + default TaskRef pn() { + return Refs.node(new TaskDef("pn", TestExampleBuilder.Pn.class)); + } + + default TaskRef t(Arg text, Arg anything, + Arg> items, Arg boxed) { + return Refs.call(new TaskDef("t", TestExampleBuilder.T.class), text, anything, items, boxed); + } + } + """, + ) + } + @Test @DisplayName("lower explicit annotation attributes into config calls") fun generateBuilderLowersConfigAttributes() { @@ -262,6 +394,13 @@ class BuilderTest { public class TestExample { @Builder.Task(retries = 2, queue = "q", retryDelay = "PT5M", retryExponentialBackoff = 1.5) public void t1() {} + + @Builder.Deps + static class Wiring implements TestExampleDeps { + void depends() { + t1(); + } + } } """, ) @@ -276,14 +415,13 @@ class BuilderTest { import java.lang.Exception; import java.lang.Override; - import java.time.Duration; import java.time.OffsetDateTime; 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.Task; - import org.apache.airflow.sdk.TaskDef; + import org.apache.airflow.sdk.internal.Refs; public final class TestExampleBuilder { public static DagDef build() { @@ -292,8 +430,7 @@ class BuilderTest { dag.config("tags", List.of("a", "b")); dag.config("catchup", true); dag.config("start_date", OffsetDateTime.parse("2026-01-01T00:00:00Z")); - dag.addTask(new TaskDef("t1", T1.class).config("retries", 2).config("queue", "q").config("retry_delay", Duration.parse("PT5M")).config("retry_exponential_backoff", 1.5)); - return dag; + return Refs.record(dag, List.of("t1"), new TestExample.Wiring()::depends); } public static final class T1 implements Task { @@ -305,6 +442,34 @@ class BuilderTest { } """, ) + assertThat(compilation) + .generatedSourceFile("org.apache.airflow.example.TestExampleDeps") + .hasSourceEquivalentTo( + "org.apache.airflow.example.TestExampleDeps", + """ + package org.apache.airflow.example; + + import java.lang.Void; + import java.time.Duration; + 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; + + /** + * Wiring view of {@link TestExample}'s task methods, for declaring its task graph. + * + *

Calling one registers its task with the Dag being built; passing the handle it + * returned into another call feeds the upstream's output into that task's parameter + * and wires the data edge. {@code before} and {@code after} wire an ordering-only edge. + */ + public interface TestExampleDeps extends Deps { + default TaskRef t1() { + return Refs.node(new TaskDef("t1", TestExampleBuilder.T1.class).config("retries", 2).config("queue", "q").config("retry_delay", Duration.parse("PT5M")).config("retry_exponential_backoff", 1.5)); + } + } + """, + ) } @Test @@ -664,6 +829,13 @@ class BuilderTest { @Builder.Dag public class TestExample { @Builder.Task(id = "foo") public void t1() {} + + @Builder.Deps + static class Wiring implements TestExampleDeps { + void depends() { + t1(); + } + } } """, ) @@ -677,17 +849,17 @@ class BuilderTest { 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.Task; - import org.apache.airflow.sdk.TaskDef; + import org.apache.airflow.sdk.internal.Refs; public final class TestExampleBuilder { public static DagDef build() { var dag = new DagDef("TestExample"); - dag.addTask(new TaskDef("foo", T1.class)); - return dag; + return Refs.record(dag, List.of("foo"), new TestExample.Wiring()::depends); } public static final class T1 implements Task { @@ -699,6 +871,140 @@ class BuilderTest { } """, ) + 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; + + /** + * Wiring view of {@link TestExample}'s task methods, for declaring its task graph. + * + *

Calling one registers its task with the Dag being built; passing the handle it + * returned into another call feeds the upstream's output into that task's parameter + * and wires the data edge. {@code before} and {@code after} wire an ordering-only edge. + */ + public interface TestExampleDeps extends Deps { + default TaskRef t1() { + return Refs.node(new TaskDef("foo", TestExampleBuilder.T1.class)); + } + } + """, + ) + } + + @Test + @DisplayName("reject wiring that feeds an incompatible upstream type") + fun rejectIncompatibleWiring() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + @Builder.Dag + public class TestExample { + @Builder.Task + public String ps() { return "x"; } + + @Builder.Task + public void t(int v) {} + + @Builder.Deps + static class Wiring implements TestExampleDeps { + void depends() { t(ps()); } + } + } + """, + ) + assertThat(compilation).failed() + assertThat(compilation).hadErrorContaining("incompatible types") + } + + @Test + @DisplayName("reject more than one wiring class") + fun rejectMultipleWiringClasses() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + @Builder.Dag + public class TestExample { + @Builder.Task public void t1() {} + + @Builder.Deps + static class One implements TestExampleDeps { + void depends() { t1(); } + } + + @Builder.Deps + static class Two implements TestExampleDeps { + void depends() {} + } + } + """, + ) + assertThat(compilation).failed() + assertThat(compilation).hadErrorContaining( + "Dag class TestExample declares more than one @Builder.Deps class: One, Two", + ) + } + + @Test + @DisplayName("reject a non-static wiring class") + fun rejectNonStaticWiringClass() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + @Builder.Dag + public class TestExample { + @Builder.Task public void t1() {} + + @Builder.Deps + class Wiring implements TestExampleDeps { + void depends() { t1(); } + } + } + """, + ) + assertThat(compilation).failed() + assertThat(compilation).hadErrorContaining( + "@Builder.Deps class 'Wiring' must be static and non-private", + ) + } + + @Test + @DisplayName("reject a wiring class with no depends() method") + fun rejectWiringClassWithoutDepends() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + @Builder.Dag + public class TestExample { + @Builder.Task public void t1() {} + + @Builder.Deps + static class Wiring implements TestExampleDeps { + void wireItUp() { t1(); } + } + } + """, + ) + assertThat(compilation).failed() + assertThat(compilation).hadErrorContaining( + "@Builder.Deps class 'Wiring' must declare a non-private, no-argument depends() method", + ) } @Test @@ -719,6 +1025,28 @@ class BuilderTest { ) } + @Test + @DisplayName("reject duplicate task ids") + fun rejectDuplicateTaskIds() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + @Builder.Dag + public class TestExample { + @Builder.Task(id = "x") + public void t1() {} + + @Builder.Task(id = "x") + public void t2() {} + } + """, + ) + assertThat(compilation).failed() + assertThat(compilation).hadErrorContaining("Tasks in Dag have duplicate ID: x") + } + @Test @DisplayName("reject a duration attribute that is not ISO-8601") fun rejectInvalidDurationAttribute() { diff --git a/java-sdk/sdk/build.gradle.kts b/java-sdk/sdk/build.gradle.kts index 5e9a74a704638..ba7b46fab0df4 100644 --- a/java-sdk/sdk/build.gradle.kts +++ b/java-sdk/sdk/build.gradle.kts @@ -504,7 +504,8 @@ abstract class GenerateDagDslTask : DefaultTask() { | * Container for the annotation-based Dag-authoring API. | * | * Annotating a class with [Dag] generates a `Builder` whose static - | * `build()` returns the [DagDef] to add to a [Bundle]. + | * `build()` returns the [DagDef] to add to a [Bundle], and — when the class + | * declares a [Deps] class — a `Deps` wiring view for it to implement. | * | * Example: | * @@ -521,9 +522,10 @@ abstract class GenerateDagDslTask : DefaultTask() { | * ``` | * | * A task method's data parameters — everything other than the injected - | * [Client] and [Context] — receive, by position, the arguments the Python - | * `@task.stub` call site bound. Keyword arguments bind by name instead - | * through a single [TaskInput] parameter. + | * [Client] and [Context] — receive, by position, the inputs the [Deps] + | * class wired. For a task the Python Dag file declares with `@task.stub`, + | * the arguments bound at that call site take their place. Keyword + | * arguments bind by name instead through a single [TaskInput] parameter. | */ |class Builder internal constructor() { | /** @@ -587,6 +589,35 @@ abstract class GenerateDagDslTask : DefaultTask() { | val dag: String, | val task: String = "", | ) + | + | /** + | * Marks the nested class that declares this Dag's task graph. + | * + | * Declare it as a `static` nested class that implements the generated + | * `Deps` wiring view and has a no-argument `depends()` method. + | * Calling a view method registers its task; passing the handle one + | * returned into another call wires a data edge; `before` and `after` + | * wire an ordering-only one: + | * + | * ```java + | * @Builder.Deps + | * static class Wiring implements EtlPipelineDeps { + | * void depends() { + | * var rows = extract(); + | * var loaded = load(transform(rows, lit(0.9))); + | * rows.before(audit()); + | * report().after(loaded, audit()); + | * } + | * } + | * ``` + | * + | * What makes the Dag native: a Dag class without one registers every + | * task with no Java-side edges, which is the task-handler shape rather + | * than a Dag Java owns. + | */ + | @Target(AnnotationTarget.CLASS) + | @MustBeDocumented + | annotation class Deps |} | """.trimMargin(), diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Arg.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Arg.kt index 556ed5a7d4a2c..bff37ad02efa3 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Arg.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Arg.kt @@ -28,7 +28,21 @@ package org.apache.airflow.sdk * * @param T Type of the value. */ -sealed class Arg +sealed class Arg { + companion object { + /** + * Wraps an inline constant as a task argument, passed to the task as a + * constant and creating no dependency edge. + * + * [Deps.lit] is the spelling a [Builder.Deps] class uses; this is the same + * call for code with no `Deps` in scope. + * + * @param value Constant to bind; may be null for a nullable parameter. + */ + @JvmStatic + fun lit(value: T?): Arg = LiteralArg(value) + } +} internal class LiteralArg( internal val value: T?, @@ -37,8 +51,10 @@ internal class LiteralArg( /** * The output of a registered task, and the task's place in the flow. * - * [Deps.Flow.before] and [Deps.Flow.after] wire an ordering-only edge from - * this task, where nothing flows but the sequence. + * Passing this handle into another task's arguments feeds this task's return + * value into that parameter and wires the data edge. [Deps.Flow.before] and + * [Deps.Flow.after] wire an ordering-only edge instead, where nothing flows + * but the sequence. * * @param T Return type of the task this handle refers to. */ 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 e1dff61889465..e26fc9608c78f 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 @@ -89,28 +89,26 @@ class Bundle( } /** - * Registers every task handler a class holds, from the ids each - * [Builder.TaskHandler] names. + * Registers what an annotated class holds — the same class the annotations + * are on, so there is no second name to keep in sync. * - * @param handlerClass A class with [Builder.TaskHandler] methods. + * A [Builder.Dag] class contributes the Dag its generated builder builds; a + * class of [Builder.TaskHandler] methods contributes each handler, bound to + * the Dag the Python file owns. + * + * @param annotated A class carrying [Builder.Dag] or [Builder.TaskHandler]. * @return This bundle, for chaining. - * @throws IllegalArgumentException if the class has no generated - * registrar, because annotation processing did not run over it. + * @throws IllegalArgumentException if the class has no generated code, + * because annotation processing did not run over it. */ - fun register(handlerClass: Class<*>): Bundle { + fun register(annotated: Class<*>): Bundle { checkOpen() - val name = registrarName(handlerClass.name) - val registrar = - try { - Class.forName(name, true, handlerClass.classLoader) - } catch (e: ClassNotFoundException) { - throw IllegalArgumentException( - "No generated registrar $name for ${handlerClass.name}; does it declare " + - "@Builder.TaskHandler methods, and is airflow-sdk-processor on the " + - "annotationProcessor path?", - e, - ) - } + annotated.getAnnotation(Builder.Dag::class.java)?.let { dag -> + val builderName = dag.to.ifBlank { "${annotated.simpleName}Builder" } + val builder = generated("${annotated.packageName}.$builderName", annotated, "builder") + return register(builder.getMethod("build").invoke(null) as DagDef) + } + val registrar = generated(registrarName(annotated.name), annotated, "registrar") registrar.getMethod("registerInto", Bundle::class.java).invoke(null, this) return this } @@ -150,6 +148,21 @@ class Bundle( taskId: String, ): TaskDef? = (dags[dagId] ?: taskHandlers[dagId])?.tasks?.get(taskId) + private fun generated( + name: String, + from: Class<*>, + what: String, + ): Class<*> = + try { + Class.forName(name, true, from.classLoader) + } catch (e: ClassNotFoundException) { + throw IllegalArgumentException( + "No generated $what $name for ${from.name}; does it carry @Builder.Dag or " + + "@Builder.TaskHandler, and is airflow-sdk-processor on the annotationProcessor path?", + e, + ) + } + /** * Ends registration, so a `register` left below `serve` is reported as the * mistake it is rather than racing the runtime. [Server] calls it when it diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Deps.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Deps.kt index 38a14d0035477..3b0fda4d0f2b2 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Deps.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Deps.kt @@ -19,7 +19,13 @@ package org.apache.airflow.sdk -/** Vocabulary for declaring a Dag's task graph in Java. */ +/** + * Vocabulary for declaring a Dag's task graph in Java, and the base of every + * generated `Deps` wiring view. + * + * A [Builder.Deps] class inherits [lit] and [Flow] by simple name, so it needs + * no import and `Flow` does not collide with `java.util.concurrent.Flow`. + */ interface Deps { /** * A point in the task graph: one task, or a set of them. @@ -81,6 +87,15 @@ interface Deps { fun of(vararg flows: Flow): Flow = FlowSet(flows.flatMap { it.nodes() }) } } + + /** + * Wraps an inline constant as a task argument — `transform(extract(), + * lit(0.9))` — passed to the task as a constant, creating no dependency + * edge. + * + * @param value Constant to bind; may be null for a nullable parameter. + */ + fun lit(value: T?): Arg = Arg.lit(value) } /** Several tasks as one point in the flow, which no single [TaskRef] can represent. */ diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/Refs.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/Refs.kt new file mode 100644 index 0000000000000..e875f9be98cc4 --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/Refs.kt @@ -0,0 +1,114 @@ +/* + * 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.Arg +import org.apache.airflow.sdk.DagDef +import org.apache.airflow.sdk.TaskDef +import org.apache.airflow.sdk.TaskRef + +/** + * @suppress + * + * The recorder behind a `@Builder.Deps` wiring class. Public so that + * processor-generated wiring views can call it; not user-facing API. + * + * A wiring view's methods are `default` methods on an interface, so they hold + * no Dag of their own. [record] puts the Dag being built in scope for exactly + * the duration of one `depends()` call, and the view's methods register into + * it. Nothing parses a syntax tree: identity travels with the [TaskRef] a + * call returns, so a result held in a local and reused just works. + */ +object Refs { + private class Recording( + val dag: DagDef, + ) { + val byTaskId = linkedMapOf>() + } + + private val recording = ThreadLocal() + + /** + * Runs one `depends()` call with [dag] in scope, then returns the Dag the + * wiring built. + * + * @throws IllegalArgumentException if the wiring left a declared task + * unregistered. + */ + @JvmStatic + fun record( + dag: DagDef, + taskIds: List, + depends: Runnable, + ): DagDef { + check(recording.get() == null) { "Dag wiring is already being recorded on this thread" } + recording.set(Recording(dag)) + try { + depends.run() + } finally { + recording.remove() + } + val missing = taskIds.filterNot { it in dag.tasks } + require(missing.isEmpty()) { + "Wiring for Dag '${dag.id}' did not register task(s) ${missing.joinToString { "'$it'" }}: " + + "every @Builder.Task method must be called in the @Builder.Deps class" + } + return dag + } + + /** + * Records a task that takes no data arguments. + * + * @return The handle representing this task, memoized by [TaskDef.id] so every call + * yields the same one. + */ + @JvmStatic + fun node(def: TaskDef): TaskRef = call(def) + + /** + * Records a task and the data edge for every [TaskRef] among [args]; a + * literal argument records a baked value and no edge. + * + * @return The handle representing this task, memoized by [TaskDef.id] so a result + * held in a local and reused refers to one node. + */ + @JvmStatic + @Suppress("UNCHECKED_CAST", "SpreadOperator") + fun call( + def: TaskDef, + vararg args: Arg<*>, + ): TaskRef { + val active = + checkNotNull(recording.get()) { + "Task '${def.id}' was wired outside a @Builder.Deps class; the wiring view's methods " + + "only record while the generated builder is running depends()" + } + active.byTaskId[def.id]?.let { existing -> + require(args.isEmpty()) { + "Task '${def.id}' is wired more than once with arguments; call it once and reuse the handle it returned" + } + return existing as TaskRef + } + args.filterIsInstance>().forEach { def.dependsOn(it.def) } + def.inputs += args + active.dag.addTask(def) + return TaskRef(def).also { active.byTaskId[def.id] = it } + } +} 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 1cc3ff3dbaef7..6733621ab0bf3 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 @@ -24,6 +24,7 @@ import org.apache.airflow.sdk.execution.comm.StartupDetails import org.apache.airflow.sdk.execution.comm.TIRunContext import org.apache.airflow.sdk.execution.comm.VariableResult import org.apache.airflow.sdk.execution.comm.XComResult +import org.apache.airflow.sdk.internal.Refs import org.apache.airflow.sdk.execution.comm.TaskInstance as CommTaskInstance /** Records getXCom calls and serves canned values keyed by task id. */ @@ -104,8 +105,8 @@ internal class NoopTask : Task { /** A context whose task was wired by its Dag with the given inputs. */ internal fun contextWiredWith(inputs: List>): Context { + val dag = DagDef("d") val def = TaskDef("t", NoopTask::class.java) - DagDef("d").addTask(def) - def.inputs += inputs + Refs.record(dag, listOf("t")) { Refs.call(def, *inputs.toTypedArray()) } return taskContext().also { it.taskDef = def } } diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/BundleTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/BundleTest.kt index 29c5fc8a99b6f..a74c062828037 100644 --- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/BundleTest.kt +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/BundleTest.kt @@ -134,14 +134,47 @@ internal class BundleTest { Assertions.assertEquals(emptySet(), bundle.dags.keys) } + @Test + @DisplayName("Should register Dags one at a time and reject a duplicate ID") + fun shouldRegisterDagsIncrementally() { + val bundle = Bundle().register(DagDef("a")).register(DagDef("b")) + + Assertions.assertEquals(setOf("a", "b"), bundle.dags.keys) + val error = + Assertions.assertThrows(IllegalArgumentException::class.java) { bundle.register(DagDef("a")) } + Assertions.assertEquals("Dags in bundle have duplicate ID: a", error.message) + } + + @Test + @DisplayName("Should create the Dag on first use when registering a stub-backed handler") + fun shouldRegisterHandlerAgainstPythonOwnedDag() { + val bundle = + Bundle() + .register("etl", "score", NoopBundleTask::class.java) + .register("etl", "report", NoopBundleTask::class.java) + + Assertions.assertEquals(setOf("etl"), bundle.taskHandlers.keys) + Assertions.assertEquals( + setOf("score", "report"), + bundle.taskHandlers + .getValue("etl") + .tasks.keys, + ) + } + @Test @DisplayName("Should find the registrar generated for a nested handler class") fun shouldFindRegistrarOfNestedHandlerClass() { val bundle = Bundle().register(Nested::class.java) - val etl = bundle.taskHandlers.getValue("etl") Assertions.assertEquals(listOf("etl"), bundle.taskHandlers.keys.toList()) - Assertions.assertEquals(listOf("score"), etl.tasks.keys.toList()) + Assertions.assertEquals( + listOf("score"), + bundle.taskHandlers + .getValue("etl") + .tasks.keys + .toList(), + ) } @Test @@ -152,10 +185,10 @@ internal class BundleTest { Bundle().register(NoOp::class.java) } - Assertions.assertTrue( - error.message!!.startsWith( - "No generated registrar org.apache.airflow.sdk.BundleTest_NoOpHandlers for ", - ), + Assertions.assertEquals( + "No generated registrar org.apache.airflow.sdk.BundleTest_NoOpHandlers for " + + "${NoOp::class.java.name}; does it carry @Builder.Dag or @Builder.TaskHandler, " + + "and is airflow-sdk-processor on the annotationProcessor path?", error.message, ) } @@ -240,6 +273,13 @@ internal class BundleTest { } } +class NoopBundleTask : Task { + override fun execute( + context: Context, + client: Client, + ) = Unit +} + /** * Stands in for the registrar the annotation processor generates beside * [BundleTest.Nested], to pin the name [Bundle.register] looks up. @@ -249,14 +289,7 @@ class BundleTest_NestedHandlers { companion object { @JvmStatic fun registerInto(bundle: Bundle) { - bundle.register("etl", "score", NoOpHandler::class.java) + bundle.register("etl", "score", NoopBundleTask::class.java) } } - - class NoOpHandler : Task { - override fun execute( - context: Context, - client: Client, - ) = Unit - } } diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/InputTaskTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/InputTaskTest.kt index ca557b9261158..fadbf6fad61d2 100644 --- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/InputTaskTest.kt +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/InputTaskTest.kt @@ -122,9 +122,9 @@ internal class InputTaskTest { @DisplayName("Should decode a TaskInput wholesale from its wired input when no bindings arrive") fun shouldDecodeTaskInputFromWiredInput() { // A native Dag has no stub call site, so there are no argument names to - // match fields against: the input the Dag wired to this task decodes into - // the whole TaskInput at once. - val context = contextWiredWith(listOf(LiteralArg(mapOf("region" to "emea", "threshold" to 0.5)))) + // match fields against: the input the wiring class fed this task decodes + // into the whole TaskInput at once. + val context = contextWiredWith(listOf(Arg.lit(mapOf("region" to "emea", "threshold" to 0.5)))) val (client, _) = clientWith(null) val task = Summarize() 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 768ad066f5c2a..7d8439634aae0 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 @@ -51,7 +51,7 @@ private class NoopArgTask : Task { ) = Unit } -/** Resolution of the inputs a Dag declared in Java, without runtime bindings. */ +/** Resolution of the inputs a `@Builder.Deps` class recorded, without runtime bindings. */ internal class ArgValuesTest { /** Upstream task ids read through the transport, in arrival order. */ private val pulls = java.util.concurrent.CopyOnWriteArrayList() @@ -119,8 +119,7 @@ internal class ArgValuesTest { .distinct() .forEach { dag.addTask(it) } val def = TaskDef("consumer", NoopArgTask::class.java) - dag.addTask(def) - def.inputs += inputs + Refs.record(dag, listOf("consumer")) { Refs.call(def, *inputs.toTypedArray()) } return contextWithoutTaskDef().also { it.taskDef = def } } 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 new file mode 100644 index 0000000000000..f941803c8708e --- /dev/null +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/RefsTest.kt @@ -0,0 +1,115 @@ +/* + * 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.Arg +import org.apache.airflow.sdk.Client +import org.apache.airflow.sdk.Context +import org.apache.airflow.sdk.DagDef +import org.apache.airflow.sdk.LiteralArg +import org.apache.airflow.sdk.Task +import org.apache.airflow.sdk.TaskDef +import org.apache.airflow.sdk.TaskRef +import org.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.Assertions.assertSame +import org.junit.jupiter.api.Assertions.assertThrows +import org.junit.jupiter.api.DisplayName +import org.junit.jupiter.api.Test + +private class NoopRefTask : Task { + override fun execute( + context: Context, + client: Client, + ) = Unit +} + +internal class RefsTest { + @Test + @DisplayName("Should register the task, record inputs, and wire handle edges") + fun shouldRegisterTaskWithInputsAndEdges() { + val dag = DagDef("d") + Refs.record(dag, listOf("p", "c")) { + val producer = Refs.node(TaskDef("p", NoopRefTask::class.java)) + Refs.call(TaskDef("c", NoopRefTask::class.java), producer, Arg.lit(5)) + } + + val consumerDef = dag.tasks.getValue("c") + assertEquals(setOf("p", "c"), dag.tasks.keys) + assertEquals(setOf(dag.tasks.getValue("p")), consumerDef.upstreams) + assertEquals(2, consumerDef.inputs.size) + assertEquals(dag.tasks.getValue("p"), (consumerDef.inputs[0] as TaskRef<*>).def) + assertEquals(5, (consumerDef.inputs[1] as LiteralArg<*>).value) + } + + @Test + @DisplayName("Should return the same handle wherever a task is wired") + fun shouldMemoizeHandleByTaskId() { + val dag = DagDef("d") + Refs.record(dag, listOf("a", "b")) { + val first = Refs.node(TaskDef("a", NoopRefTask::class.java)) + val again = Refs.node(TaskDef("a", NoopRefTask::class.java)) + assertSame(first, again) + first.before(Refs.node(TaskDef("b", NoopRefTask::class.java))) + } + + assertEquals(setOf("a", "b"), dag.tasks.keys) + assertEquals(setOf(dag.tasks.getValue("a")), dag.tasks.getValue("b").upstreams) + } + + @Test + @DisplayName("Should pass when the wiring registered every task") + fun shouldPassWhenWiringComplete() { + val dag = DagDef("d") + + Refs.record(dag, listOf("t")) { Refs.node(TaskDef("t", NoopRefTask::class.java)) } + } + + @Test + @DisplayName("Should fail naming the tasks the wiring missed") + fun shouldFailNamingMissedTasks() { + val dag = DagDef("d") + + val error = + assertThrows(IllegalArgumentException::class.java) { + Refs.record(dag, listOf("t", "x", "y")) { Refs.node(TaskDef("t", NoopRefTask::class.java)) } + } + + assertEquals( + "Wiring for Dag 'd' did not register task(s) 'x', 'y': " + + "every @Builder.Task method must be called in the @Builder.Deps class", + error.message, + ) + } + + @Test + @DisplayName("Should refuse a wiring call made outside a recording") + fun shouldRefuseWiringOutsideRecording() { + val error = + assertThrows(IllegalStateException::class.java) { + Refs.node(TaskDef("t", NoopRefTask::class.java)) + } + + assertEquals( + "Task 't' was wired outside a @Builder.Deps class; the wiring view's methods " + + "only record while the generated builder is running depends()", + error.message, + ) + } +} From 8d92b50fc44577169750b2334cd2dc3c1fc8ebc1 Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Sat, 3 Oct 2026 13:04:04 +0000 Subject: [PATCH 02/13] Java SDK: Quote wired task ids as Java string literals --- .../apache/airflow/sdk/BuilderProcessor.kt | 2 +- .../org/apache/airflow/sdk/BuilderTest.kt | 31 +++++++++++++++++++ 2 files changed, 32 insertions(+), 1 deletion(-) 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 f97f72997d0a0..d9f1371b9a34d 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 @@ -218,7 +218,7 @@ class BuilderProcessor : AbstractProcessor() { $$"return $T.record(dag, $T.of($L), new $T()::depends)", REFS_TYPE, ClassName.get(List::class.java), - declarations.joinToString { "\"${it.id}\"" }, + CodeBlock.join(declarations.map { CodeBlock.of($$"$S", it.id) }, ", "), ClassName.get(deps), ) } else { 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 d4c8a5cf4550d..45a792ccbda34 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 @@ -1089,6 +1089,37 @@ class BuilderTest { .contains("""dag.config("tags", List.of("say \"hi\"", "back\\slash"));""") } + @Test + @DisplayName("escape quotes and backslashes in the task ids the wiring must register") + fun generateBuilderEscapesWiredTaskIds() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + @Builder.Dag + public class TestExample { + @Builder.Task(id = "say \"hi\"") public void a() {} + @Builder.Task(id = "back\\slash") public void b() {} + + @Builder.Deps + static class Wiring implements TestExampleDeps { + void depends() { + a(); + b(); + } + } + } + """, + ) + + assertThat(compilation).succeeded() + assertThat(compilation) + .generatedSourceFile("org.apache.airflow.example.TestExampleBuilder") + .contentsAsUtf8String() + .contains("""List.of("say \"hi\"", "back\\slash")""") + } + @Test @DisplayName("bind a TaskInput through the shared populator") fun generateBuilderBindsTaskInputFields() { From 2e4bec76b5f3d29dd86051904c65b63375de70ad Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Sat, 3 Oct 2026 13:05:42 +0000 Subject: [PATCH 03/13] Java SDK: Validate the @Builder.Deps class at compile time --- .../apache/airflow/sdk/BuilderProcessor.kt | 64 ++++++- .../org/apache/airflow/sdk/BuilderTest.kt | 157 +++++++++++++++++- 2 files changed, 215 insertions(+), 6 deletions(-) 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 d9f1371b9a34d..3cb387bdc5ca5 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 @@ -92,7 +92,9 @@ import javax.tools.Diagnostic */ @SupportedAnnotationTypes( "org.apache.airflow.sdk.Builder.Dag", + "org.apache.airflow.sdk.Builder.Task", "org.apache.airflow.sdk.Builder.TaskHandler", + "org.apache.airflow.sdk.Builder.Deps", ) @SupportedSourceVersion(SourceVersion.RELEASE_11) class BuilderProcessor : AbstractProcessor() { @@ -101,6 +103,16 @@ class BuilderProcessor : AbstractProcessor() { roundEnv: RoundEnvironment, ): Boolean { if (annotations.isEmpty()) return false + roundEnv.getElementsAnnotatedWith(Builder.Deps::class.java).forEach { el -> + val owner = el.enclosingElement + if (owner !is TypeElement || owner.getAnnotation(Builder.Dag::class.java) == null) { + processingEnv.messager.printMessage( + Diagnostic.Kind.ERROR, + "@Builder.Deps class '${el.simpleName}' must be nested directly in a @Builder.Dag class", + el, + ) + } + } roundEnv .getElementsAnnotatedWith(Builder.TaskHandler::class.java) .mapNotNull { it.enclosingElement as? TypeElement } @@ -120,9 +132,9 @@ class BuilderProcessor : AbstractProcessor() { runCatching { val packageName = elementUtils.getPackageOf(el).qualifiedName.toString() val declarations = collectTasks(el) - val deps = findDeps(el) val builderName = ClassName.get(packageName, dagAnnotation(el).to.ifBlank { "${el.simpleName}Builder" }) val depsName = ClassName.get(packageName, "${el.simpleName}Deps") + val deps = findDeps(el, depsName) JavaFile .builder(packageName, buildBuilder(el, declarations, deps, builderName)) .build() @@ -339,8 +351,14 @@ class BuilderProcessor : AbstractProcessor() { /** * Finds and validates the class's `@Builder.Deps` wiring class. It is * optional: without one, every task registers with no Java-side edges. + * + * The generated builder runs `new Wiring()::depends`, so everything that + * expression needs is checked here, where the error can name the class. */ - private fun findDeps(el: TypeElement): TypeElement? { + private fun findDeps( + el: TypeElement, + view: ClassName, + ): TypeElement? { val classes = el.enclosedElements .filterIsInstance() @@ -352,15 +370,51 @@ class BuilderProcessor : AbstractProcessor() { "Dag class ${el.simpleName} declares more than one @Builder.Deps class: " + classes.joinToString { it.simpleName.toString() }, ) + val name = deps.simpleName + require(deps.kind == ElementKind.CLASS && Modifier.ABSTRACT !in deps.modifiers) { + "@Builder.Deps '$name' must be a concrete class" + } require(Modifier.STATIC in deps.modifiers && Modifier.PRIVATE !in deps.modifiers) { - "@Builder.Deps class '${deps.simpleName}' must be static and non-private" + "@Builder.Deps class '$name' must be static and non-private" + } + require(deps.interfaces.any { it.isView(view) }) { + "@Builder.Deps class '$name' must implement ${view.simpleName()}, the wiring view of ${el.simpleName}" } - require(deps.enclosedElements.filterIsInstance().any { it.isNoArgDepends() }) { - "@Builder.Deps class '${deps.simpleName}' must declare a non-private, no-argument depends() method" + require( + deps.enclosedElements + .filterIsInstance() + .any { it.kind == ElementKind.CONSTRUCTOR && it.parameters.isEmpty() && Modifier.PRIVATE !in it.modifiers }, + ) { + "@Builder.Deps class '$name' needs a non-private no-argument constructor" + } + val depends = + processingEnv.elementUtils + .getAllMembers(deps) + .filterIsInstance() + .firstOrNull { it.isNoArgDepends() } + ?: throw IllegalArgumentException( + "@Builder.Deps class '$name' must have a non-private, no-argument depends() method", + ) + val checked = depends.thrownTypes.filterNot { isUnchecked(it) } + require(checked.isEmpty()) { + "depends() of @Builder.Deps class '$name' must not throw checked exceptions: ${checked.joinToString()}" } return deps } + /** + * Matches the view by the name the class wrote: the view is generated in + * this same round, so javac may not have resolved it yet. + */ + private fun TypeMirror.isView(view: ClassName): Boolean = toString().let { it == view.canonicalName() || it == view.simpleName() } + + private fun isUnchecked(type: TypeMirror): Boolean = + with(processingEnv) { + listOf(RuntimeException::class.java, Error::class.java).any { + typeUtils.isAssignable(type, elementUtils.getTypeElement(it.canonicalName).asType()) + } + } + private fun ExecutableElement.isNoArgDepends(): Boolean = simpleName.contentEquals("depends") && parameters.isEmpty() && 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 45a792ccbda34..04b90bc7eb43a 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 @@ -1003,7 +1003,162 @@ class BuilderTest { ) assertThat(compilation).failed() assertThat(compilation).hadErrorContaining( - "@Builder.Deps class 'Wiring' must declare a non-private, no-argument depends() method", + "@Builder.Deps class 'Wiring' must have a non-private, no-argument depends() method", + ) + } + + @Test + @DisplayName("reject a wiring class that implements another Dag's wiring view") + fun rejectWiringClassOfAnotherDag() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + @Builder.Dag + public class TestExample { + @Builder.Task public void t1() {} + + @Builder.Deps + static class Wiring implements OtherDeps { + void depends() { t1(); } + } + } + + @Builder.Dag + class Other { + @Builder.Task public void t1() {} + + @Builder.Deps + static class Wiring implements OtherDeps { + void depends() { t1(); } + } + } + """, + ) + assertThat(compilation).failed() + assertThat(compilation).hadErrorContaining( + "@Builder.Deps class 'Wiring' must implement TestExampleDeps, the wiring view of TestExample", + ) + } + + @Test + @DisplayName("reject an abstract wiring class") + fun rejectAbstractWiringClass() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + @Builder.Dag + public class TestExample { + @Builder.Task public void t1() {} + + @Builder.Deps + abstract static class Wiring implements TestExampleDeps { + void depends() { t1(); } + } + } + """, + ) + assertThat(compilation).failed() + assertThat(compilation).hadErrorContaining("@Builder.Deps 'Wiring' must be a concrete class") + } + + @Test + @DisplayName("reject a wiring class with no no-argument constructor") + fun rejectWiringClassWithoutNoArgConstructor() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + @Builder.Dag + public class TestExample { + @Builder.Task public void t1() {} + + @Builder.Deps + static class Wiring implements TestExampleDeps { + Wiring(int unused) {} + void depends() { t1(); } + } + } + """, + ) + assertThat(compilation).failed() + assertThat(compilation).hadErrorContaining( + "@Builder.Deps class 'Wiring' needs a non-private no-argument constructor", + ) + } + + @Test + @DisplayName("reject a depends() that throws a checked exception") + fun rejectDependsThrowingCheckedException() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + @Builder.Dag + public class TestExample { + @Builder.Task public void t1() {} + + @Builder.Deps + static class Wiring implements TestExampleDeps { + void depends() throws java.io.IOException { t1(); } + } + } + """, + ) + assertThat(compilation).failed() + assertThat(compilation).hadErrorContaining( + "depends() of @Builder.Deps class 'Wiring' must not throw checked exceptions: java.io.IOException", + ) + } + + @Test + @DisplayName("accept a depends() the wiring class inherits") + fun acceptInheritedDepends() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + @Builder.Dag + public class TestExample { + @Builder.Task public void t1() {} + + static class Base implements TestExampleDeps { + void depends() { t1(); } + } + + @Builder.Deps + static class Wiring extends Base implements TestExampleDeps {} + } + """, + ) + assertThat(compilation).succeeded() + } + + @Test + @DisplayName("reject a wiring class outside a Dag class") + fun rejectMisplacedWiringClass() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + public class TestExample { + @Builder.Deps + static class Wiring { + void depends() {} + } + } + """, + ) + assertThat(compilation).failed() + assertThat(compilation).hadErrorContaining( + "@Builder.Deps class 'Wiring' must be nested directly in a @Builder.Dag class", ) } From 79e339556d5dffcaaca5e02e50bc08bbbfdb3360 Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Sat, 3 Oct 2026 13:06:14 +0000 Subject: [PATCH 04/13] Java SDK: Reject task method names that clash with the wiring view --- .../apache/airflow/sdk/BuilderProcessor.kt | 29 +++++++++++++++++++ .../org/apache/airflow/sdk/BuilderTest.kt | 26 +++++++++++++++++ 2 files changed, 55 insertions(+) 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 3cb387bdc5ca5..f76969b5a971f 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 @@ -135,6 +135,7 @@ class BuilderProcessor : AbstractProcessor() { val builderName = ClassName.get(packageName, dagAnnotation(el).to.ifBlank { "${el.simpleName}Builder" }) val depsName = ClassName.get(packageName, "${el.simpleName}Deps") val deps = findDeps(el, depsName) + if (deps != null) declarations.forEach { checkViewName(it) } JavaFile .builder(packageName, buildBuilder(el, declarations, deps, builderName)) .build() @@ -402,6 +403,19 @@ class BuilderProcessor : AbstractProcessor() { return deps } + /** + * Rejects a task method whose wiring-view twin would clash with a member + * the view or the wiring class already has: `depends`, `lit`, or a method + * of `Object`. + */ + private fun checkViewName(decl: TaskDeclaration) { + val name = decl.method.simpleName.toString() + require(name !in RESERVED_VIEW_NAMES) { + "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}\")" + } + } + /** * Matches the view by the name the class wrote: the view is generated in * this same round, so javac may not have resolved it yet. @@ -708,6 +722,21 @@ private val DEPS_TYPE = ClassName.get(Deps::class.java) private const val DAG_ANNOTATION = "org.apache.airflow.sdk.Builder.Dag" private const val TASK_ANNOTATION = "org.apache.airflow.sdk.Builder.Task" +private val RESERVED_VIEW_NAMES = + setOf( + "depends", + "lit", + "clone", + "equals", + "finalize", + "getClass", + "hashCode", + "notify", + "notifyAll", + "toString", + "wait", + ) + private val DAG_STRUCTURAL_ATTRIBUTES = setOf("id", "to") private val TASK_STRUCTURAL_ATTRIBUTES = setOf("id") 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 04b90bc7eb43a..989ae23d24813 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 @@ -1162,6 +1162,32 @@ class BuilderTest { ) } + @Test + @DisplayName("reject a task method whose name clashes with a wiring-view member") + fun rejectTaskNameClashingWithViewMember() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + @Builder.Dag + public class TestExample { + @Builder.Task(id = "wire") public void depends() {} + + @Builder.Deps + static class Wiring implements TestExampleDeps { + void depends() {} + } + } + """, + ) + assertThat(compilation).failed() + assertThat(compilation).hadErrorContaining( + "Task method 'depends' clashes with a member of the wiring view; rename the method and keep " + + "the task id with @Builder.Task(id = \"wire\")", + ) + } + @Test @DisplayName("generate builder for dag class with varargs task parameter") fun generateBuilderForDagClassWithVarArgsTaskParameter() { From 202e9c772c3a75b0baef924c541991cb8a335b3e Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Sat, 3 Oct 2026 13:07:17 +0000 Subject: [PATCH 05/13] Java SDK: Type a wired numeric input by its parameter's own type --- .../apache/airflow/sdk/BuilderProcessor.kt | 30 ++++------------ .../org/apache/airflow/sdk/BuilderTest.kt | 35 ++++++++++++++++--- 2 files changed, 37 insertions(+), 28 deletions(-) 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 f76969b5a971f..e74b24cc90ffc 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 @@ -318,23 +318,14 @@ class BuilderProcessor : AbstractProcessor() { } /** - * Maps a data parameter's declared type to its twin-input type. Numeric - * parameters accept any numeric upstream (`Arg`, widened - * at run time); `Object`, raw `Map`, and raw `List` parameters accept any - * upstream (`Arg`, decoded loosely at run time); everything else accepts - * covariant matches of the declared type (`Arg`). + * Maps a data parameter's declared type to its wiring-view input type, + * `Arg` of the boxed type. A numeric parameter therefore takes + * only its own type, so javac rejects wiring that could lose a value, such + * as a `double` upstream into a `long` parameter. An `Object` parameter + * takes any upstream, including a `void` task's handle, whose value is null. */ - private fun inType(paramType: TypeMirror): TypeName { - val boxed = TypeName.get(paramType).boxIfPossible() - val argument = - when { - isNumeric(paramType) -> WildcardTypeName.subtypeOf(TypeName.get(Number::class.java)) - else -> WildcardTypeName.subtypeOf(boxed) - } - return ParameterizedTypeName.get(ARG_TYPE, argument) - } - - private fun isNumeric(t: TypeMirror): Boolean = t.kind in NUMERIC_KINDS || TypeName.get(t) in BOXED_NUMERICS + private fun inType(paramType: TypeMirror): TypeName = + ParameterizedTypeName.get(ARG_TYPE, WildcardTypeName.subtypeOf(TypeName.get(paramType).boxIfPossible())) private fun collectTasks(el: TypeElement): List { val declarations = mutableListOf() @@ -740,13 +731,6 @@ private val RESERVED_VIEW_NAMES = private val DAG_STRUCTURAL_ATTRIBUTES = setOf("id", "to") private val TASK_STRUCTURAL_ATTRIBUTES = setOf("id") -private val NUMERIC_KINDS = - setOf(TypeKind.BYTE, TypeKind.SHORT, TypeKind.INT, TypeKind.LONG, TypeKind.FLOAT, TypeKind.DOUBLE) - -private val BOXED_NUMERICS: Set = - setOf(TypeName.BYTE, TypeName.SHORT, TypeName.INT, TypeName.LONG, TypeName.FLOAT, TypeName.DOUBLE) - .mapTo(mutableSetOf()) { it.box() } - private fun TypeName.boxIfPossible(): TypeName = if (this == TypeName.VOID || isPrimitive) box() else this private fun ProcessingEnvironment.isType( 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 989ae23d24813..7136ef7e01aee 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 @@ -135,7 +135,6 @@ class BuilderTest { package org.apache.airflow.example; import java.lang.Integer; - import java.lang.Number; import java.lang.Void; import org.apache.airflow.sdk.Arg; import org.apache.airflow.sdk.Deps; @@ -159,7 +158,7 @@ class BuilderTest { return Refs.node(new TaskDef("t2", TestExampleBuilder.T2.class)); } - default TaskRef t3(Arg value) { + default TaskRef t3(Arg value) { return Refs.call(new TaskDef("t3", TestExampleBuilder.T3.class), value); } } @@ -317,7 +316,7 @@ class BuilderTest { public long pn() { return 1L; } @Builder.Task - public void t(String text, Object anything, List items, Integer boxed) {} + public void t(String text, Object anything, List items, Long boxed) {} @Builder.Deps static class Wiring implements TestExampleDeps { @@ -338,7 +337,6 @@ class BuilderTest { package org.apache.airflow.example; import java.lang.Long; - import java.lang.Number; import java.lang.String; import java.lang.Void; import java.util.List; @@ -373,7 +371,7 @@ class BuilderTest { } default TaskRef t(Arg text, Arg anything, - Arg> items, Arg boxed) { + Arg> items, Arg boxed) { return Refs.call(new TaskDef("t", TestExampleBuilder.T.class), text, anything, items, boxed); } } @@ -927,6 +925,33 @@ class BuilderTest { assertThat(compilation).hadErrorContaining("incompatible types") } + @Test + @DisplayName("reject wiring a numeric upstream into a narrower numeric parameter") + fun rejectLossyNumericWiring() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + @Builder.Dag + public class TestExample { + @Builder.Task + public double ratio() { return 2.7; } + + @Builder.Task + public void load(long rows) {} + + @Builder.Deps + static class Wiring implements TestExampleDeps { + void depends() { load(ratio()); } + } + } + """, + ) + assertThat(compilation).failed() + assertThat(compilation).hadErrorContaining("incompatible types") + } + @Test @DisplayName("reject more than one wiring class") fun rejectMultipleWiringClasses() { From 37a076564fc2534f2a0fe210508e0877bea1affa Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Sat, 3 Oct 2026 13:09:00 +0000 Subject: [PATCH 06/13] Java SDK: Register both the Dag and handlers of a class, unwrapping builder failures --- .../apache/airflow/sdk/BuilderProcessor.kt | 7 +- .../kotlin/org/apache/airflow/sdk/Bundle.kt | 37 +++++--- .../apache/airflow/sdk/internal/Registrar.kt | 23 +++++ .../org/apache/airflow/sdk/BundleTest.kt | 85 +++++++++++++++++++ .../airflow/sdk/internal/RegistrarTest.kt | 32 +++++++ 5 files changed, 173 insertions(+), 11 deletions(-) create mode 100644 java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/RegistrarTest.kt 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 e74b24cc90ffc..11d741a6f1be1 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 @@ -57,6 +57,7 @@ import javax.lang.model.element.VariableElement import javax.lang.model.type.TypeKind import javax.lang.model.type.TypeMirror import javax.tools.Diagnostic +import org.apache.airflow.sdk.internal.builderName as generatedBuilderName /** * @suppress @@ -132,7 +133,11 @@ class BuilderProcessor : AbstractProcessor() { runCatching { val packageName = elementUtils.getPackageOf(el).qualifiedName.toString() val declarations = collectTasks(el) - val builderName = ClassName.get(packageName, dagAnnotation(el).to.ifBlank { "${el.simpleName}Builder" }) + val builderName = + ClassName.get( + packageName, + generatedBuilderName(packageName, el.simpleName.toString(), dagAnnotation(el).to).substringAfterLast('.'), + ) val depsName = ClassName.get(packageName, "${el.simpleName}Deps") val deps = findDeps(el, depsName) if (deps != null) declarations.forEach { checkViewName(it) } 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 e26fc9608c78f..605b2b46ddff9 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 @@ -19,7 +19,10 @@ package org.apache.airflow.sdk +import org.apache.airflow.sdk.internal.builderName import org.apache.airflow.sdk.internal.registrarName +import java.lang.reflect.InvocationTargetException +import java.lang.reflect.Method /** * All [DagDef]s that this JVM process can execute. @@ -89,27 +92,31 @@ class Bundle( } /** - * Registers what an annotated class holds — the same class the annotations - * are on, so there is no second name to keep in sync. + * Registers what an annotated class holds, read from the class itself so + * there is no second name to keep in sync. * * A [Builder.Dag] class contributes the Dag its generated builder builds; a * class of [Builder.TaskHandler] methods contributes each handler, bound to - * the Dag the Python file owns. + * the Dag the Python file owns. A class can carry both. * * @param annotated A class carrying [Builder.Dag] or [Builder.TaskHandler]. * @return This bundle, for chaining. * @throws IllegalArgumentException if the class has no generated code, - * because annotation processing did not run over it. + * because annotation processing did not run over it, or if the Dag's + * wiring is invalid, such as a task the `@Builder.Deps` class did not + * call. */ fun register(annotated: Class<*>): Bundle { checkOpen() - annotated.getAnnotation(Builder.Dag::class.java)?.let { dag -> - val builderName = dag.to.ifBlank { "${annotated.simpleName}Builder" } - val builder = generated("${annotated.packageName}.$builderName", annotated, "builder") - return register(builder.getMethod("build").invoke(null) as DagDef) + val dag = annotated.getAnnotation(Builder.Dag::class.java) + if (dag != null) { + val builder = generated(builderName(annotated.packageName, annotated.simpleName, dag.to), annotated, "builder") + register(invokeGenerated(builder.getMethod("build")) as DagDef) + } + if (dag == null || annotated.declaredMethods.any { it.isAnnotationPresent(Builder.TaskHandler::class.java) }) { + val registrar = generated(registrarName(annotated.name), annotated, "registrar") + invokeGenerated(registrar.getMethod("registerInto", Bundle::class.java), this) } - val registrar = generated(registrarName(annotated.name), annotated, "registrar") - registrar.getMethod("registerInto", Bundle::class.java).invoke(null, this) return this } @@ -163,6 +170,16 @@ class Bundle( ) } + private fun invokeGenerated( + method: Method, + vararg args: Any?, + ): Any? = + try { + method.invoke(null, *args) + } catch (e: InvocationTargetException) { + throw e.cause ?: e + } + /** * Ends registration, so a `register` left below `serve` is reported as the * mistake it is rather than racing the runtime. [Server] calls it when it diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/Registrar.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/Registrar.kt index 3220ae78c02aa..30edadb3060fb 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/Registrar.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/Registrar.kt @@ -34,3 +34,26 @@ package org.apache.airflow.sdk.internal * returns it. */ fun registrarName(binaryName: String): String = "${binaryName.replace('$', '_')}Handlers" + +/** + * @suppress + * + * Names the builder generated for a `@Builder.Dag` class: a top-level class + * in the Dag class's package, named by the annotation's `to` or, when that is + * blank, `Builder`. + * + * Public so the annotation processor emits the name the runtime looks up; not + * user-facing API. + * + * @param packageName Package of the Dag class; empty for the default package. + * @param simpleName Simple name of the Dag class. + * @param to The annotation's `to` attribute. + */ +fun builderName( + packageName: String, + simpleName: String, + to: String, +): String { + val name = to.ifBlank { "${simpleName}Builder" } + return if (packageName.isEmpty()) name else "$packageName.$name" +} diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/BundleTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/BundleTest.kt index a74c062828037..d1819866bcce5 100644 --- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/BundleTest.kt +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/BundleTest.kt @@ -34,6 +34,21 @@ internal class BundleTest { /** A class of handlers that the processor generated [BundleTest_NestedHandlers] for. */ class Nested + /** A Dag class that the processor generated [WiredDagBuilder] for. */ + @Builder.Dag(id = "wired") + class WiredDag + + /** A Dag class whose generated [BrokenDagBuilder] fails, as an invalid wiring would. */ + @Builder.Dag + class BrokenDag + + /** A Dag class that also holds a task handler, so the processor generated both. */ + @Builder.Dag(id = "mixed") + class MixedDag { + @Builder.TaskHandler(dag = "etl", task = "score") + fun score() = Unit + } + @Test @DisplayName("Should index dags by dagId") fun shouldIndexDagsByDagId() { @@ -162,6 +177,41 @@ internal class BundleTest { ) } + @Test + @DisplayName("Should register the Dag a @Builder.Dag class's generated builder builds") + fun shouldRegisterDagFromBuilderClass() { + val bundle = Bundle().register(WiredDag::class.java) + + Assertions.assertEquals(listOf("wired"), bundle.dags.keys.toList()) + Assertions.assertEquals(emptySet(), bundle.taskHandlers.keys) + } + + @Test + @DisplayName("Should rethrow the failure of a generated builder rather than its reflection wrapper") + fun shouldUnwrapBuilderFailure() { + val error = + Assertions.assertThrows(IllegalArgumentException::class.java) { + Bundle().register(BrokenDag::class.java) + } + + Assertions.assertEquals("wiring failed", error.message) + } + + @Test + @DisplayName("Should register both the Dag and the handlers of a class that carries both") + fun shouldRegisterDagAndHandlersOfOneClass() { + val bundle = Bundle().register(MixedDag::class.java) + + Assertions.assertEquals(listOf("mixed"), bundle.dags.keys.toList()) + Assertions.assertEquals( + listOf("score"), + bundle.taskHandlers + .getValue("etl") + .tasks.keys + .toList(), + ) + } + @Test @DisplayName("Should find the registrar generated for a nested handler class") fun shouldFindRegistrarOfNestedHandlerClass() { @@ -280,6 +330,41 @@ class NoopBundleTask : Task { ) = Unit } +/** Stands in for the builder the annotation processor generates for [BundleTest.WiredDag]. */ +class WiredDagBuilder { + companion object { + @JvmStatic + fun build() = DagDef("wired") + } +} + +/** Stands in for a generated builder whose wiring is invalid. */ +class BrokenDagBuilder { + companion object { + @JvmStatic + fun build(): DagDef = throw IllegalArgumentException("wiring failed") + } +} + +/** Stands in for the builder the annotation processor generates for [BundleTest.MixedDag]. */ +class MixedDagBuilder { + companion object { + @JvmStatic + fun build() = DagDef("mixed") + } +} + +/** Stands in for the registrar the annotation processor generates for [BundleTest.MixedDag]. */ +@Suppress("ktlint:standard:class-naming", "ClassName") +class BundleTest_MixedDagHandlers { + companion object { + @JvmStatic + fun registerInto(bundle: Bundle) { + bundle.register("etl", "score", NoopBundleTask::class.java) + } + } +} + /** * Stands in for the registrar the annotation processor generates beside * [BundleTest.Nested], to pin the name [Bundle.register] looks up. diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/RegistrarTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/RegistrarTest.kt new file mode 100644 index 0000000000000..ad32b92e0b9c0 --- /dev/null +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/internal/RegistrarTest.kt @@ -0,0 +1,32 @@ +/* + * 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.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.DisplayName +import org.junit.jupiter.api.Test + +internal class RegistrarTest { + @Test + @DisplayName("Should name the builder of a Dag class in the default package with no leading dot") + fun shouldNameBuilderInDefaultPackage() { + assertEquals("FooBuilder", builderName("", "Foo", "")) + } +} From 14c180a82785e7bbb66f46dcadc9ec4dacbe9334 Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Sat, 3 Oct 2026 13:09:33 +0000 Subject: [PATCH 07/13] Java SDK: Reject a raw null wiring argument and test the recorder's guards --- .../org/apache/airflow/sdk/internal/Refs.kt | 16 +++++-- .../apache/airflow/sdk/internal/RefsTest.kt | 43 +++++++++++++++++++ 2 files changed, 55 insertions(+), 4 deletions(-) diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/Refs.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/Refs.kt index e875f9be98cc4..39021f03b4eaa 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/Refs.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/Refs.kt @@ -88,26 +88,34 @@ object Refs { * * @return The handle representing this task, memoized by [TaskDef.id] so a result * held in a local and reused refers to one node. + * @throws IllegalArgumentException if an argument is a raw Java `null` + * rather than `lit(null)`. */ @JvmStatic @Suppress("UNCHECKED_CAST", "SpreadOperator") fun call( def: TaskDef, - vararg args: Arg<*>, + vararg args: Arg<*>?, ): TaskRef { + val inputs = + args.mapIndexed { i, arg -> + requireNotNull(arg) { + "Argument ${i + 1} of task '${def.id}' is null; wrap a null constant as lit(null)" + } + } val active = checkNotNull(recording.get()) { "Task '${def.id}' was wired outside a @Builder.Deps class; the wiring view's methods " + "only record while the generated builder is running depends()" } active.byTaskId[def.id]?.let { existing -> - require(args.isEmpty()) { + require(inputs.isEmpty()) { "Task '${def.id}' is wired more than once with arguments; call it once and reuse the handle it returned" } return existing as TaskRef } - args.filterIsInstance>().forEach { def.dependsOn(it.def) } - def.inputs += args + inputs.filterIsInstance>().forEach { def.dependsOn(it.def) } + def.inputs += inputs active.dag.addTask(def) return TaskRef(def).also { active.byTaskId[def.id] = it } } 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 f941803c8708e..ffb77fa7b40d8 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 @@ -112,4 +112,47 @@ internal class RefsTest { error.message, ) } + + @Test + @DisplayName("Should reject a raw null argument, pointing to lit(null)") + fun shouldRejectRawNullArgument() { + val error = + assertThrows(IllegalArgumentException::class.java) { + Refs.record(DagDef("d"), listOf("t")) { + Refs.call(TaskDef("t", NoopRefTask::class.java), Arg.lit(1), null) + } + } + + assertEquals("Argument 2 of task 't' is null; wrap a null constant as lit(null)", error.message) + } + + @Test + @DisplayName("Should reject a task wired a second time with arguments") + fun shouldRejectTaskWiredTwiceWithArguments() { + val error = + assertThrows(IllegalArgumentException::class.java) { + Refs.record(DagDef("d"), listOf("t")) { + Refs.node(TaskDef("t", NoopRefTask::class.java)) + Refs.call(TaskDef("t", NoopRefTask::class.java), Arg.lit(1)) + } + } + + assertEquals( + "Task 't' is wired more than once with arguments; call it once and reuse the handle it returned", + error.message, + ) + } + + @Test + @DisplayName("Should refuse to record a Dag while another is being recorded") + fun shouldRefuseNestedRecording() { + val error = + assertThrows(IllegalStateException::class.java) { + Refs.record(DagDef("outer"), emptyList()) { + Refs.record(DagDef("inner"), emptyList()) {} + } + } + + assertEquals("Dag wiring is already being recorded on this thread", error.message) + } } From 4dadc5dc96f37ba7041cf4a715053a1acd1b9116 Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Sat, 3 Oct 2026 13:10:06 +0000 Subject: [PATCH 08/13] Java SDK: Reword the native-Dag docs without em-dashes --- .../authoring-and-scheduling/language-sdks/java.rst | 4 ++-- .../kotlin/org/apache/airflow/sdk/BuilderProcessor.kt | 2 +- java-sdk/sdk/build.gradle.kts | 11 ++++++----- .../src/main/kotlin/org/apache/airflow/sdk/Deps.kt | 6 +++--- 4 files changed, 12 insertions(+), 11 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 e5edd6b8b9a0c..1c9d01d51a8cc 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst @@ -581,7 +581,7 @@ Wiring the graph with ``@Builder.Deps`` ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ For a Dag written with annotations, the graph is declared by a nested ``@Builder.Deps`` class. The -annotation processor generates a ``Deps`` interface — the *wiring view* — with one method +annotation processor generates a ``Deps`` interface, the *wiring view*, with one method per ``@Builder.Task`` method: the injected ``Client`` and ``Context`` parameters are dropped, each data parameter becomes an ``Arg``, and the return value becomes a ``TaskRef``. Calling a view method registers its task, and passing the handle one returned into another feeds the upstream's @@ -631,7 +631,7 @@ generated view, with a no-argument ``depends()`` method: } Every ``@Builder.Task`` method must be called in the wiring class; a task the wiring missed fails at -Dag-parse time. ``lit(...)`` wires an inline constant where no upstream feeds a parameter — a bare +Dag-parse time. ``lit(...)`` wires an inline constant where no upstream feeds a parameter. A bare ``double`` cannot be an ``Arg``, so a constant is wrapped. A view method that takes no arguments returns the same handle every time, so it names one node wherever it appears; one that takes arguments is called once, and the wiring fails if it is called again with arguments, so hold its 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 11d741a6f1be1..47a0bc5182e6a 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 @@ -74,7 +74,7 @@ import org.apache.airflow.sdk.internal.builderName as generatedBuilderName * method (implementing [Task]), and a static `build()` that constructs the * [DagDef], lowers every explicitly-written `@Builder.Dag` attribute into a * `DagDef.config` call, then runs the class's [Builder.Deps] class and - * verifies it registered every task — or, when the class declares none, + * verifies it registered every task. When the class declares none, it * registers every task with no Java-side edges. * - A `*Deps` wiring-view interface (only when a [Builder.Deps] class exists) * whose methods mirror the task methods: injectable parameters ([Client], diff --git a/java-sdk/sdk/build.gradle.kts b/java-sdk/sdk/build.gradle.kts index ba7b46fab0df4..9c2d552add332 100644 --- a/java-sdk/sdk/build.gradle.kts +++ b/java-sdk/sdk/build.gradle.kts @@ -504,8 +504,9 @@ abstract class GenerateDagDslTask : DefaultTask() { | * Container for the annotation-based Dag-authoring API. | * | * Annotating a class with [Dag] generates a `Builder` whose static - | * `build()` returns the [DagDef] to add to a [Bundle], and — when the class - | * declares a [Deps] class — a `Deps` wiring view for it to implement. + | * `build()` returns the [DagDef] to add to a [Bundle]. When the class + | * declares a [Deps] class, it also generates a `Deps` wiring view + | * for that class to implement. | * | * Example: | * @@ -521,9 +522,9 @@ abstract class GenerateDagDslTask : DefaultTask() { | * } | * ``` | * - | * A task method's data parameters — everything other than the injected - | * [Client] and [Context] — receive, by position, the inputs the [Deps] - | * class wired. For a task the Python Dag file declares with `@task.stub`, + | * A task method's data parameters, meaning every parameter other than the + | * injected [Client] and [Context], receive by position the inputs the + | * [Deps] class wired. For a task the Python Dag file declares with `@task.stub`, | * the arguments bound at that call site take their place. Keyword | * arguments bind by name instead through a single [TaskInput] parameter. | */ diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Deps.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Deps.kt index 3b0fda4d0f2b2..e70df9f853480 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Deps.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Deps.kt @@ -89,9 +89,9 @@ interface Deps { } /** - * Wraps an inline constant as a task argument — `transform(extract(), - * lit(0.9))` — passed to the task as a constant, creating no dependency - * edge. + * Wraps an inline constant as a task argument, as in + * `transform(extract(), lit(0.9))`. It is passed to the task as a constant + * and creates no dependency edge. * * @param value Constant to bind; may be null for a nullable parameter. */ From 0dc60cd77a3a5b6de445248e9b83db5d4b1b94d5 Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Sat, 3 Oct 2026 13:10:32 +0000 Subject: [PATCH 09/13] Java SDK: Drop Bundle tests that repeat existing coverage --- .../org/apache/airflow/sdk/BundleTest.kt | 28 ------------------- 1 file changed, 28 deletions(-) diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/BundleTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/BundleTest.kt index d1819866bcce5..4c645aab89fb1 100644 --- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/BundleTest.kt +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/BundleTest.kt @@ -149,34 +149,6 @@ internal class BundleTest { Assertions.assertEquals(emptySet(), bundle.dags.keys) } - @Test - @DisplayName("Should register Dags one at a time and reject a duplicate ID") - fun shouldRegisterDagsIncrementally() { - val bundle = Bundle().register(DagDef("a")).register(DagDef("b")) - - Assertions.assertEquals(setOf("a", "b"), bundle.dags.keys) - val error = - Assertions.assertThrows(IllegalArgumentException::class.java) { bundle.register(DagDef("a")) } - Assertions.assertEquals("Dags in bundle have duplicate ID: a", error.message) - } - - @Test - @DisplayName("Should create the Dag on first use when registering a stub-backed handler") - fun shouldRegisterHandlerAgainstPythonOwnedDag() { - val bundle = - Bundle() - .register("etl", "score", NoopBundleTask::class.java) - .register("etl", "report", NoopBundleTask::class.java) - - Assertions.assertEquals(setOf("etl"), bundle.taskHandlers.keys) - Assertions.assertEquals( - setOf("score", "report"), - bundle.taskHandlers - .getValue("etl") - .tasks.keys, - ) - } - @Test @DisplayName("Should register the Dag a @Builder.Dag class's generated builder builds") fun shouldRegisterDagFromBuilderClass() { From 23be1d072a6e06ed564e4b2454799a5890873d70 Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Mon, 5 Oct 2026 09:42:45 +0000 Subject: [PATCH 10/13] Java SDK: Require a @Builder.Deps class on every @Builder.Dag class --- .../language-sdks/java.rst | 5 +- .../apache/airflow/sdk/BuilderProcessor.kt | 50 ++++---- .../org/apache/airflow/sdk/BuilderTest.kt | 118 +++++++++++++++--- java-sdk/sdk/build.gradle.kts | 6 +- 4 files changed, 125 insertions(+), 54 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 1c9d01d51a8cc..649a1ebf5db85 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst @@ -641,8 +641,9 @@ handle in a local and reuse that. the wiring class ``Flow`` is inherited by simple name, so it needs no import and never collides with ``java.util.concurrent.Flow``. -The wiring class is what makes the Dag native. A Dag class without one registers every task with -no Java-side edges, which is the task-handler shape rather than a Dag Java owns. +Every ``@Builder.Dag`` class declares a wiring class, because the graph is what the Dag owns. A +class that supplies only task bodies, for a Dag a Python file declares, carries +``@Builder.TaskHandler`` instead and contributes no Dag. .. note:: 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 47a0bc5182e6a..711ac23d5aadf 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 @@ -74,10 +74,9 @@ import org.apache.airflow.sdk.internal.builderName as generatedBuilderName * method (implementing [Task]), and a static `build()` that constructs the * [DagDef], lowers every explicitly-written `@Builder.Dag` attribute into a * `DagDef.config` call, then runs the class's [Builder.Deps] class and - * verifies it registered every task. When the class declares none, it - * registers every task with no Java-side edges. - * - A `*Deps` wiring-view interface (only when a [Builder.Deps] class exists) - * whose methods mirror the task methods: injectable parameters ([Client], + * verifies it registered every task. + * - A `*Deps` wiring-view interface whose methods mirror the task methods: + * injectable parameters ([Client], * [Context]) are dropped, data parameters become [Arg]-typed inputs, and the * return value becomes a [TaskRef]. Calling one registers the task with its * explicitly-written `@Builder.Task` attributes lowered into `TaskDef.config` @@ -140,14 +139,12 @@ class BuilderProcessor : AbstractProcessor() { ) val depsName = ClassName.get(packageName, "${el.simpleName}Deps") val deps = findDeps(el, depsName) - if (deps != null) declarations.forEach { checkViewName(it) } + declarations.forEach { checkViewName(it) } JavaFile .builder(packageName, buildBuilder(el, declarations, deps, builderName)) .build() .writeTo(filer) - if (deps != null) { - JavaFile.builder(packageName, buildDeps(el, declarations, builderName, depsName)).build().writeTo(filer) - } + JavaFile.builder(packageName, buildDeps(el, declarations, builderName, depsName)).build().writeTo(filer) }.onFailure { e -> messager.printMessage( Diagnostic.Kind.ERROR, @@ -212,7 +209,7 @@ class BuilderProcessor : AbstractProcessor() { private fun buildBuilder( el: TypeElement, declarations: List, - deps: TypeElement?, + deps: TypeElement, builderName: ClassName, ): TypeSpec { val ann = dagAnnotation(el) @@ -231,22 +228,13 @@ class BuilderProcessor : AbstractProcessor() { explicitConfig(el, DAG_ANNOTATION, DAG_STRUCTURAL_ATTRIBUTES, SchemaFields.DAG).forEach { (key, value) -> buildMethod.addStatement($$"dag.config($S, $L)", key, value) } - if (deps != null) { - buildMethod.addStatement( - $$"return $T.record(dag, $T.of($L), new $T()::depends)", - REFS_TYPE, - ClassName.get(List::class.java), - CodeBlock.join(declarations.map { CodeBlock.of($$"$S", it.id) }, ", "), - ClassName.get(deps), - ) - } else { - // No wiring class: register every task with no Java-side edges, which - // is the task-handler shape rather than a Dag Java owns. - declarations.forEach { decl -> - buildMethod.addStatement($$"dag.addTask($L)", taskDefCode(decl, CodeBlock.of($$"$L", decl.className))) - } - buildMethod.addStatement("return dag") - } + buildMethod.addStatement( + $$"return $T.record(dag, $T.of($L), new $T()::depends)", + REFS_TYPE, + ClassName.get(List::class.java), + CodeBlock.join(declarations.map { CodeBlock.of($$"$S", it.id) }, ", "), + ClassName.get(deps), + ) builderClass.addMethod(buildMethod.build()) declarations.forEach { builderClass.addType(buildTask(it, el)) } @@ -346,8 +334,8 @@ class BuilderProcessor : AbstractProcessor() { } /** - * Finds and validates the class's `@Builder.Deps` wiring class. It is - * optional: without one, every task registers with no Java-side edges. + * Finds and validates the class's `@Builder.Deps` wiring class, which + * declares the Dag's task graph and is what makes it a Dag Java owns. * * The generated builder runs `new Wiring()::depends`, so everything that * expression needs is checked here, where the error can name the class. @@ -355,12 +343,16 @@ class BuilderProcessor : AbstractProcessor() { private fun findDeps( el: TypeElement, view: ClassName, - ): TypeElement? { + ): TypeElement { val classes = el.enclosedElements .filterIsInstance() .filter { it.getAnnotation(Builder.Deps::class.java) != null } - if (classes.isEmpty()) return null + require(classes.isNotEmpty()) { + "Dag class ${el.simpleName} must declare a @Builder.Deps class implementing ${view.simpleName()} " + + "to declare its task graph; a class of task bodies for a Dag the Python file owns carries " + + "@Builder.TaskHandler instead" + } val deps = classes.singleOrNull() ?: throw IllegalArgumentException( 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 7136ef7e01aee..9a71a4e688f59 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 @@ -180,6 +180,11 @@ class BuilderTest { public class TestExample { @Builder.Task public void t(long first, Client client, String second, Context ctx, Integer third) {} + + @Builder.Deps + static class Wiring implements TestExampleDeps { + void depends() { t(lit(1L), lit("second"), lit(3)); } + } } """, ) @@ -197,18 +202,18 @@ class BuilderTest { import java.lang.Long; import java.lang.Override; import java.lang.String; + 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.Task; - import org.apache.airflow.sdk.TaskDef; + 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"); - dag.addTask(new TaskDef("t", T.class)); - return dag; + return Refs.record(dag, List.of("t"), new TestExample.Wiring()::depends); } public static final class T implements Task { @@ -240,6 +245,11 @@ class BuilderTest { public class TestExample { @Builder.Task public void t(boolean flag, float fraction, Double boxed, List tags, Map raw) {} + + @Builder.Deps + static class Wiring implements TestExampleDeps { + void depends() { t(lit(true), lit(1f), lit(2.0), lit(List.of("a")), lit(Map.of())); } + } } """, ) @@ -264,15 +274,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.TaskDef; + 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"); - dag.addTask(new TaskDef("t", T.class)); - return dag; + return Refs.record(dag, List.of("t"), new TestExample.Wiring()::depends); } public static final class T implements Task { @@ -482,6 +491,11 @@ class BuilderTest { public class TestExample { @Builder.Task public void t(String args, int other) {} + + @Builder.Deps + static class Wiring implements TestExampleDeps { + void depends() { t(lit("a"), lit(1)); } + } } """, ) @@ -498,18 +512,18 @@ class BuilderTest { import java.lang.Integer; import java.lang.Override; import java.lang.String; + 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.Task; - import org.apache.airflow.sdk.TaskDef; + 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"); - dag.addTask(new TaskDef("t", T.class)); - return dag; + return Refs.record(dag, List.of("t"), new TestExample.Wiring()::depends); } public static final class T implements Task { @Override @@ -532,6 +546,7 @@ class BuilderTest { compile( """ package org.apache.airflow.example; + import java.util.List; import org.apache.airflow.sdk.Builder; import org.apache.airflow.sdk.TaskInput; @Builder.Dag @@ -545,6 +560,11 @@ class BuilderTest { @Builder.Task public void named(ScoreInput context) {} + + @Builder.Deps + static class Wiring implements TestExampleDeps { + void depends() { flat(lit("a"), lit(1)); named(lit(null)); } + } } """, ) @@ -561,20 +581,19 @@ class BuilderTest { import java.lang.Integer; import java.lang.Override; import java.lang.String; + 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.Task; - import org.apache.airflow.sdk.TaskDef; import org.apache.airflow.sdk.internal.ArgValues; + 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"); - dag.addTask(new TaskDef("flat", Flat.class)); - dag.addTask(new TaskDef("named", Named.class)); - return dag; + return Refs.record(dag, List.of("flat", "named"), new TestExample.Wiring()::depends); } public static final class Flat implements Task { @@ -774,7 +793,13 @@ class BuilderTest { """ package org.apache.airflow.example; import org.apache.airflow.sdk.Builder; - @Builder.Dag(id = "foo") public class TestExample {} + @Builder.Dag(id = "foo") + public class TestExample { + @Builder.Deps + static class Wiring implements TestExampleDeps { + void depends() {} + } + } """, ) assertThat(compilation) @@ -783,9 +808,14 @@ class BuilderTest { "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.Refs; public final class TestExampleBuilder { - public static DagDef build() { var dag = new DagDef("foo"); return dag; } + public static DagDef build() { + var dag = new DagDef("foo"); + return Refs.record(dag, List.of(), new TestExample.Wiring()::depends); + } } """, ) @@ -799,7 +829,13 @@ class BuilderTest { """ package org.apache.airflow.example; import org.apache.airflow.sdk.Builder; - @Builder.Dag(to = "Foo") public class TestExample {} + @Builder.Dag(to = "Foo") + public class TestExample { + @Builder.Deps + static class Wiring implements TestExampleDeps { + void depends() {} + } + } """, ) assertThat(compilation) @@ -808,9 +844,14 @@ class BuilderTest { "org.apache.airflow.example.Foo", """ package org.apache.airflow.example; + import java.util.List; import org.apache.airflow.sdk.DagDef; + import org.apache.airflow.sdk.internal.Refs; public final class Foo { - public static DagDef build() { var dag = new DagDef("TestExample"); return dag; } + public static DagDef build() { + var dag = new DagDef("TestExample"); + return Refs.record(dag, List.of(), new TestExample.Wiring()::depends); + } } """, ) @@ -952,6 +993,28 @@ class BuilderTest { assertThat(compilation).hadErrorContaining("incompatible types") } + @Test + @DisplayName("reject a dag class with no wiring class") + fun rejectDagClassWithoutWiringClass() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + @Builder.Dag + public class TestExample { + @Builder.Task public void t1() {} + } + """, + ) + assertThat(compilation).failed() + assertThat(compilation).hadErrorContaining( + "Dag class TestExample must declare a @Builder.Deps class implementing TestExampleDeps " + + "to declare its task graph; a class of task bodies for a Dag the Python file owns carries " + + "@Builder.TaskHandler instead", + ) + } + @Test @DisplayName("reject more than one wiring class") fun rejectMultipleWiringClasses() { @@ -1264,6 +1327,11 @@ class BuilderTest { @Builder.Dag public class TestExample { @Builder.Task(retryDelay = "5 minutes") public void t() {} + + @Builder.Deps + static class Wiring implements TestExampleDeps { + void depends() { t(); } + } } """, ) @@ -1284,6 +1352,11 @@ class BuilderTest { @Builder.Dag(tags = {"say \"hi\"", "back\\slash"}) public class TestExample { @Builder.Task public void t() {} + + @Builder.Deps + static class Wiring implements TestExampleDeps { + void depends() { t(); } + } } """, ) @@ -1348,6 +1421,11 @@ class BuilderTest { @Builder.Task public double score(Client client, ScoreInput input) { return input.threshold; } + + @Builder.Deps + static class Wiring implements TestExampleDeps { + void depends() { score(lit(null)); } + } } """, ) @@ -1362,18 +1440,18 @@ class BuilderTest { 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.Task; - import org.apache.airflow.sdk.TaskDef; import org.apache.airflow.sdk.internal.ArgValues; + import org.apache.airflow.sdk.internal.Refs; public final class TestExampleBuilder { public static DagDef build() { var dag = new DagDef("TestExample"); - dag.addTask(new TaskDef("score", Score.class)); - return dag; + return Refs.record(dag, List.of("score"), new TestExample.Wiring()::depends); } public static final class Score implements Task { diff --git a/java-sdk/sdk/build.gradle.kts b/java-sdk/sdk/build.gradle.kts index 9c2d552add332..d19e8f771dda1 100644 --- a/java-sdk/sdk/build.gradle.kts +++ b/java-sdk/sdk/build.gradle.kts @@ -612,9 +612,9 @@ abstract class GenerateDagDslTask : DefaultTask() { | * } | * ``` | * - | * What makes the Dag native: a Dag class without one registers every - | * task with no Java-side edges, which is the task-handler shape rather - | * than a Dag Java owns. + | * Every [Dag] class declares one, because the graph is what the Dag + | * owns. A class that supplies only task bodies, for a Dag a Python + | * file declares, carries [TaskHandler] instead. | */ | @Target(AnnotationTarget.CLASS) | @MustBeDocumented From c394f9d7589929532676ed201e9c360fc79446f8 Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Mon, 5 Oct 2026 11:05:17 +0000 Subject: [PATCH 11/13] Java SDK: Back the k8s test's stub tasks with task handlers --- .../airflow/k8sexample/CombinedExample.java | 15 +++++++-------- .../airflow/k8sexample/K8sBundleBuilder.java | 17 ++++++----------- 2 files changed, 13 insertions(+), 19 deletions(-) diff --git a/kubernetes-tests/lang_sdk/java_example/src/java/org/apache/airflow/k8sexample/CombinedExample.java b/kubernetes-tests/lang_sdk/java_example/src/java/org/apache/airflow/k8sexample/CombinedExample.java index 63aaa5b6f051d..09b7b67c8e6ac 100644 --- a/kubernetes-tests/lang_sdk/java_example/src/java/org/apache/airflow/k8sexample/CombinedExample.java +++ b/kubernetes-tests/lang_sdk/java_example/src/java/org/apache/airflow/k8sexample/CombinedExample.java @@ -24,22 +24,21 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; -// Java half of the KubernetesExecutor lang-SDK system test bundle. Registers the -// Java tasks of the shared "lang_sdk_combined" Dag; the Go tasks of the same -// dag_id live in ../go_example and the Python stub Dag in ../dags. The -// coordinator locates this jar by dag_id, so only the Java tasks are registered -// here. -@Builder.Dag(id = "lang_sdk_combined") +// Java half of the KubernetesExecutor lang-SDK system test bundle. Supplies the +// bodies of the "lang_sdk_combined" Dag's Java stub tasks; the Dag itself is +// declared by the Python file in ../dags, and the Go tasks of the same dag_id +// live in ../go_example. The coordinator locates this jar by dag_id, so only the +// Java tasks are registered here. public class CombinedExample { private static final Logger logger = LoggerFactory.getLogger(CombinedExample.class); - @Builder.Task(id = "java_extract") + @Builder.TaskHandler(dag = "lang_sdk_combined", task = "java_extract") public long extract(Client client) { logger.info("java_extract running"); return new Date().getTime(); } - @Builder.Task(id = "java_transform") + @Builder.TaskHandler(dag = "lang_sdk_combined", task = "java_transform") public void transform(Client client) { var variable = client.getVariable("my_variable"); logger.info("java_transform obtained variable {}", variable); diff --git a/kubernetes-tests/lang_sdk/java_example/src/java/org/apache/airflow/k8sexample/K8sBundleBuilder.java b/kubernetes-tests/lang_sdk/java_example/src/java/org/apache/airflow/k8sexample/K8sBundleBuilder.java index c3b511b898ced..dd5b53b79834f 100644 --- a/kubernetes-tests/lang_sdk/java_example/src/java/org/apache/airflow/k8sexample/K8sBundleBuilder.java +++ b/kubernetes-tests/lang_sdk/java_example/src/java/org/apache/airflow/k8sexample/K8sBundleBuilder.java @@ -19,21 +19,16 @@ package org.apache.airflow.k8sexample; -import java.util.List; import org.apache.airflow.sdk.*; -import org.jetbrains.annotations.NotNull; -// CombinedExampleBuilder is generated by the annotation processor from the -// @Builder.Dag class CombinedExample. -public class K8sBundleBuilder implements BundleBuilder { - @NotNull - @Override - public Iterable getDags() { - return List.of(CombinedExampleBuilder.build()); +// The bundle holds task handlers only: the "lang_sdk_combined" Dag is the Python +// file's, so nothing here declares a Dag. +public class K8sBundleBuilder { + public static Bundle build() { + return new Bundle().register(CombinedExample.class); } public static void main(String[] args) { - var bundle = new K8sBundleBuilder().build(); - Server.create(args).serve(bundle); + Server.create(args).serve(build()); } } From 162a3ffaa4e7d7fc24c3a507351928daa804dbbc Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Mon, 5 Oct 2026 11:29:15 +0000 Subject: [PATCH 12/13] Java SDK: Reject overloaded task and task-handler methods --- .../apache/airflow/sdk/BuilderProcessor.kt | 11 ++++ .../org/apache/airflow/sdk/BuilderTest.kt | 50 +++++++++++++++++++ 2 files changed, 61 insertions(+) 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 711ac23d5aadf..93e185a614eee 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 @@ -183,6 +183,7 @@ class BuilderProcessor : AbstractProcessor() { .addModifiers(Modifier.PUBLIC, Modifier.STATIC) .addParameter(BUNDLE_TYPE, "bundle") + val names = mutableSetOf() for (inner in el.enclosedElements) { if (inner !is ExecutableElement) continue val handler = inner.getAnnotation(Builder.TaskHandler::class.java) ?: continue @@ -193,6 +194,11 @@ class BuilderProcessor : AbstractProcessor() { "@Builder.TaskHandler on '${inner.simpleName}' must name the Dag the Python file declares" } val decl = TaskDeclaration(inner, handler.task.ifBlank { inner.simpleName.toString() }, collectDataParams(inner)) + require(names.add(inner.simpleName.toString())) { + "Class ${el.simpleName} overloads task-handler method '${inner.simpleName}'; a method's name is " + + "the name of its generated task class, so rename one and keep its task id with " + + "@Builder.TaskHandler(task = \"${decl.id}\")" + } registrar.addType(buildTask(decl, el)) registerInto.addStatement( $$"bundle.register($S, $S, $L.class)", @@ -328,6 +334,11 @@ class BuilderProcessor : AbstractProcessor() { if (inner.isVarArgs) throw IllegalArgumentException("Cannot create task from vararg function ${inner.simpleName}") val id = ann.id.ifBlank { inner.simpleName.toString() } require(declarations.none { it.id == id }) { "Tasks in Dag have duplicate ID: $id" } + require(declarations.none { it.method.simpleName.contentEquals(inner.simpleName) }) { + "Dag 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 = \"$id\")" + } declarations += TaskDeclaration(inner, id, collectDataParams(inner)) } return declarations 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 9a71a4e688f59..572fde64e9c8b 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 @@ -1316,6 +1316,56 @@ class BuilderTest { assertThat(compilation).hadErrorContaining("Tasks in Dag have duplicate ID: x") } + @Test + @DisplayName("reject overloaded task methods, whatever their parameters") + fun rejectOverloadedTaskMethods() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + @Builder.Dag + public class TestExample { + @Builder.Task(id = "a") public void extract() {} + @Builder.Task(id = "b") public void extract(String text) {} + + @Builder.Deps + static class Wiring implements TestExampleDeps { + void depends() {} + } + } + """, + ) + assertThat(compilation).failed() + assertThat(compilation).hadErrorContaining( + "Dag class TestExample overloads task method 'extract'; 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 = \"b\")", + ) + } + + @Test + @DisplayName("reject overloaded task-handler methods") + fun rejectOverloadedTaskHandlerMethods() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + public class TestExample { + @Builder.TaskHandler(dag = "etl", task = "a") public void score() {} + @Builder.TaskHandler(dag = "etl", task = "b") public void score(String text) {} + } + """, + ) + assertThat(compilation).failed() + assertThat(compilation).hadErrorContaining( + "Class TestExample overloads task-handler method 'score'; a method's name is the name of its " + + "generated task class, so rename one and keep its task id with " + + "@Builder.TaskHandler(task = \"b\")", + ) + } + @Test @DisplayName("reject a duration attribute that is not ISO-8601") fun rejectInvalidDurationAttribute() { From ff8ec8f1a91c00f50e8c1a3f629052ccb2c928d6 Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Tue, 6 Oct 2026 01:54:58 +0000 Subject: [PATCH 13/13] Java SDK: Fix the Dag examples a required wiring class invalidates --- java-sdk/adr/0002-native-dag-interface.md | 7 ++++--- java-sdk/sdk/build.gradle.kts | 10 +++++++--- 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/java-sdk/adr/0002-native-dag-interface.md b/java-sdk/adr/0002-native-dag-interface.md index b3ad667e57823..5ab2ffdaff326 100644 --- a/java-sdk/adr/0002-native-dag-interface.md +++ b/java-sdk/adr/0002-native-dag-interface.md @@ -65,7 +65,8 @@ public class EtlPipeline { // extends nothing of ours; your own base class stays public void audit(Client client) { /* side effect only, no data in or out */ } @Builder.Task(id = "notify") - public void notify(Client client) { /* side effect only, no data in or out */ } + // "notify" is Object.notify, so the method takes another name and keeps the task id. + public void alert(Client client) { /* side effect only, no data in or out */ } @Builder.Deps static class Wiring implements EtlPipelineDeps { @@ -77,7 +78,7 @@ public class EtlPipeline { // extends nothing of ours; your own base class stays // non-TaskFlow (ordering-only) edges: sequence with no data flowing rows.then(audit()); // extract >> audit - Flow.of(loaded, audit()).then(notify()); // [load, audit] >> notify + Flow.of(loaded, audit()).then(alert()); // [load, audit] >> notify } } @@ -101,7 +102,7 @@ interface EtlPipelineDeps extends Deps { // Deps: the shared base that nests Flo default TaskRef transform(Arg extracted, Arg threshold) { return Flow.call("transform", extracted, threshold); } default TaskRef load(Arg transformed) { return Flow.call("load", transformed); } default TaskRef audit() { return Flow.node("audit"); } - default TaskRef notify() { return Flow.node("notify"); } + default TaskRef alert() { return Flow.node("notify"); } } ``` diff --git a/java-sdk/sdk/build.gradle.kts b/java-sdk/sdk/build.gradle.kts index d19e8f771dda1..f14358b5ce959 100644 --- a/java-sdk/sdk/build.gradle.kts +++ b/java-sdk/sdk/build.gradle.kts @@ -504,9 +504,8 @@ abstract class GenerateDagDslTask : DefaultTask() { | * Container for the annotation-based Dag-authoring API. | * | * Annotating a class with [Dag] generates a `Builder` whose static - | * `build()` returns the [DagDef] to add to a [Bundle]. When the class - | * declares a [Deps] class, it also generates a `Deps` wiring view - | * for that class to implement. + | * `build()` returns the [DagDef] to add to a [Bundle], and a `Deps` + | * wiring view for the class's [Deps] class to implement. | * | * Example: | * @@ -519,6 +518,11 @@ abstract class GenerateDagDslTask : DefaultTask() { | * | * @Builder.Task(id = "transform") | * public long transform(Client client, long extracted) { ... } + | * + | * @Builder.Deps + | * static class Wiring implements MyPipelineDeps { + | * void depends() { transform(extract()); } + | * } | * } | * ``` | *