Route dag file discovery to importers

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)",

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,43 +79,51 @@ 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:
         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)
 

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",
         [

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))

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"),

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