Discover Dag definitions through the Task SDK importers

airflow-core/src/airflow/dag_processing/dagbag.py

@@ -47,7 +47,7 @@
 from airflow.sdk.importers import DagImportError, get_importer_registry
 from airflow.serialization.definitions.notset import NOTSET, ArgNotSet, is_arg_set
 from airflow.serialization.serialized_objects import LazyDeserializedDAG
-from airflow.utils.file import correct_maybe_zipped, find_enclosing_file
+from airflow.utils.file import find_enclosing_file
 from airflow.utils.log.logging_mixin import LoggingMixin
 from airflow.utils.session import NEW_SESSION, provide_session
 
@@ -276,9 +276,7 @@ def get_dag(self, dag_id, *, session: Session = NEW_SESSION):
             self.dags.pop(dag_id, None)
         if dag is None or is_expired:
             # Reprocess source file.
-            found_dags = self.process_file(
-                filepath=correct_maybe_zipped(orm_dag.fileloc), only_if_updated=False
-            )
+            found_dags = self.process_file(filepath=orm_dag.fileloc, only_if_updated=False)
 
             # If the source file no longer exports `dag_id`, delete it from self.dags
             if found_dags and dag_id in [found_dag.dag_id for found_dag in found_dags]:
@@ -287,8 +285,8 @@ def get_dag(self, dag_id, *, session: Session = NEW_SESSION):
         return self.dags.get(dag_id)
 
     def process_file(self, filepath, only_if_updated=True, safe_mode=True):
-        """Process a DAG file and return found DAGs."""
-        if filepath is None or not os.path.isfile(filepath):
+        """Process a Dag file, or a Dag definition nested in one, and return found Dags."""
+        if filepath is None or os.path.isdir(filepath):
             return []
         return [
             dag

airflow-core/src/airflow/dag_processing/manager.py

@@ -29,7 +29,6 @@
 import signal
 import sys
 import time
-import zipfile
 from collections import OrderedDict, defaultdict
 from dataclasses import dataclass, field
 from datetime import datetime, timedelta
@@ -65,9 +64,10 @@
 from airflow.models.errors import ParseImportError
 from airflow.observability.metrics import stats_utils
 from airflow.sdk import SecretCache
+from airflow.sdk.importers import DagImportError, 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 find_enclosing_file
 from airflow.utils.helpers import prune_dict
 from airflow.utils.log.logging_mixin import LoggingMixin
 from airflow.utils.net import get_hostname
@@ -84,7 +84,7 @@
 )
 
 if TYPE_CHECKING:
-    from collections.abc import Callable, Collection, Iterable, Iterator, Sequence
+    from collections.abc import Callable, Collection, Iterable, Sequence
     from socket import socket
 
     from sqlalchemy.orm import Session
@@ -133,6 +133,7 @@ class DagFileInfo:
     bundle_name: str
     bundle_path: Path | None = field(compare=False, default=None)
     bundle_version: str | None = None
+    definition_locs: frozenset[str] = field(compare=False, default=frozenset())
 
     @property
     def absolute_path(self) -> Path:
@@ -367,12 +368,22 @@ def before_run(self) -> None:
         self.log.info("Process each file at most once every %s seconds", self._file_process_interval)
         self.prepare_bundles()
         self._symlink_latest_log_directory()
+        self.warm_importers()
         # To prevent COW in forked process parsing dag file
         gc.freeze()
 
     def after_run(self) -> None:
         """Tear down state after the parsing loop exits. Default no-op; override to customize."""
 
