diff --git a/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst b/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst index 067d2d0bc6ac3..2d53b116c5499 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst @@ -409,14 +409,32 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the - Default - Description * - ``executables_root`` - - *(required)* + - *(optional)* - One or more directories scanned recursively for executable bundles. Accepts a string, - a path, or a list of strings/paths. + a path, or a list of strings/paths. When omitted, bundles are located through a Dag + bundle instead (see the note below). Explicitly setting this option to ``null`` or + an empty list is invalid. + * - ``dag_bundle_name`` + - *(auto: task's own bundle)* + - Name of a configured Dag bundle to load executable bundles from. Mutually exclusive + with ``executables_root``. * - ``task_startup_timeout`` - ``10.0`` - Seconds to wait for the bundle subprocess to connect after launch. Increase this if your bundle startup is slow (e.g. on constrained hardware). +.. note:: + + **Locating bundles.** ``executables_root`` and ``dag_bundle_name`` are mutually exclusive, + and both are optional: + + * Set ``executables_root`` to scan explicit filesystem directories you manage yourself. + * Set ``dag_bundle_name`` to load bundles from a configured Dag bundle, so they are delivered + and versioned through the same bundle machinery as your Dags. The task uses the version that + bundle is on when it starts, pinned for the whole task. + * Leave both unset (the default) to load bundles from the **task's own** Dag bundle, pinned + to the version the run was created with. + .. _go-sdk/limitations: Limitations 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 c8b73b4626a28..069c036c74911 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst @@ -663,9 +663,15 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the - Default - Description * - ``jars_root`` - - *(required)* + - *(optional)* - One or more directories scanned recursively for ``.jar`` files. Accepts a string, - a path, or a list of strings/paths. + a path, or a list of strings/paths. When omitted, JARs are located through a Dag + bundle instead (see the note below). Explicitly setting this option to ``null`` or + an empty list is invalid. + * - ``dag_bundle_name`` + - *(auto: task's own bundle)* + - Name of a configured Dag bundle to load JARs from. Mutually exclusive with + ``jars_root``. * - ``java_executable`` - ``"java"`` - Path to the ``java`` binary. Defaults to ``java`` on ``$PATH``. @@ -674,15 +680,27 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the - Extra JVM arguments such as ``["-Xmx1g", "-Dsome.property=value"]``. * - ``main_class`` - *(auto-detect)* - - Explicit entry-point class. If omitted, - :class:`~airflow.sdk.coordinators.java.JavaCoordinator` scans ``jars_root`` for a - JAR whose manifest sets ``Main-Class``. If multiple executable JARs are found the - result is non-deterministic; set ``main_class`` explicitly in that case. + - Explicit entry-point class. If omitted, the coordinator scans for a JAR whose + manifest sets ``Main-Class`` — in ``jars_root`` when set, otherwise across the + resolved Dag bundle. If multiple executable JARs match the result is + non-deterministic; set ``main_class`` explicitly in that case. * - ``task_startup_timeout`` - ``10.0`` - Seconds to wait for the JVM subprocess to connect after launch. Increase this if your JVM startup is slow (e.g. on constrained hardware or with a large classpath). +.. note:: + + **Locating JARs.** ``jars_root`` and ``dag_bundle_name`` are mutually exclusive, and both + are optional: + + * Set ``jars_root`` to scan explicit filesystem directories you manage yourself. + * Set ``dag_bundle_name`` to load JARs from a configured Dag bundle, so they are delivered and + versioned through the same bundle machinery as your Dags. The task uses the version that bundle + is on when it starts, pinned for the whole task. + * Leave both unset (the default) to load JARs from the **task's own** Dag bundle, pinned to + the version the run was created with. + .. note:: The ``[sdk]`` configuration is read at startup, so changes to ``coordinators`` or diff --git a/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst b/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst index 266aad2d0c408..df4d487140415 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst @@ -289,9 +289,15 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the - Default - Description * - ``bundles_root`` - - *(required)* + - *(optional)* - One or more directories searched, in order, for a ``bundle.mjs`` (with embedded metadata, or with an - ``airflow-metadata.yaml`` sidecar). Accepts a string, a path, or a list of strings/paths. + ``airflow-metadata.yaml`` sidecar). Accepts a string, a path, or a list of strings/paths. When + omitted, the bundle is located through a Dag bundle instead (see the note below). Explicitly setting + this option to ``null`` or an empty list is invalid. + * - ``dag_bundle_name`` + - *(auto: task's own bundle)* + - Name of a configured Dag bundle to load the ``bundle.mjs`` from. Mutually exclusive with + ``bundles_root``. * - ``node_executable`` - ``"node"`` - Path to the ``node`` binary. Defaults to ``node`` on ``$PATH``. @@ -300,6 +306,18 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the - Seconds to wait for the Node.js subprocess to connect after launch. Increase this if your bundle startup is slow (e.g. on constrained hardware). +.. note:: + + **Locating the bundle.** ``bundles_root`` and ``dag_bundle_name`` are mutually exclusive, and both + are optional: + + * Set ``bundles_root`` to scan explicit filesystem directories you manage yourself. + * Set ``dag_bundle_name`` to load the bundle from a configured Dag bundle, so it is delivered + and versioned through the same bundle machinery as your Dags. The task uses the version that + bundle is on when it starts, pinned for the whole task. + * Leave both unset (the default) to load the bundle from the **task's own** Dag bundle, pinned to the + version the run was created with. + Limitations ----------- diff --git a/airflow-core/src/airflow/dag_processing/bundles/manager.py b/airflow-core/src/airflow/dag_processing/bundles/manager.py index 292bb440a21e7..2431e014fdd84 100644 --- a/airflow-core/src/airflow/dag_processing/bundles/manager.py +++ b/airflow-core/src/airflow/dag_processing/bundles/manager.py @@ -16,13 +16,14 @@ # under the License. from __future__ import annotations +import functools import importlib import logging import os import warnings from collections import defaultdict from pathlib import Path -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any, NamedTuple, cast from itsdangerous import URLSafeSerializer from pydantic import BaseModel, ValidationError @@ -107,6 +108,34 @@ def _bundle_item_exc(msg): ) +class _BundleConfigSnapshot(NamedTuple): + """The configured Dag bundles, as names only and as full configs.""" + + configs: tuple[_ExternalBundleConfig, ...] + names: frozenset[str] + + +@functools.cache +def _load_bundle_config_snapshot() -> _BundleConfigSnapshot: + """Read and validate the configured Dag bundles, without importing their classes.""" + config_list = conf.getjson("dag_processor", "dag_bundle_config_list") + if not config_list: + return _BundleConfigSnapshot(configs=(), names=frozenset()) + if not isinstance(config_list, list): + raise AirflowConfigException( + "Section `dag_processor` key `dag_bundle_config_list` " + f"must be list but got {config_list.__class__}" + ) + bundle_config_list = _parse_bundle_config(config_list) + if conf.getboolean("core", "LOAD_EXAMPLES"): + _add_example_dag_bundle(bundle_config_list) + _add_provider_example_dags_to_bundle(bundle_config_list) + return _BundleConfigSnapshot( + configs=tuple(bundle_config_list), + names=frozenset(cfg.name for cfg in bundle_config_list), + ) + + def _parse_bundle_config(config_list) -> list[_ExternalBundleConfig]: bundles = {} for item in config_list: @@ -274,18 +303,9 @@ def parse_config(self) -> None: if self._bundle_config: return - config_list = conf.getjson("dag_processor", "dag_bundle_config_list") - if not config_list: + bundle_config_list = _load_bundle_config_snapshot().configs + if not bundle_config_list: return - if not isinstance(config_list, list): - raise AirflowConfigException( - "Section `dag_processor` key `dag_bundle_config_list` " - f"must be list but got {config_list.__class__}" - ) - bundle_config_list = _parse_bundle_config(config_list) - if conf.getboolean("core", "LOAD_EXAMPLES"): - _add_example_dag_bundle(bundle_config_list) - _add_provider_example_dags_to_bundle(bundle_config_list) for bundle_config in bundle_config_list: if bundle_config.team_name and not conf.getboolean("core", "multi_team"): @@ -653,6 +673,16 @@ def get_bundle( name=name, version=version, version_data=version_data, **cfg_bundle.kwargs ) + @classmethod + def is_bundle_configured(cls, name: str) -> bool: + """ + Return whether *name* is a configured Dag bundle. + + Deliberately reads configured names only: a caller validating a bundle name + must not depend on every *other* configured bundle's class being importable. + """ + return name in _load_bundle_config_snapshot().names + def get_all_dag_bundles(self) -> Iterable[BaseDagBundle]: """ Get all DAG bundles. diff --git a/airflow-core/tests/integration/otel/test_otel.py b/airflow-core/tests/integration/otel/test_otel.py index 39d7ee6bdab61..5be887bbdfa7d 100644 --- a/airflow-core/tests/integration/otel/test_otel.py +++ b/airflow-core/tests/integration/otel/test_otel.py @@ -330,7 +330,7 @@ def test_export_metrics_during_process_shutdown(self, capfd): # Additional detail spans are deferred to follow-up PRs; tracked # at https://linear.app/astronomer/issue/ACD-157. { - "_verify_bundle_access": "parse", + "verify_bundle_access": "parse", "parse": "startup", "get_template_context": "startup", "startup": "worker.task1", diff --git a/airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py b/airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py index bbd18ec80d85e..6d522a0813c55 100644 --- a/airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py +++ b/airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py @@ -29,7 +29,11 @@ from sqlalchemy import func, select, update from airflow.dag_processing.bundles.base import BaseDagBundle -from airflow.dag_processing.bundles.manager import DagBundlesManager, _guess_best_bundle_for_fileloc +from airflow.dag_processing.bundles.manager import ( + DagBundlesManager, + _guess_best_bundle_for_fileloc, + _load_bundle_config_snapshot, +) from airflow.exceptions import AirflowConfigException from airflow.models.dag import DagModel from airflow.models.dag_version import DagVersion @@ -167,6 +171,55 @@ def test_get_bundle(): assert bundle.version is None +@pytest.fixture(autouse=True) +def _clear_bundle_config_snapshot_cache(): + _load_bundle_config_snapshot.cache_clear() + yield + _load_bundle_config_snapshot.cache_clear() + + +def _bundle_config_env(config) -> dict[str, str]: + return { + "AIRFLOW__CORE__LOAD_EXAMPLES": "False", + "AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST": json.dumps(config), + } + + +def test_is_bundle_configured(): + """is_bundle_configured reports membership without constructing the bundle.""" + with patch.dict(os.environ, _bundle_config_env(BASIC_BUNDLE_CONFIG)): + assert DagBundlesManager.is_bundle_configured("my-test-bundle") is True + assert DagBundlesManager.is_bundle_configured("bundle-that-doesn't-exist") is False + + +def test_is_bundle_configured_does_not_import_bundle_class(): + """A validly-named bundle whose class is unimportable is still reported as configured. + + The check reads names only; the class is imported lazily at materialization, + so an unimportable classpath must not make a coordinator fail at construction. + """ + config = [{"name": "artifacts", "classpath": "does.not.exist.Bundle", "kwargs": {}}] + with patch.dict(os.environ, _bundle_config_env(config)): + assert DagBundlesManager.is_bundle_configured("artifacts") is True + + +def test_is_bundle_configured_empty_config(): + with patch.dict(os.environ, _bundle_config_env([])): + assert DagBundlesManager.is_bundle_configured("anything") is False + + +def test_is_bundle_configured_sees_implicitly_added_bundles(): + """Bundles parse_config adds implicitly are configured for the name check too.""" + with patch.dict( + os.environ, + { + "AIRFLOW__CORE__LOAD_EXAMPLES": "True", + "AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST": json.dumps(BASIC_BUNDLE_CONFIG), + }, + ): + assert DagBundlesManager.is_bundle_configured("example_dags") is True + + @pytest.fixture def clear_db(): clear_db_dag_bundles() @@ -183,9 +236,7 @@ def _get_bundle_names_and_active(): ).all() # Initial add - with patch.dict( - os.environ, {"AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST": json.dumps(BASIC_BUNDLE_CONFIG)} - ): + with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(BASIC_BUNDLE_CONFIG)}): manager = DagBundlesManager() manager.sync_bundles_to_db() assert _get_bundle_names_and_active() == [("my-test-bundle", True)] @@ -210,9 +261,7 @@ def _get_bundle_names_and_active(): assert session.scalar(select(func.count(ParseImportError.id))) == 0 # Re-enable one that reappears in config - with patch.dict( - os.environ, {"AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST": json.dumps(BASIC_BUNDLE_CONFIG)} - ): + with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(BASIC_BUNDLE_CONFIG)}): manager = DagBundlesManager() manager.sync_bundles_to_db() assert _get_bundle_names_and_active() == [ @@ -260,16 +309,12 @@ def _get_bundle_names_and_active(): ).all() # Processor A: only knows "my-test-bundle". - with patch.dict( - os.environ, {"AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST": json.dumps(BASIC_BUNDLE_CONFIG)} - ): + with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(BASIC_BUNDLE_CONFIG)}): DagBundlesManager().sync_bundles_to_db(deactivate_missing=False) assert _get_bundle_names_and_active() == [("my-test-bundle", True)] # Processor B: only knows "other-test-bundle". It must not touch A's bundle. - with patch.dict( - os.environ, {"AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST": json.dumps(OTHER_BUNDLE_CONFIG)} - ): + with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(OTHER_BUNDLE_CONFIG)}): DagBundlesManager().sync_bundles_to_db(deactivate_missing=False) assert _get_bundle_names_and_active() == [ ("my-test-bundle", True), @@ -277,9 +322,7 @@ def _get_bundle_names_and_active(): ] # Processor A runs again: still must not disable B's bundle. - with patch.dict( - os.environ, {"AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST": json.dumps(BASIC_BUNDLE_CONFIG)} - ): + with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(BASIC_BUNDLE_CONFIG)}): DagBundlesManager().sync_bundles_to_db(deactivate_missing=False) assert _get_bundle_names_and_active() == [ ("my-test-bundle", True), @@ -377,9 +420,7 @@ def test_bundle_model_render_url(clear_db, session): @conf_vars({("core", "LOAD_EXAMPLES"): "False"}) def test_template_params_update_on_sync(clear_db, session): """Test that template parameters are updated when bundle configuration changes.""" - with patch.dict( - os.environ, {"AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST": json.dumps(TEMPLATE_BUNDLE_CONFIG)} - ): + with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(TEMPLATE_BUNDLE_CONFIG)}): manager = DagBundlesManager() manager.sync_bundles_to_db() @@ -402,9 +443,7 @@ def test_template_params_update_on_sync(clear_db, session): } ] - with patch.dict( - os.environ, {"AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST": json.dumps(updated_config)} - ): + with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(updated_config)}): manager = DagBundlesManager() manager.sync_bundles_to_db() @@ -421,9 +460,7 @@ def test_template_params_update_on_sync(clear_db, session): def test_template_update_on_sync(clear_db, session): """Test that templates are updated when bundle configuration changes.""" # First, sync with initial template - with patch.dict( - os.environ, {"AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST": json.dumps(TEMPLATE_BUNDLE_CONFIG)} - ): + with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(TEMPLATE_BUNDLE_CONFIG)}): manager = DagBundlesManager() manager.sync_bundles_to_db() @@ -446,9 +483,7 @@ def test_template_update_on_sync(clear_db, session): } ] - with patch.dict( - os.environ, {"AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST": json.dumps(updated_config)} - ): + with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(updated_config)}): manager = DagBundlesManager() manager.sync_bundles_to_db() diff --git a/contributing-docs/30_new_language_sdk.rst b/contributing-docs/30_new_language_sdk.rst index 180a3e6539c0a..6b23ce7115651 100644 --- a/contributing-docs/30_new_language_sdk.rst +++ b/contributing-docs/30_new_language_sdk.rst @@ -120,6 +120,12 @@ The method returns a ``(command, subprocess_schema_version)`` pair: subprocess understands, used by the supervisor to negotiate message formats across SDK versions. See `Supervisor Schema`_ below. +Call ``self._get_scan_roots()`` to retrieve the artifact directories the base +class has already resolved from the coordinator's configured source — an +explicit filesystem root, a named Dag bundle (``dag_bundle_name``), or the +task's own bundle. Subclasses should scan those roots rather than reading the +configured root directly. + Supervisor Schema ~~~~~~~~~~~~~~~~~ diff --git a/devel-common/src/tests_common/pytest_plugin.py b/devel-common/src/tests_common/pytest_plugin.py index b5e1df054e419..0a7fda6b6871b 100644 --- a/devel-common/src/tests_common/pytest_plugin.py +++ b/devel-common/src/tests_common/pytest_plugin.py @@ -2147,6 +2147,32 @@ def reset_team_name_cache(): clear_team_name_cache() +@pytest.fixture(autouse=True) +def reset_dag_bundle_config_cache(): + """Reset the per-process Dag bundle configuration cache between tests. + + The configuration is parsed once per process, so a test that sets a different + ``[dag_processor] dag_bundle_config_list`` would otherwise be served the previous + test's bundles. ``conf_vars`` clears it too, for tests that switch config midway. + """ + if importlib.util.find_spec("airflow") is None: + yield + return + + try: + from airflow.dag_processing.bundles.manager import _load_bundle_config_snapshot + except ImportError: + # compat for airflow versions without the snapshot cache + yield + return + + _load_bundle_config_snapshot.cache_clear() + try: + yield + finally: + _load_bundle_config_snapshot.cache_clear() + + @pytest.fixture(autouse=True) def refuse_to_run_test_from_wrongly_named_files(request: pytest.FixtureRequest): filepath = request.node.path diff --git a/devel-common/src/tests_common/test_utils/config.py b/devel-common/src/tests_common/test_utils/config.py index 9a278346d3068..f4de8f598185d 100644 --- a/devel-common/src/tests_common/test_utils/config.py +++ b/devel-common/src/tests_common/test_utils/config.py @@ -98,6 +98,7 @@ def conf_vars(overrides): if "airflow.configuration" in sys.modules: settings.configure_vars() + _clear_dag_bundle_config_cache() try: yield @@ -116,6 +117,18 @@ def conf_vars(overrides): if "airflow.configuration" in sys.modules: settings.configure_vars() + _clear_dag_bundle_config_cache() + + +def _clear_dag_bundle_config_cache() -> None: + """Drop the per-process Dag bundle configuration cache so the new config is read.""" + import sys + + manager = sys.modules.get("airflow.dag_processing.bundles.manager") + # compat for airflow versions without the snapshot cache + cache = getattr(manager, "_load_bundle_config_snapshot", None) + if cache is not None: + cache.cache_clear() @overload diff --git a/generated/known_airflow_exceptions.txt b/generated/known_airflow_exceptions.txt index acfd9ae3eb825..d3f088e99e6ab 100644 --- a/generated/known_airflow_exceptions.txt +++ b/generated/known_airflow_exceptions.txt @@ -417,6 +417,6 @@ task-sdk/src/airflow/sdk/definitions/_internal/setup_teardown.py::1 task-sdk/src/airflow/sdk/definitions/connection.py::4 task-sdk/src/airflow/sdk/definitions/decorators/setup_teardown.py::4 task-sdk/src/airflow/sdk/definitions/xcom_arg.py::1 -task-sdk/src/airflow/sdk/execution_time/task_runner.py::1 +task-sdk/src/airflow/sdk/execution_time/bundles.py::1 task-sdk/tests/task_sdk/bases/test_sensor.py::1 task-sdk/tests/task_sdk/execution_time/test_task_runner.py::2 diff --git a/task-sdk/src/airflow/sdk/coordinators/_bundle_metadata.py b/task-sdk/src/airflow/sdk/coordinators/_bundle_metadata.py index e95ffe52396f0..250986175c009 100644 --- a/task-sdk/src/airflow/sdk/coordinators/_bundle_metadata.py +++ b/task-sdk/src/airflow/sdk/coordinators/_bundle_metadata.py @@ -21,7 +21,7 @@ import os import pathlib -from typing import Any +from typing import Any, Final import attrs import yaml @@ -29,6 +29,13 @@ from airflow.sdk.execution_time.schema import get_schema_version_migrator +class _UnsetArtifactRoots: + """Sentinel distinguishing an omitted artifact-roots option from an empty value.""" + + +ARTIFACT_ROOTS_NOT_CONFIGURED: Final = _UnsetArtifactRoots() + + def convert_roots( value: None | os.PathLike[str] | pathlib.Path | list[os.PathLike[str] | pathlib.Path], ) -> list[pathlib.Path]: @@ -40,6 +47,25 @@ def convert_roots( return [pathlib.Path(v).expanduser() for v in value] +def convert_configured_roots( + value: _UnsetArtifactRoots + | None + | os.PathLike[str] + | pathlib.Path + | list[os.PathLike[str] | pathlib.Path], +) -> list[pathlib.Path]: + """Normalize configured roots while rejecting explicitly empty values.""" + if isinstance(value, _UnsetArtifactRoots): + return [] + roots = convert_roots(value) + if not roots: + raise ValueError( + "Artifact roots must contain at least one path when provided; " + "omit the option to use the task's Dag bundle." + ) + return roots + + def validate_schema_version(instance, _, value) -> str: """Attrs validator resolving a bundle's supervisor schema version to a known one.""" return get_schema_version_migrator().resolve_version(str(value)) diff --git a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py index a0c3f518fb068..9e69255ba3d07 100644 --- a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py +++ b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py @@ -27,6 +27,8 @@ from __future__ import annotations +import contextlib +import enum import ipaddress import itertools import os @@ -41,18 +43,24 @@ import psutil import structlog +from airflow.dag_processing.bundles.base import BundleVersionLock, unpack_bundle_version # noqa: SDK002 +from airflow.dag_processing.bundles.manager import DagBundlesManager # noqa: SDK002 +from airflow.sdk.api.datamodels._generated import BundleInfo from airflow.sdk.configuration import conf +from airflow.sdk.execution_time.bundles import initialize_ti_bundle from airflow.sdk.execution_time.coordinator import BaseCoordinator from airflow.sdk.execution_time.supervisor import ActivitySubprocess, NeverRaised, ProcessTracker if TYPE_CHECKING: + import pathlib from collections.abc import Sequence from structlog.typing import FilteringBoundLogger from typing_extensions import Self + from airflow.dag_processing.bundles.base import BaseDagBundle # noqa: SDK002 from airflow.sdk.api.client import Client - from airflow.sdk.api.datamodels._generated import BundleInfo, TaskInstance + from airflow.sdk.api.datamodels._generated import TaskInstance Tracked = TypeVar("Tracked", socket.socket, subprocess.Popen) @@ -371,6 +379,45 @@ def wait(self) -> int: return code +def _initialize_pinned_bundle(target: BundleInfo, logger: FilteringBoundLogger) -> BaseDagBundle: + """ + Materialize *target* at a concrete version, so the tree handed to the subprocess is lockable. + + A bundle resolved without a version points at the bundle's shared, mutable + checkout: another task refreshing the same bundle resets it underneath a running + subprocess, and ``BundleVersionLock`` cannot protect it because a version-less + lock is a no-op. Re-resolving at the version current now yields a private + ``versions/`` tree that the lock does cover. + + Bundles that do not track versions have nothing to pin and keep their single path. + """ + bundle = initialize_ti_bundle(target) + if bundle.version is not None: + return bundle + + version, version_data = unpack_bundle_version(bundle.get_current_version(), bundle) + if version is None: + return bundle + logger.debug("Pinning Dag bundle to its current version", bundle=target.name, version=version) + return initialize_ti_bundle(BundleInfo(name=target.name, version=version, version_data=version_data)) + + +class _ArtifactSource(enum.Enum): + """How a subprocess coordinator locates the compiled task artifacts.""" + + EXPLICIT_ROOT = enum.auto() + """An explicit filesystem root (``jars_root`` / ``executables_root`` / ``bundles_root``).""" + NAMED_BUNDLE = enum.auto() + """``dag_bundle_name`` names a configured Dag bundle; its version current at task start is used.""" + TASK_BUNDLE = enum.auto() + """ + Neither is set: artifacts are *co-located* with the Python stub Dag. + + The task's own bundle is scanned, so the compiled artifacts ship in the same + bundle, at the same version, as the Dag that delegates to them. + """ + + @attrs.define(kw_only=True) class SubprocessCoordinator(BaseCoordinator): """ @@ -385,23 +432,129 @@ class SubprocessCoordinator(BaseCoordinator): :param task_startup_timeout: Maximum time the coordinator waits for the subprocess to connect to both servers, in seconds. The default is 10 seconds. + :param dag_bundle_name: Locate artifacts through a configured Dag bundle rather + than an explicit root. Mutually exclusive with the subclass's explicit root; + if neither is set, the task's own 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. """ task_startup_timeout: float = 10.0 + dag_bundle_name: str | None = None + + _artifact_source: _ArtifactSource = attrs.field(init=False) + # The subclass's explicit root, recorded at construction so the base can + # resolve roots without knowing the subclass field name. + _configured_roots: list[pathlib.Path] = attrs.field(init=False, factory=list) + _active_scan_roots: tuple[pathlib.Path, ...] | None = attrs.field(init=False, default=None) + + @property + def _explicit_artifact_roots(self) -> tuple[str, Sequence[pathlib.Path]]: + """ + The subclass's explicit-root kwarg name and its configured value. + + The name is only used in error messages. An empty value — the default, for a + subclass that does not override this — selects task-bundle mode rather than + failing at execute time. + """ + return "root", () + + def __attrs_post_init__(self) -> None: + self._classify_artifact_source() + + def _classify_artifact_source(self) -> None: + """ + Classify and validate how this coordinator locates artifacts (construction time). + + Rejects setting both an explicit root and ``dag_bundle_name``, fails fast + when ``dag_bundle_name`` names a bundle that is not configured, and records + the resulting :class:`_ArtifactSource` and explicit root. + """ + root_kwarg, configured = self._explicit_artifact_roots + if configured and self.dag_bundle_name is not None: + raise ValueError( + f"Set at most one of {root_kwarg!r} or 'dag_bundle_name': {root_kwarg!r} for an " + f"explicit path, 'dag_bundle_name' for a configured Dag bundle, or leave both " + f"unset to scan the task's own bundle." + ) + if configured: + source = _ArtifactSource.EXPLICIT_ROOT + self._configured_roots = list(configured) + elif self.dag_bundle_name is not None: + source = _ArtifactSource.NAMED_BUNDLE + if not DagBundlesManager.is_bundle_configured(self.dag_bundle_name): + raise ValueError( + f"Coordinator 'dag_bundle_name' references unconfigured Dag bundle " + f"{self.dag_bundle_name!r}." + ) + else: + source = _ArtifactSource.TASK_BUNDLE + + self._artifact_source = source + log.debug( + "Coordinator artifact source selected", + mode=source.name, + dag_bundle_name=self.dag_bundle_name, + configured_roots=[str(root) for root in self._configured_roots], + ) + + def _init_root_source( + self, bundle_info: BundleInfo, logger: FilteringBoundLogger + ) -> tuple[list[pathlib.Path], BaseDagBundle | None]: + """ + Resolve the directories to scan for artifacts, dispatched on the classified mode. + + Returns ``(roots, bundle)``: an explicit root yields no bundle (``None``); + a Dag-bundle mode returns the materialized path and the resolved bundle so + :meth:`execute_task` can hold a version lock over it. *logger* is the task + logger, so materialization failures surface in the task log. + """ + if self._artifact_source is _ArtifactSource.EXPLICIT_ROOT: + return self._configured_roots, None + + if self._artifact_source is _ArtifactSource.NAMED_BUNDLE: + # NAMED_BUNDLE implies dag_bundle_name is set. + target = BundleInfo(name=cast("str", self.dag_bundle_name)) + else: + target = bundle_info + + bundle = _initialize_pinned_bundle(target, logger) + path = bundle.path + if not path.exists(): + raise FileNotFoundError(f"Dag bundle {target.name!r} resolved to {path}, which does not exist.") + return [path], bundle + + def _get_scan_roots(self) -> tuple[pathlib.Path, ...]: + """Return the artifact roots resolved for the active task.""" + if self._active_scan_roots is None: + raise RuntimeError("_get_scan_roots requires an active task; call it during execute_task.") + return self._active_scan_roots def _build_execute_task_command(self, *, what: TaskInstance) -> tuple[list[str], str | None]: """ Build the subprocess command and resolve its supervisor wire-schema version for *what*. - Returns a ``(command, subprocess_schema_version)`` pair. *command* - MUST NOT include the ``--comm`` / ``--logs`` flags — those are - appended by :class:`_PopenActivitySubprocess` once the listening - sockets have been bound. A ``None`` schema version disables schema - migration; messages are then exchanged at the runtime's native wire - format. + Subclasses can retrieve the directories to scan for artifacts with + :meth:`_get_scan_roots`. + Returns a ``(command, subprocess_schema_version)`` pair. *command* MUST + NOT include the ``--comm`` / ``--logs`` flags — those are appended by + :class:`_PopenActivitySubprocess` once the listening sockets have been + bound. A ``None`` schema version disables schema migration; messages are + then exchanged at the runtime's native wire format. """ raise NotImplementedError + @contextlib.contextmanager + def _set_scan_roots(self, roots: Sequence[pathlib.Path]): + """Expose *roots* to the command builder for the duration of the task.""" + if self._active_scan_roots is not None: + raise RuntimeError("SubprocessCoordinator.execute_task is not re-entrant.") + self._active_scan_roots = tuple(roots) + try: + yield + finally: + self._active_scan_roots = None + def execute_task( self, *, @@ -414,18 +567,32 @@ def execute_task( subprocess_logs_to_stdout: bool, **kwargs, ) -> BaseCoordinator.ExecutionResult: - command, subprocess_schema_version = self._build_execute_task_command(what=what) - process = _PopenActivitySubprocess.start( - what=what, - dag_rel_path=dag_rel_path, - bundle_info=bundle_info, - client=client, - logger=logger, - subprocess_logs_to_stdout=subprocess_logs_to_stdout, - sentry_integration=sentry_integration, - command=command, - subprocess_schema_version=subprocess_schema_version, - startup_timeout=self.task_startup_timeout, - ) - exit_code = process.wait() - return self.ExecutionResult(exit_code, process.final_state) + task_logger = logger or log + with contextlib.ExitStack() as stack: + roots, resolved_bundle = self._init_root_source(bundle_info, task_logger) + if resolved_bundle is not None: + # 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. + stack.enter_context( + BundleVersionLock( + bundle_name=resolved_bundle.name, + bundle_version=resolved_bundle.version, + ) + ) + stack.enter_context(self._set_scan_roots(roots)) + command, subprocess_schema_version = self._build_execute_task_command(what=what) + process = _PopenActivitySubprocess.start( + what=what, + dag_rel_path=dag_rel_path, + bundle_info=bundle_info, + client=client, + logger=logger, + subprocess_logs_to_stdout=subprocess_logs_to_stdout, + sentry_integration=sentry_integration, + command=command, + subprocess_schema_version=subprocess_schema_version, + startup_timeout=self.task_startup_timeout, + ) + exit_code = process.wait() + return self.ExecutionResult(exit_code, process.final_state) diff --git a/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py b/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py index 70ba5490450c3..9d6b80b249f81 100644 --- a/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py +++ b/task-sdk/src/airflow/sdk/coordinators/executable/coordinator.py @@ -31,8 +31,9 @@ import structlog from airflow.sdk.coordinators._bundle_metadata import ( + ARTIFACT_ROOTS_NOT_CONFIGURED, ResolvedBundle, - convert_roots, + convert_configured_roots, extract_supervisor_schema_version, parse_metadata_mapping, ) @@ -339,16 +340,22 @@ class ExecutableCoordinator(SubprocessCoordinator): :param executables_root: A list of directories scanned for executable bundles when a Python stub DAG delegates task execution to a native - runtime. + runtime. See :class:`SubprocessCoordinator` for its interaction with + ``dag_bundle_name``. :param task_startup_timeout: Maximum time the coordinator waits for a task process to start, in seconds. The default is 10 seconds. """ executables_root: list[pathlib.Path] = attrs.field( - converter=convert_roots, - validator=attrs.validators.min_len(1), + default=ARTIFACT_ROOTS_NOT_CONFIGURED, + converter=convert_configured_roots, ) + @property + def _explicit_artifact_roots(self) -> tuple[str, list[pathlib.Path]]: + return "executables_root", self.executables_root + def _build_execute_task_command(self, *, what: TaskInstance) -> tuple[list[str], str | None]: - bundle = _Bundle.find(self.executables_root, what.dag_id) + roots = self._get_scan_roots() + bundle = _Bundle.find(roots, what.dag_id) return [str(bundle.path)], bundle.schema_version diff --git a/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py b/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py index 0ebe089162ba1..1f0f6e6efd5f9 100644 --- a/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py +++ b/task-sdk/src/airflow/sdk/coordinators/java/coordinator.py @@ -29,7 +29,11 @@ import attrs import structlog -from airflow.sdk.coordinators._bundle_metadata import convert_roots, validate_schema_version +from airflow.sdk.coordinators._bundle_metadata import ( + ARTIFACT_ROOTS_NOT_CONFIGURED, + convert_configured_roots, + validate_schema_version, +) from airflow.sdk.coordinators._subprocess import SubprocessCoordinator if TYPE_CHECKING: @@ -172,7 +176,8 @@ class JavaCoordinator(SubprocessCoordinator): :param java_executable: Path to the ``java`` command (defaults to ``"java"``, which relies on ``$PATH``). :param jvm_args: Extra arguments passed to the JVM (e.g. ``["-Xmx512m"]``). - :param jars_root: A list of directories scanned for JAR bundles. + :param jars_root: A list of directories scanned for JAR bundles. See + :class:`SubprocessCoordinator` for its interaction with ``dag_bundle_name``. :param main_class: Explicit entry point to execute with *java_executable*. :param task_startup_timeout: Maximum time the coordinator waits for a task process to start, in seconds. The default is 10 seconds. @@ -199,17 +204,24 @@ class JavaCoordinator(SubprocessCoordinator): java_executable: str = "java" jvm_args: list[str] = attrs.field(factory=list) jars_root: list[pathlib.Path] = attrs.field( - converter=convert_roots, - validator=attrs.validators.min_len(1), + default=ARTIFACT_ROOTS_NOT_CONFIGURED, + converter=convert_configured_roots, ) main_class: str = "" + @property + def _explicit_artifact_roots(self) -> tuple[str, list[pathlib.Path]]: + return "jars_root", self.jars_root + def _build_execute_task_command(self, *, what: TaskInstance) -> tuple[list[str], str | None]: - jar = _JarInfo.find(self.jars_root, self.main_class) + # Without main_class, the first executable JAR in walk order wins; tracked at + # https://github.com/apache/airflow/issues/71134 + roots = self._get_scan_roots() + jar = _JarInfo.find(roots, self.main_class) command = [ self.java_executable, "-classpath", - _calculate_classpath(self.jars_root), + _calculate_classpath(roots), *self.jvm_args, jar.main_class, ] diff --git a/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py b/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py index 38e5224336c92..63897e7aea417 100644 --- a/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py +++ b/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py @@ -28,8 +28,9 @@ import structlog from airflow.sdk.coordinators._bundle_metadata import ( + ARTIFACT_ROOTS_NOT_CONFIGURED, ResolvedBundle, - convert_roots, + convert_configured_roots, extract_supervisor_schema_version, parse_metadata_mapping, ) @@ -145,19 +146,25 @@ class NodeCoordinator(SubprocessCoordinator): TypeScript bundle. Each bundle directory must contain ``bundle.mjs`` with embedded metadata (as produced by ``airflow-ts-pack``). This is a fallback search path; it does not yet route different Dag/task pairs - to different bundles. + to different bundles. See :class:`SubprocessCoordinator` for its + interaction with ``dag_bundle_name``. :param task_startup_timeout: Maximum time the coordinator waits for a task process to start, in seconds. The default is 10 seconds. """ node_executable: str = "node" bundles_root: list[pathlib.Path] = attrs.field( - converter=convert_roots, - validator=attrs.validators.min_len(1), + default=ARTIFACT_ROOTS_NOT_CONFIGURED, + converter=convert_configured_roots, ) + @property + def _explicit_artifact_roots(self) -> tuple[str, list[pathlib.Path]]: + return "bundles_root", self.bundles_root + def _build_execute_task_command(self, *, what: TaskInstance) -> tuple[list[str], str | None]: # Multi-bundle routing should be added here by passing `what.dag_id` and # `what.task_id` into bundle selection and matching against metadata["dags"]. - bundle = _find_bundle(self.bundles_root) + roots = self._get_scan_roots() + bundle = _find_bundle(roots) return [self.node_executable, os.fspath(bundle.path)], bundle.schema_version diff --git a/task-sdk/src/airflow/sdk/execution_time/bundles.py b/task-sdk/src/airflow/sdk/execution_time/bundles.py new file mode 100644 index 0000000000000..8ab131beb64f3 --- /dev/null +++ b/task-sdk/src/airflow/sdk/execution_time/bundles.py @@ -0,0 +1,90 @@ +# 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. +""" +Materialize a task instance's Dag bundle on the worker. + +Kept apart from :mod:`airflow.sdk.execution_time.task_runner` so the supervisor +side — which needs a bundle on disk before it can launch a language-SDK +subprocess — does not have to import the task runner and everything it pulls in. +""" + +from __future__ import annotations + +import os +from getpass import getuser +from typing import TYPE_CHECKING + +from airflow.dag_processing.bundles.manager import DagBundlesManager # noqa: SDK002 +from airflow.sdk.exceptions import AirflowException +from airflow.sdk.execution_time.tracing import detail_span + +if TYPE_CHECKING: + from airflow.dag_processing.bundles.base import BaseDagBundle # noqa: SDK002 + from airflow.sdk.api.datamodels._generated import BundleInfo + +__all__ = ["initialize_ti_bundle", "verify_bundle_access"] + + +def initialize_ti_bundle(bundle_info: BundleInfo) -> BaseDagBundle: + """ + Resolve, initialize, and access-check the Dag bundle for a task instance. + + Shared by :func:`~airflow.sdk.execution_time.task_runner.parse` (Python task + path) and the subprocess coordinators (language-SDK path), which both need a + task instance's bundle materialized on disk before use. Returns the + initialized bundle so callers can read ``bundle.path``. + """ + bundle_instance = DagBundlesManager().get_bundle( + name=bundle_info.name, + version=bundle_info.version, + version_data=bundle_info.version_data, + ) + bundle_instance.initialize() + verify_bundle_access(bundle_instance) + return bundle_instance + + +@detail_span("verify_bundle_access") +def verify_bundle_access(bundle_instance: BaseDagBundle) -> None: + """ + Verify bundle is accessible by the current user. + + This is called after user impersonation (if any) to ensure the bundle + is actually accessible. Uses os.access() which works with any permission + scheme (standard Unix permissions, ACLs, SELinux, etc.). + + :param bundle_instance: The bundle instance to check + :raises AirflowException: if bundle is not accessible + """ + bundle_path = bundle_instance.path + + if not bundle_path.exists(): + # Already handled by initialize() with a warning + return + + # Check read permission (and execute for directories to list contents) + access_mode = os.R_OK + if bundle_path.is_dir(): + access_mode |= os.X_OK + + if not os.access(bundle_path, access_mode): + raise AirflowException( + f"Bundle '{bundle_instance.name}' path '{bundle_path}' is not accessible " + f"by user '{getuser()}'. When using run_as_user, ensure bundle directories " + f"are readable by the impersonated user. " + f"See: https://airflow.apache.org/docs/apache-airflow/stable/administration-and-deployment/dag-bundles.html" + ) diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py b/task-sdk/src/airflow/sdk/execution_time/task_runner.py index eeb2788062248..b9e5f68df3cea 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -43,10 +43,8 @@ from structlog.contextvars import bind_contextvars from airflow.dag_processing.bundles.base import BaseDagBundle, BundleVersionLock -from airflow.dag_processing.bundles.manager import DagBundlesManager from airflow.sdk._shared.observability.metrics import stats from airflow.sdk._shared.observability.metrics.stats import build_dag_metric_tags -from airflow.sdk._shared.observability.traces import get_task_span_detail_level from airflow.sdk._shared.template_rendering import truncate_rendered_value from airflow.sdk.api.client import get_hostname, getuser from airflow.sdk.api.datamodels._generated import ( @@ -84,6 +82,7 @@ TaskAwaitingInput, TaskDeferred, ) +from airflow.sdk.execution_time.bundles import initialize_ti_bundle from airflow.sdk.execution_time.callback_runner import create_executable_runner from airflow.sdk.execution_time.comms import ( AssetEventDagRunReferenceResult, @@ -146,6 +145,7 @@ _LegacyEmailBackendNotifier, ) from airflow.sdk.execution_time.sentry import Sentry +from airflow.sdk.execution_time.tracing import detail_span from airflow.sdk.execution_time.xcom import XCom from airflow.sdk.listener import get_listener_manager from airflow.sdk.observability.metrics import stats_utils @@ -169,38 +169,6 @@ tracer = trace.get_tracer(__name__) -class detail_span: - """Context manager and decorator that creates a child span when detail level > 1.""" - - def __init__(self, *args, **kwargs): - self._args = args - self._kwargs = kwargs - self._ctx = None - - def _make_ctx(self): - parent_span = trace.get_current_span() - config_level = get_task_span_detail_level(span=parent_span) - if config_level > 1: - return tracer.start_as_current_span(*self._args, **self._kwargs) - return trace.INVALID_SPAN - - def __enter__(self): - self._ctx = self._make_ctx() - return self._ctx.__enter__() - - def __exit__(self, *exc_info): - return self._ctx.__exit__(*exc_info) - - def __call__(self, f): - @functools.wraps(f) - def wrapper(*inner_args, **inner_kwargs): - with self._make_ctx(): - return f(*inner_args, **inner_kwargs) - - wrapper.__signature__ = inspect.signature(f) - return wrapper - - @contextmanager def _make_task_span(msg: StartupDetails): parent_context = ( @@ -1020,13 +988,7 @@ def parse(what: StartupDetails, log: Logger) -> RuntimeTaskInstance: bundle_info = what.bundle_info bundle_prepare_start = time.monotonic() - bundle_instance = DagBundlesManager().get_bundle( - name=bundle_info.name, - version=bundle_info.version, - version_data=bundle_info.version_data, - ) - bundle_instance.initialize() - _verify_bundle_access(bundle_instance, log) + bundle_instance = initialize_ti_bundle(bundle_info) bundle_prepare_ms = int((time.monotonic() - bundle_prepare_start) * 1000) dag_absolute_path = os.fspath(Path(bundle_instance.path, what.dag_rel_path)) @@ -1122,43 +1084,6 @@ def parse(what: StartupDetails, log: Logger) -> RuntimeTaskInstance: # 3. Shutdown and report status -@detail_span("_verify_bundle_access") -def _verify_bundle_access(bundle_instance: BaseDagBundle, log: Logger) -> None: - """ - Verify bundle is accessible by the current user. - - This is called after user impersonation (if any) to ensure the bundle - is actually accessible. Uses os.access() which works with any permission - scheme (standard Unix permissions, ACLs, SELinux, etc.). - - :param bundle_instance: The bundle instance to check - :param log: Logger instance - :raises AirflowException: if bundle is not accessible - """ - from getpass import getuser - - from airflow.sdk.exceptions import AirflowException - - bundle_path = bundle_instance.path - - if not bundle_path.exists(): - # Already handled by initialize() with a warning - return - - # Check read permission (and execute for directories to list contents) - access_mode = os.R_OK - if bundle_path.is_dir(): - access_mode |= os.X_OK - - if not os.access(bundle_path, access_mode): - raise AirflowException( - f"Bundle '{bundle_instance.name}' path '{bundle_path}' is not accessible " - f"by user '{getuser()}'. When using run_as_user, ensure bundle directories " - f"are readable by the impersonated user. " - f"See: https://airflow.apache.org/docs/apache-airflow/stable/administration-and-deployment/dag-bundles.html" - ) - - def get_startup_details() -> StartupDetails: # The parent sends us a StartupDetails message un-prompted. After this, every single message is only sent # in response to us sending a request. @@ -1266,8 +1191,6 @@ def _serialize_template_field( Uses the SDK secrets masker to redact secrets in the serialized output. """ - import inspect - from airflow.sdk._shared.module_loading import qualname from airflow.sdk._shared.secrets_masker import redact diff --git a/task-sdk/src/airflow/sdk/execution_time/tracing.py b/task-sdk/src/airflow/sdk/execution_time/tracing.py new file mode 100644 index 0000000000000..86a28199c61da --- /dev/null +++ b/task-sdk/src/airflow/sdk/execution_time/tracing.py @@ -0,0 +1,60 @@ +# 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. +"""Span helpers shared by the modules that make up a task run.""" + +from __future__ import annotations + +import functools +import inspect + +from opentelemetry import trace + +from airflow.sdk._shared.observability.traces import get_task_span_detail_level + +tracer = trace.get_tracer(__name__) + + +class detail_span: + """Context manager and decorator that creates a child span when detail level > 1.""" + + def __init__(self, *args, **kwargs): + self._args = args + self._kwargs = kwargs + self._ctx = None + + def _make_ctx(self): + parent_span = trace.get_current_span() + config_level = get_task_span_detail_level(span=parent_span) + if config_level > 1: + return tracer.start_as_current_span(*self._args, **self._kwargs) + return trace.INVALID_SPAN + + def __enter__(self): + self._ctx = self._make_ctx() + return self._ctx.__enter__() + + def __exit__(self, *exc_info): + return self._ctx.__exit__(*exc_info) + + def __call__(self, f): + @functools.wraps(f) + def wrapper(*inner_args, **inner_kwargs): + with self._make_ctx(): + return f(*inner_args, **inner_kwargs) + + wrapper.__signature__ = inspect.signature(f) + return wrapper diff --git a/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py b/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py index 5963e5f4892a7..2fb1ea96d097d 100644 --- a/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/executable/test_coordinator.py @@ -370,13 +370,42 @@ def test_executables_root_accepts_list(self, tmp_path): coordinator = ExecutableCoordinator(executables_root=[str(tmp_path), other]) assert coordinator.executables_root == [tmp_path, other] - def test_executables_root_required(self): - with pytest.raises(TypeError, match="executables_root"): - ExecutableCoordinator() - - def test_executables_root_must_be_non_empty(self): - with pytest.raises(ValueError, match="executables_root"): - ExecutableCoordinator(executables_root=None) + def test_executables_root_optional_defaults_to_empty(self): + coordinator = ExecutableCoordinator() + assert coordinator.executables_root == [] + assert coordinator.dag_bundle_name is None + + @pytest.mark.parametrize("executables_root", [None, []], ids=["none", "empty-list"]) + def test_explicit_empty_executables_root_raises(self, executables_root): + with pytest.raises(ValueError, match="must contain at least one path when provided"): + ExecutableCoordinator(executables_root=executables_root) + + def test_root_and_dag_bundle_name_are_mutually_exclusive(self, tmp_path): + with pytest.raises(ValueError, match="at most one of 'executables_root' or 'dag_bundle_name'"): + ExecutableCoordinator(executables_root=[tmp_path], dag_bundle_name="artifacts") + + @patch("airflow.sdk.coordinators._subprocess.DagBundlesManager") + def test_unconfigured_dag_bundle_name_raises(self, mock_manager): + mock_manager.is_bundle_configured.return_value = False + with pytest.raises(ValueError, match="unconfigured Dag bundle 'ghost'"): + ExecutableCoordinator(dag_bundle_name="ghost") + + @patch("airflow.sdk.coordinators._subprocess.DagBundlesManager") + def test_configured_dag_bundle_name_accepted(self, mock_manager): + mock_manager.is_bundle_configured.return_value = True + coordinator = ExecutableCoordinator(dag_bundle_name="artifacts") + assert coordinator.dag_bundle_name == "artifacts" + mock_manager.is_bundle_configured.assert_called_once_with("artifacts") + + def test_build_command_scans_passed_roots_in_colocated_mode(self, tmp_path): + binary = _build_bundle(tmp_path / "my_bundle", dag_ids=["tutorial_dag"]) + coordinator = ExecutableCoordinator() + with coordinator._set_scan_roots([tmp_path]): + command, schema_version = coordinator._build_execute_task_command( + what=_make_ti(dag_id="tutorial_dag") + ) + assert command == [str(binary.resolve())] + assert schema_version == "2026-06-16" class TestBuildExecuteTaskCommand: @@ -385,7 +414,8 @@ def test_returns_resolved_executable_and_schema_version(self, tmp_path): ti = _make_ti(dag_id="tutorial_dag") coordinator = ExecutableCoordinator(executables_root=[tmp_path]) - command, schema_version = coordinator._build_execute_task_command(what=ti) + with coordinator._set_scan_roots([tmp_path]): + command, schema_version = coordinator._build_execute_task_command(what=ti) assert command == [str(binary.resolve())] assert schema_version == "2026-06-16" @@ -396,7 +426,10 @@ def test_raises_when_bundle_omits_schema_version(self, tmp_path): ti = _make_ti(dag_id="tutorial_dag") coordinator = ExecutableCoordinator(executables_root=[tmp_path]) - with pytest.raises(FileNotFoundError, match="matching bundles were rejected"): + with ( + coordinator._set_scan_roots([tmp_path]), + pytest.raises(FileNotFoundError, match="matching bundles were rejected"), + ): coordinator._build_execute_task_command(what=ti) def test_raises_when_dag_id_not_found(self, tmp_path): @@ -404,7 +437,10 @@ def test_raises_when_dag_id_not_found(self, tmp_path): ti = _make_ti(dag_id="tutorial_dag") coordinator = ExecutableCoordinator(executables_root=[tmp_path]) - with pytest.raises(FileNotFoundError, match="cannot find executable bundle"): + with ( + coordinator._set_scan_roots([tmp_path]), + pytest.raises(FileNotFoundError, match="cannot find executable bundle"), + ): coordinator._build_execute_task_command(what=ti) 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 c30ecff35b9d2..5d4dab326c5fe 100644 --- a/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/java/test_coordinator.py @@ -237,6 +237,41 @@ def test_custom_kwargs(self): assert coordinator.jvm_args == ["-Xmx512m", "-Xms256m"] assert coordinator.jars_root == [pathlib.Path("/airflow/java-bundles")] + def test_jars_root_optional_defaults_to_empty(self): + # main_class stays optional: the entrypoint is auto-detected from the bundle scan. + coordinator = JavaCoordinator() + assert coordinator.jars_root == [] + assert coordinator.dag_bundle_name is None + assert coordinator.main_class == "" + + @pytest.mark.parametrize("jars_root", [None, []], ids=["none", "empty-list"]) + def test_explicit_empty_jars_root_raises(self, jars_root): + with pytest.raises(ValueError, match="must contain at least one path when provided"): + JavaCoordinator(jars_root=jars_root) + + def test_explicit_root_does_not_require_main_class(self, tmp_path): + coordinator = JavaCoordinator(jars_root=tmp_path) + assert coordinator.main_class == "" + + def test_root_and_dag_bundle_name_are_mutually_exclusive(self): + with pytest.raises(ValueError, match="at most one of 'jars_root' or 'dag_bundle_name'"): + JavaCoordinator(jars_root="/airflow/java-bundles", dag_bundle_name="artifacts") + + @patch("airflow.sdk.coordinators._subprocess.DagBundlesManager") + def test_unconfigured_dag_bundle_name_raises(self, mock_manager): + mock_manager.is_bundle_configured.return_value = False + with pytest.raises(ValueError, match="unconfigured Dag bundle 'ghost'"): + JavaCoordinator(dag_bundle_name="ghost") + + def test_build_command_scans_passed_roots_in_colocated_mode(self, tmp_path): + _make_jar(tmp_path / "app.jar", main_class="com.example.TaskRunner", schema_version="2026-06-16") + coordinator = JavaCoordinator(main_class="com.example.TaskRunner") + with coordinator._set_scan_roots([tmp_path]): + command, schema_version = coordinator._build_execute_task_command(what=_make_ti()) + assert command[0] == "java" + assert command[-1] == "com.example.TaskRunner" + assert schema_version == "2026-06-16" + @pytest.fixture def jars_root(tmp_path): diff --git a/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py b/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py index 6efe45566a6dd..dc2ce1b912a6b 100644 --- a/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py @@ -20,6 +20,7 @@ import base64 import pathlib +from unittest.mock import patch import pytest from uuid6 import uuid7 @@ -94,9 +95,33 @@ def test_custom_kwargs(self): ] assert coordinator.task_startup_timeout == 30.0 - def test_bundles_root_is_required(self): - with pytest.raises(ValueError, match="Length of 'bundles_root' must be >= 1"): - NodeCoordinator(bundles_root=None) + def test_bundles_root_optional_defaults_to_empty(self): + coordinator = NodeCoordinator() + assert coordinator.bundles_root == [] + assert coordinator.dag_bundle_name is None + + @pytest.mark.parametrize("bundles_root", [None, []], ids=["none", "empty-list"]) + def test_explicit_empty_bundles_root_raises(self, bundles_root): + with pytest.raises(ValueError, match="must contain at least one path when provided"): + NodeCoordinator(bundles_root=bundles_root) + + def test_root_and_dag_bundle_name_are_mutually_exclusive(self): + with pytest.raises(ValueError, match="at most one of 'bundles_root' or 'dag_bundle_name'"): + NodeCoordinator(bundles_root="/airflow/ts-bundles", dag_bundle_name="artifacts") + + @patch("airflow.sdk.coordinators._subprocess.DagBundlesManager") + def test_unconfigured_dag_bundle_name_raises(self, mock_manager): + mock_manager.is_bundle_configured.return_value = False + with pytest.raises(ValueError, match="unconfigured Dag bundle 'ghost'"): + NodeCoordinator(dag_bundle_name="ghost") + + def test_build_command_scans_passed_roots_in_colocated_mode(self, tmp_path): + bundle = write_bundle(tmp_path) + coordinator = NodeCoordinator() + with coordinator._set_scan_roots([tmp_path]): + command, schema_version = coordinator._build_execute_task_command(what=_make_ti()) + assert command == ["node", str(bundle)] + assert schema_version == SCHEMA_VERSION class TestNodeCoordinatorBundleSelection: @@ -217,7 +242,8 @@ def test_build_execute_task_command_returns_node_bundle_and_schema_version(self, bundles_root=tmp_path, ) - command, schema_version = coordinator._build_execute_task_command(what=_make_ti()) + with coordinator._set_scan_roots([tmp_path]): + command, schema_version = coordinator._build_execute_task_command(what=_make_ti()) assert command == ["/opt/node/bin/node", str(bundle)] assert schema_version == SCHEMA_VERSION diff --git a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py index 62b7fbcf39c17..5297f56dead4f 100644 --- a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py +++ b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py @@ -19,6 +19,7 @@ import contextlib import os +import pathlib import socket import subprocess import sys @@ -31,16 +32,19 @@ import pytest from uuid6 import uuid7 +from airflow.dag_processing.bundles.base import BundleVersion from airflow.sdk.api.client import Client, TaskInstanceOperations -from airflow.sdk.api.datamodels._generated import TaskInstance +from airflow.sdk.api.datamodels._generated import BundleInfo, TaskInstance from airflow.sdk.coordinators._subprocess import ( SubprocessCoordinator, _accept_connections, + _ArtifactSource, _connection_owned_by_process_tree, _is_connection_from_process, _PopenActivitySubprocess, _ResourceTracker, _start_server, + log, ) from airflow.sdk.execution_time.coordinator import BaseCoordinator from airflow.sdk.execution_time.supervisor import ActivitySubprocess @@ -567,12 +571,25 @@ def test_untrack_unknown_object_does_not_raise(self): @attrs.define(kw_only=True) class _StubSubprocessCoordinator(SubprocessCoordinator): - """Minimal SubprocessCoordinator subclass used to exercise the base machinery.""" + """Minimal SubprocessCoordinator subclass used to exercise the base machinery. + + ``explicit_roots`` defaults to a real path so the coordinator classifies as + EXPLICIT_ROOT and execute_task resolves without touching a Dag bundle; pass + ``explicit_roots=[]`` to exercise TASK_BUNDLE mode. Roots handed to the + command builder are recorded in ``recorded_roots`` so wiring can be asserted. + """ command: list[str] schema_version: str | None = None + explicit_roots: list[pathlib.Path] = attrs.field(factory=lambda: [pathlib.Path(".")]) + recorded_roots: list[list[pathlib.Path]] = attrs.field(init=False, factory=list) + + @property + def _explicit_artifact_roots(self) -> tuple[str, list[pathlib.Path]]: + return "explicit_roots", self.explicit_roots def _build_execute_task_command(self, *, what): + self.recorded_roots.append(list(self._get_scan_roots())) return list(self.command), self.schema_version @@ -844,3 +861,246 @@ def test_register_pipe_readers_called_with_four_sockets(self, mock_client): subprocess_logs_to_stdout=False, ) assert mock_register.mock_calls == [call(ANY, ANY, ANY, ANY, data=ANY)] + + +class TestClassifyArtifactSource: + """Construction-time classification and per-mode validation. + + Classification runs from the base ``__attrs_post_init__``, so it is exercised + through construction rather than by calling the helper directly. + """ + + def test_explicit_root_classified_and_recorded(self): + configured = [pathlib.Path("/artifacts")] + coordinator = _StubSubprocessCoordinator(command=["x"], explicit_roots=configured) + assert coordinator._artifact_source is _ArtifactSource.EXPLICIT_ROOT + assert coordinator._configured_roots == configured + + def test_task_bundle_classified_when_neither_set(self): + coordinator = _StubSubprocessCoordinator(command=["x"], explicit_roots=[]) + assert coordinator._artifact_source is _ArtifactSource.TASK_BUNDLE + + def test_subclass_without_overrides_defaults_to_task_bundle(self): + """A subclass that overrides neither hook is classified, not left to error at execute time.""" + + @attrs.define(kw_only=True) + class _Bare(SubprocessCoordinator): + def _build_execute_task_command(self, *, what): + return [], None + + assert _Bare()._artifact_source is _ArtifactSource.TASK_BUNDLE + + def test_rejects_explicit_root_and_dag_bundle_name_together(self): + with pytest.raises(ValueError, match="at most one of 'explicit_roots' or 'dag_bundle_name'"): + _StubSubprocessCoordinator( + command=["x"], explicit_roots=[pathlib.Path("/artifacts")], dag_bundle_name="artifacts" + ) + + @patch("airflow.sdk.coordinators._subprocess.DagBundlesManager") + def test_rejects_unconfigured_dag_bundle_name(self, mock_manager): + mock_manager.is_bundle_configured.return_value = False + with pytest.raises(ValueError, match="unconfigured Dag bundle 'ghost'"): + _StubSubprocessCoordinator(command=["x"], explicit_roots=[], dag_bundle_name="ghost") + + @patch("airflow.sdk.coordinators._subprocess.DagBundlesManager") + def test_accepts_configured_dag_bundle_name(self, mock_manager): + mock_manager.is_bundle_configured.return_value = True + coordinator = _StubSubprocessCoordinator( + command=["x"], explicit_roots=[], dag_bundle_name="artifacts" + ) + assert coordinator._artifact_source is _ArtifactSource.NAMED_BUNDLE + mock_manager.is_bundle_configured.assert_called_once_with("artifacts") + + +class TestInitRootSource: + """Execute-time root resolution, dispatched on the classified mode.""" + + def test_returns_configured_root_and_no_bundle_in_explicit_mode(self): + configured = [pathlib.Path("/a"), pathlib.Path("/b")] + coordinator = _StubSubprocessCoordinator(command=["x"], explicit_roots=configured) + roots, bundle = coordinator._init_root_source(MagicMock(), log) + assert roots == configured + assert bundle is None + + @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle") + def test_task_bundle_mode_uses_passed_task_bundle(self, mock_initialize, tmp_path): + resolved = MagicMock(path=tmp_path, version="v3") + mock_initialize.return_value = resolved + coordinator = _StubSubprocessCoordinator(command=["x"], explicit_roots=[]) + bundle_info = BundleInfo(name="dags", version="v3", version_data={"k": "v"}) + + roots, bundle = coordinator._init_root_source(bundle_info, log) + + assert roots == [tmp_path] + assert bundle is resolved + mock_initialize.assert_called_once_with(bundle_info) + + @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle") + def test_version_less_bundle_is_rematerialized_at_its_current_version(self, mock_initialize, tmp_path): + """A bundle resolved without a version is re-resolved at the version current now. + + Otherwise the subprocess reads the bundle's shared mutable checkout, which no + version lock can protect and which another task's refresh can reset underneath it. + """ + shared_checkout = tmp_path / "tracking_repo" + shared_checkout.mkdir() + pinned_tree = tmp_path / "versions" / "sha-abc" + pinned_tree.mkdir(parents=True) + + unpinned = MagicMock(path=shared_checkout, version=None) + unpinned.get_current_version.return_value = BundleVersion(version="sha-abc", data={"k": "v"}) + pinned = MagicMock(path=pinned_tree, version="sha-abc") + mock_initialize.side_effect = [unpinned, pinned] + + coordinator = _StubSubprocessCoordinator(command=["x"], explicit_roots=[]) + coordinator.dag_bundle_name = "artifacts" + coordinator._artifact_source = _ArtifactSource.NAMED_BUNDLE + + roots, bundle = coordinator._init_root_source(MagicMock(), log) + + assert roots == [pinned_tree] + assert bundle is pinned + assert mock_initialize.call_args_list == [ + call(BundleInfo(name="artifacts")), + call(BundleInfo(name="artifacts", version="sha-abc", version_data={"k": "v"})), + ] + + @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle") + def test_unversioned_bundle_keeps_its_single_path(self, mock_initialize, tmp_path): + resolved = MagicMock(path=tmp_path, version=None) + resolved.get_current_version.return_value = None + mock_initialize.return_value = resolved + + coordinator = _StubSubprocessCoordinator(command=["x"], explicit_roots=[]) + coordinator.dag_bundle_name = "artifacts" + coordinator._artifact_source = _ArtifactSource.NAMED_BUNDLE + + roots, bundle = coordinator._init_root_source(MagicMock(), log) + + assert roots == [tmp_path] + assert bundle is resolved + mock_initialize.assert_called_once_with(BundleInfo(name="artifacts")) + + @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle") + def test_missing_resolved_path_raises(self, mock_initialize, tmp_path): + missing = tmp_path / "nope" + mock_initialize.return_value = MagicMock(path=missing, version="v1") + coordinator = _StubSubprocessCoordinator(command=["x"], explicit_roots=[]) + coordinator.dag_bundle_name = "artifacts" + coordinator._artifact_source = _ArtifactSource.NAMED_BUNDLE + + with pytest.raises(FileNotFoundError, match="does not exist"): + coordinator._init_root_source(MagicMock(), log) + + +class TestExecuteTaskBundleWiring: + """execute_task passes the task bundle, forwards resolved roots, and holds the version lock.""" + + @patch("airflow.sdk.coordinators._subprocess.BundleVersionLock") + @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle") + @patch.object(_PopenActivitySubprocess, "start") + def test_task_bundle_mode_binds_forwards_roots_and_locks( + self, mock_start, mock_initialize, mock_lock, mock_client, tmp_path + ): + resolved = MagicMock(path=tmp_path, version="v9") + resolved.name = "dags" + mock_initialize.return_value = resolved + mock_start.return_value.wait.return_value = 0 + + coordinator = _StubSubprocessCoordinator(command=["/runtime"], explicit_roots=[]) + bundle_info = BundleInfo(name="dags", version="v9") + + coordinator.execute_task( + what=_make_ti(), + dag_rel_path="dag.py", + bundle_info=bundle_info, + client=mock_client, + subprocess_logs_to_stdout=False, + ) + + mock_initialize.assert_called_once_with(bundle_info) + assert coordinator.recorded_roots == [[tmp_path]] + mock_lock.assert_called_once_with(bundle_name="dags", bundle_version="v9") + mock_lock.return_value.__enter__.assert_called_once() + mock_lock.return_value.__exit__.assert_called_once() + + @patch("airflow.sdk.coordinators._subprocess.BundleVersionLock") + @patch.object(_PopenActivitySubprocess, "start") + def test_explicit_root_mode_forwards_roots_without_locking( + self, mock_start, mock_lock, mock_client, tmp_path + ): + mock_start.return_value.wait.return_value = 0 + coordinator = _StubSubprocessCoordinator(command=["/runtime"], explicit_roots=[tmp_path]) + + coordinator.execute_task( + what=_make_ti(), + dag_rel_path="dag.py", + bundle_info=MagicMock(), + client=mock_client, + subprocess_logs_to_stdout=False, + ) + + assert coordinator.recorded_roots == [[tmp_path]] + mock_lock.assert_not_called() + + @patch("airflow.sdk.coordinators._subprocess.BundleVersionLock") + @patch("airflow.sdk.coordinators._subprocess.initialize_ti_bundle") + @patch.object(_PopenActivitySubprocess, "start") + def test_named_bundle_mode_locks_the_tree_it_scans( + self, mock_start, mock_initialize, mock_lock, mock_client, tmp_path + ): + """The lock names the pinned version whose tree is handed to the subprocess.""" + pinned_tree = tmp_path / "versions" / "sha-abc" + pinned_tree.mkdir(parents=True) + + unpinned = MagicMock(path=tmp_path, version=None) + unpinned.get_current_version.return_value = BundleVersion(version="sha-abc", data=None) + pinned = MagicMock(path=pinned_tree, version="sha-abc") + pinned.name = "artifacts" + mock_initialize.side_effect = [unpinned, pinned] + mock_start.return_value.wait.return_value = 0 + + coordinator = _StubSubprocessCoordinator(command=["/runtime"], explicit_roots=[]) + coordinator.dag_bundle_name = "artifacts" + coordinator._artifact_source = _ArtifactSource.NAMED_BUNDLE + + coordinator.execute_task( + what=_make_ti(), + dag_rel_path="dag.py", + bundle_info=MagicMock(), + client=mock_client, + subprocess_logs_to_stdout=False, + ) + + assert coordinator.recorded_roots == [[pinned_tree]] + mock_lock.assert_called_once_with(bundle_name="artifacts", bundle_version="sha-abc") + + +class TestGetScanRoots: + def test_returns_bound_roots_and_clears_them_on_exit(self, tmp_path): + coordinator = _StubSubprocessCoordinator(command=["x"]) + + with coordinator._set_scan_roots([tmp_path]): + assert coordinator._get_scan_roots() == (tmp_path,) + + with pytest.raises(RuntimeError, match="requires an active task"): + coordinator._get_scan_roots() + + def test_rejects_nested_reentry(self, tmp_path): + coordinator = _StubSubprocessCoordinator(command=["/bin/true"]) + with coordinator._set_scan_roots([tmp_path]): + with pytest.raises(RuntimeError, match="not re-entrant"): + with coordinator._set_scan_roots([tmp_path]): + pass + + def test_execute_task_is_not_reentrant(self, mock_client, tmp_path): + coordinator = _StubSubprocessCoordinator(command=["/bin/true"]) + with coordinator._set_scan_roots([tmp_path]): + with pytest.raises(RuntimeError, match="not re-entrant"): + coordinator.execute_task( + what=_make_ti(), + dag_rel_path="dag.py", + bundle_info=MagicMock(), + client=mock_client, + subprocess_logs_to_stdout=False, + ) diff --git a/task-sdk/tests/task_sdk/execution_time/test_bundles.py b/task-sdk/tests/task_sdk/execution_time/test_bundles.py new file mode 100644 index 0000000000000..0a58b6484dc64 --- /dev/null +++ b/task-sdk/tests/task_sdk/execution_time/test_bundles.py @@ -0,0 +1,86 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +from __future__ import annotations + +from pathlib import Path +from unittest import mock +from unittest.mock import patch + +import pytest + +from airflow.sdk.api.datamodels._generated import BundleInfo +from airflow.sdk.exceptions import AirflowException +from airflow.sdk.execution_time.bundles import initialize_ti_bundle, verify_bundle_access + + +def test_verify_bundle_access_raises_when_not_accessible(tmp_path: Path): + bundle_path = tmp_path / "test_bundle" + bundle_path.mkdir() + + mock_bundle = mock.Mock() + mock_bundle.path = bundle_path + mock_bundle.name = "test-bundle" + + # Mock os.access to simulate permission denied (avoids root user issues in CI) + with patch("airflow.sdk.execution_time.bundles.os.access", return_value=False): + with pytest.raises(AirflowException) as exc_info: + verify_bundle_access(mock_bundle) + + assert "not accessible" in str(exc_info.value) + assert "test-bundle" in str(exc_info.value) + + +def test_verify_bundle_access_succeeds_when_readable(tmp_path: Path): + bundle_path = tmp_path / "accessible_bundle" + bundle_path.mkdir() + + mock_bundle = mock.Mock() + mock_bundle.path = bundle_path + mock_bundle.name = "test-bundle" + + verify_bundle_access(mock_bundle) + + +def test_verify_bundle_access_skips_nonexistent_path(tmp_path: Path): + mock_bundle = mock.Mock() + mock_bundle.path = tmp_path / "nonexistent" + mock_bundle.name = "test-bundle" + + # Should not raise - nonexistent paths are handled by initialize() + verify_bundle_access(mock_bundle) + + +def test_initialize_ti_bundle_resolves_initializes_and_verifies(tmp_path: Path): + """initialize_ti_bundle resolves the bundle from BundleInfo, initializes it, and access-checks it.""" + bundle_path = tmp_path / "bundle" + bundle_path.mkdir() + mock_bundle = mock.Mock() + mock_bundle.path = bundle_path + mock_bundle.name = "my-bundle" + + bundle_info = BundleInfo(name="my-bundle", version="v2", version_data={"k": "v"}) + + with patch("airflow.sdk.execution_time.bundles.DagBundlesManager") as mock_manager: + mock_manager.return_value.get_bundle.return_value = mock_bundle + result = initialize_ti_bundle(bundle_info) + + assert result is mock_bundle + mock_manager.return_value.get_bundle.assert_called_once_with( + name="my-bundle", version="v2", version_data={"k": "v"} + ) + mock_bundle.initialize.assert_called_once_with() diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py index cc9fb77e08921..1811e97518aa0 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py +++ b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py @@ -180,7 +180,6 @@ _run_execute_callable, _serialize_outlet_events, _xcom_push, - detail_span, finalize, get_startup_details, parse, @@ -837,56 +836,6 @@ def test_parse_module_in_bundle_root(tmp_path: Path, make_ti_context): assert ti.task.dag.dag_id == "dag_name" -def test_verify_bundle_access_raises_when_not_accessible(tmp_path: Path, make_ti_context): - """Test that _verify_bundle_access raises AirflowException when bundle path is not accessible.""" - from airflow.sdk.execution_time.task_runner import _verify_bundle_access - - # Create a directory that exists - bundle_path = tmp_path / "test_bundle" - bundle_path.mkdir() - - # Create a mock bundle instance - mock_bundle = mock.Mock() - mock_bundle.path = bundle_path - mock_bundle.name = "test-bundle" - - # Mock os.access to simulate permission denied (avoids root user issues in CI) - with patch("airflow.sdk.execution_time.task_runner.os.access", return_value=False): - with pytest.raises(AirflowException) as exc_info: - _verify_bundle_access(mock_bundle, mock.Mock()) - - assert "not accessible" in str(exc_info.value) - assert "test-bundle" in str(exc_info.value) - - -def test_verify_bundle_access_succeeds_when_readable(tmp_path: Path, make_ti_context): - """Test that _verify_bundle_access succeeds when bundle path is accessible.""" - from airflow.sdk.execution_time.task_runner import _verify_bundle_access - - # Create a directory with read permissions - bundle_path = tmp_path / "accessible_bundle" - bundle_path.mkdir() - - mock_bundle = mock.Mock() - mock_bundle.path = bundle_path - mock_bundle.name = "test-bundle" - - # Should not raise - _verify_bundle_access(mock_bundle, mock.Mock()) - - -def test_verify_bundle_access_skips_nonexistent_path(tmp_path: Path): - """Test that _verify_bundle_access does nothing when bundle path doesn't exist.""" - from airflow.sdk.execution_time.task_runner import _verify_bundle_access - - mock_bundle = mock.Mock() - mock_bundle.path = tmp_path / "nonexistent" - mock_bundle.name = "test-bundle" - - # Should not raise - nonexistent paths are handled by initialize() - _verify_bundle_access(mock_bundle, mock.Mock()) - - @pytest.mark.parametrize("use_queues", [False, True]) def test_run_deferred_basic(time_machine, create_runtime_ti, mock_supervisor_comms, use_queues: bool): """Test that a task can transition to a deferred state.""" @@ -5973,120 +5922,6 @@ def test_operator_metrics_respect_team_name( backend.incr.assert_any_call(ti_metric, tags=stats_tags) -class TestDetailSpan: - """Tests for the detail_span decorator / context manager.""" - - @pytest.fixture(autouse=True) - def _sampled_carrier_provider(self): - """Make new_dagrun_trace_carrier produce a SAMPLED carrier. - - new_dagrun_trace_carrier consults the global tracer provider's sampler to - decide the carrier's SAMPLED flag. In the test process the global provider - is a no-op ProxyTracerProvider (no sampler) -> unsampled carrier, which - would make the parent span (and its detail children) non-recording. Patch - the lookup to a real SDK provider whose default sampler - (parentbased_always_on) samples the root, mirroring "otel on" in production. - """ - provider = TracerProvider() - with mock.patch( - "airflow._shared.observability.traces.trace.get_tracer_provider", - return_value=provider, - ): - yield - - def test_level_1_no_child_span_as_context_manager(self): - """At detail level 1, entering detail_span should not create a real recorded span.""" - exporter = InMemorySpanExporter() - provider = TracerProvider() - provider.add_span_processor(SimpleSpanProcessor(exporter)) - t = provider.get_tracer("test") - carrier = new_dagrun_trace_carrier(task_span_detail_level=1) - parent_ctx = TraceContextTextMapPropagator().extract(carrier) - - with mock.patch("airflow.sdk.execution_time.task_runner.tracer", t): - with t.start_as_current_span("parent", context=parent_ctx): - with detail_span("child") as span: - assert span is trace.INVALID_SPAN - - # Only the "parent" span should be recorded; no "child". - names = [s.name for s in exporter.get_finished_spans()] - assert "child" not in names - - def test_level_2_creates_child_span_as_context_manager(self): - """At detail level 2, detail_span should create a real recorded child span.""" - exporter = InMemorySpanExporter() - provider = TracerProvider() - provider.add_span_processor(SimpleSpanProcessor(exporter)) - t = provider.get_tracer("test") - carrier = new_dagrun_trace_carrier(task_span_detail_level=2) - parent_ctx = TraceContextTextMapPropagator().extract(carrier) - - with mock.patch("airflow.sdk.execution_time.task_runner.tracer", t): - with t.start_as_current_span("parent", context=parent_ctx): - with detail_span("child"): - pass - - names = [s.name for s in exporter.get_finished_spans()] - assert "child" in names - - def test_decorator_at_level_1_does_not_create_span(self): - """@detail_span at level 1 should not produce a recorded span.""" - exporter = InMemorySpanExporter() - provider = TracerProvider() - provider.add_span_processor(SimpleSpanProcessor(exporter)) - t = provider.get_tracer("test") - carrier = new_dagrun_trace_carrier(task_span_detail_level=1) - parent_ctx = TraceContextTextMapPropagator().extract(carrier) - - @detail_span("decorated") - def my_func(): - return 42 - - with mock.patch("airflow.sdk.execution_time.task_runner.tracer", t): - with t.start_as_current_span("parent", context=parent_ctx): - result = my_func() - - assert result == 42 - names = [s.name for s in exporter.get_finished_spans()] - assert "decorated" not in names - - def test_decorator_at_level_2_creates_span_and_preserves_return_value(self): - """@detail_span at level 2 creates a span and the wrapped function's return value is preserved.""" - exporter = InMemorySpanExporter() - provider = TracerProvider() - provider.add_span_processor(SimpleSpanProcessor(exporter)) - t = provider.get_tracer("test") - carrier = new_dagrun_trace_carrier(task_span_detail_level=2) - parent_ctx = TraceContextTextMapPropagator().extract(carrier) - - @detail_span("decorated") - def my_func(x): - return x * 2 - - with mock.patch("airflow.sdk.execution_time.task_runner.tracer", t): - with t.start_as_current_span("parent", context=parent_ctx): - result = my_func(7) - - assert result == 14 - names = [s.name for s in exporter.get_finished_spans()] - assert "decorated" in names - - def test_exception_in_context_manager_propagates(self): - """Exceptions inside `with detail_span(...)` propagate normally.""" - exporter = InMemorySpanExporter() - provider = TracerProvider() - provider.add_span_processor(SimpleSpanProcessor(exporter)) - t = provider.get_tracer("test") - carrier = new_dagrun_trace_carrier(task_span_detail_level=2) - parent_ctx = TraceContextTextMapPropagator().extract(carrier) - - with mock.patch("airflow.sdk.execution_time.task_runner.tracer", t): - with t.start_as_current_span("parent", context=parent_ctx): - with pytest.raises(ValueError, match="boom"): - with detail_span("child"): - raise ValueError("boom") - - class TestRunExecuteCallable: """Tests for ``_run_execute_callable``. @@ -6168,7 +6003,7 @@ def test_emits_task_execute_span_at_detail_level_2(self): task = self._make_task() - with mock.patch("airflow.sdk.execution_time.task_runner.tracer", t): + with mock.patch("airflow.sdk.execution_time.tracing.tracer", t): with t.start_as_current_span("parent", context=parent_ctx): result = _run_execute_callable(context={}, execute=lambda context: "ok", task=task) @@ -6196,7 +6031,7 @@ def execute(context): with t.start_as_current_span("operator_child"): return "ok" - with mock.patch("airflow.sdk.execution_time.task_runner.tracer", t): + with mock.patch("airflow.sdk.execution_time.tracing.tracer", t): with t.start_as_current_span("parent", context=parent_ctx): result = _run_execute_callable(context={}, execute=execute, task=task) @@ -6215,7 +6050,7 @@ def test_no_task_execute_span_at_detail_level_1(self): task = self._make_task() - with mock.patch("airflow.sdk.execution_time.task_runner.tracer", t): + with mock.patch("airflow.sdk.execution_time.tracing.tracer", t): with t.start_as_current_span("parent", context=parent_ctx): result = _run_execute_callable(context={}, execute=lambda context: "ok", task=task) diff --git a/task-sdk/tests/task_sdk/execution_time/test_tracing.py b/task-sdk/tests/task_sdk/execution_time/test_tracing.py new file mode 100644 index 0000000000000..fcd467ba6b8e7 --- /dev/null +++ b/task-sdk/tests/task_sdk/execution_time/test_tracing.py @@ -0,0 +1,144 @@ +# 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 + +from unittest import mock + +import pytest +from opentelemetry import trace +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator + +from airflow.sdk._shared.observability.traces import new_dagrun_trace_carrier +from airflow.sdk.execution_time.tracing import detail_span + + +class TestDetailSpan: + """Tests for the detail_span decorator / context manager.""" + + @pytest.fixture(autouse=True) + def _sampled_carrier_provider(self): + """Make new_dagrun_trace_carrier produce a SAMPLED carrier. + + new_dagrun_trace_carrier consults the global tracer provider's sampler to + decide the carrier's SAMPLED flag. In the test process the global provider + is a no-op ProxyTracerProvider (no sampler) -> unsampled carrier, which + would make the parent span (and its detail children) non-recording. Patch + the lookup to a real SDK provider whose default sampler + (parentbased_always_on) samples the root, mirroring "otel on" in production. + """ + provider = TracerProvider() + with mock.patch( + "airflow._shared.observability.traces.trace.get_tracer_provider", + return_value=provider, + ): + yield + + def test_level_1_no_child_span_as_context_manager(self): + """At detail level 1, entering detail_span should not create a real recorded span.""" + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + t = provider.get_tracer("test") + carrier = new_dagrun_trace_carrier(task_span_detail_level=1) + parent_ctx = TraceContextTextMapPropagator().extract(carrier) + + with mock.patch("airflow.sdk.execution_time.tracing.tracer", t): + with t.start_as_current_span("parent", context=parent_ctx): + with detail_span("child") as span: + assert span is trace.INVALID_SPAN + + # Only the "parent" span should be recorded; no "child". + names = [s.name for s in exporter.get_finished_spans()] + assert "child" not in names + + def test_level_2_creates_child_span_as_context_manager(self): + """At detail level 2, detail_span should create a real recorded child span.""" + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + t = provider.get_tracer("test") + carrier = new_dagrun_trace_carrier(task_span_detail_level=2) + parent_ctx = TraceContextTextMapPropagator().extract(carrier) + + with mock.patch("airflow.sdk.execution_time.tracing.tracer", t): + with t.start_as_current_span("parent", context=parent_ctx): + with detail_span("child"): + pass + + names = [s.name for s in exporter.get_finished_spans()] + assert "child" in names + + def test_decorator_at_level_1_does_not_create_span(self): + """@detail_span at level 1 should not produce a recorded span.""" + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + t = provider.get_tracer("test") + carrier = new_dagrun_trace_carrier(task_span_detail_level=1) + parent_ctx = TraceContextTextMapPropagator().extract(carrier) + + @detail_span("decorated") + def my_func(): + return 42 + + with mock.patch("airflow.sdk.execution_time.tracing.tracer", t): + with t.start_as_current_span("parent", context=parent_ctx): + result = my_func() + + assert result == 42 + names = [s.name for s in exporter.get_finished_spans()] + assert "decorated" not in names + + def test_decorator_at_level_2_creates_span_and_preserves_return_value(self): + """@detail_span at level 2 creates a span and the wrapped function's return value is preserved.""" + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + t = provider.get_tracer("test") + carrier = new_dagrun_trace_carrier(task_span_detail_level=2) + parent_ctx = TraceContextTextMapPropagator().extract(carrier) + + @detail_span("decorated") + def my_func(x): + return x * 2 + + with mock.patch("airflow.sdk.execution_time.tracing.tracer", t): + with t.start_as_current_span("parent", context=parent_ctx): + result = my_func(7) + + assert result == 14 + names = [s.name for s in exporter.get_finished_spans()] + assert "decorated" in names + + def test_exception_in_context_manager_propagates(self): + """Exceptions inside `with detail_span(...)` propagate normally.""" + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + t = provider.get_tracer("test") + carrier = new_dagrun_trace_carrier(task_span_detail_level=2) + parent_ctx = TraceContextTextMapPropagator().extract(carrier) + + with mock.patch("airflow.sdk.execution_time.tracing.tracer", t): + with t.start_as_current_span("parent", context=parent_ctx): + with pytest.raises(ValueError, match="boom"): + with detail_span("child"): + raise ValueError("boom")