This is an automated email from the ASF dual-hosted git repository. jason810496 pushed a commit to branch jason/java-sdk/pack-dag-sources in repository https://gitbox.apache.org/repos/asf/airflow.git
commit 11c4637de7539827cb7a40d2b4a0a6a112376633 Author: ZHE YOU LIU <[email protected]> AuthorDate: Fri Oct 2 15:19:36 2026 +0000 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. --- .../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 | 58 +++++++-- .../org/apache/airflow/sdk/internal/DagSource.kt | 66 +++++++++++ .../apache/airflow/example/DagSourceFixtures.kt | 51 ++++++++ .../kotlin/org/apache/airflow/sdk/DagSourceTest.kt | 132 +++++++++++++++++++++ 7 files changed, 379 insertions(+), 23 deletions(-) diff --git a/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt b/java-sdk/processor/src/main/kotlin/org/apache/airflow/sdk/BuilderProcessor.kt index f97f72997d0..86aa54662ae 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.Refs @@ -209,7 +210,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) } @@ -646,6 +653,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 d4c8a5cf455..c005fa48dab 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"), new TestExample.Wiring()::depends); } @@ -203,11 +207,12 @@ class BuilderTest { import org.apache.airflow.sdk.DagDef; import org.apache.airflow.sdk.Task; import org.apache.airflow.sdk.TaskDef; + import org.apache.airflow.sdk.internal.DagSource; 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); dag.addTask(new TaskDef("t", T.class)); return dag; } @@ -266,12 +271,13 @@ class BuilderTest { import org.apache.airflow.sdk.DagDef; import org.apache.airflow.sdk.Task; import org.apache.airflow.sdk.TaskDef; + import org.apache.airflow.sdk.internal.DagSource; 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); dag.addTask(new TaskDef("t", T.class)); return dag; } @@ -421,11 +427,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); @@ -505,11 +512,12 @@ class BuilderTest { import org.apache.airflow.sdk.DagDef; import org.apache.airflow.sdk.Task; import org.apache.airflow.sdk.TaskDef; + import org.apache.airflow.sdk.internal.DagSource; 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); dag.addTask(new TaskDef("t", T.class)); return dag; } @@ -569,11 +577,12 @@ class BuilderTest { import org.apache.airflow.sdk.Task; import org.apache.airflow.sdk.TaskDef; import org.apache.airflow.sdk.internal.ArgValues; + import org.apache.airflow.sdk.internal.DagSource; 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); dag.addTask(new TaskDef("flat", Flat.class)); dag.addTask(new TaskDef("named", Named.class)); return dag; @@ -786,8 +795,9 @@ class BuilderTest { """ package org.apache.airflow.example; import org.apache.airflow.sdk.DagDef; + import org.apache.airflow.sdk.internal.DagSource; public final class TestExampleBuilder { - public static DagDef build() { var dag = new DagDef("foo"); return dag; } + public static DagDef build() { var dag = DagSource.declaredBy(new DagDef("foo"), TestExample.class); return dag; } } """, ) @@ -811,8 +821,9 @@ class BuilderTest { """ package org.apache.airflow.example; import org.apache.airflow.sdk.DagDef; + import org.apache.airflow.sdk.internal.DagSource; public final class Foo { - public static DagDef build() { var dag = new DagDef("TestExample"); return dag; } + public static DagDef build() { var dag = DagSource.declaredBy(new DagDef("TestExample"), TestExample.class); return dag; } } """, ) @@ -854,11 +865,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"), new TestExample.Wiring()::depends); } @@ -1131,10 +1143,11 @@ class BuilderTest { import org.apache.airflow.sdk.Task; import org.apache.airflow.sdk.TaskDef; import org.apache.airflow.sdk.internal.ArgValues; + import org.apache.airflow.sdk.internal.DagSource; public final class TestExampleBuilder { public static DagDef build() { - var dag = new DagDef("TestExample"); + var dag = DagSource.declaredBy(new DagDef("TestExample"), TestExample.class); dag.addTask(new TaskDef("score", Score.class)); return dag; } @@ -1313,4 +1326,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 c579ddd517b..5d9b918f614 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.SchemaFields import org.apache.airflow.sdk.internal.checkConfigValue import org.apache.airflow.sdk.internal.validateTaskInput @@ -50,6 +51,9 @@ class DagDef( internal val tasks = linkedMapOf<String, TaskDef>() internal val dagConfig = linkedMapOf<String, Any>() + /** 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..c0caccf5d47 100644 --- a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt +++ b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/Server.kt @@ -19,7 +19,10 @@ package org.apache.airflow.sdk +import com.fasterxml.jackson.databind.ObjectMapper import com.xenomachina.argparser.ArgParser +import com.xenomachina.argparser.MissingValueException +import com.xenomachina.argparser.default import io.ktor.network.selector.SelectorManager import io.ktor.network.sockets.InetSocketAddress import io.ktor.network.sockets.aSocket @@ -38,6 +41,7 @@ import org.apache.airflow.sdk.execution.comm.ErrorResponse import org.apache.airflow.sdk.execution.comm.StartupDetails import org.apache.airflow.sdk.execution.parseDags import org.apache.airflow.sdk.execution.runTask +import java.io.File import kotlin.text.substringAfterLast import kotlin.text.substringBeforeLast @@ -50,12 +54,16 @@ private class Args( address.substringAfterLast(':').toInt(), ) - val comm by parser.storing("--comm", help = "Address (host:port) to communicate with parent") { - parseAddress(this) - } - val logs by parser.storing("--logs", help = "Address (host:port) to send Airflow logs to") { - parseAddress(this) - } + val comm by parser + .storing("--comm", help = "Address (host:port) to communicate with parent") { parseAddress(this) } + .default(null) + val logs by parser + .storing("--logs", help = "Address (host:port) to send Airflow logs to") { parseAddress(this) } + .default(null) + val describeSources by parser + .storing("--describe-sources", help = "File to write each Dag's declaring class to, instead of serving") { + File(this) + }.default(null) } /** @@ -83,10 +91,13 @@ class ApiError( * The process exits when the coordinator closes the connection (normally after * one task-instance execution). */ -class Server( - private val comm: InetSocketAddress, - private val logs: InetSocketAddress, +class Server private constructor( + private val comm: InetSocketAddress?, + private val logs: InetSocketAddress?, + private val describeSources: File?, ) { + constructor(comm: InetSocketAddress, logs: InetSocketAddress) : this(comm, logs, null) + companion object { /** * Parses coordinator addresses from command-line arguments and returns a @@ -97,6 +108,10 @@ class Server( * * * `--comm host:port` address for task-execution messages. * * `--logs host:port` address for log forwarding. + * * `--describe-sources file` makes [serve] write a JSON object mapping each + * Java-declared Dag ID to the binary name of its declaring class into `file` + * and return without connecting anywhere; `--comm` and `--logs` are not + * required then. It is for the Gradle plugin, not for users. * * @param args Command-line arguments as received by `main`. * @return A configured [Server] ready to call [serve]. @@ -104,7 +119,12 @@ class Server( @JvmStatic fun create(args: Array<String>): Server { val args = ArgParser(args).parseInto(::Args) - return Server(args.comm, args.logs) + args.describeSources?.let { return Server(args.comm, args.logs, it) } + return Server( + args.comm ?: throw MissingValueException("--comm"), + args.logs ?: throw MissingValueException("--logs"), + null, + ) } } @@ -146,12 +166,16 @@ class Server( suspend fun serveAsync(bundle: Bundle) = coroutineScope { bundle.finalizeRegistration() + if (describeSources != null) { + writeSources(bundle, describeSources) + return@coroutineScope + } val deferral = CompletableDeferred<Unit>() launch { try { SelectorManager(Dispatchers.IO).use { selector -> - aSocket(selector).tcp().connect(comm).use { socket -> + aSocket(selector).tcp().connect(comm!!).use { socket -> logger.debug("Connected comm", mapOf("addr" to comm)) CoordinatorComm( socket.openReadChannel(), @@ -167,7 +191,7 @@ class Server( } launch { SelectorManager(Dispatchers.IO).use { selector -> - aSocket(selector).tcp().connect(logs).use { socket -> + aSocket(selector).tcp().connect(logs!!).use { socket -> logger.debug("Connected logs", mapOf("addr" to logs)) LogSender.configure(socket.openWriteChannel(autoFlush = true)) deferral.await() @@ -176,6 +200,16 @@ class Server( } } + private fun writeSources( + bundle: Bundle, + target: File, + ) { + val sources = linkedMapOf<String, String>() + bundle.dags.values.forEach { dag -> dag.declaringClass?.let { sources[dag.id] = it.name } } + target.absoluteFile.parentFile?.mkdirs() + ObjectMapper().writeValue(target, sources) + } + internal suspend fun dispatchTask( bundle: Bundle, coordinator: CoordinatorComm, diff --git a/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/DagSource.kt b/java-sdk/sdk/src/main/kotlin/org/apache/airflow/sdk/internal/DagSource.kt new file mode 100644 index 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()) } + } +}
