Skip to content

Commit df353a7

Browse files
jason810496jason810496
committed
Java SDK: Pack each native Dag's source file into the bundle JAR
The packDagSources task runs mainClass in describe mode, finds each Dag's source file through the class file's SourceFile attribute, and packs each unique file with a sources.json mapping under META-INF/airflow/. The JAR manifest names it with Airflow-Java-SDK-Sources.
1 parent 6dbd79a commit df353a7

10 files changed

Lines changed: 753 additions & 8 deletions

File tree

‎java-sdk/README.md‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,10 @@ Now `cd example` into the example project, and
9191
../gradlew bundle
9292
```
9393

94+
The bundle JAR also carries the source file of each Dag declared in Java
95+
(and of the entrypoint) under `META-INF/airflow/`, so Airflow can show a
96+
Dag's source.
97+
9498
* Put the [DAG with stub tasks](./example/src/resources/dags) to somewhere Airflow can find.
9599

96100
* Ensure the `java` command is available in the same environment the Airflow

‎java-sdk/plugin/build.gradle.kts‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,12 @@ plugins {
2525

2626
dependencies {
2727
implementation("com.gradleup.shadow:shadow-gradle-plugin:9.1.0") // Last supporting Java 11.
28+
29+
testImplementation(kotlin("test"))
30+
}
31+
32+
tasks.withType<Test> {
33+
useJUnitPlatform()
2834
}
2935

3036
gradlePlugin {

‎java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt‎

Lines changed: 36 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -22,11 +22,13 @@ package org.apache.airflow.sdk.plugin
2222
import com.github.jengelman.gradle.plugins.shadow.tasks.ShadowJar
2323
import org.gradle.api.Plugin
2424
import org.gradle.api.Project
25+
import org.gradle.api.plugins.JavaPluginExtension
2526
import org.gradle.api.provider.Property
2627
import org.gradle.api.tasks.Copy
2728
import org.gradle.api.tasks.Input
2829
import org.gradle.api.tasks.SourceSetContainer
2930
import org.gradle.api.tasks.bundling.Jar
31+
import org.gradle.jvm.toolchain.JavaToolchainService
3032
import java.lang.reflect.Modifier
3133
import java.net.URLClassLoader
3234
import java.util.jar.JarFile
@@ -93,6 +95,13 @@ abstract class AirflowBundleExtension {
9395
* identify which version of the Supervisor Schema it should use to communicate
9496
* with the built JAR.
9597
*
98+
* When `mainClass` is set, the bundle JAR also carries the source file of each
99+
* Dag declared in Java and of the entrypoint, so Airflow can show a Dag's source.
100+
* They are collected by the `packDagSources` task, which runs `mainClass` once in
101+
* a describe mode, and are listed in `META-INF/airflow/sources.json`, named by the
102+
* `Airflow-Java-SDK-Sources` manifest attribute. If that run fails, a warning is
103+
* logged and only the entrypoint's source is packed.
104+
*
96105
* If `fatJar` is explicitly set to `false`, the `bundle` task builds a bare JAR
97106
* containing only the Dag bundle, and collect all dependency JARs into the target
98107
* directory instead. In this mode, `Airflow-Supervisor-Schema-Version` lives in
@@ -107,21 +116,43 @@ class AirflowSdkPlugin : Plugin<Project> {
107116
ext.fatJar.convention(true)
108117

109118
project.afterEvaluate {
119+
val main = project.extensions.getByType(SourceSetContainer::class.java).getByName("main")
120+
121+
val packTask =
122+
project.tasks.register("packDagSources", PackDagSources::class.java) { task ->
123+
task.group = "build"
124+
task.description = "Collects the source file of each Dag and the entrypoint to pack into the bundle JAR."
125+
task.dependsOn(project.tasks.named("classes"))
126+
task.onlyIf { ext.mainClass.isPresent }
127+
task.mainClass.set(ext.mainClass)
128+
task.classesDirs.from(main.output.classesDirs)
129+
task.runtimeClasspath.from(main.output, main.runtimeClasspath)
130+
task.sourceDirs.from(main.allSource.srcDirs)
131+
task.launcher.convention(
132+
project.extensions
133+
.getByType(JavaToolchainService::class.java)
134+
.launcherFor(project.extensions.getByType(JavaPluginExtension::class.java).toolchain),
135+
)
136+
task.describeFile.set(project.layout.buildDirectory.file("airflow/describe-sources.json"))
137+
task.sourcesDir.set(project.layout.buildDirectory.dir("airflow/sources"))
138+
}
139+
110140
project.tasks.withType(Jar::class.java).configureEach { task ->
111141
task.doFirst {
112142
ext.mainClass.orNull?.let { className ->
113143
task.manifest.attributes(mapOf("Main-Class" to className))
114144
}
115145
}
146+
if (ext.mainClass.isPresent) {
147+
task.dependsOn(packTask)
148+
task.from(packTask.flatMap { it.sourcesDir })
149+
task.manifest.attributes(mapOf(SOURCES_MANIFEST_ATTRIBUTE to SOURCES_JSON_PATH))
150+
}
116151
}
117152

118153
val classFiles =
119154
project.objects.fileCollection().from(
120-
project.extensions
121-
.getByType(SourceSetContainer::class.java)
122-
.getByName("main")
123-
.output
124-
.classesDirs,
155+
main.output.classesDirs,
125156
project.configurations.getByName("runtimeClasspath"),
126157
)
127158

‎java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/ClassFiles.kt‎

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package org.apache.airflow.sdk.plugin
21+
22+
import java.io.ByteArrayInputStream
23+
import java.io.DataInputStream
24+
import java.io.IOException
25+
26+
private const val CLASS_MAGIC = 0xCAFEBABE.toInt()
27+
28+
/**
29+
* Reads the `SourceFile` attribute of a class file, which javac, kotlinc and scalac all
30+
* write, or returns `null` if the class has none or the bytes are not a class file.
31+
*/
32+
internal fun readSourceFile(classFile: ByteArray): String? =
33+
try {
34+
DataInputStream(ByteArrayInputStream(classFile)).use(::readSourceFile)
35+
} catch (_: IOException) {
36+
null
37+
}
38+
39+
private fun readSourceFile(input: DataInputStream): String? {
40+
if (input.readInt() != CLASS_MAGIC) return null
41+
input.skipBytes(4) // minor and major version
42+
43+
val utf8 = arrayOfNulls<String>(input.readUnsignedShort())
44+
var index = 1
45+
while (index < utf8.size) {
46+
when (val tag = input.readUnsignedByte()) {
47+
1 -> utf8[index] = input.readUTF()
48+
3, 4, 9, 10, 11, 12, 17, 18 -> input.skipBytes(4)
49+
5, 6 -> {
50+
input.skipBytes(8)
51+
index++ // long and double take two slots
52+
}
53+
7, 8, 16, 19, 20 -> input.skipBytes(2)
54+
15 -> input.skipBytes(3)
55+
else -> throw IOException("Unknown constant pool tag $tag")
56+
}
57+
index++
58+
}
59+
60+
input.skipBytes(6) // access flags, this class, super class
61+
input.skipBytes(2 * input.readUnsignedShort()) // interfaces
62+
repeat(2) { skipMembers(input) } // fields, then methods
63+
64+
repeat(input.readUnsignedShort()) {
65+
val name = utf8.getOrNull(input.readUnsignedShort())
66+
val length = input.readInt()
67+
if (name == "SourceFile" && length == 2) return utf8.getOrNull(input.readUnsignedShort())
68+
input.skipBytes(length)
69+
}
70+
return null
71+
}
72+
73+
private fun skipMembers(input: DataInputStream) {
74+
repeat(input.readUnsignedShort()) {
75+
input.skipBytes(6) // access flags, name, descriptor
76+
skipAttributes(input)
77+
}
78+
}
79+
80+
private fun skipAttributes(input: DataInputStream) {
81+
repeat(input.readUnsignedShort()) {
82+
input.skipBytes(2)
83+
input.skipBytes(input.readInt())
84+
}
85+
}

‎java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/PackDagSources.kt‎

Lines changed: 198 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,198 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one
3+
* or more contributor license agreements. See the NOTICE file
4+
* distributed with this work for additional information
5+
* regarding copyright ownership. The ASF licenses this file
6+
* to you under the Apache License, Version 2.0 (the
7+
* "License"); you may not use this file except in compliance
8+
* with the License. You may obtain a copy of the License at
9+
*
10+
* http://www.apache.org/licenses/LICENSE-2.0
11+
*
12+
* Unless required by applicable law or agreed to in writing,
13+
* software distributed under the License is distributed on an
14+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
* KIND, either express or implied. See the License for the
16+
* specific language governing permissions and limitations
17+
* under the License.
18+
*/
19+
20+
package org.apache.airflow.sdk.plugin
21+
22+
import groovy.json.JsonOutput
23+
import groovy.json.JsonSlurper
24+
import org.gradle.api.DefaultTask
25+
import org.gradle.api.file.ConfigurableFileCollection
26+
import org.gradle.api.file.DirectoryProperty
27+
import org.gradle.api.file.RegularFileProperty
28+
import org.gradle.api.provider.Property
29+
import org.gradle.api.tasks.CacheableTask
30+
import org.gradle.api.tasks.Classpath
31+
import org.gradle.api.tasks.Input
32+
import org.gradle.api.tasks.InputFiles
33+
import org.gradle.api.tasks.Nested
34+
import org.gradle.api.tasks.Optional
35+
import org.gradle.api.tasks.OutputDirectory
36+
import org.gradle.api.tasks.OutputFile
37+
import org.gradle.api.tasks.PathSensitive
38+
import org.gradle.api.tasks.PathSensitivity
39+
import org.gradle.api.tasks.TaskAction
40+
import org.gradle.jvm.toolchain.JavaLauncher
41+
import org.gradle.process.ExecOperations
42+
import java.io.ByteArrayOutputStream
43+
import java.io.File
44+
import javax.inject.Inject
45+
46+
internal const val SOURCES_MANIFEST_ATTRIBUTE = "Airflow-Java-SDK-Sources"
47+
internal const val SOURCES_JSON_PATH = "META-INF/airflow/sources.json"
48+
internal const val SOURCES_DIR_PATH = "META-INF/airflow/sources"
49+
50+
/**
51+
* Finds the source file of a class by its `SourceFile` attribute: the class
52+
* file's package path plus that name, looked up in each source directory.
53+
*/
54+
internal class SourceLocator(
55+
private val classesDirs: Iterable<File>,
56+
private val sourceDirs: Iterable<File>,
57+
) {
58+
/** Path of the source file relative to its source directory, or `null` if it cannot be found. */
59+
fun locate(className: String): String? {
60+
val packagePath = className.substringBeforeLast('.', "").replace('.', '/')
61+
val classFile =
62+
classesDirs
63+
.map { File(it, className.replace('.', '/') + ".class") }
64+
.firstOrNull { it.isFile } ?: return null
65+
val sourceFile = readSourceFile(classFile.readBytes())?.takeIf { it.isNotEmpty() && '/' !in it && '\\' !in it }
66+
val relative = listOfNotNull(packagePath.ifEmpty { null }, sourceFile ?: return null).joinToString("/")
67+
return relative.takeIf { path -> sourceDirs.any { File(it, path).isFile } }
68+
}
69+
70+
fun file(relativePath: String): File = sourceDirs.map { File(it, relativePath) }.first { it.isFile }
71+
}
72+
73+
internal fun sourcesJson(
74+
entrypointPath: String?,
75+
dagSourcePaths: Map<String, String>,
76+
): String =
77+
JsonOutput.prettyPrint(
78+
JsonOutput.toJson(
79+
linkedMapOf<String, Any>().apply {
80+
entrypointPath?.let { put("entrypoint_path", it) }
81+
put("dag_source_paths", dagSourcePaths)
82+
},
83+
),
84+
)
85+
86+
/**
87+
* Runs the bundle's `mainClass` in describe mode to learn which class declared
88+
* each Dag, then collects those classes' source files, and the entry point's,
89+
* for the bundle JAR to carry.
90+
*/
91+
@CacheableTask
92+
abstract class PackDagSources : DefaultTask() {
93+
@get:Inject
94+
abstract val execOperations: ExecOperations
95+
96+
@get:Input
97+
@get:Optional
98+
abstract val mainClass: Property<String>
99+
100+
@get:Classpath
101+
abstract val classesDirs: ConfigurableFileCollection
102+
103+
@get:Classpath
104+
abstract val runtimeClasspath: ConfigurableFileCollection
105+
106+
@get:InputFiles
107+
@get:PathSensitive(PathSensitivity.RELATIVE)
108+
abstract val sourceDirs: ConfigurableFileCollection
109+
110+
@get:Nested
111+
@get:Optional
112+
abstract val launcher: Property<JavaLauncher>
113+
114+
@get:OutputFile
115+
abstract val describeFile: RegularFileProperty
116+
117+
@get:OutputDirectory
118+
abstract val sourcesDir: DirectoryProperty
119+
120+
private var describeFailed = false
121+
122+
init {
123+
outputs.doNotCacheIf("the Dag describe run failed") { describeFailed }
124+
}
125+
126+
@TaskAction
127+
fun pack() {
128+
val describe = describeFile.get().asFile
129+
val root = sourcesDir.get().asFile
130+
describe.delete()
131+
root.deleteRecursively()
132+
describe.parentFile.mkdirs()
133+
134+
val main = mainClass.get()
135+
val declaringClasses = describeDags(main, describe)
136+
val locator = SourceLocator(classesDirs.files, sourceDirs.files)
137+
138+
val entrypoint = locator.locate(main)
139+
val dagPaths = linkedMapOf<String, String>()
140+
declaringClasses.forEach { (dagId, className) ->
141+
val path = locator.locate(className)
142+
if (path == null) {
143+
logger.info("No source file found for class {} of Dag '{}'; it falls back to the entrypoint", className, dagId)
144+
} else {
145+
dagPaths[dagId] = path
146+
}
147+
}
148+
149+
(listOfNotNull(entrypoint) + dagPaths.values).distinct().forEach { path ->
150+
val target = File(root, "$SOURCES_DIR_PATH/$path")
151+
target.parentFile.mkdirs()
152+
locator.file(path).copyTo(target)
153+
}
154+
File(root, SOURCES_JSON_PATH).apply {
155+
parentFile.mkdirs()
156+
writeText(sourcesJson(entrypoint, dagPaths) + "\n")
157+
}
158+
}
159+
160+
private fun describeDags(
161+
main: String,
162+
describe: File,
163+
): Map<String, String> {
164+
val output = ByteArrayOutputStream()
165+
val failure =
166+
try {
167+
val result =
168+
execOperations.javaexec { spec ->
169+
launcher.orNull?.let { spec.executable(it.executablePath.asFile.absolutePath) }
170+
spec.classpath(classesDirs, runtimeClasspath)
171+
spec.mainClass.set(main)
172+
spec.args("--describe-sources", describe.absolutePath)
173+
spec.isIgnoreExitValue = true
174+
spec.standardOutput = output
175+
spec.errorOutput = output
176+
}
177+
if (result.exitValue != 0) "exit code ${result.exitValue}" else null
178+
} catch (e: Exception) {
179+
e.message ?: e.javaClass.simpleName
180+
}
181+
if (failure == null && describe.isFile) {
182+
@Suppress("UNCHECKED_CAST")
183+
(runCatching { JsonSlurper().parse(describe) }.getOrNull() as? Map<String, Any?>)?.let { parsed ->
184+
return parsed.mapNotNull { (id, cls) -> (cls as? String)?.let { id to it } }.toMap(linkedMapOf())
185+
}
186+
}
187+
describeFailed = true
188+
val why = failure ?: "it wrote no valid --describe-sources file; does main call Server.serve?"
189+
val log = output.toString().trim()
190+
logger.warn(
191+
"Could not read each Dag's source from '{}' ({}); only its entrypoint source is packed.{}",
192+
main,
193+
why,
194+
if (log.isEmpty()) "" else "\n$log",
195+
)
196+
return emptyMap()
197+
}
198+
}

0 commit comments

Comments
 (0)