+    def warm_importers(self) -> None:
+        """Build each bundle's Dag importers, so parse processes forked later share them."""
+        for bundle in self._dag_bundles:
+            try:
+                get_importer_registry(bundle.name).warm_importers()
+            except Exception:
+                # The importer fails again when the bundle is listed, which reports it per refresh.
+                self.log.exception("Error loading Dag importers for bundle %s", bundle.name)
+
     def prepare_server_process_context(self) -> None:
         """
         Mark this process as running in "server" context so MetastoreBackend is available.
@@ -504,11 +515,17 @@ def deactivate_stale_dags(
             # Dag file's last_finish_time, the Dag is considered stale as has apparently been removed from the file,
             # This is especially relevant for Dag files that generate Dags in a dynamic manner.
             rel_path = Path(dag.relative_fileloc)
-            file_info = DagFileInfo(rel_path=rel_path, bundle_name=dag.bundle_name)
-            if file_info not in last_parsed:
-                # Zip-packaged dags are keyed by the archive path, not the inner file, so try the parent as well
-                file_info = DagFileInfo(rel_path=rel_path.parent, bundle_name=dag.bundle_name)
-            if last_finish_time := last_parsed.get(file_info, None):
+            # A Dag nested in a container (``archive.zip/sub/dag.py``) is parsed under the container.
+            last_finish_time = next(
+                (
+                    last_parsed[file_info]
+                    for candidate in (rel_path, *rel_path.parents)
+                    if (file_info := DagFileInfo(rel_path=candidate, bundle_name=dag.bundle_name))
+                    in last_parsed
+                ),
+                None,
+            )
+            if last_finish_time:
                 if dag.last_parsed_time + timedelta(seconds=self.stale_dag_threshold) < last_finish_time:
                     self.log.info(
                         "Deactivating stale DAG %s. Not parsed for %s seconds (last parsed: %s).",
@@ -940,6 +957,14 @@ def _refresh_dag_bundles(self, known_files: dict[str, set[DagFileInfo]]):
                 version_after_refresh = None
                 version_data_after_refresh = None
 
+            try:
+                found_files = self._find_files_in_bundle(bundle)
+            except Exception:
+                # Keep the bundle's known files and Dags, and leave its version unadvanced so the
+                # next refresh lists it again.
+                self.log.exception("Error listing Dag definitions in bundle %s", bundle.name)
+                continue
+
             # Persistence failure must not skip file scanning (bundle is already refreshed locally).
             # _bundle_versions is only advanced on success to stay consistent with the DB.
             try:
@@ -950,11 +975,6 @@ 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
 
-            found_files = {
-                DagFileInfo(rel_path=p, bundle_name=bundle.name, bundle_path=bundle.path)
-                for p in self._find_files_in_bundle(bundle)
-            }
-
             known_files[bundle.name] = found_files
 
             self.deactivate_deleted_dags(bundle_name=bundle.name, present=found_files)
@@ -970,56 +990,53 @@ 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)
-        ]
+    def _find_files_in_bundle(self, bundle: BaseDagBundle) -> set[DagFileInfo]:
+        """
+        List the files to parse in a bundle through its importers.
+
+        A file holding several Dag definitions (a zip archive, for instance) is parsed as one.
+        """
+        self.log.info("Searching for Dag definitions in %s at %s", bundle.name, bundle.path)
+        registry = get_importer_registry(bundle.name)
+        definition_locs: defaultdict[Path, set[str]] = defaultdict(set)
+        for _, item in registry.list_dag_definitions(bundle, safe_mode=self.dag_discovery_safe_mode):
+            if isinstance(item, DagImportError):
+                # Importers report a source either absolutely or relative to the bundle.
+                rel_fileloc = os.path.relpath(bundle.path / item.source_reference, bundle.path)
+            else:
+                rel_fileloc = item.get_relative_loc(bundle.path)
+            loc = Path(os.path.normpath(bundle.path / rel_fileloc))
+            if not loc.is_relative_to(bundle.path):
+                self.log.warning(
+                    "Ignoring %r listed in bundle %s: it resolves outside the bundle", item, bundle.name
+                )
+                continue
+            if (path := find_enclosing_file(loc)) is None:
+                self.log.warning(
+                    "Ignoring %r listed in bundle %s: no file in the bundle holds it", item, bundle.name
+                )
+                continue
+            definition_locs[path.relative_to(bundle.path)].add(rel_fileloc)
         self.log.info(
             "Found %s files for bundle %s (dag_discovery_safe_mode=%s)",
-            len(rel_paths),
+            len(definition_locs),
             bundle.name,
             self.dag_discovery_safe_mode,
         )
+        return {
+            DagFileInfo(
+                rel_path=rel_path,
+                bundle_name=bundle.name,
+                bundle_path=bundle.path,
+                definition_locs=frozenset(locs),
+            )
+            for rel_path, locs in definition_locs.items()
+        }
 
-        return rel_paths
-
-    def _get_observed_filelocs(self, present: set[DagFileInfo]) -> set[str]:
-        """
-        Return observed DAG source paths for bundle entries.
-
-        For regular files this includes the relative file path.
-        For ZIP archives this includes DAG-like inner paths such as
-        ``archive.zip/dag.py``.
-        """
-
-        def find_zipped_dags(abs_path: os.PathLike) -> Iterator[str]:
-            """Yield absolute paths for DAG-like files inside a ZIP archive."""
-            try:
-                with zipfile.ZipFile(abs_path) as z:
-                    for info in z.infolist():
-                        # Use the configured discovery safe mode
-                        if might_contain_dag(info.filename, self.dag_discovery_safe_mode, z, conf=conf):
-                            yield os.path.join(abs_path, info.filename)
-            except zipfile.BadZipFile:
-                self.log.exception("There was an error accessing ZIP file %s", abs_path)
-
-        observed_filelocs: set[str] = set()
-        for info in present:
-            abs_path = str(info.absolute_path)
-            if abs_path.endswith(".py") or not zipfile.is_zipfile(abs_path):
-                observed_filelocs.add(str(info.rel_path))
-            else:
-                if TYPE_CHECKING:
-                    assert info.bundle_path
-                for abs_sub_path in find_zipped_dags(abs_path=info.absolute_path):
-                    rel_sub_path = Path(abs_sub_path).relative_to(info.bundle_path)
-                    observed_filelocs.add(str(rel_sub_path))
-
-        return observed_filelocs
+    @staticmethod
+    def _get_observed_filelocs(present: set[DagFileInfo]) -> set[str]:
+        """Return the bundle-relative locations of the files and of the definitions found in them."""
+        return {loc for file in present for loc in (str(file.rel_path), *file.definition_locs)}
 
     def deactivate_deleted_dags(self, bundle_name: str, present: set[DagFileInfo]) -> None:
         """Deactivate DAGs that come from files that are no longer present in bundle."""

airflow-core/src/airflow/utils/file.py

@@ -18,7 +18,6 @@
 from __future__ import annotations
 
 import ast
-import logging
 import os
 import re
 import zipfile
@@ -32,10 +31,6 @@
     might_contain_dag as might_contain_dag,
     might_contain_dag_via_default_heuristic as might_contain_dag_via_default_heuristic,
 )
-from airflow.configuration import conf
-
-log = logging.getLogger(__name__)
-
 
 ZIP_REGEX = re.compile(rf"((.*\.zip){re.escape(os.sep)})?(.*)")
 
@@ -76,49 +71,6 @@ def open_maybe_zipped(fileloc, mode="r"):
     return open(fileloc, mode=mode)
 
 
-def list_py_file_paths(
-    directory: str | os.PathLike[str] | None,
-    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
-    """
-    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))
-    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."""
-    from airflow._shared.module_loading.file_discovery import find_path_from_directory
-
-    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)
-        except Exception:
-            log.exception("Error while examining %s", file_path)
-
-    return file_paths
-
-
 def find_enclosing_file(path: Path) -> Path | None:
     """
     Return ``path`` or its nearest ancestor that is a file, or ``None`` if there is none.

