diff --git a/airflow-core/src/airflow/dag_processing/dagbag.py b/airflow-core/src/airflow/dag_processing/dagbag.py index 295eb52536b96..2ec6ee4afe1d5 100644 --- a/airflow-core/src/airflow/dag_processing/dagbag.py +++ b/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 diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index 8c4fd0588eb87..9827cc590829c 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/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.""" diff --git a/airflow-core/src/airflow/utils/file.py b/airflow-core/src/airflow/utils/file.py index ff7d7db16623b..9dd55f462deea 100644 --- a/airflow-core/src/airflow/utils/file.py +++ b/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. diff --git a/airflow-core/tests/unit/dag_processing/test_dagbag.py b/airflow-core/tests/unit/dag_processing/test_dagbag.py index 54b06ce75aa51..5fd7d51c64609 100644 --- a/airflow-core/tests/unit/dag_processing/test_dagbag.py +++ b/airflow-core/tests/unit/dag_processing/test_dagbag.py @@ -1040,8 +1040,13 @@ 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")) @@ -1049,21 +1054,21 @@ def test_refresh_packaged_dag(self, mock_dagmodel, test_zip_path): 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 diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index a3624e1d4a143..fb15c2486c632 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -72,7 +72,7 @@ from airflow.models.team import Team from airflow.providers.standard.operators.empty import EmptyOperator from airflow.sdk import DAG as SdkDAG -from airflow.sdk.importers import DagSourceCode +from airflow.sdk.importers import DagDefinition, DagImporterRegistry, DagImportError, DagSourceCode from airflow.serialization.serialized_objects import LazyDeserializedDAG from airflow.utils.net import get_hostname from airflow.utils.session import create_session @@ -137,6 +137,13 @@ def encode_mtime_in_filename(val): return out +def _make_bundle(path: Path) -> MagicMock: + bundle = MagicMock(spec=BaseDagBundle) + bundle.name = "testing" + bundle.path = path + return bundle + + def _create_zip_bundle_with_valid_and_broken_dags(zip_path: Path) -> None: with zipfile.ZipFile(zip_path, "w") as zf: zf.writestr( @@ -400,18 +407,11 @@ def test_clear_orphaned_import_errors_keeps_zip_inner_file_errors(self, session, ) session.flush() + bundle = _make_bundle(tmp_path) manager = DagFileProcessorManager(max_runs=1) manager.clear_orphaned_import_errors( bundle_name="testing", - observed_filelocs=manager._get_observed_filelocs( - { - DagFileInfo( - bundle_name="testing", - rel_path=Path("test_zip.zip"), - bundle_path=tmp_path, - ) - } - ), + observed_filelocs=manager._get_observed_filelocs(manager._find_files_in_bundle(bundle)), session=session, ) session.flush() @@ -420,26 +420,116 @@ def test_clear_orphaned_import_errors_keeps_zip_inner_file_errors(self, session, assert len(import_errors) == 1 assert import_errors[0].filename == "test_zip.zip/broken_dag.py" - def test_get_observed_filelocs_expands_zip_inner_paths(self, tmp_path): + def test_find_files_in_bundle_parses_zip_once_and_observes_its_members(self, tmp_path): zip_path = tmp_path / "test_zip.zip" _create_zip_bundle_with_valid_and_broken_dags(zip_path) + bundle = _make_bundle(tmp_path) manager = DagFileProcessorManager(max_runs=1) - observed_filelocs = manager._get_observed_filelocs( - { - DagFileInfo( - bundle_name="testing", - rel_path=Path("test_zip.zip"), - bundle_path=tmp_path, - ) - } - ) + found_files = manager._find_files_in_bundle(bundle) - assert observed_filelocs == { + assert found_files == { + DagFileInfo(bundle_name="testing", rel_path=Path("test_zip.zip"), bundle_path=tmp_path) + } + assert manager._get_observed_filelocs(found_files) == { + "test_zip.zip", "test_zip.zip/valid_dag.py", "test_zip.zip/broken_dag.py", } + @pytest.mark.parametrize("absolute", [False, True], ids=["relative", "absolute"]) + @mock.patch("airflow.dag_processing.manager.get_importer_registry", autospec=True) + def test_find_files_in_bundle_makes_discovery_errors_bundle_relative( + self, mock_get_registry, tmp_path, absolute + ): + (tmp_path / "broken.dag").write_text("") + source_reference = os.fspath(tmp_path / "broken.dag") if absolute else "broken.dag" + mock_get_registry.return_value.list_dag_definitions.return_value = [ + (mock.sentinel.importer, DagImportError(source_reference=source_reference, message="boom")) + ] + bundle = _make_bundle(tmp_path) + + found_files = DagFileProcessorManager(max_runs=1)._find_files_in_bundle(bundle) + + assert [file.definition_locs for file in found_files] == [frozenset({"broken.dag"})] + assert found_files == { + DagFileInfo(bundle_name="testing", rel_path=Path("broken.dag"), bundle_path=tmp_path) + } + + @pytest.mark.parametrize( + "item", + [ + DagImportError(source_reference="missing.dag", message="boom"), + MagicMock(spec=DagDefinition, **{"get_relative_loc.return_value": "virtual/dag"}), + ], + ids=["error", "definition"], + ) + @mock.patch("airflow.dag_processing.manager.get_importer_registry", autospec=True) + def test_find_files_in_bundle_warns_about_items_no_file_holds( + self, mock_get_registry, tmp_path, item, caplog + ): + mock_get_registry.return_value.list_dag_definitions.return_value = [(mock.sentinel.importer, item)] + + found_files = DagFileProcessorManager(max_runs=1)._find_files_in_bundle(_make_bundle(tmp_path)) + + assert found_files == set() + assert { + "event": f"Ignoring {item!r} listed in bundle testing: no file in the bundle holds it", + "log_level": "warning", + } in caplog + + @pytest.mark.parametrize("absolute", [False, True], ids=["relative", "absolute"]) + @mock.patch("airflow.dag_processing.manager.get_importer_registry", autospec=True) + def test_find_files_in_bundle_warns_about_items_outside_the_bundle( + self, mock_get_registry, tmp_path, absolute, caplog + ): + (tmp_path / "outside.dag").write_text("") + bundle_path = tmp_path / "bundle" + bundle_path.mkdir() + loc = os.fspath(tmp_path / "outside.dag") if absolute else "../outside.dag" + item = MagicMock(spec=DagDefinition, **{"get_relative_loc.return_value": loc}) + mock_get_registry.return_value.list_dag_definitions.return_value = [(mock.sentinel.importer, item)] + + found_files = DagFileProcessorManager(max_runs=1)._find_files_in_bundle(_make_bundle(bundle_path)) + + assert found_files == set() + assert { + "event": f"Ignoring {item!r} listed in bundle testing: it resolves outside the bundle", + "log_level": "warning", + } in caplog + + @conf_vars( + { + ( + "dag_processor", + "dag_importer_configs", + ): '[{"classpath": "airflow.sdk.importers.PythonDagImporter", "extensions": [".dag"]}]' + } + ) + def test_find_files_in_bundle_uses_configured_importers(self, tmp_path): + (tmp_path / "custom.dag").write_text("from airflow.sdk import DAG\n") + bundle = _make_bundle(tmp_path) + + manager = DagFileProcessorManager(max_runs=1) + found_files = manager._find_files_in_bundle(bundle) + + assert [file.definition_locs for file in found_files] == [frozenset({"custom.dag"})] + assert found_files == { + DagFileInfo(bundle_name="testing", rel_path=Path("custom.dag"), bundle_path=tmp_path) + } + + def test_find_files_in_bundle_honours_airflowignore(self, tmp_path): + dag_source = "from airflow.sdk import DAG\n" + (tmp_path / ".airflowignore").write_text("ignored.py\nignored.zip\n") + (tmp_path / "kept.py").write_text(dag_source) + (tmp_path / "ignored.py").write_text(dag_source) + with zipfile.ZipFile(tmp_path / "ignored.zip", "w") as zf: + zf.writestr("dag.py", dag_source) + + found_files = DagFileProcessorManager(max_runs=1)._find_files_in_bundle(_make_bundle(tmp_path)) + + assert {file.rel_path for file in found_files} == {Path("kept.py")} + def test_sync_bundles_deactivates_missing_when_owning_all_bundles(self): """A processor with no bundle filter owns the full config and may deactivate missing bundles.""" manager = DagFileProcessorManager(max_runs=1) @@ -464,9 +554,7 @@ def test_sync_bundles_does_not_deactivate_missing_when_filtered(self): def test_find_files_in_bundle_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 = _make_bundle(tmp_path) with conf_vars({("core", "dag_discovery_safe_mode"): str(safe_mode)}): manager = DagFileProcessorManager(max_runs=1) @@ -474,7 +562,8 @@ def test_find_files_in_bundle_respects_dag_discovery_safe_mode(self, tmp_path, s 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_files = manager._find_files_in_bundle(bundle) + assert {file.rel_path for file in found_files} == expected @pytest.mark.parametrize( "safe_mode", @@ -483,7 +572,7 @@ def test_find_files_in_bundle_respects_dag_discovery_safe_mode(self, tmp_path, s pytest.param(True, id="safe-mode-on-filters-keywordless"), ], ) - def test_get_observed_filelocs_respects_dag_discovery_safe_mode(self, tmp_path, safe_mode): + def test_find_files_in_bundle_observed_zip_members_respect_safe_mode(self, tmp_path, safe_mode): """ZIP-member discovery used for deactivation must honor the configured safe_mode. With ``dag_discovery_safe_mode=False`` a keyword-less (wrapped) zip member is parsed and @@ -493,13 +582,13 @@ def test_get_observed_filelocs_respects_dag_discovery_safe_mode(self, tmp_path, zip_path = tmp_path / "test_zip.zip" _create_zip_bundle_with_keywordless_dag(zip_path) + bundle = _make_bundle(tmp_path) + with conf_vars({("core", "dag_discovery_safe_mode"): str(safe_mode)}): manager = DagFileProcessorManager(max_runs=1) - observed_filelocs = manager._get_observed_filelocs( - {DagFileInfo(bundle_name="testing", rel_path=Path("test_zip.zip"), bundle_path=tmp_path)} - ) + observed_filelocs = manager._get_observed_filelocs(manager._find_files_in_bundle(bundle)) - expected = {"test_zip.zip/with_keywords.py"} + expected = {"test_zip.zip", "test_zip.zip/with_keywords.py"} if not safe_mode: expected.add("test_zip.zip/no_keywords.py") assert observed_filelocs == expected @@ -1346,6 +1435,29 @@ def test_deactivate_stale_dags_marks_dags_in_inactive_bundles(self, session): ) assert is_stale_by_dag == {"dag_in_inactive_bundle": True, "dag_in_active_bundle": False} + @pytest.mark.usefixtures("testing_dag_bundle") + def test_deactivate_stale_dags_matches_dags_nested_in_a_container(self, session): + session.add( + DagModel( + dag_id="nested_dag", + bundle_name="testing", + relative_fileloc="dags.zip/sub/nested_dag.py", + last_parsed_time=timezone.utcnow(), + is_stale=False, + ) + ) + session.flush() + + manager = DagFileProcessorManager(max_runs=1, processor_timeout=10 * 60) + manager.deactivate_stale_dags( + last_parsed={ + DagFileInfo(rel_path=Path("dags.zip"), bundle_name="testing"): timezone.utcnow() + + timedelta(hours=1) + } + ) + + assert session.scalar(select(DagModel.is_stale).where(DagModel.dag_id == "nested_dag")) + @pytest.mark.usefixtures("testing_dag_bundle") def test_deactivate_stale_dags_marks_dags_with_null_bundle_name(self, session): """Dags carried over from Airflow 2.x keep a NULL bundle_name and must still be deactivated. @@ -2091,6 +2203,25 @@ def test_deactivate_deleted_dags(self, dag_maker, session): # and the DAG from test_dag2.py is deactivated assert session.get(DagModel, "test_dag2").is_stale is True + def test_deactivate_deleted_dags_keeps_dags_of_observed_definitions(self, dag_maker, session): + with dag_maker("zipped_dag") as dag: + dag.relative_fileloc = "dags.zip/zipped_dag.py" + dag_maker.sync_dagbag_to_db() + + DagFileProcessorManager(max_runs=1).deactivate_deleted_dags( + "dag_maker", + { + DagFileInfo( + bundle_name="dag_maker", + rel_path=Path("dags.zip"), + bundle_path=TEST_DAGS_FOLDER, + definition_locs=frozenset({"dags.zip/zipped_dag.py"}), + ) + }, + ) + + assert session.get(DagModel, "zipped_dag").is_stale is False + @mock.patch("airflow.dag_processing.manager.update_dag_parsing_results_in_db", autospec=True) def test_persist_parsing_result_passes_parsed_definitions_and_source_codes(self, mock_update): source_codes = {"/bundle/dags.zip/a.py": DagSourceCode(source_code="src", language="python")} @@ -3118,12 +3249,12 @@ def test_bundles_are_refreshed(self): bundleone = MagicMock() bundleone.name = "bundleone" - bundleone.path = "/dev/null" + bundleone.path = Path("/dev/null") bundleone.refresh_interval = 0 bundleone.get_current_version.return_value = None bundletwo = MagicMock() bundletwo.name = "bundletwo" - bundletwo.path = "/dev/null" + bundletwo.path = Path("/dev/null") bundletwo.refresh_interval = 300 bundletwo.get_current_version.return_value = None @@ -3175,7 +3306,7 @@ def test_bundle_refresh_check_interval(self): bundleone = MagicMock() bundleone.name = "bundleone" - bundleone.path = "/dev/null" + bundleone.path = Path("/dev/null") bundleone.refresh_interval = 0 bundleone.get_current_version.return_value = None @@ -3246,7 +3377,7 @@ def test_bundle_force_refresh(self): bundleone = MagicMock() bundleone.name = "bundleone" - bundleone.path = "/dev/null" + bundleone.path = Path("/dev/null") bundleone.refresh_interval = 0 bundleone.get_current_version.return_value = None @@ -3276,7 +3407,7 @@ def test_bundles_versions_are_stored(self, session): mybundle = MagicMock() mybundle.name = "bundleone" - mybundle.path = "/dev/null" + mybundle.path = Path("/dev/null") mybundle.refresh_interval = 0 mybundle.supports_versioning = True mybundle.get_current_version.return_value = "123" @@ -3696,7 +3827,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(manager, "_find_files_in_bundle", return_value=set()), mock.patch.object(manager, "deactivate_deleted_dags"), mock.patch.object(manager, "clear_orphaned_import_errors"), mock.patch.object(manager, "handle_removed_files"), @@ -3775,7 +3906,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(manager, "_find_files_in_bundle", return_value=set()) as mock_find, mock.patch.object(manager, "deactivate_deleted_dags"), mock.patch.object(manager, "clear_orphaned_import_errors"), mock.patch.object(manager, "handle_removed_files"), @@ -3844,7 +3975,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(manager, "_find_files_in_bundle", return_value=set()), mock.patch.object(manager, "deactivate_deleted_dags"), mock.patch.object(manager, "clear_orphaned_import_errors"), mock.patch.object(manager, "handle_removed_files"), @@ -3871,7 +4002,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(manager, "_find_files_in_bundle", return_value=set()), mock.patch.object(manager, "deactivate_deleted_dags"), mock.patch.object(manager, "clear_orphaned_import_errors"), mock.patch.object(manager, "handle_removed_files"), @@ -3887,6 +4018,98 @@ def test_refresh_dag_bundles_update_bundle_state_failure_still_scans_files(self) # iteration will see a version mismatch and re-refresh rather than skip incorrectly assert "mock_bundle" not in manager._bundle_versions + def test_refresh_dag_bundles_discovery_failure_keeps_known_files_and_dags(self): + """A failed listing keeps the bundle's known files and does not deactivate its Dags.""" + manager = DagFileProcessorManager(max_runs=1) + bundle = self._make_refresh_bundle() + manager._dag_bundles = [bundle] + known = {DagFileInfo(bundle_name="mock_bundle", rel_path=Path("dag.py"), bundle_path=bundle.path)} + known_files = {"mock_bundle": known} + + 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, "_find_files_in_bundle", side_effect=OSError("listing failed")), + mock.patch.object(manager, "deactivate_deleted_dags", autospec=True) as mock_deactivate, + mock.patch.object(manager, "clear_orphaned_import_errors", autospec=True) as mock_clear, + 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) + + mock_deactivate.assert_not_called() + mock_clear.assert_not_called() + assert known_files == {"mock_bundle": known} + + def test_refresh_dag_bundles_discovery_failure_lists_again_on_next_refresh(self): + """A failed listing must not advance the bundle version, or the next refresh skips the bundle.""" + manager = DagFileProcessorManager(max_runs=1, bundle_refresh_check_interval=0) + bundle = self._make_refresh_bundle(supports_versioning=True, current_version="v2") + manager._dag_bundles = [bundle] + manager._bundle_versions["mock_bundle"] = "v1" + found = {DagFileInfo(bundle_name="mock_bundle", rel_path=Path("dag.py"), bundle_path=bundle.path)} + + with ( + mock.patch.object( + manager, "get_bundle_state", return_value=BundleState(last_refreshed=None, version="v1") + ), + mock.patch.object(manager, "update_bundle_state") as mock_update, + mock.patch.object( + manager, "_find_files_in_bundle", side_effect=[OSError("listing failed"), found] + ) as mock_find, + 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({}) + mock_update.assert_not_called() + assert manager._bundle_versions["mock_bundle"] == "v1" + + known_files: dict[str, set[DagFileInfo]] = {} + manager._refresh_dag_bundles(known_files) + + assert mock_find.call_count == 2 + assert known_files == {"mock_bundle": found} + assert manager._bundle_versions["mock_bundle"] == "v2" + + @mock.patch("airflow.dag_processing.manager.gc", autospec=True) + def test_before_run_warms_importers_before_freezing_heap(self, mock_gc): + manager = DagFileProcessorManager(max_runs=1) + calls = MagicMock() + calls.attach_mock(mock_gc.freeze, "freeze") + with ( + mock.patch.object(manager, "prepare_server_process_context"), + mock.patch.object(manager, "prepare_process_context"), + mock.patch.object(manager, "register_exit_signals"), + mock.patch.object(manager, "prepare_bundles"), + mock.patch.object(manager, "_symlink_latest_log_directory"), + mock.patch.object(manager, "warm_importers") as mock_warm, + ): + calls.attach_mock(mock_warm, "warm_importers") + manager.before_run() + + assert calls.mock_calls == [mock.call.warm_importers(), mock.call.freeze()] + + @mock.patch("airflow.dag_processing.manager.get_importer_registry", autospec=True) + def test_warm_importers_failure_does_not_stop_other_bundles(self, mock_get_registry, caplog): + manager = DagFileProcessorManager(max_runs=1) + broken, healthy = MagicMock(spec=BaseDagBundle), MagicMock(spec=BaseDagBundle) + broken.name, healthy.name = "broken", "healthy" + manager._dag_bundles = [broken, healthy] + registries = {name: MagicMock(spec=DagImporterRegistry) for name in ("broken", "healthy")} + registries["broken"].warm_importers.side_effect = ImportError("no module") + mock_get_registry.side_effect = registries.__getitem__ + + manager.warm_importers() + + registries["healthy"].warm_importers.assert_called_once_with() + assert {"event": "Error loading Dag importers for bundle broken", "log_level": "error"} in caplog + def test_unpack_bundle_version_with_bundle_version_dataclass(self): from airflow.dag_processing.bundles.base import BundleVersion, unpack_bundle_version diff --git a/airflow-core/tests/unit/models/test_dag.py b/airflow-core/tests/unit/models/test_dag.py index e821d3efe8d9a..46edb019731f4 100644 --- a/airflow-core/tests/unit/models/test_dag.py +++ b/airflow-core/tests/unit/models/test_dag.py @@ -100,7 +100,6 @@ OnceTimetable, ) from airflow.triggers.base import TriggerEvent -from airflow.utils.file import list_py_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 +1177,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=[], ) 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..38fcd1534ac5f 100644 --- a/airflow-core/tests/unit/utils/test_file.py +++ b/airflow-core/tests/unit/utils/test_file.py @@ -29,15 +29,12 @@ from airflow.utils import file as file_utils from airflow.utils.file import ( correct_maybe_zipped, - list_py_file_paths, open_maybe_zipped, ) from tests_common.test_utils.config import conf_vars from unit.models import TEST_DAGS_FOLDER -TEST_DAG_FOLDER = os.environ["AIRFLOW__CORE__DAGS_FOLDER"] - def might_contain_dag(file_path: str, zip_file: zipfile.ZipFile | None = None): return False @@ -96,7 +93,7 @@ def test_open_maybe_zipped_archive(self, test_zip_path): assert isinstance(content, str) -class TestListPyFilesPath: +class TestDagFileHelpers: def test_find_path_from_directory_regex_ignore(self): should_ignore = [ "test_invalid_cron.py", @@ -181,38 +178,6 @@ def test_get_modules_from_invalid_file(self): assert len(modules) == 0 - def test_list_py_file_paths(self, test_zip_path): - detected_files = set() - expected_files = set() - # No_dags is empty, _invalid_ is ignored by .airflowignore - ignored_files = { - "no_dags.py", - "should_ignore_this.py", - "test_explicit_ignore.py", - "test_invalid_cron.py", - "test_invalid_dup_task.py", - "test_ignore_this.py", - "test_invalid_param.py", - "test_invalid_param2.py", - "test_invalid_param3.py", - "test_invalid_param4.py", - "test_nested_dag.py", - "test_imports.py", - "test_nested_negate_ignore.py", - "file_no_airflow_dag.py", # no_dag test case in test_zip folder - "test.py", # no_dag test case in test_zip_module folder - "__init__.py", - } - for root, _, files in os.walk(TEST_DAG_FOLDER): - for file_name in files: - 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)) - assert detected_files == expected_files, ( - f"Detected files mismatched expected files:\ndetected_files: {pformat(detected_files)}\nexpected_files: {pformat(expected_files)}" - ) - @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..8c799031129b9 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 diff --git a/task-sdk/src/airflow/sdk/importers/base.py b/task-sdk/src/airflow/sdk/importers/base.py index 82798727fb25d..afc72feee0d06 100644 --- a/task-sdk/src/airflow/sdk/importers/base.py +++ b/task-sdk/src/airflow/sdk/importers/base.py @@ -245,11 +245,13 @@ def list_dag_definitions( safe_mode: bool = True, ) -> Iterator[DefT | DagImportError]: """ - List DAG definitions in a bundle that this importer can handle (identity-only discovery). + List Dag definitions in a bundle that this importer can handle. - A yielded :class:`DagImportError` reports a discovery-time failure (e.g. an unreadable - container) for the caller to forward to a :class:`DagImportResult`; it is not a source - to import. + Apply :meth:`might_contain_dag` to each definition before yielding it; nothing + applies it after listing. A definition that cannot be read is yielded as a + :class:`DagImportError` rather than raised, so the rest of the bundle is still + listed. A yielded :class:`DagImportError` reports a discovery-time failure for the + caller to forward to a :class:`DagImportResult`; it is not a source to import. """ @abstractmethod @@ -266,11 +268,12 @@ def get_source_code(self, definition: DagDefinition) -> DagSourceCode: def might_contain_dag(self, definition: DagDefinition, safe_mode: bool) -> bool: """ - Cheap, optional pre-check for whether a discovered definition may contain a DAG. + Cheap, optional pre-check for whether a definition may contain a Dag. The default returns True (keep the definition): an importer that can only tell by - attempting the import leaves this as-is. Importers with a cheap content heuristic - override it, so obvious non-DAG sources are dropped during discovery. + attempting the import leaves this as-is. An importer with a cheap content heuristic + overrides it and applies it in :meth:`list_dag_definitions`, so obvious non-Dag + sources never get a parse process. """ return True @@ -501,10 +504,7 @@ def list_dag_definitions( it, such as the members of an archive. A :class:`DagImportError` item is a discovery-time failure rather than a source to import. """ - # A spec registered for several extensions appears once per extension, and - # materialising it drops all of them, so take one pending spec at a time. - while self._extension_specs: - self._materialise_spec(next(iter(self._extension_specs.values()))) + self.warm_importers() for importer in self._ordered_importers: for item in importer.list_dag_definitions(bundle, safe_mode=safe_mode): @@ -517,6 +517,18 @@ def list_dag_definitions( continue yield importer, item + def warm_importers(self) -> None: + """ + Instantiate every configured importer that has not been instantiated yet. + + The Dag processor calls this before freezing its heap, so forked parse processes + share the importers instead of building them. + """ + # A spec registered for several extensions appears once per extension, and + # materialising it drops all of them, so take one pending spec at a time. + while self._extension_specs: + self._materialise_spec(next(iter(self._extension_specs.values()))) + def _materialise_spec(self, spec: _ImporterSpec) -> AbstractDagImporter[Any]: """Instantiate a configured spec and take over every extension it was registered for.""" importer = self._instantiate_spec(spec) diff --git a/task-sdk/src/airflow/sdk/importers/python_importer.py b/task-sdk/src/airflow/sdk/importers/python_importer.py index 74d82ff0024a5..5a884f707d641 100644 --- a/task-sdk/src/airflow/sdk/importers/python_importer.py +++ b/task-sdk/src/airflow/sdk/importers/python_importer.py @@ -151,14 +151,24 @@ def list_dag_definitions( A lightweight content sniff (``might_contain_dag``) is applied here so files that clearly hold no Dag never become definitions -- keeping the discovered set (and the - eventual parse-process count) close to the number of real Dag files. Zip members are - discovered by :class:`..zip_importer.ZipImporter`, not here. + eventual parse-process count) close to the number of real Dag files. A file that cannot + be read is reported as a DagImportError and the remaining files are still yielded. Zip + members are discovered by :class:`..zip_importer.ZipImporter`, not here. """ for definition in find_file_dag_definitions(bundle.path, self.supported_extensions): - if self.might_contain_dag(definition, safe_mode): - yield definition - else: - log.debug("Skipping %r: no Airflow Dag markers found", definition) + try: + if not self.might_contain_dag(definition, safe_mode): + log.debug("Skipping %r: no Airflow Dag markers found", definition) + continue + except Exception as e: + log.warning("Cannot read Dag file %s: %s", definition.path, e) + yield DagImportError( + source_reference=definition.get_relative_loc(bundle.path), + message=f"Failed to read Dag file: {e}", + error_type="read_error", + ) + continue + yield definition def import_definition( self, diff --git a/task-sdk/tests/task_sdk/importers/test_python_importer.py b/task-sdk/tests/task_sdk/importers/test_python_importer.py index 4e517f2a79945..54cfdc94df664 100644 --- a/task-sdk/tests/task_sdk/importers/test_python_importer.py +++ b/task-sdk/tests/task_sdk/importers/test_python_importer.py @@ -35,6 +35,7 @@ from airflow.sdk.exceptions import AirflowConfigException from airflow.sdk.importers import ( + DagImportError, FileDagDefinition, FilesystemDagDefinition, PythonDagImporter, @@ -208,6 +209,25 @@ def test_list_dag_definitions(self, mock_bundle): defs = list(importer.list_dag_definitions(mock_bundle)) assert {d.path.name for d in defs} == {"sample_dag.py"} + def test_list_dag_definitions_reports_unreadable_file_and_lists_the_rest(self, mock_bundle): + (mock_bundle.path / "bad.py").write_text("from airflow.sdk import DAG\n") + (mock_bundle.path / "good.py").write_text("from airflow.sdk import DAG\n") + + def _might_contain_dag(self, definition, safe_mode): + if definition.path.name == "bad.py": + raise PermissionError("Permission denied") + return True + + with mock.patch.object( + PythonDagImporter, "might_contain_dag", autospec=True, side_effect=_might_contain_dag + ): + items = list(PythonDagImporter().list_dag_definitions(mock_bundle)) + + [error] = [item for item in items if isinstance(item, DagImportError)] + assert error.source_reference == "bad.py" + assert error.error_type == "read_error" + assert [item.path.name for item in items if not isinstance(item, DagImportError)] == ["good.py"] + def test_list_prefers_source_over_pyc_and_skips_pycache(self, mock_bundle): (mock_bundle.path / "foo.py").write_text("from airflow.sdk import DAG\n") (mock_bundle.path / "foo.pyc").write_bytes(b"compiled") # side-by-side -> skipped diff --git a/task-sdk/tests/task_sdk/importers/test_registry.py b/task-sdk/tests/task_sdk/importers/test_registry.py index 87e80d456a615..9ce47e0103903 100644 --- a/task-sdk/tests/task_sdk/importers/test_registry.py +++ b/task-sdk/tests/task_sdk/importers/test_registry.py @@ -296,6 +296,20 @@ def test_list_dag_definitions_materialises_a_multi_extension_spec_once(self, tmp assert repr(definition) == str(archive / "member_dag.py") assert registry.get_importer("dags.zipx") is importer + def test_warm_importers_instantiates_pending_specs(self): + LazyTestImporter.instances = 0 + reg = DagImporterRegistry(register_defaults=False) + reg.register_specs( + [{"classpath": f"{__name__}.LazyTestImporter", "extensions": [".lazy", ".lazy2"]}], + context="test", + ) + + reg.warm_importers() + + assert LazyTestImporter.instances == 1 + assert reg.get_importer("file.lazy") is reg.get_importer("file.lazy2") + assert LazyTestImporter.instances == 1 + def test_lazy_importer_instantiation(self): """Importer classes are not imported or instantiated until get_importer is called.""" LazyTestImporter.instances = 0