diff --git a/airflow-core/src/airflow/dag_processing/bundles/manager.py b/airflow-core/src/airflow/dag_processing/bundles/manager.py index 292bb440a21e7..dc54be00420ff 100644 --- a/airflow-core/src/airflow/dag_processing/bundles/manager.py +++ b/airflow-core/src/airflow/dag_processing/bundles/manager.py @@ -16,13 +16,15 @@ # under the License. from __future__ import annotations +import functools import importlib +import json import logging import os import warnings from collections import defaultdict from pathlib import Path -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any, NamedTuple, cast from itsdangerous import URLSafeSerializer from pydantic import BaseModel, ValidationError @@ -86,6 +88,7 @@ class _ExternalBundleConfig(BaseModel): classpath: str kwargs: dict team_name: str | None = None + importers: list[dict[str, Any]] | None = None class _InternalBundleConfig(BaseModel): @@ -200,6 +203,62 @@ def _add_provider_example_dags_to_bundle(bundle_config_list: list[_ExternalBundl ) +class _BundleConfigSnapshot(NamedTuple): + """The configured Dag bundles, as names only and as full configs.""" + + configs: tuple[_ExternalBundleConfig, ...] + names: frozenset[str] + + +_EMPTY_BUNDLE_CONFIG_SNAPSHOT = _BundleConfigSnapshot(configs=(), names=frozenset()) + + +@functools.cache +def _parse_bundle_config_snapshot(config_json: str, load_examples: bool) -> _BundleConfigSnapshot: + """ + Build the snapshot for one configuration, without importing any bundle class. + + Keyed on the configuration rather than cached outright, so every reader of the + same configuration shares one snapshot while a configuration change is still + picked up. + """ + bundle_config_list = _parse_bundle_config(json.loads(config_json)) + if 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 load_bundle_config_snapshot() -> _BundleConfigSnapshot: + """ + Read and validate the configured Dag bundles, without importing their classes. + + Shared by :meth:`DagBundlesManager.parse_config`, which goes on to import each + bundle class, and by callers that only need the names/configs. + """ + config_list = conf.getjson("dag_processor", "dag_bundle_config_list") + if not config_list: + return _EMPTY_BUNDLE_CONFIG_SNAPSHOT + 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_snapshot( + json.dumps(config_list, sort_keys=True), conf.getboolean("core", "LOAD_EXAMPLES") + ) + + +def _clear_bundle_config_snapshot_cache() -> None: + _parse_bundle_config_snapshot.cache_clear() + + +load_bundle_config_snapshot.cache_clear = _clear_bundle_config_snapshot_cache # type: ignore[attr-defined] + + def _is_safe_bundle_url(url: str) -> bool: """ Check if a bundle URL is safe to use. @@ -274,18 +333,9 @@ def parse_config(self) -> None: if self._bundle_config: return - config_list = conf.getjson("dag_processor", "dag_bundle_config_list") - if not config_list: + bundle_config_list = load_bundle_config_snapshot().configs + if not bundle_config_list: return - if not isinstance(config_list, list): - raise AirflowConfigException( - "Section `dag_processor` key `dag_bundle_config_list` " - f"must be list but got {config_list.__class__}" - ) - bundle_config_list = _parse_bundle_config(config_list) - if conf.getboolean("core", "LOAD_EXAMPLES"): - _add_example_dag_bundle(bundle_config_list) - _add_provider_example_dags_to_bundle(bundle_config_list) for bundle_config in bundle_config_list: if bundle_config.team_name and not conf.getboolean("core", "multi_team"): @@ -650,7 +700,10 @@ def get_bundle( 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 + name=name, + version=version, + version_data=version_data, + **cfg_bundle.kwargs, ) def get_all_dag_bundles(self) -> Iterable[BaseDagBundle]: diff --git a/airflow-core/src/airflow/dag_processing/dagbag.py b/airflow-core/src/airflow/dag_processing/dagbag.py index 584840f25c329..a18772e73cf2e 100644 --- a/airflow-core/src/airflow/dag_processing/dagbag.py +++ b/airflow-core/src/airflow/dag_processing/dagbag.py @@ -32,7 +32,6 @@ from airflow import settings from airflow._shared.timezones import timezone from airflow.configuration import conf -from airflow.dag_processing.importers import get_importer_registry from airflow.exceptions import ( AirflowClusterPolicyError, AirflowClusterPolicySkipDag, @@ -54,6 +53,7 @@ from sqlalchemy.orm import Session from airflow import DAG + from airflow.dag_processing.importers import DagImporterRegistry from airflow.models.dagwarning import DagWarning @@ -216,11 +216,18 @@ def __init__( known_pools: set[str] | None = None, bundle_path: Path | None = None, bundle_name: str | None = None, + importer_registry: DagImporterRegistry | None = None, ): super().__init__() self.bundle_path = bundle_path self.bundle_name = bundle_name + if importer_registry is None: + from airflow.dag_processing.importers import get_importer_registry + + importer_registry = get_importer_registry(bundle_name) + self.importer_registry = importer_registry + dag_folder = dag_folder or settings.DAGS_FOLDER self.dag_folder = dag_folder self.dags: dict[str, DAG] = {} @@ -314,8 +321,7 @@ def process_file(self, filepath, only_if_updated=True, safe_mode=True): self.captured_warnings.pop(filepath, None) - registry = get_importer_registry() - importer = registry.get_importer(filepath) + importer = self.importer_registry.get_importer(filepath) if importer is None: self.log.debug("No importer found for file: %s", filepath) @@ -495,8 +501,7 @@ def collect_dags( # Ensure dag_folder is a str -- it may have been a pathlib.Path dag_folder = correct_maybe_zipped(str(dag_folder)) - registry = get_importer_registry() - files_to_parse = registry.list_dag_files(dag_folder, safe_mode=safe_mode) + files_to_parse = self.importer_registry.list_dag_files(dag_folder, safe_mode=safe_mode) for filepath in files_to_parse: try: diff --git a/airflow-core/src/airflow/dag_processing/importers/__init__.py b/airflow-core/src/airflow/dag_processing/importers/__init__.py index 15fe3734bc5a2..19e653c944a49 100644 --- a/airflow-core/src/airflow/dag_processing/importers/__init__.py +++ b/airflow-core/src/airflow/dag_processing/importers/__init__.py @@ -25,6 +25,7 @@ DagImportResult, DagImportWarning, get_importer_registry, + reset_importer_registry, ) from airflow.dag_processing.importers.python_importer import PythonDagImporter @@ -36,4 +37,5 @@ "DagImportWarning", "PythonDagImporter", "get_importer_registry", + "reset_importer_registry", ] diff --git a/airflow-core/src/airflow/dag_processing/importers/base.py b/airflow-core/src/airflow/dag_processing/importers/base.py index 036dfab98d7a4..a2467b66f845a 100644 --- a/airflow-core/src/airflow/dag_processing/importers/base.py +++ b/airflow-core/src/airflow/dag_processing/importers/base.py @@ -18,14 +18,14 @@ from __future__ import annotations +import functools import logging import os -import threading from abc import ABC, abstractmethod -from collections.abc import Iterator +from collections.abc import Iterable, Iterator from dataclasses import dataclass, field from pathlib import Path -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any from airflow._shared.module_loading.file_discovery import find_path_from_directory from airflow.configuration import conf @@ -33,6 +33,7 @@ if TYPE_CHECKING: from airflow.sdk import DAG + from airflow.typing_compat import Self log = logging.getLogger(__name__) @@ -89,12 +90,20 @@ class DagImportResult: @property def success(self) -> bool: """Return True if no fatal errors occurred.""" - return len(self.errors) == 0 + return not self.errors + + +def _normalize_extensions(extensions: Iterable[str]) -> list[str]: + """Normalize file extensions to lowercase with leading dot.""" + return [ext.lower() if ext.startswith(".") else f".{ext.lower()}" for ext in extensions] class AbstractDagImporter(ABC): """Abstract base class for DAG importers.""" + def __init__(self, extensions: list[str] | None = None) -> None: + self._configured_extensions = _normalize_extensions(extensions) if extensions is not None else None + @classmethod @abstractmethod def supported_extensions(cls) -> list[str]: @@ -111,19 +120,32 @@ def import_file( ) -> DagImportResult: """Import DAGs from a file.""" + def get_supported_extensions(self) -> list[str]: + """Return active supported extensions (configured extensions if set, else class defaults).""" + if self._configured_extensions is not None: + return self._configured_extensions + if callable(self.supported_extensions): + return self.supported_extensions() + return self.supported_extensions + + def set_configured_extensions(self, extensions: list[str]) -> None: + """Assign configured extensions to this importer instance.""" + self._configured_extensions = _normalize_extensions(extensions) + def can_handle(self, file_path: str | Path) -> bool: """Check if this importer can handle the given file.""" path = Path(file_path) if isinstance(file_path, str) else file_path - return path.suffix.lower() in self.supported_extensions() + return path.suffix.lower() in self.get_supported_extensions() - def get_relative_path(self, file_path: str | Path, bundle_path: Path | None) -> str: - """Get the relative file path from the bundle root.""" - if bundle_path is None: - return str(file_path) - try: - return str(Path(file_path).relative_to(bundle_path)) - except ValueError: - return str(file_path) + def might_contain_dag(self, file_path: str | Path, safe_mode: bool = True) -> bool: + """ + Check whether a file might contain Airflow DAGs according to safe mode heuristics. + + Custom importers can override this to implement format-specific heuristics. + """ + if not safe_mode: + return True + return might_contain_dag(str(file_path), safe_mode, conf=conf) def list_dag_files( self, @@ -133,91 +155,232 @@ def list_dag_files( """ List DAG files in a directory that this importer can handle. - Override this method to customize file discovery for your importer. - The default implementation finds files matching supported_extensions() - and respects .airflowignore files. - :param directory: Directory to search for DAG files :param safe_mode: Whether to use heuristics to filter non-DAG files :return: Iterator of file paths """ ignore_file_syntax = conf.get_mandatory_value("core", "DAG_IGNORE_FILE_SYNTAX", fallback="glob") - supported_exts = [ext.lower() for ext in self.supported_extensions()] + supported_exts = [ext.lower() for ext in self.get_supported_extensions()] for file_path in find_path_from_directory(directory, ".airflowignore", ignore_file_syntax): path = Path(file_path) if not path.is_file(): continue - - # Check if this importer handles this file extension if path.suffix.lower() not in supported_exts: continue - - # Apply safe_mode heuristic if enabled - if safe_mode and not might_contain_dag(file_path, safe_mode, conf=conf): + if not self.might_contain_dag(file_path, safe_mode): continue yield file_path + def get_relative_path(self, file_path: str | Path, bundle_path: Path | None) -> str: + """Get the relative file path from the bundle root.""" + if bundle_path is None: + return str(file_path) + try: + return str(Path(file_path).relative_to(bundle_path)) + except ValueError: + return str(file_path) + + +@dataclass(frozen=True) +class _ImporterSpec: + """Declarative specification for a DAG importer.""" + + classpath: str + kwargs: dict[str, Any] = field(default_factory=dict) + extensions: list[str] | None = None + context: str = "importer configuration" + + +def _parse_importer_specs(configs: Any, context: str) -> list[_ImporterSpec]: + from airflow.exceptions import AirflowConfigException + + if not isinstance(configs, list): + raise AirflowConfigException( + f"Invalid importer configuration for {context}: expected a list of dictionaries." + ) + specs: list[_ImporterSpec] = [] + for item in configs: + if not isinstance(item, dict): + raise AirflowConfigException( + f"Invalid importer configuration for {context}: each entry must be a dictionary." + ) + classpath = item.get("classpath") + if not classpath: + raise AirflowConfigException( + f"Missing required 'classpath' in importer configuration for {context}." + ) + kwargs = item.get("kwargs", {}) + if not isinstance(kwargs, dict): + raise AirflowConfigException( + f"Field 'kwargs' must be a dictionary in importer configuration for {context}." + ) + extensions = item.get("extensions") + if extensions is not None: + if not isinstance(extensions, list) or any(not isinstance(ext, str) for ext in extensions): + raise AirflowConfigException( + f"Field 'extensions' must be a list of strings in importer configuration for {context}." + ) + extensions = _normalize_extensions(extensions) + specs.append( + _ImporterSpec( + classpath=classpath, + kwargs=kwargs, + extensions=extensions, + context=context, + ) + ) + return specs + class DagImporterRegistry: """ - Registry for DAG importers. Singleton that manages importers by file extension. + Registry for DAG importers that manages importers by file extension. Each file extension can only be handled by one importer at a time. If multiple importers claim the same extension, the last registered one wins and a warning is logged. The built-in PythonDagImporter handles .py and .zip extensions. """ - _instance: DagImporterRegistry | None = None _importers: dict[str, AbstractDagImporter] - _lock = threading.Lock() + _importer_specs: dict[str, _ImporterSpec] - def __new__(cls) -> DagImporterRegistry: - with cls._lock: - if cls._instance is None: - cls._instance = super().__new__(cls) - cls._instance._importers = {} - cls._instance._register_default_importers() - return cls._instance + def __init__(self, register_defaults: bool = True): + self._importers = {} + self._importer_specs = {} + if register_defaults: + self._register_default_importers() def _register_default_importers(self) -> None: from airflow.dag_processing.importers.python_importer import PythonDagImporter - self.register(PythonDagImporter()) + self.register(importer=PythonDagImporter()) - def register(self, importer: AbstractDagImporter) -> None: + @staticmethod + def _instantiate_spec(spec: _ImporterSpec) -> AbstractDagImporter: + from airflow._shared.module_loading import import_string + from airflow.exceptions import AirflowConfigException + + try: + importer_class = import_string(spec.classpath) + importer = importer_class(**spec.kwargs) + except Exception as err: + raise AirflowConfigException( + f"Failed to load DAG importer '{spec.classpath}' for {spec.context}: {err}" + ) from err + + if not isinstance(importer, AbstractDagImporter): + raise AirflowConfigException( + f"Configured DAG importer {type(importer).__module__}." + f"{type(importer).__qualname__} for {spec.context} must inherit " + "from AbstractDagImporter." + ) + + if spec.extensions is not None: + importer.set_configured_extensions(spec.extensions) + return importer + + def _warn_and_evict_extension(self, ext: str, new_name: str) -> None: + if ext in self._importers or ext in self._importer_specs: + existing = self._importers.get(ext) + existing_name = type(existing).__name__ if existing else self._importer_specs[ext].classpath + log.warning( + "Extension '%s' already registered by %s, overriding with %s", + ext, + existing_name, + new_name, + ) + self._importers.pop(ext, None) + self._importer_specs.pop(ext, None) + + def register(self, importer: AbstractDagImporter, extensions: list[str] | None = None) -> None: """ Register an importer for its supported extensions. Each extension can only have one importer. If an extension is already registered, the new importer will override it and a warning will be logged. """ - for ext in importer.supported_extensions(): - ext_lower = ext.lower() - if ext_lower in self._importers: - existing = self._importers[ext_lower] - log.warning( - "Extension '%s' already registered by %s, overriding with %s", - ext, - type(existing).__name__, - type(importer).__name__, - ) + if extensions is None: + extensions = importer.get_supported_extensions() + normalized_extensions = _normalize_extensions(extensions) + importer.set_configured_extensions(normalized_extensions) + + for ext_lower in normalized_extensions: + self._warn_and_evict_extension(ext_lower, type(importer).__name__) self._importers[ext_lower] = importer + def register_specs(self, configs: list[dict[str, Any]], context: str) -> None: + """Register importer specifications from configuration dictionaries.""" + for spec in _parse_importer_specs(configs, context=context): + if spec.extensions is None: + self.register(self._instantiate_spec(spec)) + continue + + for ext_lower in spec.extensions: + self._warn_and_evict_extension(ext_lower, spec.classpath) + self._importer_specs[ext_lower] = spec + + @classmethod + def from_config(cls, bundle_name: str | None = None) -> Self: + """Create and configure a DagImporterRegistry with 3-tier precedence.""" + from airflow.exceptions import AirflowConfigException + + registry = cls(register_defaults=True) + + global_importers = conf.getjson("dag_processor", "dag_importer_configs", fallback=None) + if global_importers: + if not isinstance(global_importers, list): + raise AirflowConfigException( + "Section `dag_processor` key `dag_importer_configs` must be a list " + f"but got {global_importers.__class__.__name__}" + ) + registry.register_specs(global_importers, context="global configuration") + + if bundle_name: + bundle_importers = cls._get_bundle_importers_config(bundle_name) + if bundle_importers: + registry.register_specs(bundle_importers, context=f"bundle '{bundle_name}'") + + return registry + + @staticmethod + def _get_bundle_importers_config(bundle_name: str) -> list[dict[str, Any]] | None: + """Retrieve importer configs for a specific bundle from the bundle config snapshot.""" + from airflow.dag_processing.bundles.manager import load_bundle_config_snapshot + + for cfg in load_bundle_config_snapshot().configs: + if cfg.name == bundle_name: + return cfg.importers + return None + def get_importer(self, file_path: str | Path) -> AbstractDagImporter | None: """Get the appropriate importer for a file, or None if unsupported.""" - path = Path(file_path) if isinstance(file_path, str) else file_path - return self._importers.get(path.suffix.lower()) + ext = (Path(file_path) if isinstance(file_path, str) else file_path).suffix.lower() + + if ext in self._importers: + return self._importers[ext] + + if ext in self._importer_specs: + spec = self._importer_specs[ext] + importer = self._instantiate_spec(spec) + for e, s in list(self._importer_specs.items()): + if s is spec: + self._importers[e] = importer + del self._importer_specs[e] + return importer + + return None def can_handle(self, file_path: str | Path) -> bool: """Check if any registered importer can handle this file.""" - return self.get_importer(file_path) is not None + ext = (Path(file_path) if isinstance(file_path, str) else file_path).suffix.lower() + return ext in self._importers or ext in self._importer_specs def supported_extensions(self) -> list[str]: """Return all registered file extensions.""" - return list(self._importers.keys()) + return sorted(set(self._importers) | set(self._importer_specs)) def list_dag_files( self, @@ -227,40 +390,61 @@ def list_dag_files( """ List all DAG files in a directory using all registered importers. - If directory is actually a file, returns that file if any importer can handle it. + Performs a single filesystem traversal, matching files against registered + importers by extension and delegating safe-mode validation to the importer. :param directory: Directory (or file) to search for DAG files :param safe_mode: Whether to use heuristics to filter non-DAG files :return: List of file paths (deduplicated) """ path = Path(directory) - - # If it's a file, just return it if we can handle it if path.is_file(): - if self.can_handle(path): + importer = self.get_importer(path) + if importer and importer.might_contain_dag(path, safe_mode): return [str(path)] return [] if not path.is_dir(): return [] - seen_files: set[str] = set() + ignore_file_syntax = conf.get_mandatory_value("core", "DAG_IGNORE_FILE_SYNTAX", fallback="glob") file_paths: list[str] = [] - for importer in set(self._importers.values()): - for file_path in importer.list_dag_files(directory, safe_mode): - if file_path not in seen_files: - seen_files.add(file_path) + for file_path in find_path_from_directory(directory, ".airflowignore", ignore_file_syntax): + p = Path(file_path) + try: + if not p.is_file(): + continue + + importer = self.get_importer(p) + if importer is None: + continue + + if importer.might_contain_dag(file_path, safe_mode): file_paths.append(file_path) + except Exception: + log.exception("Error while examining %s", file_path) return file_paths @classmethod def reset(cls) -> None: - """Reset the singleton (for testing).""" - cls._instance = None + """Reset the cached importer registries (for testing).""" + reset_importer_registry() + + +@functools.cache +def get_importer_registry(bundle_name: str | None = None) -> DagImporterRegistry: + """Get the cached DagImporterRegistry instance for global or bundle scope.""" + return DagImporterRegistry.from_config(bundle_name=bundle_name) + +def reset_importer_registry() -> None: + """Reset cached importer registries.""" + get_importer_registry.cache_clear() + try: + from airflow.dag_processing.bundles.manager import load_bundle_config_snapshot -def get_importer_registry() -> DagImporterRegistry: - """Get the global importer registry instance.""" - return DagImporterRegistry() + load_bundle_config_snapshot.cache_clear() # type: ignore[attr-defined] + except (ImportError, AttributeError): + pass diff --git a/airflow-core/src/airflow/dag_processing/importers/python_importer.py b/airflow-core/src/airflow/dag_processing/importers/python_importer.py index 2eedc4dff4dec..fceadc1ca8524 100644 --- a/airflow-core/src/airflow/dag_processing/importers/python_importer.py +++ b/airflow-core/src/airflow/dag_processing/importers/python_importer.py @@ -29,12 +29,10 @@ import traceback import warnings import zipfile -from collections.abc import Iterator from pathlib import Path from typing import TYPE_CHECKING, Any from airflow import settings -from airflow._shared.module_loading.file_discovery import find_path_from_directory from airflow.configuration import conf from airflow.dag_processing.importers.base import ( AbstractDagImporter, @@ -95,28 +93,6 @@ def supported_extensions(cls) -> list[str]: """Return file extensions handled by this importer (.py and .zip).""" return [".py", ".zip"] - def list_dag_files( - self, - directory: str | os.PathLike[str], - safe_mode: bool = True, - ) -> Iterator[str]: - """ - List Python DAG files in a directory. - - Handles both .py files and .zip archives containing Python DAGs. - Respects .airflowignore files in the directory tree. - """ - ignore_file_syntax = conf.get_mandatory_value("core", "DAG_IGNORE_FILE_SYNTAX", fallback="glob") - - for file_path in find_path_from_directory(directory, ".airflowignore", ignore_file_syntax): - path = Path(file_path) - try: - if path.is_file() and (path.suffix.lower() == ".py" or zipfile.is_zipfile(path)): - if might_contain_dag(file_path, safe_mode, conf=conf): - yield file_path - except Exception: - log.exception("Error while examining %s", file_path) - def import_file( self, file_path: str | Path, diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index ee71d86a9a866..7a3a96cc63522 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -55,6 +55,7 @@ ) from airflow.dag_processing.bundles.manager import DagBundlesManager from airflow.dag_processing.collection import update_dag_parsing_results_in_db +from airflow.dag_processing.importers import get_importer_registry from airflow.dag_processing.processor import DagFileParsingResult, DagFileProcessorProcess from airflow.models.asset import remove_references_to_deleted_dags from airflow.models.dag import DagModel @@ -67,7 +68,7 @@ from airflow.sdk import SecretCache from airflow.sdk.log import init_log_file, logging_processors from airflow.typing_compat import assert_never -from airflow.utils.file import list_py_file_paths, might_contain_dag +from airflow.utils.file import might_contain_dag from airflow.utils.helpers import prune_dict from airflow.utils.log.logging_mixin import LoggingMixin from airflow.utils.net import get_hostname @@ -936,10 +937,23 @@ def _refresh_dag_bundles(self, known_files: dict[str, set[DagFileInfo]]): self._bundle_versions[bundle.name] = version_after_refresh self._bundle_version_data[bundle.name] = version_data_after_refresh + self.log.info("Searching for files in %s at %s", bundle.name, bundle.path) + importer_registry = get_importer_registry(bundle.name) + dag_files = importer_registry.list_dag_files(bundle.path, safe_mode=self.dag_discovery_safe_mode) found_files = { - DagFileInfo(rel_path=p, bundle_name=bundle.name, bundle_path=bundle.path) - for p in self._find_files_in_bundle(bundle) + DagFileInfo( + rel_path=Path(p).relative_to(bundle.path), + bundle_name=bundle.name, + bundle_path=bundle.path, + ) + for p in dag_files } + self.log.info( + "Found %s files for bundle %s (dag_discovery_safe_mode=%s)", + len(found_files), + bundle.name, + self.dag_discovery_safe_mode, + ) known_files[bundle.name] = found_files @@ -956,23 +970,6 @@ 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 _find_files_in_bundle(self, bundle: BaseDagBundle) -> list[Path]: - """Get relative paths for dag files from bundle dir.""" - # Build up a list of Python files that could contain DAGs - self.log.info("Searching for files in %s at %s", bundle.name, bundle.path) - rel_paths = [ - Path(x).relative_to(bundle.path) - for x in list_py_file_paths(bundle.path, safe_mode=self.dag_discovery_safe_mode) - ] - self.log.info( - "Found %s files for bundle %s (dag_discovery_safe_mode=%s)", - len(rel_paths), - bundle.name, - self.dag_discovery_safe_mode, - ) - - return rel_paths - def _get_observed_filelocs(self, present: set[DagFileInfo]) -> set[str]: """ Return observed DAG source paths for bundle entries. diff --git a/airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py b/airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py index bbd18ec80d85e..fee44e421a7a0 100644 --- a/airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py +++ b/airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py @@ -29,7 +29,11 @@ from sqlalchemy import func, select, update from airflow.dag_processing.bundles.base import BaseDagBundle -from airflow.dag_processing.bundles.manager import DagBundlesManager, _guess_best_bundle_for_fileloc +from airflow.dag_processing.bundles.manager import ( + DagBundlesManager, + _guess_best_bundle_for_fileloc, + load_bundle_config_snapshot, +) from airflow.exceptions import AirflowConfigException from airflow.models.dag import DagModel from airflow.models.dag_version import DagVersion @@ -1791,3 +1795,37 @@ def test_skipped_row_is_recoverable_after_operator_fix( assert restored.is_stale is False assert restored.bundle_name == "configured-bundle" assert restored.relative_fileloc == "legacy.py" + + +class TestBundleConfigSnapshot: + """Tests for load_bundle_config_snapshot.""" + + def test_load_bundle_config_snapshot_includes_importers(self) -> None: + bundle_config = [ + { + "name": "custom-bundle", + "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", + "kwargs": {"path": "/tmp/test"}, + "importers": [ + { + "classpath": "airflow.dag_processing.importers.python_importer.PythonDagImporter", + "extensions": [".py"], + } + ], + } + ] + with conf_vars({("dag_processor", "dag_bundle_config_list"): json.dumps(bundle_config)}): + snapshot = load_bundle_config_snapshot() + assert "custom-bundle" in snapshot.names + cfg = next(c for c in snapshot.configs if c.name == "custom-bundle") + assert cfg.importers == [ + { + "classpath": "airflow.dag_processing.importers.python_importer.PythonDagImporter", + "extensions": [".py"], + } + ] + + def test_load_bundle_config_snapshot_is_cached(self) -> None: + snapshot1 = load_bundle_config_snapshot() + snapshot2 = load_bundle_config_snapshot() + assert snapshot1 is snapshot2 diff --git a/airflow-core/tests/unit/dag_processing/importers/test_registry.py b/airflow-core/tests/unit/dag_processing/importers/test_registry.py index 955308b4559a7..2b0aff738d657 100644 --- a/airflow-core/tests/unit/dag_processing/importers/test_registry.py +++ b/airflow-core/tests/unit/dag_processing/importers/test_registry.py @@ -18,65 +18,70 @@ from __future__ import annotations +import json +import logging from pathlib import Path +import pytest + from airflow.dag_processing.importers import ( DagImporterRegistry, PythonDagImporter, get_importer_registry, + reset_importer_registry, ) +from airflow.exceptions import AirflowConfigException +from tests_common.test_utils.config import conf_vars -class TestDagImporterRegistry: - """Test the DagImporterRegistry singleton.""" - def setup_method(self): - """Reset the registry before each test.""" - DagImporterRegistry.reset() +class GlobalDagImporter(PythonDagImporter): + pass - def teardown_method(self): - """Reset the registry after each test.""" - DagImporterRegistry.reset() - def test_singleton_pattern(self): - """Registry should return the same instance.""" - registry1 = get_importer_registry() - registry2 = get_importer_registry() - assert registry1 is registry2 +class BundleDagImporter(PythonDagImporter): + pass + + +class LazyTestImporter(PythonDagImporter): + instances = 0 + + def __init__(self, **kwargs): + super().__init__() + self.kwargs = kwargs + LazyTestImporter.instances += 1 + + +class TestDagImporterRegistry: + """Test the DagImporterRegistry.""" def test_default_importers_registered(self): """Registry should have Python importer by default.""" registry = get_importer_registry() extensions = registry.supported_extensions() - assert ".py" in extensions def test_get_importer_for_python(self): """Should return PythonDagImporter for .py files.""" registry = get_importer_registry() importer = registry.get_importer("test.py") - assert importer is not None assert isinstance(importer, PythonDagImporter) def test_get_importer_for_unknown(self): """Should return None for unknown file types.""" registry = get_importer_registry() - importer = registry.get_importer("test.txt") - - assert importer is None + assert registry.get_importer("test.txt") is None def test_can_handle_supported_files(self): """can_handle should return True for supported file types.""" registry = get_importer_registry() - assert registry.can_handle("dag.py") assert registry.can_handle(Path("subdir/dag.py")) def test_can_handle_unsupported_files(self): """can_handle should return False for unsupported file types.""" registry = get_importer_registry() - assert not registry.can_handle("readme.txt") assert not registry.can_handle("config.json") assert not registry.can_handle("script.sh") @@ -84,8 +89,6 @@ def test_can_handle_unsupported_files(self): def test_case_insensitive_extension_matching(self): """Extension matching should be case-insensitive.""" registry = get_importer_registry() - - # All these should be handled assert registry.can_handle("dag.PY") assert registry.can_handle("dag.Py") @@ -94,6 +97,199 @@ def test_reset_clears_singleton(self): registry1 = get_importer_registry() DagImporterRegistry.reset() registry2 = get_importer_registry() - - # Should be different instances after reset assert registry1 is not registry2 + + def test_independent_registry_instances(self): + """Directly instantiating DagImporterRegistry creates isolated instances.""" + assert DagImporterRegistry() is not DagImporterRegistry() + + def test_register_explicit_extensions_and_override_warning(self, caplog): + """register() with explicit extensions overrides existing and logs warning.""" + reg = DagImporterRegistry(register_defaults=True) + importer = PythonDagImporter() + with caplog.at_level(logging.WARNING): + reg.register(importer, extensions=[".py", "custom"]) + assert reg.can_handle("test.custom") + assert importer.can_handle("test.custom") + assert any("Extension '.py' already registered" in r.message for r in caplog.records) + + def test_list_dag_files_with_configured_extensions(self, tmp_path): + """list_dag_files() discovers files matching explicitly configured extensions.""" + dag_file = tmp_path / "valid_dag.custom" + dag_file.write_text("from airflow.sdk import DAG\n") + ignored_file = tmp_path / "ignored.py" + ignored_file.write_text("from airflow.sdk import DAG\n") + notes_file = tmp_path / "notes.txt" + notes_file.write_text("not a dag") + + reg = DagImporterRegistry(register_defaults=False) + reg.register(PythonDagImporter(), extensions=[".custom"]) + + files = reg.list_dag_files(tmp_path, safe_mode=True) + assert files == [str(dag_file)] + + def test_list_dag_files_single_file(self, tmp_path): + """list_dag_files() returns single file if registered importer can handle it.""" + dag_file = tmp_path / "valid_dag.custom" + dag_file.write_text("from airflow.sdk import DAG\n") + + reg = DagImporterRegistry(register_defaults=False) + reg.register(PythonDagImporter(), extensions=[".custom"]) + + assert reg.list_dag_files(dag_file, safe_mode=True) == [str(dag_file)] + + unhandled_file = tmp_path / "unhandled.txt" + unhandled_file.write_text("text") + assert reg.list_dag_files(unhandled_file, safe_mode=True) == [] + + def test_list_dag_files_respects_safe_mode(self, tmp_path): + """list_dag_files() filters non-Dag files when safe_mode=True.""" + no_keywords_file = tmp_path / "no_keywords.custom" + no_keywords_file.write_text("print('hello world')\n") + + reg = DagImporterRegistry(register_defaults=False) + reg.register(PythonDagImporter(), extensions=[".custom"]) + + assert reg.list_dag_files(tmp_path, safe_mode=True) == [] + assert reg.list_dag_files(tmp_path, safe_mode=False) == [str(no_keywords_file)] + + def test_list_dag_files_single_pass_override(self, tmp_path): + """list_dag_files() routes discovered files to overriding importer.""" + py_file = tmp_path / "dag.py" + py_file.write_text("from airflow.sdk import DAG\n") + + class OverridingImporter(PythonDagImporter): + pass + + reg = DagImporterRegistry(register_defaults=True) + overriding_importer = OverridingImporter() + reg.register(overriding_importer, extensions=[".py"]) + + files = reg.list_dag_files(tmp_path, safe_mode=True) + assert files == [str(py_file)] + assert reg.get_importer(py_file) is overriding_importer + + def test_lazy_importer_instantiation(self): + """Importer classes are not imported or instantiated until get_importer is called.""" + LazyTestImporter.instances = 0 + reg = DagImporterRegistry(register_defaults=False) + reg.register_specs( + [ + { + "classpath": f"{__name__}.LazyTestImporter", + "extensions": [".lazy", ".lazy2"], + "kwargs": {"param": "value"}, + } + ], + context="test", + ) + + assert LazyTestImporter.instances == 0 + assert reg.can_handle("file.lazy") + assert reg.can_handle("file.lazy2") + assert LazyTestImporter.instances == 0 + assert set(reg.supported_extensions()) == {".lazy", ".lazy2"} + assert LazyTestImporter.instances == 0 + + importer1 = reg.get_importer("file.lazy") + assert LazyTestImporter.instances == 1 + assert isinstance(importer1, LazyTestImporter) + assert importer1.kwargs == {"param": "value"} + + importer2 = reg.get_importer("file.lazy2") + assert importer2 is importer1 + assert LazyTestImporter.instances == 1 + + def test_from_config_three_tier_precedence(self): + """Bundle config overrides global config, which overrides defaults.""" + global_config = [ + { + "classpath": f"{__name__}.GlobalDagImporter", + "extensions": [".py", ".custom"], + } + ] + bundle_config_list = [ + { + "name": "test_bundle", + "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", + "kwargs": {"path": "/tmp/test"}, + "importers": [ + { + "classpath": f"{__name__}.BundleDagImporter", + "extensions": [".custom"], + } + ], + } + ] + + with conf_vars( + { + ("dag_processor", "dag_importer_configs"): json.dumps(global_config), + ("dag_processor", "dag_bundle_config_list"): json.dumps(bundle_config_list), + } + ): + global_reg = DagImporterRegistry.from_config() + assert isinstance(global_reg.get_importer("dag.py"), GlobalDagImporter) + assert isinstance(global_reg.get_importer("dag.custom"), GlobalDagImporter) + + bundle_reg = DagImporterRegistry.from_config("test_bundle") + assert isinstance(bundle_reg.get_importer("dag.py"), GlobalDagImporter) + assert isinstance(bundle_reg.get_importer("dag.custom"), BundleDagImporter) + + @pytest.mark.parametrize( + ("global_cfg", "bundle_cfg", "match"), + [ + (json.dumps({"invalid": "object"}), None, "key `dag_importer_configs` must be a list"), + (None, [{"extensions": [".py"]}], "Missing required 'classpath'"), + (None, [{"classpath": "invalid.path"}], "Failed to load DAG importer"), + ( + None, + [{"classpath": "builtins.dict"}], + r"Configured DAG importer builtins\.dict for bundle 'test_bundle' must inherit from AbstractDagImporter\.", + ), + ( + json.dumps([{"classpath": "builtins.dict"}]), + None, + r"Configured DAG importer builtins\.dict for global configuration must inherit from AbstractDagImporter\.", + ), + ], + ) + def test_from_config_invalid_configs(self, global_cfg, bundle_cfg, match): + """Invalid configurations raise AirflowConfigException.""" + overrides = {} + if global_cfg: + overrides[("dag_processor", "dag_importer_configs")] = global_cfg + if bundle_cfg: + overrides[("dag_processor", "dag_bundle_config_list")] = json.dumps( + [ + { + "name": "test_bundle", + "classpath": "airflow.dag_processing.bundles.local.LocalDagBundle", + "kwargs": {"path": "/tmp/test"}, + "importers": bundle_cfg, + } + ] + ) + + with conf_vars(overrides), pytest.raises(AirflowConfigException, match=match): + DagImporterRegistry.from_config("test_bundle" if bundle_cfg else None) + + def test_get_importer_registry_caching_and_isolation(self): + """get_importer_registry caches instances per bundle name and clears on reset.""" + reg1 = get_importer_registry() + reg2 = get_importer_registry() + assert reg1 is reg2 + + bundle_a = get_importer_registry("bundle_a") + bundle_a2 = get_importer_registry("bundle_a") + bundle_b = get_importer_registry("bundle_b") + + assert bundle_a is bundle_a2 + assert bundle_a is not bundle_b + assert bundle_a is not reg1 + + reset_importer_registry() + new_reg = get_importer_registry() + new_bundle_a = get_importer_registry("bundle_a") + assert new_reg is not reg1 + assert new_bundle_a is not bundle_a diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index 8dd6a8eca98c3..16b3fb916555f 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.dag_processing.bundles.manager import DagBundlesManager from airflow.dag_processing.collection import update_dag_parsing_results_in_db from airflow.dag_processing.dagbag import DagBag +from airflow.dag_processing.importers import DagImporterRegistry from airflow.dag_processing.manager import ( BundleState, DagFileInfo, @@ -436,20 +437,40 @@ def test_sync_bundles_does_not_deactivate_missing_when_filtered(self): pytest.param(True, id="safe-mode-on-filters-keywordless"), ], ) - def test_find_files_in_bundle_respects_dag_discovery_safe_mode(self, tmp_path, safe_mode): + def test_refresh_dag_bundles_respects_dag_discovery_safe_mode(self, tmp_path, safe_mode): (tmp_path / "with_keywords.py").write_text("from airflow.sdk import DAG\n") (tmp_path / "no_keywords.py").write_text("from mycompany.pipelines import flow\n") bundle = MagicMock(spec=BaseDagBundle) bundle.name = "testing" bundle.path = tmp_path + bundle.refresh_interval = 0 + bundle.supports_versioning = False + bundle.is_initialized = True with conf_vars({("core", "dag_discovery_safe_mode"): str(safe_mode)}): manager = DagFileProcessorManager(max_runs=1) + manager._dag_bundles = [bundle] + manager._force_refresh_bundles = set() + known_files: dict[str, set[DagFileInfo]] = {} + + with ( + mock.patch.object( + manager, "get_bundle_state", return_value=BundleState(last_refreshed=None, version=None) + ), + mock.patch.object(manager, "update_bundle_state"), + mock.patch.object(manager, "deactivate_deleted_dags"), + mock.patch.object(manager, "clear_orphaned_import_errors"), + mock.patch.object(manager, "handle_removed_files"), + mock.patch.object(manager, "_resort_file_queue"), + mock.patch.object(manager, "_add_new_files_to_queue"), + ): + manager._refresh_dag_bundles(known_files) expected = {Path("with_keywords.py")} if not safe_mode: expected.add(Path("no_keywords.py")) - assert set(manager._find_files_in_bundle(bundle)) == expected + found_rel_paths = {f.rel_path for f in known_files.get("testing", set())} + assert found_rel_paths == expected @pytest.mark.parametrize( "safe_mode", @@ -3361,7 +3382,7 @@ def _refresh_with_mocked_state(self, manager, bundle, initial_state): with ( mock_get as patched_get, mock_update as patched_update, - mock.patch.object(manager, "_find_files_in_bundle", return_value=[]), + mock.patch.object(DagImporterRegistry, "list_dag_files", return_value=[]), mock.patch.object(manager, "deactivate_deleted_dags"), mock.patch.object(manager, "clear_orphaned_import_errors"), mock.patch.object(manager, "handle_removed_files"), @@ -3440,7 +3461,7 @@ def test_refresh_dag_bundles_versioned_version_unchanged_persist_failure(self): manager, "get_bundle_state", return_value=BundleState(last_refreshed=None, version="v1") ), mock.patch.object(manager, "update_bundle_state", side_effect=Exception("DB error")), - mock.patch.object(manager, "_find_files_in_bundle", return_value=[]) as mock_find, + mock.patch.object(DagImporterRegistry, "list_dag_files", return_value=[]) as mock_list_files, mock.patch.object(manager, "deactivate_deleted_dags"), mock.patch.object(manager, "clear_orphaned_import_errors"), mock.patch.object(manager, "handle_removed_files"), @@ -3451,7 +3472,7 @@ def test_refresh_dag_bundles_versioned_version_unchanged_persist_failure(self): bundle.refresh.assert_called_once() # Short-circuit continues to next bundle — no file scanning - mock_find.assert_not_called() + mock_list_files.assert_not_called() assert "mock_bundle" not in known_files # _bundle_versions unchanged assert manager._bundle_versions["mock_bundle"] == "v1" @@ -3509,7 +3530,7 @@ def test_refresh_dag_bundles_initialize_non_airflow_exception_skips_bundle(self) manager, "get_bundle_state", return_value=BundleState(last_refreshed=None, version=None) ), mock.patch.object(manager, "update_bundle_state"), - mock.patch.object(manager, "_find_files_in_bundle", return_value=[]), + mock.patch.object(DagImporterRegistry, "list_dag_files", return_value=[]), mock.patch.object(manager, "deactivate_deleted_dags"), mock.patch.object(manager, "clear_orphaned_import_errors"), mock.patch.object(manager, "handle_removed_files"), @@ -3536,7 +3557,7 @@ def test_refresh_dag_bundles_update_bundle_state_failure_still_scans_files(self) manager, "get_bundle_state", return_value=BundleState(last_refreshed=None, version=None) ), mock.patch.object(manager, "update_bundle_state", side_effect=Exception("API error")), - mock.patch.object(manager, "_find_files_in_bundle", return_value=[]), + mock.patch.object(DagImporterRegistry, "list_dag_files", return_value=[]), mock.patch.object(manager, "deactivate_deleted_dags"), mock.patch.object(manager, "clear_orphaned_import_errors"), mock.patch.object(manager, "handle_removed_files"), diff --git a/devel-common/src/tests_common/pytest_plugin.py b/devel-common/src/tests_common/pytest_plugin.py index e2d9913205f11..cdb91deaeb337 100644 --- a/devel-common/src/tests_common/pytest_plugin.py +++ b/devel-common/src/tests_common/pytest_plugin.py @@ -2147,6 +2147,28 @@ def reset_team_name_cache(): clear_team_name_cache() +@pytest.fixture(autouse=True) +def reset_dag_bundle_and_importer_config_cache(): + """Reset the per-process Dag bundle and importer configuration cache between tests.""" + if importlib.util.find_spec("airflow") is None: + yield + return + + def _clear() -> None: + try: + from airflow.dag_processing.importers import reset_importer_registry + + reset_importer_registry() + except (ImportError, AttributeError): + pass + + _clear() + try: + yield + finally: + _clear() + + @pytest.fixture(autouse=True) def refuse_to_run_test_from_wrongly_named_files(request: pytest.FixtureRequest): filepath = request.node.path diff --git a/devel-common/src/tests_common/test_utils/config.py b/devel-common/src/tests_common/test_utils/config.py index 9893da665a4a7..9b60393925ad7 100644 --- a/devel-common/src/tests_common/test_utils/config.py +++ b/devel-common/src/tests_common/test_utils/config.py @@ -98,6 +98,7 @@ def conf_vars(overrides): if "airflow.configuration" in sys.modules: settings.configure_vars() + _clear_dag_bundle_and_importer_caches() try: yield @@ -116,6 +117,25 @@ def conf_vars(overrides): if "airflow.configuration" in sys.modules: settings.configure_vars() + _clear_dag_bundle_and_importer_caches() + + +def _clear_dag_bundle_and_importer_caches() -> None: + """Drop cached Dag bundle snapshot and importer registry so new config is read.""" + import sys + + importers = sys.modules.get("airflow.dag_processing.importers.base") + if importers is not None: + reset_fn = getattr(importers, "reset_importer_registry", None) + if reset_fn is not None: + reset_fn() + return + + manager = sys.modules.get("airflow.dag_processing.bundles.manager") + if manager is not None: + cache = getattr(manager, "load_bundle_config_snapshot", None) + if cache is not None: + cache.cache_clear() @overload