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