From 9eeda412da4d466aa7c67f35bb0adc0701056195 Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Fri, 2 Oct 2026 15:19:36 +0000 Subject: [PATCH 01/10] Java SDK: Record the class that declares each Dag and describe it on request DagDef remembers the outermost class that built it, and a generated @Builder.Dag builder names the annotated class instead of itself. Server accepts --describe-sources to write each Java-declared Dag's declaring class and return without connecting, for the Gradle plugin. --- .../apache/airflow/sdk/BuilderProcessor.kt | 10 +- .../org/apache/airflow/sdk/BuilderTest.kt | 81 +++++++++-- .../kotlin/org/apache/airflow/sdk/DagDef.kt | 4 + .../kotlin/org/apache/airflow/sdk/Server.kt | 58 ++++++-- .../apache/airflow/sdk/internal/DagSource.kt | 66 +++++++++ .../airflow/example/DagSourceFixtures.kt | 51 +++++++ .../org/apache/airflow/sdk/DagSourceTest.kt | 132 ++++++++++++++++++ 7 files changed, 379 insertions(+), 23 deletions(-) create mode 100644 java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/DagSource.kt create mode 100644 java-sdk/sdk/src/test/kotlin/org/apache/airflow/example/DagSourceFixtures.kt create mode 100644 java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/DagSourceTest.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 cc0d9c594f89e..145ee4714c76a 100644 --- a/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt +++ b/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt @@ -30,6 +30,7 @@ import com.squareup.javapoet.TypeName import com.squareup.javapoet.TypeSpec import com.squareup.javapoet.WildcardTypeName import org.apache.airflow.sdk.internal.ArgValues +import org.apache.airflow.sdk.internal.DagSource import org.apache.airflow.sdk.internal.Field import org.apache.airflow.sdk.internal.FieldType import org.apache.airflow.sdk.internal.GROUP_ID @@ -249,7 +250,13 @@ class BuilderProcessor : AbstractProcessor() { .methodBuilder("build") .addModifiers(Modifier.PUBLIC, Modifier.STATIC) .returns(DAG_DEF_TYPE) - .addStatement($$"var dag = new $T($S)", DAG_DEF_TYPE, ann.id.ifBlank { el.simpleName }) + .addStatement( + $$"var dag = $T.declaredBy(new $T($S), $T.class)", + DAG_SOURCE_TYPE, + DAG_DEF_TYPE, + ann.id.ifBlank { el.simpleName }, + ClassName.get(el), + ) explicitConfig(el, DAG_ANNOTATION, DAG_STRUCTURAL_ATTRIBUTES, SchemaFields.DAG).forEach { (key, value) -> buildMethod.addStatement($$"dag.config($S, $L)", key, value) } @@ -913,6 +920,7 @@ private val TASK_INPUT_TYPE = ClassName.get(TaskInput::class.java) private val TASK_ARGS_TYPE = ClassName.get(TaskArgs::class.java) private val TYPE_REF_TYPE = ClassName.get(TypeRef::class.java) private val ARG_VALUES_TYPE = ClassName.get(ArgValues::class.java) +private val DAG_SOURCE_TYPE = ClassName.get(DagSource::class.java) private val REFS_TYPE = ClassName.get(Refs::class.java) private val ARG_TYPE = ClassName.get(Arg::class.java) private val TASK_HANDLE_TYPE = ClassName.get(TaskRef::class.java) diff --git a/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt b/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt index c3759fa19985d..a30dbacf00843 100644 --- a/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt +++ b/java-sdk/processor/src/test/kotlin/org/apache/airflow/sdk/BuilderTest.kt @@ -23,8 +23,11 @@ import com.google.testing.compile.CompilationSubject.assertThat import com.google.testing.compile.Compiler import com.google.testing.compile.JavaFileObjectSubject import com.google.testing.compile.JavaFileObjects +import org.junit.jupiter.api.Assertions import org.junit.jupiter.api.DisplayName import org.junit.jupiter.api.Test +import java.nio.file.Files +import javax.tools.JavaFileObject private fun compile(source: String) = Compiler.javac().withProcessors(BuilderProcessor()).compile( @@ -93,12 +96,13 @@ class BuilderTest { import org.apache.airflow.sdk.Context; import org.apache.airflow.sdk.DagDef; import org.apache.airflow.sdk.Task; + import org.apache.airflow.sdk.internal.DagSource; import org.apache.airflow.sdk.internal.Refs; import org.apache.airflow.sdk.internal.TaskArgs; public final class TestExampleBuilder { public static DagDef build() { - var dag = new DagDef("TestExample"); + var dag = DagSource.declaredBy(new DagDef("TestExample"), TestExample.class); return Refs.record(dag, List.of("t1", "t2", "t3"), List.of(), new TestExample.Wiring()::depends); } @@ -208,12 +212,13 @@ class BuilderTest { import org.apache.airflow.sdk.Context; import org.apache.airflow.sdk.DagDef; import org.apache.airflow.sdk.Task; + import org.apache.airflow.sdk.internal.DagSource; import org.apache.airflow.sdk.internal.Refs; import org.apache.airflow.sdk.internal.TaskArgs; public final class TestExampleBuilder { public static DagDef build() { - var dag = new DagDef("TestExample"); + var dag = DagSource.declaredBy(new DagDef("TestExample"), TestExample.class); return Refs.record(dag, List.of("t"), List.of(), new TestExample.Wiring()::depends); } @@ -275,13 +280,14 @@ class BuilderTest { import org.apache.airflow.sdk.Context; import org.apache.airflow.sdk.DagDef; import org.apache.airflow.sdk.Task; + import org.apache.airflow.sdk.internal.DagSource; import org.apache.airflow.sdk.internal.Refs; import org.apache.airflow.sdk.internal.TaskArgs; import org.apache.airflow.sdk.internal.TypeRef; public final class TestExampleBuilder { public static DagDef build() { - var dag = new DagDef("TestExample"); + var dag = DagSource.declaredBy(new DagDef("TestExample"), TestExample.class); return Refs.record(dag, List.of("t"), List.of(), new TestExample.Wiring()::depends); } @@ -429,11 +435,12 @@ class BuilderTest { import org.apache.airflow.sdk.Context; import org.apache.airflow.sdk.DagDef; import org.apache.airflow.sdk.Task; + import org.apache.airflow.sdk.internal.DagSource; import org.apache.airflow.sdk.internal.Refs; public final class TestExampleBuilder { public static DagDef build() { - var dag = new DagDef("cfg"); + var dag = DagSource.declaredBy(new DagDef("cfg"), TestExample.class); dag.config("schedule", "@daily"); dag.config("tags", List.of("a", "b")); dag.config("catchup", true); @@ -518,12 +525,13 @@ class BuilderTest { import org.apache.airflow.sdk.Context; import org.apache.airflow.sdk.DagDef; import org.apache.airflow.sdk.Task; + import org.apache.airflow.sdk.internal.DagSource; import org.apache.airflow.sdk.internal.Refs; import org.apache.airflow.sdk.internal.TaskArgs; public final class TestExampleBuilder { public static DagDef build() { - var dag = new DagDef("TestExample"); + var dag = DagSource.declaredBy(new DagDef("TestExample"), TestExample.class); return Refs.record(dag, List.of("t"), List.of(), new TestExample.Wiring()::depends); } public static final class T implements Task { @@ -588,12 +596,13 @@ class BuilderTest { import org.apache.airflow.sdk.DagDef; import org.apache.airflow.sdk.Task; import org.apache.airflow.sdk.internal.ArgValues; + import org.apache.airflow.sdk.internal.DagSource; import org.apache.airflow.sdk.internal.Refs; import org.apache.airflow.sdk.internal.TaskArgs; public final class TestExampleBuilder { public static DagDef build() { - var dag = new DagDef("TestExample"); + var dag = DagSource.declaredBy(new DagDef("TestExample"), TestExample.class); return Refs.record(dag, List.of("flat", "named"), List.of(), new TestExample.Wiring()::depends); } @@ -811,10 +820,11 @@ class BuilderTest { package org.apache.airflow.example; import java.util.List; import org.apache.airflow.sdk.DagDef; + import org.apache.airflow.sdk.internal.DagSource; import org.apache.airflow.sdk.internal.Refs; public final class TestExampleBuilder { public static DagDef build() { - var dag = new DagDef("foo"); + var dag = DagSource.declaredBy(new DagDef("foo"), TestExample.class); return Refs.record(dag, List.of(), List.of(), new TestExample.Wiring()::depends); } } @@ -847,10 +857,11 @@ class BuilderTest { package org.apache.airflow.example; import java.util.List; import org.apache.airflow.sdk.DagDef; + import org.apache.airflow.sdk.internal.DagSource; import org.apache.airflow.sdk.internal.Refs; public final class Foo { public static DagDef build() { - var dag = new DagDef("TestExample"); + var dag = DagSource.declaredBy(new DagDef("TestExample"), TestExample.class); return Refs.record(dag, List.of(), List.of(), new TestExample.Wiring()::depends); } } @@ -894,11 +905,12 @@ class BuilderTest { import org.apache.airflow.sdk.Context; import org.apache.airflow.sdk.DagDef; import org.apache.airflow.sdk.Task; + import org.apache.airflow.sdk.internal.DagSource; import org.apache.airflow.sdk.internal.Refs; public final class TestExampleBuilder { public static DagDef build() { - var dag = new DagDef("TestExample"); + var dag = DagSource.declaredBy(new DagDef("TestExample"), TestExample.class); return Refs.record(dag, List.of("foo"), List.of(), new TestExample.Wiring()::depends); } @@ -1965,11 +1977,12 @@ class BuilderTest { import org.apache.airflow.sdk.DagDef; import org.apache.airflow.sdk.Task; import org.apache.airflow.sdk.internal.ArgValues; + import org.apache.airflow.sdk.internal.DagSource; import org.apache.airflow.sdk.internal.Refs; public final class TestExampleBuilder { public static DagDef build() { - var dag = new DagDef("TestExample"); + var dag = DagSource.declaredBy(new DagDef("TestExample"), TestExample.class); return Refs.record(dag, List.of("score"), List.of(), new TestExample.Wiring()::depends); } @@ -2147,4 +2160,52 @@ class BuilderTest { "Nested class 'Inner' holding @Builder.TaskHandler methods must be static", ) } + + @Test + @DisplayName("map the Dag to the annotated class, not its generated builder") + fun dagDeclaredByAnnotatedClass() { + val compilation = + compile( + """ + package org.apache.airflow.example; + import org.apache.airflow.sdk.Builder; + @Builder.Dag(id = "orders") + public class TestExample { + @Builder.Task + public void t1() {} + + @Builder.Deps + static class Wiring implements TestExampleDeps { + void depends() { + t1(); + } + } + } + """, + ) + assertThat(compilation).succeeded() + + val classes = + compilation + .generatedFiles() + .filter { it.kind == JavaFileObject.Kind.CLASS } + .associate { + it.name + .removePrefix("/CLASS_OUTPUT/") + .removeSuffix(".class") + .replace('/', '.') to it.openInputStream().readBytes() + } + val loader = + object : ClassLoader(javaClass.classLoader) { + override fun findClass(name: String): Class<*> { + val bytes = classes[name] ?: throw ClassNotFoundException(name) + return defineClass(name, bytes, 0, bytes.size) + } + } + val bundle = Bundle().register(loader.loadClass("org.apache.airflow.example.TestExample")) + val target = Files.createTempFile("sources", ".json").toFile() + Server.create(arrayOf("--describe-sources", target.path)).serve(bundle) + + Assertions.assertEquals("""{"orders":"org.apache.airflow.example.TestExample"}""", target.readText()) + } } diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/DagDef.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/DagDef.kt index ff6f740af180e..a269dff8c2c2a 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/DagDef.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/DagDef.kt @@ -19,6 +19,7 @@ package org.apache.airflow.sdk +import org.apache.airflow.sdk.internal.DagSource import org.apache.airflow.sdk.internal.GROUP_ID import org.apache.airflow.sdk.internal.SchemaFields import org.apache.airflow.sdk.internal.checkConfigValue @@ -57,6 +58,9 @@ class DagDef( /** Edges with a task group at either end, in the order drawn. */ internal val groupEdges = linkedSetOf>() + /** Outermost class that declared this Dag, or `null` if it could not be told. */ + internal var declaringClass: Class<*>? = DagSource.capture() + /** * Sets one Dag-level configuration value. * diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt index 47b44662b9969..c0caccf5d4717 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt @@ -19,7 +19,10 @@ package org.apache.airflow.sdk +import com.fasterxml.jackson.databind.ObjectMapper import com.xenomachina.argparser.ArgParser +import com.xenomachina.argparser.MissingValueException +import com.xenomachina.argparser.default import io.ktor.network.selector.SelectorManager import io.ktor.network.sockets.InetSocketAddress import io.ktor.network.sockets.aSocket @@ -38,6 +41,7 @@ import org.apache.airflow.sdk.execution.comm.ErrorResponse import org.apache.airflow.sdk.execution.comm.StartupDetails import org.apache.airflow.sdk.execution.parseDags import org.apache.airflow.sdk.execution.runTask +import java.io.File import kotlin.text.substringAfterLast import kotlin.text.substringBeforeLast @@ -50,12 +54,16 @@ private class Args( address.substringAfterLast(':').toInt(), ) - val comm by parser.storing("--comm", help = "Address (host:port) to communicate with parent") { - parseAddress(this) - } - val logs by parser.storing("--logs", help = "Address (host:port) to send Airflow logs to") { - parseAddress(this) - } + val comm by parser + .storing("--comm", help = "Address (host:port) to communicate with parent") { parseAddress(this) } + .default(null) + val logs by parser + .storing("--logs", help = "Address (host:port) to send Airflow logs to") { parseAddress(this) } + .default(null) + val describeSources by parser + .storing("--describe-sources", help = "File to write each Dag's declaring class to, instead of serving") { + File(this) + }.default(null) } /** @@ -83,10 +91,13 @@ class ApiError( * The process exits when the coordinator closes the connection (normally after * one task-instance execution). */ -class Server( - private val comm: InetSocketAddress, - private val logs: InetSocketAddress, +class Server private constructor( + private val comm: InetSocketAddress?, + private val logs: InetSocketAddress?, + private val describeSources: File?, ) { + constructor(comm: InetSocketAddress, logs: InetSocketAddress) : this(comm, logs, null) + companion object { /** * Parses coordinator addresses from command-line arguments and returns a @@ -97,6 +108,10 @@ class Server( * * * `--comm host:port` address for task-execution messages. * * `--logs host:port` address for log forwarding. + * * `--describe-sources file` makes [serve] write a JSON object mapping each + * Java-declared Dag ID to the binary name of its declaring class into `file` + * and return without connecting anywhere; `--comm` and `--logs` are not + * required then. It is for the Gradle plugin, not for users. * * @param args Command-line arguments as received by `main`. * @return A configured [Server] ready to call [serve]. @@ -104,7 +119,12 @@ class Server( @JvmStatic fun create(args: Array): Server { val args = ArgParser(args).parseInto(::Args) - return Server(args.comm, args.logs) + args.describeSources?.let { return Server(args.comm, args.logs, it) } + return Server( + args.comm ?: throw MissingValueException("--comm"), + args.logs ?: throw MissingValueException("--logs"), + null, + ) } } @@ -146,12 +166,16 @@ class Server( suspend fun serveAsync(bundle: Bundle) = coroutineScope { bundle.finalizeRegistration() + if (describeSources != null) { + writeSources(bundle, describeSources) + return@coroutineScope + } val deferral = CompletableDeferred() launch { try { SelectorManager(Dispatchers.IO).use { selector -> - aSocket(selector).tcp().connect(comm).use { socket -> + aSocket(selector).tcp().connect(comm!!).use { socket -> logger.debug("Connected comm", mapOf("addr" to comm)) CoordinatorComm( socket.openReadChannel(), @@ -167,7 +191,7 @@ class Server( } launch { SelectorManager(Dispatchers.IO).use { selector -> - aSocket(selector).tcp().connect(logs).use { socket -> + aSocket(selector).tcp().connect(logs!!).use { socket -> logger.debug("Connected logs", mapOf("addr" to logs)) LogSender.configure(socket.openWriteChannel(autoFlush = true)) deferral.await() @@ -176,6 +200,16 @@ class Server( } } + private fun writeSources( + bundle: Bundle, + target: File, + ) { + val sources = linkedMapOf() + bundle.dags.values.forEach { dag -> dag.declaringClass?.let { sources[dag.id] = it.name } } + target.absoluteFile.parentFile?.mkdirs() + ObjectMapper().writeValue(target, sources) + } + internal suspend fun dispatchTask( bundle: Bundle, coordinator: CoordinatorComm, diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/DagSource.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/DagSource.kt new file mode 100644 index 0000000000000..809dd818a3d3f --- /dev/null +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/DagSource.kt @@ -0,0 +1,66 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.airflow.sdk.internal + +import org.apache.airflow.sdk.DagDef + +/** + * @suppress + * + * Tracks which class declared a [DagDef], so the bundle can ship that class's + * source file. Public so that processor-generated builders can call [declaredBy]; + * not user-facing API. + */ +object DagSource { + private const val SDK_PACKAGE = "org.apache.airflow.sdk." + private val IGNORED_PREFIXES = listOf(SDK_PACKAGE, "java.", "javax.", "jdk.", "sun.", "kotlin.", "kotlinx.") + + private val walker = StackWalker.getInstance(StackWalker.Option.RETAIN_CLASS_REFERENCE) + + /** + * Names [declaring] as the class that declared [dag], replacing what was + * captured at construction. A generated builder calls this so the Dag points + * at the annotated class rather than the builder. + */ + @JvmStatic + fun declaredBy( + dag: DagDef, + declaring: Class<*>, + ): DagDef { + dag.declaringClass = outermost(declaring) + return dag + } + + /** The outermost class of the first caller outside the SDK and the standard libraries. */ + internal fun capture(): Class<*>? = + walker.walk { frames -> + frames + .map { it.declaringClass } + .filter { c -> IGNORED_PREFIXES.none { c.name.startsWith(it) } } + .findFirst() + .map { outermost(it) } + .orElse(null) + } + + internal fun outermost(cls: Class<*>): Class<*> { + var current = cls + while (true) current = current.enclosingClass ?: return current + } +} diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/example/DagSourceFixtures.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/example/DagSourceFixtures.kt new file mode 100644 index 0000000000000..bdcf697060b4a --- /dev/null +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/example/DagSourceFixtures.kt @@ -0,0 +1,51 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.airflow.example + +import org.apache.airflow.sdk.DagDef + +// Lives outside org.apache.airflow.sdk, as user code does, so the SDK does not skip its frames. +class DagSourceFixtures { + fun plain() = DagDef("plain") + + fun fromLambda(): DagDef { + lateinit var dag: DagDef + Runnable { dag = DagDef("lambda") }.run() + return dag + } + + fun fromAnonymous(): DagDef { + lateinit var dag: DagDef + object : Runnable { + override fun run() { + dag = DagDef("anonymous") + } + }.run() + return dag + } + + class Nested { + fun make() = DagDef("nested") + + class Deeper { + fun make() = DagDef("deeper") + } + } +} diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/DagSourceTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/DagSourceTest.kt new file mode 100644 index 0000000000000..39a3b740acab6 --- /dev/null +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/DagSourceTest.kt @@ -0,0 +1,132 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.airflow.sdk + +import com.fasterxml.jackson.databind.ObjectMapper +import com.xenomachina.argparser.SystemExitException +import org.apache.airflow.example.DagSourceFixtures +import org.apache.airflow.sdk.internal.DagSource +import org.junit.jupiter.api.Assertions +import org.junit.jupiter.api.DisplayName +import org.junit.jupiter.api.Test +import org.junit.jupiter.api.io.TempDir +import java.io.File + +internal class DagSourceTest { + private class NoOp : Task { + override fun execute( + context: Context, + client: Client, + ) = Unit + } + + @Test + @DisplayName("Should record the class that constructed the Dag") + fun recordsConstructingClass() { + Assertions.assertEquals(DagSourceFixtures::class.java, DagSourceFixtures().plain().declaringClass) + } + + @Test + @DisplayName("Should record the outermost class for nested, lambda and anonymous callers") + fun recordsOutermostClass() { + val fixtures = DagSourceFixtures() + Assertions.assertEquals(DagSourceFixtures::class.java, DagSourceFixtures.Nested().make().declaringClass) + Assertions.assertEquals( + DagSourceFixtures::class.java, + DagSourceFixtures.Nested + .Deeper() + .make() + .declaringClass, + ) + Assertions.assertEquals(DagSourceFixtures::class.java, fixtures.fromLambda().declaringClass) + Assertions.assertEquals(DagSourceFixtures::class.java, fixtures.fromAnonymous().declaringClass) + } + + @Test + @DisplayName("Should let a generated builder name the annotated class instead of itself") + fun declaredByOverridesCapturedClass() { + val dag = DagSource.declaredBy(DagDef("dag"), DagSourceFixtures.Nested::class.java) + Assertions.assertEquals(DagSourceFixtures::class.java, dag.declaringClass) + } + + private fun describe( + dir: File, + bundle: Bundle, + ): Map<*, *> { + val target = File(dir, "out/sources.json") + Server.create(arrayOf("--describe-sources", target.path)).serve(bundle) + return ObjectMapper().readValue(target, Map::class.java) + } + + @Test + @DisplayName("Should write each Java-declared Dag's declaring class and skip task-handler Dags") + fun describeSourcesWritesDeclaringClasses( + @TempDir dir: File, + ) { + val fixtures = DagSourceFixtures() + val bundle = + Bundle(listOf(fixtures.plain(), DagSourceFixtures.Nested().make())) + .register("python_owned", "t", NoOp::class.java) + + Assertions.assertEquals( + mapOf("plain" to DagSourceFixtures::class.java.name, "nested" to DagSourceFixtures::class.java.name), + describe(dir, bundle), + ) + } + + @Test + @DisplayName("Should omit a Dag whose declaring class is unknown") + fun describeSourcesOmitsUnknownDeclaringClass( + @TempDir dir: File, + ) { + val unknown = DagDef("unknown").also { it.declaringClass = null } + val bundle = Bundle(listOf(DagSourceFixtures().plain(), unknown)) + + Assertions.assertEquals(mapOf("plain" to DagSourceFixtures::class.java.name), describe(dir, bundle)) + } + + @Test + @DisplayName("Should write an empty object when no Dag is declared in Java") + fun describeSourcesWithoutJavaDags( + @TempDir dir: File, + ) { + Assertions.assertEquals(emptyMap(), describe(dir, Bundle())) + } + + @Test + @DisplayName("Should use the binary name of the declaring class") + fun describeSourcesUsesBinaryName( + @TempDir dir: File, + ) { + val dag = DagDef("dag").also { it.declaringClass = DagSourceFixtures.Nested::class.java } + Assertions.assertEquals( + mapOf("dag" to "org.apache.airflow.example.DagSourceFixtures\$Nested"), + describe(dir, Bundle(listOf(dag))), + ) + } + + @Test + @DisplayName("Should still require --comm and --logs without --describe-sources") + fun createStillRequiresAddresses() { + Assertions.assertThrows(SystemExitException::class.java) { Server.create(arrayOf("--comm", "localhost:1")) } + Assertions.assertThrows(SystemExitException::class.java) { Server.create(arrayOf("--logs", "localhost:1")) } + Assertions.assertThrows(SystemExitException::class.java) { Server.create(emptyArray()) } + } +} From 030a0702bd729de838a0357620d11dcc62819b1c Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Fri, 2 Oct 2026 15:19:36 +0000 Subject: [PATCH 02/10] Java SDK: Pack each native Dag's source file into the bundle JAR The packDagSources task runs mainClass in describe mode, finds each Dag's source file through the class file's SourceFile attribute, and packs each unique file with a sources.json mapping under META-INF/airflow/. The JAR manifest names it with Airflow-Java-SDK-Sources. --- java-sdk/README.md | 4 + java-sdk/plugin/build.gradle.kts | 6 + .../airflow/sdk/plugin/AirflowSdkPlugin.kt | 41 +++- .../apache/airflow/sdk/plugin/ClassFiles.kt | 85 +++++++ .../airflow/sdk/plugin/PackDagSources.kt | 198 +++++++++++++++++ .../sdk/plugin/AirflowSdkPluginTest.kt | 207 ++++++++++++++++++ .../airflow/sdk/plugin/ClassFilesTest.kt | 80 +++++++ .../airflow/sdk/plugin/PackDagSourcesTest.kt | 89 ++++++++ .../apache/airflow/sdk/plugin/TestSupport.kt | 45 ++++ 9 files changed, 750 insertions(+), 5 deletions(-) create mode 100644 java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/ClassFiles.kt create mode 100644 java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/PackDagSources.kt create mode 100644 java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPluginTest.kt create mode 100644 java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/ClassFilesTest.kt create mode 100644 java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/PackDagSourcesTest.kt create mode 100644 java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/TestSupport.kt diff --git a/java-sdk/README.md b/java-sdk/README.md index 6b4ba3522fd44..8a3284403b827 100644 --- a/java-sdk/README.md +++ b/java-sdk/README.md @@ -91,6 +91,10 @@ Now `cd example` into the example project, and ../gradlew bundle ``` + The bundle JAR also carries the source file of each Dag declared in Java + (and of the entrypoint) under `META-INF/airflow/`, so Airflow can show a + Dag's source. + * Put the [DAG with stub tasks](./example/src/resources/dags) to somewhere Airflow can find. * Ensure the `java` command is available in the same environments the Airflow diff --git a/java-sdk/plugin/build.gradle.kts b/java-sdk/plugin/build.gradle.kts index edf930bb6fc84..7efcdc36e612c 100644 --- a/java-sdk/plugin/build.gradle.kts +++ b/java-sdk/plugin/build.gradle.kts @@ -25,6 +25,12 @@ plugins { dependencies { implementation("com.gradleup.shadow:shadow-gradle-plugin:9.1.0") // Last supporting Java 11. + + testImplementation(kotlin("test")) +} + +tasks.withType { + useJUnitPlatform() } gradlePlugin { diff --git a/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt index 3756ef688e174..2ecb0acb733d2 100644 --- a/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt +++ b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt @@ -22,11 +22,13 @@ package org.apache.airflow.sdk.plugin import com.github.jengelman.gradle.plugins.shadow.tasks.ShadowJar import org.gradle.api.Plugin import org.gradle.api.Project +import org.gradle.api.plugins.JavaPluginExtension import org.gradle.api.provider.Property import org.gradle.api.tasks.Copy import org.gradle.api.tasks.Input import org.gradle.api.tasks.SourceSetContainer import org.gradle.api.tasks.bundling.Jar +import org.gradle.jvm.toolchain.JavaToolchainService import java.lang.reflect.Modifier import java.net.URLClassLoader import java.util.jar.JarFile @@ -93,6 +95,13 @@ abstract class AirflowBundleExtension { * identify which version of the Supervisor Schema it should use to communicate * with the built JAR. * + * When `mainClass` is set, the bundle JAR also carries the source file of each + * Dag declared in Java and of the entrypoint, so Airflow can show a Dag's source. + * They are collected by the `packDagSources` task, which runs `mainClass` once in + * a describe mode, and are listed in `META-INF/airflow/sources.json`, named by the + * `Airflow-Java-SDK-Sources` manifest attribute. If that run fails, a warning is + * logged and only the entrypoint's source is packed. + * * If `fatJar` is explicitly set to `false`, the `bundle` task builds a bare JAR * containing only the Dag bundle, and collect all dependency JARs into the target * directory instead. In this mode, `Airflow-Supervisor-Schema-Version` lives in @@ -107,21 +116,43 @@ class AirflowSdkPlugin : Plugin { ext.fatJar.convention(true) project.afterEvaluate { + val main = project.extensions.getByType(SourceSetContainer::class.java).getByName("main") + + val packTask = + project.tasks.register("packDagSources", PackDagSources::class.java) { task -> + task.group = "build" + task.description = "Collects the source file of each Dag and the entrypoint to pack into the bundle JAR." + task.dependsOn(project.tasks.named("classes")) + task.onlyIf { ext.mainClass.isPresent } + task.mainClass.set(ext.mainClass) + task.classesDirs.from(main.output.classesDirs) + task.runtimeClasspath.from(main.output, main.runtimeClasspath) + task.sourceDirs.from(main.allSource.srcDirs) + task.launcher.convention( + project.extensions + .getByType(JavaToolchainService::class.java) + .launcherFor(project.extensions.getByType(JavaPluginExtension::class.java).toolchain), + ) + task.describeFile.set(project.layout.buildDirectory.file("airflow/describe-sources.json")) + task.sourcesDir.set(project.layout.buildDirectory.dir("airflow/sources")) + } + project.tasks.withType(Jar::class.java).configureEach { task -> task.doFirst { ext.mainClass.orNull?.let { className -> task.manifest.attributes(mapOf("Main-Class" to className)) } } + if (ext.mainClass.isPresent) { + task.dependsOn(packTask) + task.from(packTask.flatMap { it.sourcesDir }) + task.manifest.attributes(mapOf(SOURCES_MANIFEST_ATTRIBUTE to SOURCES_JSON_PATH)) + } } val classFiles = project.objects.fileCollection().from( - project.extensions - .getByType(SourceSetContainer::class.java) - .getByName("main") - .output - .classesDirs, + main.output.classesDirs, project.configurations.getByName("runtimeClasspath"), ) diff --git a/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/ClassFiles.kt b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/ClassFiles.kt new file mode 100644 index 0000000000000..ef693559501f9 --- /dev/null +++ b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/ClassFiles.kt @@ -0,0 +1,85 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.airflow.sdk.plugin + +import java.io.ByteArrayInputStream +import java.io.DataInputStream +import java.io.IOException + +private const val CLASS_MAGIC = 0xCAFEBABE.toInt() + +/** + * Reads the `SourceFile` attribute of a class file, which javac, kotlinc and scalac all + * write, or returns `null` if the class has none or the bytes are not a class file. + */ +internal fun readSourceFile(classFile: ByteArray): String? = + try { + DataInputStream(ByteArrayInputStream(classFile)).use(::readSourceFile) + } catch (_: IOException) { + null + } + +private fun readSourceFile(input: DataInputStream): String? { + if (input.readInt() != CLASS_MAGIC) return null + input.skipBytes(4) // minor and major version + + val utf8 = arrayOfNulls(input.readUnsignedShort()) + var index = 1 + while (index < utf8.size) { + when (val tag = input.readUnsignedByte()) { + 1 -> utf8[index] = input.readUTF() + 3, 4, 9, 10, 11, 12, 17, 18 -> input.skipBytes(4) + 5, 6 -> { + input.skipBytes(8) + index++ // long and double take two slots + } + 7, 8, 16, 19, 20 -> input.skipBytes(2) + 15 -> input.skipBytes(3) + else -> throw IOException("Unknown constant pool tag $tag") + } + index++ + } + + input.skipBytes(6) // access flags, this class, super class + input.skipBytes(2 * input.readUnsignedShort()) // interfaces + repeat(2) { skipMembers(input) } // fields, then methods + + repeat(input.readUnsignedShort()) { + val name = utf8.getOrNull(input.readUnsignedShort()) + val length = input.readInt() + if (name == "SourceFile" && length == 2) return utf8.getOrNull(input.readUnsignedShort()) + input.skipBytes(length) + } + return null +} + +private fun skipMembers(input: DataInputStream) { + repeat(input.readUnsignedShort()) { + input.skipBytes(6) // access flags, name, descriptor + skipAttributes(input) + } +} + +private fun skipAttributes(input: DataInputStream) { + repeat(input.readUnsignedShort()) { + input.skipBytes(2) + input.skipBytes(input.readInt()) + } +} diff --git a/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/PackDagSources.kt b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/PackDagSources.kt new file mode 100644 index 0000000000000..ac271f495adc6 --- /dev/null +++ b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/PackDagSources.kt @@ -0,0 +1,198 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.airflow.sdk.plugin + +import groovy.json.JsonOutput +import groovy.json.JsonSlurper +import org.gradle.api.DefaultTask +import org.gradle.api.file.ConfigurableFileCollection +import org.gradle.api.file.DirectoryProperty +import org.gradle.api.file.RegularFileProperty +import org.gradle.api.provider.Property +import org.gradle.api.tasks.CacheableTask +import org.gradle.api.tasks.Classpath +import org.gradle.api.tasks.Input +import org.gradle.api.tasks.InputFiles +import org.gradle.api.tasks.Nested +import org.gradle.api.tasks.Optional +import org.gradle.api.tasks.OutputDirectory +import org.gradle.api.tasks.OutputFile +import org.gradle.api.tasks.PathSensitive +import org.gradle.api.tasks.PathSensitivity +import org.gradle.api.tasks.TaskAction +import org.gradle.jvm.toolchain.JavaLauncher +import org.gradle.process.ExecOperations +import java.io.ByteArrayOutputStream +import java.io.File +import javax.inject.Inject + +internal const val SOURCES_MANIFEST_ATTRIBUTE = "Airflow-Java-SDK-Sources" +internal const val SOURCES_JSON_PATH = "META-INF/airflow/sources.json" +internal const val SOURCES_DIR_PATH = "META-INF/airflow/sources" + +/** + * Finds the source file of a class by its `SourceFile` attribute: the class + * file's package path plus that name, looked up in each source directory. + */ +internal class SourceLocator( + private val classesDirs: Iterable, + private val sourceDirs: Iterable, +) { + /** Path of the source file relative to its source directory, or `null` if it cannot be found. */ + fun locate(className: String): String? { + val packagePath = className.substringBeforeLast('.', "").replace('.', '/') + val classFile = + classesDirs + .map { File(it, className.replace('.', '/') + ".class") } + .firstOrNull { it.isFile } ?: return null + val sourceFile = readSourceFile(classFile.readBytes())?.takeIf { it.isNotEmpty() && '/' !in it && '\\' !in it } + val relative = listOfNotNull(packagePath.ifEmpty { null }, sourceFile ?: return null).joinToString("/") + return relative.takeIf { path -> sourceDirs.any { File(it, path).isFile } } + } + + fun file(relativePath: String): File = sourceDirs.map { File(it, relativePath) }.first { it.isFile } +} + +internal fun sourcesJson( + entrypointPath: String?, + dagSourcePaths: Map, +): String = + JsonOutput.prettyPrint( + JsonOutput.toJson( + linkedMapOf().apply { + entrypointPath?.let { put("entrypoint_path", it) } + put("dag_source_paths", dagSourcePaths) + }, + ), + ) + +/** + * Runs the bundle's `mainClass` in describe mode to learn which class declared + * each Dag, then collects those classes' source files, and the entry point's, + * for the bundle JAR to carry. + */ +@CacheableTask +abstract class PackDagSources : DefaultTask() { + @get:Inject + abstract val execOperations: ExecOperations + + @get:Input + @get:Optional + abstract val mainClass: Property + + @get:Classpath + abstract val classesDirs: ConfigurableFileCollection + + @get:Classpath + abstract val runtimeClasspath: ConfigurableFileCollection + + @get:InputFiles + @get:PathSensitive(PathSensitivity.RELATIVE) + abstract val sourceDirs: ConfigurableFileCollection + + @get:Nested + @get:Optional + abstract val launcher: Property + + @get:OutputFile + abstract val describeFile: RegularFileProperty + + @get:OutputDirectory + abstract val sourcesDir: DirectoryProperty + + private var describeFailed = false + + init { + outputs.doNotCacheIf("the Dag describe run failed") { describeFailed } + } + + @TaskAction + fun pack() { + val describe = describeFile.get().asFile + val root = sourcesDir.get().asFile + describe.delete() + root.deleteRecursively() + describe.parentFile.mkdirs() + + val main = mainClass.get() + val declaringClasses = describeDags(main, describe) + val locator = SourceLocator(classesDirs.files, sourceDirs.files) + + val entrypoint = locator.locate(main) + val dagPaths = linkedMapOf() + declaringClasses.forEach { (dagId, className) -> + val path = locator.locate(className) + if (path == null) { + logger.info("No source file found for class {} of Dag '{}'; it falls back to the entrypoint", className, dagId) + } else { + dagPaths[dagId] = path + } + } + + (listOfNotNull(entrypoint) + dagPaths.values).distinct().forEach { path -> + val target = File(root, "$SOURCES_DIR_PATH/$path") + target.parentFile.mkdirs() + locator.file(path).copyTo(target) + } + File(root, SOURCES_JSON_PATH).apply { + parentFile.mkdirs() + writeText(sourcesJson(entrypoint, dagPaths) + "\n") + } + } + + private fun describeDags( + main: String, + describe: File, + ): Map { + val output = ByteArrayOutputStream() + val failure = + try { + val result = + execOperations.javaexec { spec -> + launcher.orNull?.let { spec.executable(it.executablePath.asFile.absolutePath) } + spec.classpath(classesDirs, runtimeClasspath) + spec.mainClass.set(main) + spec.args("--describe-sources", describe.absolutePath) + spec.isIgnoreExitValue = true + spec.standardOutput = output + spec.errorOutput = output + } + if (result.exitValue != 0) "exit code ${result.exitValue}" else null + } catch (e: Exception) { + e.message ?: e.javaClass.simpleName + } + if (failure == null && describe.isFile) { + @Suppress("UNCHECKED_CAST") + (runCatching { JsonSlurper().parse(describe) }.getOrNull() as? Map)?.let { parsed -> + return parsed.mapNotNull { (id, cls) -> (cls as? String)?.let { id to it } }.toMap(linkedMapOf()) + } + } + describeFailed = true + val why = failure ?: "it wrote no valid --describe-sources file; does main call Server.serve?" + val log = output.toString().trim() + logger.warn( + "Could not read each Dag's source from '{}' ({}); only its entrypoint source is packed.{}", + main, + why, + if (log.isEmpty()) "" else "\n$log", + ) + return emptyMap() + } +} diff --git a/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPluginTest.kt b/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPluginTest.kt new file mode 100644 index 0000000000000..91fb5c5d5ed45 --- /dev/null +++ b/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPluginTest.kt @@ -0,0 +1,207 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.airflow.sdk.plugin + +import groovy.json.JsonSlurper +import org.gradle.testkit.runner.BuildResult +import org.gradle.testkit.runner.GradleRunner +import org.gradle.testkit.runner.TaskOutcome +import org.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.Assertions.assertFalse +import org.junit.jupiter.api.Assertions.assertNull +import org.junit.jupiter.api.Assertions.assertTrue +import org.junit.jupiter.api.Test +import org.junit.jupiter.api.io.TempDir +import java.io.File +import java.util.jar.JarFile + +class AirflowSdkPluginTest { + private fun main(describeBody: String) = + """ + package com.example; + + public class Main { + public static void main(String[] args) throws Exception { + if (args.length == 2 && args[0].equals("--describe-sources")) { + $describeBody + } + } + } + """.trimIndent() + + private val describeDags = + """ + java.nio.file.Files.write(java.nio.file.Paths.get(args[1]), ( + "{\"orders\":\"com.example.Main\"," + + "\"reports\":\"com.example.dags.Reports\"," + + "\"reports_backfill\":\"com.example.dags.Reports\"," + + "\"ghost\":\"com.example.Ghost\"}").getBytes()); + """.trimIndent() + + private fun project( + dir: File, + mainBody: String = describeDags, + mainClass: String? = "com.example.Main", + ): File { + dir.write("settings.gradle", "rootProject.name = 'bundle-test'\n") + dir.write( + "build.gradle", + """ + plugins { id 'org.apache.airflow.sdk' } + airflowBundle { + ${mainClass?.let { "mainClass = '$it'" } ?: ""} + fatJar = false + } + sourceSets { main { java.srcDir 'src/extra' } } + """.trimIndent(), + ) + dir.write("src/main/java/com/example/Main.java", main(mainBody)) + dir.write("src/extra/com/example/dags/Reports.java", "package com.example.dags; public class Reports {}") + return dir + } + + private fun gradle( + dir: File, + vararg tasks: String, + ): BuildResult = + GradleRunner + .create() + .withPluginClasspath() + .withProjectDir(dir) + .withArguments(*tasks, "--stacktrace") + .build() + + private fun entries(jar: File): List = + JarFile(jar).use { + it + .entries() + .asSequence() + .map { e -> e.name } + .toList() + } + + @Suppress("UNCHECKED_CAST") + private fun sourcesJson(jar: File): Map = + JarFile(jar).use { + JsonSlurper().parse(it.getInputStream(it.getJarEntry("META-INF/airflow/sources.json"))) as Map + } + + @Test + fun packsEachDagSourceOnceAndTheEntrypoint( + @TempDir dir: File, + ) { + project(dir) + val result = gradle(dir, "jar") + assertEquals(TaskOutcome.SUCCESS, result.task(":packDagSources")!!.outcome) + + val jar = File(dir, "build/libs/bundle-test.jar") + assertEquals( + listOf( + "META-INF/airflow/sources.json", + "META-INF/airflow/sources/com/example/Main.java", + "META-INF/airflow/sources/com/example/dags/Reports.java", + ), + entries(jar).filter { it.startsWith("META-INF/airflow/") && !it.endsWith("/") }.sorted(), + ) + assertEquals( + mapOf( + "entrypoint_path" to "com/example/Main.java", + "dag_source_paths" to + mapOf( + "orders" to "com/example/Main.java", + "reports" to "com/example/dags/Reports.java", + "reports_backfill" to "com/example/dags/Reports.java", + ), + ), + sourcesJson(jar), + ) + JarFile(jar).use { + assertEquals("META-INF/airflow/sources.json", it.manifest.mainAttributes.getValue("Airflow-Java-SDK-Sources")) + assertEquals("com.example.Main", it.manifest.mainAttributes.getValue("Main-Class")) + } + } + + @Test + fun isUpToDateOnSecondRunAndRerunsWhenASourceChanges( + @TempDir dir: File, + ) { + project(dir) + gradle(dir, "jar") + assertEquals(TaskOutcome.UP_TO_DATE, gradle(dir, "jar").task(":packDagSources")!!.outcome) + + dir.write("src/extra/com/example/dags/Reports.java", "package com.example.dags; public class Reports { int v; }") + assertEquals(TaskOutcome.SUCCESS, gradle(dir, "jar").task(":packDagSources")!!.outcome) + val packed = File(dir, "build/airflow/sources/META-INF/airflow/sources/com/example/dags/Reports.java") + assertTrue(packed.readText().contains("int v")) + } + + @Test + fun worksWithConfigurationCache( + @TempDir dir: File, + ) { + project(dir) + gradle(dir, "jar", "--configuration-cache") + val again = gradle(dir, "jar", "--configuration-cache") + assertTrue(again.output.contains("Reusing configuration cache")) + } + + @Test + fun fallsBackToEntrypointOnlyWhenTheDescribeRunFails( + @TempDir dir: File, + ) { + project(dir, mainBody = """throw new IllegalStateException("boom");""") + val result = gradle(dir, "jar") + + assertEquals(TaskOutcome.SUCCESS, result.task(":jar")!!.outcome) + assertTrue(result.output.contains("only its entrypoint source is packed"), result.output) + assertTrue(result.output.contains("boom"), result.output) + assertEquals( + mapOf("entrypoint_path" to "com/example/Main.java", "dag_source_paths" to emptyMap()), + sourcesJson(File(dir, "build/libs/bundle-test.jar")), + ) + } + + @Test + fun fallsBackToEntrypointOnlyWhenMainWritesNothing( + @TempDir dir: File, + ) { + project(dir, mainBody = "") + val result = gradle(dir, "jar") + + assertTrue(result.output.contains("does main call Server.serve?"), result.output) + assertEquals( + mapOf("entrypoint_path" to "com/example/Main.java", "dag_source_paths" to emptyMap()), + sourcesJson(File(dir, "build/libs/bundle-test.jar")), + ) + } + + @Test + fun skipsEverythingWithoutMainClass( + @TempDir dir: File, + ) { + project(dir, mainClass = null) + val result = gradle(dir, "jar", "packDagSources") + + assertEquals(TaskOutcome.SKIPPED, result.task(":packDagSources")!!.outcome) + val jar = File(dir, "build/libs/bundle-test.jar") + assertFalse(entries(jar).any { it.startsWith("META-INF/airflow") }) + JarFile(jar).use { assertNull(it.manifest.mainAttributes.getValue("Airflow-Java-SDK-Sources")) } + } +} diff --git a/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/ClassFilesTest.kt b/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/ClassFilesTest.kt new file mode 100644 index 0000000000000..693a8a21c12c6 --- /dev/null +++ b/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/ClassFilesTest.kt @@ -0,0 +1,80 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.airflow.sdk.plugin + +import org.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.Assertions.assertNull +import org.junit.jupiter.api.Test +import org.junit.jupiter.api.io.TempDir +import java.io.File + +class ClassFilesTest { + private class Nested + + private fun bytesOf(cls: Class<*>): ByteArray = + cls.getResourceAsStream("/" + cls.name.replace('.', '/') + ".class")!!.use { it.readBytes() } + + @Test + fun readsSourceFileOfKotlinClass() { + assertEquals("ClassFilesTest.kt", readSourceFile(bytesOf(ClassFilesTest::class.java))) + } + + @Test + fun readsSourceFileOfNestedClassFromItsOutermostFile() { + assertEquals("ClassFilesTest.kt", readSourceFile(bytesOf(Nested::class.java))) + } + + @Test + fun readsSourceFileOfJdkClassesWithLongAndDoubleConstants() { + // These carry every constant pool kind, including long, double and invokedynamic entries. + assertEquals("String.java", readSourceFile(bytesOf(String::class.java))) + assertEquals("Long.java", readSourceFile(bytesOf(Long::class.javaObjectType))) + assertEquals("Double.java", readSourceFile(bytesOf(Double::class.javaObjectType))) + assertEquals("HashMap.java", readSourceFile(bytesOf(java.util.HashMap::class.java))) + } + + @Test + fun readsSourceFileOfJavaClassWithSeveralTopLevelTypes( + @TempDir dir: File, + ) { + dir.write("p/Main.java", "package p; public class Main { static final long L = 1L; } class Other {}") + compileJava(dir, File(dir, "out"), "p/Main.java") + + assertEquals("Main.java", readSourceFile(File(dir, "out/p/Main.class").readBytes())) + assertEquals("Main.java", readSourceFile(File(dir, "out/p/Other.class").readBytes())) + } + + @Test + fun returnsNullWithoutSourceFileAttribute( + @TempDir dir: File, + ) { + dir.write("p/Main.java", "package p; public class Main {}") + compileJava(dir, File(dir, "out"), "p/Main.java", options = listOf("-g:none")) + + assertNull(readSourceFile(File(dir, "out/p/Main.class").readBytes())) + } + + @Test + fun returnsNullForBytesThatAreNotAClassFile() { + assertNull(readSourceFile(ByteArray(0))) + assertNull(readSourceFile("not a class file".toByteArray())) + assertNull(readSourceFile(bytesOf(ClassFilesTest::class.java).copyOf(40))) + } +} diff --git a/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/PackDagSourcesTest.kt b/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/PackDagSourcesTest.kt new file mode 100644 index 0000000000000..bca13df2396b0 --- /dev/null +++ b/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/PackDagSourcesTest.kt @@ -0,0 +1,89 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.airflow.sdk.plugin + +import org.junit.jupiter.api.Assertions.assertEquals +import org.junit.jupiter.api.Assertions.assertNull +import org.junit.jupiter.api.Test +import org.junit.jupiter.api.io.TempDir +import java.io.File + +class PackDagSourcesTest { + private fun project(dir: File): Pair { + val src = File(dir, "src") + val extra = File(dir, "extra") + src.write("com/example/Main.java", "package com.example; public class Main {}") + src.write("com/example/Same.java", "package com.example; public class Same { class Inner {} }") + extra.write("com/example/dags/Reports.java", "package com.example.dags; public class Reports {} class Helper {}") + src.write("com/example/NoSource.java", "package com.example; public class NoSource {}") + val classes = File(dir, "classes") + compileJava( + src, + classes, + "com/example/Main.java", + "com/example/Same.java", + "com/example/NoSource.java", + ) + compileJava(extra, classes, "com/example/dags/Reports.java") + File(src, "com/example/NoSource.java").delete() + return src to classes + } + + @Test + fun locatesSourceByPackagePathAndSourceFileAttribute( + @TempDir dir: File, + ) { + val (src, classes) = project(dir) + val locator = SourceLocator(listOf(classes), listOf(src, File(dir, "extra"))) + + assertEquals("com/example/Main.java", locator.locate("com.example.Main")) + assertEquals("com/example/Same.java", locator.locate("com.example.Same\$Inner")) + assertEquals("com/example/dags/Reports.java", locator.locate("com.example.dags.Reports")) + assertEquals("com/example/dags/Reports.java", locator.locate("com.example.dags.Helper")) + assertEquals(File(dir, "extra/com/example/dags/Reports.java"), locator.file("com/example/dags/Reports.java")) + } + + @Test + fun locatesNothingForUnknownClassOrMissingSource( + @TempDir dir: File, + ) { + val (src, classes) = project(dir) + val locator = SourceLocator(listOf(classes), listOf(src)) + + assertNull(locator.locate("com.example.Missing")) + assertNull(locator.locate("com.example.NoSource")) + assertNull(locator.locate("com.example.dags.Reports")) + } + + @Test + fun rendersSourcesJson() { + assertEquals( + """{"entrypoint_path":"com/example/Main.java","dag_source_paths":{"orders":"com/example/Main.java"}}""", + compact(sourcesJson("com/example/Main.java", mapOf("orders" to "com/example/Main.java"))), + ) + } + + @Test + fun omitsEntrypointPathWhenUnresolved() { + assertEquals("""{"dag_source_paths":{}}""", compact(sourcesJson(null, emptyMap()))) + } + + private fun compact(json: String) = (groovy.json.JsonSlurper().parseText(json) as Map<*, *>).let(groovy.json.JsonOutput::toJson) +} diff --git a/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/TestSupport.kt b/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/TestSupport.kt new file mode 100644 index 0000000000000..d8acc86fcfc67 --- /dev/null +++ b/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/TestSupport.kt @@ -0,0 +1,45 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.airflow.sdk.plugin + +import java.io.File +import javax.tools.ToolProvider + +/** Compiles [sources] (paths under [srcDir]) into [classesDir] with the JDK's own compiler. */ +internal fun compileJava( + srcDir: File, + classesDir: File, + vararg sources: String, + options: List = emptyList(), +) { + classesDir.mkdirs() + val compiler = checkNotNull(ToolProvider.getSystemJavaCompiler()) { "Tests need a JDK" } + val args = options + listOf("-d", classesDir.path) + sources.map { File(srcDir, it).path } + check(compiler.run(null, null, null, *args.toTypedArray()) == 0) { "javac failed for ${sources.toList()}" } +} + +internal fun File.write( + relativePath: String, + text: String, +): File = + File(this, relativePath).apply { + parentFile.mkdirs() + writeText(text) + } From d36ca2cd9a61046a4cc98019e1b857429db5c0df Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Tue, 6 Oct 2026 05:50:35 +0000 Subject: [PATCH 03/10] Java SDK: Bound the describe run and keep its payload to the bundle JAR Every `Jar` task gained a dependency on `packDagSources`, so a sources or javadoc JAR carried the Dag sources and claimed the `Airflow-Java-SDK-Sources` attribute that the Python side picks a bundle JAR out of a directory by. Only `jar` and `shadowJar` carry them now. The describe run had no time bound, so a `main` that does not route through `Server.create(args)` and starts serving instead held the build open with no way out. It now runs under a timeout and falls back to packing the entrypoint's source alone, which is what every other describe failure already does. Say in `serve` and `serveAsync` that a `--describe-sources` server returns without connecting, and record in `sources.json` that its paths are relative to the source directory, which is also where they sit under `META-INF/airflow/sources/`. --- .../airflow/sdk/plugin/AirflowSdkPlugin.kt | 10 ++- .../airflow/sdk/plugin/PackDagSources.kt | 62 +++++++++++++------ .../sdk/plugin/AirflowSdkPluginTest.kt | 23 +++++++ .../kotlin/org/apache/airflow/sdk/Server.kt | 6 ++ 4 files changed, 81 insertions(+), 20 deletions(-) diff --git a/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt index 2ecb0acb733d2..c3bdbc4f8fa98 100644 --- a/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt +++ b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt @@ -65,6 +65,9 @@ abstract class AirflowBundleExtension { abstract val fatJar: Property } +/** The JAR tasks a deployable bundle is assembled from, fat or thin. */ +private val BUNDLE_JAR_TASKS = setOf("jar", "shadowJar") + /** * Gradle plugin for building Apache Airflow Java SDK bundles. * @@ -108,6 +111,7 @@ abstract class AirflowBundleExtension { * the `airflow-sdk` JAR instead. The bundle JAR still contains `Main-Class`. * */ + class AirflowSdkPlugin : Plugin { override fun apply(project: Project) { project.plugins.apply("java") @@ -143,7 +147,11 @@ class AirflowSdkPlugin : Plugin { task.manifest.attributes(mapOf("Main-Class" to className)) } } - if (ext.mainClass.isPresent) { + // Only the JARs a bundle is assembled from carry the Dag sources. A + // sources or javadoc JAR would claim the attribute without being a + // bundle, and the Python side picks the bundle JAR out of a directory + // by that attribute. + if (ext.mainClass.isPresent && task.name in BUNDLE_JAR_TASKS) { task.dependsOn(packTask) task.from(packTask.flatMap { it.sourcesDir }) task.manifest.attributes(mapOf(SOURCES_MANIFEST_ATTRIBUTE to SOURCES_JSON_PATH)) diff --git a/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/PackDagSources.kt b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/PackDagSources.kt index ac271f495adc6..3ba8e98fc15f5 100644 --- a/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/PackDagSources.kt +++ b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/PackDagSources.kt @@ -38,15 +38,22 @@ import org.gradle.api.tasks.PathSensitive import org.gradle.api.tasks.PathSensitivity import org.gradle.api.tasks.TaskAction import org.gradle.jvm.toolchain.JavaLauncher -import org.gradle.process.ExecOperations -import java.io.ByteArrayOutputStream import java.io.File -import javax.inject.Inject +import java.util.concurrent.TimeUnit internal const val SOURCES_MANIFEST_ATTRIBUTE = "Airflow-Java-SDK-Sources" internal const val SOURCES_JSON_PATH = "META-INF/airflow/sources.json" internal const val SOURCES_DIR_PATH = "META-INF/airflow/sources" +/** + * How long the describe run gets before the build gives up on it and packs + * the entrypoint's source alone. A `main` that does not reach + * `Server.create(args)` never answers `--describe-sources`, and one that + * starts serving instead would otherwise hold the build open. + */ +private const val DESCRIBE_TIMEOUT_SECONDS = 120L +private const val DRAIN_TIMEOUT_MILLIS = 2_000L + /** * Finds the source file of a class by its `SourceFile` attribute: the class * file's package path plus that name, looked up in each source directory. @@ -70,6 +77,14 @@ internal class SourceLocator( fun file(relativePath: String): File = sourceDirs.map { File(it, relativePath) }.first { it.isFile } } +/** + * The `META-INF/airflow/sources.json` body: `entrypoint_path` and each Dag ID + * in `dag_source_paths`. + * + * Every path is relative to the source directory the file was found in, which + * is also where it sits under `META-INF/airflow/sources/` in the JAR, so a + * reader resolves one by prefixing that directory. + */ internal fun sourcesJson( entrypointPath: String?, dagSourcePaths: Map, @@ -90,9 +105,6 @@ internal fun sourcesJson( */ @CacheableTask abstract class PackDagSources : DefaultTask() { - @get:Inject - abstract val execOperations: ExecOperations - @get:Input @get:Optional abstract val mainClass: Property @@ -161,20 +173,32 @@ abstract class PackDagSources : DefaultTask() { main: String, describe: File, ): Map { - val output = ByteArrayOutputStream() + var output = "" val failure = try { - val result = - execOperations.javaexec { spec -> - launcher.orNull?.let { spec.executable(it.executablePath.asFile.absolutePath) } - spec.classpath(classesDirs, runtimeClasspath) - spec.mainClass.set(main) - spec.args("--describe-sources", describe.absolutePath) - spec.isIgnoreExitValue = true - spec.standardOutput = output - spec.errorOutput = output - } - if (result.exitValue != 0) "exit code ${result.exitValue}" else null + val java = + launcher.orNull + ?.executablePath + ?.asFile + ?.absolutePath ?: "java" + val classpath = (classesDirs + runtimeClasspath).joinToString(File.pathSeparator) { it.absolutePath } + val process = + ProcessBuilder(java, "-cp", classpath, main, "--describe-sources", describe.absolutePath) + .redirectErrorStream(true) + .start() + // Drained on its own thread: a main that writes more than the pipe + // holds would otherwise block before the timeout could fire. + val drain = Thread { output = process.inputStream.bufferedReader().readText() } + drain.isDaemon = true + drain.start() + val finished = process.waitFor(DESCRIBE_TIMEOUT_SECONDS, TimeUnit.SECONDS) + if (!finished) process.destroyForcibly() + drain.join(DRAIN_TIMEOUT_MILLIS) + when { + !finished -> "it did not finish within $DESCRIBE_TIMEOUT_SECONDS seconds" + process.exitValue() != 0 -> "exit code ${process.exitValue()}" + else -> null + } } catch (e: Exception) { e.message ?: e.javaClass.simpleName } @@ -186,7 +210,7 @@ abstract class PackDagSources : DefaultTask() { } describeFailed = true val why = failure ?: "it wrote no valid --describe-sources file; does main call Server.serve?" - val log = output.toString().trim() + val log = output.trim() logger.warn( "Could not read each Dag's source from '{}' ({}); only its entrypoint source is packed.{}", main, diff --git a/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPluginTest.kt b/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPluginTest.kt index 91fb5c5d5ed45..05ed412887594 100644 --- a/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPluginTest.kt +++ b/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPluginTest.kt @@ -192,6 +192,29 @@ class AirflowSdkPluginTest { ) } + @Test + fun leavesTheSourcesPayloadOutOfANonBundleJar( + @TempDir dir: File, + ) { + project(dir) + dir.write( + "build.gradle", + File(dir, "build.gradle").readText() + "\njava { withSourcesJar() }\n", + ) + + gradle(dir, "jar", "sourcesJar") + + val sources = File(dir, "build/libs/bundle-test-sources.jar") + assertFalse(entries(sources).any { it.startsWith("META-INF/airflow") }) + JarFile(sources).use { assertNull(it.manifest.mainAttributes.getValue("Airflow-Java-SDK-Sources")) } + JarFile(File(dir, "build/libs/bundle-test.jar")).use { + assertEquals( + "META-INF/airflow/sources.json", + it.manifest.mainAttributes.getValue("Airflow-Java-SDK-Sources"), + ) + } + } + @Test fun skipsEverythingWithoutMainClass( @TempDir dir: File, diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt index c0caccf5d4717..2d8447af324fe 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt @@ -139,6 +139,9 @@ class Server private constructor( * The call returns when the coordinator closes the connection (normally after * one task-instance execution). * + * A [Server] built from `--describe-sources` writes that file and returns + * without connecting anywhere. + * * @param bundle Bundle containing all Dags this process can execute. * * @see [serveAsync] @@ -156,6 +159,9 @@ class Server private constructor( * coordinator closes the connection (normally after one task-instance * execution). The coroutine returns once both channels have been closed. * + * A [Server] built from `--describe-sources` writes that file and returns + * without connecting anywhere. + * * Use this variant when calling from an existing coroutine scope; use the * blocking [serve] from a plain `main` method. * From 10a695499b0621afa05dd6c09a32a2973908c8b5 Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Wed, 7 Oct 2026 16:41:15 +0000 Subject: [PATCH 04/10] Java SDK: Show each Dag's own source and find it outside the package layout The Java importer now passes the Dag ID to the source lookup, so the Code view reads dag_source_paths instead of always showing the entrypoint. The source locator falls back to a unique match on the SourceFile name when the package does not match the directory, using an index built only when the package path misses. An unresolvable entrypoint or Dag class is now logged as a warning. The describe task is no longer up to date after a failed run, so a transient failure does not stick until a source changes. Its hint now points at Server.create(args), which is what answers --describe-sources. Add a shadowJar test, correct the comment on which JARs carry the sources payload, and document that a Dag maps to the class that constructed it. --- .../0015-per-dag-source-in-bundle-artifact.md | 2 + .../language-sdks/java.rst | 4 + .../airflow/sdk/plugin/AirflowSdkPlugin.kt | 6 +- .../airflow/sdk/plugin/PackDagSources.kt | 32 ++++-- .../sdk/plugin/AirflowSdkPluginTest.kt | 105 +++++++++++++++++- .../sdk/coordinators/java/_dag_importer.py | 4 +- .../coordinators/java/test_dag_importer.py | 19 ++++ 7 files changed, 158 insertions(+), 14 deletions(-) diff --git a/airflow-core/adr/lang-sdk/0015-per-dag-source-in-bundle-artifact.md b/airflow-core/adr/lang-sdk/0015-per-dag-source-in-bundle-artifact.md index 602066d958f8c..902a4d43b9ce0 100644 --- a/airflow-core/adr/lang-sdk/0015-per-dag-source-in-bundle-artifact.md +++ b/airflow-core/adr/lang-sdk/0015-per-dag-source-in-bundle-artifact.md @@ -94,6 +94,8 @@ TypeScript is the reference ([#73723](https://github.com/apache/airflow/pull/737 packers adopt the same shape — per-Dag source resolution, de-duplication by path, and an always-embedded entrypoint fallback — so a bundle reader treats every language the same way and the Code tab behaves identically regardless of which SDK produced the artifact. +A JVM SDK has no module body, so it resolves a Dag by the class that declared it, and a Dag built by a +factory maps to the factory's class. Go has no module evaluation, so the Go SDK records the file that calls `airflow.Dag`. A Dag built in a factory function maps to the factory's file, not to the file that calls the factory. diff --git a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst index 12b65d85580f9..793eda354121a 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst @@ -1144,6 +1144,10 @@ into the bundle JAR, so the Airflow UI can show the source of a native Java Dag :ref:`java-sdk/native-dag-parsing`). To find those classes, the plugin runs ``mainClass`` once at build time. If that run fails, the build logs a warning and packs only the ``mainClass`` source. +A Dag is recorded against the class that constructed it. A Dag built by a factory therefore maps to the +factory's class, and one built by a factory in a dependency JAR has no source in the project and falls +back to the ``mainClass`` source. + .. note:: You only need the ``annotationProcessor`` entry if you use the annotation-based API. It is not needed for diff --git a/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt index c3bdbc4f8fa98..ed78aedec9400 100644 --- a/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt +++ b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt @@ -147,10 +147,8 @@ class AirflowSdkPlugin : Plugin { task.manifest.attributes(mapOf("Main-Class" to className)) } } - // Only the JARs a bundle is assembled from carry the Dag sources. A - // sources or javadoc JAR would claim the attribute without being a - // bundle, and the Python side picks the bundle JAR out of a directory - // by that attribute. + // The sources payload belongs only in the JARs a bundle is assembled from, + // so a sources or javadoc JAR does not carry it. if (ext.mainClass.isPresent && task.name in BUNDLE_JAR_TASKS) { task.dependsOn(packTask) task.from(packTask.flatMap { it.sourcesDir }) diff --git a/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/PackDagSources.kt b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/PackDagSources.kt index 3ba8e98fc15f5..f422fb7671bb5 100644 --- a/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/PackDagSources.kt +++ b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/PackDagSources.kt @@ -56,22 +56,35 @@ private const val DRAIN_TIMEOUT_MILLIS = 2_000L /** * Finds the source file of a class by its `SourceFile` attribute: the class - * file's package path plus that name, looked up in each source directory. + * file's package path plus that name, looked up in each source directory, or + * else the one source file of that name when the layout does not match the + * package. */ internal class SourceLocator( private val classesDirs: Iterable, private val sourceDirs: Iterable, ) { + /** Every source file under the source dirs by file name, walked once and only if a lookup misses. */ + private val pathsByName: Map> by lazy { + sourceDirs + .flatMap { dir -> dir.walkTopDown().filter(File::isFile).map { it.relativeTo(dir).invariantSeparatorsPath } } + .groupBy { it.substringAfterLast('/') } + } + /** Path of the source file relative to its source directory, or `null` if it cannot be found. */ fun locate(className: String): String? { - val packagePath = className.substringBeforeLast('.', "").replace('.', '/') val classFile = classesDirs .map { File(it, className.replace('.', '/') + ".class") } .firstOrNull { it.isFile } ?: return null - val sourceFile = readSourceFile(classFile.readBytes())?.takeIf { it.isNotEmpty() && '/' !in it && '\\' !in it } - val relative = listOfNotNull(packagePath.ifEmpty { null }, sourceFile ?: return null).joinToString("/") - return relative.takeIf { path -> sourceDirs.any { File(it, path).isFile } } + val sourceFile = + readSourceFile(classFile.readBytes())?.takeIf { it.isNotEmpty() && '/' !in it && '\\' !in it } ?: return null + val packagePath = className.substringBeforeLast('.', "").replace('.', '/') + val byPackage = listOfNotNull(packagePath.ifEmpty { null }, sourceFile).joinToString("/") + if (sourceDirs.any { File(it, byPackage).isFile }) return byPackage + // Neither javac nor Kotlin requires the package to match the directory, so fall back to the + // file name when exactly one source tree holds it. An ambiguous name is left unresolved. + return pathsByName[sourceFile]?.distinct()?.singleOrNull() } fun file(relativePath: String): File = sourceDirs.map { File(it, relativePath) }.first { it.isFile } @@ -133,6 +146,8 @@ abstract class PackDagSources : DefaultTask() { init { outputs.doNotCacheIf("the Dag describe run failed") { describeFailed } + // A failed run leaves describeFile missing, which Gradle would otherwise read as unchanged. + outputs.upToDateWhen { describeFile.get().asFile.isFile } } @TaskAction @@ -148,11 +163,14 @@ abstract class PackDagSources : DefaultTask() { val locator = SourceLocator(classesDirs.files, sourceDirs.files) val entrypoint = locator.locate(main) + if (entrypoint == null) { + logger.warn("No source file found for entrypoint class {}; the Code view will show no source for this JAR", main) + } val dagPaths = linkedMapOf() declaringClasses.forEach { (dagId, className) -> val path = locator.locate(className) if (path == null) { - logger.info("No source file found for class {} of Dag '{}'; it falls back to the entrypoint", className, dagId) + logger.warn("No source file found for class {} of Dag '{}'; it falls back to the entrypoint", className, dagId) } else { dagPaths[dagId] = path } @@ -209,7 +227,7 @@ abstract class PackDagSources : DefaultTask() { } } describeFailed = true - val why = failure ?: "it wrote no valid --describe-sources file; does main call Server.serve?" + val why = failure ?: "it wrote no valid --describe-sources file; does main pass its args to Server.create?" val log = output.trim() logger.warn( "Could not read each Dag's source from '{}' ({}); only its entrypoint source is packed.{}", diff --git a/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPluginTest.kt b/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPluginTest.kt index 05ed412887594..b00c23b0fc932 100644 --- a/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPluginTest.kt +++ b/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPluginTest.kt @@ -25,12 +25,16 @@ import org.gradle.testkit.runner.GradleRunner import org.gradle.testkit.runner.TaskOutcome import org.junit.jupiter.api.Assertions.assertEquals import org.junit.jupiter.api.Assertions.assertFalse +import org.junit.jupiter.api.Assertions.assertNotEquals import org.junit.jupiter.api.Assertions.assertNull import org.junit.jupiter.api.Assertions.assertTrue import org.junit.jupiter.api.Test import org.junit.jupiter.api.io.TempDir import java.io.File +import java.util.jar.Attributes import java.util.jar.JarFile +import java.util.jar.JarOutputStream +import java.util.jar.Manifest class AirflowSdkPluginTest { private fun main(describeBody: String) = @@ -185,13 +189,112 @@ class AirflowSdkPluginTest { project(dir, mainBody = "") val result = gradle(dir, "jar") - assertTrue(result.output.contains("does main call Server.serve?"), result.output) + assertTrue(result.output.contains("does main pass its args to Server.create?"), result.output) assertEquals( mapOf("entrypoint_path" to "com/example/Main.java", "dag_source_paths" to emptyMap()), sourcesJson(File(dir, "build/libs/bundle-test.jar")), ) } + @Test + fun findsASourceWhosePackageDoesNotMatchItsDirectory( + @TempDir dir: File, + ) { + project(dir) + File(dir, "src/extra/com/example/dags/Reports.java").delete() + dir.write("src/main/java/dags/Reports.java", "package com.example.dags; public class Reports {}") + gradle(dir, "jar") + + assertEquals( + mapOf( + "entrypoint_path" to "com/example/Main.java", + "dag_source_paths" to + mapOf( + "orders" to "com/example/Main.java", + "reports" to "dags/Reports.java", + "reports_backfill" to "dags/Reports.java", + ), + ), + sourcesJson(File(dir, "build/libs/bundle-test.jar")), + ) + } + + @Test + fun warnsWhenTheEntrypointSourceCannotBeFound( + @TempDir dir: File, + ) { + project(dir) + dir.write( + "build.gradle", + File(dir, "build.gradle").readText() + + "\ncompileJava.doLast { delete 'src/main/java/com/example/Main.java' }\n", + ) + val result = gradle(dir, "jar") + + assertTrue(result.output.contains("No source file found for entrypoint class com.example.Main"), result.output) + assertNull(sourcesJson(File(dir, "build/libs/bundle-test.jar"))["entrypoint_path"]) + } + + @Test + fun rerunsAfterAFailedDescribeRun( + @TempDir dir: File, + ) { + project(dir, mainBody = "") + gradle(dir, "jar") + val again = gradle(dir, "jar") + + assertNotEquals(TaskOutcome.UP_TO_DATE, again.task(":packDagSources")!!.outcome) + assertTrue(again.output.contains("only its entrypoint source is packed"), again.output) + } + + @Test + fun shadowJarCarriesTheSourcesPayloadByDefault( + @TempDir dir: File, + ) { + dir.write("settings.gradle", "rootProject.name = 'bundle-test'\n") + dir.write( + "build.gradle", + """ + plugins { id 'org.apache.airflow.sdk' } + repositories { maven { url = uri('repo') } } + dependencies { implementation 'org.apache.airflow:airflow-sdk:1.0' } + airflowBundle { mainClass = 'com.example.Main' } + """.trimIndent(), + ) + dir.write("src/main/java/com/example/Main.java", main(describeDags)) + dir.write("src/main/java/com/example/dags/Reports.java", "package com.example.dags; public class Reports {}") + stubSdk(File(dir, "repo/org/apache/airflow/airflow-sdk/1.0")) + + gradle(dir, "shadowJar") + + val jar = File(dir, "build/libs/bundle-test-all.jar") + assertTrue(entries(jar).contains("META-INF/airflow/sources.json")) + JarFile(jar).use { + assertEquals("META-INF/airflow/sources.json", it.manifest.mainAttributes.getValue("Airflow-Java-SDK-Sources")) + assertEquals("1", it.manifest.mainAttributes.getValue("Airflow-Supervisor-Schema-Version")) + } + } + + private fun stubSdk(dir: File) { + dir.mkdirs() + File(dir, "airflow-sdk-1.0.pom").writeText( + """ + + 4.0.0 + org.apache.airflow + airflow-sdk + 1.0 + + """.trimIndent(), + ) + val manifest = + Manifest().apply { + mainAttributes[Attributes.Name.MANIFEST_VERSION] = "1.0" + mainAttributes.putValue("Airflow-Supervisor-Schema-Version", "1") + } + JarOutputStream(File(dir, "airflow-sdk-1.0.jar").outputStream(), manifest).close() + } + @Test fun leavesTheSourcesPayloadOutOfANonBundleJar( @TempDir dir: File, diff --git a/task-sdk/src/airflow/sdk/coordinators/java/_dag_importer.py b/task-sdk/src/airflow/sdk/coordinators/java/_dag_importer.py index 6a20f3d60182b..b1e9283d26e4f 100644 --- a/task-sdk/src/airflow/sdk/coordinators/java/_dag_importer.py +++ b/task-sdk/src/airflow/sdk/coordinators/java/_dag_importer.py @@ -74,9 +74,9 @@ def might_contain_dag(self, definition: DagDefinition, safe_mode: bool) -> bool: return not wanted or main_class == wanted def get_source_code(self, definition: DagDefinition, dag_id: str | None = None) -> DagSourceCode: - """Return the entrypoint source the JAR embeds, or a notice when it embeds none.""" + """Return the source the JAR embeds for *dag_id*, else the entrypoint's, else a notice.""" with definition.as_file() as path, zipfile.ZipFile(path) as zf: - info = _find_source_entry(zf) + info = _find_source_entry(zf, dag_id) if info is None: return DagSourceCode(source_code=_NO_SOURCE, language="java") if info.file_size > _MAX_SOURCE_BYTES: diff --git a/task-sdk/tests/task_sdk/coordinators/java/test_dag_importer.py b/task-sdk/tests/task_sdk/coordinators/java/test_dag_importer.py index 45ecda314eedb..e3952299fbfaa 100644 --- a/task-sdk/tests/task_sdk/coordinators/java/test_dag_importer.py +++ b/task-sdk/tests/task_sdk/coordinators/java/test_dag_importer.py @@ -227,6 +227,25 @@ def test_maps_a_dag_to_its_file_with_the_entrypoint_as_fallback(self, tmp_path, with zipfile.ZipFile(jar) as zf: assert _find_source_entry(zf, dag_id).filename == f"META-INF/airflow/sources/{expected}" + @pytest.mark.parametrize( + ("dag_id", "expected"), + [ + pytest.param("reports", "class Reports", id="mapped"), + pytest.param("orders", "class Dags", id="unmapped"), + pytest.param(None, "class Dags", id="no-dag-id"), + ], + ) + def test_get_source_code_reads_the_source_of_the_dag(self, tmp_path, dag_id, expected): + jar = self._jar( + tmp_path, + index={"entrypoint_path": ENTRYPOINT, "dag_source_paths": {"reports": REPORTS}}, + files={ENTRYPOINT: "class Dags", REPORTS: "class Reports"}, + ) + + result = _importer().get_source_code(_definition(jar), dag_id) + + assert result.source_code == expected + @pytest.mark.parametrize( ("attributes", "index", "files"), [ From b18c0aa45c93d8e5090e03c37835c153edf9e2be Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Thu, 8 Oct 2026 11:10:33 +0000 Subject: [PATCH 05/10] Java SDK: Use a describe run's output even when the run itself failed --- .../airflow/sdk/plugin/PackDagSources.kt | 22 ++++++++++++++----- .../sdk/plugin/AirflowSdkPluginTest.kt | 18 +++++++++++++++ 2 files changed, 35 insertions(+), 5 deletions(-) diff --git a/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/PackDagSources.kt b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/PackDagSources.kt index f422fb7671bb5..eeddfab75211d 100644 --- a/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/PackDagSources.kt +++ b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/PackDagSources.kt @@ -220,15 +220,27 @@ abstract class PackDagSources : DefaultTask() { } catch (e: Exception) { e.message ?: e.javaClass.simpleName } - if (failure == null && describe.isFile) { - @Suppress("UNCHECKED_CAST") - (runCatching { JsonSlurper().parse(describe) }.getOrNull() as? Map)?.let { parsed -> - return parsed.mapNotNull { (id, cls) -> (cls as? String)?.let { id to it } }.toMap(linkedMapOf()) + val log = output.trim() + val parsed = + describe + .takeIf { it.isFile } + ?.let { + @Suppress("UNCHECKED_CAST") + (runCatching { JsonSlurper().parse(it) }.getOrNull() as? Map) + } + if (parsed != null) { + if (failure != null) { + logger.warn( + "'{}' wrote each Dag's source but {}; its sources are used anyway.{}", + main, + failure, + if (log.isEmpty()) "" else "\n$log", + ) } + return parsed.mapNotNull { (id, cls) -> (cls as? String)?.let { id to it } }.toMap(linkedMapOf()) } describeFailed = true val why = failure ?: "it wrote no valid --describe-sources file; does main pass its args to Server.create?" - val log = output.trim() logger.warn( "Could not read each Dag's source from '{}' ({}); only its entrypoint source is packed.{}", main, diff --git a/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPluginTest.kt b/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPluginTest.kt index b00c23b0fc932..574761838e0ba 100644 --- a/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPluginTest.kt +++ b/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPluginTest.kt @@ -182,6 +182,24 @@ class AirflowSdkPluginTest { ) } + @Test + fun usesTheDescribeFileEvenWhenTheRunExitsNonZero( + @TempDir dir: File, + ) { + project(dir, mainBody = describeDags + "\nSystem.exit(3);") + val result = gradle(dir, "jar") + + assertTrue(result.output.contains("its sources are used anyway"), result.output) + assertEquals( + mapOf( + "orders" to "com/example/Main.java", + "reports" to "com/example/dags/Reports.java", + "reports_backfill" to "com/example/dags/Reports.java", + ), + sourcesJson(File(dir, "build/libs/bundle-test.jar"))["dag_source_paths"], + ) + } + @Test fun fallsBackToEntrypointOnlyWhenMainWritesNothing( @TempDir dir: File, From 88934351785a2e1e65723a8a3fbd376a2c5a0384 Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Thu, 8 Oct 2026 11:11:26 +0000 Subject: [PATCH 06/10] Java SDK: Serve and describe sources from their own Server classes --- .../kotlin/org/apache/airflow/sdk/Server.kt | 115 +++++++++--------- .../org/apache/airflow/sdk/ServerTest.kt | 6 +- 2 files changed, 61 insertions(+), 60 deletions(-) diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt index 2d8447af324fe..5847f5b7326ef 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt @@ -91,45 +91,7 @@ class ApiError( * The process exits when the coordinator closes the connection (normally after * one task-instance execution). */ -class Server private constructor( - private val comm: InetSocketAddress?, - private val logs: InetSocketAddress?, - private val describeSources: File?, -) { - constructor(comm: InetSocketAddress, logs: InetSocketAddress) : this(comm, logs, null) - - companion object { - /** - * Parses coordinator addresses from command-line arguments and returns a - * ready-to-use [Server]. - * - * The arguments are supplied automatically by Airflow and are not intended - * to be constructed by hand: - * - * * `--comm host:port` address for task-execution messages. - * * `--logs host:port` address for log forwarding. - * * `--describe-sources file` makes [serve] write a JSON object mapping each - * Java-declared Dag ID to the binary name of its declaring class into `file` - * and return without connecting anywhere; `--comm` and `--logs` are not - * required then. It is for the Gradle plugin, not for users. - * - * @param args Command-line arguments as received by `main`. - * @return A configured [Server] ready to call [serve]. - */ - @JvmStatic - fun create(args: Array): Server { - val args = ArgParser(args).parseInto(::Args) - args.describeSources?.let { return Server(args.comm, args.logs, it) } - return Server( - args.comm ?: throw MissingValueException("--comm"), - args.logs ?: throw MissingValueException("--logs"), - null, - ) - } - } - - private val logger = Logger(Server::class) - +sealed class Server { /** * Blocking entry point: connects to the coordinator and serves task-execution * requests from the given [bundle]. @@ -139,7 +101,7 @@ class Server private constructor( * The call returns when the coordinator closes the connection (normally after * one task-instance execution). * - * A [Server] built from `--describe-sources` writes that file and returns + * A [Server] created from `--describe-sources` writes that file and returns * without connecting anywhere. * * @param bundle Bundle containing all Dags this process can execute. @@ -159,7 +121,7 @@ class Server private constructor( * coordinator closes the connection (normally after one task-instance * execution). The coroutine returns once both channels have been closed. * - * A [Server] built from `--describe-sources` writes that file and returns + * A [Server] created from `--describe-sources` writes that file and returns * without connecting anywhere. * * Use this variant when calling from an existing coroutine scope; use the @@ -169,19 +131,67 @@ class Server private constructor( * * @see [serve] */ - suspend fun serveAsync(bundle: Bundle) = + abstract suspend fun serveAsync(bundle: Bundle) + + companion object { + /** + * Parses coordinator addresses from command-line arguments and returns a + * ready-to-use [Server]. + * + * The arguments are supplied automatically by Airflow and are not intended + * to be constructed by hand: + * + * * `--comm host:port` address for task-execution messages. + * * `--logs host:port` address for log forwarding. + * * `--describe-sources file` makes [serve] write a JSON object mapping each + * Java-declared Dag ID to the binary name of its declaring class into `file` + * and return without connecting anywhere; `--comm` and `--logs` are not + * required then. It is for the Gradle plugin, not for users. + * + * @param args Command-line arguments as received by `main`. + * @return A configured [Server] ready to call [serve]. + */ + @JvmStatic + fun create(args: Array): Server { + val args = ArgParser(args).parseInto(::Args) + args.describeSources?.let { return SourceDescriber(it) } + return CoordinatorServer( + args.comm ?: throw MissingValueException("--comm"), + args.logs ?: throw MissingValueException("--logs"), + ) + } + } +} + +/** Writes each Java-declared Dag's declaring class to [target] instead of serving. */ +internal class SourceDescriber( + private val target: File, +) : Server() { + override suspend fun serveAsync(bundle: Bundle) { + bundle.finalizeRegistration() + val sources = linkedMapOf() + bundle.dags.values.forEach { dag -> dag.declaringClass?.let { sources[dag.id] = it.name } } + target.absoluteFile.parentFile?.mkdirs() + ObjectMapper().writeValue(target, sources) + } +} + +/** Serves task-execution requests over the coordinator's [comm] and [logs] sockets. */ +class CoordinatorServer( + private val comm: InetSocketAddress, + private val logs: InetSocketAddress, +) : Server() { + private val logger = Logger(CoordinatorServer::class) + + override suspend fun serveAsync(bundle: Bundle) { coroutineScope { bundle.finalizeRegistration() - if (describeSources != null) { - writeSources(bundle, describeSources) - return@coroutineScope - } val deferral = CompletableDeferred() launch { try { SelectorManager(Dispatchers.IO).use { selector -> - aSocket(selector).tcp().connect(comm!!).use { socket -> + aSocket(selector).tcp().connect(comm).use { socket -> logger.debug("Connected comm", mapOf("addr" to comm)) CoordinatorComm( socket.openReadChannel(), @@ -197,7 +207,7 @@ class Server private constructor( } launch { SelectorManager(Dispatchers.IO).use { selector -> - aSocket(selector).tcp().connect(logs!!).use { socket -> + aSocket(selector).tcp().connect(logs).use { socket -> logger.debug("Connected logs", mapOf("addr" to logs)) LogSender.configure(socket.openWriteChannel(autoFlush = true)) deferral.await() @@ -205,15 +215,6 @@ class Server private constructor( } } } - - private fun writeSources( - bundle: Bundle, - target: File, - ) { - val sources = linkedMapOf() - bundle.dags.values.forEach { dag -> dag.declaringClass?.let { sources[dag.id] = it.name } } - target.absoluteFile.parentFile?.mkdirs() - ObjectMapper().writeValue(target, sources) } internal suspend fun dispatchTask( diff --git a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ServerTest.kt b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ServerTest.kt index d6e93fe211421..7c7ab91b2ab22 100644 --- a/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ServerTest.kt +++ b/java-sdk/sdk/src/test/kotlin/org/apache/airflow/sdk/ServerTest.kt @@ -76,7 +76,7 @@ class ServerTest { val toServer = ByteChannel(autoFlush = true) val fromServer = ByteChannel(autoFlush = true) val comm = CoordinatorComm(toServer, fromServer) - val server = Server(InetSocketAddress("localhost", 0), InetSocketAddress("localhost", 0)) + val server = CoordinatorServer(InetSocketAddress("localhost", 0), InetSocketAddress("localhost", 0)) val reported = ArrayBlockingQueue(1) val supervisor = @@ -110,7 +110,7 @@ class ServerTest { val toServer = ByteChannel(autoFlush = true) val fromServer = ByteChannel(autoFlush = true) val comm = CoordinatorComm(toServer, fromServer) - val server = Server(InetSocketAddress("localhost", 0), InetSocketAddress("localhost", 0)) + val server = CoordinatorServer(InetSocketAddress("localhost", 0), InetSocketAddress("localhost", 0)) Assertions.assertThrows(ApiError::class.java) { runBlocking { @@ -128,7 +128,7 @@ class ServerTest { val toServer = ByteChannel(autoFlush = true) val fromServer = ByteChannel(autoFlush = true) val comm = CoordinatorComm(toServer, fromServer) - val server = Server(InetSocketAddress("localhost", 0), InetSocketAddress("localhost", 0)) + val server = CoordinatorServer(InetSocketAddress("localhost", 0), InetSocketAddress("localhost", 0)) val bundle = Bundle(listOf(DagDef("parsed_dag").addTask(TaskDef("t", ServerNoopTask::class.java)))) val reported = ArrayBlockingQueue(1) From 62dd96e93bbdace2cc6cf5acc59271871ca3eec5 Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Thu, 8 Oct 2026 11:12:00 +0000 Subject: [PATCH 07/10] Java SDK: Pack the sources into the one JAR the bundle is built from --- .../airflow/sdk/plugin/AirflowSdkPlugin.kt | 21 ++++++++++--------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt index ed78aedec9400..fa0e01ec8e543 100644 --- a/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt +++ b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt @@ -65,9 +65,6 @@ abstract class AirflowBundleExtension { abstract val fatJar: Property } -/** The JAR tasks a deployable bundle is assembled from, fat or thin. */ -private val BUNDLE_JAR_TASKS = setOf("jar", "shadowJar") - /** * Gradle plugin for building Apache Airflow Java SDK bundles. * @@ -147,13 +144,6 @@ class AirflowSdkPlugin : Plugin { task.manifest.attributes(mapOf("Main-Class" to className)) } } - // The sources payload belongs only in the JARs a bundle is assembled from, - // so a sources or javadoc JAR does not carry it. - if (ext.mainClass.isPresent && task.name in BUNDLE_JAR_TASKS) { - task.dependsOn(packTask) - task.from(packTask.flatMap { it.sourcesDir }) - task.manifest.attributes(mapOf(SOURCES_MANIFEST_ATTRIBUTE to SOURCES_JSON_PATH)) - } } val classFiles = @@ -199,8 +189,17 @@ class AirflowSdkPlugin : Plugin { } } + fun packSourcesInto(jarTask: String) { + if (!ext.mainClass.isPresent) return + project.tasks.named(jarTask, Jar::class.java).configure { task -> + task.from(packTask.flatMap { it.sourcesDir }) + task.manifest.attributes(mapOf(SOURCES_MANIFEST_ATTRIBUTE to SOURCES_JSON_PATH)) + } + } + if (ext.fatJar.get()) { project.plugins.apply("com.gradleup.shadow") + packSourcesInto("shadowJar") val schemaVersionProvider = project.providers.provider { @@ -242,6 +241,8 @@ class AirflowSdkPlugin : Plugin { task.into(project.layout.buildDirectory.dir("bundle")) } } else { + packSourcesInto("jar") + // bundle copies the thin JAR and all runtime dependency JARs into // build/bundle/, mirroring what installDist puts in lib/. project.tasks.register("bundle", Copy::class.java) { task -> From cb3dc2d48f2db255f72802b697f926e99d1556f6 Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Thu, 8 Oct 2026 11:12:09 +0000 Subject: [PATCH 08/10] Java SDK: Attach the plugin's documentation to the plugin class --- .../kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt | 1 - 1 file changed, 1 deletion(-) diff --git a/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt index fa0e01ec8e543..3117dcdfa6e79 100644 --- a/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt +++ b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt @@ -108,7 +108,6 @@ abstract class AirflowBundleExtension { * the `airflow-sdk` JAR instead. The bundle JAR still contains `Main-Class`. * */ - class AirflowSdkPlugin : Plugin { override fun apply(project: Project) { project.plugins.apply("java") From ba08f8e4fab54a0b636e50845787290615177607 Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Thu, 8 Oct 2026 11:12:15 +0000 Subject: [PATCH 09/10] Java SDK: Drop the duplicated classes dir from the describe classpath --- .../kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt index 3117dcdfa6e79..3b715b6011b97 100644 --- a/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt +++ b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt @@ -126,7 +126,7 @@ class AirflowSdkPlugin : Plugin { task.onlyIf { ext.mainClass.isPresent } task.mainClass.set(ext.mainClass) task.classesDirs.from(main.output.classesDirs) - task.runtimeClasspath.from(main.output, main.runtimeClasspath) + task.runtimeClasspath.from(main.runtimeClasspath) task.sourceDirs.from(main.allSource.srcDirs) task.launcher.convention( project.extensions From a073f81ca682fc2b740c792f595b706014832185 Mon Sep 17 00:00:00 2001 From: ZHE YOU LIU Date: Thu, 8 Oct 2026 11:13:03 +0000 Subject: [PATCH 10/10] Java SDK: Let a bundle set how long the describe run may take --- .../airflow/sdk/plugin/AirflowSdkPlugin.kt | 8 +++++++ .../airflow/sdk/plugin/PackDagSources.kt | 22 +++++++++++-------- .../sdk/plugin/AirflowSdkPluginTest.kt | 14 ++++++++++++ 3 files changed, 35 insertions(+), 9 deletions(-) diff --git a/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt index 3b715b6011b97..5e3a41b5ce00a 100644 --- a/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt +++ b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt @@ -31,6 +31,7 @@ import org.gradle.api.tasks.bundling.Jar import org.gradle.jvm.toolchain.JavaToolchainService import java.lang.reflect.Modifier import java.net.URLClassLoader +import java.time.Duration import java.util.jar.JarFile import kotlin.jvm.java @@ -41,6 +42,7 @@ import kotlin.jvm.java * airflowBundle { * mainClass = "com.example.ExampleBundleBuilder" * // fatJar = false // opt out of shadow JAR creation + * // describeTimeout = Duration.ofSeconds(60) * } * ``` */ @@ -63,6 +65,10 @@ abstract class AirflowBundleExtension { */ @get:Input abstract val fatJar: Property + + /** How long the describe run may take before the build gives up and packs only the entrypoint's source. */ + @get:Input + abstract val describeTimeout: Property } /** @@ -114,6 +120,7 @@ class AirflowSdkPlugin : Plugin { val ext = project.extensions.create("airflowBundle", AirflowBundleExtension::class.java) ext.fatJar.convention(true) + ext.describeTimeout.convention(Duration.ofSeconds(30)) project.afterEvaluate { val main = project.extensions.getByType(SourceSetContainer::class.java).getByName("main") @@ -125,6 +132,7 @@ class AirflowSdkPlugin : Plugin { task.dependsOn(project.tasks.named("classes")) task.onlyIf { ext.mainClass.isPresent } task.mainClass.set(ext.mainClass) + task.describeTimeout.set(ext.describeTimeout) task.classesDirs.from(main.output.classesDirs) task.runtimeClasspath.from(main.runtimeClasspath) task.sourceDirs.from(main.allSource.srcDirs) diff --git a/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/PackDagSources.kt b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/PackDagSources.kt index eeddfab75211d..63bafb05de94b 100644 --- a/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/PackDagSources.kt +++ b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/PackDagSources.kt @@ -39,19 +39,13 @@ import org.gradle.api.tasks.PathSensitivity import org.gradle.api.tasks.TaskAction import org.gradle.jvm.toolchain.JavaLauncher import java.io.File +import java.time.Duration import java.util.concurrent.TimeUnit internal const val SOURCES_MANIFEST_ATTRIBUTE = "Airflow-Java-SDK-Sources" internal const val SOURCES_JSON_PATH = "META-INF/airflow/sources.json" internal const val SOURCES_DIR_PATH = "META-INF/airflow/sources" -/** - * How long the describe run gets before the build gives up on it and packs - * the entrypoint's source alone. A `main` that does not reach - * `Server.create(args)` never answers `--describe-sources`, and one that - * starts serving instead would otherwise hold the build open. - */ -private const val DESCRIBE_TIMEOUT_SECONDS = 120L private const val DRAIN_TIMEOUT_MILLIS = 2_000L /** @@ -122,6 +116,15 @@ abstract class PackDagSources : DefaultTask() { @get:Optional abstract val mainClass: Property + /** + * How long the describe run gets before the build gives up on it and packs + * the entrypoint's source alone. A `main` that does not reach + * `Server.create(args)` never answers `--describe-sources`, and one that + * starts serving instead would otherwise hold the build open. + */ + @get:Input + abstract val describeTimeout: Property + @get:Classpath abstract val classesDirs: ConfigurableFileCollection @@ -209,11 +212,12 @@ abstract class PackDagSources : DefaultTask() { val drain = Thread { output = process.inputStream.bufferedReader().readText() } drain.isDaemon = true drain.start() - val finished = process.waitFor(DESCRIBE_TIMEOUT_SECONDS, TimeUnit.SECONDS) + val timeout = describeTimeout.get() + val finished = process.waitFor(timeout.toMillis(), TimeUnit.MILLISECONDS) if (!finished) process.destroyForcibly() drain.join(DRAIN_TIMEOUT_MILLIS) when { - !finished -> "it did not finish within $DESCRIBE_TIMEOUT_SECONDS seconds" + !finished -> "it did not finish within ${timeout.seconds} seconds" process.exitValue() != 0 -> "exit code ${process.exitValue()}" else -> null } diff --git a/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPluginTest.kt b/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPluginTest.kt index 574761838e0ba..3297ae940db79 100644 --- a/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPluginTest.kt +++ b/java-sdk/plugin/src/test/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPluginTest.kt @@ -200,6 +200,20 @@ class AirflowSdkPluginTest { ) } + @Test + fun givesUpOnADescribeRunThatOutlastsTheConfiguredTimeout( + @TempDir dir: File, + ) { + project(dir, mainBody = "Thread.sleep(60000);") + dir.write( + "build.gradle", + File(dir, "build.gradle").readText() + "\nairflowBundle { describeTimeout = java.time.Duration.ofSeconds(2) }\n", + ) + val result = gradle(dir, "jar") + + assertTrue(result.output.contains("it did not finish within 2 seconds"), result.output) + } + @Test fun fallsBackToEntrypointOnlyWhenMainWritesNothing( @TempDir dir: File,