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 eafd542df6955..e08ea14d0a57f 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst @@ -286,6 +286,32 @@ the Dag processor resolve ``task_handler_bundle_name`` through it, and wherever it is rejected if the name is missing there. The Dag processor still receives ``sales_pipeline.py`` through the separate Dag delivery process. +A Dag processor with this ``[sdk]`` configuration also parses the executable JARs of every Dag bundle, +and needs a JRE to do so (see :ref:`java-sdk/native-dag-parsing`). Keep it from parsing a bundle's JARs +with that bundle's ``.airflowignore``, whose pattern syntax follows ``[core] dag_ignore_file_syntax``. + +With ``task_handler_bundle_name`` set to its own bundle, as above, ``java-task-handlers`` holds only the +JARs the Python Dag's tasks run, so ignoring everything in it is safe: + +.. code-block:: bash + + # [core] dag_ignore_file_syntax = glob (the default) + echo '*' > /opt/airflow/jars/sales-pipeline/.airflowignore + + # [core] dag_ignore_file_syntax = regexp; a bare * is not a valid pattern and is dropped + echo '.' > /opt/airflow/jars/sales-pipeline/.airflowignore + +With ``task_handler_bundle_name`` unset, the JAR sits in the same Dag bundle as ``sales_pipeline.py``, so +ignore only the JAR, not the whole bundle: + +.. code-block:: bash + + # [core] dag_ignore_file_syntax = glob (the default) + echo '*.jar' >> /opt/airflow/dags/.airflowignore + + # [core] dag_ignore_file_syntax = regexp + echo '\.jar$' >> /opt/airflow/dags/.airflowignore + After Airflow has parsed the Dag, trigger it from the UI or command line: .. code-block:: bash @@ -606,6 +632,69 @@ 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`. +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"} + +Once a ``JavaCoordinator`` is configured, the Dag processor parses the executable JARs of every Dag bundle, +so it needs this ``[sdk]`` configuration and a JRE. With one ``JavaCoordinator``, it parses them all. +With several, map each Dag bundle that holds native Java Dags to one of them in ``[sdk] dag_bundle_to_coordinator``. +A JAR in a bundle that has no entry, or an entry that names no ``JavaCoordinator``, fails to parse with an import error: + +.. code-block:: ini + + [sdk] + dag_bundle_to_coordinator = {"dags-folder": "java-native"} + +A Dag bundle that holds only the JARs that Python Dags' tasks run should list ``*`` in its ``.airflowignore`` +either way: with one ``JavaCoordinator`` its JARs are still parsed on every loop, and a JAR built with a Java +SDK older than schema ``2026-10-30``, every released one today, fails to parse as an import error. With +several Java coordinators, an unmapped bundle's JARs fail to parse too. + +* 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. A JAR without ``Main-Class`` is + skipped, but many dependency JARs set one (the PostgreSQL JDBC driver and H2 do, for example), so a + thin bundle should set ``main_class`` or list its dependency JARs in ``.airflowignore``. +* Parsing needs a JAR built with a Java SDK whose supervisor schema version (the + ``Airflow-Supervisor-Schema-Version`` manifest attribute) is ``2026-10-30`` or later. An older JAR + fails to parse, so list it in ``.airflowignore``; its tasks still find it, because tasks do not read + ``.airflowignore``. +* 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``. +* Two JARs in the bundle that set the same ``Main-Class`` fail to parse, because the JVM would load + the classes of only one of them. Keep one in the bundle. +* 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. +* A task runs the JAR of its Dag on the ``JavaCoordinator`` that its queue routes to, + which need not be the one that parsed the JAR. For example, a queue can route to a coordinator that uses another JDK. + A task whose queue routes to another kind of coordinator fails without retries. A task whose JAR is missing, + or that the coordinator cannot run (for example, because ``main_class`` does not match the JAR's ``Main-Class``), + fails and retries while it has retries left. +* The Code view shows the source of the JAR's main class, which the Gradle plugin packs into the JAR. +* 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`` and ``tasks render`` refuse a native Java Dag. + ``airflow tasks list`` lists its tasks by running the JAR's main class, so it needs a JRE. + .. _java-sdk/logging: Logging @@ -873,6 +962,11 @@ The ``build/bundle/`` directory contains all required JAR(s). Copy or mount it i :class:`~airflow.sdk.coordinators.java.JavaCoordinator` scans that Dag bundle recursively and builds the classpath automatically. +The plugin also packs the source file of ``mainClass``, and of each class that declares a Dag in Java, +into the bundle JAR, so the Airflow UI can show the source of a native Java Dag (see +: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. + .. note:: You only need the ``annotationProcessor`` entry if you use the annotation-based API. It is not needed for @@ -1046,6 +1140,32 @@ directory into the Dag bundle named by ``task_handler_bundle_name``. 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 main class's source file under +``META-INF/airflow/sources/`` with an index that names it, ``src/main/resources/META-INF/airflow/sources.json``: + +.. code-block:: json + + {"entrypoint_path": "com/example/Main.java"} + +.. code-block:: xml + + + + + + src/main/resources + + + src/main/java/com/example + Main.java + META-INF/airflow/sources/com/example + + + + +Then add ``META-INF/airflow/sources.json`` to the +``manifestEntries`` of ``maven-shade-plugin`` or ``maven-jar-plugin`` shown above. + .. _java-sdk/coordinator-config: :class:`~airflow.sdk.coordinators.java.JavaCoordinator` configuration @@ -1077,8 +1197,10 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the * - ``main_class`` - *(auto-detect)* - Explicit entry-point class. If omitted, the coordinator scans the Dag bundle for a JAR - whose manifest sets ``Main-Class``. If more than one JAR in that Dag bundle sets it, which - one runs is non-deterministic, so set ``main_class`` explicitly in that case. + whose manifest sets ``Main-Class``. If more than one JAR in that Dag bundle sets it, the first + by path is used, so set ``main_class`` explicitly in that case. + A task of a native Java Dag runs the JAR the Dag was parsed from. + 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 @@ -1098,6 +1220,10 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the dependencies. To isolate conflicting dependency versions, put the handlers in a second Dag bundle served by a second coordinator on its own queue. + A task of a native Java Dag ignores ``task_handler_bundle_name``: + it runs the JAR of its Dag from the Dag's own bundle, at the version the run was created with. + See :ref:`java-sdk/native-dag-parsing`. + .. note:: The ``[sdk]`` configuration is read at startup, so changes to ``coordinators`` or diff --git a/task-sdk/src/airflow/sdk/coordinators/_dag_importer.py b/task-sdk/src/airflow/sdk/coordinators/_dag_importer.py index 54517ee38310d..f339546934a69 100644 --- a/task-sdk/src/airflow/sdk/coordinators/_dag_importer.py +++ b/task-sdk/src/airflow/sdk/coordinators/_dag_importer.py @@ -41,7 +41,9 @@ from airflow.sdk.coordinators._subprocess import SubprocessCoordinator from airflow.sdk.importers.base import DagDefinition -COORDINATOR_DAG_IMPORTERS: Final[tuple[str, ...]] = () +COORDINATOR_DAG_IMPORTERS: Final[tuple[str, ...]] = ( + "airflow.sdk.coordinators.java._dag_importer.JavaDagImporter", +) """ The classpaths of the :class:`CoordinatorDagImporter` subclasses a Dag bundle's registry may hold. diff --git a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py index 97d9caccae417..feed0da740b04 100644 --- a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py +++ b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py @@ -34,8 +34,10 @@ import selectors import signal import socket +import stat import subprocess import time +from pathlib import PurePosixPath from typing import TYPE_CHECKING, NoReturn, TypeVar, cast import attrs @@ -43,10 +45,12 @@ import structlog from airflow.dag_processing.bundles.base import BundleVersionLock, unpack_bundle_version # noqa: SDK002 +from airflow.sdk._shared.module_loading import import_string from airflow.sdk.api.datamodels._generated import BundleInfo from airflow.sdk.configuration import conf +from airflow.sdk.coordinators._dag_importer import find_claiming_importer from airflow.sdk.execution_time.bundles import initialize_ti_bundle -from airflow.sdk.execution_time.coordinator import BaseCoordinator +from airflow.sdk.execution_time.coordinator import BaseCoordinator, TaskLaunchError from airflow.sdk.execution_time.schema import get_schema_version_migrator from airflow.sdk.execution_time.supervisor import ActivitySubprocess, NeverRaised, ProcessTracker @@ -427,6 +431,54 @@ def _initialize_pinned_bundle(target: BundleInfo, logger: FilteringBoundLogger) return initialize_ti_bundle(BundleInfo(name=target.name, version=version, version_data=version_data)) +def _resolve_task_bundle(bundle_info: BundleInfo, logger: FilteringBoundLogger) -> BaseDagBundle: + """ + Materialize the Dag bundle of a task, pinned to a concrete version. + + *logger* is the task logger, so materialization failures surface in the task log. + + :raises TaskLaunchError: when the bundle cannot be read. + """ + cannot_read = f"Dag bundle {bundle_info.name!r} cannot be read" + try: + bundle = _initialize_pinned_bundle(bundle_info, logger) + except Exception as e: + raise TaskLaunchError(f"{cannot_read}: {e}") from e + try: + bundle.path.stat() + except (FileNotFoundError, NotADirectoryError): + raise TaskLaunchError(f"{cannot_read}: it resolved to {bundle.path}, which does not exist.") from None + except OSError as e: + raise TaskLaunchError(f"{cannot_read}: {e}") from e + return bundle + + +def _is_file_in_bundle(bundle: BaseDagBundle, rel_path: str) -> bool: + """ + Return whether *rel_path* names a file inside *bundle*. + + *rel_path* must be a relative path inside the bundle, with no ``..`` part. A symlink in the + bundle is followed, as the Dag processor follows it when it lists the bundle's files. + """ + path = PurePosixPath(rel_path) + if path.is_absolute() or ".." in path.parts: + return False + try: + return stat.S_ISREG((bundle.path / path).stat().st_mode) + except (FileNotFoundError, NotADirectoryError): + return False + + +def _require_supported_schema_version(version: str) -> None: + """Raise ``ValueError`` unless the supervisor of this worker's Task SDK knows *version*.""" + try: + get_schema_version_migrator().resolve_version(version) + except ValueError as e: + raise ValueError( + f"uses supervisor schema version {version!r}, which this worker's Task SDK does not support" + ) from e + + @attrs.define(kw_only=True) class SubprocessCoordinator(BaseCoordinator): """ @@ -446,6 +498,7 @@ class SubprocessCoordinator(BaseCoordinator): If unset, the task's own Dag bundle is used. A named bundle resolves to the version current when the task starts; the task's own bundle uses the run's version. Either way the resolved version is pinned for the whole task. + A task of a native Dag does not read this bundle: it runs its own Dag file. """ task_startup_timeout: float = 10.0 @@ -496,6 +549,18 @@ def _build_execute_task_command(self, *, what: TaskInstance) -> tuple[list[str], """ raise NotImplementedError + def _build_dag_file_command( + self, *, what: TaskInstance, path: pathlib.Path + ) -> tuple[list[str], str | None]: + """ + Build the command that runs *what* from the native Dag file at *path*. + + Subclasses can retrieve the root of the Dag bundle holding *path* with + :meth:`_get_scan_roots`. The contract is that of :meth:`_build_execute_task_command`. + Raise an exception when the file cannot run *what*. + """ + raise NotImplementedError(f"{type(self).__name__} cannot run the tasks of a native Dag") + def _build_parse_dag_command(self, *, path: pathlib.Path) -> tuple[list[str], str | None]: """ Build the command that parses the Dag file at *path* and resolve its wire-schema version. @@ -557,6 +622,72 @@ def _set_scan_roots(self, roots: Sequence[pathlib.Path]): finally: self._active_scan_roots = None + def _is_native_dag_file(self, rel_path: str, bundle_name: str) -> bool: + """ + Return whether *rel_path* is a native Dag file, which this coordinator must run. + + :raises TaskLaunchError: when the bundle's Dag importers cannot be built, the coordinator class + of the file's importer cannot be loaded, or the file is a native Dag file of a runtime that + this coordinator does not run. None of these retries. + """ + try: + importer = find_claiming_importer(rel_path, bundle_name) + except Exception as e: + raise TaskLaunchError( + f"Cannot tell whether {rel_path!r} is a native Dag file, because the Dag importers of " + f"Dag bundle {bundle_name!r} cannot be built: {e.__cause__ or e}", + retryable=False, + ) from e + if importer is None: + return False + try: + runtime = import_string(importer.coordinator_classpath) + is_runtime = isinstance(self, runtime) + except Exception as e: + raise TaskLaunchError( + f"{rel_path!r} is a native Dag file, but its coordinator class " + f"{importer.coordinator_classpath!r} cannot be loaded: {e}", + retryable=False, + ) from e + if not is_runtime: + raise TaskLaunchError( + f"{rel_path!r} is a native Dag file that a {runtime.__name__} runs, but the task's queue " + f"routes it to a {type(self).__name__}. Route the queue to a {runtime.__name__} in " + "[sdk] queue_to_coordinator.", + retryable=False, + ) + return True + + def _build_native_task_command( + self, *, what: TaskInstance, bundle: BaseDagBundle, rel_path: str + ) -> tuple[list[str], str | None]: + """ + Check that *rel_path* is a file of *bundle* that this coordinator can run, and return its command. + + :raises TaskLaunchError: when the file is missing or cannot be read, cannot run *what*, or needs + a supervisor schema version that this worker's Task SDK does not support. + """ + version = f" at version {bundle.version!r}" if bundle.version is not None else "" + try: + is_file = _is_file_in_bundle(bundle, rel_path) + except OSError as e: + raise TaskLaunchError( + f"Dag file {rel_path!r} in Dag bundle {bundle.name!r}{version} cannot be read: {e}" + ) from e + if not is_file: + raise TaskLaunchError( + f"Dag file {rel_path!r} is not a file in Dag bundle {bundle.name!r}{version}" + ) + try: + command, schema_version = self._build_dag_file_command(what=what, path=bundle.path / rel_path) + if schema_version is not None: + _require_supported_schema_version(schema_version) + except Exception as e: + raise TaskLaunchError( + f"Dag file {rel_path!r} in Dag bundle {bundle.name!r}{version} cannot run: {e}" + ) from e + return command, schema_version + def execute_task( self, *, @@ -569,7 +700,23 @@ def execute_task( subprocess_logs_to_stdout: bool, **kwargs, ) -> BaseCoordinator.ExecutionResult: - bundle = self._resolve_artifact_bundle(bundle_info, logger or log) + """ + Run *what*. + + A task of a native Dag runs its own Dag file, *dag_rel_path* in the Dag bundle of + *bundle_info* at the version of the run, and never another Dag file. Its runtime may still + read other files of the bundle, such as the JARs on a Java classpath. Any other task runs the + artifact found by scanning the bundle that holds the artifacts. + + :raises TaskLaunchError: when a native Dag file cannot run, before the runtime starts. + """ + task_logger = logger or log + rel_path = os.fspath(dag_rel_path) + is_native = self._is_native_dag_file(rel_path, bundle_info.name) + if is_native: + bundle = _resolve_task_bundle(bundle_info, task_logger) + else: + bundle = self._resolve_artifact_bundle(bundle_info, task_logger) # Hold the version lock across start()/wait() so bundle cleanup cannot # rmtree a version this task is still reading from, mirroring # task_runner.main() for the Python task path. @@ -577,7 +724,12 @@ def execute_task( BundleVersionLock(bundle_name=bundle.name, bundle_version=bundle.version), self._set_scan_roots([bundle.path]), ): - command, subprocess_schema_version = self._build_execute_task_command(what=what) + if is_native: + command, subprocess_schema_version = self._build_native_task_command( + what=what, bundle=bundle, rel_path=rel_path + ) + else: + command, subprocess_schema_version = self._build_execute_task_command(what=what) process = _PopenActivitySubprocess.start( what=what, dag_rel_path=dag_rel_path, 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..dfe81ba9be3be --- /dev/null +++ b/task-sdk/src/airflow/sdk/coordinators/java/_dag_importer.py @@ -0,0 +1,122 @@ +# +# 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 json +import zipfile +from typing import TYPE_CHECKING, Any, ClassVar, Final, cast + +from airflow.sdk.coordinators._dag_importer import CoordinatorDagImporter +from airflow.sdk.coordinators.java._jar_manifest import MAIN_CLASS, SOURCES, 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 + +_SOURCES_DIR: Final = "META-INF/airflow/sources/" +_MAX_SOURCE_BYTES: Final = 1024 * 1024 +_NO_SOURCE: Final = ( + "// This JAR embeds no Dag source. Build it with the Airflow Java SDK Gradle plugin to show the\n" + "// source here.\n" +) +_SOURCE_TOO_LARGE: Final = "// This Dag source file is over 1 MiB, so it is not shown.\n" + + +class JavaDagImporter(CoordinatorDagImporter): + """ + Claim the native Dags of Java bundle JARs. + + A :class:`~airflow.sdk.coordinators.java.JavaCoordinator` parses them. Only a JAR whose manifest sets + ``Main-Class`` (matching the coordinator's ``main_class`` if set) is parsed. + """ + + coordinator_classpath: ClassVar[str] = "airflow.sdk.coordinators.java.JavaCoordinator" + artifact_suffix: ClassVar[str] = ".jar" + supported_extensions = [".jar"] + + def might_contain_dag(self, definition: DagDefinition, safe_mode: bool) -> bool: + """ + Return whether the JAR sets ``Main-Class``, matching the parsing coordinator's ``main_class`` if 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. So is every JAR when no coordinator + can parse the bundle, so that each parse reports why. + """ + try: + with definition.as_file() as path, zipfile.ZipFile(path) as zf: + attributes = read_main_attributes(zf) or {} + except Exception: + return True + if not (main_class := attributes.get(MAIN_CLASS)): + return False + try: + wanted = cast("JavaCoordinator", self.get_parsing_coordinator()).main_class + except Exception: + return True + return not wanted or main_class == wanted + + def get_source_code(self, definition: DagDefinition) -> DagSourceCode: + """Return the entrypoint source the JAR embeds, or a notice when it embeds none.""" + with definition.as_file() as path, zipfile.ZipFile(path) as zf: + info = _find_source_entry(zf) + if info 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, dag_id: str | None = None) -> zipfile.ZipInfo | None: + """ + Return the JAR entry of the source embedded for *dag_id*, or ``None`` when there is none. + + The Gradle plugin packs each Dag's source file once, maps every Java-declared Dag to its file in + ``dag_source_paths``, and always packs the entrypoint (the ``Main-Class`` source). A Dag that is + not mapped, or no *dag_id*, falls back to the entrypoint. + """ + index = _read_source_index(zf) + paths = index.get("dag_source_paths") + source_path = paths.get(dag_id) if dag_id is not None and isinstance(paths, dict) else None + if not isinstance(source_path, str): + source_path = index.get("entrypoint_path") + if not isinstance(source_path, str): + return None + try: + return zf.getinfo(_SOURCES_DIR + source_path) + except KeyError: + return None + + +def _read_source_index(zf: zipfile.ZipFile) -> dict[str, Any]: + if not (entry := (read_main_attributes(zf) or {}).get(SOURCES)): + return {} + try: + info = zf.getinfo(entry) + except KeyError: + return {} + if info.file_size > _MAX_SOURCE_BYTES: + return {} + try: + index = json.loads(zf.read(info)) + except ValueError: + return {} + return index if isinstance(index, dict) else {} 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..12f381af54f79 --- /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" +SOURCES: Final = "airflow-java-sdk-sources" + +_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 5d483110580f2..c295709dd5c2d 100644 --- a/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py +++ b/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py @@ -19,18 +19,23 @@ from __future__ import annotations -import email import os import pathlib import stat import zipfile -from typing import TYPE_CHECKING +import zlib +from typing import TYPE_CHECKING, Final import attrs import structlog from airflow.sdk.coordinators._bundle_metadata import validate_schema_version from airflow.sdk.coordinators._subprocess import SubprocessCoordinator +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 @@ -42,6 +47,9 @@ log: FilteringBoundLogger = structlog.get_logger(logger_name="coordinators.java") +# The first supervisor schema version whose runtime can answer a Dag parse request. +_DAG_PARSING_SCHEMA_VERSION: Final = "2026-10-30" + def _find_jars(items: Iterable[pathlib.Path]) -> Iterator[pathlib.Path]: """ @@ -76,37 +84,54 @@ def _walk_jars(items: Iterable[pathlib.Path], seen_dirs: set[tuple[int, int]]) - def _iter_dir(directory: pathlib.Path) -> Iterator[pathlib.Path]: # iterdir() is lazy, so an unreadable directory raises only once iteration # starts; swallow it here so a single bad directory does not abort the scan. + # Sorted, so the JAR a scan picks does not depend on filesystem order. try: - yield from directory.iterdir() + children = sorted(directory.iterdir()) except OSError: return + yield from children -def _calculate_classpath(roots: Sequence[pathlib.Path]) -> str: - jars = (p.as_posix() for p in _find_jars(roots)) - return os.pathsep.join(sorted(jars)) # Keep output deterministic. +def _calculate_classpath(roots: Sequence[pathlib.Path], *, primary: pathlib.Path | None = None) -> str: + """Return the classpath of every JAR under *roots*. *primary*, when given, goes first.""" + jars = sorted(_find_jars(roots), key=lambda p: p.as_posix()) + if primary is not None: + target = primary.resolve() + jars = [primary, *(p for p in jars if p.resolve() != target)] + return os.pathsep.join(p.as_posix() for p in jars) @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 (FileNotFoundError, IsADirectoryError, zipfile.BadZipFile, zlib.error, NotImplementedError): 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)) + + +def _read_executable_jar(path: pathlib.Path) -> tuple[str, str | None]: + """Return the Main-Class and schema version of the JAR at *path*, raising ``ValueError`` if it cannot run.""" + try: + with zipfile.ZipFile(path) as zf: + attributes = read_main_attributes(zf) + except zipfile.BadZipFile as e: + raise ValueError(f"{path} is not a valid JAR: {e}") from e + if attributes is None: + raise ValueError(f"{path} has no META-INF/MANIFEST.MF") + if not (main_class := attributes.get(MAIN_CLASS)): + raise ValueError(f"{path} is not an executable JAR: its manifest sets no Main-Class") + return main_class, attributes.get(SUPERVISOR_SCHEMA_VERSION) @attrs.define @@ -124,6 +149,37 @@ def collect(self) -> _JarInfo | None: return None return _JarInfo(self.main_class, self.schema_version) + @classmethod + def for_jar( + cls, roots: Sequence[pathlib.Path], jar: pathlib.Path, main_class: str, schema_version: str | None + ) -> _JarInfo: + """ + Return how to run *jar*, whose manifest sets *main_class* and *schema_version*. + + Another JAR under *roots* that sets the same Main-Class is rejected, because the JVM would + load the classes of whichever comes first on the classpath. Without its own schema version, + *jar* takes the first one another JAR sets. + """ + target = jar.resolve() + same_main_class = [jar] + for p in _find_jars(roots): + if p.resolve() == target or (metadata := _JarMetadata.from_jar(p)) is None: + continue + if metadata.main_class == main_class: + same_main_class.append(p) + schema_version = schema_version or metadata.schema_version + if len(same_main_class) > 1: + paths = ", ".join(os.fspath(p) for p in same_main_class) + raise ValueError( + f"These JARs all set Main-Class {main_class!r}: {paths}. Keep one in the bundle." + ) + if schema_version is None: + raise FileNotFoundError( + "cannot find a JAR with Airflow-Supervisor-Schema-Version metadata in " + + os.pathsep.join(os.fspath(p.resolve()) for p in roots) + ) + return cls(main_class, schema_version) + @classmethod def find(cls, roots: Sequence[pathlib.Path], main_class: str) -> _JarInfo: log.debug("Finding JARs recursively", roots=roots) @@ -131,7 +187,11 @@ def find(cls, roots: Sequence[pathlib.Path], main_class: str) -> _JarInfo: for p in _find_jars(roots): if (metadata := _JarMetadata.from_jar(p)) is None: continue - if metadata.main_class and ((main_class == metadata.main_class) or not main_class): + if ( + progress.main_class is None + and metadata.main_class + and ((main_class == metadata.main_class) or not main_class) + ): log.debug("JAR located with Main-Class metadata", path=p, main_class=metadata.main_class) progress.main_class = metadata.main_class if metadata.schema_version: @@ -185,9 +245,9 @@ class JavaCoordinator(SubprocessCoordinator): If *main_class* is not explicitly set, JavaCoordinator scans the bundle to find an executable JAR (one with Main-Class set in its metadata). If more - than one executable JAR is found, it may be nondeterministic which one ends - up being executed, so set *main_class* when more than one JAR in the bundle - declares Main-Class. + than one executable JAR is found, the first by path is executed, so set *main_class* + when more than one JAR in the bundle declares Main-Class. A task of a native Java Dag + does not scan: it runs the JAR the Dag was parsed from. A JAR containing metadata *Airflow-Supervisor-Schema-Version* should also be available to specify the wire schema version. The JAR containing the Java @@ -195,6 +255,13 @@ 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. + The coordinator also parses native Java Dags: every JAR whose manifest sets + Main-Class (matching *main_class* when that is set) is run to list the Dags + its main class declares. With one JavaCoordinator configured, it parses the + JARs of every Dag bundle. With several, ``[sdk] dag_bundle_to_coordinator`` + picks the one that parses a bundle. A JAR whose Main-Class another JAR in the + bundle also sets is rejected. + 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 @@ -207,16 +274,52 @@ class JavaCoordinator(SubprocessCoordinator): jvm_args: list[str] = attrs.field(factory=list) main_class: str = "" - 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 = [ + def _build_command( + self, roots: Sequence[pathlib.Path], main_class: str, *, primary: pathlib.Path | None = None + ) -> list[str]: + return [ self.java_executable, "-classpath", - _calculate_classpath(roots), + _calculate_classpath(roots, primary=primary), *self.jvm_args, - jar.main_class, + main_class, ] - return command, jar.schema_version + + def _build_execute_task_command(self, *, what: TaskInstance) -> tuple[list[str], str | None]: + # Without main_class, the first executable JAR in path order wins; tracked at + # https://github.com/apache/airflow/issues/71134 + roots = self._get_scan_roots() + jar = _JarInfo.find(roots, self.main_class) + return self._build_command(roots, jar.main_class), jar.schema_version + + def _build_jar_command(self, path: pathlib.Path) -> tuple[list[str], str]: + """ + Build the command that runs the executable JAR at *path*, and return its schema version. + + The JAR's own ``Main-Class`` runs, so *main_class* must match it when set. The bundle's other + JARs go on the classpath, so another JAR that sets the same ``Main-Class`` is rejected. + + :raises ValueError: when the JAR cannot run. + """ + main_class, schema_version = _read_executable_jar(path) + if self.main_class and main_class != self.main_class: + raise ValueError( + f"{path} runs {main_class!r}, but this coordinator's main_class is {self.main_class!r}" + ) + roots = self._get_scan_roots() + jar = _JarInfo.for_jar(roots, path, main_class, schema_version) + return self._build_command(roots, jar.main_class, primary=path), jar.schema_version + + def _build_dag_file_command( + self, *, what: TaskInstance, path: pathlib.Path + ) -> tuple[list[str], str | None]: + return self._build_jar_command(path) + + def _build_parse_dag_command(self, *, path: pathlib.Path) -> tuple[list[str], str | None]: + command, schema_version = self._build_jar_command(path) + if schema_version < _DAG_PARSING_SCHEMA_VERSION: + raise ValueError( + f"{path} uses supervisor schema {schema_version}, which cannot parse Dags; " + "rebuild it with a newer Java SDK or list it in .airflowignore" + ) + return command, schema_version diff --git a/task-sdk/src/airflow/sdk/execution_time/coordinator.py b/task-sdk/src/airflow/sdk/execution_time/coordinator.py index 53b239fb856d9..311011f15f652 100644 --- a/task-sdk/src/airflow/sdk/execution_time/coordinator.py +++ b/task-sdk/src/airflow/sdk/execution_time/coordinator.py @@ -58,6 +58,7 @@ __all__ = [ "BaseCoordinator", "CoordinatorManager", + "TaskLaunchError", "get_coordinator_manager", "reset_coordinator_manager", ] @@ -65,6 +66,21 @@ log = structlog.get_logger(__name__) +class TaskLaunchError(Exception): + """ + A coordinator cannot launch the runtime of a task. + + A coordinator raises this from :meth:`BaseCoordinator.execute_task` before it starts the runtime. + The supervisor then fails the task instance and writes the message to the task log. A *retryable* + error leaves the task to retry while it has retries left. One that is not, such as a routing + mistake that retrying cannot fix, fails the task at once. + """ + + def __init__(self, message: str, *, retryable: bool = True) -> None: + super().__init__(message) + self.retryable = retryable + + class BaseCoordinator: """ Base coordinator for runtime-specific DAG file processing and task execution. @@ -98,6 +114,9 @@ def execute_task( Start task execution. This should execute the task and return a result. + + :raises TaskLaunchError: when the runtime cannot be launched. Raise it before the runtime + starts, so the supervisor can fail the task instance with the message in its task log. """ raise NotImplementedError diff --git a/task-sdk/src/airflow/sdk/execution_time/supervisor.py b/task-sdk/src/airflow/sdk/execution_time/supervisor.py index f2373b3bb835e..40a43208269c5 100644 --- a/task-sdk/src/airflow/sdk/execution_time/supervisor.py +++ b/task-sdk/src/airflow/sdk/execution_time/supervisor.py @@ -146,7 +146,7 @@ _RequestFrame, _ResponseFrame, ) -from airflow.sdk.execution_time.coordinator import get_coordinator_manager +from airflow.sdk.execution_time.coordinator import TaskLaunchError, get_coordinator_manager from airflow.sdk.execution_time.request_handlers import ( handle_delete_variable, handle_delete_xcom, @@ -1594,6 +1594,22 @@ def _remote_logging_conn(client: Client): del client +_START_REJECTED_AS_CLEARED_MESSAGE = ( + "Server rejected task start because the task was cleared. Task process stopped." +) + + +def _is_start_rejected_as_cleared(error: Exception) -> bool: + """Return whether *error* is the server refusing to start a task that was cleared while it was queued.""" + return ( + isinstance(error, ServerResponseError) + and error.response.status_code == HTTPStatus.CONFLICT + and isinstance(error.detail, dict) + and error.detail.get("reason") == "invalid_state" + and error.detail.get("previous_state") == "restarting" + ) + + @attrs.define(kw_only=True) class ActivitySubprocess(WatchedSubprocess): client: Client @@ -1680,16 +1696,8 @@ def _on_child_started( except Exception as e: # On any error kill that subprocess! self.kill(signal.SIGKILL) - if ( - isinstance(e, ServerResponseError) - and e.response.status_code == HTTPStatus.CONFLICT - and isinstance(e.detail, dict) - and e.detail.get("reason") == "invalid_state" - and e.detail.get("previous_state") == "restarting" - ): - self.process_log.info( - "Server rejected task start because the task was cleared. Task process stopped." - ) + if _is_start_rejected_as_cleared(e): + self.process_log.info(_START_REJECTED_AS_CLEARED_MESSAGE) self._terminal_state = SERVER_TERMINATED return raise @@ -2898,6 +2906,63 @@ def _configure_logging(log_path: str, client: Client) -> tuple[FilteringBoundLog return logger, log_file_descriptor +def _fail_task_before_runtime( + *, ti: TaskInstance, client: Client, logger: FilteringBoundLogger, error: TaskLaunchError +) -> int: + """ + Fail *ti* with *error*'s message, for a task whose runtime could not be launched. + + Starts the task instance, writes the message to the task log and the supervisor log, reports the + task as up for retry or failed with the message as its reason, and uploads the task log. A + retryable *error* puts the task up for retry while it has retries left. Returns the process exit + code, 1. + """ + from airflow.sdk._shared.secrets_masker import redact + from airflow.sdk.log import upload_to_remote + + try: + ti_context = client.task_instances.start(ti.id, os.getpid(), datetime.now(tz=timezone.utc)) + except Exception as e: + if _is_start_rejected_as_cleared(e): + logger.info(_START_REJECTED_AS_CLEARED_MESSAGE) + return 1 + raise + + reason = str(redact(str(error))) + log.warning("Cannot launch the task's runtime", ti_id=str(ti.id), reason=reason) + logger.error( + "Cannot launch the task's runtime", + reason=reason, + exc_info=error if error.__cause__ is not None else None, + ) + try: + end_date = datetime.now(tz=timezone.utc) + try: + if error.retryable and ti_context.should_retry: + client.task_instances.retry( + ti.id, end_date=end_date, rendered_map_index=None, retry_reason=reason[:500] + ) + else: + client.task_instances.finish( + id=ti.id, + state=TerminalStateNonSuccess.FAILED, + when=end_date, + rendered_map_index=None, + retry_reason=reason[:500], + ) + except ServerResponseError as e: + if e.response.status_code != HTTPStatus.CONFLICT: + raise + logger.info("Server rejected task outcome with a conflict. Discarding task outcome.") + finally: + try: + with _remote_logging_conn(client): + upload_to_remote(logger, ti) # type: ignore[arg-type] + except Exception: + logger.exception("Failed to upload remote logs", ti_id=ti.id) + return 1 + + def supervise_task( *, ti: TaskInstance, @@ -3018,6 +3083,13 @@ def supervise_task( final_state=result.final_state, ) return result.exit_code + except TaskLaunchError as e: + return _fail_task_before_runtime( + ti=ti, + client=client, + logger=logger or structlog.get_logger(logger_name="task").bind(), + error=e, + ) finally: if log_path and log_file_descriptor: log_file_descriptor.close() 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 56e21d970e612..36e7616089bc7 100644 --- a/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py @@ -24,21 +24,25 @@ import re import socket import subprocess -import zipfile +import zlib 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 +from airflow.sdk.api.datamodels._generated import BundleInfo, TaskInstance +from airflow.sdk.coordinators._subprocess import _PopenActivitySubprocess from airflow.sdk.coordinators.java.coordinator import ( JavaCoordinator, _calculate_classpath, _JarInfo, + _JarMetadata, _walk_jars, ) -from airflow.sdk.execution_time.coordinator import BaseCoordinator +from airflow.sdk.execution_time.coordinator import BaseCoordinator, TaskLaunchError from airflow.sdk.execution_time.supervisor import ActivitySubprocess +from airflow.sdk.importers import 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 @@ -67,15 +71,58 @@ 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 + + def test_permission_error_propagates(self, tmp_path): + jar = tmp_path / "locked.jar" + with ( + patch( + "airflow.sdk.coordinators.java.coordinator.zipfile.ZipFile", + side_effect=PermissionError(13, "Permission denied", str(jar)), + ), + pytest.raises(PermissionError, match="locked.jar"), + ): + _JarMetadata.from_jar(jar) + + @pytest.mark.parametrize( + "error", [zlib.error("incomplete or truncated stream"), NotImplementedError("compression type")] + ) + def test_a_corrupt_manifest_entry_gives_none(self, tmp_path, error): + jar = make_jar(tmp_path / "corrupt.jar", attributes={"Main-Class": "com.example.Main"}) + with patch("airflow.sdk.coordinators.java.coordinator.read_main_attributes", side_effect=error): + assert _JarMetadata.from_jar(jar) is None class TestCalculateClasspath: @@ -109,6 +156,16 @@ def test_empty_directory_returns_empty_string(self, tmp_path): result = _calculate_classpath([tmp_path]) assert result == "" + def test_primary_jar_goes_first(self, tmp_path): + a = tmp_path.joinpath("a.jar") + a.write_bytes(b"") + b = tmp_path.joinpath("b.jar") + b.write_bytes(b"") + + result = _calculate_classpath([tmp_path], primary=b) + + assert result.split(os.pathsep) == [b.as_posix(), a.as_posix()] + class TestMainJar: def test_returns_main_class_from_jar(self, tmp_path): @@ -178,6 +235,13 @@ def test_symlink_cycle_does_not_infinite_recurse(self, tmp_path): result = _JarInfo.find([tmp_path], "com.example.Loop") assert result == _JarInfo("com.example.Loop", "2026-06-16") + def test_the_first_executable_jar_wins_even_when_the_schema_version_is_elsewhere(self, tmp_path): + _make_jar(tmp_path.joinpath("a-etl.jar"), main_class="com.example.Etl") + _make_jar(tmp_path.joinpath("b-tools.jar"), main_class="com.example.Tools") + _make_jar(tmp_path.joinpath("c-sdk.jar"), main_class=None, schema_version="2026-06-16") + + assert _JarInfo.find([tmp_path], "") == _JarInfo("com.example.Etl", "2026-06-16") + class TestWalkJars: def test_skips_directory_whose_key_is_already_in_seen_dirs(self, tmp_path): @@ -413,3 +477,233 @@ def test_returns_execution_result(self, java_task_handlers, mock_client): assert isinstance(result, BaseCoordinator.ExecutionResult) assert result.exit_code == 0 + + +class TestJavaCoordinatorExecuteNativeDag: + """With a JavaCoordinator configured, a task of a native Java Dag runs the JAR of its Dag.""" + + @pytest.fixture(autouse=True) + def _java_coordinator_config(self): + coordinators = {"java": {"classpath": "airflow.sdk.coordinators.java.JavaCoordinator"}} + reset_importer_registry() + with conf_vars({("sdk", "coordinators"): json.dumps(coordinators)}): + yield + reset_importer_registry() + + @pytest.fixture + def mock_start(self, tmp_path): + _make_jar(tmp_path / "a.jar", main_class="com.example.A", schema_version="2026-06-16") + _make_jar(tmp_path / "b.jar", main_class="com.example.B", schema_version="2026-10-30") + bundle = MagicMock(path=tmp_path, version="v1") + bundle.name = "dags-folder" + with ( + patch("airflow.sdk.coordinators._subprocess._initialize_pinned_bundle", return_value=bundle), + patch("airflow.sdk.coordinators._subprocess.BundleVersionLock"), + patch.object(_PopenActivitySubprocess, "start") as mock_start, + ): + mock_start.return_value.wait.return_value = 0 + yield mock_start + + def _execute(self, rel_path: str, client, coordinator: JavaCoordinator | None = None): + return (coordinator or JavaCoordinator()).execute_task( + what=_make_ti(), + dag_rel_path=rel_path, + bundle_info=BundleInfo(name="dags-folder", version="v1"), + client=client, + subprocess_logs_to_stdout=False, + ) + + def test_runs_the_jar_of_the_dag_and_not_the_first_jar_by_path(self, mock_start, mock_client): + self._execute("b.jar", mock_client) + + assert mock_start.call_args.kwargs["command"][-1] == "com.example.B" + assert mock_start.call_args.kwargs["subprocess_schema_version"] == "2026-10-30" + + def test_a_task_of_a_python_dag_still_scans_the_bundle(self, mock_start, mock_client): + self._execute("dag.py", mock_client) + + assert mock_start.call_args.kwargs["command"][-1] == "com.example.A" + + def test_fails_without_starting_the_jvm_for_a_jar_the_coordinator_cannot_run( + self, mock_start, mock_client + ): + with pytest.raises( + TaskLaunchError, match="runs 'com.example.B', but .* main_class is 'com.example.A'" + ): + self._execute("b.jar", mock_client, JavaCoordinator(main_class="com.example.A")) + + mock_start.assert_not_called() + + +class TestBuildDagFileCommand: + def _build(self, coordinator: JavaCoordinator, root: pathlib.Path, jar: pathlib.Path): + with coordinator._set_scan_roots([root]): + return coordinator._build_dag_file_command(what=_make_ti(), path=jar) + + @pytest.mark.parametrize("parsed", ["a.jar", "b.jar"]) + def test_runs_the_main_class_of_the_jar_the_dag_was_parsed_from(self, tmp_path, parsed): + _make_jar(tmp_path / "a.jar", main_class="com.example.A", schema_version="2026-10-30") + _make_jar(tmp_path / "b.jar", main_class="com.example.B", schema_version="2026-10-30") + coordinator = JavaCoordinator() + + command, schema_version = self._build(coordinator, tmp_path, tmp_path / parsed) + + expected = {"a.jar": "com.example.A", "b.jar": "com.example.B"}[parsed] + assert command[-1] == expected + assert schema_version == "2026-10-30" + with coordinator._set_scan_roots([tmp_path]): + assert coordinator._build_parse_dag_command(path=tmp_path / parsed) == (command, schema_version) + + def test_runs_a_jar_whose_schema_version_is_too_old_to_parse(self, tmp_path): + jar = _make_jar(tmp_path / "app.jar", main_class="com.example.Dags", schema_version="2026-06-16") + + _, schema_version = self._build(JavaCoordinator(), tmp_path, jar) + + assert schema_version == "2026-06-16" + + def test_rejects_a_jar_whose_main_class_differs_from_the_pinned_one(self, tmp_path): + jar = _make_jar(tmp_path / "app.jar", main_class="com.example.Other", schema_version="2026-10-30") + + with pytest.raises(ValueError, match="main_class is 'com.example.Dags'"): + self._build(JavaCoordinator(main_class="com.example.Dags"), tmp_path, jar) + + def test_rejects_a_main_class_another_jar_sets(self, tmp_path): + new = _make_jar(tmp_path / "etl-new.jar", main_class="com.example.A", schema_version="2026-10-30") + _make_jar(tmp_path / "etl-old.jar", main_class="com.example.A", schema_version="2026-10-30") + + with pytest.raises(ValueError, match="all set Main-Class 'com.example.A'"): + self._build(JavaCoordinator(), tmp_path, new) + + +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-10-30") + 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-10-30" + + 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-10-30") + + 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-10-30" + + def test_the_parsed_jar_is_first_on_the_classpath(self, tmp_path): + """A JAR that shades the same classes as another JAR in the bundle must load its own.""" + sibling = _make_jar(tmp_path / "a-shaded.jar", main_class=None, schema_version="2026-06-16") + jar = _make_jar(tmp_path / "z-app.jar", main_class="com.example.Dags", schema_version="2026-10-30") + + command, _ = self._build(JavaCoordinator(), tmp_path, jar) + + assert command[2].split(os.pathsep) == [jar.as_posix(), sibling.as_posix()] + + def test_matches_the_execute_command(self, tmp_path): + jar = _make_jar(tmp_path / "app.jar", main_class="com.example.Dags", schema_version="2026-10-30") + _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 + + def test_takes_the_schema_version_of_the_parsed_jar(self, tmp_path): + _make_jar(tmp_path / "a.jar", main_class="com.example.A", schema_version="2026-06-16") + jar = _make_jar(tmp_path / "b.jar", main_class="com.example.B", schema_version="2026-10-30") + + _, schema_version = self._build(JavaCoordinator(), tmp_path, jar) + + assert schema_version == "2026-10-30" + + @pytest.mark.parametrize("own", [True, False], ids=["own", "sibling"]) + def test_rejects_a_schema_version_that_cannot_parse(self, tmp_path, own): + jar = _make_jar( + tmp_path / "app.jar", main_class="com.example.Dags", schema_version="2026-06-16" if own else None + ) + if not own: + _make_jar(tmp_path / "sdk.jar", main_class=None, schema_version="2026-06-16") + + with pytest.raises( + ValueError, + match=re.escape( + f"{jar} uses supervisor schema 2026-06-16, which cannot parse Dags; " + "rebuild it with a newer Java SDK or list it in .airflowignore" + ), + ): + self._build(JavaCoordinator(), tmp_path, jar) + + def test_rejects_a_main_class_another_jar_sets(self, tmp_path): + new = _make_jar(tmp_path / "etl-new.jar", main_class="com.example.A", schema_version="2026-10-30") + old = _make_jar(tmp_path / "etl-old.jar", main_class="com.example.A", schema_version="2026-10-30") + + with pytest.raises(ValueError, match="all set Main-Class 'com.example.A'") as excinfo: + self._build(JavaCoordinator(), tmp_path, new) + + assert os.fspath(new) in str(excinfo.value) + assert os.fspath(old) in str(excinfo.value) + assert "Keep one in the bundle" in str(excinfo.value) + + @pytest.mark.parametrize( + ("attributes", "match"), + [ + pytest.param(None, "has no META-INF/MANIFEST.MF", 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="broken.jar is not a valid JAR: File is not a zip file"): + self._build(JavaCoordinator(), tmp_path, jar) + + def test_rejects_a_jar_deleted_after_discovery(self, tmp_path): + with pytest.raises(FileNotFoundError, match="gone.jar"): + self._build(JavaCoordinator(), tmp_path, tmp_path / "gone.jar") + + @patch("airflow.sdk.coordinators._subprocess._set_close_on_exec_above_stderr", autospec=True) + @patch("airflow.sdk.coordinators._subprocess.signal.signal", autospec=True) + @patch("airflow.sdk.coordinators._subprocess.os.execvpe", autospec=True, side_effect=OSError("exec")) + def test_parse_dag_execs_the_jvm(self, mock_execvpe, mock_signal, mock_close_on_exec, tmp_path): + jar = _make_jar(tmp_path / "app.jar", main_class="com.example.Dags", schema_version="2026-10-30") + 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-10-30"] + 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..45ecda314eedb --- /dev/null +++ b/task-sdk/tests/task_sdk/coordinators/java/test_dag_importer.py @@ -0,0 +1,308 @@ +# 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 +import zlib +from unittest.mock import patch + +import pytest +from task_sdk.coordinators.java._jar_test_utils import SCHEMA_VERSION, make_jar + +from airflow.sdk.coordinators._dag_importer import find_claiming_importer +from airflow.sdk.coordinators.java import JavaCoordinator +from airflow.sdk.coordinators.java._dag_importer import JavaDagImporter, _find_source_entry +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" +SOURCES_INDEX = "META-INF/airflow/sources.json" +ENTRYPOINT = "com/example/Dags.java" +REPORTS = "com/example/dags/Reports.kt" +LONG_MAIN_CLASS = "org.apache.airflow.example.nativedag.generated.VeryLongNativeDagBundleBuilder" +JAVA_COORDINATOR = "airflow.sdk.coordinators.java.JavaCoordinator" + + +class _JavaCoordinatorSubclass(JavaCoordinator): + pass + + +def _importer() -> JavaDagImporter: + return JavaDagImporter(bundle_name="dags-folder") + + +def _coordinators(*kwargs: dict, mapping: dict[str, str] | None = None) -> dict[tuple[str, str], str]: + """Config with one JavaCoordinator per *kwargs*, keyed "java-0", "java-1" and so on.""" + specs = {f"java-{i}": {"classpath": JAVA_COORDINATOR, "kwargs": kw} for i, kw in enumerate(kwargs)} + config = {("sdk", "coordinators"): json.dumps(specs)} + if mapping is not None: + config[("sdk", "dag_bundle_to_coordinator")] = json.dumps(mapping) + return config + + +def _definition(path) -> FilesystemDagDefinition: + return FilesystemDagDefinition(path=path) + + +class _Bundle: + def __init__(self, path): + self.path = path + self.name = "dags-folder" + + +@pytest.fixture(autouse=True) +def _clean_registry(): + reset_importer_registry() + yield + reset_importer_registry() + + +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") + + @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""}) + + with conf_vars(_coordinators({"main_class": pin})): + assert _importer().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( + "error", [zlib.error("incomplete or truncated stream"), NotImplementedError("compression type")] + ) + @patch("airflow.sdk.coordinators.java._dag_importer.read_main_attributes") + def test_keeps_a_jar_whose_manifest_cannot_be_read(self, mock_read, tmp_path, error): + mock_read.side_effect = error + jar = make_jar(tmp_path / "app.jar", attributes={"Main-Class": MAIN_CLASS}) + + assert _importer().might_contain_dag(_definition(jar), safe_mode=True) is True + + @pytest.mark.parametrize(("mapped", "expected"), [("java-1", True), ("java-0", False)]) + def test_filters_by_the_main_class_of_the_coordinator_mapped_to_its_bundle( + self, tmp_path, mapped, expected + ): + jar = make_jar(tmp_path / "app.jar", attributes={"Main-Class": MAIN_CLASS}) + coordinators = _coordinators( + {"main_class": "com.example.Other"}, {"main_class": MAIN_CLASS}, mapping={"dags-folder": mapped} + ) + + with conf_vars(coordinators): + assert _importer().might_contain_dag(_definition(jar), safe_mode=True) is expected + + @pytest.mark.parametrize( + "config", + [ + pytest.param(_coordinators({}, {}), id="several-and-no-entry"), + pytest.param( + _coordinators({}, {}, mapping={"dags-folder": "gone"}), id="entry-names-no-coordinator" + ), + pytest.param({("sdk", "coordinators"): "{}"}, id="no-coordinator"), + pytest.param( + { + ("sdk", "coordinators"): json.dumps({"java": {"classpath": "no.such.Coordinator"}}), + }, + id="coordinator-that-cannot-be-loaded", + ), + ], + ) + def test_keeps_a_jar_when_no_coordinator_can_parse_its_bundle(self, tmp_path, config): + jar = make_jar(tmp_path / "app.jar", attributes={"Main-Class": MAIN_CLASS}) + bare = make_jar(tmp_path / "dep.jar", entries={"dep/Dep.class": b""}) + + with conf_vars(config): + importer = _importer() + + assert importer.might_contain_dag(_definition(jar), safe_mode=True) is True + assert importer.might_contain_dag(_definition(bare), safe_mode=True) is False + + @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}, + ) + + with conf_vars(_coordinators({})): + listed = list(_importer().list_dag_definitions(_Bundle(tmp_path), safe_mode=safe_mode)) + + assert [d.path for d in listed] == [app] + + +def _sources(index, files): + """JAR entries for a source index and the files it packs.""" + raw_index = index if isinstance(index, str) else json.dumps(index) + return {SOURCES_INDEX: raw_index, **{f"META-INF/airflow/sources/{p}": c for p, c in files.items()}} + + +class TestGetSourceCode: + def _jar(self, tmp_path, *, index=None, files=None, attributes=None): + attributes = { + "Main-Class": MAIN_CLASS, + "Airflow-Java-SDK-Sources": SOURCES_INDEX, + **(attributes or {}), + } + entries = _sources(index, files or {}) if index is not None else {} + return make_jar(tmp_path / "app.jar", attributes=attributes, entries=entries) + + def test_returns_the_entrypoint_source(self, tmp_path): + source = "public class Dags {}\n" + jar = self._jar( + tmp_path, + index={"entrypoint_path": ENTRYPOINT, "dag_source_paths": {"reports": REPORTS}}, + files={ENTRYPOINT: source, REPORTS: "class Reports"}, + attributes={"Main-Class": LONG_MAIN_CLASS}, + ) + 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( + ("dag_id", "expected"), + [ + pytest.param("reports", REPORTS, id="mapped"), + pytest.param("orders", ENTRYPOINT, id="unmapped"), + pytest.param(None, ENTRYPOINT, id="no-dag-id"), + ], + ) + def test_maps_a_dag_to_its_file_with_the_entrypoint_as_fallback(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"}, + ) + + with zipfile.ZipFile(jar) as zf: + assert _find_source_entry(zf, dag_id).filename == f"META-INF/airflow/sources/{expected}" + + @pytest.mark.parametrize( + ("attributes", "index", "files"), + [ + pytest.param( + {"Airflow-Java-SDK-Sources": ""}, + {"entrypoint_path": ENTRYPOINT}, + {ENTRYPOINT: "x"}, + id="no-attribute", + ), + pytest.param({}, None, {}, id="no-index"), + pytest.param({}, "{not json", {}, id="bad-index"), + pytest.param({}, ["not", "an", "object"], {}, id="index-not-an-object"), + pytest.param({}, {"dag_source_paths": {}}, {}, id="no-entrypoint"), + pytest.param({}, {"entrypoint_path": ENTRYPOINT}, {}, id="no-entry"), + pytest.param({}, {"entrypoint_path": ENTRYPOINT}, {ENTRYPOINT: ""}, id="empty-entry"), + ], + ) + def test_placeholder_when_no_source_is_embedded(self, tmp_path, attributes, index, files): + jar = self._jar(tmp_path, index=index, files=files, attributes=attributes) + + 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=_sources({"entrypoint_path": ENTRYPOINT}, {ENTRYPOINT: "x"}) + ) + + assert "embeds no Dag source" in _importer().get_source_code(_definition(jar)).source_code + + def test_placeholder_when_the_source_file_is_too_large(self, tmp_path): + jar = self._jar(tmp_path, index={"entrypoint_path": ENTRYPOINT}, files={ENTRYPOINT: "x" * 200}) + + with patch("airflow.sdk.coordinators.java._dag_importer._MAX_SOURCE_BYTES", 100): + 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, index={"entrypoint_path": ENTRYPOINT}, files={ENTRYPOINT: b"a\xffb"}) + + assert _importer().get_source_code(_definition(jar)).source_code == "a\ufffdb" + + 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: + def test_a_java_coordinator_registers_the_importer_in_every_bundle(self): + bundles = [ + {"name": name, "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", "kwargs": {}} + for name in ("dags-folder", "java-jars") + ] + coordinators = _coordinators({"task_handler_bundle_name": "java-jars"}) + with conf_vars({**coordinators, ("dag_processor", "dag_bundle_config_list"): json.dumps(bundles)}): + in_named = get_importer_registry("java-jars").get_importer("java-jars/app.jar") + in_other = get_importer_registry("dags-folder").get_importer("dags/app.jar") + + assert isinstance(in_named, JavaDagImporter) + assert isinstance(in_other, JavaDagImporter) + assert in_other.bundle_name == "dags-folder" + + def test_without_a_java_coordinator_no_jar_is_claimed(self): + with conf_vars({("sdk", "coordinators"): "{}"}): + assert get_importer_registry("dags-folder").get_importer("dags/app.jar") is None + assert find_claiming_importer("dags/app.jar", "dags-folder") is None + + def test_a_subclass_of_the_java_coordinator_registers_the_importer(self): + subclass = {"java": {"classpath": f"{__name__}._JavaCoordinatorSubclass"}} + with conf_vars({("sdk", "coordinators"): json.dumps(subclass)}): + claiming = find_claiming_importer("dags/app.jar", "dags-folder") + + assert isinstance(claiming, JavaDagImporter) 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..44839855a1244 --- /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 = "org.apache.airflow.example.nativedag.generated.VeryLongNativeDagBundleBuilderMainClass" + + +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-Sources": LONG_VALUE, "Main-Class": twice}, newline=newline) + assert data.count(newline + b" ") == 3 + + attributes = parse_main_attributes(data) + + assert attributes["airflow-java-sdk-sources"] == 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 diff --git a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py index 22fbb3acb14e0..d5fd0689064fd 100644 --- a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py +++ b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py @@ -18,6 +18,7 @@ from __future__ import annotations import contextlib +import json import os import pathlib import signal @@ -36,19 +37,22 @@ from airflow.dag_processing.bundles.base import BaseDagBundle, BundleVersion from airflow.sdk.api.client import Client, TaskInstanceOperations from airflow.sdk.api.datamodels._generated import BundleInfo, TaskInstance +from airflow.sdk.coordinators._dag_importer import CoordinatorDagImporter from airflow.sdk.coordinators._subprocess import ( SubprocessCoordinator, _accept_connections, _connection_owned_by_process_tree, _is_connection_from_pid, _is_connection_from_process, + _is_file_in_bundle, _PopenActivitySubprocess, _ResourceTracker, _start_server, log, ) -from airflow.sdk.execution_time.coordinator import BaseCoordinator +from airflow.sdk.execution_time.coordinator import BaseCoordinator, TaskLaunchError from airflow.sdk.execution_time.supervisor import ActivitySubprocess +from airflow.sdk.importers import DagSourceCode, 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 @@ -638,6 +642,13 @@ class _Plain(SubprocessCoordinator): with pytest.raises(NotImplementedError): _Plain()._build_execute_task_command(what=_make_ti()) + def test_build_dag_file_command_default_raises(self, tmp_path): + class _Plain(SubprocessCoordinator): + pass + + with pytest.raises(NotImplementedError, match="_Plain cannot run the tasks of a native Dag"): + _Plain()._build_dag_file_command(what=_make_ti(), path=tmp_path / "dag.native") + @pytest.mark.usefixtures("artifact_bundle") class TestSubprocessCoordinatorExecuteTask: @@ -1036,6 +1047,380 @@ def test_named_bundle_mode_locks_the_tree_it_scans( mock_lock.assert_called_once_with(bundle_name="artifacts", bundle_version="sha-abc") +@attrs.define(kw_only=True) +class _NativeStubCoordinator(_StubSubprocessCoordinator): + """Runs native Dag files, recording the roots and the Dag files its command builder is given.""" + + recorded_dag_files: list[pathlib.Path] = attrs.field(init=False, factory=list) + scans: int = attrs.field(init=False, default=0) + + def _build_execute_task_command(self, *, what): + self.scans += 1 + return super()._build_execute_task_command(what=what) + + def _build_dag_file_command(self, *, what, path): + self.recorded_roots.append(list(self._get_scan_roots())) + self.recorded_dag_files.append(path) + return [*self.command, os.fspath(path)], self.schema_version + + +@attrs.define(kw_only=True) +class _BrokenNativeCoordinator(_NativeStubCoordinator): + def _build_dag_file_command(self, *, what, path): + raise ValueError("no Main-Class") + + +@attrs.define(kw_only=True) +class _OtherStubCoordinator(_StubSubprocessCoordinator): + """A coordinator of another class than the one that runs native Dag files.""" + + +class _NativeDagImporter(CoordinatorDagImporter): + coordinator_classpath = f"{__name__}._NativeStubCoordinator" + artifact_suffix = ".native" + supported_extensions = [".native"] + + def get_source_code(self, definition) -> DagSourceCode: + return DagSourceCode("", "native") + + +class _UnloadableDagImporter(_NativeDagImporter): + coordinator_classpath = "nonexistent.module.Coordinator" + + +@contextlib.contextmanager +def _native_dag_files(*keys: str, mapping: dict[str, str] | None = None): + """Make ``.native`` files native Dag files, with a _NativeStubCoordinator configured for each key.""" + config = { + ("sdk", "coordinators"): json.dumps( + { + key: {"classpath": f"{__name__}._NativeStubCoordinator", "kwargs": {"command": ["/runtime"]}} + for key in keys + } + ) + } + if mapping is not None: + config[("sdk", "dag_bundle_to_coordinator")] = json.dumps(mapping) + with conf_vars(config): + reset_importer_registry() + yield + + +@patch( + "airflow.sdk.coordinators._dag_importer.COORDINATOR_DAG_IMPORTERS", (f"{__name__}._NativeDagImporter",) +) +class TestExecuteTaskNativeDagFile: + """A task of a native Dag runs its own Dag file at the version of the run, with no scan.""" + + BUNDLE_INFO = BundleInfo(name="dags", version="v9") + + @pytest.fixture(autouse=True) + def _clean_registry(self): + reset_importer_registry() + yield + reset_importer_registry() + + @pytest.fixture + def bundle(self, tmp_path): + """The task's Dag bundle, resolved at version v9 to a directory that holds one native Dag file.""" + (tmp_path / "sub").mkdir() + (tmp_path / "sub" / "dag.native").write_text("") + (tmp_path / "dir.native").mkdir() + return _make_bundle(tmp_path, version="v9") + + @pytest.fixture + def mock_initialize(self, bundle): + with patch( + "airflow.sdk.coordinators._subprocess.initialize_ti_bundle", autospec=True, return_value=bundle + ) as m: + yield m + + @pytest.fixture + def mock_lock(self): + with patch("airflow.sdk.coordinators._subprocess.BundleVersionLock", autospec=True) as mock_lock: + yield mock_lock + + @pytest.fixture + def mock_start(self): + with patch.object(_PopenActivitySubprocess, "start", autospec=True) as mock_start: + mock_start.return_value.wait.return_value = 0 + yield mock_start + + def _execute(self, coordinator, client, *, rel_path="sub/dag.native", bundle_info=BUNDLE_INFO): + return coordinator.execute_task( + what=_make_ti(), + dag_rel_path=rel_path, + bundle_info=bundle_info, + client=client, + subprocess_logs_to_stdout=False, + ) + + @pytest.mark.usefixtures("mock_initialize") + def test_runs_the_dag_file_at_the_version_of_the_run_without_a_scan( + self, mock_start, mock_lock, mock_initialize, mock_client, tmp_path + ): + coordinator = _NativeStubCoordinator(command=["/runtime"], schema_version="2026-06-16") + + with _native_dag_files("native"): + result = self._execute(coordinator, mock_client) + + dag_file = tmp_path / "sub" / "dag.native" + mock_initialize.assert_called_once_with(self.BUNDLE_INFO) + assert coordinator.recorded_dag_files == [dag_file] + assert coordinator.recorded_roots == [[tmp_path]] + assert coordinator.scans == 0 + assert coordinator._active_scan_roots is None + mock_lock.assert_called_once_with(bundle_name="dags", bundle_version="v9") + mock_lock.return_value.__exit__.assert_called_once() + assert mock_start.call_args.kwargs["command"] == ["/runtime", os.fspath(dag_file)] + assert mock_start.call_args.kwargs["subprocess_schema_version"] == "2026-06-16" + assert result.exit_code == 0 + + @pytest.mark.usefixtures("mock_start", "mock_lock") + def test_does_not_read_the_artifact_bundle_of_the_coordinator(self, mock_initialize, mock_client): + coordinator = _NativeStubCoordinator(command=["/runtime"], task_handler_bundle_name="artifacts") + + with _native_dag_files("native"): + self._execute(coordinator, mock_client) + + mock_initialize.assert_called_once_with(self.BUNDLE_INFO) + + @pytest.mark.usefixtures("mock_start", "mock_lock") + def test_pins_a_run_without_a_version_to_the_current_version( + self, mock_initialize, mock_lock, mock_client, tmp_path + ): + pinned_tree = tmp_path / "versions" / "sha-abc" + (pinned_tree / "sub").mkdir(parents=True) + (pinned_tree / "sub" / "dag.native").write_text("") + unpinned = _make_bundle(tmp_path, version=None) + unpinned.get_current_version.return_value = BundleVersion(version="sha-abc", data=None) + pinned = _make_bundle(pinned_tree, version="sha-abc") + mock_initialize.side_effect = [unpinned, pinned] + coordinator = _NativeStubCoordinator(command=["/runtime"]) + + with _native_dag_files("native"): + self._execute(coordinator, mock_client, bundle_info=BundleInfo(name="dags")) + + assert coordinator.recorded_dag_files == [pinned_tree / "sub" / "dag.native"] + mock_lock.assert_called_once_with(bundle_name="dags", bundle_version="sha-abc") + + @pytest.mark.usefixtures("mock_initialize", "mock_lock") + @pytest.mark.parametrize( + "mapping", + [ + pytest.param(None, id="no-entry-for-the-bundle"), + pytest.param({"dags": "first"}, id="entry-for-another-coordinator"), + ], + ) + @patch.object(CoordinatorDagImporter, "get_parsing_coordinator", autospec=True) + def test_runs_on_a_coordinator_that_does_not_parse_the_bundle( + self, mock_get_parsing_coordinator, mock_start, mock_client, tmp_path, mapping + ): + queue_coordinator = _NativeStubCoordinator(command=["/runtime"]) + + with _native_dag_files("first", "second", mapping=mapping): + self._execute(queue_coordinator, mock_client) + + mock_get_parsing_coordinator.assert_not_called() + assert queue_coordinator.recorded_dag_files == [tmp_path / "sub" / "dag.native"] + mock_start.assert_called_once() + + @pytest.mark.usefixtures("mock_start", "mock_lock") + def test_fails_without_a_retry_for_a_coordinator_of_another_class(self, mock_initialize, mock_client): + coordinator = _OtherStubCoordinator(command=["/runtime"]) + + with _native_dag_files("native"), pytest.raises(TaskLaunchError) as raised: + self._execute(coordinator, mock_client) + + assert raised.value.retryable is False + assert str(raised.value) == ( + "'sub/dag.native' is a native Dag file that a _NativeStubCoordinator runs, but the task's queue " + "routes it to a _OtherStubCoordinator. Route the queue to a _NativeStubCoordinator in " + "[sdk] queue_to_coordinator." + ) + mock_initialize.assert_not_called() + + @pytest.mark.usefixtures("mock_start", "mock_lock") + @patch( + "airflow.sdk.coordinators._subprocess.find_claiming_importer", + return_value=_UnloadableDagImporter(bundle_name="dags"), + ) + def test_fails_without_a_retry_when_the_coordinator_class_cannot_be_loaded( + self, _, mock_initialize, mock_client + ): + coordinator = _NativeStubCoordinator(command=["/runtime"]) + + with pytest.raises(TaskLaunchError) as raised: + self._execute(coordinator, mock_client) + + assert raised.value.retryable is False + assert str(raised.value).startswith( + "'sub/dag.native' is a native Dag file, but its coordinator class " + "'nonexistent.module.Coordinator' cannot be loaded: " + ) + assert isinstance(raised.value.__cause__, ImportError) + mock_initialize.assert_not_called() + + @pytest.mark.usefixtures("mock_start", "mock_lock") + @pytest.mark.parametrize("rel_path", ["sub/dag.native", "dag.py"]) + def test_fails_without_a_retry_when_the_dag_importers_cannot_be_built( + self, mock_initialize, mock_client, rel_path + ): + coordinator = _NativeStubCoordinator(command=["/runtime"]) + + with ( + patch( + "airflow.sdk.coordinators._dag_importer.COORDINATOR_DAG_IMPORTERS", + ("nonexistent.module.Importer",), + ), + _native_dag_files("native"), + pytest.raises(TaskLaunchError) as raised, + ): + self._execute(coordinator, mock_client, rel_path=rel_path) + + assert raised.value.retryable is False + assert str(raised.value).startswith( + f"Cannot tell whether {rel_path!r} is a native Dag file, because the Dag importers of " + "Dag bundle 'dags' cannot be built: " + ) + assert isinstance(raised.value.__cause__, RuntimeError) + assert coordinator.scans == 0 + mock_initialize.assert_not_called() + + @pytest.mark.usefixtures("mock_start", "mock_lock", "mock_initialize") + @pytest.mark.parametrize( + "rel_path", ["sub/gone.native", "dir.native", "/sub/dag.native", "../dag.native"] + ) + def test_fails_with_a_retry_for_a_file_that_is_not_in_the_bundle(self, mock_client, rel_path): + coordinator = _NativeStubCoordinator(command=["/runtime"]) + + with _native_dag_files("native"), pytest.raises(TaskLaunchError) as raised: + self._execute(coordinator, mock_client, rel_path=rel_path) + + assert raised.value.retryable is True + assert ( + str(raised.value) == f"Dag file {rel_path!r} is not a file in Dag bundle 'dags' at version 'v9'" + ) + assert coordinator.recorded_dag_files == [] + + @pytest.mark.usefixtures("mock_start", "mock_lock", "mock_initialize") + @patch("airflow.sdk.coordinators._subprocess._is_file_in_bundle", side_effect=PermissionError("denied")) + def test_fails_with_a_retry_when_the_dag_file_cannot_be_read(self, _, mock_client): + coordinator = _NativeStubCoordinator(command=["/runtime"]) + + with _native_dag_files("native"), pytest.raises(TaskLaunchError) as raised: + self._execute(coordinator, mock_client) + + assert raised.value.retryable is True + assert str(raised.value) == ( + "Dag file 'sub/dag.native' in Dag bundle 'dags' at version 'v9' cannot be read: denied" + ) + assert coordinator.recorded_dag_files == [] + + @pytest.mark.usefixtures("mock_start", "mock_lock", "mock_initialize") + def test_fails_with_a_retry_when_the_dag_file_cannot_run(self, mock_client): + coordinator = _BrokenNativeCoordinator(command=["/runtime"]) + + with _native_dag_files("native"), pytest.raises(TaskLaunchError) as raised: + self._execute(coordinator, mock_client) + + assert raised.value.retryable is True + assert str(raised.value) == ( + "Dag file 'sub/dag.native' in Dag bundle 'dags' at version 'v9' cannot run: no Main-Class" + ) + assert isinstance(raised.value.__cause__, ValueError) + + @pytest.mark.usefixtures("mock_start", "mock_lock", "mock_initialize") + def test_fails_with_a_retry_for_a_schema_version_the_worker_does_not_support(self, mock_client): + coordinator = _NativeStubCoordinator(command=["/runtime"], schema_version="1999-01-01") + + with _native_dag_files("native"), pytest.raises(TaskLaunchError) as raised: + self._execute(coordinator, mock_client) + + assert raised.value.retryable is True + assert str(raised.value).startswith( + "Dag file 'sub/dag.native' in Dag bundle 'dags' at version 'v9' cannot run: " + "uses supervisor schema version '1999-01-01', which this worker's Task SDK does not support" + ) + + @pytest.mark.usefixtures("mock_start", "mock_lock") + @pytest.mark.parametrize( + ("resolve", "expected"), + [ + pytest.param(OSError("denied"), "Dag bundle 'dags' cannot be read: denied", id="not-resolved"), + pytest.param( + None, + "Dag bundle 'dags' cannot be read: it resolved to {missing}, which does not exist.", + id="gone", + ), + ], + ) + def test_fails_with_a_retry_when_the_bundle_cannot_be_read( + self, mock_initialize, mock_client, tmp_path, resolve, expected + ): + missing = tmp_path / "gone" + if resolve is not None: + mock_initialize.side_effect = resolve + else: + mock_initialize.return_value = MagicMock(path=missing, version="v9") + coordinator = _NativeStubCoordinator(command=["/runtime"]) + + with _native_dag_files("native"), pytest.raises(TaskLaunchError) as raised: + self._execute(coordinator, mock_client) + + assert raised.value.retryable is True + assert str(raised.value) == expected.format(missing=missing) + + @pytest.mark.usefixtures("mock_start", "mock_lock") + def test_fails_with_a_retry_when_the_bundle_path_cannot_be_inspected(self, mock_initialize, mock_client): + path = MagicMock(spec=pathlib.Path) + path.stat.side_effect = PermissionError("denied") + mock_initialize.return_value = MagicMock(path=path, version="v9") + coordinator = _NativeStubCoordinator(command=["/runtime"]) + + with _native_dag_files("native"), pytest.raises(TaskLaunchError) as raised: + self._execute(coordinator, mock_client) + + assert raised.value.retryable is True + assert str(raised.value) == "Dag bundle 'dags' cannot be read: denied" + + @pytest.mark.usefixtures("mock_lock", "artifact_bundle") + @pytest.mark.parametrize("rel_path", ["dag.py", "sub/dag.native.py", "sub/dag"]) + def test_a_file_that_is_not_native_scans_the_bundle(self, mock_start, mock_client, rel_path): + coordinator = _NativeStubCoordinator(command=["/runtime"]) + + with _native_dag_files("native"): + self._execute(coordinator, mock_client, rel_path=rel_path) + + assert coordinator.scans == 1 + assert coordinator.recorded_dag_files == [] + assert mock_start.call_args.kwargs["command"] == ["/runtime"] + + +class TestIsFileInBundle: + @pytest.fixture + def bundle(self, tmp_path): + (tmp_path / "sub").mkdir() + (tmp_path / "sub" / "dag.native").write_text("") + (tmp_path / "link.native").symlink_to(tmp_path / "sub" / "dag.native") + return MagicMock(path=tmp_path) + + @pytest.mark.parametrize( + ("rel_path", "expected"), + [ + pytest.param("sub/dag.native", True, id="file"), + pytest.param("link.native", True, id="symlink-to-a-file"), + pytest.param("sub", False, id="directory"), + pytest.param("sub/gone.native", False, id="missing"), + pytest.param("sub/dag.native/more", False, id="under-a-file"), + pytest.param("/sub/dag.native", False, id="absolute"), + pytest.param("sub/../sub/dag.native", False, id="parent-part"), + ], + ) + def test_names_a_file_inside_the_bundle(self, bundle, rel_path, expected): + assert _is_file_in_bundle(bundle, rel_path) is expected + + class TestGetScanRoots: def test_returns_bound_roots_and_clears_them_on_exit(self, tmp_path): coordinator = _StubSubprocessCoordinator(command=["x"]) diff --git a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py index 12daa4173d58f..7ea6f112db0b5 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_supervisor.py +++ b/task-sdk/tests/task_sdk/execution_time/test_supervisor.py @@ -57,6 +57,7 @@ from airflow.executors.workloads import BundleInfo from airflow.sdk import DAG, BaseOperator, timezone +from airflow.sdk._shared.secrets_masker import mask_secret from airflow.sdk.api import client as sdk_client from airflow.sdk.api.client import ServerResponseError from airflow.sdk.api.datamodels._generated import ( @@ -163,6 +164,7 @@ _RequestFrame, _ResponseFrame, ) +from airflow.sdk.execution_time.coordinator import BaseCoordinator, TaskLaunchError from airflow.sdk.execution_time.supervisor import ( SERVER_TERMINATED, ActivitySubprocess, @@ -184,6 +186,8 @@ from tests_common.test_utils.config import conf_vars if TYPE_CHECKING: + from collections.abc import Callable + import kgb from airflow.sdk.definitions.context import Context @@ -270,6 +274,258 @@ def test_supervise( supervise_task(**kw) +def _response_error(status: int, detail: dict[str, Any]) -> ServerResponseError: + error = ServerResponseError.from_response( + httpx.Response(status, json={"detail": detail}, request=httpx.Request("PATCH", "http://server/x")) + ) + assert error is not None + return error + + +class TestSuperviseTaskLaunchError: + """A task whose runtime cannot be launched fails its try in the supervisor, with the reason.""" + + LAUNCH_EVENT = "Cannot launch the task's runtime" + + @pytest.fixture + def ti(self): + return TaskInstance( + id=uuid7(), + task_id="extract", + dag_id="etl", + run_id="r", + try_number=1, + dag_version_id=uuid7(), + queue="java", + ) + + @pytest.fixture + def client(self, make_ti_context): + client = MagicMock(spec=sdk_client.Client) + client.task_instances.start.return_value = make_ti_context(should_retry=False) + return client + + @pytest.fixture + def upload_to_remote(self): + with patch("airflow.sdk.log.upload_to_remote", autospec=True) as mock_upload: + yield mock_upload + + def _supervise( + self, client, ti, error: BaseException | Callable[..., Any], *, log_path: str | None = None + ) -> int: + coordinator = MagicMock(spec=BaseCoordinator) + coordinator.execute_task.side_effect = error + with patch.object(supervisor, "get_coordinator_manager", autospec=True) as mock_manager: + mock_manager.return_value.for_queue.return_value = coordinator + return supervise_task( + ti=ti, + bundle_info=BundleInfo(name="dags", version="v1"), + dag_rel_path="etl.jar", + token="", + client=client, + log_path=log_path, + ) + + def _read_task_log(self, tmp_path, log_path: str) -> list[dict[str, Any]]: + return [json.loads(line) for line in (tmp_path / log_path).read_text().splitlines() if line] + + def test_fails_the_try_with_the_reason(self, client, ti, upload_to_remote): + exit_code = self._supervise(client, ti, TaskLaunchError("no jar for 'extract'")) + + assert exit_code == 1 + client.task_instances.start.assert_called_once_with(ti.id, os.getpid(), mock.ANY) + client.task_instances.finish.assert_called_once_with( + id=ti.id, + state=TaskInstanceState.FAILED, + when=mock.ANY, + rendered_map_index=None, + retry_reason="no jar for 'extract'", + ) + client.task_instances.retry.assert_not_called() + + def test_retries_the_try_when_the_task_has_retries_left( + self, client, ti, make_ti_context, upload_to_remote + ): + client.task_instances.start.return_value = make_ti_context(should_retry=True, max_tries=2) + + exit_code = self._supervise(client, ti, TaskLaunchError("no jar for 'extract'")) + + assert exit_code == 1 + client.task_instances.retry.assert_called_once_with( + ti.id, end_date=mock.ANY, rendered_map_index=None, retry_reason="no jar for 'extract'" + ) + client.task_instances.finish.assert_not_called() + + def test_fails_the_try_of_an_error_that_is_not_retryable_even_with_retries_left( + self, client, ti, make_ti_context, upload_to_remote + ): + client.task_instances.start.return_value = make_ti_context(should_retry=True, max_tries=2) + + exit_code = self._supervise(client, ti, TaskLaunchError("wrong coordinator", retryable=False)) + + assert exit_code == 1 + client.task_instances.finish.assert_called_once_with( + id=ti.id, + state=TaskInstanceState.FAILED, + when=mock.ANY, + rendered_map_index=None, + retry_reason="wrong coordinator", + ) + client.task_instances.retry.assert_not_called() + + def test_writes_the_reason_to_the_task_log(self, client, ti, tmp_path, upload_to_remote): + with conf_vars({("logging", "base_log_folder"): str(tmp_path)}): + self._supervise(client, ti, TaskLaunchError("no jar for 'extract'"), log_path="etl/extract.log") + + entries = self._read_task_log(tmp_path, "etl/extract.log") + assert [(e["level"], e["event"], e["reason"]) for e in entries if "reason" in e] == [ + ("error", self.LAUNCH_EVENT, "no jar for 'extract'") + ] + + def test_logs_the_cause_of_the_error(self, client, ti, tmp_path, upload_to_remote): + try: + try: + raise OSError("disk gone") + except OSError as cause: + raise TaskLaunchError("cannot read the bundle") from cause + except TaskLaunchError as e: + error = e + + with conf_vars({("logging", "base_log_folder"): str(tmp_path)}): + self._supervise(client, ti, error, log_path="etl/extract.log") + + (entry,) = [e for e in self._read_task_log(tmp_path, "etl/extract.log") if "reason" in e] + assert "disk gone" in json.dumps(entry["exception"]) + + def test_truncates_the_reason_for_the_server_but_not_the_log( + self, client, ti, tmp_path, upload_to_remote + ): + reason = "x" * 700 + + with conf_vars({("logging", "base_log_folder"): str(tmp_path)}): + self._supervise(client, ti, TaskLaunchError(reason), log_path="etl/extract.log") + + assert client.task_instances.finish.call_args.kwargs["retry_reason"] == "x" * 500 + (entry,) = [e for e in self._read_task_log(tmp_path, "etl/extract.log") if "reason" in e] + assert entry["reason"] == reason + + @pytest.mark.enable_redact + @pytest.mark.parametrize( + ("message", "expected"), + [ + pytest.param( + "cannot read https://user:hunter2-hunter2@git/repo", + "cannot read https://user:***@git/repo", + id="short", + ), + pytest.param("x" * 490 + "hunter2-hunter2", "x" * 490 + "***", id="across-the-cut"), + ], + ) + def test_redacts_a_masked_secret_in_the_reason_before_cutting_it( + self, client, ti, tmp_path, upload_to_remote, message, expected + ): + def fail_with_a_secret(**kwargs): + mask_secret("hunter2-hunter2") + raise TaskLaunchError(message) + + with conf_vars({("logging", "base_log_folder"): str(tmp_path)}): + self._supervise(client, ti, fail_with_a_secret, log_path="etl/extract.log") + + assert client.task_instances.finish.call_args.kwargs["retry_reason"] == expected + assert "hunter2" not in (tmp_path / "etl/extract.log").read_text() + + def test_uploads_the_task_log(self, client, ti, upload_to_remote): + self._supervise(client, ti, TaskLaunchError("no jar")) + + upload_to_remote.assert_called_once_with(mock.ANY, ti) + + def test_uploads_the_task_log_when_reporting_the_state_fails(self, client, ti, upload_to_remote): + client.task_instances.finish.side_effect = _response_error(500, {"reason": "boom"}) + + with pytest.raises(ServerResponseError): + self._supervise(client, ti, TaskLaunchError("no jar")) + + upload_to_remote.assert_called_once() + + def test_an_upload_failure_is_logged_and_not_raised(self, client, ti, tmp_path, upload_to_remote): + upload_to_remote.side_effect = OSError("disk gone") + + with conf_vars({("logging", "base_log_folder"): str(tmp_path)}): + exit_code = self._supervise(client, ti, TaskLaunchError("no jar"), log_path="etl/extract.log") + + assert exit_code == 1 + client.task_instances.finish.assert_called_once() + events = [e["event"] for e in self._read_task_log(tmp_path, "etl/extract.log")] + assert "Failed to upload remote logs" in events + + @pytest.mark.parametrize( + "should_retry", [pytest.param(False, id="finish"), pytest.param(True, id="retry")] + ) + def test_a_conflict_on_the_state_is_logged_and_the_try_is_left_to_the_server( + self, client, ti, make_ti_context, tmp_path, upload_to_remote, should_retry + ): + client.task_instances.start.return_value = make_ti_context( + should_retry=should_retry, max_tries=2 if should_retry else 0 + ) + reported = client.task_instances.retry if should_retry else client.task_instances.finish + reported.side_effect = _response_error(409, {"reason": "invalid_state"}) + + with conf_vars({("logging", "base_log_folder"): str(tmp_path)}): + exit_code = self._supervise(client, ti, TaskLaunchError("no jar"), log_path="etl/extract.log") + + assert exit_code == 1 + reported.assert_called_once() + events = [e["event"] for e in self._read_task_log(tmp_path, "etl/extract.log")] + assert "Server rejected task outcome with a conflict. Discarding task outcome." in events + upload_to_remote.assert_called_once() + + def test_a_cleared_task_reports_no_state_and_logs_it(self, client, ti, tmp_path, upload_to_remote): + client.task_instances.start.side_effect = _response_error( + 409, {"reason": "invalid_state", "previous_state": "restarting"} + ) + + with conf_vars({("logging", "base_log_folder"): str(tmp_path)}): + exit_code = self._supervise(client, ti, TaskLaunchError("no jar"), log_path="etl/extract.log") + + assert exit_code == 1 + client.task_instances.finish.assert_not_called() + client.task_instances.retry.assert_not_called() + events = [e["event"] for e in self._read_task_log(tmp_path, "etl/extract.log")] + assert "Server rejected task start because the task was cleared. Task process stopped." in events + + def test_logs_the_reason_to_the_supervisor_log(self, client, ti, captured_logs, upload_to_remote): + self._supervise(client, ti, TaskLaunchError("no jar for 'extract'")) + + supervisor_entries = [e for e in captured_logs if e["event"] == self.LAUNCH_EVENT and "ti_id" in e] + assert [(e["level"], e["ti_id"], e["reason"]) for e in supervisor_entries] == [ + ("warning", str(ti.id), "no jar for 'extract'") + ] + + def test_a_task_already_running_elsewhere_propagates(self, client, ti, upload_to_remote): + client.task_instances.start.side_effect = TaskAlreadyRunningError("already running") + + with pytest.raises(TaskAlreadyRunningError): + self._supervise(client, ti, TaskLaunchError("no jar")) + + client.task_instances.finish.assert_not_called() + + def test_another_start_error_propagates(self, client, ti, upload_to_remote): + client.task_instances.start.side_effect = _response_error(500, {"reason": "boom"}) + + with pytest.raises(ServerResponseError): + self._supervise(client, ti, TaskLaunchError("no jar")) + + client.task_instances.finish.assert_not_called() + + def test_other_errors_from_the_coordinator_are_not_reported_as_a_launch_failure( + self, client, ti, upload_to_remote + ): + with pytest.raises(RuntimeError, match="boom"): + self._supervise(client, ti, RuntimeError("boom")) + + client.task_instances.start.assert_not_called() + + @pytest.mark.usefixtures("disable_capturing") class TestWatchedSubprocess: @pytest.fixture(autouse=True)