diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index 8c4fd0588eb87..3bdbabb2814e4 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -65,9 +65,10 @@ from airflow.models.errors import ParseImportError from airflow.observability.metrics import stats_utils from airflow.sdk import SecretCache +from airflow.sdk.importers import get_importer_registry 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 list_dag_file_paths, 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 @@ -972,11 +973,12 @@ def _refresh_dag_bundles(self, known_files: dict[str, set[DagFileInfo]]): 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 + # Discovery Dag files with the bundle's Dag importer registry. + registry = get_importer_registry(bundle.name) 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) + for x in list_dag_file_paths(bundle.path, registry, safe_mode=self.dag_discovery_safe_mode) ] self.log.info( "Found %s files for bundle %s (dag_discovery_safe_mode=%s)", diff --git a/airflow-core/src/airflow/utils/file.py b/airflow-core/src/airflow/utils/file.py index ff7d7db16623b..5549f84968f98 100644 --- a/airflow-core/src/airflow/utils/file.py +++ b/airflow-core/src/airflow/utils/file.py @@ -25,7 +25,7 @@ from collections.abc import Generator from io import TextIOWrapper from pathlib import Path -from typing import overload +from typing import TYPE_CHECKING, overload from airflow._shared.module_loading import ( get_unique_dag_module_name as get_unique_dag_module_name, @@ -34,6 +34,9 @@ ) from airflow.configuration import conf +if TYPE_CHECKING: + from airflow.sdk.importers import DagImporterRegistry + log = logging.getLogger(__name__) @@ -76,19 +79,22 @@ def open_maybe_zipped(fileloc, mode="r"): return open(fileloc, mode=mode) -def list_py_file_paths( +def list_dag_file_paths( directory: str | os.PathLike[str] | None, + registry: DagImporterRegistry, + *, safe_mode: bool = conf.getboolean("core", "DAG_DISCOVERY_SAFE_MODE", fallback=True), ) -> list[str]: """ - Traverse a directory and look for Python files. - - :param directory: the directory to traverse - :param safe_mode: whether to use a heuristic to determine whether a file - contains Airflow DAG definitions. If not provided, use the - core.DAG_DISCOVERY_SAFE_MODE configuration setting. If not set, default - to safe. - :return: a list of paths to Python files in the specified directory + Traverse a directory and return paths of Dag files. + + :param directory: The directory to traverse. + :param registry: The Dag importer registry for the bundle. A file is + discovered only when an importer in the registry claims its extension, + and that importer's own ``might_contain_dag`` accepts it. + :param safe_mode: Passed to each importer for Dag file detection heuristic. + If not provided, use the core.DAG_DISCOVERY_SAFE_MODE configuration. + :return: A list of paths to candidate Dag files in the specified directory. """ file_paths: list[str] = [] if directory is None: @@ -96,23 +102,28 @@ def list_py_file_paths( elif os.path.isfile(directory): file_paths = [str(directory)] elif os.path.isdir(directory): - file_paths.extend(find_dag_file_paths(directory, safe_mode)) + file_paths.extend(find_dag_file_paths(directory, registry, safe_mode)) return file_paths -def find_dag_file_paths(directory: str | os.PathLike[str], safe_mode: bool) -> list[str]: - """Find file paths of all DAG files.""" +def find_dag_file_paths( + directory: str | os.PathLike[str], + registry: DagImporterRegistry, + safe_mode: bool, +) -> list[str]: + """Find paths of candidate Dag files under *directory*.""" from airflow._shared.module_loading.file_discovery import find_path_from_directory + from airflow.sdk.importers import FilesystemDagDefinition file_paths = [] 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 == ".py" or zipfile.is_zipfile(path)): - if might_contain_dag(file_path, safe_mode, conf=conf): - file_paths.append(file_path) + if not (path := Path(file_path)).is_file() or (importer := registry.get_importer(path)) is None: + continue + if importer.might_contain_dag(FilesystemDagDefinition(path), safe_mode): + file_paths.append(file_path) except Exception: log.exception("Error while examining %s", file_path) diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index a3624e1d4a143..8719d8b4f355d 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -284,6 +284,15 @@ def _statement_breakdown(counts: Counter[tuple[str, str]]) -> str: SWEEP_CALLS = 4 +class _FakeYamlImporter: + """Minimal duck-typed Dag importer: claims .yaml/.yml and detects dags its own way.""" + + supported_extensions = [".yaml", ".yml"] + + def might_contain_dag(self, definition, safe_mode): + return not safe_mode or b"dag_id" in definition.read_bytes() + + class LateResolvingBundle(BaseDagBundle): supports_versioning = True @@ -476,6 +485,40 @@ def test_find_files_in_bundle_respects_dag_discovery_safe_mode(self, tmp_path, s expected.add(Path("no_keywords.py")) assert set(manager._find_files_in_bundle(bundle)) == expected + def test_find_files_in_bundle_uses_registry_importers(self, tmp_path, monkeypatch): + """Discovery gates each file by its importer's might_contain_dag; a YAML importer finds .yaml. + + Closes the discovery half of the AIP-85 importer gap: a registered YAML importer makes the + manager hand .yaml files to parsing, accepted by the importer's own heuristic (here in + safe mode, with no "airflow" marker). Default bundles discover only the built-in + importers' extensions. + """ + from airflow.sdk.importers import DagImporterRegistry + + (tmp_path / "pipeline.py").write_text("from airflow.sdk import DAG\n") + (tmp_path / "pipeline.yaml").write_text("dag_id: d\n") # no "airflow" marker + (tmp_path / "notes.txt").write_text("ignore me\n") + bundle = MagicMock(spec=BaseDagBundle) + bundle.name = "testing" + bundle.path = tmp_path + + def _make(yaml): + reg = DagImporterRegistry() # built-in Python/Zip importers + if yaml: + reg.register(_FakeYamlImporter(), extensions=[".yaml", ".yml"]) + return reg + + # Built-in importers only: .yaml is not discovered. + monkeypatch.setattr("airflow.sdk.importers.get_importer_registry", lambda name=None: _make(False)) + manager = DagFileProcessorManager(max_runs=1) + assert set(manager._find_files_in_bundle(bundle)) == {Path("pipeline.py")} + + # With a YAML importer: pipeline.yaml is discovered via its own might_contain_dag; + # notes.txt has no importer and is ignored. + monkeypatch.setattr("airflow.sdk.importers.get_importer_registry", lambda name=None: _make(True)) + manager = DagFileProcessorManager(max_runs=1) + assert set(manager._find_files_in_bundle(bundle)) == {Path("pipeline.py"), Path("pipeline.yaml")} + @pytest.mark.parametrize( "safe_mode", [ diff --git a/airflow-core/tests/unit/models/test_dag.py b/airflow-core/tests/unit/models/test_dag.py index e821d3efe8d9a..2d8bd0171fb76 100644 --- a/airflow-core/tests/unit/models/test_dag.py +++ b/airflow-core/tests/unit/models/test_dag.py @@ -89,6 +89,7 @@ from airflow.sdk.definitions.param import Param from airflow.sdk.exceptions import TaskAwaitingInput from airflow.sdk.execution_time.hitl import upsert_hitl_detail +from airflow.sdk.importers import get_importer_registry from airflow.serialization.definitions.dag import SerializedDAG from airflow.serialization.encoders import coerce_to_core_timetable from airflow.serialization.serialized_objects import LazyDeserializedDAG @@ -100,7 +101,7 @@ OnceTimetable, ) from airflow.triggers.base import TriggerEvent -from airflow.utils.file import list_py_file_paths +from airflow.utils.file import list_dag_file_paths from airflow.utils.session import create_session from airflow.utils.state import DagRunState, DagSchedulingState, State, TaskInstanceState from airflow.utils.types import DagRunTriggeredByType, DagRunType @@ -1178,7 +1179,7 @@ def test_dag_is_deactivated_upon_dagfile_deletion(self, dag_maker): DagModel.deactivate_deleted_dags( bundle_name=orm_dag.bundle_name, - rel_filelocs=list_py_file_paths(settings.DAGS_FOLDER), + rel_filelocs=list_dag_file_paths(settings.DAGS_FOLDER, get_importer_registry()), ) orm_dag = session.scalar(select(DagModel).where(DagModel.dag_id == dag_id)) diff --git a/airflow-core/tests/unit/utils/test_file.py b/airflow-core/tests/unit/utils/test_file.py index 115027771284e..a195e6618a0f1 100644 --- a/airflow-core/tests/unit/utils/test_file.py +++ b/airflow-core/tests/unit/utils/test_file.py @@ -19,6 +19,7 @@ import os import zipfile +from pathlib import Path from pprint import pformat from unittest import mock @@ -29,7 +30,8 @@ from airflow.utils import file as file_utils from airflow.utils.file import ( correct_maybe_zipped, - list_py_file_paths, + find_dag_file_paths, + list_dag_file_paths, open_maybe_zipped, ) @@ -43,6 +45,26 @@ def might_contain_dag(file_path: str, zip_file: zipfile.ZipFile | None = None): return False +class _FakeYamlImporter: + """Minimal duck-typed Dag importer: claims .yaml/.yml and detects dags its own way.""" + + supported_extensions = [".yaml", ".yml"] + + def might_contain_dag(self, definition, safe_mode): + # A YAML-specific heuristic -- note it does NOT require the "airflow" marker the Python + # importer looks for, so it accepts a plain `dag_id:` file in safe mode. + return not safe_mode or b"dag_id" in definition.read_bytes() + + +def _registry(*, yaml=False, defaults=True): + from airflow.sdk.importers import DagImporterRegistry + + reg = DagImporterRegistry(register_defaults=defaults) + if yaml: + reg.register(_FakeYamlImporter(), extensions=[".yaml", ".yml"]) + return reg + + class TestCorrectMaybeZipped: @mock.patch("zipfile.is_zipfile") def test_correct_maybe_zipped_normal_file(self, mocked_is_zipfile): @@ -181,7 +203,7 @@ def test_get_modules_from_invalid_file(self): assert len(modules) == 0 - def test_list_py_file_paths(self, test_zip_path): + def test_list_dag_file_paths(self, test_zip_path): detected_files = set() expected_files = set() # No_dags is empty, _invalid_ is ignored by .airflowignore @@ -208,11 +230,39 @@ def test_list_py_file_paths(self, test_zip_path): if file_name.endswith((".py", ".zip")): if file_name not in ignored_files: expected_files.add(f"{root}/{file_name}") - detected_files = set(list_py_file_paths(TEST_DAG_FOLDER)) + detected_files = set(list_dag_file_paths(TEST_DAG_FOLDER, _registry())) assert detected_files == expected_files, ( f"Detected files mismatched expected files:\ndetected_files: {pformat(detected_files)}\nexpected_files: {pformat(expected_files)}" ) + def test_find_dag_file_paths_uses_each_importers_might_contain_dag(self, tmp_path): + """Each file is gated by ITS importer's might_contain_dag, not a single shared sniff.""" + (tmp_path / "a.py").write_text("from airflow.sdk import DAG\n") # airflow + dag markers + (tmp_path / "b.yaml").write_text("dag_id: d\n") # NO "airflow" marker + (tmp_path / "c.txt").write_text("not a dag\n") + + # No YAML importer registered: only .py is discovered (via the Python importer). + assert {Path(x).name for x in find_dag_file_paths(tmp_path, _registry(), safe_mode=True)} == {"a.py"} + + # With a YAML importer: its own heuristic accepts the .yaml in safe mode even though it + # lacks the "airflow" marker the Python importer requires. .txt has no importer -> skipped. + assert { + Path(x).name for x in find_dag_file_paths(tmp_path, _registry(yaml=True), safe_mode=True) + } == {"a.py", "b.yaml"} + + def test_find_dag_file_paths_follows_registered_importers(self, tmp_path): + """Dropping the built-in importers drops their extensions -- nothing is hardcoded.""" + (tmp_path / "a.py").write_text("from airflow.sdk import DAG\n") + (tmp_path / "b.yaml").write_text("dag_id: d\n") + reg = _registry(yaml=True, defaults=False) # only the YAML importer, no Python/Zip + assert {Path(x).name for x in find_dag_file_paths(tmp_path, reg, safe_mode=False)} == {"b.yaml"} + + def test_list_dag_file_paths_forwards_registry(self, tmp_path): + (tmp_path / "a.py").write_text("from airflow.sdk import DAG\n") + (tmp_path / "b.yaml").write_text("dag_id: d\n") + names = {Path(x).name for x in list_dag_file_paths(tmp_path, _registry(yaml=True), safe_mode=False)} + assert names == {"a.py", "b.yaml"} + @pytest.mark.parametrize( ("edge_filename", "expected_modification"), diff --git a/generated/known_sdk_imports_in_core.txt b/generated/known_sdk_imports_in_core.txt index b906ba258554b..c487454fb6912 100644 --- a/generated/known_sdk_imports_in_core.txt +++ b/generated/known_sdk_imports_in_core.txt @@ -4,7 +4,7 @@ airflow-core/src/airflow/cli/commands/task_command.py::7 airflow-core/src/airflow/cli/commands/triggerer_command.py::1 airflow-core/src/airflow/configuration.py::1 airflow-core/src/airflow/dag_processing/dagbag.py::3 -airflow-core/src/airflow/dag_processing/manager.py::4 +airflow-core/src/airflow/dag_processing/manager.py::5 airflow-core/src/airflow/dag_processing/processor.py::14 airflow-core/src/airflow/exceptions.py::1 airflow-core/src/airflow/executors/base_executor.py::3 @@ -38,5 +38,6 @@ airflow-core/src/airflow/utils/dag_cycle_tester.py::1 airflow-core/src/airflow/utils/dag_edges.py::1 airflow-core/src/airflow/utils/dot_renderer.py::3 airflow-core/src/airflow/utils/edgemodifier.py::1 +airflow-core/src/airflow/utils/file.py::2 airflow-core/src/airflow/utils/json.py::4 airflow-core/src/airflow/utils/operator_helpers.py::2