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 3f69772e4d92d..c5c61a67f0950 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst @@ -589,6 +589,46 @@ Durations and date-times are ISO-8601 strings in annotations (``retryDelay = "PT ``java.time.OffsetDateTime`` values in ``config`` calls. An unknown key or a mismatched value type fails the build for an annotation, and the ``config`` call itself for an object. +.. _java-sdk/native-dag-parsing: + +Parsing native Java Dags +~~~~~~~~~~~~~~~~~~~~~~~~ + +To have Airflow parse the Dags a bundle JAR declares, put the JAR in a Dag bundle and configure a +:class:`~airflow.sdk.coordinators.java.JavaCoordinator` that reads that bundle: leave ``jars_root`` +unset to read each Dag's own bundle, or set ``dag_bundle_name``. The Dag processor runs the JAR's main +class to list its Dags, so it needs a Java executable, as the workers do: + +.. code-block:: ini + + [sdk] + coordinators = { + "java-native": { + "classpath": "airflow.sdk.coordinators.java.JavaCoordinator", + "kwargs": {"java_executable": "/usr/lib/jvm/java-17-openjdk/bin/java"} + } + } + queue_to_coordinator = {"java-native": "java-native"} + +A coordinator with ``jars_root`` set only runs tasks; it parses no Dags. + +* Every JAR in the bundle whose manifest sets ``Main-Class`` is parsed. Each Dag its main class + declares, through ``Bundle.register`` of a ``DagDef`` or an ``@Builder.Dag`` class, is stored with + that JAR as its file. Task handlers for a Python Dag are not Dags, and a JAR without ``Main-Class``, + such as a dependency of a thin bundle, is skipped. +* Do not declare a Dag in Java that a Python file in the same bundle also defines. +* Keep one executable JAR per bundle, or set ``main_class``, so that only JARs with that + ``Main-Class`` are parsed. List JARs that should not be parsed in ``.airflowignore``. +* Set ``queue`` on every task, with ``@Builder.Task(queue = "java-native")`` or + ``TaskDef.config("queue", "java-native")``, so it runs on the coordinator's queue. There is no + Dag-level queue yet. +* The Code view shows the source file the Gradle plugin packs into the JAR: the main class by + default, or the file set with ``airflowBundle { dagSource = file("...") }``. +* Cluster policies (``dag_policy``, ``task_policy``) are not applied to a native Java Dag. +* ``airflow dags reserialize`` does not store the Dags of a JAR, which only the Dag processor stores. + ``airflow dags test``, ``tasks test``, ``tasks render`` and ``tasks list`` refuse a native Java + Dag. + .. _java-sdk/logging: Logging @@ -855,6 +895,17 @@ The ``build/bundle/`` directory contains all required JAR(s). Copy or mount it i by ``jars_root`` in the coordinator configuration. :class:`~airflow.sdk.coordinators.java.JavaCoordinator` scans ``jars_root`` recursively and builds the classpath automatically. +The plugin also packs the source file of ``mainClass`` into the bundle JAR, so the Airflow UI can show +the source of a native Java Dag (see :ref:`java-sdk/native-dag-parsing`). Set ``dagSource`` in +``airflowBundle`` to pack another file instead: + +.. code-block:: groovy + + airflowBundle { + mainClass = "com.example.Main" + dagSource = file("src/main/java/com/example/MyDag.java") + } + .. note:: You only need the ``annotationProcessor`` entry if you use the annotation-based API. It is not needed for @@ -1029,6 +1080,28 @@ this directory. Unlike the Gradle plugin, Maven has no equivalent of the ``verifyBundleMainClass`` validation step. A wrong ```` value will not be caught until runtime. +To show the source of a native Java Dag in the Airflow UI, pack the source file into the bundle JAR and +name its entry with the ``Airflow-Java-SDK-Dag-Code`` manifest attribute: + +.. code-block:: xml + + + + + + src/main/resources + + + src/main/java/com/example + Main.java + META-INF/airflow/dag-code/com/example + + + + +Then add ``META-INF/airflow/dag-code/com/example/Main.java`` +to the ``manifestEntries`` of ``maven-shade-plugin`` or ``maven-jar-plugin`` shown above. + .. _java-sdk/coordinator-config: :class:`~airflow.sdk.coordinators.java.JavaCoordinator` configuration @@ -1065,7 +1138,9 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the - Explicit entry-point class. If omitted, the coordinator scans for a JAR whose manifest sets ``Main-Class`` — in ``jars_root`` when set, otherwise across the resolved Dag bundle. If multiple executable JARs match the result is - non-deterministic; set ``main_class`` explicitly in that case. + non-deterministic; set ``main_class`` explicitly in that case. When the + coordinator parses native Java Dags, only JARs with this ``Main-Class`` are + parsed. * - ``task_startup_timeout`` - ``10.0`` - Seconds to wait for the JVM subprocess to connect after launch. Increase this if your 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 95836c12fc5fe..ce8d11185fd47 100644 --- a/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt +++ b/java-sdk/plugin/src/main/kotlin/org/apache/airflow/sdk/plugin/AirflowSdkPlugin.kt @@ -22,11 +22,15 @@ 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.file.RegularFileProperty import org.gradle.api.provider.Property import org.gradle.api.tasks.Copy import org.gradle.api.tasks.Input +import org.gradle.api.tasks.InputFile +import org.gradle.api.tasks.Optional import org.gradle.api.tasks.SourceSetContainer import org.gradle.api.tasks.bundling.Jar +import java.io.File import java.lang.reflect.Modifier import java.net.URLClassLoader import java.util.jar.JarFile @@ -61,8 +65,22 @@ abstract class AirflowBundleExtension { */ @get:Input abstract val fatJar: Property + + /** + * The Dag source file the Airflow UI shows for this bundle. + * + * It defaults to the `.java` file of [mainClass] in the `main` source set. The file is packed + * into the bundle JAR under `META-INF/airflow/dag-code/`, and the `Airflow-Java-SDK-Dag-Code` + * manifest attribute names it. + */ + @get:InputFile + @get:Optional + abstract val dagSource: RegularFileProperty } +private const val DAG_CODE_ATTRIBUTE = "Airflow-Java-SDK-Dag-Code" +private const val DAG_CODE_DIR = "META-INF/airflow/dag-code" + /** * Gradle plugin for building Apache Airflow Java SDK bundles. * @@ -98,6 +116,9 @@ abstract class AirflowBundleExtension { * directory instead. In this mode, `Airflow-Supervisor-Schema-Version` lives in * the `airflow-sdk` JAR instead. The bundle JAR still contains `Main-Class`. * + * In both modes the bundle JAR also carries the Dag source (see + * [AirflowBundleExtension.dagSource]), so the Airflow UI can show the source of a + * native Java Dag. */ class AirflowSdkPlugin : Plugin { override fun apply(project: Project) { @@ -107,12 +128,17 @@ class AirflowSdkPlugin : Plugin { ext.fatJar.convention(true) project.afterEvaluate { + val dagCode = resolveDagCode(project, ext) project.tasks.withType(Jar::class.java).configureEach { task -> task.doFirst { ext.mainClass.orNull?.let { className -> task.manifest.attributes(mapOf("Main-Class" to className)) } } + dagCode?.let { (file, entryDir) -> + task.from(file) { spec -> spec.into(entryDir) } + task.manifest.attributes(mapOf(DAG_CODE_ATTRIBUTE to "$entryDir/${file.name}")) + } } val classFiles = @@ -219,3 +245,35 @@ class AirflowSdkPlugin : Plugin { } } } + +/** + * Returns the Dag source file to pack and its directory in the JAR, or `null` when there is none. + * + * The directory mirrors the package of `mainClass`, so the entry reads like a source path. + */ +private fun resolveDagCode( + project: Project, + ext: AirflowBundleExtension, +): Pair? { + val mainClass = ext.mainClass.orNull ?: return null + val className = mainClass.substringBefore('$') + val packagePath = className.substringBeforeLast('.', "").replace('.', '/') + val entryDir = if (packagePath.isEmpty()) DAG_CODE_DIR else "$DAG_CODE_DIR/$packagePath" + ext.dagSource.orNull?.let { return it.asFile to entryDir } + + val relativePath = className.replace('.', '/') + ".java" + val source = + project.extensions + .getByType(SourceSetContainer::class.java) + .getByName("main") + .java.srcDirs + .map { it.resolve(relativePath) } + .firstOrNull { it.isFile } + if (source == null) { + project.logger.info( + "airflowBundle: no $relativePath in the main source set, so the bundle JAR carries no Dag " + + "source for the Airflow UI. Set airflowBundle.dagSource to choose the file.", + ) + } + return source?.let { it to entryDir } +} diff --git a/task-sdk/src/airflow/sdk/coordinators/java/_dag_importer.py b/task-sdk/src/airflow/sdk/coordinators/java/_dag_importer.py new file mode 100644 index 0000000000000..a173f3a1183a0 --- /dev/null +++ b/task-sdk/src/airflow/sdk/coordinators/java/_dag_importer.py @@ -0,0 +1,86 @@ +# +# 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. +"""The Dag importer of :class:`~airflow.sdk.coordinators.java.JavaCoordinator`.""" + +from __future__ import annotations + +import zipfile +from typing import TYPE_CHECKING, ClassVar, Final + +from airflow.sdk.coordinators._dag_importer import CoordinatorDagImporter +from airflow.sdk.coordinators.java._jar_manifest import DAG_CODE, MAIN_CLASS, read_main_attributes +from airflow.sdk.importers.base import DagSourceCode + +if TYPE_CHECKING: + from airflow.sdk.coordinators.java.coordinator import JavaCoordinator + from airflow.sdk.importers.base import DagDefinition + +_MAX_SOURCE_BYTES: Final = 1024 * 1024 +_NO_SOURCE: Final = ( + "// This JAR embeds no Dag source. Build it with the Airflow Java SDK Gradle plugin, or set\n" + "// airflowBundle.dagSource, to show the source here.\n" +) +_SOURCE_TOO_LARGE: Final = "// The Dag source this JAR embeds is over 1 MiB, so it is not shown.\n" + + +class JavaDagImporter(CoordinatorDagImporter): + """ + Parse the native Dags of Java bundle JARs with the coordinator's JVM. + + Only a JAR whose manifest sets ``Main-Class`` is parsed, so the dependency JARs of a thin bundle + are skipped. + """ + + artifact_suffix: ClassVar[str] = ".jar" + supported_extensions = [".jar"] + coordinator: JavaCoordinator + + def might_contain_dag(self, definition: DagDefinition, safe_mode: bool) -> bool: + """ + Return whether the JAR sets ``Main-Class``, matching the coordinator's ``main_class`` if that is set. + + ``safe_mode`` does not apply, because a JAR without ``Main-Class`` cannot run at all. A JAR that + cannot be read is kept, so that parsing it reports the error. + """ + try: + with definition.as_file() as path, zipfile.ZipFile(path) as zf: + attributes = read_main_attributes(zf) or {} + except (OSError, zipfile.BadZipFile): + return True + if not (main_class := attributes.get(MAIN_CLASS)): + return False + return not self.coordinator.main_class or main_class == self.coordinator.main_class + + def get_source_code(self, definition: DagDefinition) -> DagSourceCode: + """Return the Dag source the JAR embeds, or a placeholder when it embeds none.""" + with definition.as_file() as path, zipfile.ZipFile(path) as zf: + if (info := _find_source_entry(zf)) is None: + return DagSourceCode(source_code=_NO_SOURCE, language="java") + if info.file_size > _MAX_SOURCE_BYTES: + return DagSourceCode(source_code=_SOURCE_TOO_LARGE, language="java") + source = zf.read(info).decode("utf-8", errors="replace") + return DagSourceCode(source_code=source or _NO_SOURCE, language="java") + + +def _find_source_entry(zf: zipfile.ZipFile) -> zipfile.ZipInfo | None: + if not (entry := (read_main_attributes(zf) or {}).get(DAG_CODE)): + return None + try: + return zf.getinfo(entry) + except KeyError: + return None diff --git a/task-sdk/src/airflow/sdk/coordinators/java/_jar_manifest.py b/task-sdk/src/airflow/sdk/coordinators/java/_jar_manifest.py new file mode 100644 index 0000000000000..d2f84395a3e55 --- /dev/null +++ b/task-sdk/src/airflow/sdk/coordinators/java/_jar_manifest.py @@ -0,0 +1,71 @@ +# +# 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. +"""Read JAR manifests as the JAR File Specification defines them.""" + +from __future__ import annotations + +import re +from typing import TYPE_CHECKING, Final + +if TYPE_CHECKING: + import zipfile + +MANIFEST_NAME: Final = "META-INF/MANIFEST.MF" + +# Attribute keys as parse_main_attributes returns them, lower-cased. +MAIN_CLASS: Final = "main-class" +SUPERVISOR_SCHEMA_VERSION: Final = "airflow-supervisor-schema-version" +DAG_CODE: Final = "airflow-java-sdk-dag-code" + +_LINE_END = re.compile(rb"\r\n|\r|\n") + + +def parse_main_attributes(data: bytes) -> dict[str, str]: + """ + Return the main-section attributes of a JAR manifest, keyed by lower-cased name. + + Attribute names are case-insensitive. Lines end with CRLF, LF or CR, and a line that starts with + one space continues the previous value. The main section ends at the first blank line. + """ + headers: list[bytearray] = [] + for line in _LINE_END.split(data): + if not line: + break + if line.startswith(b" "): + # Lines are folded by bytes, which can split a multi-byte character, so unfold + # before decoding. + if headers: + headers[-1] += line[1:] + continue + headers.append(bytearray(line)) + + attributes: dict[str, str] = {} + for header in headers: + name, sep, value = header.decode("utf-8", errors="replace").partition(":") + if sep: + attributes[name.strip().lower()] = value.removeprefix(" ") + return attributes + + +def read_main_attributes(zf: zipfile.ZipFile) -> dict[str, str] | None: + """Return the main attributes of the JAR open in *zf*, or ``None`` when it has no manifest.""" + try: + info = zf.getinfo(MANIFEST_NAME) + except KeyError: + return None + return parse_main_attributes(zf.read(info)) diff --git a/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py b/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py index f57dacaec558c..97a99ad7b2efa 100644 --- a/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py +++ b/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py @@ -19,7 +19,6 @@ from __future__ import annotations -import email import os import pathlib import stat @@ -35,6 +34,12 @@ validate_schema_version, ) from airflow.sdk.coordinators._subprocess import SubprocessCoordinator +from airflow.sdk.coordinators.java._dag_importer import JavaDagImporter +from airflow.sdk.coordinators.java._jar_manifest import ( + MAIN_CLASS, + SUPERVISOR_SCHEMA_VERSION, + read_main_attributes, +) if TYPE_CHECKING: from collections.abc import Iterable, Iterator, Sequence @@ -93,24 +98,21 @@ def _calculate_classpath(jars_root: Sequence[pathlib.Path]) -> str: @attrs.define class _JarMetadata: - main_class: str - schema_version: str + main_class: str | None + schema_version: str | None @classmethod def from_jar(cls, path: pathlib.Path) -> Self | None: try: with zipfile.ZipFile(path) as zf: - try: - manifest_info = zf.getinfo("META-INF/MANIFEST.MF") - except KeyError: - log.debug("JAR does not contain META-INF/MANIFEST.MF; ignored", path=path) - return None - with zf.open(manifest_info) as f: - manifest = email.message_from_binary_file(f) - return cls(manifest["Main-Class"], manifest["Airflow-Supervisor-Schema-Version"]) - except zipfile.BadZipFile: + attributes = read_main_attributes(zf) + except (OSError, zipfile.BadZipFile): log.exception("Cannot read JAR; ignored", path=path) return None + if attributes is None: + log.debug("JAR does not contain META-INF/MANIFEST.MF; ignored", path=path) + return None + return cls(attributes.get(MAIN_CLASS), attributes.get(SUPERVISOR_SCHEMA_VERSION)) @attrs.define @@ -193,6 +195,11 @@ class JavaCoordinator(SubprocessCoordinator): dependency JARs are deployed as-is. If you repackage the dependencies, however, you must also reproduce the metadata entry in one of the JARs. + Without *jars_root*, the coordinator also parses native Java Dags: every JAR + in the Dag bundle it reads whose manifest sets Main-Class (matching + *main_class* when that is set) is run to list the Dags its main class + declares. + The default *task_startup_timeout* should plenty long enough since a task- containing JAR is not supposed to consume significant time to perform setup (it should happen in individual tasks instead). However, if the launch time @@ -211,16 +218,32 @@ class JavaCoordinator(SubprocessCoordinator): _explicit_root_kwarg = "jars_root" + @classmethod + def get_dag_importer_class(cls) -> type[JavaDagImporter]: + return JavaDagImporter + + def _build_command(self, roots: Sequence[pathlib.Path], main_class: str) -> list[str]: + return [self.java_executable, "-classpath", _calculate_classpath(roots), *self.jvm_args, main_class] + def _build_execute_task_command(self, *, what: TaskInstance) -> tuple[list[str], str | None]: # Without main_class, the first executable JAR in walk order wins; tracked at # https://github.com/apache/airflow/issues/71134 roots = self._get_scan_roots() jar = _JarInfo.find(roots, self.main_class) - command = [ - self.java_executable, - "-classpath", - _calculate_classpath(roots), - *self.jvm_args, - jar.main_class, - ] - return command, jar.schema_version + return self._build_command(roots, jar.main_class), jar.schema_version + + def _build_parse_dag_command(self, *, path: pathlib.Path) -> tuple[list[str], str | None]: + # Same command shape as execution. With one executable JAR per bundle, or main_class set, + # a parse runs the class a task runs. + meta = _JarMetadata.from_jar(path) + if meta is None: + raise ValueError(f"Cannot read the manifest of {path}") + if not meta.main_class: + raise ValueError(f"{path} is not an executable JAR: its manifest sets no Main-Class") + if self.main_class and meta.main_class != self.main_class: + raise ValueError( + f"{path} runs {meta.main_class!r}, but this coordinator's main_class is {self.main_class!r}" + ) + roots = self._get_scan_roots() + jar = _JarInfo.find(roots, meta.main_class) + return self._build_command(roots, jar.main_class), jar.schema_version diff --git a/task-sdk/tests/task_sdk/coordinators/java/_jar_test_utils.py b/task-sdk/tests/task_sdk/coordinators/java/_jar_test_utils.py new file mode 100644 index 0000000000000..0964a4ddf7428 --- /dev/null +++ b/task-sdk/tests/task_sdk/coordinators/java/_jar_test_utils.py @@ -0,0 +1,53 @@ +# 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. + +from __future__ import annotations + +import pathlib +import zipfile + +SCHEMA_VERSION = "2026-06-16" + + +def write_manifest(attributes: dict[str, str], *, newline: bytes = b"\r\n") -> bytes: + """Write *attributes* as a manifest main section, folded at 72 bytes as ``java.util.jar.Manifest`` does.""" + out = bytearray() + for name, value in attributes.items(): + line = f"{name}: {value}".encode() + out += line[:72] + for start in range(72, len(line), 71): + out += newline + b" " + line[start : start + 71] + out += newline + return bytes(out + newline) + + +def make_jar( + path: pathlib.Path, + *, + attributes: dict[str, str] | None = None, + entries: dict[str, bytes | str] | None = None, + manifest: bytes | None = None, +) -> pathlib.Path: + """Write a JAR whose manifest holds *attributes* (or the raw *manifest*), plus *entries*.""" + if manifest is None and attributes is not None: + manifest = write_manifest(attributes) + with zipfile.ZipFile(path, "w") as zf: + if manifest is not None: + zf.writestr("META-INF/MANIFEST.MF", manifest) + for name, data in (entries or {}).items(): + zf.writestr(name, data) + return path diff --git a/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py b/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py index c55a2cc891253..b3843cd3d8ffd 100644 --- a/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py @@ -23,10 +23,10 @@ import re import socket import subprocess -import zipfile from unittest.mock import MagicMock, patch import pytest +from task_sdk.coordinators.java._jar_test_utils import make_jar from uuid6 import uuid7 from airflow.sdk.api.datamodels._generated import TaskInstance @@ -34,6 +34,7 @@ JavaCoordinator, _calculate_classpath, _JarInfo, + _JarMetadata, _walk_jars, ) from airflow.sdk.execution_time.coordinator import BaseCoordinator @@ -65,15 +66,39 @@ def _make_jar( schema_version: str | None = None, ) -> pathlib.Path: """Write a minimal JAR with (optionally) a Main-Class manifest entry.""" - lines = ["Manifest-Version: 1.0"] + attributes = {"Manifest-Version": "1.0"} if main_class: - lines.append(f"Main-Class: {main_class}") + attributes["Main-Class"] = main_class if schema_version: - lines.append(f"Airflow-Supervisor-Schema-Version: {schema_version}") - manifest = "\n".join(lines) + "\n\n" - with zipfile.ZipFile(path, "w") as zf: - zf.writestr("META-INF/MANIFEST.MF", manifest) - return path + attributes["Airflow-Supervisor-Schema-Version"] = schema_version + return make_jar(path, attributes=attributes) + + +class TestJarMetadata: + def test_reads_a_folded_main_class(self, tmp_path): + main_class = "org.apache.airflow.example.nativedag.generated.VeryLongNativeDagBundleBuilder" + jar = make_jar(tmp_path / "app.jar", attributes={"Main-Class": main_class}) + + assert _JarMetadata.from_jar(jar) == _JarMetadata(main_class, None) + + def test_jar_without_manifest_gives_none(self, tmp_path): + jar = make_jar(tmp_path / "lib.jar", entries={"com/example/Lib.class": b"\xca\xfe"}) + + assert _JarMetadata.from_jar(jar) is None + + def test_non_zip_gives_none(self, tmp_path): + jar = tmp_path / "broken.jar" + jar.write_bytes(b"not a zip") + + assert _JarMetadata.from_jar(jar) is None + + @pytest.mark.parametrize("kind", ["missing", "directory"]) + def test_unreadable_jar_gives_none(self, tmp_path, kind): + jar = tmp_path / "gone.jar" + if kind == "directory": + jar.mkdir() + + assert _JarMetadata.from_jar(jar) is None class TestCalculateClasspath: @@ -425,3 +450,90 @@ def test_returns_execution_result(self, jars_root, mock_client): assert isinstance(result, BaseCoordinator.ExecutionResult) assert result.exit_code == 0 + + +class TestBuildParseDagCommand: + def _build(self, coordinator: JavaCoordinator, root: pathlib.Path, jar: pathlib.Path): + with coordinator._set_scan_roots([root]): + return coordinator._build_parse_dag_command(path=jar) + + def test_fat_jar(self, tmp_path): + jar = _make_jar(tmp_path / "app.jar", main_class="com.example.Dags", schema_version="2026-06-16") + coordinator = JavaCoordinator(java_executable="/opt/java/bin/java", jvm_args=["-Xmx256m"]) + + command, schema_version = self._build(coordinator, tmp_path, jar) + + assert command == ["/opt/java/bin/java", "-classpath", jar.as_posix(), "-Xmx256m", "com.example.Dags"] + assert schema_version == "2026-06-16" + + def test_thin_jar_takes_the_schema_version_from_a_sibling(self, tmp_path): + jar = _make_jar(tmp_path / "app.jar", main_class="com.example.Dags") + (tmp_path / "libs").mkdir() + sdk = _make_jar(tmp_path / "libs" / "airflow-sdk.jar", main_class=None, schema_version="2026-06-16") + + command, schema_version = self._build(JavaCoordinator(), tmp_path, jar) + + assert command[2].split(os.pathsep) == [jar.as_posix(), sdk.as_posix()] + assert command[-1] == "com.example.Dags" + assert schema_version == "2026-06-16" + + def test_matches_the_execute_command(self, tmp_path): + jar = _make_jar(tmp_path / "app.jar", main_class="com.example.Dags", schema_version="2026-06-16") + _make_jar(tmp_path / "dep.jar", main_class=None) + coordinator = JavaCoordinator(jvm_args=["-Xmx256m"]) + + parse = self._build(coordinator, tmp_path, jar) + with coordinator._set_scan_roots([tmp_path]): + execute = coordinator._build_execute_task_command(what=_make_ti()) + + assert parse == execute + + @pytest.mark.parametrize( + ("attributes", "match"), + [ + pytest.param(None, "Cannot read the manifest", id="no-manifest"), + pytest.param({"Manifest-Version": "1.0"}, "sets no Main-Class", id="no-main-class"), + pytest.param({"Main-Class": "com.example.Other"}, "main_class is 'com.example.Dags'", id="pin"), + ], + ) + def test_rejects_a_jar_it_cannot_run(self, tmp_path, attributes, match): + jar = make_jar(tmp_path / "app.jar", attributes=attributes, entries={"a.class": b""}) + + with pytest.raises(ValueError, match=match): + self._build(JavaCoordinator(main_class="com.example.Dags"), tmp_path, jar) + + def test_rejects_a_non_zip(self, tmp_path): + jar = tmp_path / "broken.jar" + jar.write_bytes(b"not a zip") + + with pytest.raises(ValueError, match="Cannot read the manifest"): + self._build(JavaCoordinator(), tmp_path, jar) + + def test_rejects_a_jar_deleted_after_discovery(self, tmp_path): + with pytest.raises(ValueError, match="Cannot read the manifest"): + self._build(JavaCoordinator(), tmp_path, tmp_path / "gone.jar") + + @patch("airflow.sdk.coordinators._subprocess.os.execvpe", autospec=True, side_effect=OSError("exec")) + def test_parse_dag_execs_the_jvm(self, mock_execvpe, tmp_path): + jar = _make_jar(tmp_path / "app.jar", main_class="com.example.Dags", schema_version="2026-06-16") + reported: list[str | None] = [] + + with pytest.raises(OSError, match="exec"): + JavaCoordinator().parse_dag( + path=jar, + bundle_path=tmp_path, + comm_address=("127.0.0.1", 1001), + logs_address=("127.0.0.1", 1002), + report_schema_version=reported.append, + ) + + assert reported == ["2026-06-16"] + argv = mock_execvpe.call_args.args[1] + assert argv == [ + "java", + "-classpath", + jar.as_posix(), + "com.example.Dags", + "--comm=127.0.0.1:1001", + "--logs=127.0.0.1:1002", + ] 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 new file mode 100644 index 0000000000000..a2e990f6e268e --- /dev/null +++ b/task-sdk/tests/task_sdk/coordinators/java/test_dag_importer.py @@ -0,0 +1,201 @@ +# 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. + +from __future__ import annotations + +import json +import zipfile +from unittest.mock import patch + +import pytest +from task_sdk.coordinators.java._jar_test_utils import SCHEMA_VERSION, make_jar + +from airflow.sdk.coordinators.java import JavaCoordinator +from airflow.sdk.coordinators.java._dag_importer import JavaDagImporter +from airflow.sdk.importers import FilesystemDagDefinition, get_importer_registry, reset_importer_registry + +from tests_common.test_utils.config import conf_vars +from tests_common.test_utils.version_compat import AIRFLOW_V_3_3_PLUS + +if not AIRFLOW_V_3_3_PLUS: + pytest.skip("Coordinator is only compatible with Airflow >= 3.3.0", allow_module_level=True) + +MAIN_CLASS = "com.example.Dags" +SOURCE_ENTRY = "META-INF/airflow/dag-code/com/example/Dags.java" +LONG_MAIN_CLASS = "org.apache.airflow.example.nativedag.generated.VeryLongNativeDagBundleBuilder" + + +def _importer(main_class: str = "") -> JavaDagImporter: + return JavaDagImporter(coordinator=JavaCoordinator(main_class=main_class)) + + +def _definition(path) -> FilesystemDagDefinition: + return FilesystemDagDefinition(path=path) + + +class _Bundle: + def __init__(self, path): + self.path = path + self.name = "dags-folder" + + +class TestJavaDagImporter: + def test_claims_jar_files(self): + importer = _importer() + + assert importer.supported_extensions == [".jar"] + assert importer.can_handle("dags/app.jar") + assert not importer.can_handle("dags/app.zip") + + def test_coordinator_hands_out_an_importer_bound_to_it(self): + coordinator = JavaCoordinator() + + importer = coordinator.get_dag_importer() + + assert JavaCoordinator.get_dag_importer_class() is JavaDagImporter + assert isinstance(importer, JavaDagImporter) + assert importer.coordinator is coordinator + + @pytest.mark.parametrize( + ("attributes", "pin", "expected"), + [ + pytest.param( + {"Main-Class": MAIN_CLASS, "Airflow-Supervisor-Schema-Version": SCHEMA_VERSION}, + "", + True, + id="fat", + ), + pytest.param({"Main-Class": MAIN_CLASS}, "", True, id="thin-app"), + pytest.param({"Airflow-Supervisor-Schema-Version": SCHEMA_VERSION}, "", False, id="dependency"), + pytest.param(None, "", False, id="no-manifest"), + pytest.param({"Main-Class": LONG_MAIN_CLASS}, LONG_MAIN_CLASS, True, id="folded-main-class"), + pytest.param({"Main-Class": MAIN_CLASS}, MAIN_CLASS, True, id="pin-match"), + pytest.param({"Main-Class": MAIN_CLASS}, "com.example.Other", False, id="pin-mismatch"), + ], + ) + def test_might_contain_dag(self, tmp_path, attributes, pin, expected): + jar = make_jar(tmp_path / "app.jar", attributes=attributes, entries={"a.class": b""}) + + assert _importer(pin).might_contain_dag(_definition(jar), safe_mode=True) is expected + + def test_keeps_a_jar_it_cannot_read(self, tmp_path): + jar = tmp_path / "partial.jar" + jar.write_bytes(b"PK\x03\x04 truncated") + + assert _importer().might_contain_dag(_definition(jar), safe_mode=True) is True + + @pytest.mark.parametrize("safe_mode", [True, False]) + def test_lists_only_executable_jars(self, tmp_path, safe_mode): + app = make_jar(tmp_path / "app.jar", attributes={"Main-Class": MAIN_CLASS}) + (tmp_path / "libs").mkdir() + make_jar(tmp_path / "libs" / "dep.jar", entries={"dep/Dep.class": b""}) + make_jar( + tmp_path / "libs" / "airflow-sdk.jar", + attributes={"Airflow-Supervisor-Schema-Version": SCHEMA_VERSION}, + ) + + listed = list(_importer().list_dag_definitions(_Bundle(tmp_path), safe_mode=safe_mode)) + + assert [d.path for d in listed] == [app] + + +class TestGetSourceCode: + def _jar(self, tmp_path, *, attributes=None, entries=None): + attributes = {"Main-Class": MAIN_CLASS, **(attributes or {})} + return make_jar(tmp_path / "app.jar", attributes=attributes, entries=entries) + + def test_returns_the_embedded_source(self, tmp_path): + source = "public class Dags {}\n" + jar = self._jar( + tmp_path, attributes={"Airflow-Java-SDK-Dag-Code": SOURCE_ENTRY}, entries={SOURCE_ENTRY: source} + ) + assert b"\r\n " in zipfile.ZipFile(jar).read("META-INF/MANIFEST.MF") + + result = _importer().get_source_code(_definition(jar)) + + assert (result.source_code, result.language) == (source, "java") + + @pytest.mark.parametrize( + ("attributes", "entries"), + [ + pytest.param({}, {SOURCE_ENTRY: "class Dags {}"}, id="no-attribute"), + pytest.param({"Airflow-Java-SDK-Dag-Code": SOURCE_ENTRY}, {}, id="no-entry"), + pytest.param({"Airflow-Java-SDK-Dag-Code": SOURCE_ENTRY}, {SOURCE_ENTRY: ""}, id="empty-entry"), + ], + ) + def test_placeholder_when_no_source_is_embedded(self, tmp_path, attributes, entries): + jar = self._jar(tmp_path, attributes=attributes, entries=entries) + + result = _importer().get_source_code(_definition(jar)) + + assert "embeds no Dag source" in result.source_code + assert result.language == "java" + + def test_placeholder_for_a_jar_without_manifest(self, tmp_path): + jar = make_jar(tmp_path / "lib.jar", entries={SOURCE_ENTRY: "class Dags {}"}) + + assert "embeds no Dag source" in _importer().get_source_code(_definition(jar)).source_code + + def test_placeholder_when_the_source_is_too_large(self, tmp_path): + jar = self._jar( + tmp_path, attributes={"Airflow-Java-SDK-Dag-Code": SOURCE_ENTRY}, entries={SOURCE_ENTRY: "x" * 9} + ) + + with patch("airflow.sdk.coordinators.java._dag_importer._MAX_SOURCE_BYTES", 8): + result = _importer().get_source_code(_definition(jar)) + + assert "is not shown" in result.source_code + + def test_replaces_invalid_utf8(self, tmp_path): + jar = self._jar( + tmp_path, + attributes={"Airflow-Java-SDK-Dag-Code": SOURCE_ENTRY}, + entries={SOURCE_ENTRY: b"a\xffb"}, + ) + + assert _importer().get_source_code(_definition(jar)).source_code == "a�b" + + def test_bad_zip_raises(self, tmp_path): + jar = tmp_path / "broken.jar" + jar.write_bytes(b"not a zip") + + with pytest.raises(zipfile.BadZipFile): + _importer().get_source_code(_definition(jar)) + + +class TestRegistry: + @pytest.fixture(autouse=True) + def _clean_registry(self): + reset_importer_registry() + yield + reset_importer_registry() + + def _coordinators(self, kwargs: dict) -> dict[tuple[str, str], str]: + spec = {"java": {"classpath": "airflow.sdk.coordinators.java.JavaCoordinator", "kwargs": kwargs}} + return {("sdk", "coordinators"): json.dumps(spec)} + + def test_a_bundle_backed_coordinator_registers_its_importer(self): + with conf_vars(self._coordinators({})): + importer = get_importer_registry("dags-folder").get_importer("dags/app.jar") + + assert isinstance(importer, JavaDagImporter) + + def test_an_explicit_root_coordinator_registers_none(self, tmp_path): + with conf_vars(self._coordinators({"jars_root": [str(tmp_path)]})): + importer = get_importer_registry("dags-folder").get_importer("dags/app.jar") + + assert importer is None diff --git a/task-sdk/tests/task_sdk/coordinators/java/test_jar_manifest.py b/task-sdk/tests/task_sdk/coordinators/java/test_jar_manifest.py new file mode 100644 index 0000000000000..cf9a724e8083b --- /dev/null +++ b/task-sdk/tests/task_sdk/coordinators/java/test_jar_manifest.py @@ -0,0 +1,94 @@ +# 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. + +from __future__ import annotations + +import zipfile + +import pytest +from task_sdk.coordinators.java._jar_test_utils import make_jar, write_manifest + +from airflow.sdk.coordinators.java._jar_manifest import parse_main_attributes, read_main_attributes + +LONG_VALUE = "META-INF/airflow/dag-code/org/apache/airflow/example/nativedag/InterfaceExample.java" + + +class TestParseMainAttributes: + @pytest.mark.parametrize("newline", [b"\r\n", b"\n", b"\r"], ids=["crlf", "lf", "cr"]) + def test_splits_on_every_line_ending(self, newline): + data = newline.join([b"Manifest-Version: 1.0", b"Main-Class: com.example.Main", b"", b""]) + + assert parse_main_attributes(data) == {"manifest-version": "1.0", "main-class": "com.example.Main"} + + @pytest.mark.parametrize("newline", [b"\r\n", b"\n"], ids=["crlf", "lf"]) + def test_unfolds_a_value_folded_once_and_twice(self, newline): + twice = "com.example." + "x" * 150 + data = write_manifest({"Airflow-Java-SDK-Dag-Code": LONG_VALUE, "Main-Class": twice}, newline=newline) + assert data.count(newline + b" ") == 3 + + attributes = parse_main_attributes(data) + + assert attributes["airflow-java-sdk-dag-code"] == LONG_VALUE + assert attributes["main-class"] == twice + + def test_decodes_a_character_the_fold_splits(self): + # "Main-Class: " is 12 bytes, so the two-byte "é" straddles the 72-byte fold. + value = "c" * 59 + "é" + "d" * 10 + data = write_manifest({"Main-Class": value}) + assert "é".encode() not in data + + assert parse_main_attributes(data)["main-class"] == value + + def test_stops_at_the_end_of_the_main_section(self): + data = ( + b"Main-Class: com.example.Main\r\n\r\nName: com/example/Main.class\r\nMain-Class: other.Main\r\n" + ) + + assert parse_main_attributes(data) == {"main-class": "com.example.Main"} + + def test_reads_a_manifest_without_a_trailing_newline(self): + assert parse_main_attributes(b"Manifest-Version: 1.0\r\nMain-Class: com.example.Main") == { + "manifest-version": "1.0", + "main-class": "com.example.Main", + } + + def test_names_are_case_insensitive_and_a_later_duplicate_wins(self): + data = b"MAIN-CLASS: first.Main\r\nmain-class: second.Main\r\n" + + assert parse_main_attributes(data) == {"main-class": "second.Main"} + + def test_ignores_a_stray_continuation_and_a_line_without_a_colon(self): + data = b" orphan\r\nnot a header\r\nMain-Class: com.example.Main\r\n" + + assert parse_main_attributes(data) == {"main-class": "com.example.Main"} + + def test_empty_manifest(self): + assert parse_main_attributes(b"") == {} + + +class TestReadMainAttributes: + def test_reads_the_manifest_entry(self, tmp_path): + jar = make_jar(tmp_path / "app.jar", attributes={"Main-Class": "com.example.Main"}) + + with zipfile.ZipFile(jar) as zf: + assert read_main_attributes(zf) == {"main-class": "com.example.Main"} + + def test_returns_none_without_a_manifest(self, tmp_path): + jar = make_jar(tmp_path / "lib.jar", entries={"com/example/Lib.class": b"\xca\xfe"}) + + with zipfile.ZipFile(jar) as zf: + assert read_main_attributes(zf) is None