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/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..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 @@ -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 + + /** 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 } /** @@ -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 { 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 { 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 { } } + 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 { 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 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..63bafb05de94b --- /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, + 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 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 = + 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:Input + @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 + + @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 } + // 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() + 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 { + 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) + } + 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 0000000000000..3297ae940db79 --- /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 = + 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 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()), + 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, + ) { + 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 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) + } 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..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 @@ -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): 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): 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() val deferral = CompletableDeferred() @@ -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 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()) } + } +} 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) 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"), [