From d33ac32d36bf92b750cb5d5292a60dd6cfaa4a77 Mon Sep 17 00:00:00 2001 From: Sameer Raj Date: Tue, 15 Sep 2026 22:01:13 +0530 Subject: [PATCH 01/34] Allow custom Dag bundle configuration providers Deployments need to load Dag bundle configuration from sources other than the static Airflow configuration. A provider boundary lets the Dag processor select that source without coupling it to bundle reconciliation. --- .../src/airflow/cli/commands/team_command.py | 6 +- .../src/airflow/config_templates/config.yml | 13 + .../airflow/dag_processing/bundles/manager.py | 289 +++++------------- .../dag_processing/bundles/provider.py | 251 +++++++++++++++ .../src/airflow/dag_processing/dagbag.py | 30 +- .../src/airflow/dag_processing/manager.py | 2 +- .../bundles/test_dag_bundle_manager.py | 144 ++++++++- .../tests/unit/dag_processing/test_dagbag.py | 38 +-- 8 files changed, 515 insertions(+), 258 deletions(-) create mode 100644 airflow-core/src/airflow/dag_processing/bundles/provider.py diff --git a/airflow-core/src/airflow/cli/commands/team_command.py b/airflow-core/src/airflow/cli/commands/team_command.py index 630141cfd86d2..70655fd57ad39 100644 --- a/airflow-core/src/airflow/cli/commands/team_command.py +++ b/airflow-core/src/airflow/cli/commands/team_command.py @@ -218,7 +218,7 @@ def team_sync(args, *, session=NEW_SESSION): dag_bundle_teams = { bundle.team_name - for bundle in DagBundlesManager()._bundle_config.values() + for bundle in DagBundlesManager().get_all_bundle_configurations() if bundle.team_name is not None } existing_teams = Team.get_all_team_names(session=session) @@ -311,9 +311,9 @@ def team_verify(args, *, session=NEW_SESSION): existing_teams = {team.name for team in teams} - for bundle_name, bundle in DagBundlesManager()._bundle_config.items(): + for bundle in DagBundlesManager().get_all_bundle_configurations(): if bundle.team_name and bundle.team_name not in existing_teams: - issues.append(f"DAG bundle '{bundle_name}' references unknown team '{bundle.team_name}'.") + issues.append(f"DAG bundle '{bundle.name}' references unknown team '{bundle.team_name}'.") if issues: print("Verification failed.\n") diff --git a/airflow-core/src/airflow/config_templates/config.yml b/airflow-core/src/airflow/config_templates/config.yml index 873b519576fab..6cf2041f8a601 100644 --- a/airflow-core/src/airflow/config_templates/config.yml +++ b/airflow-core/src/airflow/config_templates/config.yml @@ -3277,6 +3277,19 @@ dag_processor: example: "/tmp/some-place" default: ~ + dag_bundle_provider: + description: | + Import path for the class that provides Dag bundle configurations and constructs Dag bundles. + The class must extend ``airflow.dag_processing.bundles.provider.DagBundleProvider`` and have a + no-argument constructor. + + The default provider reads ``dag_bundle_config_list``. A custom provider can read bundle + configuration from another source and must return the complete set of active bundles. + version_added: 3.4.0 + type: string + example: ~ + default: "airflow.dag_processing.bundles.provider.ConfigDagBundleProvider" + dag_bundle_config_list: description: | List of backend configs. Must supply name, classpath, and kwargs for each backend. diff --git a/airflow-core/src/airflow/dag_processing/bundles/manager.py b/airflow-core/src/airflow/dag_processing/bundles/manager.py index fe74ffff7ab9f..5f42c1e768b9b 100644 --- a/airflow-core/src/airflow/dag_processing/bundles/manager.py +++ b/airflow-core/src/airflow/dag_processing/bundles/manager.py @@ -16,26 +16,23 @@ # 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, NamedTuple, cast +from typing import TYPE_CHECKING, Any, cast from itsdangerous import URLSafeSerializer -from pydantic import BaseModel, ValidationError from sqlalchemy import and_, delete, exists, or_, select, update -from airflow._shared.module_loading import import_string from airflow.configuration import conf from airflow.dag_processing.bundles.base import BaseDagBundle # noqa: TC001 +from airflow.dag_processing.bundles.provider import ( + DagBundleConfiguration, + DagBundleProvider, +) from airflow.exceptions import AirflowConfigException from airflow.models.dagbundle import DagBundleModel from airflow.models.team import Team -from airflow.providers_manager import ProvidersManager from airflow.utils.log.logging_mixin import LoggingMixin from airflow.utils.session import NEW_SESSION, create_session, provide_session @@ -45,10 +42,6 @@ from sqlalchemy.engine import CursorResult from sqlalchemy.orm import Session -log = logging.getLogger(__name__) - -_example_dag_bundle_name = "example_dags" - # Chunk size for the one-time startup repair of unconfigured bundles. _REASSIGN_BATCH_SIZE = 1000 @@ -80,178 +73,13 @@ def _guess_best_bundle_for_fileloc( return None -class _ExternalBundleConfig(BaseModel): - """Schema defining the user-specified configuration for a DAG bundle.""" - - name: str - classpath: str - kwargs: dict - team_name: str | None = None - - -class _InternalBundleConfig(BaseModel): - """ - Schema used internally (in this file) to define the configuration for a DAG bundle. - - Configuration defined by users when read must match ``_ExternalBundleConfig``. - This configuration is then parsed and converted to ``_InternalBundleConfig`` to be used across this file. - """ - - bundle_class: type[BaseDagBundle] - kwargs: dict - team_name: str | None = None - - -def _bundle_item_exc(msg): - return AirflowConfigException( - "Invalid config for section `dag_processor` key `dag_bundle_config_list`. " + 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: - if not isinstance(item, dict): - raise _bundle_item_exc(f"Expected dict but got {item.__class__}") - - try: - cfg = _ExternalBundleConfig(**item) - except ValidationError as e: - raise _bundle_item_exc(f"Item {item} failed validation: {e}") - - if cfg.name == _example_dag_bundle_name: - raise AirflowConfigException( - f"Bundle name '{_example_dag_bundle_name}' is a reserved name. Please choose another name for your bundle." - " Example DAGs can be enabled with the '[core] load_examples' config." - ) - - bundles[cfg.name] = cfg - if len(bundles.keys()) != len(config_list): - raise _bundle_item_exc("One or more bundle names appeared multiple times") - return list(bundles.values()) - - -def _read_bundle_config_list() -> list[_ExternalBundleConfig]: - config_list = conf.getjson("dag_processor", "dag_bundle_config_list") - if not config_list: - return [] - if not isinstance(config_list, list): - raise AirflowConfigException( - "Section `dag_processor` key `dag_bundle_config_list` " - f"must be list but got {config_list.__class__}" - ) - return _parse_bundle_config(config_list) - - -def get_configured_bundle_names() -> frozenset[str]: - """Get the names of the configured Dag bundles, without importing their classes.""" - return _load_bundle_config_snapshot().names - - -def get_configured_bundle_team_names() -> dict[str, str | None]: - """ - Get the team owning each explicitly configured Dag bundle. - - This reads the config rather than going through ``DagBundlesManager`` so that callers who only - need the declared bundle partition neither import every bundle class nor see the example-Dag - bundles that ``DagBundlesManager.parse_config`` injects when ``[core] load_examples`` is set -- - those are added by Airflow, not declared by the deployment. - - :return: mapping of bundle name to team name, ``None`` for bundles that are not team scoped. - """ - return {cfg.name: cfg.team_name for cfg in _read_bundle_config_list()} - - -def _add_example_dag_bundle(bundle_config_list: list[_ExternalBundleConfig]): - from airflow import example_dags - - example_dag_folder = next(iter(example_dags.__path__)) - bundle_config_list.append( - _ExternalBundleConfig( - name=_example_dag_bundle_name, - classpath="airflow.dag_processing.bundles.local.LocalDagBundle", - kwargs={ - "path": example_dag_folder, - }, - ) - ) - - -def _add_provider_example_dags_to_bundle(bundle_config_list: list[_ExternalBundleConfig]): - """ - Add an ``example_dags`` folder of every installed provider as a bundle. - - Provider locations are resolved through ``ProvidersManager`` instead of - walking ``airflow.providers.__path__`` so that: - - - nested providers (e.g. ``apache-airflow-providers-common-sql`` whose - module path is ``airflow.providers.common.sql``) are discovered; - - providers installed outside the ``airflow.providers`` namespace package - are discovered via their entry point. - """ - # Dedup on the resolved on-disk folder rather than the bundle name: distributions - # under ``airflow.providers.common.*`` use ``pkgutil.extend_path``, so when several - # ``common-*`` packages are installed ``airflow.providers.common.__path__`` has - # multiple entries and the inner loop iterates more than once. Path-based dedup - # only skips when the same folder is seen twice; distinct folders are preserved. - seen: set[str] = set() - - for package_name in ProvidersManager().providers: - # Heuristic: derive the import path from the canonical - # ``apache-airflow-providers-*`` distribution name. Tracked as a follow-up - # to record the provider module path on ``ProviderInfo`` (see - # https://github.com/apache/airflow/issues/66305). - if package_name.startswith("apache-airflow-providers-"): - suffix = package_name[len("apache-airflow-providers-") :] - module_name = "airflow.providers." + suffix.replace("-", ".") - else: - module_name = package_name.replace("-", "_") - try: - module = importlib.import_module(module_name) - module_paths = list(getattr(module, "__path__", [])) - except Exception: - log.exception("Could not load provider module %s for example DAG discovery", module_name) - continue - - for module_path in module_paths: - example_dag_folder = os.path.join(module_path, "example_dags") - if not os.path.isdir(example_dag_folder): - continue - if example_dag_folder in seen: - continue - seen.add(example_dag_folder) - bundle_name = f"{package_name}-example-dags" - bundle_config_list.append( - _ExternalBundleConfig( - name=bundle_name, - classpath="airflow.dag_processing.bundles.local.LocalDagBundle", - kwargs={ - "path": example_dag_folder, - }, - ) - ) +def _get_configured_bundle_team_names() -> dict[str, str | None]: + """Get the team owning each Dag bundle returned by the configured provider.""" + manager = DagBundlesManager() + return { + configuration.name: configuration.team_name + for configuration in manager._bundle_provider.get_all_bundle_configurations() + } def _is_safe_bundle_url(url: str) -> bool: @@ -314,38 +142,57 @@ class DagBundlesManager(LoggingMixin): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - self._bundle_config: dict[str, _InternalBundleConfig] = {} - self.parse_config() + self._bundle_provider = self._load_bundle_provider() + self._bundle_config: dict[str, DagBundleConfiguration] = {} + self.refresh_bundle_configurations() - def parse_config(self) -> None: - """ - Get all DAG bundle configurations and store in instance variable. + @staticmethod + def _load_bundle_provider() -> DagBundleProvider: + provider_class = conf.getimport("dag_processor", "dag_bundle_provider") + if not isinstance(provider_class, type) or not issubclass(provider_class, DagBundleProvider): + raise AirflowConfigException( + "The class configured in section `dag_processor` key `dag_bundle_provider` " + "must extend DagBundleProvider." + ) + try: + return provider_class() + except AirflowConfigException: + raise + except Exception as e: + raise AirflowConfigException( + "Could not initialize the Dag bundle provider configured in section " + "`dag_processor` key `dag_bundle_provider`." + ) from e - If a bundle class for a given name has already been imported, it will not be imported again. + def refresh_bundle_configurations(self) -> None: + """ + Refresh the active Dag bundle configurations from the provider. :meta private: """ if self._bundle_config: return - bundle_config_list = _load_bundle_config_snapshot().configs - if not bundle_config_list: - return - - for bundle_config in bundle_config_list: + bundle_configurations: dict[str, DagBundleConfiguration] = {} + for bundle_config in self._bundle_provider.get_all_bundle_configurations(): + if not isinstance(bundle_config, DagBundleConfiguration): + raise AirflowConfigException( + "Dag bundle providers must return DagBundleConfiguration objects from " + "get_all_bundle_configurations()." + ) + if bundle_config.name in bundle_configurations: + raise AirflowConfigException( + f"Dag bundle provider returned duplicate bundle name '{bundle_config.name}'." + ) if bundle_config.team_name and not conf.getboolean("core", "multi_team"): raise AirflowConfigException( - "Section `dag_processor` key `dag_bundle_config_list` " - "cannot have a team name when multi-team mode is disabled." - "To enable multi-team, you need to update section `core` key `multi_team` in your config." + "Dag bundle configurations cannot have a team name when multi-team mode is disabled. " + "To enable multi-team, update section `core` key `multi_team` in your config." ) - - class_ = import_string(bundle_config.classpath) - self._bundle_config[bundle_config.name] = _InternalBundleConfig( - bundle_class=class_, - kwargs=bundle_config.kwargs, - team_name=bundle_config.team_name, - ) + bundle_configurations[bundle_config.name] = bundle_config + if not bundle_configurations: + return + self._bundle_config = bundle_configurations self.log.info("DAG bundles loaded: %s", ", ".join(self._bundle_config.keys())) @provide_session @@ -398,7 +245,9 @@ def _extract_and_sign_template(bundle_name: str) -> tuple[str | None, dict]: if config.team_name: team = session.scalars(select(Team).where(Team.name == config.team_name)).one_or_none() if not team: - raise _bundle_item_exc(f"Team '{config.team_name}' does not exist") + raise AirflowConfigException( + f"Team '{config.team_name}' configured for Dag bundle '{name}' does not exist." + ) try: new_template, new_params = _extract_and_sign_template(name) @@ -617,8 +466,8 @@ def reassign_dags_with_unconfigured_bundles(self) -> int: self.log.warning( "Skipped %d legacy Dag(s) whose fileloc is not under any configured bundle; " "triggering them will keep raising \"Requested bundle '{name}' is not configured.\" " - "until a bundle whose path contains the fileloc is added to " - "dag_bundle_config_list. The next parse will then restore them automatically, " + "until the configured Dag bundle provider returns a bundle whose path contains the " + "fileloc. The next parse will then restore them automatically, " "or run `airflow dags reserialize` to force the parse path to rewrite " "bundle_name and relative_fileloc immediately.", total_skipped, @@ -691,12 +540,18 @@ def get_bundle( :return: The DAG bundle. """ - cfg_bundle = self._bundle_config.get(name) - if not cfg_bundle: - raise ValueError(f"Requested bundle '{name}' is not configured.") - return cfg_bundle.bundle_class( - name=name, version=version, version_data=version_data, **cfg_bundle.kwargs - ) + return self._bundle_provider.get_bundle(name=name, version=version, version_data=version_data) + + def get_bundle_configuration(self, name: str) -> DagBundleConfiguration: + """Get the active configuration for a Dag bundle.""" + try: + return self._bundle_config[name] + except KeyError: + raise ValueError(f"Requested bundle '{name}' is not configured.") from None + + def get_all_bundle_configurations(self) -> tuple[DagBundleConfiguration, ...]: + """Get all active Dag bundle configurations.""" + return tuple(self._bundle_config.values()) @classmethod def is_bundle_configured(cls, name: str) -> bool: @@ -706,7 +561,7 @@ def is_bundle_configured(cls, name: str) -> bool: 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 + return name in cls().get_all_bundle_names() def get_all_dag_bundles(self) -> Iterable[BaseDagBundle]: """ @@ -714,9 +569,9 @@ def get_all_dag_bundles(self) -> Iterable[BaseDagBundle]: :return: list of DAG bundles. """ - for name, cfg in self._bundle_config.items(): + for name in self._bundle_config: try: - yield cfg.bundle_class(name=name, version=None, **cfg.kwargs) + yield self.get_bundle(name=name) except Exception as e: self.log.exception("Error creating bundle '%s': %s", name, e) # Skip this bundle and continue with others diff --git a/airflow-core/src/airflow/dag_processing/bundles/provider.py b/airflow-core/src/airflow/dag_processing/bundles/provider.py new file mode 100644 index 0000000000000..216296a67ea4e --- /dev/null +++ b/airflow-core/src/airflow/dag_processing/bundles/provider.py @@ -0,0 +1,251 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import importlib +import logging +import os +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +from pydantic import BaseModel, ValidationError + +from airflow._shared.module_loading import import_string +from airflow.configuration import conf +from airflow.dag_processing.bundles.base import BaseDagBundle # noqa: TC001 +from airflow.exceptions import AirflowConfigException +from airflow.providers_manager import ProvidersManager + +if TYPE_CHECKING: + from collections.abc import Sequence + +log = logging.getLogger(__name__) + +_example_dag_bundle_name = "example_dags" + + +@dataclass(frozen=True) +class DagBundleConfiguration: + """Configuration used by Airflow to manage a Dag bundle.""" + + name: str + team_name: str | None = None + + +class DagBundleProvider(ABC): + """ + Provide Dag bundle configurations and construct Dag bundles. + + The configuration list is the complete set of active bundles. A bundle name + may remain resolvable through ``get_bundle`` after it leaves that list because + retained Dag runs can still need an older version. A name must continue to + identify the same bundle implementation; use a new name for a different one. + """ + + @abstractmethod + def get_all_bundle_configurations(self) -> Sequence[DagBundleConfiguration]: + """Return the complete set of active Dag bundle configurations.""" + + @abstractmethod + def get_bundle( + self, + name: str, + version: str | None = None, + version_data: dict[str, Any] | None = None, + ) -> BaseDagBundle: + """Construct a Dag bundle by name and optional version information.""" + + +class _ExternalBundleConfig(BaseModel): + """Schema defining the user-specified configuration for a Dag bundle.""" + + name: str + classpath: str + kwargs: dict + team_name: str | None = None + + +class _InternalBundleConfig(BaseModel): + """ + Schema used internally (in this file) to define the configuration for a Dag bundle. + + Configuration defined by users when read must match ``_ExternalBundleConfig``. + This configuration is then parsed and converted to ``_InternalBundleConfig`` to be used across this file. + """ + + bundle_class: type[BaseDagBundle] + kwargs: dict + team_name: str | None = None + + +def _bundle_item_exc(msg): + return AirflowConfigException( + "Invalid config for section `dag_processor` key `dag_bundle_config_list`. " + msg + ) + + +def _parse_bundle_config(config_list) -> list[_ExternalBundleConfig]: + bundles = {} + for item in config_list: + if not isinstance(item, dict): + raise _bundle_item_exc(f"Expected dict but got {item.__class__}") + + try: + cfg = _ExternalBundleConfig(**item) + except ValidationError as e: + raise _bundle_item_exc(f"Item {item} failed validation: {e}") + + if cfg.name == _example_dag_bundle_name: + raise AirflowConfigException( + f"Bundle name '{_example_dag_bundle_name}' is a reserved name. Please choose another name for your bundle." + " Example Dags can be enabled with the '[core] load_examples' config." + ) + + bundles[cfg.name] = cfg + if len(bundles.keys()) != len(config_list): + raise _bundle_item_exc("One or more bundle names appeared multiple times") + return list(bundles.values()) + + +def _add_example_dag_bundle(bundle_config_list: list[_ExternalBundleConfig]): + from airflow import example_dags + + example_dag_folder = next(iter(example_dags.__path__)) + bundle_config_list.append( + _ExternalBundleConfig( + name=_example_dag_bundle_name, + classpath="airflow.dag_processing.bundles.local.LocalDagBundle", + kwargs={ + "path": example_dag_folder, + }, + ) + ) + + +def _add_provider_example_dags_to_bundle(bundle_config_list: list[_ExternalBundleConfig]): + """ + Add an ``example_dags`` folder of every installed provider as a bundle. + + Provider locations are resolved through ``ProvidersManager`` instead of + walking ``airflow.providers.__path__`` so that: + + - nested providers (e.g. ``apache-airflow-providers-common-sql`` whose + module path is ``airflow.providers.common.sql``) are discovered; + - providers installed outside the ``airflow.providers`` namespace package + are discovered via their entry point. + """ + # Dedup on the resolved on-disk folder rather than the bundle name: distributions + # under ``airflow.providers.common.*`` use ``pkgutil.extend_path``, so when several + # ``common-*`` packages are installed ``airflow.providers.common.__path__`` has + # multiple entries and the inner loop iterates more than once. Path-based dedup + # only skips when the same folder is seen twice; distinct folders are preserved. + seen: set[str] = set() + + for package_name in ProvidersManager().providers: + # Heuristic: derive the import path from the canonical + # ``apache-airflow-providers-*`` distribution name. Tracked as a follow-up + # to record the provider module path on ``ProviderInfo`` (see + # https://github.com/apache/airflow/issues/66305). + if package_name.startswith("apache-airflow-providers-"): + suffix = package_name[len("apache-airflow-providers-") :] + module_name = "airflow.providers." + suffix.replace("-", ".") + else: + module_name = package_name.replace("-", "_") + try: + module = importlib.import_module(module_name) + module_paths = list(getattr(module, "__path__", [])) + except Exception: + log.exception("Could not load provider module %s for example Dag discovery", module_name) + continue + + for module_path in module_paths: + example_dag_folder = os.path.join(module_path, "example_dags") + if not os.path.isdir(example_dag_folder): + continue + if example_dag_folder in seen: + continue + seen.add(example_dag_folder) + bundle_name = f"{package_name}-example-dags" + bundle_config_list.append( + _ExternalBundleConfig( + name=bundle_name, + classpath="airflow.dag_processing.bundles.local.LocalDagBundle", + kwargs={ + "path": example_dag_folder, + }, + ) + ) + + +class ConfigDagBundleProvider(DagBundleProvider): + """Provide Dag bundles configured by ``dag_bundle_config_list``.""" + + def __init__(self) -> None: + self._bundle_config: dict[str, _InternalBundleConfig] = {} + self._load_configured_bundles() + + def _load_configured_bundles(self) -> None: + """ + Load Dag bundles from ``dag_bundle_config_list``. + + If a bundle class for a given name has already been imported, it will not be imported again. + + :meta private: + """ + if self._bundle_config: + return + + config_list = conf.getjson("dag_processor", "dag_bundle_config_list") + if not config_list: + return + if not isinstance(config_list, list): + raise AirflowConfigException( + "Section `dag_processor` key `dag_bundle_config_list` " + f"must be list but got {config_list.__class__}" + ) + bundle_config_list = _parse_bundle_config(config_list) + if conf.getboolean("core", "LOAD_EXAMPLES"): + _add_example_dag_bundle(bundle_config_list) + _add_provider_example_dags_to_bundle(bundle_config_list) + + for bundle_config in bundle_config_list: + class_ = import_string(bundle_config.classpath) + self._bundle_config[bundle_config.name] = _InternalBundleConfig( + bundle_class=class_, + kwargs=bundle_config.kwargs, + team_name=bundle_config.team_name, + ) + + def get_all_bundle_configurations(self) -> Sequence[DagBundleConfiguration]: + return [ + DagBundleConfiguration(name=name, team_name=config.team_name) + for name, config in self._bundle_config.items() + ] + + def get_bundle( + self, + name: str, + version: str | None = None, + version_data: dict[str, Any] | None = None, + ) -> BaseDagBundle: + cfg_bundle = self._bundle_config.get(name) + if not cfg_bundle: + raise ValueError(f"Requested bundle '{name}' is not configured.") + return cfg_bundle.bundle_class( + name=name, version=version, version_data=version_data, **cfg_bundle.kwargs + ) diff --git a/airflow-core/src/airflow/dag_processing/dagbag.py b/airflow-core/src/airflow/dag_processing/dagbag.py index dc80433676ca7..595936b519bb0 100644 --- a/airflow-core/src/airflow/dag_processing/dagbag.py +++ b/airflow-core/src/airflow/dag_processing/dagbag.py @@ -128,9 +128,20 @@ def _validate_executor_fields(dag: DAG | SerializedDAG, bundle_name: str | None log = logging.getLogger(__name__) - dag_team_name = _get_bundle_team_name(bundle_name) - if dag_team_name: - log.debug("Found team '%s' for DAG '%s' via bundle '%s'", dag_team_name, dag.dag_id, bundle_name) + # Check if multi team is available by reading the multi_team configuration (which is boolean) + if conf.getboolean("core", "multi_team"): + # Get team name from bundle configuration if available + if bundle_name: + from airflow.dag_processing.bundles.manager import DagBundlesManager + + bundle_manager = DagBundlesManager() + bundle_config = bundle_manager.get_bundle_configuration(bundle_name) + + dag_team_name = bundle_config.team_name + if dag_team_name: + log.debug( + "Found team '%s' for DAG '%s' via bundle '%s'", dag_team_name, dag.dag_id, bundle_name + ) for task in dag.tasks: if not task.executor: @@ -157,7 +168,18 @@ def _assign_default_team_pools( bundle_name: str | None = None, ) -> None: """Assign the default team pool to tasks that do not explicitly specify a pool.""" - if not (dag_team_name := _get_bundle_team_name(bundle_name)): + dag_team_name = None + + if conf.getboolean("core", "multi_team"): + if bundle_name: + from airflow.dag_processing.bundles.manager import DagBundlesManager + + bundle_manager = DagBundlesManager() + bundle_config = bundle_manager.get_bundle_configuration(bundle_name) + + dag_team_name = bundle_config.team_name + + if not dag_team_name: return for task in dag.tasks: diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index 343d72c5e002d..e5043104dba99 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -557,7 +557,7 @@ def deactivate_stale_dags( if stuck_legacy_rows: # Surface how many legacy rows the startup repair could not route; # each one keeps raising "Requested bundle is not configured." until - # a matching bundle is added to dag_bundle_config_list. + # the configured Dag bundle provider returns a matching bundle. self.log.info( "Skipped stale check for %d legacy Dag(s) with NULL relative_fileloc.", stuck_legacy_rows, 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 8938d2d33b0fc..dd8892c77c47d 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 @@ -36,6 +36,7 @@ get_configured_bundle_names, get_configured_bundle_team_names, ) +from airflow.dag_processing.bundles.provider import DagBundleConfiguration, DagBundleProvider from airflow.exceptions import AirflowConfigException from airflow.models.dag import DagModel from airflow.models.dag_version import DagVersion @@ -124,6 +125,44 @@ def path(self) -> Path: return Path("/__basic_bundle_unmatched__") +class CustomDagBundleProvider(DagBundleProvider): + def get_all_bundle_configurations(self): + return [DagBundleConfiguration(name="active-bundle")] + + def get_bundle(self, name, version=None, version_data=None): + if name not in {"active-bundle", "retired-bundle"}: + raise ValueError(f"Unknown test bundle: {name}") + return BasicBundle( + name=name, + version=version, + version_data=version_data, + refresh_interval=1, + ) + + +class DuplicateConfigurationDagBundleProvider(CustomDagBundleProvider): + def get_all_bundle_configurations(self): + return [ + DagBundleConfiguration(name="duplicate"), + DagBundleConfiguration(name="duplicate"), + ] + + +class NonConfigurationDagBundleProvider(CustomDagBundleProvider): + def get_all_bundle_configurations(self): + return [object()] + + +class MissingTeamDagBundleProvider(CustomDagBundleProvider): + def get_all_bundle_configurations(self): + return [DagBundleConfiguration(name="active-bundle", team_name="missing-team")] + + +class FailingInitializationDagBundleProvider(CustomDagBundleProvider): + def __init__(self): + raise RuntimeError("Provider initialization failed") + + BASIC_BUNDLE_CONFIG = [ { "name": "my-test-bundle", @@ -221,13 +260,6 @@ 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", @@ -270,6 +302,78 @@ def test_is_bundle_configured_sees_implicitly_added_bundles(): assert DagBundlesManager.is_bundle_configured("example_dags") is True +@conf_vars( + { + ( + "dag_processor", + "dag_bundle_provider", + ): "unit.dag_processing.bundles.test_dag_bundle_manager.CustomDagBundleProvider" + } +) +def test_custom_bundle_provider_resolves_active_and_retired_bundles(): + manager = DagBundlesManager() + + assert manager.get_all_bundle_configurations() == (DagBundleConfiguration(name="active-bundle"),) + assert manager.get_bundle_configuration("active-bundle") == DagBundleConfiguration(name="active-bundle") + assert manager.get_all_bundle_names() == ["active-bundle"] + with pytest.raises(ValueError, match="'unknown-bundle' is not configured"): + manager.get_bundle_configuration("unknown-bundle") + + active_bundle = manager.get_bundle("active-bundle") + retired_bundle = manager.get_bundle( + "retired-bundle", + version="v1", + version_data={"manifest": "retired.json"}, + ) + + assert active_bundle.name == "active-bundle" + assert retired_bundle.name == "retired-bundle" + assert retired_bundle.version == "v1" + assert retired_bundle.version_data == {"manifest": "retired.json"} + + +@pytest.mark.parametrize( + ("provider_class", "expected_message"), + [ + pytest.param( + "unit.dag_processing.bundles.test_dag_bundle_manager.DuplicateConfigurationDagBundleProvider", + "duplicate bundle name 'duplicate'", + id="duplicate-name", + ), + pytest.param( + "unit.dag_processing.bundles.test_dag_bundle_manager.NonConfigurationDagBundleProvider", + "must return DagBundleConfiguration objects", + id="wrong-configuration-type", + ), + ], +) +def test_custom_bundle_provider_rejects_invalid_configurations(provider_class, expected_message): + with conf_vars({("dag_processor", "dag_bundle_provider"): provider_class}): + with pytest.raises(AirflowConfigException, match=expected_message): + DagBundlesManager() + + +@pytest.mark.parametrize( + ("provider_class", "expected_message"), + [ + pytest.param( + "unit.dag_processing.bundles.test_dag_bundle_manager.BasicBundle", + "must extend DagBundleProvider", + id="wrong-base-class", + ), + pytest.param( + "unit.dag_processing.bundles.test_dag_bundle_manager.FailingInitializationDagBundleProvider", + "Could not initialize the Dag bundle provider", + id="initialization-failure", + ), + ], +) +def test_invalid_custom_bundle_provider(provider_class, expected_message): + with conf_vars({("dag_processor", "dag_bundle_provider"): provider_class}): + with pytest.raises(AirflowConfigException, match=expected_message): + DagBundlesManager() + + @pytest.fixture def clear_db(): clear_db_dag_bundles() @@ -320,6 +424,26 @@ def _get_bundle_names_and_active(): ] +@pytest.mark.db_test +@conf_vars( + { + ("core", "multi_team"): "True", + ( + "dag_processor", + "dag_bundle_provider", + ): "unit.dag_processing.bundles.test_dag_bundle_manager.MissingTeamDagBundleProvider", + } +) +def test_sync_bundles_to_db_reports_missing_team_from_provider(clear_db): + manager = DagBundlesManager() + + with pytest.raises( + AirflowConfigException, + match="Team 'missing-team' configured for Dag bundle 'active-bundle' does not exist", + ): + manager.sync_bundles_to_db() + + @pytest.mark.db_test @conf_vars({("core", "LOAD_EXAMPLES"): "False"}) def test_sync_bundles_to_db_does_not_log_removing_none_team(clear_db, caplog): @@ -558,12 +682,12 @@ def test_dag_bundle_model_render_url_with_invalid_template(): def test_example_dags_bundle_added(): manager = DagBundlesManager() - manager.parse_config() + manager.refresh_bundle_configurations() assert "example_dags" in manager._bundle_config with conf_vars({("core", "LOAD_EXAMPLES"): "False"}): manager = DagBundlesManager() - manager.parse_config() + manager.refresh_bundle_configurations() assert "example_dags" not in manager._bundle_config @@ -571,7 +695,7 @@ def test_example_dags_name_is_reserved(): reserved_name_config = [{"name": "example_dags", "classpath": "yo face", "kwargs": {}}] with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(reserved_name_config)}): with pytest.raises(AirflowConfigException, match="Bundle name 'example_dags' is a reserved name."): - DagBundlesManager().parse_config() + DagBundlesManager().refresh_bundle_configurations() class FailingBundle(BaseDagBundle): diff --git a/airflow-core/tests/unit/dag_processing/test_dagbag.py b/airflow-core/tests/unit/dag_processing/test_dagbag.py index d24443f709fe3..97a84c0220c02 100644 --- a/airflow-core/tests/unit/dag_processing/test_dagbag.py +++ b/airflow-core/tests/unit/dag_processing/test_dagbag.py @@ -36,6 +36,7 @@ from sqlalchemy import select from airflow import settings +from airflow.dag_processing.bundles.provider import DagBundleConfiguration from airflow.dag_processing.dagbag import ( BundleDagBag, DagBag, @@ -135,11 +136,10 @@ def test_multi_team_disabled_ignores_bundle_name(self, mock_lookup): def test_multi_team_enabled_bundle_exists_with_team(self, mock_lookup, mock_manager_class): """Test successful team lookup when bundle exists and has team_name.""" # Setup mock bundle manager - mock_bundle_config = mock.MagicMock() - mock_bundle_config.team_name = "test_team" + bundle_config = DagBundleConfiguration(name="test_bundle", team_name="test_team") mock_manager = mock_manager_class.return_value - mock_manager._bundle_config = {"test_bundle": mock_bundle_config} + mock_manager.get_bundle_configuration.return_value = bundle_config with DAG("test-dag", schedule=None) as dag: BaseOperator(task_id="t1", executor="team.executor") @@ -154,11 +154,10 @@ def test_multi_team_enabled_bundle_exists_with_team(self, mock_lookup, mock_mana @patch.object(ExecutorLoader, "lookup_executor_name_by_str") def test_multi_team_enabled_bundle_exists_no_team(self, mock_lookup, mock_manager_class): """Test when bundle exists but has no team_name (None or empty).""" - mock_bundle_config = mock.MagicMock() - mock_bundle_config.team_name = None # No team associated + bundle_config = DagBundleConfiguration(name="test_bundle") mock_manager = mock_manager_class.return_value - mock_manager._bundle_config = {"test_bundle": mock_bundle_config} + mock_manager.get_bundle_configuration.return_value = bundle_config with DAG("test-dag", schedule=None) as dag: BaseOperator(task_id="t1", executor="test.executor") @@ -188,11 +187,10 @@ def test_multiple_tasks_with_executors(self, mock_lookup): @patch.object(ExecutorLoader, "lookup_executor_name_by_str") def test_executor_validation_failure_with_team(self, mock_lookup, mock_manager_class): """Test executor validation failure when team is associated (team-specific error).""" - mock_bundle_config = mock.MagicMock() - mock_bundle_config.team_name = "test_team" + bundle_config = DagBundleConfiguration(name="test_bundle", team_name="test_team") mock_manager = mock_manager_class.return_value - mock_manager._bundle_config = {"test_bundle": mock_bundle_config} + mock_manager.get_bundle_configuration.return_value = bundle_config # ExecutorLoader raises exception mock_lookup.side_effect = UnknownExecutorException("Executor not found") @@ -236,11 +234,10 @@ def test_executor_validation_failure_no_team(self, mock_lookup): @patch.object(ExecutorLoader, "lookup_executor_name_by_str") def test_global_executor_fallback_success(self, mock_lookup, mock_manager_class): """Test that team-specific executor failure falls back to global executor successfully.""" - mock_bundle_config = mock.MagicMock() - mock_bundle_config.team_name = "test_team" + bundle_config = DagBundleConfiguration(name="test_bundle", team_name="test_team") mock_manager = mock_manager_class.return_value - mock_manager._bundle_config = {"test_bundle": mock_bundle_config} + mock_manager.get_bundle_configuration.return_value = bundle_config # First call (team-specific) fails, second call (global) succeeds mock_lookup.side_effect = [UnknownExecutorException("Team executor not found"), None] @@ -261,11 +258,10 @@ def test_global_executor_fallback_success(self, mock_lookup, mock_manager_class) @patch.object(ExecutorLoader, "lookup_executor_name_by_str") def test_global_executor_fallback_failure(self, mock_lookup, mock_manager_class): """Test that when both team-specific and global executors fail, appropriate error is raised.""" - mock_bundle_config = mock.MagicMock() - mock_bundle_config.team_name = "test_team" + bundle_config = DagBundleConfiguration(name="test_bundle", team_name="test_team") mock_manager = mock_manager_class.return_value - mock_manager._bundle_config = {"test_bundle": mock_bundle_config} + mock_manager.get_bundle_configuration.return_value = bundle_config # Both calls fail mock_lookup.side_effect = UnknownExecutorException("Executor not found") @@ -295,11 +291,10 @@ def test_global_executor_fallback_failure(self, mock_lookup, mock_manager_class) @patch.object(ExecutorLoader, "lookup_executor_name_by_str") def test_team_specific_executor_success_no_fallback(self, mock_lookup, mock_manager_class): """Test that when team-specific executor succeeds, global fallback is not attempted.""" - mock_bundle_config = mock.MagicMock() - mock_bundle_config.team_name = "test_team" + bundle_config = DagBundleConfiguration(name="test_bundle", team_name="test_team") mock_manager = mock_manager_class.return_value - mock_manager._bundle_config = {"test_bundle": mock_bundle_config} + mock_manager.get_bundle_configuration.return_value = bundle_config # First call (team-specific) succeeds mock_lookup.return_value = None @@ -470,11 +465,8 @@ def test_default_pool_replaced_with_team_pool( operator_args, expected_pool, ): - mock_bundle = mock.MagicMock() - mock_bundle.team_name = team_name - mock_manager.return_value._bundle_config = { - "test_bundle": mock_bundle, - } + bundle_config = DagBundleConfiguration(name="test_bundle", team_name=team_name) + mock_manager.return_value.get_bundle_configuration.return_value = bundle_config dag_file = tmp_path / "test_dag.py" dag_file.write_text( From 8e1166ee235c5dd242a60c64b4e31c4beb757a58 Mon Sep 17 00:00:00 2001 From: Sameer Raj Date: Thu, 17 Sep 2026 11:40:20 +0530 Subject: [PATCH 02/34] Clarify Dag bundle provider configuration Airflow also uses provider to describe installable integrations. Clear configuration wording prevents this extension class from being mistaken for a new provider package. --- airflow-core/src/airflow/config_templates/config.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/airflow-core/src/airflow/config_templates/config.yml b/airflow-core/src/airflow/config_templates/config.yml index 6cf2041f8a601..c52c7fa8c7dfa 100644 --- a/airflow-core/src/airflow/config_templates/config.yml +++ b/airflow-core/src/airflow/config_templates/config.yml @@ -3283,8 +3283,8 @@ dag_processor: The class must extend ``airflow.dag_processing.bundles.provider.DagBundleProvider`` and have a no-argument constructor. - The default provider reads ``dag_bundle_config_list``. A custom provider can read bundle - configuration from another source and must return the complete set of active bundles. + The default implementation reads ``dag_bundle_config_list``. A custom implementation can read + bundle configuration from another source and must return the complete set of active bundles. version_added: 3.4.0 type: string example: ~ From 9554e461761d6a87b339c7e0727e8213c5c28031 Mon Sep 17 00:00:00 2001 From: Sameer Raj Date: Mon, 21 Sep 2026 15:45:51 +0530 Subject: [PATCH 03/34] Simplify Dag bundle provider responsibilities Keeping Airflow-owned validation and example Dag handling in the manager gives custom providers one consistent contract and avoids requiring each implementation to reproduce manager behavior. --- .../airflow/dag_processing/bundles/manager.py | 105 +++++++++++++----- .../dag_processing/bundles/provider.py | 90 +-------------- .../bundles/test_dag_bundle_manager.py | 38 +++++-- 3 files changed, 112 insertions(+), 121 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/bundles/manager.py b/airflow-core/src/airflow/dag_processing/bundles/manager.py index 5f42c1e768b9b..138412bcff467 100644 --- a/airflow-core/src/airflow/dag_processing/bundles/manager.py +++ b/airflow-core/src/airflow/dag_processing/bundles/manager.py @@ -16,8 +16,12 @@ # under the License. from __future__ import annotations +import importlib +import logging +import os import warnings from collections import defaultdict +from functools import cached_property from pathlib import Path from typing import TYPE_CHECKING, Any, cast @@ -26,6 +30,7 @@ from airflow.configuration import conf from airflow.dag_processing.bundles.base import BaseDagBundle # noqa: TC001 +from airflow.dag_processing.bundles.local import LocalDagBundle from airflow.dag_processing.bundles.provider import ( DagBundleConfiguration, DagBundleProvider, @@ -33,6 +38,7 @@ from airflow.exceptions import AirflowConfigException from airflow.models.dagbundle import DagBundleModel from airflow.models.team import Team +from airflow.providers_manager import ProvidersManager from airflow.utils.log.logging_mixin import LoggingMixin from airflow.utils.session import NEW_SESSION, create_session, provide_session @@ -42,10 +48,48 @@ from sqlalchemy.engine import CursorResult from sqlalchemy.orm import Session +log = logging.getLogger(__name__) + +_example_dag_bundle_name = "example_dags" + # Chunk size for the one-time startup repair of unconfigured bundles. _REASSIGN_BATCH_SIZE = 1000 +def _discover_example_dag_bundle_paths() -> dict[str, str]: + """Return the example Dag bundle names and paths available in this installation.""" + from airflow import example_dags + + bundle_paths = {_example_dag_bundle_name: next(iter(example_dags.__path__))} + seen: set[str] = set() + + for package_name in ProvidersManager().providers: + # Heuristic: derive the import path from the canonical + # ``apache-airflow-providers-*`` distribution name. Tracked as a follow-up + # to record the provider module path on ``ProviderInfo`` (see + # https://github.com/apache/airflow/issues/66305). + if package_name.startswith("apache-airflow-providers-"): + suffix = package_name[len("apache-airflow-providers-") :] + module_name = "airflow.providers." + suffix.replace("-", ".") + else: + module_name = package_name.replace("-", "_") + try: + module = importlib.import_module(module_name) + module_paths = list(getattr(module, "__path__", [])) + except Exception: + log.exception("Could not load provider module %s for example Dag discovery", module_name) + continue + + for module_path in module_paths: + example_dag_folder = os.path.join(module_path, "example_dags") + if not os.path.isdir(example_dag_folder) or example_dag_folder in seen: + continue + seen.add(example_dag_folder) + bundle_paths[f"{package_name}-example-dags"] = example_dag_folder + + return bundle_paths + + def _guess_best_bundle_for_fileloc( fileloc: str, descending_bundle_paths: dict[str, Path] ) -> tuple[str, str] | None: @@ -143,8 +187,12 @@ class DagBundlesManager(LoggingMixin): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self._bundle_provider = self._load_bundle_provider() - self._bundle_config: dict[str, DagBundleConfiguration] = {} - self.refresh_bundle_configurations() + + @cached_property + def _example_dag_bundle_paths(self) -> dict[str, str]: + if not conf.getboolean("core", "LOAD_EXAMPLES"): + return {} + return _discover_example_dag_bundle_paths() @staticmethod def _load_bundle_provider() -> DagBundleProvider: @@ -164,15 +212,8 @@ def _load_bundle_provider() -> DagBundleProvider: "`dag_processor` key `dag_bundle_provider`." ) from e - def refresh_bundle_configurations(self) -> None: - """ - Refresh the active Dag bundle configurations from the provider. - - :meta private: - """ - if self._bundle_config: - return - + def get_all_bundle_configurations(self) -> tuple[DagBundleConfiguration, ...]: + """Get all active Dag bundle configurations.""" bundle_configurations: dict[str, DagBundleConfiguration] = {} for bundle_config in self._bundle_provider.get_all_bundle_configurations(): if not isinstance(bundle_config, DagBundleConfiguration): @@ -180,6 +221,11 @@ def refresh_bundle_configurations(self) -> None: "Dag bundle providers must return DagBundleConfiguration objects from " "get_all_bundle_configurations()." ) + if bundle_config.name == _example_dag_bundle_name: + raise AirflowConfigException( + f"Bundle name '{_example_dag_bundle_name}' is a reserved name. Please choose another name for your bundle." + " Example Dags can be enabled with the '[core] load_examples' config." + ) if bundle_config.name in bundle_configurations: raise AirflowConfigException( f"Dag bundle provider returned duplicate bundle name '{bundle_config.name}'." @@ -190,10 +236,13 @@ def refresh_bundle_configurations(self) -> None: "To enable multi-team, update section `core` key `multi_team` in your config." ) bundle_configurations[bundle_config.name] = bundle_config - if not bundle_configurations: - return - self._bundle_config = bundle_configurations - self.log.info("DAG bundles loaded: %s", ", ".join(self._bundle_config.keys())) + + for name in self._example_dag_bundle_paths: + if name in bundle_configurations: + raise AirflowConfigException(f"Bundle name '{name}' is reserved for example Dags.") + bundle_configurations[name] = DagBundleConfiguration(name=name) + + return tuple(bundle_configurations.values()) @provide_session def sync_bundles_to_db(self, *, deactivate_missing: bool = True, session: Session = NEW_SESSION) -> None: @@ -216,6 +265,8 @@ def sync_bundles_to_db(self, *, deactivate_missing: bool = True, session: Sessio """ self.log.debug("Syncing DAG bundles to the database") + bundle_configurations = {config.name: config for config in self.get_all_bundle_configurations()} + def _extract_and_sign_template(bundle_name: str) -> tuple[str | None, dict]: bundle_instance = self.get_bundle(name) new_template_ = bundle_instance.view_url_template() @@ -240,7 +291,7 @@ def _extract_and_sign_template(bundle_name: str) -> tuple[str | None, dict]: for bundle in stored.values() } - for name, config in self._bundle_config.items(): + for name, config in bundle_configurations.items(): team: Team | None = None if config.team_name: team = session.scalars(select(Team).where(Team.name == config.team_name)).one_or_none() @@ -540,18 +591,16 @@ def get_bundle( :return: The DAG bundle. """ + if path := self._example_dag_bundle_paths.get(name): + return LocalDagBundle(name=name, path=path, version=version, version_data=version_data) return self._bundle_provider.get_bundle(name=name, version=version, version_data=version_data) def get_bundle_configuration(self, name: str) -> DagBundleConfiguration: """Get the active configuration for a Dag bundle.""" - try: - return self._bundle_config[name] - except KeyError: - raise ValueError(f"Requested bundle '{name}' is not configured.") from None - - def get_all_bundle_configurations(self) -> tuple[DagBundleConfiguration, ...]: - """Get all active Dag bundle configurations.""" - return tuple(self._bundle_config.values()) + for configuration in self.get_all_bundle_configurations(): + if configuration.name == name: + return configuration + raise ValueError(f"Requested bundle '{name}' is not configured.") @classmethod def is_bundle_configured(cls, name: str) -> bool: @@ -569,11 +618,11 @@ def get_all_dag_bundles(self) -> Iterable[BaseDagBundle]: :return: list of DAG bundles. """ - for name in self._bundle_config: + for configuration in self.get_all_bundle_configurations(): try: - yield self.get_bundle(name=name) + yield self.get_bundle(name=configuration.name) except Exception as e: - self.log.exception("Error creating bundle '%s': %s", name, e) + self.log.exception("Error creating bundle '%s': %s", configuration.name, e) # Skip this bundle and continue with others continue @@ -583,7 +632,7 @@ def get_all_bundle_names(self) -> Iterable[str]: :return: sorted list of bundle names. """ - return sorted(self._bundle_config.keys()) + return sorted(configuration.name for configuration in self.get_all_bundle_configurations()) def view_url(self, name: str, version: str | None = None) -> str | None: warnings.warn( diff --git a/airflow-core/src/airflow/dag_processing/bundles/provider.py b/airflow-core/src/airflow/dag_processing/bundles/provider.py index 216296a67ea4e..ff497edd6944b 100644 --- a/airflow-core/src/airflow/dag_processing/bundles/provider.py +++ b/airflow-core/src/airflow/dag_processing/bundles/provider.py @@ -16,9 +16,6 @@ # under the License. from __future__ import annotations -import importlib -import logging -import os from abc import ABC, abstractmethod from dataclasses import dataclass from typing import TYPE_CHECKING, Any @@ -29,15 +26,10 @@ from airflow.configuration import conf from airflow.dag_processing.bundles.base import BaseDagBundle # noqa: TC001 from airflow.exceptions import AirflowConfigException -from airflow.providers_manager import ProvidersManager if TYPE_CHECKING: from collections.abc import Sequence -log = logging.getLogger(__name__) - -_example_dag_bundle_name = "example_dags" - @dataclass(frozen=True) class DagBundleConfiguration: @@ -49,7 +41,7 @@ class DagBundleConfiguration: class DagBundleProvider(ABC): """ - Provide Dag bundle configurations and construct Dag bundles. + Provide Dag bundle configurations and Dag bundles. The configuration list is the complete set of active bundles. A bundle name may remain resolvable through ``get_bundle`` after it leaves that list because @@ -68,7 +60,7 @@ def get_bundle( version: str | None = None, version_data: dict[str, Any] | None = None, ) -> BaseDagBundle: - """Construct a Dag bundle by name and optional version information.""" + """Return a Dag bundle by name and optional version information.""" class _ExternalBundleConfig(BaseModel): @@ -110,9 +102,9 @@ def _parse_bundle_config(config_list) -> list[_ExternalBundleConfig]: except ValidationError as e: raise _bundle_item_exc(f"Item {item} failed validation: {e}") - if cfg.name == _example_dag_bundle_name: + if cfg.name == "example_dags": raise AirflowConfigException( - f"Bundle name '{_example_dag_bundle_name}' is a reserved name. Please choose another name for your bundle." + "Bundle name 'example_dags' is a reserved name. Please choose another name for your bundle." " Example Dags can be enabled with the '[core] load_examples' config." ) @@ -122,76 +114,6 @@ def _parse_bundle_config(config_list) -> list[_ExternalBundleConfig]: return list(bundles.values()) -def _add_example_dag_bundle(bundle_config_list: list[_ExternalBundleConfig]): - from airflow import example_dags - - example_dag_folder = next(iter(example_dags.__path__)) - bundle_config_list.append( - _ExternalBundleConfig( - name=_example_dag_bundle_name, - classpath="airflow.dag_processing.bundles.local.LocalDagBundle", - kwargs={ - "path": example_dag_folder, - }, - ) - ) - - -def _add_provider_example_dags_to_bundle(bundle_config_list: list[_ExternalBundleConfig]): - """ - Add an ``example_dags`` folder of every installed provider as a bundle. - - Provider locations are resolved through ``ProvidersManager`` instead of - walking ``airflow.providers.__path__`` so that: - - - nested providers (e.g. ``apache-airflow-providers-common-sql`` whose - module path is ``airflow.providers.common.sql``) are discovered; - - providers installed outside the ``airflow.providers`` namespace package - are discovered via their entry point. - """ - # Dedup on the resolved on-disk folder rather than the bundle name: distributions - # under ``airflow.providers.common.*`` use ``pkgutil.extend_path``, so when several - # ``common-*`` packages are installed ``airflow.providers.common.__path__`` has - # multiple entries and the inner loop iterates more than once. Path-based dedup - # only skips when the same folder is seen twice; distinct folders are preserved. - seen: set[str] = set() - - for package_name in ProvidersManager().providers: - # Heuristic: derive the import path from the canonical - # ``apache-airflow-providers-*`` distribution name. Tracked as a follow-up - # to record the provider module path on ``ProviderInfo`` (see - # https://github.com/apache/airflow/issues/66305). - if package_name.startswith("apache-airflow-providers-"): - suffix = package_name[len("apache-airflow-providers-") :] - module_name = "airflow.providers." + suffix.replace("-", ".") - else: - module_name = package_name.replace("-", "_") - try: - module = importlib.import_module(module_name) - module_paths = list(getattr(module, "__path__", [])) - except Exception: - log.exception("Could not load provider module %s for example Dag discovery", module_name) - continue - - for module_path in module_paths: - example_dag_folder = os.path.join(module_path, "example_dags") - if not os.path.isdir(example_dag_folder): - continue - if example_dag_folder in seen: - continue - seen.add(example_dag_folder) - bundle_name = f"{package_name}-example-dags" - bundle_config_list.append( - _ExternalBundleConfig( - name=bundle_name, - classpath="airflow.dag_processing.bundles.local.LocalDagBundle", - kwargs={ - "path": example_dag_folder, - }, - ) - ) - - class ConfigDagBundleProvider(DagBundleProvider): """Provide Dag bundles configured by ``dag_bundle_config_list``.""" @@ -219,10 +141,6 @@ def _load_configured_bundles(self) -> None: f"must be list but got {config_list.__class__}" ) bundle_config_list = _parse_bundle_config(config_list) - if conf.getboolean("core", "LOAD_EXAMPLES"): - _add_example_dag_bundle(bundle_config_list) - _add_provider_example_dags_to_bundle(bundle_config_list) - for bundle_config in bundle_config_list: class_ = import_string(bundle_config.classpath) self._bundle_config[bundle_config.name] = _InternalBundleConfig( 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 dd8892c77c47d..0a5d09028fa3f 100644 --- a/airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py +++ b/airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py @@ -29,6 +29,7 @@ from sqlalchemy import func, select, update from airflow.dag_processing.bundles.base import BaseDagBundle +from airflow.dag_processing.bundles.local import LocalDagBundle from airflow.dag_processing.bundles.manager import ( DagBundlesManager, _guess_best_bundle_for_fileloc, @@ -126,7 +127,11 @@ def path(self) -> Path: class CustomDagBundleProvider(DagBundleProvider): + def __init__(self): + self.configuration_requests = 0 + def get_all_bundle_configurations(self): + self.configuration_requests += 1 return [DagBundleConfiguration(name="active-bundle")] def get_bundle(self, name, version=None, version_data=None): @@ -304,16 +309,21 @@ def test_is_bundle_configured_sees_implicitly_added_bundles(): @conf_vars( { + ("core", "LOAD_EXAMPLES"): "False", ( "dag_processor", "dag_bundle_provider", - ): "unit.dag_processing.bundles.test_dag_bundle_manager.CustomDagBundleProvider" + ): "unit.dag_processing.bundles.test_dag_bundle_manager.CustomDagBundleProvider", } ) def test_custom_bundle_provider_resolves_active_and_retired_bundles(): manager = DagBundlesManager() + provider = manager._bundle_provider + assert isinstance(provider, CustomDagBundleProvider) + assert provider.configuration_requests == 0 assert manager.get_all_bundle_configurations() == (DagBundleConfiguration(name="active-bundle"),) + assert provider.configuration_requests == 1 assert manager.get_bundle_configuration("active-bundle") == DagBundleConfiguration(name="active-bundle") assert manager.get_all_bundle_names() == ["active-bundle"] with pytest.raises(ValueError, match="'unknown-bundle' is not configured"): @@ -350,7 +360,7 @@ def test_custom_bundle_provider_resolves_active_and_retired_bundles(): def test_custom_bundle_provider_rejects_invalid_configurations(provider_class, expected_message): with conf_vars({("dag_processor", "dag_bundle_provider"): provider_class}): with pytest.raises(AirflowConfigException, match=expected_message): - DagBundlesManager() + DagBundlesManager().get_all_bundle_configurations() @pytest.mark.parametrize( @@ -682,20 +692,34 @@ def test_dag_bundle_model_render_url_with_invalid_template(): def test_example_dags_bundle_added(): manager = DagBundlesManager() - manager.refresh_bundle_configurations() - assert "example_dags" in manager._bundle_config + assert "example_dags" in manager.get_all_bundle_names() + assert isinstance(manager.get_bundle("example_dags"), LocalDagBundle) with conf_vars({("core", "LOAD_EXAMPLES"): "False"}): manager = DagBundlesManager() - manager.refresh_bundle_configurations() - assert "example_dags" not in manager._bundle_config + assert "example_dags" not in manager.get_all_bundle_names() + + +@conf_vars( + { + ( + "dag_processor", + "dag_bundle_provider", + ): "unit.dag_processing.bundles.test_dag_bundle_manager.CustomDagBundleProvider" + } +) +def test_example_dags_bundle_added_for_custom_provider(): + manager = DagBundlesManager() + + assert "example_dags" in manager.get_all_bundle_names() + assert isinstance(manager.get_bundle("example_dags"), LocalDagBundle) def test_example_dags_name_is_reserved(): reserved_name_config = [{"name": "example_dags", "classpath": "yo face", "kwargs": {}}] with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(reserved_name_config)}): with pytest.raises(AirflowConfigException, match="Bundle name 'example_dags' is a reserved name."): - DagBundlesManager().refresh_bundle_configurations() + DagBundlesManager().get_all_bundle_configurations() class FailingBundle(BaseDagBundle): From 564fb68b9034cffba01de4ea84cbf0fe3140bb7d Mon Sep 17 00:00:00 2001 From: Sameer Raj Date: Mon, 21 Sep 2026 17:51:04 +0530 Subject: [PATCH 04/34] Clarify the Dag bundle provider contract --- .../src/airflow/cli/commands/team_command.py | 4 +- .../src/airflow/config_templates/config.yml | 2 +- .../airflow/dag_processing/bundles/manager.py | 68 +++++++++---------- .../dag_processing/bundles/provider.py | 16 ++--- .../src/airflow/dag_processing/dagbag.py | 10 +-- .../bundles/test_dag_bundle_manager.py | 50 +++++++------- .../tests/unit/dag_processing/test_dagbag.py | 30 ++++---- 7 files changed, 90 insertions(+), 90 deletions(-) diff --git a/airflow-core/src/airflow/cli/commands/team_command.py b/airflow-core/src/airflow/cli/commands/team_command.py index 70655fd57ad39..f9cf80a960f18 100644 --- a/airflow-core/src/airflow/cli/commands/team_command.py +++ b/airflow-core/src/airflow/cli/commands/team_command.py @@ -218,7 +218,7 @@ def team_sync(args, *, session=NEW_SESSION): dag_bundle_teams = { bundle.team_name - for bundle in DagBundlesManager().get_all_bundle_configurations() + for bundle in DagBundlesManager().get_active_bundle_metadata() if bundle.team_name is not None } existing_teams = Team.get_all_team_names(session=session) @@ -311,7 +311,7 @@ def team_verify(args, *, session=NEW_SESSION): existing_teams = {team.name for team in teams} - for bundle in DagBundlesManager().get_all_bundle_configurations(): + for bundle in DagBundlesManager().get_active_bundle_metadata(): if bundle.team_name and bundle.team_name not in existing_teams: issues.append(f"DAG bundle '{bundle.name}' references unknown team '{bundle.team_name}'.") diff --git a/airflow-core/src/airflow/config_templates/config.yml b/airflow-core/src/airflow/config_templates/config.yml index c52c7fa8c7dfa..2a70a02f0ef2a 100644 --- a/airflow-core/src/airflow/config_templates/config.yml +++ b/airflow-core/src/airflow/config_templates/config.yml @@ -3279,7 +3279,7 @@ dag_processor: dag_bundle_provider: description: | - Import path for the class that provides Dag bundle configurations and constructs Dag bundles. + Import path for the class that provides active Dag bundle metadata and constructs Dag bundles. The class must extend ``airflow.dag_processing.bundles.provider.DagBundleProvider`` and have a no-argument constructor. diff --git a/airflow-core/src/airflow/dag_processing/bundles/manager.py b/airflow-core/src/airflow/dag_processing/bundles/manager.py index 138412bcff467..59a9f45a81eec 100644 --- a/airflow-core/src/airflow/dag_processing/bundles/manager.py +++ b/airflow-core/src/airflow/dag_processing/bundles/manager.py @@ -32,7 +32,7 @@ from airflow.dag_processing.bundles.base import BaseDagBundle # noqa: TC001 from airflow.dag_processing.bundles.local import LocalDagBundle from airflow.dag_processing.bundles.provider import ( - DagBundleConfiguration, + DagBundleMetadata, DagBundleProvider, ) from airflow.exceptions import AirflowConfigException @@ -121,8 +121,8 @@ def _get_configured_bundle_team_names() -> dict[str, str | None]: """Get the team owning each Dag bundle returned by the configured provider.""" manager = DagBundlesManager() return { - configuration.name: configuration.team_name - for configuration in manager._bundle_provider.get_all_bundle_configurations() + metadata.name: metadata.team_name + for metadata in manager._bundle_provider.get_active_bundle_metadata() } @@ -212,37 +212,37 @@ def _load_bundle_provider() -> DagBundleProvider: "`dag_processor` key `dag_bundle_provider`." ) from e - def get_all_bundle_configurations(self) -> tuple[DagBundleConfiguration, ...]: - """Get all active Dag bundle configurations.""" - bundle_configurations: dict[str, DagBundleConfiguration] = {} - for bundle_config in self._bundle_provider.get_all_bundle_configurations(): - if not isinstance(bundle_config, DagBundleConfiguration): + def get_active_bundle_metadata(self) -> tuple[DagBundleMetadata, ...]: + """Get metadata for all active Dag bundles.""" + bundle_metadata: dict[str, DagBundleMetadata] = {} + for metadata in self._bundle_provider.get_active_bundle_metadata(): + if not isinstance(metadata, DagBundleMetadata): raise AirflowConfigException( - "Dag bundle providers must return DagBundleConfiguration objects from " - "get_all_bundle_configurations()." + "Dag bundle providers must return DagBundleMetadata objects from " + "get_active_bundle_metadata()." ) - if bundle_config.name == _example_dag_bundle_name: + if metadata.name == _example_dag_bundle_name: raise AirflowConfigException( f"Bundle name '{_example_dag_bundle_name}' is a reserved name. Please choose another name for your bundle." " Example Dags can be enabled with the '[core] load_examples' config." ) - if bundle_config.name in bundle_configurations: + if metadata.name in bundle_metadata: raise AirflowConfigException( - f"Dag bundle provider returned duplicate bundle name '{bundle_config.name}'." + f"Dag bundle provider returned duplicate bundle name '{metadata.name}'." ) - if bundle_config.team_name and not conf.getboolean("core", "multi_team"): + if metadata.team_name and not conf.getboolean("core", "multi_team"): raise AirflowConfigException( - "Dag bundle configurations cannot have a team name when multi-team mode is disabled. " + "Dag bundle metadata cannot have a team name when multi-team mode is disabled. " "To enable multi-team, update section `core` key `multi_team` in your config." ) - bundle_configurations[bundle_config.name] = bundle_config + bundle_metadata[metadata.name] = metadata for name in self._example_dag_bundle_paths: - if name in bundle_configurations: + if name in bundle_metadata: raise AirflowConfigException(f"Bundle name '{name}' is reserved for example Dags.") - bundle_configurations[name] = DagBundleConfiguration(name=name) + bundle_metadata[name] = DagBundleMetadata(name=name) - return tuple(bundle_configurations.values()) + return tuple(bundle_metadata.values()) @provide_session def sync_bundles_to_db(self, *, deactivate_missing: bool = True, session: Session = NEW_SESSION) -> None: @@ -265,7 +265,7 @@ def sync_bundles_to_db(self, *, deactivate_missing: bool = True, session: Sessio """ self.log.debug("Syncing DAG bundles to the database") - bundle_configurations = {config.name: config for config in self.get_all_bundle_configurations()} + bundle_metadata = {metadata.name: metadata for metadata in self.get_active_bundle_metadata()} def _extract_and_sign_template(bundle_name: str) -> tuple[str | None, dict]: bundle_instance = self.get_bundle(name) @@ -291,13 +291,13 @@ def _extract_and_sign_template(bundle_name: str) -> tuple[str | None, dict]: for bundle in stored.values() } - for name, config in bundle_configurations.items(): + for name, metadata in bundle_metadata.items(): team: Team | None = None - if config.team_name: - team = session.scalars(select(Team).where(Team.name == config.team_name)).one_or_none() + if metadata.team_name: + team = session.scalars(select(Team).where(Team.name == metadata.team_name)).one_or_none() if not team: raise AirflowConfigException( - f"Team '{config.team_name}' configured for Dag bundle '{name}' does not exist." + f"Team '{metadata.team_name}' configured for Dag bundle '{name}' does not exist." ) try: @@ -322,7 +322,7 @@ def _extract_and_sign_template(bundle_name: str) -> tuple[str | None, dict]: session.add(bundle) self.log.info("Added new DAG bundle %s to the database", name) - if team and bundle_to_team.get(name) != config.team_name: + if team and bundle_to_team.get(name) != metadata.team_name: # Change of team. It can be associating a team to a dag bundle that did not have one or # swapping a team for another bundle.teams = [team] @@ -595,11 +595,11 @@ def get_bundle( return LocalDagBundle(name=name, path=path, version=version, version_data=version_data) return self._bundle_provider.get_bundle(name=name, version=version, version_data=version_data) - def get_bundle_configuration(self, name: str) -> DagBundleConfiguration: - """Get the active configuration for a Dag bundle.""" - for configuration in self.get_all_bundle_configurations(): - if configuration.name == name: - return configuration + def get_bundle_metadata(self, name: str) -> DagBundleMetadata: + """Get metadata for an active Dag bundle.""" + for metadata in self.get_active_bundle_metadata(): + if metadata.name == name: + return metadata raise ValueError(f"Requested bundle '{name}' is not configured.") @classmethod @@ -618,11 +618,11 @@ def get_all_dag_bundles(self) -> Iterable[BaseDagBundle]: :return: list of DAG bundles. """ - for configuration in self.get_all_bundle_configurations(): + for metadata in self.get_active_bundle_metadata(): try: - yield self.get_bundle(name=configuration.name) + yield self.get_bundle(name=metadata.name) except Exception as e: - self.log.exception("Error creating bundle '%s': %s", configuration.name, e) + self.log.exception("Error creating bundle '%s': %s", metadata.name, e) # Skip this bundle and continue with others continue @@ -632,7 +632,7 @@ def get_all_bundle_names(self) -> Iterable[str]: :return: sorted list of bundle names. """ - return sorted(configuration.name for configuration in self.get_all_bundle_configurations()) + return sorted(metadata.name for metadata in self.get_active_bundle_metadata()) def view_url(self, name: str, version: str | None = None) -> str | None: warnings.warn( diff --git a/airflow-core/src/airflow/dag_processing/bundles/provider.py b/airflow-core/src/airflow/dag_processing/bundles/provider.py index ff497edd6944b..3d8552aad533a 100644 --- a/airflow-core/src/airflow/dag_processing/bundles/provider.py +++ b/airflow-core/src/airflow/dag_processing/bundles/provider.py @@ -32,8 +32,8 @@ @dataclass(frozen=True) -class DagBundleConfiguration: - """Configuration used by Airflow to manage a Dag bundle.""" +class DagBundleMetadata: + """Metadata used by Airflow to manage a Dag bundle.""" name: str team_name: str | None = None @@ -41,17 +41,17 @@ class DagBundleConfiguration: class DagBundleProvider(ABC): """ - Provide Dag bundle configurations and Dag bundles. + Provide active Dag bundle metadata and construct Dag bundles. - The configuration list is the complete set of active bundles. A bundle name + The metadata list is the complete set of active bundles. A bundle name may remain resolvable through ``get_bundle`` after it leaves that list because retained Dag runs can still need an older version. A name must continue to identify the same bundle implementation; use a new name for a different one. """ @abstractmethod - def get_all_bundle_configurations(self) -> Sequence[DagBundleConfiguration]: - """Return the complete set of active Dag bundle configurations.""" + def get_active_bundle_metadata(self) -> Sequence[DagBundleMetadata]: + """Return metadata for the complete set of active Dag bundles.""" @abstractmethod def get_bundle( @@ -149,9 +149,9 @@ def _load_configured_bundles(self) -> None: team_name=bundle_config.team_name, ) - def get_all_bundle_configurations(self) -> Sequence[DagBundleConfiguration]: + def get_active_bundle_metadata(self) -> Sequence[DagBundleMetadata]: return [ - DagBundleConfiguration(name=name, team_name=config.team_name) + DagBundleMetadata(name=name, team_name=config.team_name) for name, config in self._bundle_config.items() ] diff --git a/airflow-core/src/airflow/dag_processing/dagbag.py b/airflow-core/src/airflow/dag_processing/dagbag.py index 595936b519bb0..b8968c54267de 100644 --- a/airflow-core/src/airflow/dag_processing/dagbag.py +++ b/airflow-core/src/airflow/dag_processing/dagbag.py @@ -130,14 +130,14 @@ def _validate_executor_fields(dag: DAG | SerializedDAG, bundle_name: str | None # Check if multi team is available by reading the multi_team configuration (which is boolean) if conf.getboolean("core", "multi_team"): - # Get team name from bundle configuration if available + # Get team name from bundle metadata if available if bundle_name: from airflow.dag_processing.bundles.manager import DagBundlesManager bundle_manager = DagBundlesManager() - bundle_config = bundle_manager.get_bundle_configuration(bundle_name) + bundle_metadata = bundle_manager.get_bundle_metadata(bundle_name) - dag_team_name = bundle_config.team_name + dag_team_name = bundle_metadata.team_name if dag_team_name: log.debug( "Found team '%s' for DAG '%s' via bundle '%s'", dag_team_name, dag.dag_id, bundle_name @@ -175,9 +175,9 @@ def _assign_default_team_pools( from airflow.dag_processing.bundles.manager import DagBundlesManager bundle_manager = DagBundlesManager() - bundle_config = bundle_manager.get_bundle_configuration(bundle_name) + bundle_metadata = bundle_manager.get_bundle_metadata(bundle_name) - dag_team_name = bundle_config.team_name + dag_team_name = bundle_metadata.team_name if not dag_team_name: return 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 0a5d09028fa3f..0252af6b7f6b7 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 @@ -37,7 +37,7 @@ get_configured_bundle_names, get_configured_bundle_team_names, ) -from airflow.dag_processing.bundles.provider import DagBundleConfiguration, DagBundleProvider +from airflow.dag_processing.bundles.provider import DagBundleMetadata, DagBundleProvider from airflow.exceptions import AirflowConfigException from airflow.models.dag import DagModel from airflow.models.dag_version import DagVersion @@ -128,11 +128,11 @@ def path(self) -> Path: class CustomDagBundleProvider(DagBundleProvider): def __init__(self): - self.configuration_requests = 0 + self.metadata_requests = 0 - def get_all_bundle_configurations(self): - self.configuration_requests += 1 - return [DagBundleConfiguration(name="active-bundle")] + def get_active_bundle_metadata(self): + self.metadata_requests += 1 + return [DagBundleMetadata(name="active-bundle")] def get_bundle(self, name, version=None, version_data=None): if name not in {"active-bundle", "retired-bundle"}: @@ -145,22 +145,22 @@ def get_bundle(self, name, version=None, version_data=None): ) -class DuplicateConfigurationDagBundleProvider(CustomDagBundleProvider): - def get_all_bundle_configurations(self): +class DuplicateMetadataDagBundleProvider(CustomDagBundleProvider): + def get_active_bundle_metadata(self): return [ - DagBundleConfiguration(name="duplicate"), - DagBundleConfiguration(name="duplicate"), + DagBundleMetadata(name="duplicate"), + DagBundleMetadata(name="duplicate"), ] -class NonConfigurationDagBundleProvider(CustomDagBundleProvider): - def get_all_bundle_configurations(self): +class NonMetadataDagBundleProvider(CustomDagBundleProvider): + def get_active_bundle_metadata(self): return [object()] class MissingTeamDagBundleProvider(CustomDagBundleProvider): - def get_all_bundle_configurations(self): - return [DagBundleConfiguration(name="active-bundle", team_name="missing-team")] + def get_active_bundle_metadata(self): + return [DagBundleMetadata(name="active-bundle", team_name="missing-team")] class FailingInitializationDagBundleProvider(CustomDagBundleProvider): @@ -320,14 +320,14 @@ def test_custom_bundle_provider_resolves_active_and_retired_bundles(): manager = DagBundlesManager() provider = manager._bundle_provider assert isinstance(provider, CustomDagBundleProvider) - assert provider.configuration_requests == 0 + assert provider.metadata_requests == 0 - assert manager.get_all_bundle_configurations() == (DagBundleConfiguration(name="active-bundle"),) - assert provider.configuration_requests == 1 - assert manager.get_bundle_configuration("active-bundle") == DagBundleConfiguration(name="active-bundle") + assert manager.get_active_bundle_metadata() == (DagBundleMetadata(name="active-bundle"),) + assert provider.metadata_requests == 1 + assert manager.get_bundle_metadata("active-bundle") == DagBundleMetadata(name="active-bundle") assert manager.get_all_bundle_names() == ["active-bundle"] with pytest.raises(ValueError, match="'unknown-bundle' is not configured"): - manager.get_bundle_configuration("unknown-bundle") + manager.get_bundle_metadata("unknown-bundle") active_bundle = manager.get_bundle("active-bundle") retired_bundle = manager.get_bundle( @@ -346,21 +346,21 @@ def test_custom_bundle_provider_resolves_active_and_retired_bundles(): ("provider_class", "expected_message"), [ pytest.param( - "unit.dag_processing.bundles.test_dag_bundle_manager.DuplicateConfigurationDagBundleProvider", + "unit.dag_processing.bundles.test_dag_bundle_manager.DuplicateMetadataDagBundleProvider", "duplicate bundle name 'duplicate'", id="duplicate-name", ), pytest.param( - "unit.dag_processing.bundles.test_dag_bundle_manager.NonConfigurationDagBundleProvider", - "must return DagBundleConfiguration objects", - id="wrong-configuration-type", + "unit.dag_processing.bundles.test_dag_bundle_manager.NonMetadataDagBundleProvider", + "must return DagBundleMetadata objects", + id="wrong-metadata-type", ), ], ) -def test_custom_bundle_provider_rejects_invalid_configurations(provider_class, expected_message): +def test_custom_bundle_provider_rejects_invalid_metadata(provider_class, expected_message): with conf_vars({("dag_processor", "dag_bundle_provider"): provider_class}): with pytest.raises(AirflowConfigException, match=expected_message): - DagBundlesManager().get_all_bundle_configurations() + DagBundlesManager().get_active_bundle_metadata() @pytest.mark.parametrize( @@ -719,7 +719,7 @@ def test_example_dags_name_is_reserved(): reserved_name_config = [{"name": "example_dags", "classpath": "yo face", "kwargs": {}}] with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(reserved_name_config)}): with pytest.raises(AirflowConfigException, match="Bundle name 'example_dags' is a reserved name."): - DagBundlesManager().get_all_bundle_configurations() + DagBundlesManager().get_active_bundle_metadata() class FailingBundle(BaseDagBundle): diff --git a/airflow-core/tests/unit/dag_processing/test_dagbag.py b/airflow-core/tests/unit/dag_processing/test_dagbag.py index 97a84c0220c02..4b776ff14b0f3 100644 --- a/airflow-core/tests/unit/dag_processing/test_dagbag.py +++ b/airflow-core/tests/unit/dag_processing/test_dagbag.py @@ -36,7 +36,7 @@ from sqlalchemy import select from airflow import settings -from airflow.dag_processing.bundles.provider import DagBundleConfiguration +from airflow.dag_processing.bundles.provider import DagBundleMetadata from airflow.dag_processing.dagbag import ( BundleDagBag, DagBag, @@ -136,10 +136,10 @@ def test_multi_team_disabled_ignores_bundle_name(self, mock_lookup): def test_multi_team_enabled_bundle_exists_with_team(self, mock_lookup, mock_manager_class): """Test successful team lookup when bundle exists and has team_name.""" # Setup mock bundle manager - bundle_config = DagBundleConfiguration(name="test_bundle", team_name="test_team") + bundle_metadata = DagBundleMetadata(name="test_bundle", team_name="test_team") mock_manager = mock_manager_class.return_value - mock_manager.get_bundle_configuration.return_value = bundle_config + mock_manager.get_bundle_metadata.return_value = bundle_metadata with DAG("test-dag", schedule=None) as dag: BaseOperator(task_id="t1", executor="team.executor") @@ -154,10 +154,10 @@ def test_multi_team_enabled_bundle_exists_with_team(self, mock_lookup, mock_mana @patch.object(ExecutorLoader, "lookup_executor_name_by_str") def test_multi_team_enabled_bundle_exists_no_team(self, mock_lookup, mock_manager_class): """Test when bundle exists but has no team_name (None or empty).""" - bundle_config = DagBundleConfiguration(name="test_bundle") + bundle_metadata = DagBundleMetadata(name="test_bundle") mock_manager = mock_manager_class.return_value - mock_manager.get_bundle_configuration.return_value = bundle_config + mock_manager.get_bundle_metadata.return_value = bundle_metadata with DAG("test-dag", schedule=None) as dag: BaseOperator(task_id="t1", executor="test.executor") @@ -187,10 +187,10 @@ def test_multiple_tasks_with_executors(self, mock_lookup): @patch.object(ExecutorLoader, "lookup_executor_name_by_str") def test_executor_validation_failure_with_team(self, mock_lookup, mock_manager_class): """Test executor validation failure when team is associated (team-specific error).""" - bundle_config = DagBundleConfiguration(name="test_bundle", team_name="test_team") + bundle_metadata = DagBundleMetadata(name="test_bundle", team_name="test_team") mock_manager = mock_manager_class.return_value - mock_manager.get_bundle_configuration.return_value = bundle_config + mock_manager.get_bundle_metadata.return_value = bundle_metadata # ExecutorLoader raises exception mock_lookup.side_effect = UnknownExecutorException("Executor not found") @@ -234,10 +234,10 @@ def test_executor_validation_failure_no_team(self, mock_lookup): @patch.object(ExecutorLoader, "lookup_executor_name_by_str") def test_global_executor_fallback_success(self, mock_lookup, mock_manager_class): """Test that team-specific executor failure falls back to global executor successfully.""" - bundle_config = DagBundleConfiguration(name="test_bundle", team_name="test_team") + bundle_metadata = DagBundleMetadata(name="test_bundle", team_name="test_team") mock_manager = mock_manager_class.return_value - mock_manager.get_bundle_configuration.return_value = bundle_config + mock_manager.get_bundle_metadata.return_value = bundle_metadata # First call (team-specific) fails, second call (global) succeeds mock_lookup.side_effect = [UnknownExecutorException("Team executor not found"), None] @@ -258,10 +258,10 @@ def test_global_executor_fallback_success(self, mock_lookup, mock_manager_class) @patch.object(ExecutorLoader, "lookup_executor_name_by_str") def test_global_executor_fallback_failure(self, mock_lookup, mock_manager_class): """Test that when both team-specific and global executors fail, appropriate error is raised.""" - bundle_config = DagBundleConfiguration(name="test_bundle", team_name="test_team") + bundle_metadata = DagBundleMetadata(name="test_bundle", team_name="test_team") mock_manager = mock_manager_class.return_value - mock_manager.get_bundle_configuration.return_value = bundle_config + mock_manager.get_bundle_metadata.return_value = bundle_metadata # Both calls fail mock_lookup.side_effect = UnknownExecutorException("Executor not found") @@ -291,10 +291,10 @@ def test_global_executor_fallback_failure(self, mock_lookup, mock_manager_class) @patch.object(ExecutorLoader, "lookup_executor_name_by_str") def test_team_specific_executor_success_no_fallback(self, mock_lookup, mock_manager_class): """Test that when team-specific executor succeeds, global fallback is not attempted.""" - bundle_config = DagBundleConfiguration(name="test_bundle", team_name="test_team") + bundle_metadata = DagBundleMetadata(name="test_bundle", team_name="test_team") mock_manager = mock_manager_class.return_value - mock_manager.get_bundle_configuration.return_value = bundle_config + mock_manager.get_bundle_metadata.return_value = bundle_metadata # First call (team-specific) succeeds mock_lookup.return_value = None @@ -465,8 +465,8 @@ def test_default_pool_replaced_with_team_pool( operator_args, expected_pool, ): - bundle_config = DagBundleConfiguration(name="test_bundle", team_name=team_name) - mock_manager.return_value.get_bundle_configuration.return_value = bundle_config + bundle_metadata = DagBundleMetadata(name="test_bundle", team_name=team_name) + mock_manager.return_value.get_bundle_metadata.return_value = bundle_metadata dag_file = tmp_path / "test_dag.py" dag_file.write_text( From ee27d99c3257e34520b2a0db3e546ebd2ab01b56 Mon Sep 17 00:00:00 2001 From: Sameer Raj Date: Mon, 21 Sep 2026 18:42:19 +0530 Subject: [PATCH 05/34] Clarify Dag bundle metadata ownership --- .../src/airflow/api/common/airflow_health.py | 9 ++++--- .../airflow/dag_processing/bundles/manager.py | 19 +++++++-------- .../dag_processing/bundles/provider.py | 6 ----- .../bundles/test_dag_bundle_manager.py | 24 +++++++++++++------ 4 files changed, 31 insertions(+), 27 deletions(-) diff --git a/airflow-core/src/airflow/api/common/airflow_health.py b/airflow-core/src/airflow/api/common/airflow_health.py index 728fba5a2a01e..903ec3ef8dda7 100644 --- a/airflow-core/src/airflow/api/common/airflow_health.py +++ b/airflow-core/src/airflow/api/common/airflow_health.py @@ -23,7 +23,7 @@ from sqlalchemy import select from airflow.configuration import conf -from airflow.dag_processing.bundles.manager import get_configured_bundle_team_names +from airflow.dag_processing.bundles.manager import DagBundlesManager from airflow.jobs.dag_processor_job_runner import DagProcessorJobRunner from airflow.jobs.job import Job from airflow.jobs.scheduler_job_runner import SchedulerJobRunner @@ -112,7 +112,7 @@ def _dag_processor_instance_health(job: Job) -> dict[str, Any]: # behind forever and a restarted replica adds a second one. The denominator is therefore taken from # the declared work partition instead, which is unaffected by how replicas come and go: # -# * Dag processor -- the bundles in ``[dag_processor] dag_bundle_config_list``. +# * Dag processor -- the bundles returned by the configured Dag bundle provider. # * Triggerer -- the team scopes those bundles declare, since a triggerer only picks up triggers for # its own team (see ``Trigger.ids_for_triggerer``). # * Scheduler -- schedulers are symmetric, so there is no partition and no partial state to report. @@ -121,7 +121,10 @@ def _dag_processor_instance_health(job: Job) -> dict[str, Any]: def _configured_bundle_teams() -> dict[str, str | None]: """Map every configured Dag bundle to the team owning it, empty when the config is unreadable.""" try: - return get_configured_bundle_team_names() + return { + metadata.name: metadata.team_name + for metadata in DagBundlesManager().get_configured_bundle_metadata() + } except Exception: # A health probe must not fail on malformed bundle config; callers fall back to liveness only. log.warning("Could not read the Dag bundle configuration", exc_info=True) diff --git a/airflow-core/src/airflow/dag_processing/bundles/manager.py b/airflow-core/src/airflow/dag_processing/bundles/manager.py index 59a9f45a81eec..203cda66dad97 100644 --- a/airflow-core/src/airflow/dag_processing/bundles/manager.py +++ b/airflow-core/src/airflow/dag_processing/bundles/manager.py @@ -117,15 +117,6 @@ def _guess_best_bundle_for_fileloc( return None -def _get_configured_bundle_team_names() -> dict[str, str | None]: - """Get the team owning each Dag bundle returned by the configured provider.""" - manager = DagBundlesManager() - return { - metadata.name: metadata.team_name - for metadata in manager._bundle_provider.get_active_bundle_metadata() - } - - def _is_safe_bundle_url(url: str) -> bool: """ Check if a bundle URL is safe to use. @@ -212,8 +203,8 @@ def _load_bundle_provider() -> DagBundleProvider: "`dag_processor` key `dag_bundle_provider`." ) from e - def get_active_bundle_metadata(self) -> tuple[DagBundleMetadata, ...]: - """Get metadata for all active Dag bundles.""" + def get_configured_bundle_metadata(self) -> tuple[DagBundleMetadata, ...]: + """Get validated metadata returned by the configured Dag bundle provider.""" bundle_metadata: dict[str, DagBundleMetadata] = {} for metadata in self._bundle_provider.get_active_bundle_metadata(): if not isinstance(metadata, DagBundleMetadata): @@ -237,6 +228,12 @@ def get_active_bundle_metadata(self) -> tuple[DagBundleMetadata, ...]: ) bundle_metadata[metadata.name] = metadata + return tuple(bundle_metadata.values()) + + def get_active_bundle_metadata(self) -> tuple[DagBundleMetadata, ...]: + """Get metadata for all active Dag bundles, including example bundles.""" + bundle_metadata = {metadata.name: metadata for metadata in self.get_configured_bundle_metadata()} + for name in self._example_dag_bundle_paths: if name in bundle_metadata: raise AirflowConfigException(f"Bundle name '{name}' is reserved for example Dags.") diff --git a/airflow-core/src/airflow/dag_processing/bundles/provider.py b/airflow-core/src/airflow/dag_processing/bundles/provider.py index 3d8552aad533a..410a221f40450 100644 --- a/airflow-core/src/airflow/dag_processing/bundles/provider.py +++ b/airflow-core/src/airflow/dag_processing/bundles/provider.py @@ -102,12 +102,6 @@ def _parse_bundle_config(config_list) -> list[_ExternalBundleConfig]: except ValidationError as e: raise _bundle_item_exc(f"Item {item} failed validation: {e}") - if cfg.name == "example_dags": - raise AirflowConfigException( - "Bundle name 'example_dags' is a reserved name. Please choose another name for your bundle." - " Example Dags can be enabled with the '[core] load_examples' config." - ) - bundles[cfg.name] = cfg if len(bundles.keys()) != len(config_list): raise _bundle_item_exc("One or more bundle names appeared multiple times") 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 0252af6b7f6b7..a4caf0aff1d70 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 @@ -35,7 +35,6 @@ _guess_best_bundle_for_fileloc, _load_bundle_config_snapshot, get_configured_bundle_names, - get_configured_bundle_team_names, ) from airflow.dag_processing.bundles.provider import DagBundleMetadata, DagBundleProvider from airflow.exceptions import AirflowConfigException @@ -207,7 +206,7 @@ def __init__(self): @pytest.mark.parametrize("load_examples", ["False", "True"]) -def test_get_configured_bundle_team_names(load_examples): +def test_get_configured_bundle_metadata(load_examples): with conf_vars( { ("core", "load_examples"): load_examples, @@ -215,7 +214,10 @@ def test_get_configured_bundle_team_names(load_examples): ("dag_processor", "dag_bundle_config_list"): json.dumps(TEAM_BUNDLE_CONFIG), } ): - assert get_configured_bundle_team_names() == { + assert { + metadata.name: metadata.team_name + for metadata in DagBundlesManager().get_configured_bundle_metadata() + } == { "team-bundle": "team-a", "unscoped-bundle": None, } @@ -236,8 +238,8 @@ def test_get_configured_bundle_names(): @conf_vars({("dag_processor", "dag_bundle_config_list"): "[]"}) -def test_get_configured_bundle_team_names_without_config(): - assert get_configured_bundle_team_names() == {} +def test_get_configured_bundle_metadata_without_config(): + assert DagBundlesManager().get_configured_bundle_metadata() == () def test_get_bundle(): @@ -322,8 +324,10 @@ def test_custom_bundle_provider_resolves_active_and_retired_bundles(): assert isinstance(provider, CustomDagBundleProvider) assert provider.metadata_requests == 0 - assert manager.get_active_bundle_metadata() == (DagBundleMetadata(name="active-bundle"),) + assert manager.get_configured_bundle_metadata() == (DagBundleMetadata(name="active-bundle"),) assert provider.metadata_requests == 1 + assert manager.get_active_bundle_metadata() == (DagBundleMetadata(name="active-bundle"),) + assert provider.metadata_requests == 2 assert manager.get_bundle_metadata("active-bundle") == DagBundleMetadata(name="active-bundle") assert manager.get_all_bundle_names() == ["active-bundle"] with pytest.raises(ValueError, match="'unknown-bundle' is not configured"): @@ -716,7 +720,13 @@ def test_example_dags_bundle_added_for_custom_provider(): def test_example_dags_name_is_reserved(): - reserved_name_config = [{"name": "example_dags", "classpath": "yo face", "kwargs": {}}] + reserved_name_config = [ + { + "name": "example_dags", + "classpath": "unit.dag_processing.bundles.test_dag_bundle_manager.BasicBundle", + "kwargs": {"refresh_interval": 1}, + } + ] with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(reserved_name_config)}): with pytest.raises(AirflowConfigException, match="Bundle name 'example_dags' is a reserved name."): DagBundlesManager().get_active_bundle_metadata() From a6c89d1586cc310dc972683f9be3b52d5117f697 Mon Sep 17 00:00:00 2001 From: Sameer Raj Date: Mon, 21 Sep 2026 20:32:56 +0530 Subject: [PATCH 06/34] Clarify configured Dag bundle metadata --- airflow-core/src/airflow/config_templates/config.yml | 4 ++-- .../src/airflow/dag_processing/bundles/manager.py | 5 +++-- .../src/airflow/dag_processing/bundles/provider.py | 10 +++++----- .../dag_processing/bundles/test_dag_bundle_manager.py | 8 ++++---- 4 files changed, 14 insertions(+), 13 deletions(-) diff --git a/airflow-core/src/airflow/config_templates/config.yml b/airflow-core/src/airflow/config_templates/config.yml index 2a70a02f0ef2a..978b12ef0511d 100644 --- a/airflow-core/src/airflow/config_templates/config.yml +++ b/airflow-core/src/airflow/config_templates/config.yml @@ -3279,12 +3279,12 @@ dag_processor: dag_bundle_provider: description: | - Import path for the class that provides active Dag bundle metadata and constructs Dag bundles. + Import path for the class that provides configured Dag bundle metadata and constructs Dag bundles. The class must extend ``airflow.dag_processing.bundles.provider.DagBundleProvider`` and have a no-argument constructor. The default implementation reads ``dag_bundle_config_list``. A custom implementation can read - bundle configuration from another source and must return the complete set of active bundles. + bundle configuration from another source and must return the complete set of configured bundles. version_added: 3.4.0 type: string example: ~ diff --git a/airflow-core/src/airflow/dag_processing/bundles/manager.py b/airflow-core/src/airflow/dag_processing/bundles/manager.py index 203cda66dad97..ad99265187524 100644 --- a/airflow-core/src/airflow/dag_processing/bundles/manager.py +++ b/airflow-core/src/airflow/dag_processing/bundles/manager.py @@ -206,11 +206,11 @@ def _load_bundle_provider() -> DagBundleProvider: def get_configured_bundle_metadata(self) -> tuple[DagBundleMetadata, ...]: """Get validated metadata returned by the configured Dag bundle provider.""" bundle_metadata: dict[str, DagBundleMetadata] = {} - for metadata in self._bundle_provider.get_active_bundle_metadata(): + for metadata in self._bundle_provider.get_configured_bundle_metadata(): if not isinstance(metadata, DagBundleMetadata): raise AirflowConfigException( "Dag bundle providers must return DagBundleMetadata objects from " - "get_active_bundle_metadata()." + "get_configured_bundle_metadata()." ) if metadata.name == _example_dag_bundle_name: raise AirflowConfigException( @@ -234,6 +234,7 @@ def get_active_bundle_metadata(self) -> tuple[DagBundleMetadata, ...]: """Get metadata for all active Dag bundles, including example bundles.""" bundle_metadata = {metadata.name: metadata for metadata in self.get_configured_bundle_metadata()} + # Airflow owns example Dag bundles; configured providers do not need to include them. for name in self._example_dag_bundle_paths: if name in bundle_metadata: raise AirflowConfigException(f"Bundle name '{name}' is reserved for example Dags.") diff --git a/airflow-core/src/airflow/dag_processing/bundles/provider.py b/airflow-core/src/airflow/dag_processing/bundles/provider.py index 410a221f40450..c676a9106a5f8 100644 --- a/airflow-core/src/airflow/dag_processing/bundles/provider.py +++ b/airflow-core/src/airflow/dag_processing/bundles/provider.py @@ -41,17 +41,17 @@ class DagBundleMetadata: class DagBundleProvider(ABC): """ - Provide active Dag bundle metadata and construct Dag bundles. + Provide configured Dag bundle metadata and construct Dag bundles. - The metadata list is the complete set of active bundles. A bundle name + The metadata list is the complete set of configured bundles. A bundle name may remain resolvable through ``get_bundle`` after it leaves that list because retained Dag runs can still need an older version. A name must continue to identify the same bundle implementation; use a new name for a different one. """ @abstractmethod - def get_active_bundle_metadata(self) -> Sequence[DagBundleMetadata]: - """Return metadata for the complete set of active Dag bundles.""" + def get_configured_bundle_metadata(self) -> Sequence[DagBundleMetadata]: + """Return metadata for the complete set of configured Dag bundles.""" @abstractmethod def get_bundle( @@ -143,7 +143,7 @@ def _load_configured_bundles(self) -> None: team_name=bundle_config.team_name, ) - def get_active_bundle_metadata(self) -> Sequence[DagBundleMetadata]: + def get_configured_bundle_metadata(self) -> Sequence[DagBundleMetadata]: return [ DagBundleMetadata(name=name, team_name=config.team_name) for name, config in self._bundle_config.items() 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 a4caf0aff1d70..1fde18ef97d78 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 @@ -129,7 +129,7 @@ class CustomDagBundleProvider(DagBundleProvider): def __init__(self): self.metadata_requests = 0 - def get_active_bundle_metadata(self): + def get_configured_bundle_metadata(self): self.metadata_requests += 1 return [DagBundleMetadata(name="active-bundle")] @@ -145,7 +145,7 @@ def get_bundle(self, name, version=None, version_data=None): class DuplicateMetadataDagBundleProvider(CustomDagBundleProvider): - def get_active_bundle_metadata(self): + def get_configured_bundle_metadata(self): return [ DagBundleMetadata(name="duplicate"), DagBundleMetadata(name="duplicate"), @@ -153,12 +153,12 @@ def get_active_bundle_metadata(self): class NonMetadataDagBundleProvider(CustomDagBundleProvider): - def get_active_bundle_metadata(self): + def get_configured_bundle_metadata(self): return [object()] class MissingTeamDagBundleProvider(CustomDagBundleProvider): - def get_active_bundle_metadata(self): + def get_configured_bundle_metadata(self): return [DagBundleMetadata(name="active-bundle", team_name="missing-team")] From 5f87ffc21764f6a1053c9e0413331208cfc3e682 Mon Sep 17 00:00:00 2001 From: Sameer Raj Date: Mon, 21 Sep 2026 21:02:55 +0530 Subject: [PATCH 07/34] Simplify Dag bundle manager metadata access --- .../src/airflow/api/common/airflow_health.py | 5 +---- .../src/airflow/dag_processing/bundles/manager.py | 8 ++++++-- .../bundles/test_dag_bundle_manager.py | 15 +++++---------- 3 files changed, 12 insertions(+), 16 deletions(-) diff --git a/airflow-core/src/airflow/api/common/airflow_health.py b/airflow-core/src/airflow/api/common/airflow_health.py index 903ec3ef8dda7..f7b86a1d2c44e 100644 --- a/airflow-core/src/airflow/api/common/airflow_health.py +++ b/airflow-core/src/airflow/api/common/airflow_health.py @@ -121,10 +121,7 @@ def _dag_processor_instance_health(job: Job) -> dict[str, Any]: def _configured_bundle_teams() -> dict[str, str | None]: """Map every configured Dag bundle to the team owning it, empty when the config is unreadable.""" try: - return { - metadata.name: metadata.team_name - for metadata in DagBundlesManager().get_configured_bundle_metadata() - } + return DagBundlesManager().get_declared_bundle_teams() except Exception: # A health probe must not fail on malformed bundle config; callers fall back to liveness only. log.warning("Could not read the Dag bundle configuration", exc_info=True) diff --git a/airflow-core/src/airflow/dag_processing/bundles/manager.py b/airflow-core/src/airflow/dag_processing/bundles/manager.py index ad99265187524..70bc24b8a71ef 100644 --- a/airflow-core/src/airflow/dag_processing/bundles/manager.py +++ b/airflow-core/src/airflow/dag_processing/bundles/manager.py @@ -203,7 +203,7 @@ def _load_bundle_provider() -> DagBundleProvider: "`dag_processor` key `dag_bundle_provider`." ) from e - def get_configured_bundle_metadata(self) -> tuple[DagBundleMetadata, ...]: + def _get_provider_bundle_metadata(self) -> tuple[DagBundleMetadata, ...]: """Get validated metadata returned by the configured Dag bundle provider.""" bundle_metadata: dict[str, DagBundleMetadata] = {} for metadata in self._bundle_provider.get_configured_bundle_metadata(): @@ -232,7 +232,7 @@ def get_configured_bundle_metadata(self) -> tuple[DagBundleMetadata, ...]: def get_active_bundle_metadata(self) -> tuple[DagBundleMetadata, ...]: """Get metadata for all active Dag bundles, including example bundles.""" - bundle_metadata = {metadata.name: metadata for metadata in self.get_configured_bundle_metadata()} + bundle_metadata = {metadata.name: metadata for metadata in self._get_provider_bundle_metadata()} # Airflow owns example Dag bundles; configured providers do not need to include them. for name in self._example_dag_bundle_paths: @@ -242,6 +242,10 @@ def get_active_bundle_metadata(self) -> tuple[DagBundleMetadata, ...]: return tuple(bundle_metadata.values()) + def get_declared_bundle_teams(self) -> dict[str, str | None]: + """Get the team owning each bundle declared by the configured provider.""" + return {metadata.name: metadata.team_name for metadata in self._get_provider_bundle_metadata()} + @provide_session def sync_bundles_to_db(self, *, deactivate_missing: bool = True, session: Session = NEW_SESSION) -> None: """ 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 1fde18ef97d78..4221ccee3ac33 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 @@ -206,7 +206,7 @@ def __init__(self): @pytest.mark.parametrize("load_examples", ["False", "True"]) -def test_get_configured_bundle_metadata(load_examples): +def test_get_declared_bundle_teams_excludes_example_bundles(load_examples): with conf_vars( { ("core", "load_examples"): load_examples, @@ -214,10 +214,7 @@ def test_get_configured_bundle_metadata(load_examples): ("dag_processor", "dag_bundle_config_list"): json.dumps(TEAM_BUNDLE_CONFIG), } ): - assert { - metadata.name: metadata.team_name - for metadata in DagBundlesManager().get_configured_bundle_metadata() - } == { + assert DagBundlesManager().get_declared_bundle_teams() == { "team-bundle": "team-a", "unscoped-bundle": None, } @@ -238,8 +235,8 @@ def test_get_configured_bundle_names(): @conf_vars({("dag_processor", "dag_bundle_config_list"): "[]"}) -def test_get_configured_bundle_metadata_without_config(): - assert DagBundlesManager().get_configured_bundle_metadata() == () +def test_get_declared_bundle_teams_without_config(): + assert DagBundlesManager().get_declared_bundle_teams() == {} def test_get_bundle(): @@ -324,10 +321,8 @@ def test_custom_bundle_provider_resolves_active_and_retired_bundles(): assert isinstance(provider, CustomDagBundleProvider) assert provider.metadata_requests == 0 - assert manager.get_configured_bundle_metadata() == (DagBundleMetadata(name="active-bundle"),) - assert provider.metadata_requests == 1 assert manager.get_active_bundle_metadata() == (DagBundleMetadata(name="active-bundle"),) - assert provider.metadata_requests == 2 + assert provider.metadata_requests == 1 assert manager.get_bundle_metadata("active-bundle") == DagBundleMetadata(name="active-bundle") assert manager.get_all_bundle_names() == ["active-bundle"] with pytest.raises(ValueError, match="'unknown-bundle' is not configured"): From 64ea20f640a4122c20b66f3c4e74c38f9cc4e313 Mon Sep 17 00:00:00 2001 From: Sameer Raj Date: Mon, 21 Sep 2026 21:18:08 +0530 Subject: [PATCH 08/34] Preserve configured bundle team helper --- airflow-core/src/airflow/api/common/airflow_health.py | 4 ++-- .../src/airflow/dag_processing/bundles/manager.py | 10 ++++++---- .../dag_processing/bundles/test_dag_bundle_manager.py | 9 +++++---- 3 files changed, 13 insertions(+), 10 deletions(-) diff --git a/airflow-core/src/airflow/api/common/airflow_health.py b/airflow-core/src/airflow/api/common/airflow_health.py index f7b86a1d2c44e..b135990fc46c1 100644 --- a/airflow-core/src/airflow/api/common/airflow_health.py +++ b/airflow-core/src/airflow/api/common/airflow_health.py @@ -23,7 +23,7 @@ from sqlalchemy import select from airflow.configuration import conf -from airflow.dag_processing.bundles.manager import DagBundlesManager +from airflow.dag_processing.bundles.manager import _get_configured_bundle_team_names from airflow.jobs.dag_processor_job_runner import DagProcessorJobRunner from airflow.jobs.job import Job from airflow.jobs.scheduler_job_runner import SchedulerJobRunner @@ -121,7 +121,7 @@ def _dag_processor_instance_health(job: Job) -> dict[str, Any]: def _configured_bundle_teams() -> dict[str, str | None]: """Map every configured Dag bundle to the team owning it, empty when the config is unreadable.""" try: - return DagBundlesManager().get_declared_bundle_teams() + return _get_configured_bundle_team_names() except Exception: # A health probe must not fail on malformed bundle config; callers fall back to liveness only. log.warning("Could not read the Dag bundle configuration", exc_info=True) diff --git a/airflow-core/src/airflow/dag_processing/bundles/manager.py b/airflow-core/src/airflow/dag_processing/bundles/manager.py index 70bc24b8a71ef..c5b3c258249f9 100644 --- a/airflow-core/src/airflow/dag_processing/bundles/manager.py +++ b/airflow-core/src/airflow/dag_processing/bundles/manager.py @@ -117,6 +117,12 @@ def _guess_best_bundle_for_fileloc( return None +def _get_configured_bundle_team_names() -> dict[str, str | None]: + """Get the team owning each bundle declared by the configured provider.""" + manager = DagBundlesManager() + return {metadata.name: metadata.team_name for metadata in manager._get_provider_bundle_metadata()} + + def _is_safe_bundle_url(url: str) -> bool: """ Check if a bundle URL is safe to use. @@ -242,10 +248,6 @@ def get_active_bundle_metadata(self) -> tuple[DagBundleMetadata, ...]: return tuple(bundle_metadata.values()) - def get_declared_bundle_teams(self) -> dict[str, str | None]: - """Get the team owning each bundle declared by the configured provider.""" - return {metadata.name: metadata.team_name for metadata in self._get_provider_bundle_metadata()} - @provide_session def sync_bundles_to_db(self, *, deactivate_missing: bool = True, session: Session = NEW_SESSION) -> None: """ 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 4221ccee3ac33..b040baadad9af 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 @@ -32,6 +32,7 @@ from airflow.dag_processing.bundles.local import LocalDagBundle from airflow.dag_processing.bundles.manager import ( DagBundlesManager, + _get_configured_bundle_team_names, _guess_best_bundle_for_fileloc, _load_bundle_config_snapshot, get_configured_bundle_names, @@ -206,7 +207,7 @@ def __init__(self): @pytest.mark.parametrize("load_examples", ["False", "True"]) -def test_get_declared_bundle_teams_excludes_example_bundles(load_examples): +def test_get_configured_bundle_team_names_excludes_example_bundles(load_examples): with conf_vars( { ("core", "load_examples"): load_examples, @@ -214,7 +215,7 @@ def test_get_declared_bundle_teams_excludes_example_bundles(load_examples): ("dag_processor", "dag_bundle_config_list"): json.dumps(TEAM_BUNDLE_CONFIG), } ): - assert DagBundlesManager().get_declared_bundle_teams() == { + assert _get_configured_bundle_team_names() == { "team-bundle": "team-a", "unscoped-bundle": None, } @@ -235,8 +236,8 @@ def test_get_configured_bundle_names(): @conf_vars({("dag_processor", "dag_bundle_config_list"): "[]"}) -def test_get_declared_bundle_teams_without_config(): - assert DagBundlesManager().get_declared_bundle_teams() == {} +def test_get_configured_bundle_team_names_without_config(): + assert _get_configured_bundle_team_names() == {} def test_get_bundle(): From 49db0e3f4b0812d64087095f18d83342439d3c1a Mon Sep 17 00:00:00 2001 From: Sameer Raj Date: Mon, 21 Sep 2026 21:32:42 +0530 Subject: [PATCH 09/34] Expose configured Dag bundle teams through manager --- airflow-core/src/airflow/api/common/airflow_health.py | 4 ++-- .../src/airflow/dag_processing/bundles/manager.py | 10 ++++------ .../dag_processing/bundles/test_dag_bundle_manager.py | 5 ++--- 3 files changed, 8 insertions(+), 11 deletions(-) diff --git a/airflow-core/src/airflow/api/common/airflow_health.py b/airflow-core/src/airflow/api/common/airflow_health.py index b135990fc46c1..a3571b3ce3fbf 100644 --- a/airflow-core/src/airflow/api/common/airflow_health.py +++ b/airflow-core/src/airflow/api/common/airflow_health.py @@ -23,7 +23,7 @@ from sqlalchemy import select from airflow.configuration import conf -from airflow.dag_processing.bundles.manager import _get_configured_bundle_team_names +from airflow.dag_processing.bundles.manager import DagBundlesManager from airflow.jobs.dag_processor_job_runner import DagProcessorJobRunner from airflow.jobs.job import Job from airflow.jobs.scheduler_job_runner import SchedulerJobRunner @@ -121,7 +121,7 @@ def _dag_processor_instance_health(job: Job) -> dict[str, Any]: def _configured_bundle_teams() -> dict[str, str | None]: """Map every configured Dag bundle to the team owning it, empty when the config is unreadable.""" try: - return _get_configured_bundle_team_names() + return DagBundlesManager().get_configured_bundle_team_names() except Exception: # A health probe must not fail on malformed bundle config; callers fall back to liveness only. log.warning("Could not read the Dag bundle configuration", exc_info=True) diff --git a/airflow-core/src/airflow/dag_processing/bundles/manager.py b/airflow-core/src/airflow/dag_processing/bundles/manager.py index c5b3c258249f9..2ac04b0f8d782 100644 --- a/airflow-core/src/airflow/dag_processing/bundles/manager.py +++ b/airflow-core/src/airflow/dag_processing/bundles/manager.py @@ -117,12 +117,6 @@ def _guess_best_bundle_for_fileloc( return None -def _get_configured_bundle_team_names() -> dict[str, str | None]: - """Get the team owning each bundle declared by the configured provider.""" - manager = DagBundlesManager() - return {metadata.name: metadata.team_name for metadata in manager._get_provider_bundle_metadata()} - - def _is_safe_bundle_url(url: str) -> bool: """ Check if a bundle URL is safe to use. @@ -248,6 +242,10 @@ def get_active_bundle_metadata(self) -> tuple[DagBundleMetadata, ...]: return tuple(bundle_metadata.values()) + def get_configured_bundle_team_names(self) -> dict[str, str | None]: + """Get the team owning each bundle declared by the configured provider.""" + return {metadata.name: metadata.team_name for metadata in self._get_provider_bundle_metadata()} + @provide_session def sync_bundles_to_db(self, *, deactivate_missing: bool = True, session: Session = NEW_SESSION) -> None: """ 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 b040baadad9af..ac8ef8419430e 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 @@ -32,7 +32,6 @@ from airflow.dag_processing.bundles.local import LocalDagBundle from airflow.dag_processing.bundles.manager import ( DagBundlesManager, - _get_configured_bundle_team_names, _guess_best_bundle_for_fileloc, _load_bundle_config_snapshot, get_configured_bundle_names, @@ -215,7 +214,7 @@ def test_get_configured_bundle_team_names_excludes_example_bundles(load_examples ("dag_processor", "dag_bundle_config_list"): json.dumps(TEAM_BUNDLE_CONFIG), } ): - assert _get_configured_bundle_team_names() == { + assert DagBundlesManager().get_configured_bundle_team_names() == { "team-bundle": "team-a", "unscoped-bundle": None, } @@ -237,7 +236,7 @@ def test_get_configured_bundle_names(): @conf_vars({("dag_processor", "dag_bundle_config_list"): "[]"}) def test_get_configured_bundle_team_names_without_config(): - assert _get_configured_bundle_team_names() == {} + assert DagBundlesManager().get_configured_bundle_team_names() == {} def test_get_bundle(): From 6931e38bd2a3586fd90e5353e36b21c430cf6e7a Mon Sep 17 00:00:00 2001 From: Sameer Raj Date: Tue, 22 Sep 2026 18:36:09 +0530 Subject: [PATCH 10/34] Keep bundle discovery outside Dag parsing --- .../airflow/dag_processing/bundles/manager.py | 2 + .../src/airflow/dag_processing/dagbag.py | 50 +++---------- .../src/airflow/dag_processing/manager.py | 7 +- .../src/airflow/dag_processing/processor.py | 9 ++- .../bundles/test_dag_bundle_manager.py | 10 +++ .../tests/unit/dag_processing/test_dagbag.py | 74 ++++--------------- .../tests/unit/dag_processing/test_manager.py | 11 ++- 7 files changed, 60 insertions(+), 103 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/bundles/manager.py b/airflow-core/src/airflow/dag_processing/bundles/manager.py index 2ac04b0f8d782..99b038f0aefdf 100644 --- a/airflow-core/src/airflow/dag_processing/bundles/manager.py +++ b/airflow-core/src/airflow/dag_processing/bundles/manager.py @@ -233,6 +233,8 @@ def _get_provider_bundle_metadata(self) -> tuple[DagBundleMetadata, ...]: def get_active_bundle_metadata(self) -> tuple[DagBundleMetadata, ...]: """Get metadata for all active Dag bundles, including example bundles.""" bundle_metadata = {metadata.name: metadata for metadata in self._get_provider_bundle_metadata()} + if not bundle_metadata: + return () # Airflow owns example Dag bundles; configured providers do not need to include them. for name in self._example_dag_bundle_paths: diff --git a/airflow-core/src/airflow/dag_processing/dagbag.py b/airflow-core/src/airflow/dag_processing/dagbag.py index b8968c54267de..b5f3f0d5bc341 100644 --- a/airflow-core/src/airflow/dag_processing/dagbag.py +++ b/airflow-core/src/airflow/dag_processing/dagbag.py @@ -113,35 +113,16 @@ def _executor_exists(executor_name: str, team_name: str | None) -> bool: return False -def _get_bundle_team_name(bundle_name: str | None) -> str | None: - """Return the team that owns *bundle_name* when multi-team is on, else ``None``.""" - if not conf.getboolean("core", "multi_team") or not bundle_name: - return None - from airflow.dag_processing.bundles.manager import DagBundlesManager - - return DagBundlesManager()._bundle_config[bundle_name].team_name - - -def _validate_executor_fields(dag: DAG | SerializedDAG, bundle_name: str | None = None) -> None: +def _validate_executor_fields( + dag: DAG | SerializedDAG, bundle_name: str | None = None, team_name: str | None = None +) -> None: """Validate that executors specified in tasks are available and owned by the same team as the dag bundle.""" import logging log = logging.getLogger(__name__) - - # Check if multi team is available by reading the multi_team configuration (which is boolean) - if conf.getboolean("core", "multi_team"): - # Get team name from bundle metadata if available - if bundle_name: - from airflow.dag_processing.bundles.manager import DagBundlesManager - - bundle_manager = DagBundlesManager() - bundle_metadata = bundle_manager.get_bundle_metadata(bundle_name) - - dag_team_name = bundle_metadata.team_name - if dag_team_name: - log.debug( - "Found team '%s' for DAG '%s' via bundle '%s'", dag_team_name, dag.dag_id, bundle_name - ) + dag_team_name = team_name if conf.getboolean("core", "multi_team") else None + if dag_team_name: + log.debug("Found team '%s' for DAG '%s' via bundle '%s'", dag_team_name, dag.dag_id, bundle_name) for task in dag.tasks: if not task.executor: @@ -165,19 +146,10 @@ def _validate_executor_fields(dag: DAG | SerializedDAG, bundle_name: str | None def _assign_default_team_pools( dag: DAG, - bundle_name: str | None = None, + team_name: str | None = None, ) -> None: """Assign the default team pool to tasks that do not explicitly specify a pool.""" - dag_team_name = None - - if conf.getboolean("core", "multi_team"): - if bundle_name: - from airflow.dag_processing.bundles.manager import DagBundlesManager - - bundle_manager = DagBundlesManager() - bundle_metadata = bundle_manager.get_bundle_metadata(bundle_name) - - dag_team_name = bundle_metadata.team_name + dag_team_name = team_name if conf.getboolean("core", "multi_team") else None if not dag_team_name: return @@ -221,11 +193,13 @@ def __init__( bundle_path: Path | None = None, bundle_name: str | None = None, parse_lang_sdk_files: bool = False, + team_name: str | None = None, ): super().__init__() self.bundle_path = bundle_path self.bundle_name = bundle_name self.parse_lang_sdk_files = parse_lang_sdk_files + self.team_name = team_name dag_folder = dag_folder or settings.DAGS_FOLDER self.dag_folder = dag_folder @@ -402,8 +376,8 @@ def _import_sdk_definition(self, importer: AbstractDagImporter, definition: DagD dag.relative_fileloc = self._get_relative_fileloc(fileloc) dag.bundle_name = self.bundle_name dag.validate() - _validate_executor_fields(dag, self.bundle_name) - _assign_default_team_pools(dag, self.bundle_name) + _validate_executor_fields(dag, self.bundle_name, self.team_name) + _assign_default_team_pools(dag, self.team_name) self.bag_dag(dag=dag) bagged_dags.append(dag) except AirflowClusterPolicySkipDag: diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index e5043104dba99..4bc185fe67f58 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -1474,7 +1474,9 @@ def client(self) -> Client: client.base_url = "http://in-process.invalid./" return client - def _create_process(self, dag_file: DagFileInfo) -> BaseDagFileProcessorProcess: + def _create_process( + self, dag_file: DagFileInfo, team_name: str | None = None + ) -> BaseDagFileProcessorProcess: id = uuid7() callback_to_execute_for_file = self._callback_to_execute.pop(dag_file, []) @@ -1500,6 +1502,7 @@ def _create_process(self, dag_file: DagFileInfo) -> BaseDagFileProcessorProcess: path=dag_file.absolute_path, bundle_path=cast("Path", dag_file.bundle_path), bundle_name=dag_file.bundle_name, + team_name=team_name, dag_file_rel_path=str(dag_file.rel_path), callbacks=callback_to_execute_for_file, selector=self.selector, @@ -1519,7 +1522,7 @@ def _start_new_processes(self): if file in self._processors: continue - processor = self._create_process(file) + processor = self._create_process(file, team_name=bundle_to_team.get(file.bundle_name)) stats.incr( "dag_processing.processes", tags=prune_dict( diff --git a/airflow-core/src/airflow/dag_processing/processor.py b/airflow-core/src/airflow/dag_processing/processor.py index 9530d7be8c9cd..c424efd239bfc 100644 --- a/airflow-core/src/airflow/dag_processing/processor.py +++ b/airflow-core/src/airflow/dag_processing/processor.py @@ -112,6 +112,9 @@ class DagFileParseRequest(BaseModel): bundle_name: str """Bundle name for team-specific executor validation.""" + team_name: str | None = None + """Team resolved by the Dag processor manager for parser-side validation.""" + callback_requests: list[CallbackRequest] = Field(default_factory=list) type: Literal["DagFileParseRequest"] = "DagFileParseRequest" @@ -245,6 +248,7 @@ def _parse_file(msg: DagFileParseRequest, log: FilteringBoundLogger) -> DagFileP dag_folder=msg.file, bundle_path=msg.bundle_path, bundle_name=msg.bundle_name, + team_name=msg.team_name, load_op_links=False, ) @@ -726,6 +730,7 @@ def start( # type: ignore[override] bundle_name: str, dag_file_rel_path: str, callbacks: list[CallbackRequest], + team_name: str | None = None, target: Callable[[], None] = _parse_file_entrypoint, client: Client, **kwargs, @@ -753,7 +758,7 @@ def start( # type: ignore[override] **kwargs, ) proc.had_callbacks = bool(callbacks) # Track if this process had callbacks - proc._on_child_started(callbacks, path, bundle_path, bundle_name) + proc._on_child_started(callbacks, path, bundle_path, bundle_name, team_name) return proc def _on_child_started( @@ -762,11 +767,13 @@ def _on_child_started( path: str | os.PathLike[str], bundle_path: Path, bundle_name: str, + team_name: str | None = None, ) -> None: msg = DagFileParseRequest( file=os.fspath(path), bundle_path=bundle_path, bundle_name=bundle_name, + team_name=team_name, callback_requests=callbacks, ) self.send_msg(msg, request_id=0) 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 ac8ef8419430e..0fa5339fc9eac 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 @@ -699,6 +699,16 @@ def test_example_dags_bundle_added(): assert "example_dags" not in manager.get_all_bundle_names() +@conf_vars( + { + ("core", "load_examples"): "True", + ("dag_processor", "dag_bundle_config_list"): "[]", + } +) +def test_example_dags_not_added_without_configured_bundles(): + assert DagBundlesManager().get_all_bundle_names() == [] + + @conf_vars( { ( diff --git a/airflow-core/tests/unit/dag_processing/test_dagbag.py b/airflow-core/tests/unit/dag_processing/test_dagbag.py index 4b776ff14b0f3..1cc949c089392 100644 --- a/airflow-core/tests/unit/dag_processing/test_dagbag.py +++ b/airflow-core/tests/unit/dag_processing/test_dagbag.py @@ -36,7 +36,6 @@ from sqlalchemy import select from airflow import settings -from airflow.dag_processing.bundles.provider import DagBundleMetadata from airflow.dag_processing.dagbag import ( BundleDagBag, DagBag, @@ -120,45 +119,32 @@ class TestValidateExecutorFields: """Comprehensive tests for _validate_executor_fields function.""" @patch.object(ExecutorLoader, "lookup_executor_name_by_str") - def test_multi_team_disabled_ignores_bundle_name(self, mock_lookup): - """Test that when multi_team is disabled, bundle_name is ignored and no team lookup occurs.""" + def test_multi_team_disabled_ignores_team_name(self, mock_lookup): + """Test that team information is ignored when multi-team mode is disabled.""" with DAG("test-dag", schedule=None) as dag: BaseOperator(task_id="t1", executor="test.executor") # multi_team disabled by default, no need to add conf_vars - _validate_executor_fields(dag, bundle_name="some_bundle") + _validate_executor_fields(dag, bundle_name="some_bundle", team_name="test_team") # Should call ExecutorLoader without team_name (defaults to None) mock_lookup.assert_called_once_with("test.executor", team_name=None, validate_teams=False) - @patch("airflow.dag_processing.bundles.manager.DagBundlesManager") @patch.object(ExecutorLoader, "lookup_executor_name_by_str") - def test_multi_team_enabled_bundle_exists_with_team(self, mock_lookup, mock_manager_class): - """Test successful team lookup when bundle exists and has team_name.""" - # Setup mock bundle manager - bundle_metadata = DagBundleMetadata(name="test_bundle", team_name="test_team") - - mock_manager = mock_manager_class.return_value - mock_manager.get_bundle_metadata.return_value = bundle_metadata - + def test_multi_team_enabled_bundle_exists_with_team(self, mock_lookup): + """Test successful executor validation with a resolved team name.""" with DAG("test-dag", schedule=None) as dag: BaseOperator(task_id="t1", executor="team.executor") with conf_vars({("core", "multi_team"): "True"}): - _validate_executor_fields(dag, bundle_name="test_bundle") + _validate_executor_fields(dag, bundle_name="test_bundle", team_name="test_team") # Should call ExecutorLoader with team from bundle config mock_lookup.assert_called_once_with("team.executor", team_name="test_team", validate_teams=False) - @patch("airflow.dag_processing.bundles.manager.DagBundlesManager") @patch.object(ExecutorLoader, "lookup_executor_name_by_str") - def test_multi_team_enabled_bundle_exists_no_team(self, mock_lookup, mock_manager_class): + def test_multi_team_enabled_bundle_exists_no_team(self, mock_lookup): """Test when bundle exists but has no team_name (None or empty).""" - bundle_metadata = DagBundleMetadata(name="test_bundle") - - mock_manager = mock_manager_class.return_value - mock_manager.get_bundle_metadata.return_value = bundle_metadata - with DAG("test-dag", schedule=None) as dag: BaseOperator(task_id="t1", executor="test.executor") @@ -183,15 +169,9 @@ def test_multiple_tasks_with_executors(self, mock_lookup): mock_lookup.assert_any_call("executor1", team_name=None, validate_teams=False) mock_lookup.assert_any_call("executor2", team_name=None, validate_teams=False) - @patch("airflow.dag_processing.bundles.manager.DagBundlesManager") @patch.object(ExecutorLoader, "lookup_executor_name_by_str") - def test_executor_validation_failure_with_team(self, mock_lookup, mock_manager_class): + def test_executor_validation_failure_with_team(self, mock_lookup): """Test executor validation failure when team is associated (team-specific error).""" - bundle_metadata = DagBundleMetadata(name="test_bundle", team_name="test_team") - - mock_manager = mock_manager_class.return_value - mock_manager.get_bundle_metadata.return_value = bundle_metadata - # ExecutorLoader raises exception mock_lookup.side_effect = UnknownExecutorException("Executor not found") @@ -209,7 +189,7 @@ def test_executor_validation_failure_with_team(self, mock_lookup, mock_manager_c "configured executors for team 'test_team' or available global executors." ), ): - _validate_executor_fields(dag, bundle_name="test_bundle") + _validate_executor_fields(dag, bundle_name="test_bundle", team_name="test_team") @patch.object(ExecutorLoader, "lookup_executor_name_by_str") def test_executor_validation_failure_no_team(self, mock_lookup): @@ -230,15 +210,9 @@ def test_executor_validation_failure_no_team(self, mock_lookup): ): _validate_executor_fields(dag) # No bundle_name - @patch("airflow.dag_processing.bundles.manager.DagBundlesManager") @patch.object(ExecutorLoader, "lookup_executor_name_by_str") - def test_global_executor_fallback_success(self, mock_lookup, mock_manager_class): + def test_global_executor_fallback_success(self, mock_lookup): """Test that team-specific executor failure falls back to global executor successfully.""" - bundle_metadata = DagBundleMetadata(name="test_bundle", team_name="test_team") - - mock_manager = mock_manager_class.return_value - mock_manager.get_bundle_metadata.return_value = bundle_metadata - # First call (team-specific) fails, second call (global) succeeds mock_lookup.side_effect = [UnknownExecutorException("Team executor not found"), None] @@ -247,22 +221,16 @@ def test_global_executor_fallback_success(self, mock_lookup, mock_manager_class) with conf_vars({("core", "multi_team"): "True"}): # Should not raise exception due to global fallback - _validate_executor_fields(dag, bundle_name="test_bundle") + _validate_executor_fields(dag, bundle_name="test_bundle", team_name="test_team") # Should call lookup twice: first for team, then for global assert mock_lookup.call_count == 2 mock_lookup.assert_any_call("global.executor", team_name="test_team", validate_teams=False) mock_lookup.assert_any_call("global.executor", team_name=None, validate_teams=False) - @patch("airflow.dag_processing.bundles.manager.DagBundlesManager") @patch.object(ExecutorLoader, "lookup_executor_name_by_str") - def test_global_executor_fallback_failure(self, mock_lookup, mock_manager_class): + def test_global_executor_fallback_failure(self, mock_lookup): """Test that when both team-specific and global executors fail, appropriate error is raised.""" - bundle_metadata = DagBundleMetadata(name="test_bundle", team_name="test_team") - - mock_manager = mock_manager_class.return_value - mock_manager.get_bundle_metadata.return_value = bundle_metadata - # Both calls fail mock_lookup.side_effect = UnknownExecutorException("Executor not found") @@ -280,22 +248,16 @@ def test_global_executor_fallback_failure(self, mock_lookup, mock_manager_class) "configured executors for team 'test_team' or available global executors." ), ): - _validate_executor_fields(dag, bundle_name="test_bundle") + _validate_executor_fields(dag, bundle_name="test_bundle", team_name="test_team") # Should call lookup twice: first for team, then for global fallback assert mock_lookup.call_count == 2 mock_lookup.assert_any_call("unknown.executor", team_name="test_team", validate_teams=False) mock_lookup.assert_any_call("unknown.executor", team_name=None, validate_teams=False) - @patch("airflow.dag_processing.bundles.manager.DagBundlesManager") @patch.object(ExecutorLoader, "lookup_executor_name_by_str") - def test_team_specific_executor_success_no_fallback(self, mock_lookup, mock_manager_class): + def test_team_specific_executor_success_no_fallback(self, mock_lookup): """Test that when team-specific executor succeeds, global fallback is not attempted.""" - bundle_metadata = DagBundleMetadata(name="test_bundle", team_name="test_team") - - mock_manager = mock_manager_class.return_value - mock_manager.get_bundle_metadata.return_value = bundle_metadata - # First call (team-specific) succeeds mock_lookup.return_value = None @@ -303,7 +265,7 @@ def test_team_specific_executor_success_no_fallback(self, mock_lookup, mock_mana BaseOperator(task_id="task1", executor="team.executor") with conf_vars({("core", "multi_team"): "True"}): - _validate_executor_fields(dag, bundle_name="test_bundle") + _validate_executor_fields(dag, bundle_name="test_bundle", team_name="test_team") # Should only call lookup once for team-specific executor mock_lookup.assert_called_once_with("team.executor", team_name="test_team", validate_teams=False) @@ -456,18 +418,13 @@ def test_dag_with_bundle_name(self, tmp_path): ), ], ) - @patch("airflow.dag_processing.bundles.manager.DagBundlesManager") def test_default_pool_replaced_with_team_pool( self, - mock_manager, tmp_path, team_name, operator_args, expected_pool, ): - bundle_metadata = DagBundleMetadata(name="test_bundle", team_name=team_name) - mock_manager.return_value.get_bundle_metadata.return_value = bundle_metadata - dag_file = tmp_path / "test_dag.py" dag_file.write_text( textwrap.dedent( @@ -489,6 +446,7 @@ def my_dag(): dagbag = DagBag( dag_folder=os.fspath(tmp_path), bundle_name="test_bundle", + team_name=team_name, ) dag = dagbag.get_dag("my_dag") diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index 2a92835edba39..372c317dc4827 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -2215,6 +2215,7 @@ def test_collect_results_tolerates_stale_file_handle_on_close(self): "file": "/opt/airflow/dags/test_dag.py", "bundle_path": "/opt/airflow/dags", "bundle_name": "testing", + "team_name": None, "callback_requests": [], "type": "DagFileParseRequest", }, @@ -2236,6 +2237,7 @@ def test_collect_results_tolerates_stale_file_handle_on_close(self): "file": "/opt/airflow/dags/dag_callback_dag.py", "bundle_path": "/opt/airflow/dags", "bundle_name": "testing", + "team_name": None, "callback_requests": [ { "filepath": "dag_callback_dag.py", @@ -3799,10 +3801,10 @@ def test_bundles_with_team(self, session): assert len(team2.dag_bundles) == 0 @mock.patch.object(DagFileProcessorProcess, "start") - def test_create_process_passes_bundle_name_to_process_start( + def test_create_process_passes_bundle_metadata_to_process_start( self, mock_process_start, configure_testing_dag_bundle ): - """Test that DagFileProcessorManager._create_process() passes bundle_name to DagFileProcessorProcess.start()""" + """Test that process creation passes the bundle name and resolved team.""" with configure_testing_dag_bundle("/tmp"): manager = DagFileProcessorManager(max_runs=1) manager._dag_bundles = list(DagBundlesManager().get_all_dag_bundles()) @@ -3815,13 +3817,13 @@ def test_create_process_passes_bundle_name_to_process_start( # Mock the process creation mock_process_start.return_value = self.mock_processor()[0] - # Call _create_process (only takes one parameter: dag_file) - manager._create_process(file_info) + manager._create_process(file_info, team_name="team-a") # Verify DagFileProcessorProcess.start was called with correct bundle_name mock_process_start.assert_called_once() call_kwargs = mock_process_start.call_args.kwargs assert call_kwargs["bundle_name"] == "testing" + assert call_kwargs["team_name"] == "team-a" @mock.patch("airflow.dag_processing.manager.stats.initialize") def test_stats_initialize_called_on_run(self, stats_init_mock, tmp_path, configure_testing_dag_bundle): @@ -4658,6 +4660,7 @@ def test_start_new_processes_includes_team_name(self, mock_incr, mock_get_team_n manager._start_new_processes() + manager._create_process.assert_called_once_with(dag_file, team_name="team_alpha") mock_incr.assert_any_call( "dag_processing.processes", tags={ From 26cade399fee6ec1fc73d0286bbd3deb0ff3ac20 Mon Sep 17 00:00:00 2001 From: Sameer Raj Date: Tue, 22 Sep 2026 23:27:03 +0530 Subject: [PATCH 11/34] Remove unused bundle metadata lookup --- airflow-core/src/airflow/dag_processing/bundles/manager.py | 7 ------- .../unit/dag_processing/bundles/test_dag_bundle_manager.py | 3 --- 2 files changed, 10 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/bundles/manager.py b/airflow-core/src/airflow/dag_processing/bundles/manager.py index 99b038f0aefdf..7cf4f3ee96041 100644 --- a/airflow-core/src/airflow/dag_processing/bundles/manager.py +++ b/airflow-core/src/airflow/dag_processing/bundles/manager.py @@ -599,13 +599,6 @@ def get_bundle( return LocalDagBundle(name=name, path=path, version=version, version_data=version_data) return self._bundle_provider.get_bundle(name=name, version=version, version_data=version_data) - def get_bundle_metadata(self, name: str) -> DagBundleMetadata: - """Get metadata for an active Dag bundle.""" - for metadata in self.get_active_bundle_metadata(): - if metadata.name == name: - return metadata - raise ValueError(f"Requested bundle '{name}' is not configured.") - @classmethod def is_bundle_configured(cls, name: str) -> bool: """ 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 0fa5339fc9eac..4b02d5bfa4fb8 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 @@ -323,10 +323,7 @@ def test_custom_bundle_provider_resolves_active_and_retired_bundles(): assert manager.get_active_bundle_metadata() == (DagBundleMetadata(name="active-bundle"),) assert provider.metadata_requests == 1 - assert manager.get_bundle_metadata("active-bundle") == DagBundleMetadata(name="active-bundle") assert manager.get_all_bundle_names() == ["active-bundle"] - with pytest.raises(ValueError, match="'unknown-bundle' is not configured"): - manager.get_bundle_metadata("unknown-bundle") active_bundle = manager.get_bundle("active-bundle") retired_bundle = manager.get_bundle( From c6fac1811228038b09cae9de39b39ac8ce0c0084 Mon Sep 17 00:00:00 2001 From: Sameer Raj Date: Wed, 23 Sep 2026 12:28:36 +0530 Subject: [PATCH 12/34] Update callback process expectation --- airflow-core/tests/unit/dag_processing/test_manager.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index 372c317dc4827..95f3b8130e65e 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -2989,6 +2989,7 @@ def test_callback_queue(self, mock_get_logger, configure_testing_dag_bundle): path=Path(dag2_path.bundle_path, dag2_path.rel_path), bundle_path=dag2_path.bundle_path, bundle_name="testing", + team_name=None, dag_file_rel_path=str(dag2_path.rel_path), callbacks=[dag2_req1], selector=mock.ANY, @@ -3002,6 +3003,7 @@ def test_callback_queue(self, mock_get_logger, configure_testing_dag_bundle): path=Path(dag1_path.bundle_path, dag1_path.rel_path), bundle_path=dag1_path.bundle_path, bundle_name="testing", + team_name=None, dag_file_rel_path=str(dag1_path.rel_path), callbacks=[dag1_req1, dag1_req2], selector=mock.ANY, From ba8ebde1e296e83a7e7e8860821bbb7723326bff Mon Sep 17 00:00:00 2001 From: Sameer Raj Date: Wed, 23 Sep 2026 12:45:06 +0530 Subject: [PATCH 13/34] Version Dag processor team metadata --- .../sdk/execution_time/schema/schema.json | 12 +++ .../schema/versions/__init__.py | 2 + .../schema/versions/v2026_10_30.py | 12 ++- .../execution_time/schema/test_migrator.py | 78 ++++--------------- 4 files changed, 38 insertions(+), 66 deletions(-) diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json index 41a6ab4c6d394..386b6700a1184 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -933,6 +933,18 @@ "title": "Bundle Name", "type": "string" }, + "team_name": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Team Name" + }, "callback_requests": { "items": { "discriminator": { diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py b/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py index 0f35e9f0e6f2b..8f02b44cf01c6 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py +++ b/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py @@ -44,6 +44,7 @@ def get_bundle() -> VersionBundle: AddLoopRuntime, AddRegionCoordinatesToSupervisorTaskInstance, AddRetryReasonToTaskState, + AddTeamNameToDagFileParseRequest, ) return VersionBundle( @@ -56,6 +57,7 @@ def get_bundle() -> VersionBundle: AddDagDefinitionsToDagFileParsingResult, AddRegionCoordinatesToSupervisorTaskInstance, AddFirstTaskRescheduleStartDateToSupervisorTIRunContext, + AddTeamNameToDagFileParseRequest, ), Version("2026-06-16"), ) diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py b/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py index c9a80625c3ea0..fbbe632284873 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py +++ b/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py @@ -28,7 +28,7 @@ from cadwyn import VersionChange, schema -from airflow.dag_processing.processor import DagFileParsingResult # noqa: SDK002 +from airflow.dag_processing.processor import DagFileParseRequest, DagFileParsingResult # noqa: SDK002 from airflow.sdk.api.datamodels._generated import ( PreviousTIResponse, TaskInstance, @@ -119,3 +119,13 @@ class AddFirstTaskRescheduleStartDateToSupervisorTIRunContext(VersionChange): instructions_to_migrate_to_previous_version = ( schema(TIRunContext).field("first_task_reschedule_start_date").didnt_exist, ) + + +class AddTeamNameToDagFileParseRequest(VersionChange): + """Add the team name to Dag file parse requests.""" + + description = __doc__ + + instructions_to_migrate_to_previous_version = ( + schema(DagFileParseRequest).field("team_name").didnt_exist, + ) diff --git a/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py index e8c4a16aff248..40c78f3d2dcd0 100644 --- a/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py +++ b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py @@ -444,18 +444,7 @@ class TestRealBundleArgBindingsDowngrade: @pytest.fixture def startup_details(self): - import datetime - import uuid - - from airflow.sdk.api.datamodels._generated import ( - BundleInfo, - DagRun, - DagRunState, - DagRunType, - TaskInstance, - TIRunContext, - ) - from airflow.sdk.execution_time.comms import StartupDetails + now = datetime.datetime.now(datetime.UTC) return StartupDetails( @@ -594,61 +583,20 @@ def test_upgrade_fills_missing_retry_reason_with_none(self, real_migrator): assert out["retry_reason"] is None -class TestRealBundleFirstTaskRescheduleStartDate: - """ - Drive the *real* supervisor bundle through the ``first_task_reschedule_start_date`` migration. - - ``AddFirstTaskRescheduleStartDateToSupervisorTIRunContext`` is ``didnt_exist``-only, so the - downgrade re-validation is the only place it is visible: dropping it from the bundle, or - registering it under the wrong version, leaves the field on the wire for a pinned runtime. - """ +class TestRealBundleDagFileParseRequestDowngrade: + def test_downgrade_strips_team_name_for_previous_version(self): + from pathlib import Path - FIRST_RESCHEDULE = datetime.datetime(2026, 1, 1, tzinfo=datetime.UTC) + from airflow.dag_processing.processor import DagFileParseRequest - @pytest.fixture - def startup_details(self): - now = datetime.datetime.now(datetime.UTC) - return StartupDetails( - ti=TaskInstance( - id=uuid.uuid4(), - task_id="wait", - dag_id="d", - run_id="r", - try_number=1, - dag_version_id=uuid.uuid4(), - ), - dag_rel_path="d.py", - bundle_info=BundleInfo(name="b", version=None), - start_date=now, - ti_context=TIRunContext( - dag_run=DagRun( - dag_id="d", - run_id="r", - logical_date=now, - data_interval_start=None, - data_interval_end=None, - start_date=now, - end_date=None, - run_type=DagRunType.MANUAL, - state=DagRunState.RUNNING, - run_after=now, - consumed_asset_events=[], - partition_key=None, - ), - max_tries=1, - first_task_reschedule_start_date=self.FIRST_RESCHEDULE, - ), - sentry_integration="", + request = DagFileParseRequest( + file="/bundle/dags/example.py", + bundle_path=Path("/bundle/dags"), + bundle_name="test-bundle", + team_name="test-team", + callback_requests=[], ) - @pytest.fixture - def real_migrator(self) -> SchemaVersionMigrator: - return get_schema_version_migrator() - - def test_downgrade_strips_field_for_previous_version(self, real_migrator, startup_details): - out = real_migrator.downgrade(startup_details, "2026-06-16").model_dump() - assert "first_task_reschedule_start_date" not in out["ti_context"] + out = get_schema_version_migrator().downgrade(request, "2026-06-16").model_dump() - def test_head_version_keeps_field(self, real_migrator, startup_details): - out = real_migrator.downgrade(startup_details, "2026-10-30") - assert out.ti_context.first_task_reschedule_start_date == self.FIRST_RESCHEDULE + assert "team_name" not in out From 59c8ff113c021c05a6e0f2299879486166d0a0a5 Mon Sep 17 00:00:00 2001 From: Sameer Raj Date: Sat, 10 Oct 2026 23:20:56 +0530 Subject: [PATCH 14/34] Keep custom Dag bundles compatible with language SDK parsing Bundle discovery and team ownership must work consistently for coordinators and both parser paths after upstream language SDK support. --- .../cli/commands/dag_processor_command.py | 4 +- .../dag_processing/bundles/provider.py | 64 +++++----------- .../src/airflow/dag_processing/dagbag.py | 1 + .../dag_processing/lang_sdk_processor.py | 12 ++- .../src/airflow/dag_processing/manager.py | 1 + .../commands/test_dag_processor_command.py | 23 +++++- .../bundles/test_dag_bundle_manager.py | 31 +++++++- .../tests/unit/dag_processing/test_dagbag.py | 8 +- .../dag_processing/test_lang_sdk_processor.py | 49 +++++++++++-- .../tests/unit/dag_processing/test_manager.py | 3 +- .../src/tests_common/pytest_plugin.py | 2 +- .../src/tests_common/test_utils/config.py | 4 +- go-sdk/pkg/execution/genmodels/models.gen.go | 3 + go-sdk/schema/supervisor-schema.json | 12 +++ .../airflow/sdk/execution_time/coordinator.py | 2 +- .../execution_time/schema/test_migrator.py | 73 ++++++++++++++++++- ts-sdk/schema/supervisor-schema.json | 12 +++ ts-sdk/src/generated/supervisor.ts | 6 +- 18 files changed, 241 insertions(+), 69 deletions(-) diff --git a/airflow-core/src/airflow/cli/commands/dag_processor_command.py b/airflow-core/src/airflow/cli/commands/dag_processor_command.py index b8195bead2895..4e65cd0008f20 100644 --- a/airflow-core/src/airflow/cli/commands/dag_processor_command.py +++ b/airflow-core/src/airflow/cli/commands/dag_processor_command.py @@ -23,7 +23,7 @@ from airflow.cli.commands.daemon_utils import run_command_with_daemon_option from airflow.configuration import conf -from airflow.dag_processing.bundles.manager import get_configured_bundle_team_names +from airflow.dag_processing.bundles.manager import DagBundlesManager from airflow.dag_processing.manager import DagFileProcessorManager from airflow.jobs.dag_processor_job_runner import DagProcessorJobRunner from airflow.jobs.job import Job, run_job @@ -52,7 +52,7 @@ def _get_team_names(bundle_names: list[str] | None) -> list[str]: if not conf.getboolean("core", "multi_team"): return [] - configured = get_configured_bundle_team_names() + configured = DagBundlesManager().get_configured_bundle_team_names() names = bundle_names or list(configured) return sorted({team for name in names if (team := configured.get(name)) is not None}) diff --git a/airflow-core/src/airflow/dag_processing/bundles/provider.py b/airflow-core/src/airflow/dag_processing/bundles/provider.py index c676a9106a5f8..caacba1aeea65 100644 --- a/airflow-core/src/airflow/dag_processing/bundles/provider.py +++ b/airflow-core/src/airflow/dag_processing/bundles/provider.py @@ -18,6 +18,7 @@ from abc import ABC, abstractmethod from dataclasses import dataclass +from functools import cache from typing import TYPE_CHECKING, Any from pydantic import BaseModel, ValidationError @@ -72,19 +73,6 @@ class _ExternalBundleConfig(BaseModel): team_name: str | None = None -class _InternalBundleConfig(BaseModel): - """ - Schema used internally (in this file) to define the configuration for a Dag bundle. - - Configuration defined by users when read must match ``_ExternalBundleConfig``. - This configuration is then parsed and converted to ``_InternalBundleConfig`` to be used across this file. - """ - - bundle_class: type[BaseDagBundle] - kwargs: dict - team_name: str | None = None - - def _bundle_item_exc(msg): return AirflowConfigException( "Invalid config for section `dag_processor` key `dag_bundle_config_list`. " + msg @@ -108,40 +96,26 @@ def _parse_bundle_config(config_list) -> list[_ExternalBundleConfig]: return list(bundles.values()) +@cache +def _load_bundle_config_snapshot() -> tuple[_ExternalBundleConfig, ...]: + """Read the static bundle configuration without importing bundle classes.""" + config_list = conf.getjson("dag_processor", "dag_bundle_config_list") + if not config_list: + return () + if not isinstance(config_list, list): + raise AirflowConfigException( + "Section `dag_processor` key `dag_bundle_config_list` " + f"must be list but got {config_list.__class__}" + ) + return tuple(_parse_bundle_config(config_list)) + + class ConfigDagBundleProvider(DagBundleProvider): """Provide Dag bundles configured by ``dag_bundle_config_list``.""" def __init__(self) -> None: - self._bundle_config: dict[str, _InternalBundleConfig] = {} - self._load_configured_bundles() - - def _load_configured_bundles(self) -> None: - """ - Load Dag bundles from ``dag_bundle_config_list``. - - If a bundle class for a given name has already been imported, it will not be imported again. - - :meta private: - """ - if self._bundle_config: - return - - config_list = conf.getjson("dag_processor", "dag_bundle_config_list") - if not config_list: - return - if not isinstance(config_list, list): - raise AirflowConfigException( - "Section `dag_processor` key `dag_bundle_config_list` " - f"must be list but got {config_list.__class__}" - ) - bundle_config_list = _parse_bundle_config(config_list) - for bundle_config in bundle_config_list: - class_ = import_string(bundle_config.classpath) - self._bundle_config[bundle_config.name] = _InternalBundleConfig( - bundle_class=class_, - kwargs=bundle_config.kwargs, - team_name=bundle_config.team_name, - ) + self._bundle_config = {config.name: config for config in _load_bundle_config_snapshot()} + self._bundle_classes: dict[str, type[BaseDagBundle]] = {} def get_configured_bundle_metadata(self) -> Sequence[DagBundleMetadata]: return [ @@ -158,6 +132,8 @@ def get_bundle( cfg_bundle = self._bundle_config.get(name) if not cfg_bundle: raise ValueError(f"Requested bundle '{name}' is not configured.") - return cfg_bundle.bundle_class( + if name not in self._bundle_classes: + self._bundle_classes[name] = import_string(cfg_bundle.classpath) + return self._bundle_classes[name]( name=name, version=version, version_data=version_data, **cfg_bundle.kwargs ) diff --git a/airflow-core/src/airflow/dag_processing/dagbag.py b/airflow-core/src/airflow/dag_processing/dagbag.py index b5f3f0d5bc341..c12655efdd687 100644 --- a/airflow-core/src/airflow/dag_processing/dagbag.py +++ b/airflow-core/src/airflow/dag_processing/dagbag.py @@ -404,6 +404,7 @@ def _import_lang_sdk_definition(self, definition: DagDefinition) -> list[LangSDK path=fileloc, bundle_path=self._bundle.path, bundle_name=self._bundle.name, + team_name=self.team_name, dag_file_rel_path=relative_loc, logger=structlog.get_logger(logger_name=__name__), ) diff --git a/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py b/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py index 9e753e61059c1..44db76ac07dc4 100644 --- a/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py +++ b/airflow-core/src/airflow/dag_processing/lang_sdk_processor.py @@ -36,7 +36,7 @@ from airflow import settings from airflow.configuration import conf -from airflow.dag_processing.dagbag import _get_bundle_team_name, _validate_executor_fields +from airflow.dag_processing.dagbag import _validate_executor_fields from airflow.dag_processing.importer_routing import get_claiming_importer from airflow.dag_processing.processor import ( BaseDagFileProcessorProcess, @@ -187,6 +187,7 @@ def start( # type: ignore[override] bundle_path: Path, bundle_name: str, dag_file_rel_path: str, + team_name: str | None = None, **kwargs, ) -> Self: listeners: dict[_Channel, socket] = {"comm": _start_server(), "logs": _start_server()} @@ -194,7 +195,7 @@ def start( # type: ignore[override] for listener in listeners.values(): listener.setblocking(False) parse_request = DagFileParseRequest( - file=os.fspath(path), bundle_path=bundle_path, bundle_name=bundle_name + file=os.fspath(path), bundle_path=bundle_path, bundle_name=bundle_name, team_name=team_name ) proc = super().start( target=_start_runtime_entrypoint, @@ -239,6 +240,7 @@ def run( bundle_name: str, dag_file_rel_path: str, logger: FilteringBoundLogger, + team_name: str | None = None, ) -> DagFileParsingResult: """ Parse *path* outside the Dag processor and wait for the result. @@ -254,6 +256,7 @@ def run( path=path, bundle_path=bundle_path, bundle_name=bundle_name, + team_name=team_name, dag_file_rel_path=dag_file_rel_path, selector=selector, logger=logger, @@ -460,8 +463,9 @@ def _apply_team_rules(self, data: dict, dag: SerializedDAG) -> None: :raises UnknownExecutorException: if a task's executor is not available to the team or globally. """ - _validate_executor_fields(dag, self.bundle_name) - if not (team_name := _get_bundle_team_name(self.bundle_name)): + _validate_executor_fields(dag, self.bundle_name, self._parse_request.team_name) + team_name = self._parse_request.team_name if conf.getboolean("core", "multi_team") else None + if not team_name: return tasks = {task.task_id: task for task in dag.tasks} for encoded in data["dag"]["tasks"]: diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index 4bc185fe67f58..33e81c33f5f95 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -1489,6 +1489,7 @@ def _create_process( path=dag_file.absolute_path, bundle_path=cast("Path", dag_file.bundle_path), bundle_name=dag_file.bundle_name, + team_name=team_name, dag_file_rel_path=str(dag_file.rel_path), selector=self.selector, logger=logger, diff --git a/airflow-core/tests/unit/cli/commands/test_dag_processor_command.py b/airflow-core/tests/unit/cli/commands/test_dag_processor_command.py index 77de6f928d33e..42f1c5a16e34c 100644 --- a/airflow-core/tests/unit/cli/commands/test_dag_processor_command.py +++ b/airflow-core/tests/unit/cli/commands/test_dag_processor_command.py @@ -24,6 +24,7 @@ from airflow.cli import cli_parser from airflow.cli.commands import dag_processor_command +from airflow.dag_processing.bundles.provider import DagBundleMetadata, DagBundleProvider from tests_common.test_utils.config import conf_vars @@ -44,7 +45,10 @@ def dag_bundles_with_teams(): with ( conf_vars({("core", "multi_team"): "True"}), mock.patch.object( - dag_processor_command, "get_configured_bundle_team_names", return_value=dict(BUNDLE_TEAMS) + dag_processor_command.DagBundlesManager, + "get_configured_bundle_team_names", + autospec=True, + return_value=dict(BUNDLE_TEAMS), ), ): yield @@ -94,9 +98,24 @@ def test_bundle_names_passed(self, mock_runner, configure_testing_dag_bundle): def test_get_team_names(self, bundle_names, expected_team_names): assert dag_processor_command._get_team_names(bundle_names) == expected_team_names + @conf_vars({("core", "multi_team"): "True"}) + @mock.patch.object(dag_processor_command.DagBundlesManager, "_load_bundle_provider", autospec=True) + def test_get_team_names_uses_custom_provider_metadata(self, mock_load): + provider = mock.create_autospec(DagBundleProvider, instance=True) + provider.get_configured_bundle_metadata.return_value = [ + DagBundleMetadata(name="custom-bundle", team_name="custom-team") + ] + mock_load.return_value = provider + + assert dag_processor_command._get_team_names(["custom-bundle"]) == ["custom-team"] + provider.get_bundle.assert_not_called() + @conf_vars({("core", "multi_team"): "False"}) @mock.patch.object( - dag_processor_command, "get_configured_bundle_team_names", return_value=dict(BUNDLE_TEAMS) + dag_processor_command.DagBundlesManager, + "get_configured_bundle_team_names", + autospec=True, + return_value=dict(BUNDLE_TEAMS), ) def test_get_team_names_returns_empty_outside_multi_team(self, mock_configured): assert dag_processor_command._get_team_names(["bundle_a"]) == [] 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 4b02d5bfa4fb8..6ad8e46a3093d 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 @@ -36,7 +36,11 @@ _load_bundle_config_snapshot, get_configured_bundle_names, ) -from airflow.dag_processing.bundles.provider import DagBundleMetadata, DagBundleProvider +from airflow.dag_processing.bundles.provider import ( + ConfigDagBundleProvider, + DagBundleMetadata, + DagBundleProvider, +) from airflow.exceptions import AirflowConfigException from airflow.models.dag import DagModel from airflow.models.dag_version import DagVersion @@ -264,6 +268,31 @@ def test_get_bundle(): assert bundle.version is None +@conf_vars({("core", "load_examples"): "False"}) +@patch.object(DagBundlesManager, "_load_bundle_provider", autospec=True) +def test_is_bundle_configured_uses_current_provider_metadata(mock_load): + provider = mock.create_autospec(DagBundleProvider, instance=True) + mock_load.return_value = provider + provider.get_configured_bundle_metadata.return_value = [DagBundleMetadata(name="custom-bundle")] + + assert DagBundlesManager.is_bundle_configured("custom-bundle") is True + provider.get_configured_bundle_metadata.return_value = [] + assert DagBundlesManager.is_bundle_configured("custom-bundle") is False + provider.get_bundle.assert_not_called() + + +@conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(BASIC_BUNDLE_CONFIG)}) +@patch("airflow.dag_processing.bundles.provider.import_string", autospec=True, return_value=BasicBundle) +def test_config_provider_imports_only_requested_bundle(mock_import): + provider = ConfigDagBundleProvider() + assert provider.get_configured_bundle_metadata() == [DagBundleMetadata(name="my-test-bundle")] + mock_import.assert_not_called() + + assert isinstance(provider.get_bundle("my-test-bundle"), BasicBundle) + assert isinstance(provider.get_bundle("my-test-bundle", version="v1"), BasicBundle) + mock_import.assert_called_once_with(BASIC_BUNDLE_CONFIG[0]["classpath"]) + + def _bundle_config_env(config) -> dict[str, str]: return { "AIRFLOW__CORE__LOAD_EXAMPLES": "False", diff --git a/airflow-core/tests/unit/dag_processing/test_dagbag.py b/airflow-core/tests/unit/dag_processing/test_dagbag.py index 1cc949c089392..a50dc1bdb9204 100644 --- a/airflow-core/tests/unit/dag_processing/test_dagbag.py +++ b/airflow-core/tests/unit/dag_processing/test_dagbag.py @@ -1626,7 +1626,7 @@ def _native_runtime(**results_by_file: list[str] | DagFileParsingResult): A list names the Dags to serialize; a ``DagFileParsingResult`` is returned as it is. """ - def run(*, path, bundle_path, bundle_name, dag_file_rel_path, logger): + def run(*, path, bundle_path, bundle_name, dag_file_rel_path, logger, team_name=None): result = results_by_file[Path(path).stem] if isinstance(result, DagFileParsingResult): return result @@ -1645,11 +1645,12 @@ def run(*, path, bundle_path, bundle_name, dag_file_rel_path, logger): yield mock_run -def _native_bag(tmp_path, dag_folder=None) -> DagBag: +def _native_bag(tmp_path, dag_folder=None, team_name=None) -> DagBag: return DagBag( dag_folder=os.fspath(dag_folder or tmp_path), bundle_path=tmp_path, bundle_name="testing", + team_name=team_name, parse_lang_sdk_files=True, ) @@ -1707,12 +1708,13 @@ def test_the_runtime_parses_the_file_in_its_bundle(self, tmp_path): native = write_native_file(tmp_path / "sub" / "dags.native") with _native_runtime(dags=["native_a"]) as mock_run: - _native_bag(tmp_path) + _native_bag(tmp_path, team_name="team_a") mock_run.assert_called_once_with( path=os.fspath(native), bundle_path=tmp_path, bundle_name="testing", + team_name="team_a", dag_file_rel_path="sub/dags.native", logger=mock.ANY, ) diff --git a/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py b/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py index 8be3221b37874..9b8b0e8574a28 100644 --- a/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_lang_sdk_processor.py @@ -121,12 +121,15 @@ def _coordinator(): yield -def _start(tmp_path, selector, *, client: Client | None = None, **spec) -> LangSDKDagFileProcessorProcess: +def _start( + tmp_path, selector, *, client: Client | None = None, team_name: str | None = None, **spec +) -> LangSDKDagFileProcessorProcess: return LangSDKDagFileProcessorProcess.start( id=uuid.uuid4(), path=write_native_file(tmp_path / "dag.native", **spec), bundle_path=tmp_path, bundle_name="testing", + team_name=team_name, dag_file_rel_path="dag.native", selector=selector, logger=structlog.get_logger(), @@ -180,6 +183,22 @@ def test_parses_the_dags_the_runtime_returns(self, mock_parse_dag, parse, tmp_pa assert proc._subprocess_schema_version == OLDEST_SCHEMA_VERSION assert "Parsing the bundle" in cap_structlog + @conf_vars({("core", "multi_team"): "True"}) + @pytest.mark.parametrize("schema_version", [None, OLDEST_SCHEMA_VERSION]) + @patch.object(FakeCoordinator, "parse_dag", autospec=True) + def test_team_ownership_is_preserved_across_runtime_versions(self, mock_parse_dag, parse, schema_version): + mock_parse_dag.side_effect = play_runtime( + _reply_with(_serialize_dag("team_dag")), schema_version=schema_version + ) + + proc = parse(team_name="team_a") + + assert proc._parse_request.team_name == "team_a" + assert proc.parsing_result.import_errors is None + [stored] = proc.parsing_result.serialized_dags + dag = DagSerialization.from_dict(copy.deepcopy(stored.data)) + assert dag.task_dict["extract"].pool == "default_pool_team_a" + @patch.object(FakeCoordinatorDagImporter, "get_source_code", autospec=True) @patch.object(FakeCoordinator, "parse_dag", autospec=True) def test_each_dag_gets_its_own_source_keyed_by_dag_id(self, mock_parse_dag, mock_get_source_code, parse): @@ -660,7 +679,7 @@ def test_only_a_positive_import_timeout_applies(mock_timeout, configured, expect mock_timeout.assert_called_once_with("/b/dag.native") -def _make_process(**kwargs) -> LangSDKDagFileProcessorProcess: +def _make_process(*, team_name: str | None = None, **kwargs) -> LangSDKDagFileProcessorProcess: kwargs.setdefault("process", MagicMock(spec=PsutilTracker)) kwargs.setdefault("logger_filehandle", MagicMock(spec=BinaryIO)) return LangSDKDagFileProcessorProcess( @@ -673,7 +692,7 @@ def _make_process(**kwargs) -> LangSDKDagFileProcessorProcess: dag_file_rel_path="dag.native", listeners={}, parse_request=DagFileParseRequest( - file="/b/dag.native", bundle_path=Path("/b"), bundle_name="testing" + file="/b/dag.native", bundle_path=Path("/b"), bundle_name="testing", team_name=team_name ), **kwargs, ) @@ -845,7 +864,6 @@ def test_a_dag_with_an_unavailable_executor_is_an_import_error(mock_lookup): @conf_vars({("core", "multi_team"): "True"}) @patch("airflow.dag_processing.bundles.manager.DagBundlesManager", autospec=True) def test_tasks_in_the_default_pool_move_to_the_teams_pool(mock_bundles_manager): - mock_bundles_manager.return_value._bundle_config = {"testing": MagicMock(team_name="team_a")} with DAG("team_dag", schedule=None) as dag: BaseOperator(task_id="extract") BaseOperator(task_id="load", pool="custom") @@ -854,7 +872,7 @@ def test_tasks_in_the_default_pool_move_to_the_teams_pool(mock_bundles_manager): def fan_out(x): ... fan_out.expand(x=[1, 2]) - proc = _make_process() + proc = _make_process(team_name="team_a") proc._handle_request( DagFileParsingResult( @@ -872,3 +890,24 @@ def fan_out(x): ... "load": "custom", "fan_out": "default_pool_team_a", } + + mock_bundles_manager.assert_not_called() + + +@conf_vars({("core", "multi_team"): "True"}) +@patch.object(ExecutorLoader, "lookup_executor_name_by_str", autospec=True) +def test_executor_validation_uses_the_resolved_team(mock_lookup): + with DAG("team_dag", schedule=None) as dag: + BaseOperator(task_id="extract", executor="team.Executor") + proc = _make_process(team_name="team_a") + + proc._handle_request( + DagFileParsingResult( + fileloc="/b/dag.native", serialized_dags=[LazyDeserializedDAG(data=DagSerialization.to_dict(dag))] + ), + MagicMock(spec=socket.socket), + 1, + ) + + assert proc.parsing_result.import_errors is None + mock_lookup.assert_called_once_with("team.Executor", team_name="team_a", validate_teams=False) diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index 95f3b8130e65e..9cc5627d4051e 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -1691,11 +1691,12 @@ def test_create_process_parses_a_coordinator_file_with_its_runtime(self, mock_ge mock.patch.object(LangSDKDagFileProcessorProcess, "start", autospec=True) as mock_start, ): manager = DagFileProcessorManager(max_runs=1) - manager._create_process(dag_file) + manager._create_process(dag_file, team_name="team_a") kwargs = mock_start.call_args.kwargs assert (kwargs["path"], kwargs["dag_file_rel_path"]) == (tmp_path / "dags.native", "dags.native") assert kwargs["client"] is manager.client + assert kwargs["team_name"] == "team_a" @pytest.mark.parametrize("rel_path", ["my_dag.py", "dags.fake"]) @mock.patch.object(DagFileProcessorManager, "_get_logger_for_dag_file", autospec=True) diff --git a/devel-common/src/tests_common/pytest_plugin.py b/devel-common/src/tests_common/pytest_plugin.py index c82b2644551c0..9c9ebdb70738c 100644 --- a/devel-common/src/tests_common/pytest_plugin.py +++ b/devel-common/src/tests_common/pytest_plugin.py @@ -2216,7 +2216,7 @@ def reset_dag_bundle_config_cache(): return try: - from airflow.dag_processing.bundles.manager import _load_bundle_config_snapshot + from airflow.dag_processing.bundles.provider import _load_bundle_config_snapshot except ImportError: # compat for airflow versions without the snapshot cache yield diff --git a/devel-common/src/tests_common/test_utils/config.py b/devel-common/src/tests_common/test_utils/config.py index e5724ea65251c..c6c28f87ff777 100644 --- a/devel-common/src/tests_common/test_utils/config.py +++ b/devel-common/src/tests_common/test_utils/config.py @@ -125,9 +125,9 @@ 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") + provider = sys.modules.get("airflow.dag_processing.bundles.provider") # compat for airflow versions without the snapshot cache - cache = getattr(manager, "_load_bundle_config_snapshot", None) + cache = getattr(provider, "_load_bundle_config_snapshot", None) if cache is not None: cache.cache_clear() diff --git a/go-sdk/pkg/execution/genmodels/models.gen.go b/go-sdk/pkg/execution/genmodels/models.gen.go index 0a5dba694e675..c44ff38652a80 100644 --- a/go-sdk/pkg/execution/genmodels/models.gen.go +++ b/go-sdk/pkg/execution/genmodels/models.gen.go @@ -397,6 +397,9 @@ type DagFileParseRequest struct { // File corresponds to the JSON schema field "file". File string `msgpack:"file"` + // TeamName corresponds to the JSON schema field "team_name". + TeamName interface{} `msgpack:"team_name,omitempty"` + // Type corresponds to the JSON schema field "type". Type string `msgpack:"type,omitempty"` } diff --git a/go-sdk/schema/supervisor-schema.json b/go-sdk/schema/supervisor-schema.json index 41a6ab4c6d394..386b6700a1184 100644 --- a/go-sdk/schema/supervisor-schema.json +++ b/go-sdk/schema/supervisor-schema.json @@ -933,6 +933,18 @@ "title": "Bundle Name", "type": "string" }, + "team_name": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Team Name" + }, "callback_requests": { "items": { "discriminator": { diff --git a/task-sdk/src/airflow/sdk/execution_time/coordinator.py b/task-sdk/src/airflow/sdk/execution_time/coordinator.py index 78d13441b5fd7..95198d90a5629 100644 --- a/task-sdk/src/airflow/sdk/execution_time/coordinator.py +++ b/task-sdk/src/airflow/sdk/execution_time/coordinator.py @@ -294,7 +294,7 @@ def from_config(cls) -> Self: ): raise InvalidCoordinatorError( f"[sdk] coordinators {key!r} sets task_handler_bundle_name={bundle_name!r}, " - f"which is not a bundle in [dag_processor] dag_bundle_config_list" + f"which is not a configured Dag bundle" ) return cls(coordinator_specs=coordinator_specs, queue_to_coordinator=queue_to_coordinator) diff --git a/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py index 40c78f3d2dcd0..e3ff6cc9b9d7d 100644 --- a/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py +++ b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py @@ -444,7 +444,18 @@ class TestRealBundleArgBindingsDowngrade: @pytest.fixture def startup_details(self): - + import datetime + import uuid + + from airflow.sdk.api.datamodels._generated import ( + BundleInfo, + DagRun, + DagRunState, + DagRunType, + TaskInstance, + TIRunContext, + ) + from airflow.sdk.execution_time.comms import StartupDetails now = datetime.datetime.now(datetime.UTC) return StartupDetails( @@ -583,6 +594,66 @@ def test_upgrade_fills_missing_retry_reason_with_none(self, real_migrator): assert out["retry_reason"] is None +class TestRealBundleFirstTaskRescheduleStartDate: + """ + Drive the *real* supervisor bundle through the ``first_task_reschedule_start_date`` migration. + + ``AddFirstTaskRescheduleStartDateToSupervisorTIRunContext`` is ``didnt_exist``-only, so the + downgrade re-validation is the only place it is visible: dropping it from the bundle, or + registering it under the wrong version, leaves the field on the wire for a pinned runtime. + """ + + FIRST_RESCHEDULE = datetime.datetime(2026, 1, 1, tzinfo=datetime.UTC) + + @pytest.fixture + def startup_details(self): + now = datetime.datetime.now(datetime.UTC) + return StartupDetails( + ti=TaskInstance( + id=uuid.uuid4(), + task_id="wait", + dag_id="d", + run_id="r", + try_number=1, + dag_version_id=uuid.uuid4(), + ), + dag_rel_path="d.py", + bundle_info=BundleInfo(name="b", version=None), + start_date=now, + ti_context=TIRunContext( + dag_run=DagRun( + dag_id="d", + run_id="r", + logical_date=now, + data_interval_start=None, + data_interval_end=None, + start_date=now, + end_date=None, + run_type=DagRunType.MANUAL, + state=DagRunState.RUNNING, + run_after=now, + consumed_asset_events=[], + partition_key=None, + ), + max_tries=1, + first_task_reschedule_start_date=self.FIRST_RESCHEDULE, + ), + sentry_integration="", + ) + + @pytest.fixture + def real_migrator(self) -> SchemaVersionMigrator: + return get_schema_version_migrator() + + def test_downgrade_strips_field_for_previous_version(self, real_migrator, startup_details): + out = real_migrator.downgrade(startup_details, "2026-06-16").model_dump() + assert "first_task_reschedule_start_date" not in out["ti_context"] + + def test_head_version_keeps_field(self, real_migrator, startup_details): + out = real_migrator.downgrade(startup_details, "2026-10-30") + assert out.ti_context.first_task_reschedule_start_date == self.FIRST_RESCHEDULE + + class TestRealBundleDagFileParseRequestDowngrade: def test_downgrade_strips_team_name_for_previous_version(self): from pathlib import Path diff --git a/ts-sdk/schema/supervisor-schema.json b/ts-sdk/schema/supervisor-schema.json index 41a6ab4c6d394..386b6700a1184 100644 --- a/ts-sdk/schema/supervisor-schema.json +++ b/ts-sdk/schema/supervisor-schema.json @@ -933,6 +933,18 @@ "title": "Bundle Name", "type": "string" }, + "team_name": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Team Name" + }, "callback_requests": { "items": { "discriminator": { diff --git a/ts-sdk/src/generated/supervisor.ts b/ts-sdk/src/generated/supervisor.ts index 8b018e485149d..06619024fed99 100644 --- a/ts-sdk/src/generated/supervisor.ts +++ b/ts-sdk/src/generated/supervisor.ts @@ -202,6 +202,7 @@ export type Type12 = "DagCallbackRequest"; export type File = string; export type BundlePath = string; export type BundleName1 = string; +export type TeamName1 = string | null; export type Filepath1 = string; export type BundleName2 = string; export type BundleVersion1 = string | null; @@ -315,7 +316,7 @@ export type ConsumedAssetEvents1 = AssetEventDagRunReference[]; export type PartitionKey4 = string | null; export type PartitionDate1 = string | null; export type Note1 = string | null; -export type TeamName1 = string | null; +export type TeamName2 = string | null; export type Type18 = "DagRunResult"; export type Type19 = "DagRunStateResult"; export type State2 = "deferred"; @@ -1019,6 +1020,7 @@ export interface DagFileParseRequest { file: File; bundle_path: BundlePath; bundle_name: BundleName1; + team_name?: TeamName1; callback_requests?: CallbackRequests; type?: Type15; } @@ -1220,7 +1222,7 @@ export interface DagRunResult { partition_key: PartitionKey4; partition_date?: PartitionDate1; note?: Note1; - team_name?: TeamName1; + team_name?: TeamName2; type?: Type18; } /** From 23fcbbeb139ed77febb79854fd451abb9a99ce43 Mon Sep 17 00:00:00 2001 From: Sameer Raj Date: Sun, 11 Oct 2026 00:04:28 +0530 Subject: [PATCH 15/34] Keep team ownership when workers parse Dags Team-owned tasks must retain their executor and pool during worker parsing. The provider refactor otherwise leaves them without the ownership already resolved for their Dag run. --- .../src/airflow/dag_processing/dagbag.py | 9 +++++++- .../tests/unit/dag_processing/test_dagbag.py | 22 +++++++++++++++++++ .../airflow/sdk/execution_time/task_runner.py | 1 + .../execution_time/test_task_runner.py | 10 +++++---- 4 files changed, 37 insertions(+), 5 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/dagbag.py b/airflow-core/src/airflow/dag_processing/dagbag.py index c12655efdd687..f3a7fd60c5291 100644 --- a/airflow-core/src/airflow/dag_processing/dagbag.py +++ b/airflow-core/src/airflow/dag_processing/dagbag.py @@ -193,12 +193,19 @@ def __init__( bundle_path: Path | None = None, bundle_name: str | None = None, parse_lang_sdk_files: bool = False, - team_name: str | None = None, + team_name: str | None | ArgNotSet = NOTSET, ): super().__init__() self.bundle_path = bundle_path self.bundle_name = bundle_name self.parse_lang_sdk_files = parse_lang_sdk_files + # Explicit ownership, including None, avoids provider I/O in isolated processes. + if not is_arg_set(team_name): + team_name = None + if bundle_name and conf.getboolean("core", "multi_team"): + from airflow.dag_processing.bundles.manager import DagBundlesManager + + team_name = DagBundlesManager().get_configured_bundle_team_names().get(bundle_name) self.team_name = team_name dag_folder = dag_folder or settings.DAGS_FOLDER diff --git a/airflow-core/tests/unit/dag_processing/test_dagbag.py b/airflow-core/tests/unit/dag_processing/test_dagbag.py index a50dc1bdb9204..fa06500259a07 100644 --- a/airflow-core/tests/unit/dag_processing/test_dagbag.py +++ b/airflow-core/tests/unit/dag_processing/test_dagbag.py @@ -1872,3 +1872,25 @@ def test_runtime_failures_are_import_errors(self, tmp_path, spec, reply, message assert dagbag.dags == {} assert dagbag.import_errors == {"dags.native": message} + + +@conf_vars({("core", "multi_team"): "True"}) +@mock.patch("airflow.dag_processing.bundles.manager.DagBundlesManager", autospec=True) +def test_dagbag_resolves_omitted_team_once(mock_manager): + mock_manager.return_value.get_configured_bundle_team_names.return_value = {"team_bundle": "team_a"} + + bag = DagBag(bundle_name="team_bundle", collect_dags=False) + + assert bag.team_name == "team_a" + mock_manager.assert_called_once_with() + mock_manager.return_value.get_configured_bundle_team_names.assert_called_once_with() + + +@conf_vars({("core", "multi_team"): "True"}) +@pytest.mark.parametrize("team_name", [None, "team_a"]) +@mock.patch("airflow.dag_processing.bundles.manager.DagBundlesManager", autospec=True) +def test_dagbag_does_not_discover_explicit_team(mock_manager, team_name): + bag = DagBag(bundle_name="team_bundle", team_name=team_name, collect_dags=False) + + assert bag.team_name == team_name + mock_manager.assert_not_called() 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 bfa6e493434e5..3998b85923ee2 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -1135,6 +1135,7 @@ def parse(what: StartupDetails, log: Logger) -> RuntimeTaskInstance: load_op_links=False, bundle_path=bundle_instance.path, bundle_name=bundle_info.name, + team_name=what.ti_context.dag_run.team_name, ) dag_file_parse_ms = int((time.monotonic() - dag_file_parse_start) * 1000) if TYPE_CHECKING: 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 bf64c087a3e72..0e132f8c8db8f 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 @@ -290,11 +290,11 @@ def test_parse(test_dags_dir: Path, make_ti_context, coordinates, map_index): assert ti.map_index == map_index -@mock.patch("airflow.dag_processing.dagbag.BundleDagBag") -def test_parse_dag_bag(mock_dagbag, test_dags_dir: Path, make_ti_context): +@pytest.mark.parametrize("team_name", [None, "team_a"]) +@mock.patch("airflow.dag_processing.dagbag.BundleDagBag", autospec=True) +def test_parse_dag_bag(mock_dagbag, test_dags_dir: Path, make_ti_context, team_name): """Test that checks that the BundleDagBag is constructed as expected during parsing""" - mock_bag_instance = mock.Mock() - mock_dagbag.return_value = mock_bag_instance + mock_bag_instance = mock_dagbag.return_value mock_dag = mock.Mock(spec=DAG) mock_task = mock.Mock(spec=BaseOperator) # The worker walks dag.tasks to register declared deserialization classes; @@ -322,6 +322,7 @@ def test_parse_dag_bag(mock_dagbag, test_dags_dir: Path, make_ti_context): start_date=timezone.utcnow(), sentry_integration="", ) + what.ti_context.dag_run.team_name = team_name with patch.dict( os.environ, @@ -345,6 +346,7 @@ def test_parse_dag_bag(mock_dagbag, test_dags_dir: Path, make_ti_context): load_op_links=False, bundle_path=test_dags_dir, bundle_name="my-bundle", + team_name=team_name, ) From 0050a7e33884aff62d36fce677aaf3b1bdca4d30 Mon Sep 17 00:00:00 2001 From: Sameer Raj Date: Tue, 13 Oct 2026 00:00:10 +0530 Subject: [PATCH 16/34] Clarify Dag bundle discovery and preserve active bundles Deployments need an unambiguous configuration source contract. A temporary loading failure must not disable an existing configured bundle or discard its ownership. --- .../dag-bundles.rst | 53 +++++ .../src/airflow/config_templates/config.yml | 6 +- .../airflow/dag_processing/bundles/manager.py | 55 ++--- .../bundles/{provider.py => source.py} | 4 +- .../commands/test_dag_processor_command.py | 14 +- .../bundles/test_dag_bundle_manager.py | 201 +++++++++++++----- .../src/tests_common/pytest_plugin.py | 2 +- .../src/tests_common/test_utils/config.py | 4 +- 8 files changed, 249 insertions(+), 90 deletions(-) rename airflow-core/src/airflow/dag_processing/bundles/{provider.py => source.py} (98%) diff --git a/airflow-core/docs/administration-and-deployment/dag-bundles.rst b/airflow-core/docs/administration-and-deployment/dag-bundles.rst index 67e14fd01fc4c..47d7389213eab 100644 --- a/airflow-core/docs/administration-and-deployment/dag-bundles.rst +++ b/airflow-core/docs/administration-and-deployment/dag-bundles.rst @@ -64,6 +64,59 @@ Configuring Dag bundles Dag bundles are configured in :ref:`config:dag_processor__dag_bundle_config_list`. You can add one or more Dag bundles here. +Writing a bundle source +~~~~~~~~~~~~~~~~~~~~~~~ + +Set :ref:`config:dag_processor__dag_bundle_source` to the import path of a class that extends +``DagBundleSource``. The default, ``ConfigDagBundleSource``, reads ``dag_bundle_config_list``. +This is a class extension point; it does not require a new Airflow provider package. + +The class must support construction without arguments and implement two methods: + +* ``get_configured_bundle_metadata()`` returns the complete list of configured bundle names and, + in multi-team mode, their team names. Airflow manages validation, example bundles, and database state. +* ``get_bundle(name, version, version_data)`` constructs the requested bundle. The source chooses + its class and settings and forwards both version arguments to the bundle. + +For example, save this minimal source in an importable module named ``my_deployment.bundles``: + +.. code-block:: python + + from airflow.dag_processing.bundles.local import LocalDagBundle + from airflow.dag_processing.bundles.source import DagBundleMetadata, DagBundleSource + + + class DeploymentDagBundleSource(DagBundleSource): + def get_configured_bundle_metadata(self): + return [DagBundleMetadata(name="team-dags")] + + def get_bundle(self, name, version=None, version_data=None): + if name != "team-dags": + raise ValueError(f"Unknown bundle: {name}") + return LocalDagBundle( + name=name, + version=version, + version_data=version_data, + path="/opt/airflow/team-dags", + ) + +Configure it with: + +.. code-block:: ini + + [dag_processor] + dag_bundle_source = my_deployment.bundles.DeploymentDagBundleSource + +Install the module in each Airflow process that resolves bundles, including Dag processors and workers. +A source can read deployment-controlled configuration from a file or service instead of using the fixed +values above. Raise an exception when discovery fails; return an empty list only when no bundles are +configured. Do not include Airflow's example bundles in that list. + +Keep a bundle name tied to the same implementation and construction settings. If a bundle leaves the +configured list, its name must remain resolvable while retained Dag runs can still need its code. +Keep access to the configuration source under Deployment Manager control and reference credentials +through Airflow Connections rather than embedding them in bundle settings. + .. warning:: Reference credentials through a Connection — do not inline them Bundle ``kwargs`` are stored in the ``[dag_processor] dag_bundle_config_list`` diff --git a/airflow-core/src/airflow/config_templates/config.yml b/airflow-core/src/airflow/config_templates/config.yml index 978b12ef0511d..0d73925f56604 100644 --- a/airflow-core/src/airflow/config_templates/config.yml +++ b/airflow-core/src/airflow/config_templates/config.yml @@ -3277,10 +3277,10 @@ dag_processor: example: "/tmp/some-place" default: ~ - dag_bundle_provider: + dag_bundle_source: description: | Import path for the class that provides configured Dag bundle metadata and constructs Dag bundles. - The class must extend ``airflow.dag_processing.bundles.provider.DagBundleProvider`` and have a + The class must extend ``airflow.dag_processing.bundles.source.DagBundleSource`` and have a no-argument constructor. The default implementation reads ``dag_bundle_config_list``. A custom implementation can read @@ -3288,7 +3288,7 @@ dag_processor: version_added: 3.4.0 type: string example: ~ - default: "airflow.dag_processing.bundles.provider.ConfigDagBundleProvider" + default: "airflow.dag_processing.bundles.source.ConfigDagBundleSource" dag_bundle_config_list: description: | diff --git a/airflow-core/src/airflow/dag_processing/bundles/manager.py b/airflow-core/src/airflow/dag_processing/bundles/manager.py index 7cf4f3ee96041..bef4e250f59d1 100644 --- a/airflow-core/src/airflow/dag_processing/bundles/manager.py +++ b/airflow-core/src/airflow/dag_processing/bundles/manager.py @@ -31,9 +31,9 @@ from airflow.configuration import conf from airflow.dag_processing.bundles.base import BaseDagBundle # noqa: TC001 from airflow.dag_processing.bundles.local import LocalDagBundle -from airflow.dag_processing.bundles.provider import ( +from airflow.dag_processing.bundles.source import ( DagBundleMetadata, - DagBundleProvider, + DagBundleSource, ) from airflow.exceptions import AirflowConfigException from airflow.models.dagbundle import DagBundleModel @@ -177,7 +177,7 @@ class DagBundlesManager(LoggingMixin): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - self._bundle_provider = self._load_bundle_provider() + self._bundle_source = self._load_bundle_source() @cached_property def _example_dag_bundle_paths(self) -> dict[str, str]: @@ -186,30 +186,30 @@ def _example_dag_bundle_paths(self) -> dict[str, str]: return _discover_example_dag_bundle_paths() @staticmethod - def _load_bundle_provider() -> DagBundleProvider: - provider_class = conf.getimport("dag_processor", "dag_bundle_provider") - if not isinstance(provider_class, type) or not issubclass(provider_class, DagBundleProvider): + def _load_bundle_source() -> DagBundleSource: + source_class = conf.getimport("dag_processor", "dag_bundle_source") + if not isinstance(source_class, type) or not issubclass(source_class, DagBundleSource): raise AirflowConfigException( - "The class configured in section `dag_processor` key `dag_bundle_provider` " - "must extend DagBundleProvider." + "The class configured in section `dag_processor` key `dag_bundle_source` " + "must extend DagBundleSource." ) try: - return provider_class() + return source_class() except AirflowConfigException: raise except Exception as e: raise AirflowConfigException( - "Could not initialize the Dag bundle provider configured in section " - "`dag_processor` key `dag_bundle_provider`." + "Could not initialize the Dag bundle source configured in section " + "`dag_processor` key `dag_bundle_source`." ) from e - def _get_provider_bundle_metadata(self) -> tuple[DagBundleMetadata, ...]: - """Get validated metadata returned by the configured Dag bundle provider.""" + def _get_source_bundle_metadata(self) -> tuple[DagBundleMetadata, ...]: + """Get validated metadata returned by the configured Dag bundle source.""" bundle_metadata: dict[str, DagBundleMetadata] = {} - for metadata in self._bundle_provider.get_configured_bundle_metadata(): + for metadata in self._bundle_source.get_configured_bundle_metadata(): if not isinstance(metadata, DagBundleMetadata): raise AirflowConfigException( - "Dag bundle providers must return DagBundleMetadata objects from " + "Dag bundle sources must return DagBundleMetadata objects from " "get_configured_bundle_metadata()." ) if metadata.name == _example_dag_bundle_name: @@ -219,7 +219,7 @@ def _get_provider_bundle_metadata(self) -> tuple[DagBundleMetadata, ...]: ) if metadata.name in bundle_metadata: raise AirflowConfigException( - f"Dag bundle provider returned duplicate bundle name '{metadata.name}'." + f"Dag bundle source returned duplicate bundle name '{metadata.name}'." ) if metadata.team_name and not conf.getboolean("core", "multi_team"): raise AirflowConfigException( @@ -232,11 +232,10 @@ def _get_provider_bundle_metadata(self) -> tuple[DagBundleMetadata, ...]: def get_active_bundle_metadata(self) -> tuple[DagBundleMetadata, ...]: """Get metadata for all active Dag bundles, including example bundles.""" - bundle_metadata = {metadata.name: metadata for metadata in self._get_provider_bundle_metadata()} + bundle_metadata = {metadata.name: metadata for metadata in self._get_source_bundle_metadata()} if not bundle_metadata: return () - # Airflow owns example Dag bundles; configured providers do not need to include them. for name in self._example_dag_bundle_paths: if name in bundle_metadata: raise AirflowConfigException(f"Bundle name '{name}' is reserved for example Dags.") @@ -245,8 +244,8 @@ def get_active_bundle_metadata(self) -> tuple[DagBundleMetadata, ...]: return tuple(bundle_metadata.values()) def get_configured_bundle_team_names(self) -> dict[str, str | None]: - """Get the team owning each bundle declared by the configured provider.""" - return {metadata.name: metadata.team_name for metadata in self._get_provider_bundle_metadata()} + """Get the team owning each bundle declared by the configured source.""" + return {metadata.name: metadata.team_name for metadata in self._get_source_bundle_metadata()} @provide_session def sync_bundles_to_db(self, *, deactivate_missing: bool = True, session: Session = NEW_SESSION) -> None: @@ -304,18 +303,22 @@ def _extract_and_sign_template(bundle_name: str) -> tuple[str | None, dict]: f"Team '{metadata.team_name}' configured for Dag bundle '{name}' does not exist." ) + bundle = stored.pop(name, None) + update_template = True try: new_template, new_params = _extract_and_sign_template(name) except Exception as e: self.log.exception("Error creating bundle '%s': %s", name, e) - continue + if bundle is None: + continue + update_template = False - if bundle := stored.pop(name, None): + if bundle is not None: bundle.active = True - if new_template != bundle.signed_url_template: + if update_template and new_template != bundle.signed_url_template: bundle.signed_url_template = new_template self.log.debug("Updated URL template for bundle %s", name) - if new_params != bundle.template_params: + if update_template and new_params != bundle.template_params: bundle.template_params = new_params self.log.debug("Updated template parameters for bundle %s", name) else: @@ -521,7 +524,7 @@ def reassign_dags_with_unconfigured_bundles(self) -> int: self.log.warning( "Skipped %d legacy Dag(s) whose fileloc is not under any configured bundle; " "triggering them will keep raising \"Requested bundle '{name}' is not configured.\" " - "until the configured Dag bundle provider returns a bundle whose path contains the " + "until the configured Dag bundle source returns a bundle whose path contains the " "fileloc. The next parse will then restore them automatically, " "or run `airflow dags reserialize` to force the parse path to rewrite " "bundle_name and relative_fileloc immediately.", @@ -597,7 +600,7 @@ def get_bundle( """ if path := self._example_dag_bundle_paths.get(name): return LocalDagBundle(name=name, path=path, version=version, version_data=version_data) - return self._bundle_provider.get_bundle(name=name, version=version, version_data=version_data) + return self._bundle_source.get_bundle(name=name, version=version, version_data=version_data) @classmethod def is_bundle_configured(cls, name: str) -> bool: diff --git a/airflow-core/src/airflow/dag_processing/bundles/provider.py b/airflow-core/src/airflow/dag_processing/bundles/source.py similarity index 98% rename from airflow-core/src/airflow/dag_processing/bundles/provider.py rename to airflow-core/src/airflow/dag_processing/bundles/source.py index caacba1aeea65..1edc0f12502ff 100644 --- a/airflow-core/src/airflow/dag_processing/bundles/provider.py +++ b/airflow-core/src/airflow/dag_processing/bundles/source.py @@ -40,7 +40,7 @@ class DagBundleMetadata: team_name: str | None = None -class DagBundleProvider(ABC): +class DagBundleSource(ABC): """ Provide configured Dag bundle metadata and construct Dag bundles. @@ -110,7 +110,7 @@ def _load_bundle_config_snapshot() -> tuple[_ExternalBundleConfig, ...]: return tuple(_parse_bundle_config(config_list)) -class ConfigDagBundleProvider(DagBundleProvider): +class ConfigDagBundleSource(DagBundleSource): """Provide Dag bundles configured by ``dag_bundle_config_list``.""" def __init__(self) -> None: diff --git a/airflow-core/tests/unit/cli/commands/test_dag_processor_command.py b/airflow-core/tests/unit/cli/commands/test_dag_processor_command.py index 42f1c5a16e34c..3101be8e63264 100644 --- a/airflow-core/tests/unit/cli/commands/test_dag_processor_command.py +++ b/airflow-core/tests/unit/cli/commands/test_dag_processor_command.py @@ -24,7 +24,7 @@ from airflow.cli import cli_parser from airflow.cli.commands import dag_processor_command -from airflow.dag_processing.bundles.provider import DagBundleMetadata, DagBundleProvider +from airflow.dag_processing.bundles.source import DagBundleMetadata, DagBundleSource from tests_common.test_utils.config import conf_vars @@ -99,16 +99,16 @@ def test_get_team_names(self, bundle_names, expected_team_names): assert dag_processor_command._get_team_names(bundle_names) == expected_team_names @conf_vars({("core", "multi_team"): "True"}) - @mock.patch.object(dag_processor_command.DagBundlesManager, "_load_bundle_provider", autospec=True) - def test_get_team_names_uses_custom_provider_metadata(self, mock_load): - provider = mock.create_autospec(DagBundleProvider, instance=True) - provider.get_configured_bundle_metadata.return_value = [ + @mock.patch.object(dag_processor_command.DagBundlesManager, "_load_bundle_source", autospec=True) + def test_get_team_names_uses_custom_source_metadata(self, mock_load): + source = mock.create_autospec(DagBundleSource, instance=True) + source.get_configured_bundle_metadata.return_value = [ DagBundleMetadata(name="custom-bundle", team_name="custom-team") ] - mock_load.return_value = provider + mock_load.return_value = source assert dag_processor_command._get_team_names(["custom-bundle"]) == ["custom-team"] - provider.get_bundle.assert_not_called() + source.get_bundle.assert_not_called() @conf_vars({("core", "multi_team"): "False"}) @mock.patch.object( 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 6ad8e46a3093d..1908736e3e2f1 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 @@ -36,16 +36,17 @@ _load_bundle_config_snapshot, get_configured_bundle_names, ) -from airflow.dag_processing.bundles.provider import ( - ConfigDagBundleProvider, +from airflow.dag_processing.bundles.source import ( + ConfigDagBundleSource, DagBundleMetadata, - DagBundleProvider, + DagBundleSource, ) from airflow.exceptions import AirflowConfigException from airflow.models.dag import DagModel from airflow.models.dag_version import DagVersion from airflow.models.dagbundle import DagBundleModel from airflow.models.errors import ParseImportError +from airflow.models.team import Team from airflow.utils.session import create_session from tests_common.test_utils.config import conf_vars @@ -129,7 +130,7 @@ def path(self) -> Path: return Path("/__basic_bundle_unmatched__") -class CustomDagBundleProvider(DagBundleProvider): +class CustomDagBundleSource(DagBundleSource): def __init__(self): self.metadata_requests = 0 @@ -148,7 +149,7 @@ def get_bundle(self, name, version=None, version_data=None): ) -class DuplicateMetadataDagBundleProvider(CustomDagBundleProvider): +class DuplicateMetadataDagBundleSource(CustomDagBundleSource): def get_configured_bundle_metadata(self): return [ DagBundleMetadata(name="duplicate"), @@ -156,19 +157,19 @@ def get_configured_bundle_metadata(self): ] -class NonMetadataDagBundleProvider(CustomDagBundleProvider): +class NonMetadataDagBundleSource(CustomDagBundleSource): def get_configured_bundle_metadata(self): return [object()] -class MissingTeamDagBundleProvider(CustomDagBundleProvider): +class MissingTeamDagBundleSource(CustomDagBundleSource): def get_configured_bundle_metadata(self): return [DagBundleMetadata(name="active-bundle", team_name="missing-team")] -class FailingInitializationDagBundleProvider(CustomDagBundleProvider): +class FailingInitializationDagBundleSource(CustomDagBundleSource): def __init__(self): - raise RuntimeError("Provider initialization failed") + raise RuntimeError("Source initialization failed") BASIC_BUNDLE_CONFIG = [ @@ -269,27 +270,27 @@ def test_get_bundle(): @conf_vars({("core", "load_examples"): "False"}) -@patch.object(DagBundlesManager, "_load_bundle_provider", autospec=True) -def test_is_bundle_configured_uses_current_provider_metadata(mock_load): - provider = mock.create_autospec(DagBundleProvider, instance=True) - mock_load.return_value = provider - provider.get_configured_bundle_metadata.return_value = [DagBundleMetadata(name="custom-bundle")] +@patch.object(DagBundlesManager, "_load_bundle_source", autospec=True) +def test_is_bundle_configured_uses_current_source_metadata(mock_load): + source = mock.create_autospec(DagBundleSource, instance=True) + mock_load.return_value = source + source.get_configured_bundle_metadata.return_value = [DagBundleMetadata(name="custom-bundle")] assert DagBundlesManager.is_bundle_configured("custom-bundle") is True - provider.get_configured_bundle_metadata.return_value = [] + source.get_configured_bundle_metadata.return_value = [] assert DagBundlesManager.is_bundle_configured("custom-bundle") is False - provider.get_bundle.assert_not_called() + source.get_bundle.assert_not_called() @conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(BASIC_BUNDLE_CONFIG)}) -@patch("airflow.dag_processing.bundles.provider.import_string", autospec=True, return_value=BasicBundle) -def test_config_provider_imports_only_requested_bundle(mock_import): - provider = ConfigDagBundleProvider() - assert provider.get_configured_bundle_metadata() == [DagBundleMetadata(name="my-test-bundle")] +@patch("airflow.dag_processing.bundles.source.import_string", autospec=True, return_value=BasicBundle) +def test_config_source_imports_only_requested_bundle(mock_import): + source = ConfigDagBundleSource() + assert source.get_configured_bundle_metadata() == [DagBundleMetadata(name="my-test-bundle")] mock_import.assert_not_called() - assert isinstance(provider.get_bundle("my-test-bundle"), BasicBundle) - assert isinstance(provider.get_bundle("my-test-bundle", version="v1"), BasicBundle) + assert isinstance(source.get_bundle("my-test-bundle"), BasicBundle) + assert isinstance(source.get_bundle("my-test-bundle", version="v1"), BasicBundle) mock_import.assert_called_once_with(BASIC_BUNDLE_CONFIG[0]["classpath"]) @@ -340,18 +341,18 @@ def test_is_bundle_configured_sees_implicitly_added_bundles(): ("core", "LOAD_EXAMPLES"): "False", ( "dag_processor", - "dag_bundle_provider", - ): "unit.dag_processing.bundles.test_dag_bundle_manager.CustomDagBundleProvider", + "dag_bundle_source", + ): "unit.dag_processing.bundles.test_dag_bundle_manager.CustomDagBundleSource", } ) -def test_custom_bundle_provider_resolves_active_and_retired_bundles(): +def test_custom_bundle_source_resolves_active_and_retired_bundles(): manager = DagBundlesManager() - provider = manager._bundle_provider - assert isinstance(provider, CustomDagBundleProvider) - assert provider.metadata_requests == 0 + source = manager._bundle_source + assert isinstance(source, CustomDagBundleSource) + assert source.metadata_requests == 0 assert manager.get_active_bundle_metadata() == (DagBundleMetadata(name="active-bundle"),) - assert provider.metadata_requests == 1 + assert source.metadata_requests == 1 assert manager.get_all_bundle_names() == ["active-bundle"] active_bundle = manager.get_bundle("active-bundle") @@ -368,43 +369,43 @@ def test_custom_bundle_provider_resolves_active_and_retired_bundles(): @pytest.mark.parametrize( - ("provider_class", "expected_message"), + ("source_class", "expected_message"), [ pytest.param( - "unit.dag_processing.bundles.test_dag_bundle_manager.DuplicateMetadataDagBundleProvider", + "unit.dag_processing.bundles.test_dag_bundle_manager.DuplicateMetadataDagBundleSource", "duplicate bundle name 'duplicate'", id="duplicate-name", ), pytest.param( - "unit.dag_processing.bundles.test_dag_bundle_manager.NonMetadataDagBundleProvider", + "unit.dag_processing.bundles.test_dag_bundle_manager.NonMetadataDagBundleSource", "must return DagBundleMetadata objects", id="wrong-metadata-type", ), ], ) -def test_custom_bundle_provider_rejects_invalid_metadata(provider_class, expected_message): - with conf_vars({("dag_processor", "dag_bundle_provider"): provider_class}): +def test_custom_bundle_source_rejects_invalid_metadata(source_class, expected_message): + with conf_vars({("dag_processor", "dag_bundle_source"): source_class}): with pytest.raises(AirflowConfigException, match=expected_message): DagBundlesManager().get_active_bundle_metadata() @pytest.mark.parametrize( - ("provider_class", "expected_message"), + ("source_class", "expected_message"), [ pytest.param( "unit.dag_processing.bundles.test_dag_bundle_manager.BasicBundle", - "must extend DagBundleProvider", + "must extend DagBundleSource", id="wrong-base-class", ), pytest.param( - "unit.dag_processing.bundles.test_dag_bundle_manager.FailingInitializationDagBundleProvider", - "Could not initialize the Dag bundle provider", + "unit.dag_processing.bundles.test_dag_bundle_manager.FailingInitializationDagBundleSource", + "Could not initialize the Dag bundle source", id="initialization-failure", ), ], ) -def test_invalid_custom_bundle_provider(provider_class, expected_message): - with conf_vars({("dag_processor", "dag_bundle_provider"): provider_class}): +def test_invalid_custom_bundle_source(source_class, expected_message): + with conf_vars({("dag_processor", "dag_bundle_source"): source_class}): with pytest.raises(AirflowConfigException, match=expected_message): DagBundlesManager() @@ -465,11 +466,11 @@ def _get_bundle_names_and_active(): ("core", "multi_team"): "True", ( "dag_processor", - "dag_bundle_provider", - ): "unit.dag_processing.bundles.test_dag_bundle_manager.MissingTeamDagBundleProvider", + "dag_bundle_source", + ): "unit.dag_processing.bundles.test_dag_bundle_manager.MissingTeamDagBundleSource", } ) -def test_sync_bundles_to_db_reports_missing_team_from_provider(clear_db): +def test_sync_bundles_to_db_reports_missing_team_from_source(clear_db): manager = DagBundlesManager() with pytest.raises( @@ -582,6 +583,81 @@ def path(self): ] +@pytest.mark.db_test +@conf_vars({("core", "load_examples"): "False"}) +@pytest.mark.parametrize("active", [True, False]) +@pytest.mark.parametrize( + "classpath", + [ + pytest.param("airflow.dag_processing.bundles.local.MissingBundle", id="invalid-classpath"), + pytest.param( + "unit.dag_processing.bundles.test_dag_bundle_manager.FailingBundle", + id="constructor-failure", + ), + pytest.param( + "unit.dag_processing.bundles.test_dag_bundle_manager.BundleWithTemplate", + id="template-failure", + ), + ], +) +@mock.patch.object( + BundleWithTemplate, "view_url_template", autospec=True, side_effect=ValueError("Template unavailable") +) +def test_sync_bundles_to_db_preserves_existing_bundle_on_load_failure( + mock_view_url_template, clear_db, session, active, classpath +): + bundle = DagBundleModel(name="existing-bundle") + bundle.active = active + bundle.signed_url_template = "previous-signed-template" + bundle.template_params = {"subdir": "dags"} + import_error = ParseImportError( + bundle_name=bundle.name, filename="dag.py", stacktrace="Previous parsing error" + ) + session.add_all([bundle, import_error]) + session.flush() + config = [{"name": bundle.name, "classpath": classpath, "kwargs": {"refresh_interval": 1}}] + + with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(config)}): + DagBundlesManager().sync_bundles_to_db(session=session) + session.flush() + session.expire_all() + + assert bundle.active is True + assert bundle.signed_url_template == "previous-signed-template" + assert bundle.template_params == {"subdir": "dags"} + assert session.scalar(select(ParseImportError.id).where(ParseImportError.bundle_name == bundle.name)) == ( + import_error.id + ) + + +@pytest.mark.db_test +@conf_vars({("core", "load_examples"): "False", ("core", "multi_team"): "True"}) +@pytest.mark.parametrize("team_name", ["team-a", "team-b", None]) +def test_sync_bundles_to_db_reconciles_team_on_load_failure(clear_db, session, team_name): + team_a = Team(name="team-a") + team_b = Team(name="team-b") + bundle = DagBundleModel(name="existing-bundle") + bundle.teams = [team_a] + session.add_all([bundle, team_b]) + session.flush() + config = [ + { + "name": bundle.name, + "classpath": "airflow.dag_processing.bundles.local.MissingBundle", + "kwargs": {}, + "team_name": team_name, + } + ] + + with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(config)}): + DagBundlesManager().sync_bundles_to_db(session=session) + session.flush() + session.expire_all() + + assert bundle.active is True + assert [team.name for team in bundle.teams] == ([team_name] if team_name else []) + + @pytest.mark.db_test @conf_vars({("core", "LOAD_EXAMPLES"): "False"}) def test_sync_bundles_to_db_with_template(clear_db, session): @@ -739,15 +815,42 @@ def test_example_dags_not_added_without_configured_bundles(): { ( "dag_processor", - "dag_bundle_provider", - ): "unit.dag_processing.bundles.test_dag_bundle_manager.CustomDagBundleProvider" + "dag_bundle_source", + ): "unit.dag_processing.bundles.test_dag_bundle_manager.CustomDagBundleSource" } ) -def test_example_dags_bundle_added_for_custom_provider(): - manager = DagBundlesManager() +@pytest.mark.parametrize("load_examples", [True, False]) +def test_example_dags_bundle_added_for_custom_source(load_examples): + with conf_vars({("core", "load_examples"): str(load_examples)}): + manager = DagBundlesManager() - assert "example_dags" in manager.get_all_bundle_names() - assert isinstance(manager.get_bundle("example_dags"), LocalDagBundle) + assert ("example_dags" in manager.get_all_bundle_names()) is load_examples + if load_examples: + assert isinstance(manager.get_bundle("example_dags"), LocalDagBundle) + + +@conf_vars({("core", "load_examples"): "True"}) +@patch.object(DagBundlesManager, "_load_bundle_source", autospec=True) +def test_example_dags_not_added_without_custom_bundles(mock_load): + source = mock.create_autospec(DagBundleSource, instance=True) + mock_load.return_value = source + source.get_configured_bundle_metadata.return_value = [] + + assert DagBundlesManager().get_active_bundle_metadata() == () + source.get_bundle.assert_not_called() + + +@pytest.mark.parametrize("load_examples", [True, False]) +@patch.object(DagBundlesManager, "_load_bundle_source", autospec=True) +def test_custom_source_cannot_declare_example_dags(mock_load, load_examples): + source = mock.create_autospec(DagBundleSource, instance=True) + mock_load.return_value = source + source.get_configured_bundle_metadata.return_value = [DagBundleMetadata(name="example_dags")] + + with conf_vars({("core", "load_examples"): str(load_examples)}): + with pytest.raises(AirflowConfigException, match="Bundle name 'example_dags' is a reserved name"): + DagBundlesManager().get_active_bundle_metadata() + source.get_bundle.assert_not_called() def test_example_dags_name_is_reserved(): diff --git a/devel-common/src/tests_common/pytest_plugin.py b/devel-common/src/tests_common/pytest_plugin.py index 9c9ebdb70738c..4809cfd267d54 100644 --- a/devel-common/src/tests_common/pytest_plugin.py +++ b/devel-common/src/tests_common/pytest_plugin.py @@ -2216,7 +2216,7 @@ def reset_dag_bundle_config_cache(): return try: - from airflow.dag_processing.bundles.provider import _load_bundle_config_snapshot + from airflow.dag_processing.bundles.source import _load_bundle_config_snapshot except ImportError: # compat for airflow versions without the snapshot cache yield diff --git a/devel-common/src/tests_common/test_utils/config.py b/devel-common/src/tests_common/test_utils/config.py index c6c28f87ff777..7b47b3835a624 100644 --- a/devel-common/src/tests_common/test_utils/config.py +++ b/devel-common/src/tests_common/test_utils/config.py @@ -125,9 +125,9 @@ def _clear_dag_bundle_config_cache() -> None: """Drop the per-process Dag bundle configuration cache so the new config is read.""" import sys - provider = sys.modules.get("airflow.dag_processing.bundles.provider") + source = sys.modules.get("airflow.dag_processing.bundles.source") # compat for airflow versions without the snapshot cache - cache = getattr(provider, "_load_bundle_config_snapshot", None) + cache = getattr(source, "_load_bundle_config_snapshot", None) if cache is not None: cache.cache_clear() From b47940c4ac9a406c7400c5a2c3be5ad05afdb058 Mon Sep 17 00:00:00 2001 From: Sameer Raj Date: Tue, 13 Oct 2026 00:13:43 +0530 Subject: [PATCH 17/34] Mark Dag bundle source support for Airflow 3.5 Deployments need the configuration reference to identify the first intended release for bundle source support. --- airflow-core/src/airflow/config_templates/config.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow-core/src/airflow/config_templates/config.yml b/airflow-core/src/airflow/config_templates/config.yml index 0d73925f56604..63ea8ec2f401e 100644 --- a/airflow-core/src/airflow/config_templates/config.yml +++ b/airflow-core/src/airflow/config_templates/config.yml @@ -3285,7 +3285,7 @@ dag_processor: The default implementation reads ``dag_bundle_config_list``. A custom implementation can read bundle configuration from another source and must return the complete set of configured bundles. - version_added: 3.4.0 + version_added: 3.5.0 type: string example: ~ default: "airflow.dag_processing.bundles.source.ConfigDagBundleSource" From 6f68636075cc8e4cf15bc0f9c3412aaf72e9521b Mon Sep 17 00:00:00 2001 From: Sameer Raj Date: Wed, 14 Oct 2026 22:23:58 +0530 Subject: [PATCH 18/34] Keep custom Dag bundle sources compatible with processor credentials Processor credential provisioning and Job registration need the names and teams declared by the deployment source without constructing bundles. --- .../src/airflow/api/common/airflow_health.py | 6 +-- .../cli/commands/dag_processor_command.py | 4 +- .../airflow/dag_processing/bundles/manager.py | 10 +++++ .../execution_api/versions/head/test_jobs.py | 23 ++++++++++ .../commands/test_dag_processor_command.py | 7 +-- .../test_dag_processor_token_command.py | 24 +++++++++- .../bundles/test_dag_bundle_manager.py | 44 +++++++++++++++++-- 7 files changed, 106 insertions(+), 12 deletions(-) diff --git a/airflow-core/src/airflow/api/common/airflow_health.py b/airflow-core/src/airflow/api/common/airflow_health.py index a3571b3ce3fbf..c0f68df1b076e 100644 --- a/airflow-core/src/airflow/api/common/airflow_health.py +++ b/airflow-core/src/airflow/api/common/airflow_health.py @@ -23,7 +23,7 @@ from sqlalchemy import select from airflow.configuration import conf -from airflow.dag_processing.bundles.manager import DagBundlesManager +from airflow.dag_processing.bundles.manager import get_configured_bundle_team_names from airflow.jobs.dag_processor_job_runner import DagProcessorJobRunner from airflow.jobs.job import Job from airflow.jobs.scheduler_job_runner import SchedulerJobRunner @@ -112,7 +112,7 @@ def _dag_processor_instance_health(job: Job) -> dict[str, Any]: # behind forever and a restarted replica adds a second one. The denominator is therefore taken from # the declared work partition instead, which is unaffected by how replicas come and go: # -# * Dag processor -- the bundles returned by the configured Dag bundle provider. +# * Dag processor -- the bundles returned by the configured Dag bundle source. # * Triggerer -- the team scopes those bundles declare, since a triggerer only picks up triggers for # its own team (see ``Trigger.ids_for_triggerer``). # * Scheduler -- schedulers are symmetric, so there is no partition and no partial state to report. @@ -121,7 +121,7 @@ def _dag_processor_instance_health(job: Job) -> dict[str, Any]: def _configured_bundle_teams() -> dict[str, str | None]: """Map every configured Dag bundle to the team owning it, empty when the config is unreadable.""" try: - return DagBundlesManager().get_configured_bundle_team_names() + return get_configured_bundle_team_names() except Exception: # A health probe must not fail on malformed bundle config; callers fall back to liveness only. log.warning("Could not read the Dag bundle configuration", exc_info=True) diff --git a/airflow-core/src/airflow/cli/commands/dag_processor_command.py b/airflow-core/src/airflow/cli/commands/dag_processor_command.py index 4e65cd0008f20..b8195bead2895 100644 --- a/airflow-core/src/airflow/cli/commands/dag_processor_command.py +++ b/airflow-core/src/airflow/cli/commands/dag_processor_command.py @@ -23,7 +23,7 @@ from airflow.cli.commands.daemon_utils import run_command_with_daemon_option from airflow.configuration import conf -from airflow.dag_processing.bundles.manager import DagBundlesManager +from airflow.dag_processing.bundles.manager import get_configured_bundle_team_names from airflow.dag_processing.manager import DagFileProcessorManager from airflow.jobs.dag_processor_job_runner import DagProcessorJobRunner from airflow.jobs.job import Job, run_job @@ -52,7 +52,7 @@ def _get_team_names(bundle_names: list[str] | None) -> list[str]: if not conf.getboolean("core", "multi_team"): return [] - configured = DagBundlesManager().get_configured_bundle_team_names() + configured = get_configured_bundle_team_names() names = bundle_names or list(configured) return sorted({team for name in names if (team := configured.get(name)) is not None}) diff --git a/airflow-core/src/airflow/dag_processing/bundles/manager.py b/airflow-core/src/airflow/dag_processing/bundles/manager.py index bef4e250f59d1..0c937177054fb 100644 --- a/airflow-core/src/airflow/dag_processing/bundles/manager.py +++ b/airflow-core/src/airflow/dag_processing/bundles/manager.py @@ -56,6 +56,16 @@ _REASSIGN_BATCH_SIZE = 1000 +def get_configured_bundle_names() -> frozenset[str]: + """Get configured Dag bundle names without constructing their bundles.""" + return frozenset(DagBundlesManager().get_all_bundle_names()) + + +def get_configured_bundle_team_names() -> dict[str, str | None]: + """Get the teams owning explicitly configured Dag bundles.""" + return DagBundlesManager().get_configured_bundle_team_names() + + def _discover_example_dag_bundle_paths() -> dict[str, str]: """Return the example Dag bundle names and paths available in this installation.""" from airflow import example_dags diff --git a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_jobs.py b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_jobs.py index 5bd477dd05e70..3d5354e5e0808 100644 --- a/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_jobs.py +++ b/airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_jobs.py @@ -18,6 +18,7 @@ import json from datetime import UTC, datetime, timedelta +from unittest import mock from uuid import UUID import jwt @@ -31,6 +32,8 @@ ExecutionToken, ) from airflow.api_fastapi.execution_api.security import require_auth +from airflow.dag_processing.bundles.manager import DagBundlesManager +from airflow.dag_processing.bundles.source import DagBundleMetadata, DagBundleSource from airflow.jobs.job import Job, JobState from airflow.models.team import Team @@ -122,6 +125,26 @@ def _decode(token: str) -> dict: @pytest.mark.usefixtures("as_session") class TestRegisterJob: + @conf_vars({("core", "multi_team"): "True", ("dag_processor", "dag_bundle_config_list"): "[]"}) + @mock.patch.object(DagBundlesManager, "_load_bundle_source", autospec=True) + def test_resolves_job_teams_from_custom_source_without_constructing_bundles( + self, mock_load, client, session + ): + source = mock.create_autospec(DagBundleSource, instance=True) + source.get_configured_bundle_metadata.return_value = [ + DagBundleMetadata(name="bundle_a", team_name="team_a"), + DagBundleMetadata(name="bundle_b"), + ] + mock_load.return_value = source + session.add(Team(name="team_a")) + session.commit() + + response = _register(client) + + assert response.status_code == 201, response.json() + assert session.get(Job, response.json()["job_id"]).team_names == ["team_a"] + source.get_bundle.assert_not_called() + def test_registers_a_running_job_for_every_granted_bundle(self, client, session): response = _register(client, unixname="airflow") diff --git a/airflow-core/tests/unit/cli/commands/test_dag_processor_command.py b/airflow-core/tests/unit/cli/commands/test_dag_processor_command.py index 3101be8e63264..32dee38eb7694 100644 --- a/airflow-core/tests/unit/cli/commands/test_dag_processor_command.py +++ b/airflow-core/tests/unit/cli/commands/test_dag_processor_command.py @@ -24,6 +24,7 @@ from airflow.cli import cli_parser from airflow.cli.commands import dag_processor_command +from airflow.dag_processing.bundles.manager import DagBundlesManager from airflow.dag_processing.bundles.source import DagBundleMetadata, DagBundleSource from tests_common.test_utils.config import conf_vars @@ -45,7 +46,7 @@ def dag_bundles_with_teams(): with ( conf_vars({("core", "multi_team"): "True"}), mock.patch.object( - dag_processor_command.DagBundlesManager, + dag_processor_command, "get_configured_bundle_team_names", autospec=True, return_value=dict(BUNDLE_TEAMS), @@ -99,7 +100,7 @@ def test_get_team_names(self, bundle_names, expected_team_names): assert dag_processor_command._get_team_names(bundle_names) == expected_team_names @conf_vars({("core", "multi_team"): "True"}) - @mock.patch.object(dag_processor_command.DagBundlesManager, "_load_bundle_source", autospec=True) + @mock.patch.object(DagBundlesManager, "_load_bundle_source", autospec=True) def test_get_team_names_uses_custom_source_metadata(self, mock_load): source = mock.create_autospec(DagBundleSource, instance=True) source.get_configured_bundle_metadata.return_value = [ @@ -112,7 +113,7 @@ def test_get_team_names_uses_custom_source_metadata(self, mock_load): @conf_vars({("core", "multi_team"): "False"}) @mock.patch.object( - dag_processor_command.DagBundlesManager, + dag_processor_command, "get_configured_bundle_team_names", autospec=True, return_value=dict(BUNDLE_TEAMS), diff --git a/airflow-core/tests/unit/cli/commands/test_dag_processor_token_command.py b/airflow-core/tests/unit/cli/commands/test_dag_processor_token_command.py index b29e94603379d..f06008f042f6d 100644 --- a/airflow-core/tests/unit/cli/commands/test_dag_processor_token_command.py +++ b/airflow-core/tests/unit/cli/commands/test_dag_processor_token_command.py @@ -27,7 +27,12 @@ from airflow.cli import cli_parser from airflow.cli.commands import dag_processor_token_command -from airflow.dag_processing.bundles.manager import _load_bundle_config_snapshot +from airflow.dag_processing.bundles.manager import DagBundlesManager +from airflow.dag_processing.bundles.source import ( + DagBundleMetadata, + DagBundleSource, + _load_bundle_config_snapshot, +) from tests_common.test_utils.config import conf_vars @@ -91,6 +96,23 @@ def test_grants_the_requested_bundles(self, tmp_path): assert claims["dag_bundles"] == ["bundle_b"] assert claims["exp"] - claims["iat"] == 120 + @pytest.mark.parametrize("requested", [False, True]) + @conf_vars({("dag_processor", "dag_bundle_config_list"): "[]"}) + @mock.patch.object(DagBundlesManager, "_load_bundle_source", autospec=True) + def test_grants_custom_source_bundles_without_constructing_them(self, mock_load, tmp_path, requested): + source = mock.create_autospec(DagBundleSource, instance=True) + source.get_configured_bundle_metadata.return_value = [DagBundleMetadata(name="custom-bundle")] + mock_load.return_value = source + token_file = tmp_path / "token" + args = ["--token-file", str(token_file)] + if requested: + args.extend(["-B", "custom-bundle"]) + + self._run(*args) + + assert _decode(token_file.read_text())["dag_bundles"] == ["custom-bundle"] + source.get_bundle.assert_not_called() + def test_rejects_an_unknown_bundle(self, tmp_path): token_file = tmp_path / "token" 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 1908736e3e2f1..55e7fc3a7d41e 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,13 +33,14 @@ from airflow.dag_processing.bundles.manager import ( DagBundlesManager, _guess_best_bundle_for_fileloc, - _load_bundle_config_snapshot, get_configured_bundle_names, + get_configured_bundle_team_names, ) from airflow.dag_processing.bundles.source import ( ConfigDagBundleSource, DagBundleMetadata, DagBundleSource, + _load_bundle_config_snapshot, ) from airflow.exceptions import AirflowConfigException from airflow.models.dag import DagModel @@ -219,7 +220,7 @@ def test_get_configured_bundle_team_names_excludes_example_bundles(load_examples ("dag_processor", "dag_bundle_config_list"): json.dumps(TEAM_BUNDLE_CONFIG), } ): - assert DagBundlesManager().get_configured_bundle_team_names() == { + assert get_configured_bundle_team_names() == { "team-bundle": "team-a", "unscoped-bundle": None, } @@ -228,6 +229,7 @@ def test_get_configured_bundle_team_names_excludes_example_bundles(load_examples @conf_vars( { ("core", "load_examples"): "False", + ("core", "multi_team"): "True", ("dag_processor", "dag_bundle_config_list"): json.dumps(TEAM_BUNDLE_CONFIG), } ) @@ -241,7 +243,43 @@ def test_get_configured_bundle_names(): @conf_vars({("dag_processor", "dag_bundle_config_list"): "[]"}) def test_get_configured_bundle_team_names_without_config(): - assert DagBundlesManager().get_configured_bundle_team_names() == {} + assert get_configured_bundle_team_names() == {} + + +@pytest.mark.parametrize( + ("load_examples", "metadata", "expected_names"), + [ + ("False", [DagBundleMetadata(name="custom-bundle", team_name="custom-team")], {"custom-bundle"}), + ( + "True", + [DagBundleMetadata(name="custom-bundle", team_name="custom-team")], + {"custom-bundle", "example_dags"}, + ), + ("True", [], set()), + ], +) +@conf_vars({("core", "multi_team"): "True", ("dag_processor", "dag_bundle_config_list"): "[]"}) +@mock.patch( + "airflow.dag_processing.bundles.manager._discover_example_dag_bundle_paths", + autospec=True, + return_value={"example_dags": "/examples"}, +) +@mock.patch.object(DagBundlesManager, "_load_bundle_source", autospec=True) +def test_metadata_helpers_use_custom_source_without_constructing_bundles( + mock_load, _mock_examples, load_examples, metadata, expected_names +): + source = mock.create_autospec(DagBundleSource, instance=True) + source.get_configured_bundle_metadata.return_value = metadata + mock_load.return_value = source + + with conf_vars({("core", "load_examples"): load_examples}): + names = get_configured_bundle_names() + teams = get_configured_bundle_team_names() + + assert isinstance(names, frozenset) + assert names == expected_names + assert teams == ({"custom-bundle": "custom-team"} if metadata else {}) + source.get_bundle.assert_not_called() def test_get_bundle(): From 2b23ab900b78aeef4f0f0d9207556b23aef6bbb1 Mon Sep 17 00:00:00 2001 From: Sameer Raj Date: Wed, 14 Oct 2026 23:11:07 +0530 Subject: [PATCH 19/34] Preserve existing example Dag discovery logging --- airflow-core/src/airflow/dag_processing/bundles/manager.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow-core/src/airflow/dag_processing/bundles/manager.py b/airflow-core/src/airflow/dag_processing/bundles/manager.py index 0c937177054fb..5fb326cc9f0a2 100644 --- a/airflow-core/src/airflow/dag_processing/bundles/manager.py +++ b/airflow-core/src/airflow/dag_processing/bundles/manager.py @@ -87,7 +87,7 @@ def _discover_example_dag_bundle_paths() -> dict[str, str]: module = importlib.import_module(module_name) module_paths = list(getattr(module, "__path__", [])) except Exception: - log.exception("Could not load provider module %s for example Dag discovery", module_name) + log.exception("Could not load provider module %s for example DAG discovery", module_name) continue for module_path in module_paths: From 468a1abbaf51ca788c58d5ab598bf4b177b5a77a Mon Sep 17 00:00:00 2001 From: Sameer Raj Date: Wed, 14 Oct 2026 23:43:04 +0530 Subject: [PATCH 20/34] Preserve custom Dag bundle refresh intervals in standalone Custom bundle sources can supply remote bundles even when the static list remains at its default, so the local discovery shortcut must not change their refresh cadence. --- .../src/airflow/cli/commands/standalone_command.py | 4 ++++ .../tests/unit/cli/commands/test_standalone_command.py | 7 ++++++- 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/airflow-core/src/airflow/cli/commands/standalone_command.py b/airflow-core/src/airflow/cli/commands/standalone_command.py index 15a9f5edd1dd5..cf3035f3a6e0f 100644 --- a/airflow-core/src/airflow/cli/commands/standalone_command.py +++ b/airflow-core/src/airflow/cli/commands/standalone_command.py @@ -215,6 +215,10 @@ def uses_default_dag_discovery(self) -> bool: unset there; an environment variable counts as a choice whatever its value. Other bundles are left alone, since refreshing a remote bundle fetches from it. """ + if conf.get("dag_processor", "dag_bundle_source") != conf.get_default_value( + "dag_processor", "dag_bundle_source" + ): + return False if REFRESH_INTERVAL_ENV_VAR in os.environ: return False if conf.get("dag_processor", "refresh_interval") != conf.get_default_value( diff --git a/airflow-core/tests/unit/cli/commands/test_standalone_command.py b/airflow-core/tests/unit/cli/commands/test_standalone_command.py index 5cfed400c4326..37bfedcb4ec03 100644 --- a/airflow-core/tests/unit/cli/commands/test_standalone_command.py +++ b/airflow-core/tests/unit/cli/commands/test_standalone_command.py @@ -42,7 +42,7 @@ def default_dag_discovery_config(): """Dag discovery options at their defaults, as the generated airflow.cfg writes them.""" return { ("dag_processor", key): conf.get_default_value("dag_processor", key) - for key in ("refresh_interval", "dag_bundle_config_list") + for key in ("refresh_interval", "dag_bundle_config_list", "dag_bundle_source") } @@ -106,6 +106,11 @@ def test_calculate_env_refreshes_dags_folder_on_every_bundle_check_by_default(se @pytest.mark.parametrize( ("config", "environ"), [ + pytest.param( + {("dag_processor", "dag_bundle_source"): "my_deployment.bundles.DeploymentDagBundleSource"}, + {}, + id="custom-bundle-source", + ), pytest.param({("dag_processor", "refresh_interval"): "60"}, {}, id="interval-in-airflow-cfg"), pytest.param({}, {"AIRFLOW__DAG_PROCESSOR__REFRESH_INTERVAL": "60"}, id="interval-in-env"), pytest.param( From 371c21ee2ff6c8854d1cb4bfd0020cf745726a83 Mon Sep 17 00:00:00 2001 From: Sameer Raj Date: Thu, 15 Oct 2026 01:19:29 +0530 Subject: [PATCH 21/34] Fix Dag bundle source documentation reference Configuration documentation omits settings introduced after the current development version, so the bundle source guide must not depend on that generated link. --- airflow-core/docs/administration-and-deployment/dag-bundles.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow-core/docs/administration-and-deployment/dag-bundles.rst b/airflow-core/docs/administration-and-deployment/dag-bundles.rst index 47d7389213eab..e6b55071cf607 100644 --- a/airflow-core/docs/administration-and-deployment/dag-bundles.rst +++ b/airflow-core/docs/administration-and-deployment/dag-bundles.rst @@ -67,7 +67,7 @@ Dag bundles are configured in :ref:`config:dag_processor__dag_bundle_config_list Writing a bundle source ~~~~~~~~~~~~~~~~~~~~~~~ -Set :ref:`config:dag_processor__dag_bundle_source` to the import path of a class that extends +Set ``[dag_processor] dag_bundle_source`` to the import path of a class that extends ``DagBundleSource``. The default, ``ConfigDagBundleSource``, reads ``dag_bundle_config_list``. This is a class extension point; it does not require a new Airflow provider package. From 067bbaa70ef2f50ad933841349cc790f0895c55d Mon Sep 17 00:00:00 2001 From: Sameer Raj Date: Fri, 16 Oct 2026 22:48:26 +0530 Subject: [PATCH 22/34] Remove redundant comment about Dag ownership --- airflow-core/src/airflow/dag_processing/dagbag.py | 1 - 1 file changed, 1 deletion(-) diff --git a/airflow-core/src/airflow/dag_processing/dagbag.py b/airflow-core/src/airflow/dag_processing/dagbag.py index f3a7fd60c5291..3fce1e6c8d62c 100644 --- a/airflow-core/src/airflow/dag_processing/dagbag.py +++ b/airflow-core/src/airflow/dag_processing/dagbag.py @@ -199,7 +199,6 @@ def __init__( self.bundle_path = bundle_path self.bundle_name = bundle_name self.parse_lang_sdk_files = parse_lang_sdk_files - # Explicit ownership, including None, avoids provider I/O in isolated processes. if not is_arg_set(team_name): team_name = None if bundle_name and conf.getboolean("core", "multi_team"): From 1548011fe41d721c8ff74b260342323dc1cfbec4 Mon Sep 17 00:00:00 2001 From: Sameer Raj Date: Tue, 15 Sep 2026 23:07:16 +0530 Subject: [PATCH 23/34] Reload Dag bundle configuration without restarts Dag processors currently read bundle configuration only at startup. Large deployments need additions, updates, and removals to converge across processors without restarting each process. --- .../dag-bundles.rst | 25 +++++ airflow-core/newsfragments/73217.feature.rst | 1 + .../src/airflow/config_templates/config.yml | 10 +- .../airflow/dag_processing/bundles/manager.py | 4 + .../airflow/dag_processing/bundles/source.py | 19 +++- .../src/airflow/dag_processing/manager.py | 87 +++++++++++++- .../bundles/test_dag_bundle_manager.py | 50 ++++++++- .../tests/unit/dag_processing/test_manager.py | 106 ++++++++++++++++++ 8 files changed, 293 insertions(+), 9 deletions(-) create mode 100644 airflow-core/newsfragments/73217.feature.rst diff --git a/airflow-core/docs/administration-and-deployment/dag-bundles.rst b/airflow-core/docs/administration-and-deployment/dag-bundles.rst index e6b55071cf607..936c3272a58a0 100644 --- a/airflow-core/docs/administration-and-deployment/dag-bundles.rst +++ b/airflow-core/docs/administration-and-deployment/dag-bundles.rst @@ -64,6 +64,31 @@ Configuring Dag bundles Dag bundles are configured in :ref:`config:dag_processor__dag_bundle_config_list`. You can add one or more Dag bundles here. +Dynamic bundle configuration +~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Set :ref:`config:dag_processor__dag_bundle_provider` to use a custom source for bundle configuration. +The provider must return the complete list of active bundles. The Dag processor requests the list every +:ref:`config:dag_processor__bundle_refresh_check_interval` seconds and applies additions and removals without +a restart. A bundle that is absent from the list becomes inactive. + +The provider must raise an exception when it cannot read its source. It must return an empty list only when +no bundles are active. On an error, the Dag processor keeps the last valid list and tries again on the next +refresh cycle. + +Airflow also reconciles the active list with the metadata database on later refresh cycles. This makes the +result converge when Dag processors observe a change at different times. For example, an older Dag processor +can temporarily reactivate a removed bundle. A later cycle from a processor with the current list makes it +inactive again. This process does not use a shared configuration revision or a separate database table. + +A bundle name must keep the same construction settings. Use a new name when settings such as its bundle +class, repository, branch, connection, or refresh interval change. The provider must continue to resolve +the previous name while retained Dag runs can still need it. + +The default ``ConfigDagBundleProvider`` reads ``dag_bundle_config_list``. Its values stay static while the Dag +processor runs. When filtered Dag processors use partial static configuration, they do not deactivate bundles +that are absent from their local configuration. + Writing a bundle source ~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/airflow-core/newsfragments/73217.feature.rst b/airflow-core/newsfragments/73217.feature.rst new file mode 100644 index 0000000000000..f8f41681fc40e --- /dev/null +++ b/airflow-core/newsfragments/73217.feature.rst @@ -0,0 +1 @@ +Reload Dag bundle configuration without restarting Dag processors. diff --git a/airflow-core/src/airflow/config_templates/config.yml b/airflow-core/src/airflow/config_templates/config.yml index 63ea8ec2f401e..2a820f4396439 100644 --- a/airflow-core/src/airflow/config_templates/config.yml +++ b/airflow-core/src/airflow/config_templates/config.yml @@ -3285,6 +3285,9 @@ dag_processor: The default implementation reads ``dag_bundle_config_list``. A custom implementation can read bundle configuration from another source and must return the complete set of configured bundles. + Dag processors request this list every ``bundle_refresh_check_interval`` seconds. A missing + bundle is inactive. A bundle name must keep the same construction settings. Use a new name + when settings such as its class, repository, branch, connection, or refresh interval change. version_added: 3.5.0 type: string example: ~ @@ -3462,9 +3465,10 @@ dag_processor: default: "False" bundle_refresh_check_interval: description: | - How often the DAG processor should check if any DAG bundles are ready for a refresh, either by hitting - the bundles refresh_interval or because another DAG processor has seen a newer version of the bundle. - A low value means we check more frequently, and have a smaller window of time where DAG processors are + How often the Dag processor should request the active bundle configuration and check if any + Dag bundles are ready for a refresh, either by hitting the bundles refresh_interval or because + another Dag processor has seen a newer version of the bundle. + A low value means we check more frequently, and have a smaller window of time where Dag processors are out of sync with each other, parsing different versions of the same bundle. version_added: ~ type: integer diff --git a/airflow-core/src/airflow/dag_processing/bundles/manager.py b/airflow-core/src/airflow/dag_processing/bundles/manager.py index 5fb326cc9f0a2..815ee215853b7 100644 --- a/airflow-core/src/airflow/dag_processing/bundles/manager.py +++ b/airflow-core/src/airflow/dag_processing/bundles/manager.py @@ -622,6 +622,10 @@ def is_bundle_configured(cls, name: str) -> bool: """ return name in cls().get_all_bundle_names() + @property + def provides_complete_configuration(self) -> bool: + return self._bundle_provider.provides_complete_configuration + def get_all_dag_bundles(self) -> Iterable[BaseDagBundle]: """ Get all DAG bundles. diff --git a/airflow-core/src/airflow/dag_processing/bundles/source.py b/airflow-core/src/airflow/dag_processing/bundles/source.py index 1edc0f12502ff..7914b94f06841 100644 --- a/airflow-core/src/airflow/dag_processing/bundles/source.py +++ b/airflow-core/src/airflow/dag_processing/bundles/source.py @@ -47,12 +47,23 @@ class DagBundleSource(ABC): The metadata list is the complete set of configured bundles. A bundle name may remain resolvable through ``get_bundle`` after it leaves that list because retained Dag runs can still need an older version. A name must continue to - identify the same bundle implementation; use a new name for a different one. + identify the same bundle construction settings. Use a new name when settings + such as the bundle class, repository, branch, connection, or refresh interval change. """ + @property + def provides_complete_bundle_list(self) -> bool: + """Return whether each provider instance sees the complete active bundle list.""" + return True + @abstractmethod def get_configured_bundle_metadata(self) -> Sequence[DagBundleMetadata]: - """Return metadata for the complete set of configured Dag bundles.""" + """ + Return metadata for the complete set of configured Dag bundles. + + Return an empty sequence only when no bundles are configured. Raise an exception when + the current metadata cannot be read so Airflow keeps the last valid bundle list. + """ @abstractmethod def get_bundle( @@ -113,6 +124,10 @@ def _load_bundle_config_snapshot() -> tuple[_ExternalBundleConfig, ...]: class ConfigDagBundleSource(DagBundleSource): """Provide Dag bundles configured by ``dag_bundle_config_list``.""" + @property + def provides_complete_bundle_list(self) -> bool: + return False + def __init__(self) -> None: self._bundle_config = {config.name: config for config in _load_bundle_config_snapshot()} self._bundle_classes: dict[str, type[BaseDagBundle]] = {} diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index 33e81c33f5f95..af4b063d3bc83 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -254,6 +254,7 @@ class DagFileProcessorManager(LoggingMixin): ) _dag_bundles: list[BaseDagBundle] = attrs.field(factory=list, init=False) + _dag_bundles_manager: DagBundlesManager | None = attrs.field(default=None, init=False) _bundle_versions: dict[str, str | None] = attrs.field(factory=dict, init=False) _bundle_version_data: dict[str, dict | None] = attrs.field(factory=dict, init=False) _multi_team: bool = attrs.field(factory=lambda: conf.getboolean("core", "multi_team"), init=False) @@ -330,8 +331,10 @@ def sync_bundles(self) -> None: """Sync configured DAG bundles to the metadata database.""" # When this processor only parses a subset of bundles, it does not see the full # bundle configuration and must not deactivate bundles owned by other processors. - dag_bundle_manager = DagBundlesManager() - dag_bundle_manager.sync_bundles_to_db(deactivate_missing=not self.bundle_names_to_parse) + dag_bundle_manager = self._get_dag_bundles_manager() + dag_bundle_manager.sync_bundles_to_db( + deactivate_missing=self._can_deactivate_missing_bundles(dag_bundle_manager) + ) # Best-effort legacy repair: a failure here must not crash DFP startup. # Affected Dags self-heal on the next successful parse. try: @@ -341,7 +344,15 @@ def sync_bundles(self) -> None: def get_all_bundles(self) -> list[BaseDagBundle]: """Return configured DAG bundles filtered by ``bundle_names_to_parse`` if provided.""" - return list(DagBundlesManager().get_all_dag_bundles()) + return list(self._get_dag_bundles_manager().get_all_dag_bundles()) + + def _get_dag_bundles_manager(self) -> DagBundlesManager: + if self._dag_bundles_manager is None: + self._dag_bundles_manager = DagBundlesManager() + return self._dag_bundles_manager + + def _can_deactivate_missing_bundles(self, dag_bundle_manager: DagBundlesManager) -> bool: + return not self.bundle_names_to_parse or dag_bundle_manager.provides_complete_configuration def run(self): """ @@ -892,6 +903,7 @@ def _refresh_dag_bundles(self, known_files: dict[str, set[DagFileInfo]]): return self._bundles_last_refreshed = now_seconds + self._reconcile_bundle_configurations(known_files=known_files) any_refreshed = False for bundle in self._dag_bundles: @@ -1003,6 +1015,75 @@ def _refresh_dag_bundles(self, known_files: dict[str, set[DagFileInfo]]): self._resort_file_queue() self._add_new_files_to_queue(known_files=known_files) + def _reconcile_bundle_configurations(self, known_files: dict[str, set[DagFileInfo]]) -> None: + dag_bundle_manager = self._get_dag_bundles_manager() + try: + dag_bundle_manager.refresh_bundle_configurations() + except Exception: + self.log.exception("Error refreshing Dag bundle configuration") + + try: + dag_bundle_manager.sync_bundles_to_db( + deactivate_missing=self._can_deactivate_missing_bundles(dag_bundle_manager) + ) + except Exception: + self.log.exception("Error reconciling Dag bundle configuration") + else: + self._bundle_name_to_team_name.clear() + + bundle_configurations = dag_bundle_manager.get_all_bundle_configurations() + if self.bundle_names_to_parse: + bundle_configurations = tuple( + configuration + for configuration in bundle_configurations + if configuration.name in self.bundle_names_to_parse + ) + + loaded_bundles_by_name = {bundle.name: bundle for bundle in self._dag_bundles} + loaded_bundle_names = set(loaded_bundles_by_name) + configured_bundle_names = {configuration.name for configuration in bundle_configurations} + # Bundle construction settings are immutable for a name; changes require a new name. + if configured_bundle_names == loaded_bundle_names: + return + + removed_bundle_names = loaded_bundle_names - configured_bundle_names + for bundle_name in removed_bundle_names: + del loaded_bundles_by_name[bundle_name] + + added_bundle_names: set[str] = set() + for configuration in bundle_configurations: + if configuration.name in loaded_bundles_by_name: + continue + try: + bundle = dag_bundle_manager.get_bundle(configuration.name) + except Exception as e: + self.log.exception("Error creating bundle '%s': %s", configuration.name, e) + continue + loaded_bundles_by_name[configuration.name] = bundle + added_bundle_names.add(configuration.name) + + self._dag_bundles = [ + loaded_bundles_by_name[configuration.name] + for configuration in bundle_configurations + if configuration.name in loaded_bundles_by_name + ] + self._force_refresh_bundles.update(added_bundle_names) + + for bundle_name in removed_bundle_names: + self._bundle_versions.pop(bundle_name, None) + self._bundle_version_data.pop(bundle_name, None) + known_files.pop(bundle_name, None) + self._force_refresh_bundles.discard(bundle_name) + + if removed_bundle_names: + self.handle_removed_files(known_files=known_files) + + self.log.info( + "Refreshed Dag bundle configuration: added=%s, removed=%s", + sorted(added_bundle_names), + sorted(removed_bundle_names), + ) + def _find_files_in_bundle(self, bundle: BaseDagBundle) -> set[DagFileInfo]: """ List the files to parse in a bundle through its importers. 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 55e7fc3a7d41e..405e1c93e09d5 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 @@ -173,6 +173,17 @@ def __init__(self): raise RuntimeError("Source initialization failed") +class ReloadingDagBundleProvider(CustomDagBundleProvider): + def __init__(self): + self.metadata = [DagBundleMetadata(name="active-bundle")] + self.error: Exception | None = None + + def get_configured_bundle_metadata(self): + if self.error: + raise self.error + return self.metadata + + BASIC_BUNDLE_CONFIG = [ { "name": "my-test-bundle", @@ -288,6 +299,7 @@ def test_get_bundle(): os.environ, {"AIRFLOW__DAG_PROCESSOR__DAG_BUNDLE_CONFIG_LIST": json.dumps(BASIC_BUNDLE_CONFIG)} ): bundle_manager = DagBundlesManager() + assert bundle_manager.provides_complete_bundle_list is False with pytest.raises(ValueError, match="'bundle-that-doesn't-exist' is not configured"): bundle_manager.get_bundle(name="bundle-that-doesn't-exist", version="hello") @@ -389,6 +401,7 @@ def test_custom_bundle_source_resolves_active_and_retired_bundles(): assert isinstance(source, CustomDagBundleSource) assert source.metadata_requests == 0 + assert manager.provides_complete_bundle_list is True assert manager.get_active_bundle_metadata() == (DagBundleMetadata(name="active-bundle"),) assert source.metadata_requests == 1 assert manager.get_all_bundle_names() == ["active-bundle"] @@ -406,6 +419,37 @@ def test_custom_bundle_source_resolves_active_and_retired_bundles(): assert retired_bundle.version_data == {"manifest": "retired.json"} +@conf_vars( + { + ("core", "LOAD_EXAMPLES"): "False", + ( + "dag_processor", + "dag_bundle_provider", + ): "unit.dag_processing.bundles.test_dag_bundle_manager.ReloadingDagBundleProvider", + } +) +def test_custom_bundle_provider_returns_current_metadata(): + manager = DagBundlesManager() + provider = manager._bundle_provider + assert isinstance(provider, ReloadingDagBundleProvider) + + assert manager.get_active_bundle_metadata() == tuple(provider.metadata) + + provider.metadata = [ + DagBundleMetadata(name="active-bundle"), + DagBundleMetadata(name="added-bundle"), + ] + assert manager.get_active_bundle_metadata() == tuple(provider.metadata) + + provider.error = RuntimeError("source unavailable") + with pytest.raises(RuntimeError, match="source unavailable"): + manager.get_active_bundle_metadata() + provider.error = None + + provider.metadata = [] + assert manager.get_active_bundle_metadata() == () + + @pytest.mark.parametrize( ("source_class", "expected_message"), [ @@ -457,7 +501,7 @@ def clear_db(): @pytest.mark.db_test @conf_vars({("core", "LOAD_EXAMPLES"): "False"}) -def test_sync_bundles_to_db(clear_db, session): +def test_sync_bundles_to_db(clear_db, session, caplog): def _get_bundle_names_and_active(): return session.execute( select(DagBundleModel.name, DagBundleModel.active).order_by(DagBundleModel.name) @@ -488,6 +532,10 @@ def _get_bundle_names_and_active(): # Since my-test-bundle is inactive, the associated import errors should be deleted assert session.scalar(select(func.count(ParseImportError.id))) == 0 + caplog.clear() + DagBundlesManager().sync_bundles_to_db() + assert "DAG bundle my-test-bundle is no longer found in config and has been disabled" not in caplog + # Re-enable one that reappears in config with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(BASIC_BUNDLE_CONFIG)}): manager = DagBundlesManager() diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index 9cc5627d4051e..3ea77d4964fd3 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -50,6 +50,7 @@ from airflow.callbacks.callback_requests import DagCallbackRequest, TaskCallbackRequest from airflow.dag_processing.bundles.base import BaseDagBundle, BundleVersion from airflow.dag_processing.bundles.manager import DagBundlesManager +from airflow.dag_processing.bundles.provider import DagBundleConfiguration from airflow.dag_processing.collection import update_dag_parsing_results_in_db from airflow.dag_processing.dagbag import DagBag from airflow.dag_processing.lang_sdk_processor import LangSDKDagFileProcessorProcess @@ -377,6 +378,110 @@ def mock_processor(self, start_time: float | None = None) -> tuple[DagFileProces ret._open_sockets.clear() return ret, read_end + def test_reconcile_unchanged_bundle_configurations_clears_team_cache_after_sync(self): + manager = DagFileProcessorManager(max_runs=1) + manager._bundle_name_to_team_name = {"existing": "old-team"} + bundle_manager = MagicMock(spec=DagBundlesManager) + manager._dag_bundles_manager = bundle_manager + + manager._reconcile_bundle_configurations(known_files={}) + + assert manager._bundle_name_to_team_name == {} + bundle_manager.sync_bundles_to_db.assert_called_once_with(deactivate_missing=True) + + def test_reconcile_complete_provider_when_filtered(self): + manager = DagFileProcessorManager(max_runs=1, bundle_names_to_parse=["owned"]) + owned_bundle = MagicMock(spec=BaseDagBundle) + owned_bundle.name = "owned" + bundle_manager = MagicMock(spec=DagBundlesManager) + bundle_manager.provides_complete_configuration = True + bundle_manager.get_all_bundle_configurations.return_value = ( + DagBundleConfiguration(name="owned"), + DagBundleConfiguration(name="other"), + ) + bundle_manager.get_bundle.return_value = owned_bundle + manager._dag_bundles_manager = bundle_manager + + manager._reconcile_bundle_configurations(known_files={}) + + assert manager._dag_bundles == [owned_bundle] + bundle_manager.get_bundle.assert_called_once_with("owned") + bundle_manager.sync_bundles_to_db.assert_called_once_with(deactivate_missing=True) + + def test_reconcile_last_valid_configuration_after_provider_error(self): + manager = DagFileProcessorManager(max_runs=1) + existing_bundle = MagicMock(spec=BaseDagBundle) + existing_bundle.name = "existing" + existing_configuration = DagBundleConfiguration(name="existing") + manager._dag_bundles = [existing_bundle] + + bundle_manager = MagicMock(spec=DagBundlesManager) + bundle_manager.refresh_bundle_configurations.side_effect = RuntimeError("source unavailable") + bundle_manager.get_all_bundle_configurations.return_value = (existing_configuration,) + manager._dag_bundles_manager = bundle_manager + + manager._reconcile_bundle_configurations(known_files={}) + + assert manager._dag_bundles == [existing_bundle] + bundle_manager.sync_bundles_to_db.assert_called_once_with(deactivate_missing=True) + + def test_reconcile_adds_and_removes_bundles_without_reloading_existing(self): + manager = DagFileProcessorManager(max_runs=1) + retained_bundle = MagicMock(spec=BaseDagBundle) + retained_bundle.name = "retained" + team_updated_bundle = MagicMock(spec=BaseDagBundle) + team_updated_bundle.name = "updated" + removed_bundle = MagicMock(spec=BaseDagBundle) + removed_bundle.name = "removed" + added_bundle = MagicMock(spec=BaseDagBundle) + added_bundle.name = "added" + + retained_configuration = DagBundleConfiguration(name="retained") + manager._dag_bundles = [retained_bundle, team_updated_bundle, removed_bundle] + manager._bundle_versions = {"updated": "old", "removed": "old"} + manager._bundle_version_data = {"updated": {"old": True}, "removed": {"old": True}} + manager._force_refresh_bundles = {"removed"} + + bundle_manager = MagicMock(spec=DagBundlesManager) + bundle_manager.get_all_bundle_configurations.return_value = ( + retained_configuration, + DagBundleConfiguration(name="updated", team_name="new-team"), + DagBundleConfiguration(name="added"), + ) + bundle_manager.get_bundle.return_value = added_bundle + manager._dag_bundles_manager = bundle_manager + + removed_file = DagFileInfo( + bundle_name="removed", bundle_path=Path("/removed"), rel_path=Path("dag.py") + ) + known_files = {"retained": set(), "removed": {removed_file}} + + with mock.patch.object(manager, "handle_removed_files", autospec=True) as handle_removed_files: + manager._reconcile_bundle_configurations(known_files=known_files) + + assert manager._dag_bundles == [retained_bundle, team_updated_bundle, added_bundle] + assert manager._force_refresh_bundles == {"added"} + assert manager._bundle_versions == {"updated": "old"} + assert manager._bundle_version_data == {"updated": {"old": True}} + assert known_files == {"retained": set()} + handle_removed_files.assert_called_once_with(known_files=known_files) + bundle_manager.get_bundle.assert_called_once_with("added") + bundle_manager.sync_bundles_to_db.assert_called_once_with(deactivate_missing=True) + + def test_reconcile_retries_failed_bundle_addition(self): + manager = DagFileProcessorManager(max_runs=1) + added_configuration = DagBundleConfiguration(name="added") + + bundle_manager = MagicMock(spec=DagBundlesManager) + bundle_manager.get_all_bundle_configurations.return_value = (added_configuration,) + bundle_manager.get_bundle.side_effect = RuntimeError("cannot construct bundle") + manager._dag_bundles_manager = bundle_manager + + manager._reconcile_bundle_configurations(known_files={}) + + assert manager._dag_bundles == [] + bundle_manager.get_bundle.assert_called_once_with("added") + @pytest.fixture def clear_parse_import_errors(self): clear_db_import_errors() @@ -561,6 +666,7 @@ def test_sync_bundles_does_not_deactivate_missing_when_filtered(self): """A processor started with ``--bundle-name`` owns a subset and must not deactivate others.""" manager = DagFileProcessorManager(max_runs=1, bundle_names_to_parse=["only-mine"]) with mock.patch("airflow.dag_processing.manager.DagBundlesManager") as mock_bundles_manager: + mock_bundles_manager.return_value.provides_complete_configuration = False manager.sync_bundles() mock_bundles_manager.return_value.sync_bundles_to_db.assert_called_once_with(deactivate_missing=False) From 58a0936b1cba5d18968a865e075131f495f9046d Mon Sep 17 00:00:00 2001 From: Sameer Raj Date: Wed, 16 Sep 2026 23:39:09 +0530 Subject: [PATCH 24/34] Avoid repeated cleanup for inactive Dag bundles Repeated synchronization otherwise retries cleanup and emits the same warning on every cycle after a bundle has already been disabled. --- airflow-core/src/airflow/dag_processing/bundles/manager.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/airflow-core/src/airflow/dag_processing/bundles/manager.py b/airflow-core/src/airflow/dag_processing/bundles/manager.py index 815ee215853b7..be57b82f5bdbe 100644 --- a/airflow-core/src/airflow/dag_processing/bundles/manager.py +++ b/airflow-core/src/airflow/dag_processing/bundles/manager.py @@ -368,6 +368,9 @@ def _extract_and_sign_template(bundle_name: str) -> tuple[str | None, dict]: from airflow.models.errors import ParseImportError for name, bundle in stored.items(): + # This sync repeats, so only clean up when the bundle becomes inactive. + if bundle.active is False: + continue bundle.active = False bundle.teams = [] self.log.warning("DAG bundle %s is no longer found in config and has been disabled", name) From 2e309e0a84da034b14d2b16efc68e5f9a47e5577 Mon Sep 17 00:00:00 2001 From: Sameer Raj Date: Thu, 17 Sep 2026 00:46:16 +0530 Subject: [PATCH 25/34] Preserve Dag bundle refresh test coverage Existing refresh tests use controlled bundle objects to verify refresh and version behavior. Dynamic configuration reconciliation must not replace those fixtures before their assertions run. --- .../tests/unit/dag_processing/test_manager.py | 26 +++++++++++++++---- 1 file changed, 21 insertions(+), 5 deletions(-) diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index 3ea77d4964fd3..f4a6cf4784e17 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -3607,7 +3607,10 @@ def test_bundles_are_refreshed(self): with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(config)}): DagBundlesManager().sync_bundles_to_db() with mock.patch("airflow.dag_processing.manager.DagBundlesManager") as mock_bundle_manager: - mock_bundle_manager.return_value._bundle_config = {"bundleone": None, "bundletwo": None} + mock_bundle_manager.return_value.get_all_bundle_configurations.return_value = ( + DagBundleConfiguration(name="bundleone"), + DagBundleConfiguration(name="bundletwo"), + ) mock_bundle_manager.return_value.get_all_dag_bundles.return_value = [bundleone, bundletwo] # We should refresh bundleone twice, but bundletwo only once - it has a long refresh_interval @@ -3761,7 +3764,9 @@ def test_bundles_versions_are_stored(self, session): with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(config)}): DagBundlesManager().sync_bundles_to_db() with mock.patch("airflow.dag_processing.manager.DagBundlesManager") as mock_bundle_manager: - mock_bundle_manager.return_value._bundle_config = {"bundleone": None} + mock_bundle_manager.return_value.get_all_bundle_configurations.return_value = ( + DagBundleConfiguration(name="bundleone"), + ) mock_bundle_manager.return_value.get_all_dag_bundles.return_value = [mybundle] manager = DagFileProcessorManager(max_runs=1) manager.run() @@ -3788,7 +3793,9 @@ def test_non_versioned_bundle_get_version_not_called(self): with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(config)}): DagBundlesManager().sync_bundles_to_db() with mock.patch("airflow.dag_processing.manager.DagBundlesManager") as mock_bundle_manager: - mock_bundle_manager.return_value._bundle_config = {"bundleone": None} + mock_bundle_manager.return_value.get_all_bundle_configurations.return_value = ( + DagBundleConfiguration(name="bundleone"), + ) mock_bundle_manager.return_value.get_all_dag_bundles.return_value = [bundleone] manager = DagFileProcessorManager(max_runs=1) manager.run() @@ -3816,7 +3823,9 @@ def test_versioned_bundle_get_version_called_once(self): with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(config)}): DagBundlesManager().sync_bundles_to_db() with mock.patch("airflow.dag_processing.manager.DagBundlesManager") as mock_bundle_manager: - mock_bundle_manager.return_value._bundle_config = {"bundleone": None} + mock_bundle_manager.return_value.get_all_bundle_configurations.return_value = ( + DagBundleConfiguration(name="bundleone"), + ) mock_bundle_manager.return_value.get_all_dag_bundles.return_value = [bundleone] manager = DagFileProcessorManager(max_runs=1) manager.run() # run it once to warm up @@ -4179,6 +4188,7 @@ def _refresh_with_mocked_state(self, manager, bundle, initial_state): mock.patch.object(manager, "handle_removed_files"), mock.patch.object(manager, "_resort_file_queue"), mock.patch.object(manager, "_add_new_files_to_queue"), + mock.patch.object(manager, "_reconcile_bundle_configurations"), ): manager._refresh_dag_bundles({}) return patched_get, patched_update @@ -4258,6 +4268,7 @@ def test_refresh_dag_bundles_versioned_version_unchanged_persist_failure(self): mock.patch.object(manager, "handle_removed_files"), mock.patch.object(manager, "_resort_file_queue"), mock.patch.object(manager, "_add_new_files_to_queue"), + mock.patch.object(manager, "_reconcile_bundle_configurations"), ): manager._refresh_dag_bundles(known_files) @@ -4294,7 +4305,10 @@ def test_refresh_dag_bundles_get_bundle_state_failure_skips_bundle(self): bundle = self._make_refresh_bundle() manager._dag_bundles = [bundle] - with mock.patch.object(manager, "get_bundle_state", side_effect=Exception("API error")): + with ( + mock.patch.object(manager, "get_bundle_state", side_effect=Exception("API error")), + mock.patch.object(manager, "_reconcile_bundle_configurations"), + ): manager._refresh_dag_bundles({}) bundle.refresh.assert_not_called() @@ -4327,6 +4341,7 @@ def test_refresh_dag_bundles_initialize_non_airflow_exception_skips_bundle(self) mock.patch.object(manager, "handle_removed_files"), mock.patch.object(manager, "_resort_file_queue"), mock.patch.object(manager, "_add_new_files_to_queue"), + mock.patch.object(manager, "_reconcile_bundle_configurations"), ): manager._refresh_dag_bundles({}) @@ -4354,6 +4369,7 @@ def test_refresh_dag_bundles_update_bundle_state_failure_still_scans_files(self) mock.patch.object(manager, "handle_removed_files"), mock.patch.object(manager, "_resort_file_queue"), mock.patch.object(manager, "_add_new_files_to_queue"), + mock.patch.object(manager, "_reconcile_bundle_configurations"), ): manager._refresh_dag_bundles(known_files) From 2f118a40f3967707db64706a72b9fcfb2a282142 Mon Sep 17 00:00:00 2001 From: Sameer Raj Date: Thu, 17 Sep 2026 11:45:50 +0530 Subject: [PATCH 26/34] Clarify dynamic Dag bundle configuration docs Provider also describes installable Airflow integrations. Referring to the configured class and implementation makes the extension contract clear to deployment managers. --- .../dag-bundles.rst | 21 ++++++++++--------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/airflow-core/docs/administration-and-deployment/dag-bundles.rst b/airflow-core/docs/administration-and-deployment/dag-bundles.rst index 936c3272a58a0..63ce76ef68330 100644 --- a/airflow-core/docs/administration-and-deployment/dag-bundles.rst +++ b/airflow-core/docs/administration-and-deployment/dag-bundles.rst @@ -67,14 +67,15 @@ Dag bundles are configured in :ref:`config:dag_processor__dag_bundle_config_list Dynamic bundle configuration ~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -Set :ref:`config:dag_processor__dag_bundle_provider` to use a custom source for bundle configuration. -The provider must return the complete list of active bundles. The Dag processor requests the list every +Set :ref:`config:dag_processor__dag_bundle_provider` to the import path for the class Airflow uses to load +Dag bundle configuration. The configured implementation must return the complete list of active bundles. +The Dag processor requests the list every :ref:`config:dag_processor__bundle_refresh_check_interval` seconds and applies additions and removals without a restart. A bundle that is absent from the list becomes inactive. -The provider must raise an exception when it cannot read its source. It must return an empty list only when -no bundles are active. On an error, the Dag processor keeps the last valid list and tries again on the next -refresh cycle. +The configured implementation must raise an exception when it cannot read its source. It must return an empty +list only when no bundles are active. On an error, the Dag processor keeps the last valid list and tries again +on the next refresh cycle. Airflow also reconciles the active list with the metadata database on later refresh cycles. This makes the result converge when Dag processors observe a change at different times. For example, an older Dag processor @@ -82,12 +83,12 @@ can temporarily reactivate a removed bundle. A later cycle from a processor with inactive again. This process does not use a shared configuration revision or a separate database table. A bundle name must keep the same construction settings. Use a new name when settings such as its bundle -class, repository, branch, connection, or refresh interval change. The provider must continue to resolve -the previous name while retained Dag runs can still need it. +class, repository, branch, connection, or refresh interval change. The configured implementation must continue +to resolve the previous name while retained Dag runs can still need it. -The default ``ConfigDagBundleProvider`` reads ``dag_bundle_config_list``. Its values stay static while the Dag -processor runs. When filtered Dag processors use partial static configuration, they do not deactivate bundles -that are absent from their local configuration. +The default implementation, ``ConfigDagBundleProvider``, reads ``dag_bundle_config_list``. Its values stay static +while the Dag processor runs. When filtered Dag processors use partial static configuration, they do not +deactivate bundles that are absent from their local configuration. Writing a bundle source ~~~~~~~~~~~~~~~~~~~~~~~ From d23f5924d251fcce5fe89118142c2f410006071e Mon Sep 17 00:00:00 2001 From: Sameer Raj Date: Mon, 21 Sep 2026 15:46:17 +0530 Subject: [PATCH 27/34] Keep dynamic Dag bundle reconciliation consistent Using one provider snapshot keeps runtime reconciliation and stored bundle state aligned when configurations change while the Dag processor is running. --- .../airflow/dag_processing/bundles/manager.py | 20 +++++++++---- .../src/airflow/dag_processing/manager.py | 9 +++--- .../tests/unit/dag_processing/test_manager.py | 28 ++++++++++++------- 3 files changed, 37 insertions(+), 20 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/bundles/manager.py b/airflow-core/src/airflow/dag_processing/bundles/manager.py index be57b82f5bdbe..caa86660df48c 100644 --- a/airflow-core/src/airflow/dag_processing/bundles/manager.py +++ b/airflow-core/src/airflow/dag_processing/bundles/manager.py @@ -43,7 +43,7 @@ from airflow.utils.session import NEW_SESSION, create_session, provide_session if TYPE_CHECKING: - from collections.abc import Iterable + from collections.abc import Iterable, Sequence from sqlalchemy.engine import CursorResult from sqlalchemy.orm import Session @@ -258,7 +258,13 @@ def get_configured_bundle_team_names(self) -> dict[str, str | None]: return {metadata.name: metadata.team_name for metadata in self._get_source_bundle_metadata()} @provide_session - def sync_bundles_to_db(self, *, deactivate_missing: bool = True, session: Session = NEW_SESSION) -> None: + def sync_bundles_to_db( + self, + *, + bundle_metadata: Sequence[DagBundleMetadata] | None = None, + deactivate_missing: bool = True, + session: Session = NEW_SESSION, + ) -> None: """ Persist the configured DAG bundles into ``DagBundleModel`` rows. @@ -278,7 +284,9 @@ def sync_bundles_to_db(self, *, deactivate_missing: bool = True, session: Sessio """ self.log.debug("Syncing DAG bundles to the database") - bundle_metadata = {metadata.name: metadata for metadata in self.get_active_bundle_metadata()} + if bundle_metadata is None: + bundle_metadata = self.get_active_bundle_metadata() + bundle_metadata_by_name = {metadata.name: metadata for metadata in bundle_metadata} def _extract_and_sign_template(bundle_name: str) -> tuple[str | None, dict]: bundle_instance = self.get_bundle(name) @@ -304,7 +312,7 @@ def _extract_and_sign_template(bundle_name: str) -> tuple[str | None, dict]: for bundle in stored.values() } - for name, metadata in bundle_metadata.items(): + for name, metadata in bundle_metadata_by_name.items(): team: Team | None = None if metadata.team_name: team = session.scalars(select(Team).where(Team.name == metadata.team_name)).one_or_none() @@ -626,8 +634,8 @@ def is_bundle_configured(cls, name: str) -> bool: return name in cls().get_all_bundle_names() @property - def provides_complete_configuration(self) -> bool: - return self._bundle_provider.provides_complete_configuration + def provides_complete_bundle_list(self) -> bool: + return self._bundle_provider.provides_complete_bundle_list def get_all_dag_bundles(self) -> Iterable[BaseDagBundle]: """ diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index af4b063d3bc83..6bc33219e7ffe 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -1018,20 +1018,21 @@ def _refresh_dag_bundles(self, known_files: dict[str, set[DagFileInfo]]): def _reconcile_bundle_configurations(self, known_files: dict[str, set[DagFileInfo]]) -> None: dag_bundle_manager = self._get_dag_bundles_manager() try: - dag_bundle_manager.refresh_bundle_configurations() + bundle_configurations = dag_bundle_manager.get_all_bundle_configurations() except Exception: - self.log.exception("Error refreshing Dag bundle configuration") + self.log.exception("Error reading Dag bundle configuration") + return try: dag_bundle_manager.sync_bundles_to_db( - deactivate_missing=self._can_deactivate_missing_bundles(dag_bundle_manager) + bundle_configurations=bundle_configurations, + deactivate_missing=self._can_deactivate_missing_bundles(dag_bundle_manager), ) except Exception: self.log.exception("Error reconciling Dag bundle configuration") else: self._bundle_name_to_team_name.clear() - bundle_configurations = dag_bundle_manager.get_all_bundle_configurations() if self.bundle_names_to_parse: bundle_configurations = tuple( configuration diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index f4a6cf4784e17..06999d4565796 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -382,12 +382,16 @@ def test_reconcile_unchanged_bundle_configurations_clears_team_cache_after_sync( manager = DagFileProcessorManager(max_runs=1) manager._bundle_name_to_team_name = {"existing": "old-team"} bundle_manager = MagicMock(spec=DagBundlesManager) + bundle_configurations = () + bundle_manager.get_all_bundle_configurations.return_value = bundle_configurations manager._dag_bundles_manager = bundle_manager manager._reconcile_bundle_configurations(known_files={}) assert manager._bundle_name_to_team_name == {} - bundle_manager.sync_bundles_to_db.assert_called_once_with(deactivate_missing=True) + bundle_manager.sync_bundles_to_db.assert_called_once_with( + bundle_configurations=bundle_configurations, deactivate_missing=True + ) def test_reconcile_complete_provider_when_filtered(self): manager = DagFileProcessorManager(max_runs=1, bundle_names_to_parse=["owned"]) @@ -395,10 +399,11 @@ def test_reconcile_complete_provider_when_filtered(self): owned_bundle.name = "owned" bundle_manager = MagicMock(spec=DagBundlesManager) bundle_manager.provides_complete_configuration = True - bundle_manager.get_all_bundle_configurations.return_value = ( + bundle_configurations = ( DagBundleConfiguration(name="owned"), DagBundleConfiguration(name="other"), ) + bundle_manager.get_all_bundle_configurations.return_value = bundle_configurations bundle_manager.get_bundle.return_value = owned_bundle manager._dag_bundles_manager = bundle_manager @@ -406,24 +411,24 @@ def test_reconcile_complete_provider_when_filtered(self): assert manager._dag_bundles == [owned_bundle] bundle_manager.get_bundle.assert_called_once_with("owned") - bundle_manager.sync_bundles_to_db.assert_called_once_with(deactivate_missing=True) + bundle_manager.sync_bundles_to_db.assert_called_once_with( + bundle_configurations=bundle_configurations, deactivate_missing=True + ) - def test_reconcile_last_valid_configuration_after_provider_error(self): + def test_reconcile_keeps_loaded_bundles_after_provider_error(self): manager = DagFileProcessorManager(max_runs=1) existing_bundle = MagicMock(spec=BaseDagBundle) existing_bundle.name = "existing" - existing_configuration = DagBundleConfiguration(name="existing") manager._dag_bundles = [existing_bundle] bundle_manager = MagicMock(spec=DagBundlesManager) - bundle_manager.refresh_bundle_configurations.side_effect = RuntimeError("source unavailable") - bundle_manager.get_all_bundle_configurations.return_value = (existing_configuration,) + bundle_manager.get_all_bundle_configurations.side_effect = RuntimeError("source unavailable") manager._dag_bundles_manager = bundle_manager manager._reconcile_bundle_configurations(known_files={}) assert manager._dag_bundles == [existing_bundle] - bundle_manager.sync_bundles_to_db.assert_called_once_with(deactivate_missing=True) + bundle_manager.sync_bundles_to_db.assert_not_called() def test_reconcile_adds_and_removes_bundles_without_reloading_existing(self): manager = DagFileProcessorManager(max_runs=1) @@ -443,11 +448,12 @@ def test_reconcile_adds_and_removes_bundles_without_reloading_existing(self): manager._force_refresh_bundles = {"removed"} bundle_manager = MagicMock(spec=DagBundlesManager) - bundle_manager.get_all_bundle_configurations.return_value = ( + bundle_configurations = ( retained_configuration, DagBundleConfiguration(name="updated", team_name="new-team"), DagBundleConfiguration(name="added"), ) + bundle_manager.get_all_bundle_configurations.return_value = bundle_configurations bundle_manager.get_bundle.return_value = added_bundle manager._dag_bundles_manager = bundle_manager @@ -466,7 +472,9 @@ def test_reconcile_adds_and_removes_bundles_without_reloading_existing(self): assert known_files == {"retained": set()} handle_removed_files.assert_called_once_with(known_files=known_files) bundle_manager.get_bundle.assert_called_once_with("added") - bundle_manager.sync_bundles_to_db.assert_called_once_with(deactivate_missing=True) + bundle_manager.sync_bundles_to_db.assert_called_once_with( + bundle_configurations=bundle_configurations, deactivate_missing=True + ) def test_reconcile_retries_failed_bundle_addition(self): manager = DagFileProcessorManager(max_runs=1) From 7e0bb54fc65cefda537629b233dd5f5c434e9ae9 Mon Sep 17 00:00:00 2001 From: Sameer Raj Date: Mon, 21 Sep 2026 17:54:08 +0530 Subject: [PATCH 28/34] Clarify dynamic Dag bundle reconciliation --- .../dag-bundles.rst | 4 +- .../src/airflow/dag_processing/manager.py | 48 ++++++----- .../tests/unit/dag_processing/test_manager.py | 82 +++++++++---------- 3 files changed, 66 insertions(+), 68 deletions(-) diff --git a/airflow-core/docs/administration-and-deployment/dag-bundles.rst b/airflow-core/docs/administration-and-deployment/dag-bundles.rst index 63ce76ef68330..7e269cefda335 100644 --- a/airflow-core/docs/administration-and-deployment/dag-bundles.rst +++ b/airflow-core/docs/administration-and-deployment/dag-bundles.rst @@ -67,8 +67,8 @@ Dag bundles are configured in :ref:`config:dag_processor__dag_bundle_config_list Dynamic bundle configuration ~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -Set :ref:`config:dag_processor__dag_bundle_provider` to the import path for the class Airflow uses to load -Dag bundle configuration. The configured implementation must return the complete list of active bundles. +Set :ref:`config:dag_processor__dag_bundle_provider` to the import path for the class Airflow uses to discover +and construct Dag bundles. The configured implementation must return metadata for the complete list of active bundles. The Dag processor requests the list every :ref:`config:dag_processor__bundle_refresh_check_interval` seconds and applies additions and removals without a restart. A bundle that is absent from the list becomes inactive. diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index 6bc33219e7ffe..d817d374b2c10 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -330,7 +330,7 @@ def _exit_gracefully(self, signum, frame): def sync_bundles(self) -> None: """Sync configured DAG bundles to the metadata database.""" # When this processor only parses a subset of bundles, it does not see the full - # bundle configuration and must not deactivate bundles owned by other processors. + # bundle list and must not deactivate bundles owned by other processors. dag_bundle_manager = self._get_dag_bundles_manager() dag_bundle_manager.sync_bundles_to_db( deactivate_missing=self._can_deactivate_missing_bundles(dag_bundle_manager) @@ -352,7 +352,7 @@ def _get_dag_bundles_manager(self) -> DagBundlesManager: return self._dag_bundles_manager def _can_deactivate_missing_bundles(self, dag_bundle_manager: DagBundlesManager) -> bool: - return not self.bundle_names_to_parse or dag_bundle_manager.provides_complete_configuration + return not self.bundle_names_to_parse or dag_bundle_manager.provides_complete_bundle_list def run(self): """ @@ -903,7 +903,7 @@ def _refresh_dag_bundles(self, known_files: dict[str, set[DagFileInfo]]): return self._bundles_last_refreshed = now_seconds - self._reconcile_bundle_configurations(known_files=known_files) + self._reconcile_bundles(known_files=known_files) any_refreshed = False for bundle in self._dag_bundles: @@ -1015,58 +1015,56 @@ def _refresh_dag_bundles(self, known_files: dict[str, set[DagFileInfo]]): self._resort_file_queue() self._add_new_files_to_queue(known_files=known_files) - def _reconcile_bundle_configurations(self, known_files: dict[str, set[DagFileInfo]]) -> None: + def _reconcile_bundles(self, known_files: dict[str, set[DagFileInfo]]) -> None: dag_bundle_manager = self._get_dag_bundles_manager() try: - bundle_configurations = dag_bundle_manager.get_all_bundle_configurations() + bundle_metadata = dag_bundle_manager.get_active_bundle_metadata() except Exception: - self.log.exception("Error reading Dag bundle configuration") + self.log.exception("Error reading active Dag bundle metadata") return try: dag_bundle_manager.sync_bundles_to_db( - bundle_configurations=bundle_configurations, + bundle_metadata=bundle_metadata, deactivate_missing=self._can_deactivate_missing_bundles(dag_bundle_manager), ) except Exception: - self.log.exception("Error reconciling Dag bundle configuration") + self.log.exception("Error reconciling Dag bundles") else: self._bundle_name_to_team_name.clear() if self.bundle_names_to_parse: - bundle_configurations = tuple( - configuration - for configuration in bundle_configurations - if configuration.name in self.bundle_names_to_parse + bundle_metadata = tuple( + metadata for metadata in bundle_metadata if metadata.name in self.bundle_names_to_parse ) loaded_bundles_by_name = {bundle.name: bundle for bundle in self._dag_bundles} loaded_bundle_names = set(loaded_bundles_by_name) - configured_bundle_names = {configuration.name for configuration in bundle_configurations} + active_bundle_names = {metadata.name for metadata in bundle_metadata} # Bundle construction settings are immutable for a name; changes require a new name. - if configured_bundle_names == loaded_bundle_names: + if active_bundle_names == loaded_bundle_names: return - removed_bundle_names = loaded_bundle_names - configured_bundle_names + removed_bundle_names = loaded_bundle_names - active_bundle_names for bundle_name in removed_bundle_names: del loaded_bundles_by_name[bundle_name] added_bundle_names: set[str] = set() - for configuration in bundle_configurations: - if configuration.name in loaded_bundles_by_name: + for metadata in bundle_metadata: + if metadata.name in loaded_bundles_by_name: continue try: - bundle = dag_bundle_manager.get_bundle(configuration.name) + bundle = dag_bundle_manager.get_bundle(metadata.name) except Exception as e: - self.log.exception("Error creating bundle '%s': %s", configuration.name, e) + self.log.exception("Error creating bundle '%s': %s", metadata.name, e) continue - loaded_bundles_by_name[configuration.name] = bundle - added_bundle_names.add(configuration.name) + loaded_bundles_by_name[metadata.name] = bundle + added_bundle_names.add(metadata.name) self._dag_bundles = [ - loaded_bundles_by_name[configuration.name] - for configuration in bundle_configurations - if configuration.name in loaded_bundles_by_name + loaded_bundles_by_name[metadata.name] + for metadata in bundle_metadata + if metadata.name in loaded_bundles_by_name ] self._force_refresh_bundles.update(added_bundle_names) @@ -1080,7 +1078,7 @@ def _reconcile_bundle_configurations(self, known_files: dict[str, set[DagFileInf self.handle_removed_files(known_files=known_files) self.log.info( - "Refreshed Dag bundle configuration: added=%s, removed=%s", + "Reconciled Dag bundles: added=%s, removed=%s", sorted(added_bundle_names), sorted(removed_bundle_names), ) diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index 06999d4565796..19ac628b96b0d 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -50,7 +50,7 @@ from airflow.callbacks.callback_requests import DagCallbackRequest, TaskCallbackRequest from airflow.dag_processing.bundles.base import BaseDagBundle, BundleVersion from airflow.dag_processing.bundles.manager import DagBundlesManager -from airflow.dag_processing.bundles.provider import DagBundleConfiguration +from airflow.dag_processing.bundles.provider import DagBundleMetadata from airflow.dag_processing.collection import update_dag_parsing_results_in_db from airflow.dag_processing.dagbag import DagBag from airflow.dag_processing.lang_sdk_processor import LangSDKDagFileProcessorProcess @@ -378,19 +378,19 @@ def mock_processor(self, start_time: float | None = None) -> tuple[DagFileProces ret._open_sockets.clear() return ret, read_end - def test_reconcile_unchanged_bundle_configurations_clears_team_cache_after_sync(self): + def test_reconcile_unchanged_bundle_metadata_clears_team_cache_after_sync(self): manager = DagFileProcessorManager(max_runs=1) manager._bundle_name_to_team_name = {"existing": "old-team"} bundle_manager = MagicMock(spec=DagBundlesManager) - bundle_configurations = () - bundle_manager.get_all_bundle_configurations.return_value = bundle_configurations + bundle_metadata = () + bundle_manager.get_active_bundle_metadata.return_value = bundle_metadata manager._dag_bundles_manager = bundle_manager - manager._reconcile_bundle_configurations(known_files={}) + manager._reconcile_bundles(known_files={}) assert manager._bundle_name_to_team_name == {} bundle_manager.sync_bundles_to_db.assert_called_once_with( - bundle_configurations=bundle_configurations, deactivate_missing=True + bundle_metadata=bundle_metadata, deactivate_missing=True ) def test_reconcile_complete_provider_when_filtered(self): @@ -398,21 +398,21 @@ def test_reconcile_complete_provider_when_filtered(self): owned_bundle = MagicMock(spec=BaseDagBundle) owned_bundle.name = "owned" bundle_manager = MagicMock(spec=DagBundlesManager) - bundle_manager.provides_complete_configuration = True - bundle_configurations = ( - DagBundleConfiguration(name="owned"), - DagBundleConfiguration(name="other"), + bundle_manager.provides_complete_bundle_list = True + bundle_metadata = ( + DagBundleMetadata(name="owned"), + DagBundleMetadata(name="other"), ) - bundle_manager.get_all_bundle_configurations.return_value = bundle_configurations + bundle_manager.get_active_bundle_metadata.return_value = bundle_metadata bundle_manager.get_bundle.return_value = owned_bundle manager._dag_bundles_manager = bundle_manager - manager._reconcile_bundle_configurations(known_files={}) + manager._reconcile_bundles(known_files={}) assert manager._dag_bundles == [owned_bundle] bundle_manager.get_bundle.assert_called_once_with("owned") bundle_manager.sync_bundles_to_db.assert_called_once_with( - bundle_configurations=bundle_configurations, deactivate_missing=True + bundle_metadata=bundle_metadata, deactivate_missing=True ) def test_reconcile_keeps_loaded_bundles_after_provider_error(self): @@ -422,10 +422,10 @@ def test_reconcile_keeps_loaded_bundles_after_provider_error(self): manager._dag_bundles = [existing_bundle] bundle_manager = MagicMock(spec=DagBundlesManager) - bundle_manager.get_all_bundle_configurations.side_effect = RuntimeError("source unavailable") + bundle_manager.get_active_bundle_metadata.side_effect = RuntimeError("source unavailable") manager._dag_bundles_manager = bundle_manager - manager._reconcile_bundle_configurations(known_files={}) + manager._reconcile_bundles(known_files={}) assert manager._dag_bundles == [existing_bundle] bundle_manager.sync_bundles_to_db.assert_not_called() @@ -441,19 +441,19 @@ def test_reconcile_adds_and_removes_bundles_without_reloading_existing(self): added_bundle = MagicMock(spec=BaseDagBundle) added_bundle.name = "added" - retained_configuration = DagBundleConfiguration(name="retained") + retained_metadata = DagBundleMetadata(name="retained") manager._dag_bundles = [retained_bundle, team_updated_bundle, removed_bundle] manager._bundle_versions = {"updated": "old", "removed": "old"} manager._bundle_version_data = {"updated": {"old": True}, "removed": {"old": True}} manager._force_refresh_bundles = {"removed"} bundle_manager = MagicMock(spec=DagBundlesManager) - bundle_configurations = ( - retained_configuration, - DagBundleConfiguration(name="updated", team_name="new-team"), - DagBundleConfiguration(name="added"), + bundle_metadata = ( + retained_metadata, + DagBundleMetadata(name="updated", team_name="new-team"), + DagBundleMetadata(name="added"), ) - bundle_manager.get_all_bundle_configurations.return_value = bundle_configurations + bundle_manager.get_active_bundle_metadata.return_value = bundle_metadata bundle_manager.get_bundle.return_value = added_bundle manager._dag_bundles_manager = bundle_manager @@ -463,7 +463,7 @@ def test_reconcile_adds_and_removes_bundles_without_reloading_existing(self): known_files = {"retained": set(), "removed": {removed_file}} with mock.patch.object(manager, "handle_removed_files", autospec=True) as handle_removed_files: - manager._reconcile_bundle_configurations(known_files=known_files) + manager._reconcile_bundles(known_files=known_files) assert manager._dag_bundles == [retained_bundle, team_updated_bundle, added_bundle] assert manager._force_refresh_bundles == {"added"} @@ -473,19 +473,19 @@ def test_reconcile_adds_and_removes_bundles_without_reloading_existing(self): handle_removed_files.assert_called_once_with(known_files=known_files) bundle_manager.get_bundle.assert_called_once_with("added") bundle_manager.sync_bundles_to_db.assert_called_once_with( - bundle_configurations=bundle_configurations, deactivate_missing=True + bundle_metadata=bundle_metadata, deactivate_missing=True ) def test_reconcile_retries_failed_bundle_addition(self): manager = DagFileProcessorManager(max_runs=1) - added_configuration = DagBundleConfiguration(name="added") + added_metadata = DagBundleMetadata(name="added") bundle_manager = MagicMock(spec=DagBundlesManager) - bundle_manager.get_all_bundle_configurations.return_value = (added_configuration,) + bundle_manager.get_active_bundle_metadata.return_value = (added_metadata,) bundle_manager.get_bundle.side_effect = RuntimeError("cannot construct bundle") manager._dag_bundles_manager = bundle_manager - manager._reconcile_bundle_configurations(known_files={}) + manager._reconcile_bundles(known_files={}) assert manager._dag_bundles == [] bundle_manager.get_bundle.assert_called_once_with("added") @@ -674,7 +674,7 @@ def test_sync_bundles_does_not_deactivate_missing_when_filtered(self): """A processor started with ``--bundle-name`` owns a subset and must not deactivate others.""" manager = DagFileProcessorManager(max_runs=1, bundle_names_to_parse=["only-mine"]) with mock.patch("airflow.dag_processing.manager.DagBundlesManager") as mock_bundles_manager: - mock_bundles_manager.return_value.provides_complete_configuration = False + mock_bundles_manager.return_value.provides_complete_bundle_list = False manager.sync_bundles() mock_bundles_manager.return_value.sync_bundles_to_db.assert_called_once_with(deactivate_missing=False) @@ -3615,9 +3615,9 @@ def test_bundles_are_refreshed(self): with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(config)}): DagBundlesManager().sync_bundles_to_db() with mock.patch("airflow.dag_processing.manager.DagBundlesManager") as mock_bundle_manager: - mock_bundle_manager.return_value.get_all_bundle_configurations.return_value = ( - DagBundleConfiguration(name="bundleone"), - DagBundleConfiguration(name="bundletwo"), + mock_bundle_manager.return_value.get_active_bundle_metadata.return_value = ( + DagBundleMetadata(name="bundleone"), + DagBundleMetadata(name="bundletwo"), ) mock_bundle_manager.return_value.get_all_dag_bundles.return_value = [bundleone, bundletwo] @@ -3772,8 +3772,8 @@ def test_bundles_versions_are_stored(self, session): with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(config)}): DagBundlesManager().sync_bundles_to_db() with mock.patch("airflow.dag_processing.manager.DagBundlesManager") as mock_bundle_manager: - mock_bundle_manager.return_value.get_all_bundle_configurations.return_value = ( - DagBundleConfiguration(name="bundleone"), + mock_bundle_manager.return_value.get_active_bundle_metadata.return_value = ( + DagBundleMetadata(name="bundleone"), ) mock_bundle_manager.return_value.get_all_dag_bundles.return_value = [mybundle] manager = DagFileProcessorManager(max_runs=1) @@ -3801,8 +3801,8 @@ def test_non_versioned_bundle_get_version_not_called(self): with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(config)}): DagBundlesManager().sync_bundles_to_db() with mock.patch("airflow.dag_processing.manager.DagBundlesManager") as mock_bundle_manager: - mock_bundle_manager.return_value.get_all_bundle_configurations.return_value = ( - DagBundleConfiguration(name="bundleone"), + mock_bundle_manager.return_value.get_active_bundle_metadata.return_value = ( + DagBundleMetadata(name="bundleone"), ) mock_bundle_manager.return_value.get_all_dag_bundles.return_value = [bundleone] manager = DagFileProcessorManager(max_runs=1) @@ -3831,8 +3831,8 @@ def test_versioned_bundle_get_version_called_once(self): with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(config)}): DagBundlesManager().sync_bundles_to_db() with mock.patch("airflow.dag_processing.manager.DagBundlesManager") as mock_bundle_manager: - mock_bundle_manager.return_value.get_all_bundle_configurations.return_value = ( - DagBundleConfiguration(name="bundleone"), + mock_bundle_manager.return_value.get_active_bundle_metadata.return_value = ( + DagBundleMetadata(name="bundleone"), ) mock_bundle_manager.return_value.get_all_dag_bundles.return_value = [bundleone] manager = DagFileProcessorManager(max_runs=1) @@ -4196,7 +4196,7 @@ def _refresh_with_mocked_state(self, manager, bundle, initial_state): mock.patch.object(manager, "handle_removed_files"), mock.patch.object(manager, "_resort_file_queue"), mock.patch.object(manager, "_add_new_files_to_queue"), - mock.patch.object(manager, "_reconcile_bundle_configurations"), + mock.patch.object(manager, "_reconcile_bundles"), ): manager._refresh_dag_bundles({}) return patched_get, patched_update @@ -4276,7 +4276,7 @@ def test_refresh_dag_bundles_versioned_version_unchanged_persist_failure(self): mock.patch.object(manager, "handle_removed_files"), mock.patch.object(manager, "_resort_file_queue"), mock.patch.object(manager, "_add_new_files_to_queue"), - mock.patch.object(manager, "_reconcile_bundle_configurations"), + mock.patch.object(manager, "_reconcile_bundles"), ): manager._refresh_dag_bundles(known_files) @@ -4315,7 +4315,7 @@ def test_refresh_dag_bundles_get_bundle_state_failure_skips_bundle(self): with ( mock.patch.object(manager, "get_bundle_state", side_effect=Exception("API error")), - mock.patch.object(manager, "_reconcile_bundle_configurations"), + mock.patch.object(manager, "_reconcile_bundles"), ): manager._refresh_dag_bundles({}) @@ -4349,7 +4349,7 @@ def test_refresh_dag_bundles_initialize_non_airflow_exception_skips_bundle(self) mock.patch.object(manager, "handle_removed_files"), mock.patch.object(manager, "_resort_file_queue"), mock.patch.object(manager, "_add_new_files_to_queue"), - mock.patch.object(manager, "_reconcile_bundle_configurations"), + mock.patch.object(manager, "_reconcile_bundles"), ): manager._refresh_dag_bundles({}) @@ -4377,7 +4377,7 @@ def test_refresh_dag_bundles_update_bundle_state_failure_still_scans_files(self) mock.patch.object(manager, "handle_removed_files"), mock.patch.object(manager, "_resort_file_queue"), mock.patch.object(manager, "_add_new_files_to_queue"), - mock.patch.object(manager, "_reconcile_bundle_configurations"), + mock.patch.object(manager, "_reconcile_bundles"), ): manager._refresh_dag_bundles(known_files) From 4ee379e49306d6014438db1c73bf711cca185fdf Mon Sep 17 00:00:00 2001 From: Sameer Raj Date: Mon, 21 Sep 2026 18:00:57 +0530 Subject: [PATCH 29/34] Isolate callback recovery from bundle reconciliation --- airflow-core/tests/unit/dag_processing/test_manager.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index 19ac628b96b0d..d224f9d2ce339 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -2885,8 +2885,11 @@ def test_fetch_callbacks_filters_by_bundle_before_limit(self, configure_testing_ "other-bundle-b", } + @mock.patch.object(DagFileProcessorManager, "_reconcile_bundles", autospec=True) @mock.patch.object(LateResolvingBundle, "initialize", autospec=True) - def test_callback_executes_after_bundle_initialization_recovers(self, mock_initialize, tmp_path): + def test_callback_executes_after_bundle_initialization_recovers( + self, mock_initialize, _mock_reconcile_bundles, tmp_path + ): dag_file = tmp_path / "callback_recovery.py" dag_file.write_text( textwrap.dedent( From d8154763394cfae84c23632d0e4b464d0f56f80d Mon Sep 17 00:00:00 2001 From: Sameer Raj Date: Mon, 21 Sep 2026 20:36:49 +0530 Subject: [PATCH 30/34] Align dynamic Dag bundle terminology --- .../docs/administration-and-deployment/dag-bundles.rst | 5 +++-- airflow-core/src/airflow/dag_processing/bundles/source.py | 2 +- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/airflow-core/docs/administration-and-deployment/dag-bundles.rst b/airflow-core/docs/administration-and-deployment/dag-bundles.rst index 7e269cefda335..b902b6949b98f 100644 --- a/airflow-core/docs/administration-and-deployment/dag-bundles.rst +++ b/airflow-core/docs/administration-and-deployment/dag-bundles.rst @@ -68,13 +68,14 @@ Dynamic bundle configuration ~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Set :ref:`config:dag_processor__dag_bundle_provider` to the import path for the class Airflow uses to discover -and construct Dag bundles. The configured implementation must return metadata for the complete list of active bundles. +and construct Dag bundles. The configured implementation must return metadata for the complete list of currently +configured bundles. The Dag processor requests the list every :ref:`config:dag_processor__bundle_refresh_check_interval` seconds and applies additions and removals without a restart. A bundle that is absent from the list becomes inactive. The configured implementation must raise an exception when it cannot read its source. It must return an empty -list only when no bundles are active. On an error, the Dag processor keeps the last valid list and tries again +list only when no bundles are configured. On an error, the Dag processor keeps the last valid list and tries again on the next refresh cycle. Airflow also reconciles the active list with the metadata database on later refresh cycles. This makes the diff --git a/airflow-core/src/airflow/dag_processing/bundles/source.py b/airflow-core/src/airflow/dag_processing/bundles/source.py index 7914b94f06841..3e2f18f543f1d 100644 --- a/airflow-core/src/airflow/dag_processing/bundles/source.py +++ b/airflow-core/src/airflow/dag_processing/bundles/source.py @@ -53,7 +53,7 @@ class DagBundleSource(ABC): @property def provides_complete_bundle_list(self) -> bool: - """Return whether each provider instance sees the complete active bundle list.""" + """Return whether each provider instance sees the complete configured bundle list.""" return True @abstractmethod From 73ea84cbe164e4e9ff394a707043784d21f32d84 Mon Sep 17 00:00:00 2001 From: Sameer Raj Date: Tue, 22 Sep 2026 18:43:41 +0530 Subject: [PATCH 31/34] Keep runtime bundle reconciliation lightweight --- .../airflow/dag_processing/bundles/manager.py | 36 +++++++++----- .../src/airflow/dag_processing/manager.py | 38 +++++++++----- .../bundles/test_dag_bundle_manager.py | 24 +++++++++ .../tests/unit/dag_processing/test_manager.py | 49 +++++++++++++++---- 4 files changed, 111 insertions(+), 36 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/bundles/manager.py b/airflow-core/src/airflow/dag_processing/bundles/manager.py index caa86660df48c..4517a16fb111c 100644 --- a/airflow-core/src/airflow/dag_processing/bundles/manager.py +++ b/airflow-core/src/airflow/dag_processing/bundles/manager.py @@ -43,7 +43,7 @@ from airflow.utils.session import NEW_SESSION, create_session, provide_session if TYPE_CHECKING: - from collections.abc import Iterable, Sequence + from collections.abc import Iterable, Mapping, Sequence from sqlalchemy.engine import CursorResult from sqlalchemy.orm import Session @@ -262,6 +262,7 @@ def sync_bundles_to_db( self, *, bundle_metadata: Sequence[DagBundleMetadata] | None = None, + bundle_instances: Mapping[str, BaseDagBundle] | None = None, deactivate_missing: bool = True, session: Session = NEW_SESSION, ) -> None: @@ -281,6 +282,8 @@ def sync_bundles_to_db( with a partial config (e.g. one bundle per processor), pass ``False`` so a processor does not disable bundles owned by other processors -- otherwise the processors repeatedly deactivate each other and never converge. + :param bundle_instances: Bundles already constructed by the runtime. When provided, + synchronization does not construct any other bundles. """ self.log.debug("Syncing DAG bundles to the database") @@ -288,8 +291,11 @@ def sync_bundles_to_db( bundle_metadata = self.get_active_bundle_metadata() bundle_metadata_by_name = {metadata.name: metadata for metadata in bundle_metadata} - def _extract_and_sign_template(bundle_name: str) -> tuple[str | None, dict]: - bundle_instance = self.get_bundle(name) + def _extract_and_sign_template( + bundle_name: str, bundle_instance: BaseDagBundle | None = None + ) -> tuple[str | None, dict]: + if bundle_instance is None: + bundle_instance = self.get_bundle(bundle_name) new_template_ = bundle_instance.view_url_template() new_params_ = self._extract_template_params(bundle_instance) if new_template_: @@ -322,14 +328,17 @@ def _extract_and_sign_template(bundle_name: str) -> tuple[str | None, dict]: ) bundle = stored.pop(name, None) - update_template = True - try: - new_template, new_params = _extract_and_sign_template(name) - except Exception as e: - self.log.exception("Error creating bundle '%s': %s", name, e) - if bundle is None: - continue - update_template = False + update_template = bundle_instances is None or name in bundle_instances + if update_template: + try: + new_template, new_params = _extract_and_sign_template( + name, bundle_instances.get(name) if bundle_instances is not None else None + ) + except Exception as e: + self.log.exception("Error creating bundle '%s': %s", name, e) + if bundle_instances is None and bundle is None: + continue + update_template = False if bundle is not None: bundle.active = True @@ -341,8 +350,9 @@ def _extract_and_sign_template(bundle_name: str) -> tuple[str | None, dict]: self.log.debug("Updated template parameters for bundle %s", name) else: bundle = DagBundleModel(name=name) - bundle.signed_url_template = new_template - bundle.template_params = new_params + if update_template: + bundle.signed_url_template = new_template + bundle.template_params = new_params session.add(bundle) self.log.info("Added new DAG bundle %s to the database", name) diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index d817d374b2c10..efbd9f9ab9338 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -100,6 +100,7 @@ from airflow.callbacks.callback_requests import CallbackRequest from airflow.dag_processing.bundles.base import BaseDagBundle + from airflow.dag_processing.bundles.provider import DagBundleMetadata from airflow.sdk.api.client import Client @@ -255,6 +256,7 @@ class DagFileProcessorManager(LoggingMixin): _dag_bundles: list[BaseDagBundle] = attrs.field(factory=list, init=False) _dag_bundles_manager: DagBundlesManager | None = attrs.field(default=None, init=False) + _bundle_metadata: dict[str, DagBundleMetadata] = attrs.field(factory=dict, init=False) _bundle_versions: dict[str, str | None] = attrs.field(factory=dict, init=False) _bundle_version_data: dict[str, dict | None] = attrs.field(factory=dict, init=False) _multi_team: bool = attrs.field(factory=lambda: conf.getboolean("core", "multi_team"), init=False) @@ -332,9 +334,12 @@ def sync_bundles(self) -> None: # When this processor only parses a subset of bundles, it does not see the full # bundle list and must not deactivate bundles owned by other processors. dag_bundle_manager = self._get_dag_bundles_manager() + bundle_metadata = dag_bundle_manager.get_active_bundle_metadata() dag_bundle_manager.sync_bundles_to_db( - deactivate_missing=self._can_deactivate_missing_bundles(dag_bundle_manager) + bundle_metadata=bundle_metadata, + deactivate_missing=self._can_deactivate_missing_bundles(dag_bundle_manager), ) + self._bundle_metadata = {metadata.name: metadata for metadata in bundle_metadata} # Best-effort legacy repair: a failure here must not crash DFP startup. # Affected Dags self-heal on the next successful parse. try: @@ -1023,15 +1028,8 @@ def _reconcile_bundles(self, known_files: dict[str, set[DagFileInfo]]) -> None: self.log.exception("Error reading active Dag bundle metadata") return - try: - dag_bundle_manager.sync_bundles_to_db( - bundle_metadata=bundle_metadata, - deactivate_missing=self._can_deactivate_missing_bundles(dag_bundle_manager), - ) - except Exception: - self.log.exception("Error reconciling Dag bundles") - else: - self._bundle_name_to_team_name.clear() + current_bundle_metadata = {metadata.name: metadata for metadata in bundle_metadata} + metadata_changed = current_bundle_metadata != self._bundle_metadata if self.bundle_names_to_parse: bundle_metadata = tuple( @@ -1041,10 +1039,8 @@ def _reconcile_bundles(self, known_files: dict[str, set[DagFileInfo]]) -> None: loaded_bundles_by_name = {bundle.name: bundle for bundle in self._dag_bundles} loaded_bundle_names = set(loaded_bundles_by_name) active_bundle_names = {metadata.name for metadata in bundle_metadata} + runtime_bundles_changed = active_bundle_names != loaded_bundle_names # Bundle construction settings are immutable for a name; changes require a new name. - if active_bundle_names == loaded_bundle_names: - return - removed_bundle_names = loaded_bundle_names - active_bundle_names for bundle_name in removed_bundle_names: del loaded_bundles_by_name[bundle_name] @@ -1061,6 +1057,22 @@ def _reconcile_bundles(self, known_files: dict[str, set[DagFileInfo]]) -> None: loaded_bundles_by_name[metadata.name] = bundle added_bundle_names.add(metadata.name) + try: + dag_bundle_manager.sync_bundles_to_db( + bundle_metadata=tuple(current_bundle_metadata.values()), + bundle_instances={name: loaded_bundles_by_name[name] for name in added_bundle_names}, + deactivate_missing=self._can_deactivate_missing_bundles(dag_bundle_manager), + ) + except Exception: + self.log.exception("Error reconciling Dag bundles") + else: + self._bundle_metadata = current_bundle_metadata + if metadata_changed: + self._bundle_name_to_team_name.clear() + + if not runtime_bundles_changed: + return + self._dag_bundles = [ loaded_bundles_by_name[metadata.name] for metadata in bundle_metadata 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 405e1c93e09d5..00596e760edf4 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 @@ -546,6 +546,30 @@ def _get_bundle_names_and_active(): ] +@pytest.mark.db_test +@pytest.mark.parametrize("row_exists", [False, True]) +@conf_vars({("core", "load_examples"): "False"}) +def test_sync_bundles_to_db_with_runtime_instances_keeps_declared_bundle_active( + clear_db, session, row_exists +): + if row_exists: + bundle = DagBundleModel(name="declared") + bundle.active = False + session.add(bundle) + session.flush() + + manager = DagBundlesManager() + with patch.object(manager, "get_bundle", side_effect=RuntimeError("must not construct")) as get_bundle: + manager.sync_bundles_to_db( + bundle_metadata=(DagBundleMetadata(name="declared"),), + bundle_instances={}, + session=session, + ) + + get_bundle.assert_not_called() + assert session.get(DagBundleModel, "declared").active is True + + @pytest.mark.db_test @conf_vars( { diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index d224f9d2ce339..78f84250d0870 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -378,7 +378,7 @@ def mock_processor(self, start_time: float | None = None) -> tuple[DagFileProces ret._open_sockets.clear() return ret, read_end - def test_reconcile_unchanged_bundle_metadata_clears_team_cache_after_sync(self): + def test_reconcile_unchanged_bundle_metadata_uses_lightweight_sync(self): manager = DagFileProcessorManager(max_runs=1) manager._bundle_name_to_team_name = {"existing": "old-team"} bundle_manager = MagicMock(spec=DagBundlesManager) @@ -388,9 +388,11 @@ def test_reconcile_unchanged_bundle_metadata_clears_team_cache_after_sync(self): manager._reconcile_bundles(known_files={}) - assert manager._bundle_name_to_team_name == {} + assert manager._bundle_name_to_team_name == {"existing": "old-team"} bundle_manager.sync_bundles_to_db.assert_called_once_with( - bundle_metadata=bundle_metadata, deactivate_missing=True + bundle_metadata=bundle_metadata, + bundle_instances={}, + deactivate_missing=True, ) def test_reconcile_complete_provider_when_filtered(self): @@ -412,7 +414,9 @@ def test_reconcile_complete_provider_when_filtered(self): assert manager._dag_bundles == [owned_bundle] bundle_manager.get_bundle.assert_called_once_with("owned") bundle_manager.sync_bundles_to_db.assert_called_once_with( - bundle_metadata=bundle_metadata, deactivate_missing=True + bundle_metadata=bundle_metadata, + bundle_instances={"owned": owned_bundle}, + deactivate_missing=True, ) def test_reconcile_keeps_loaded_bundles_after_provider_error(self): @@ -473,22 +477,41 @@ def test_reconcile_adds_and_removes_bundles_without_reloading_existing(self): handle_removed_files.assert_called_once_with(known_files=known_files) bundle_manager.get_bundle.assert_called_once_with("added") bundle_manager.sync_bundles_to_db.assert_called_once_with( - bundle_metadata=bundle_metadata, deactivate_missing=True + bundle_metadata=bundle_metadata, + bundle_instances={ + "added": added_bundle, + }, + deactivate_missing=True, ) def test_reconcile_retries_failed_bundle_addition(self): manager = DagFileProcessorManager(max_runs=1) added_metadata = DagBundleMetadata(name="added") + added_bundle = MagicMock(spec=BaseDagBundle) + added_bundle.name = "added" bundle_manager = MagicMock(spec=DagBundlesManager) bundle_manager.get_active_bundle_metadata.return_value = (added_metadata,) - bundle_manager.get_bundle.side_effect = RuntimeError("cannot construct bundle") + bundle_manager.get_bundle.side_effect = [RuntimeError("cannot construct bundle"), added_bundle] manager._dag_bundles_manager = bundle_manager + manager._reconcile_bundles(known_files={}) manager._reconcile_bundles(known_files={}) - assert manager._dag_bundles == [] - bundle_manager.get_bundle.assert_called_once_with("added") + assert manager._dag_bundles == [added_bundle] + assert bundle_manager.get_bundle.call_args_list == [mock.call("added"), mock.call("added")] + assert bundle_manager.sync_bundles_to_db.call_args_list == [ + mock.call( + bundle_metadata=(added_metadata,), + bundle_instances={}, + deactivate_missing=True, + ), + mock.call( + bundle_metadata=(added_metadata,), + bundle_instances={"added": added_bundle}, + deactivate_missing=True, + ), + ] @pytest.fixture def clear_parse_import_errors(self): @@ -667,16 +690,22 @@ def test_sync_bundles_deactivates_missing_when_owning_all_bundles(self): """A processor with no bundle filter owns the full config and may deactivate missing bundles.""" manager = DagFileProcessorManager(max_runs=1) with mock.patch("airflow.dag_processing.manager.DagBundlesManager") as mock_bundles_manager: + mock_bundles_manager.return_value.get_active_bundle_metadata.return_value = () manager.sync_bundles() - mock_bundles_manager.return_value.sync_bundles_to_db.assert_called_once_with(deactivate_missing=True) + mock_bundles_manager.return_value.sync_bundles_to_db.assert_called_once_with( + bundle_metadata=(), deactivate_missing=True + ) def test_sync_bundles_does_not_deactivate_missing_when_filtered(self): """A processor started with ``--bundle-name`` owns a subset and must not deactivate others.""" manager = DagFileProcessorManager(max_runs=1, bundle_names_to_parse=["only-mine"]) with mock.patch("airflow.dag_processing.manager.DagBundlesManager") as mock_bundles_manager: mock_bundles_manager.return_value.provides_complete_bundle_list = False + mock_bundles_manager.return_value.get_active_bundle_metadata.return_value = () manager.sync_bundles() - mock_bundles_manager.return_value.sync_bundles_to_db.assert_called_once_with(deactivate_missing=False) + mock_bundles_manager.return_value.sync_bundles_to_db.assert_called_once_with( + bundle_metadata=(), deactivate_missing=False + ) @pytest.mark.parametrize( "safe_mode", From c85afb72a06aa973560b94a421636fcb43501a20 Mon Sep 17 00:00:00 2001 From: Sameer Raj Date: Wed, 23 Sep 2026 12:36:08 +0530 Subject: [PATCH 32/34] Fix runtime bundle synchronization test --- .../tests/unit/dag_processing/bundles/test_dag_bundle_manager.py | 1 + 1 file changed, 1 insertion(+) 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 00596e760edf4..8f837e9e5c08a 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 @@ -567,6 +567,7 @@ def test_sync_bundles_to_db_with_runtime_instances_keeps_declared_bundle_active( ) get_bundle.assert_not_called() + session.flush() assert session.get(DagBundleModel, "declared").active is True From 240ac4ce7ca0518c6f38dafda57491bd9f87bd40 Mon Sep 17 00:00:00 2001 From: Sameer Raj Date: Sat, 10 Oct 2026 23:33:15 +0530 Subject: [PATCH 33/34] Preserve Dag discovery failure coverage with hot reload A file-discovery failure must remain distinct from a provider removing a bundle. The upstream discovery tests need to isolate that behavior when periodic bundle reconciliation is enabled. --- airflow-core/tests/unit/dag_processing/test_manager.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index 78f84250d0870..d220990993a17 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -4420,7 +4420,8 @@ def test_refresh_dag_bundles_update_bundle_state_failure_still_scans_files(self) # iteration will see a version mismatch and re-refresh rather than skip incorrectly assert "mock_bundle" not in manager._bundle_versions - def test_refresh_dag_bundles_discovery_failure_keeps_known_files_and_dags(self): + @mock.patch.object(DagFileProcessorManager, "_reconcile_bundles", autospec=True) + def test_refresh_dag_bundles_discovery_failure_keeps_known_files_and_dags(self, _mock_reconcile): """A failed listing keeps the bundle's known files and does not deactivate its Dags.""" manager = DagFileProcessorManager(max_runs=1) bundle = self._make_refresh_bundle() @@ -4446,7 +4447,8 @@ def test_refresh_dag_bundles_discovery_failure_keeps_known_files_and_dags(self): mock_clear.assert_not_called() assert known_files == {"mock_bundle": known} - def test_refresh_dag_bundles_discovery_failure_lists_again_on_next_refresh(self): + @mock.patch.object(DagFileProcessorManager, "_reconcile_bundles", autospec=True) + def test_refresh_dag_bundles_discovery_failure_lists_again_on_next_refresh(self, _mock_reconcile): """A failed listing must not advance the bundle version, or the next refresh skips the bundle.""" manager = DagFileProcessorManager(max_runs=1, bundle_refresh_check_interval=0) bundle = self._make_refresh_bundle(supports_versioning=True, current_version="v2") From b38fc7a2084181ffb04d53222771e51ba0a6e7c3 Mon Sep 17 00:00:00 2001 From: Sameer Raj Date: Tue, 13 Oct 2026 00:00:21 +0530 Subject: [PATCH 34/34] Keep runtime Dag bundle discovery consistent Runtime discovery must follow the same source contract as startup. Temporary loading failures must preserve configured bundles while processors continue to reconcile changes. --- .../dag-bundles.rst | 4 ++-- .../airflow/dag_processing/bundles/manager.py | 2 +- .../airflow/dag_processing/bundles/source.py | 2 +- .../src/airflow/dag_processing/manager.py | 4 ++-- .../bundles/test_dag_bundle_manager.py | 24 +++++++++---------- .../tests/unit/dag_processing/test_manager.py | 6 ++--- 6 files changed, 21 insertions(+), 21 deletions(-) diff --git a/airflow-core/docs/administration-and-deployment/dag-bundles.rst b/airflow-core/docs/administration-and-deployment/dag-bundles.rst index b902b6949b98f..5b517c31c07bc 100644 --- a/airflow-core/docs/administration-and-deployment/dag-bundles.rst +++ b/airflow-core/docs/administration-and-deployment/dag-bundles.rst @@ -67,7 +67,7 @@ Dag bundles are configured in :ref:`config:dag_processor__dag_bundle_config_list Dynamic bundle configuration ~~~~~~~~~~~~~~~~~~~~~~~~~~~~ -Set :ref:`config:dag_processor__dag_bundle_provider` to the import path for the class Airflow uses to discover +Set :ref:`config:dag_processor__dag_bundle_source` to the import path for the class Airflow uses to discover and construct Dag bundles. The configured implementation must return metadata for the complete list of currently configured bundles. The Dag processor requests the list every @@ -87,7 +87,7 @@ A bundle name must keep the same construction settings. Use a new name when sett class, repository, branch, connection, or refresh interval change. The configured implementation must continue to resolve the previous name while retained Dag runs can still need it. -The default implementation, ``ConfigDagBundleProvider``, reads ``dag_bundle_config_list``. Its values stay static +The default implementation, ``ConfigDagBundleSource``, reads ``dag_bundle_config_list``. Its values stay static while the Dag processor runs. When filtered Dag processors use partial static configuration, they do not deactivate bundles that are absent from their local configuration. diff --git a/airflow-core/src/airflow/dag_processing/bundles/manager.py b/airflow-core/src/airflow/dag_processing/bundles/manager.py index 4517a16fb111c..ed6b00357ce23 100644 --- a/airflow-core/src/airflow/dag_processing/bundles/manager.py +++ b/airflow-core/src/airflow/dag_processing/bundles/manager.py @@ -645,7 +645,7 @@ def is_bundle_configured(cls, name: str) -> bool: @property def provides_complete_bundle_list(self) -> bool: - return self._bundle_provider.provides_complete_bundle_list + return self._bundle_source.provides_complete_bundle_list def get_all_dag_bundles(self) -> Iterable[BaseDagBundle]: """ diff --git a/airflow-core/src/airflow/dag_processing/bundles/source.py b/airflow-core/src/airflow/dag_processing/bundles/source.py index 3e2f18f543f1d..59ef2e26550f1 100644 --- a/airflow-core/src/airflow/dag_processing/bundles/source.py +++ b/airflow-core/src/airflow/dag_processing/bundles/source.py @@ -53,7 +53,7 @@ class DagBundleSource(ABC): @property def provides_complete_bundle_list(self) -> bool: - """Return whether each provider instance sees the complete configured bundle list.""" + """Return whether each source instance sees the complete configured bundle list.""" return True @abstractmethod diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index efbd9f9ab9338..e406814a57716 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -100,7 +100,7 @@ from airflow.callbacks.callback_requests import CallbackRequest from airflow.dag_processing.bundles.base import BaseDagBundle - from airflow.dag_processing.bundles.provider import DagBundleMetadata + from airflow.dag_processing.bundles.source import DagBundleMetadata from airflow.sdk.api.client import Client @@ -573,7 +573,7 @@ def deactivate_stale_dags( if stuck_legacy_rows: # Surface how many legacy rows the startup repair could not route; # each one keeps raising "Requested bundle is not configured." until - # the configured Dag bundle provider returns a matching bundle. + # the configured Dag bundle source returns a matching bundle. self.log.info( "Skipped stale check for %d legacy Dag(s) with NULL relative_fileloc.", stuck_legacy_rows, 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 8f837e9e5c08a..76efd07224b0c 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 @@ -173,7 +173,7 @@ def __init__(self): raise RuntimeError("Source initialization failed") -class ReloadingDagBundleProvider(CustomDagBundleProvider): +class ReloadingDagBundleSource(CustomDagBundleSource): def __init__(self): self.metadata = [DagBundleMetadata(name="active-bundle")] self.error: Exception | None = None @@ -424,29 +424,29 @@ def test_custom_bundle_source_resolves_active_and_retired_bundles(): ("core", "LOAD_EXAMPLES"): "False", ( "dag_processor", - "dag_bundle_provider", - ): "unit.dag_processing.bundles.test_dag_bundle_manager.ReloadingDagBundleProvider", + "dag_bundle_source", + ): "unit.dag_processing.bundles.test_dag_bundle_manager.ReloadingDagBundleSource", } ) -def test_custom_bundle_provider_returns_current_metadata(): +def test_custom_bundle_source_returns_current_metadata(): manager = DagBundlesManager() - provider = manager._bundle_provider - assert isinstance(provider, ReloadingDagBundleProvider) + source = manager._bundle_source + assert isinstance(source, ReloadingDagBundleSource) - assert manager.get_active_bundle_metadata() == tuple(provider.metadata) + assert manager.get_active_bundle_metadata() == tuple(source.metadata) - provider.metadata = [ + source.metadata = [ DagBundleMetadata(name="active-bundle"), DagBundleMetadata(name="added-bundle"), ] - assert manager.get_active_bundle_metadata() == tuple(provider.metadata) + assert manager.get_active_bundle_metadata() == tuple(source.metadata) - provider.error = RuntimeError("source unavailable") + source.error = RuntimeError("source unavailable") with pytest.raises(RuntimeError, match="source unavailable"): manager.get_active_bundle_metadata() - provider.error = None + source.error = None - provider.metadata = [] + source.metadata = [] assert manager.get_active_bundle_metadata() == () diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index d220990993a17..a1a950e5049e7 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -50,7 +50,7 @@ from airflow.callbacks.callback_requests import DagCallbackRequest, TaskCallbackRequest from airflow.dag_processing.bundles.base import BaseDagBundle, BundleVersion from airflow.dag_processing.bundles.manager import DagBundlesManager -from airflow.dag_processing.bundles.provider import DagBundleMetadata +from airflow.dag_processing.bundles.source import DagBundleMetadata from airflow.dag_processing.collection import update_dag_parsing_results_in_db from airflow.dag_processing.dagbag import DagBag from airflow.dag_processing.lang_sdk_processor import LangSDKDagFileProcessorProcess @@ -395,7 +395,7 @@ def test_reconcile_unchanged_bundle_metadata_uses_lightweight_sync(self): deactivate_missing=True, ) - def test_reconcile_complete_provider_when_filtered(self): + def test_reconcile_complete_source_when_filtered(self): manager = DagFileProcessorManager(max_runs=1, bundle_names_to_parse=["owned"]) owned_bundle = MagicMock(spec=BaseDagBundle) owned_bundle.name = "owned" @@ -419,7 +419,7 @@ def test_reconcile_complete_provider_when_filtered(self): deactivate_missing=True, ) - def test_reconcile_keeps_loaded_bundles_after_provider_error(self): + def test_reconcile_keeps_loaded_bundles_after_source_error(self): manager = DagFileProcessorManager(max_runs=1) existing_bundle = MagicMock(spec=BaseDagBundle) existing_bundle.name = "existing"