From 1a67bf46b5560bbe9ad30cbbb40ac96f0a8ce9f0 Mon Sep 17 00:00:00 2001 From: "Jason(Zhe-You) Liu" <68415893+jason810496@users.noreply.github.com> Date: Wed, 30 Sep 2026 16:10:25 +0800 Subject: [PATCH] Support DagBundle based deployment for Coordinator (#70805) * Add dag_bundle_name option to subprocess coordinators Language-SDK coordinators previously required a hand-managed filesystem root (jars_root/executables_root/bundles_root) for compiled artifacts, kept outside Airflow's Dag-bundle machinery and prone to drifting from the Dags they serve. Allowing artifacts to be located through a configured Dag bundle -- or the task's own bundle when none is named -- lets them be delivered and versioned the same way as everything else. * Document dag_bundle_name option for language-SDK coordinators * Harden subprocess coordinator artifact bundle resolution Language-SDK subprocess coordinators materialized artifact bundles without the version-lock protection the Python task path holds, so bundle cleanup could delete a version a running JVM was still reading from. Resolution now holds the version lock for the subprocess lifetime, mirroring the Python task path, and materialization failures surface in the task log rather than the worker log. Classification moves to the base coordinator so a subclass that adds no explicit-root kwarg falls back to the task-bundle default instead of failing every task, the Java coordinator now requires an explicit entrypoint when it has to scan a whole bundle, and configured-bundle lookup checks names only so an unrelated unimportable bundle no longer breaks coordinator construction. * Make JavaCoordinator main_class optional in co-located mode When a subprocess coordinator locates artifacts through a Dag bundle rather than an explicit root, JavaCoordinator previously rejected a missing main_class outright. That made the co-located default unusable without extra configuration and left the entrypoint no more determinate than an auto-detected scan would be. Auto-detect over the resolved bundle now handles that case, with the remaining cross-bundle ambiguity deferred to an IMPORT_ERROR-stage check under AIP-85. A named Dag bundle resolves to its latest version, so its version was unset and the version lock held around the subprocess silently did nothing, letting bundle cleanup delete artifacts a running task was still reading. The concrete current version is now pinned before locking. * Resolve coordinator Dag bundles without importing the task runner A language-SDK task start needs a bundle materialized on disk before the subprocess can be launched, but reaching that through the task runner made every such start pay for the operator tree, pydantic, and OTel plumbing that the supervisor side never uses. The bundle lifecycle now lives on its own, so the coordinator depends on the Dag-bundle machinery alone. The bundle-name check a coordinator runs at construction and the class-importing parse the manager runs at startup previously read the configuration independently, so the two could disagree about what is configured. Keying the parsed configuration on the configuration itself lets both share one snapshot while a configuration change is still picked up. Subclasses had to declare the name of their explicit-root kwarg separately from its value, which is two things to keep in step for one piece of information. * Read the Dag bundle configuration through a single entry point Splitting the read across a keyed builder and a wrapper meant two names for one idea, and the empty case leaked out as a module-level constant that existed only to be returned. One function now answers "what is configured", parsed once per process so both readers of it see the same answer. Parsing once per process means a later configuration change is not observed. That only arises in tests, which switch configuration freely, so overriding configuration there drops the memo rather than silently serving a stale one. * Preserve custom subprocess coordinator compatibility External language SDK coordinators can implement the documented command hook. Keeping that contract stable avoids unexpected-keyword failures when upgrading while allowing Dag-bundle-backed artifact discovery. * Make task bundle access dependencies explicit Module-level imports keep dependency analysis reliable and follow Airflow's import-placement contract without changing bundle access behavior. * Reject concurrent subprocess coordinator execution consistently Bundle resolution does not need mutable instance state, while overwriting scan roots can make a command build against another task's artifacts. * Reject empty subprocess coordinator artifact roots An explicitly empty root is likely a configuration or templating error. Treating it as omitted can silently scan the task's Dag bundle and run unintended artifacts. * Reject blank subprocess coordinator artifact roots An empty string normalizes to the current directory, so it bypassed the empty-roots check and scanned the worker's working directory. --- .../language-sdks/go.rst | 22 +- .../language-sdks/java.rst | 30 +- .../language-sdks/typescript.rst | 22 +- .../airflow/dag_processing/bundles/manager.py | 40 ++- .../tests/integration/otel/test_otel.py | 2 +- .../bundles/test_dag_bundle_manager.py | 86 ++++-- contributing-docs/30_new_language_sdk.rst | 6 + .../src/tests_common/pytest_plugin.py | 26 ++ .../src/tests_common/test_utils/config.py | 13 + generated/known_airflow_exceptions.txt | 2 +- .../sdk/coordinators/_bundle_metadata.py | 28 +- .../airflow/sdk/coordinators/_subprocess.py | 211 ++++++++++++-- .../coordinators/executable/coordinator.py | 17 +- .../sdk/coordinators/java/coordinator.py | 24 +- .../sdk/coordinators/node/coordinator.py | 21 +- .../src/airflow/sdk/execution_time/bundles.py | 90 ++++++ .../airflow/sdk/execution_time/task_runner.py | 83 +----- .../src/airflow/sdk/execution_time/tracing.py | 60 ++++ .../executable/test_coordinator.py | 58 +++- .../coordinators/java/test_coordinator.py | 39 +++ .../coordinators/node/test_coordinator.py | 51 +++- .../task_sdk/coordinators/test_subprocess.py | 264 +++++++++++++++++- .../task_sdk/execution_time/test_bundles.py | 86 ++++++ .../execution_time/test_task_runner.py | 171 +----------- .../task_sdk/execution_time/test_tracing.py | 144 ++++++++++ 25 files changed, 1250 insertions(+), 346 deletions(-) create mode 100644 task-sdk/src/airflow/sdk/execution_time/bundles.py create mode 100644 task-sdk/src/airflow/sdk/execution_time/tracing.py create mode 100644 task-sdk/tests/task_sdk/execution_time/test_bundles.py create mode 100644 task-sdk/tests/task_sdk/execution_time/test_tracing.py 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 72ea37bf665..bd2a3e4ac76 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/go.rst @@ -507,14 +507,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 dbd5d8996ce..3f69772e4d9 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/java.rst @@ -1045,9 +1045,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``. @@ -1056,15 +1062,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 b89173b80d8..ed8532859ce 100644 --- a/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst +++ b/airflow-core/docs/authoring-and-scheduling/language-sdks/typescript.rst @@ -482,9 +482,15 @@ All ``kwargs`` in the ``coordinators`` config entry are passed to the - Default - Description * - ``bundles_root`` - - *(required)* + - *(optional)* - One or more directories searched recursively, in order, for an integrity-verified ``*.min.mjs`` - bundle that declares the requested Dag. Accepts a string, a path, or a list of strings/paths. + bundle that declares the requested Dag. 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 ``*.min.mjs`` bundle from. Mutually exclusive with + ``bundles_root``. * - ``node_executable`` - ``"node"`` - Path to the ``node`` binary. Defaults to ``node`` on ``$PATH``. @@ -493,6 +499,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 2c20f923739..12b3b50fbca 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,28 @@ 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.""" + bundle_config_list = _read_bundle_config_list() + if not bundle_config_list: + return _BundleConfigSnapshot(configs=(), names=frozenset()) + 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: @@ -300,12 +323,9 @@ class DagBundlesManager(LoggingMixin): if self._bundle_config: return - bundle_config_list = _read_bundle_config_list() + bundle_config_list = _load_bundle_config_snapshot().configs if not bundle_config_list: return - 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"): @@ -673,6 +693,16 @@ class DagBundlesManager(LoggingMixin): 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 39d7ee6bdab..5be887bbdfa 100644 --- a/airflow-core/tests/integration/otel/test_otel.py +++ b/airflow-core/tests/integration/otel/test_otel.py @@ -330,7 +330,7 @@ class TestOtelIntegration: # 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 7d23af9afef..6eac0c575fb 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 @@ -33,6 +33,7 @@ from airflow.dag_processing.bundles.manager import ( DagBundlesManager, _get_configured_bundle_team_names, _guess_best_bundle_for_fileloc, + _load_bundle_config_snapshot, ) from airflow.exceptions import AirflowConfigException from airflow.models.dag import DagModel @@ -205,6 +206,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() @@ -221,9 +271,7 @@ def test_sync_bundles_to_db(clear_db, session): ).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)] @@ -248,9 +296,7 @@ def test_sync_bundles_to_db(clear_db, session): 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() == [ @@ -298,16 +344,12 @@ def test_sync_bundles_to_db_partial_config_does_not_disable_other_bundles(clear_ ).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), @@ -315,9 +357,7 @@ def test_sync_bundles_to_db_partial_config_does_not_disable_other_bundles(clear_ ] # 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), @@ -416,9 +456,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() @@ -441,9 +479,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() @@ -460,9 +496,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() @@ -485,9 +519,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 80fbe9e4235..34c4cdbb4a2 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 445581cd2b8..bb85fe60c2d 100644 --- a/devel-common/src/tests_common/pytest_plugin.py +++ b/devel-common/src/tests_common/pytest_plugin.py @@ -2186,6 +2186,32 @@ def clear_current_task_instance_session(): task_instance_session.__current_task_instance_session = None +@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 9893da665a4..619a88a6930 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 cb860b5cd71..ca2b46365cb 100644 --- a/generated/known_airflow_exceptions.txt +++ b/generated/known_airflow_exceptions.txt @@ -407,6 +407,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 9a0ee5e5b32..b60aa94144b 100644 --- a/task-sdk/src/airflow/sdk/coordinators/_bundle_metadata.py +++ b/task-sdk/src/airflow/sdk/coordinators/_bundle_metadata.py @@ -22,7 +22,7 @@ from __future__ import annotations import os import pathlib import stat -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Final import attrs import structlog @@ -38,6 +38,13 @@ if TYPE_CHECKING: log: FilteringBoundLogger = structlog.get_logger(logger_name="coordinators") +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]: @@ -96,6 +103,25 @@ def _sorted_children(directory: pathlib.Path) -> list[pathlib.Path]: return [] +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 [] + values = [value] if isinstance(value, (str, os.PathLike)) else (value or []) + if not values or any(not os.fspath(v).strip() for v in values): + raise ValueError( + "Artifact roots must contain at least one path when provided, and each path must be non-empty; " + "omit the option to use the task's Dag bundle." + ) + return convert_roots(value) + + 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 a0c3f518fb0..9e69255ba3d 100644 --- a/task-sdk/src/airflow/sdk/coordinators/_subprocess.py +++ b/task-sdk/src/airflow/sdk/coordinators/_subprocess.py @@ -27,6 +27,8 @@ draining machinery in this module rather than re-implementing it. from __future__ import annotations +import contextlib +import enum import ipaddress import itertools import os @@ -41,18 +43,24 @@ import attrs 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 @@ class _PopenActivitySubprocess(ActivitySubprocess): 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 @@ class SubprocessCoordinator(BaseCoordinator): 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 70ba5490450..9d6b80b249f 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 attrs 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 0ebe089162b..1f0f6e6efd5 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 @@ from typing import TYPE_CHECKING 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 cd352f3871b..dd2a3f6f041 100644 --- a/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py +++ b/task-sdk/src/airflow/sdk/coordinators/node/coordinator.py @@ -26,7 +26,12 @@ from typing import TYPE_CHECKING import attrs import structlog -from airflow.sdk.coordinators._bundle_metadata import ResolvedBundle, convert_roots, walk_files +from airflow.sdk.coordinators._bundle_metadata import ( + ARTIFACT_ROOTS_NOT_CONFIGURED, + ResolvedBundle, + convert_configured_roots, + walk_files, +) from airflow.sdk.coordinators._subprocess import SubprocessCoordinator from airflow.sdk.coordinators.node._bundle_reader import read_bundle @@ -111,17 +116,23 @@ class NodeCoordinator(SubprocessCoordinator): :param node_executable: Path to the ``node`` binary (defaults to ``"node"``, which relies on ``$PATH``). :param bundles_root: Directories searched recursively, in order, for the first verified - ``*.min.mjs`` bundle declaring the task instance's Dag. + ``*.min.mjs`` bundle declaring the task instance's Dag. 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]: - bundle = _Bundle.find(self.bundles_root, what.dag_id) + roots = self._get_scan_roots() + bundle = _Bundle.find(roots, what.dag_id) 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 00000000000..8ab131beb64 --- /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 fba2ac56829..4e2b5bc7354 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -44,10 +44,8 @@ from pydantic import AwareDatetime, ConfigDict, Field, JsonValue, TypeAdapter 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 ( @@ -86,6 +84,7 @@ from airflow.sdk.exceptions import ( 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, @@ -148,6 +147,7 @@ from airflow.sdk.execution_time.email_backend import ( _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 @@ -171,38 +171,6 @@ log = structlog.get_logger("task") 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 = ( @@ -1034,13 +1002,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)) @@ -1136,43 +1098,6 @@ SUPERVISOR_COMMS: CommsDecoder[ToTask, ToSupervisor] # 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. @@ -1280,8 +1205,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 00000000000..86a28199c61 --- /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 5963e5f4892..c20913e26ea 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,46 @@ class TestExecutableCoordinatorAttributes: 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_optional_defaults_to_empty(self): + coordinator = ExecutableCoordinator() + assert coordinator.executables_root == [] + assert coordinator.dag_bundle_name is None - def test_executables_root_must_be_non_empty(self): - with pytest.raises(ValueError, match="executables_root"): - ExecutableCoordinator(executables_root=None) + @pytest.mark.parametrize( + "executables_root", + [None, [], "", " ", [""]], + ids=["none", "empty-list", "empty-str", "blank-str", "list-of-empty-str"], + ) + def test_explicit_empty_executables_root_raises(self, executables_root): + with pytest.raises(ValueError, match="and each path must be non-empty"): + 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 +418,8 @@ class TestBuildExecuteTaskCommand: 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 +430,10 @@ class TestBuildExecuteTaskCommand: 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 +441,10 @@ class TestBuildExecuteTaskCommand: 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 c30ecff35b9..c55a2cc8912 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,45 @@ class TestJavaCoordinatorAttributes: 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", "empty-str", "blank-str", "list-of-empty-str"], + ) + def test_explicit_empty_jars_root_raises(self, jars_root): + with pytest.raises(ValueError, match="and each path must be non-empty"): + 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 492c410c8aa..146cb91c1d9 100644 --- a/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py +++ b/task-sdk/tests/task_sdk/coordinators/node/test_coordinator.py @@ -21,6 +21,7 @@ from __future__ import annotations import json import pathlib from unittest import mock +from unittest.mock import patch import pytest from task_sdk.coordinators.node._bundle_test_utils import ( @@ -79,9 +80,37 @@ class TestNodeCoordinatorAttributes: ] 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", "empty-str", "blank-str", "list-of-empty-str"], + ) + def test_explicit_empty_bundles_root_raises(self, bundles_root): + with pytest.raises(ValueError, match="and each path must be non-empty"): + 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, "test_dag") + 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 TestNodeCoordinatorExecuteTaskCommand: @@ -92,11 +121,25 @@ class TestNodeCoordinatorExecuteTaskCommand: bundles_root=tmp_path, ) - command, schema_version = coordinator._build_execute_task_command(what=_make_ti(dag_id="sales")) + with coordinator._set_scan_roots([tmp_path]): + command, schema_version = coordinator._build_execute_task_command(what=_make_ti(dag_id="sales")) assert command == ["/opt/node/bin/node", str(selected)] assert schema_version == SCHEMA_VERSION + def test_build_execute_task_command_returns_node_bundle_and_schema_version(self, tmp_path): + bundle = write_bundle(tmp_path, "test_dag") + coordinator = NodeCoordinator( + node_executable="/opt/node/bin/node", + bundles_root=tmp_path, + ) + + 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 + class TestBundleFind: @pytest.mark.parametrize( diff --git a/task-sdk/tests/task_sdk/coordinators/test_subprocess.py b/task-sdk/tests/task_sdk/coordinators/test_subprocess.py index 62b7fbcf39c..5297f56dead 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 @@ from __future__ import annotations import contextlib import os +import pathlib import socket import subprocess import sys @@ -31,16 +32,19 @@ import psutil 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 @@ class TestResourceTracker: @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 @@ class TestPopenActivitySubprocessStart: 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 00000000000..0a58b6484dc --- /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 b14e3e1bb69..42e7c045449 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 @@ -181,7 +181,6 @@ from airflow.sdk.execution_time.task_runner import ( _run_execute_callable, _serialize_outlet_events, _xcom_push, - detail_span, finalize, get_startup_details, parse, @@ -838,56 +837,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.""" @@ -6188,120 +6137,6 @@ class TestTaskInstanceMetrics: 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``. @@ -6383,7 +6218,7 @@ class TestRunExecuteCallable: 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) @@ -6411,7 +6246,7 @@ class TestRunExecuteCallable: 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) @@ -6430,7 +6265,7 @@ class TestRunExecuteCallable: 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 00000000000..fcd467ba6b8 --- /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")