airflow-core/tests/unit/dag_processing/test_dagbag.py

@@ -1040,30 +1040,35 @@ def _process_definition(self, importer, definition, *, only_if_updated):
     @patch.object(DagModel, "get_current")
     def test_refresh_packaged_dag(self, mock_dagmodel, test_zip_path):
         """
-        Test that we can refresh a packaged DAG
+        Test that refreshing a packaged Dag re-imports only the archive member that defines it
         """
+        with zipfile.ZipFile(test_zip_path, "a") as zf:
+            zf.writestr(
+                "other_dag.py",
+                "from airflow.sdk import DAG\n\ndag = DAG(dag_id='other_dag', schedule=None)\n",
+            )
         dag_id = "test_zip_dag"
         fileloc = os.path.realpath(os.path.join(test_zip_path, "test_zip.py"))
 
         mock_dagmodel.return_value = DagModel()
         mock_dagmodel.return_value.last_expired = datetime.max.replace(tzinfo=timezone.utc)
         mock_dagmodel.return_value.fileloc = fileloc
 
-        class _TestDagBag(DagBag):
-            import_calls = 0
+        processed: list[str] = []
 
+        class _TestDagBag(DagBag):
             def _process_definition(self, importer, definition, *, only_if_updated):
-                if repr(definition) == fileloc:
-                    _TestDagBag.import_calls += 1
+                processed.append(repr(definition))
                 return super()._process_definition(importer, definition, only_if_updated=only_if_updated)
 
         dagbag = _TestDagBag(dag_folder=os.path.realpath(test_zip_path))
+        assert "other_dag" in dagbag.dags
+        processed.clear()
 
-        assert dagbag.import_calls == 1
         dag = dagbag.get_dag(dag_id)
-        assert dag is not None
-        assert dag_id == dag.dag_id
-        assert dagbag.import_calls == 2
+
+        assert dag.dag_id == dag_id
+        assert processed == [fileloc]
 
     def process_dag(self, create_dag, tmp_path):
         """
@@ -1125,6 +1130,13 @@ def basic_cycle():
         self.validate_dags(test_dag, found_dags, dagbag, should_be_found=False)
         assert file_path in dagbag.import_errors
 
+    def test_process_file_nested_definition(self, tmp_path, test_zip_path):
+        dagbag = DagBag(dag_folder=os.fspath(tmp_path), collect_dags=False)
+
+        found_dags = dagbag.process_file(os.path.join(test_zip_path, "test_zip.py"))
+
+        assert sorted(dag.dag_id for dag in found_dags) == ["test_zip_autoregister", "test_zip_dag"]
+
     def test_process_file_with_none(self, tmp_path):
         """
         test that process_file can handle Nones