This is an automated email from the ASF dual-hosted git repository.
uranusjr pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 4803c0b3f79 Discover Dag definitions through the Task SDK importers
(#74020)
4803c0b3f79 is described below
commit 4803c0b3f797fc57b92b1d5c0e44a8b19917f7a6
Author: Dilnaz Amanzholova <[email protected]>
AuthorDate: Fri Oct 2 17:03:30 2026 +0200
Discover Dag definitions through the Task SDK importers (#74020)
---
airflow-core/src/airflow/dag_processing/dagbag.py | 10 +-
airflow-core/src/airflow/dag_processing/manager.py | 131 +++++----
airflow-core/src/airflow/utils/file.py | 48 ----
.../tests/unit/dag_processing/test_dagbag.py | 30 +-
.../tests/unit/dag_processing/test_manager.py | 301 ++++++++++++++++++---
airflow-core/tests/unit/models/test_dag.py | 3 +-
airflow-core/tests/unit/utils/test_file.py | 37 +--
generated/known_sdk_imports_in_core.txt | 2 +-
task-sdk/src/airflow/sdk/importers/base.py | 34 ++-
.../src/airflow/sdk/importers/python_importer.py | 22 +-
.../task_sdk/importers/test_python_importer.py | 20 ++
task-sdk/tests/task_sdk/importers/test_registry.py | 14 +
12 files changed, 437 insertions(+), 215 deletions(-)
diff --git a/airflow-core/src/airflow/dag_processing/dagbag.py
b/airflow-core/src/airflow/dag_processing/dagbag.py
index 295eb52536b..2ec6ee4afe1 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.models.pool import Pool
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 @@ class DagBag(LoggingMixin):
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 @@ class DagBag(LoggingMixin):
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 8c4fd0588eb..9827cc59082 100644
--- a/airflow-core/src/airflow/dag_processing/manager.py
+++ b/airflow-core/src/airflow/dag_processing/manager.py
@@ -29,7 +29,6 @@ import selectors
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.db_callback_request import
DbCallbackRequest
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 @@ from airflow.utils.sqlalchemy import (
)
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 @@ class DagFileProcessorManager(LoggingMixin):
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 @@ class DagFileProcessorManager(LoggingMixin):
# 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 @@ class DagFileProcessorManager(LoggingMixin):
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 @@ class DagFileProcessorManager(LoggingMixin):
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 @@ class DagFileProcessorManager(LoggingMixin):
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 ff7d7db1662..9dd55f462de 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 @@ from airflow._shared.module_loading import (
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 54b06ce75aa..5fd7d51c646 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 @@ class TestDagBag:
@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 @@ class TestDagBag:
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 @@ class TestDagBag:
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 a3624e1d4a1..fb15c2486c6 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.serialized_dag import SerializedDagModel
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 @@ class TestDagFileProcessorManager:
)
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 @@ class TestDagFileProcessorManager:
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 @@ class TestDagFileProcessorManager:
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 @@ class TestDagFileProcessorManager:
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 @@ class TestDagFileProcessorManager:
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 @@ class TestDagFileProcessorManager:
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 @@ class TestDagFileProcessorManager:
)
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 @@ class TestDagFileProcessorManager:
# 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 @@ class TestDagFileProcessorManager:
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 @@ class TestDagFileProcessorManager:
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 @@ class TestDagFileProcessorManager:
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 @@ class TestDagFileProcessorManager:
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 @@ class TestDagFileProcessorManager:
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 @@ class TestDagFileProcessorManager:
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 @@ class TestDagFileProcessorManager:
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 @@ class TestDagFileProcessorManager:
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 @@ class TestDagFileProcessorManager:
# 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 fadf893d0de..96a7ccd9725 100644
--- a/airflow-core/tests/unit/models/test_dag.py
+++ b/airflow-core/tests/unit/models/test_dag.py
@@ -100,7 +100,6 @@ from airflow.timetables.simple import (
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
@@ -1180,7 +1179,7 @@ class TestDag:
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 11502777128..38fcd1534ac 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.configuration import conf
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 @@ class TestOpenMaybeZipped:
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 @@ class TestListPyFilesPath:
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 b906ba25855..8c799031129 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 82798727fb2..afc72feee0d 100644
--- a/task-sdk/src/airflow/sdk/importers/base.py
+++ b/task-sdk/src/airflow/sdk/importers/base.py
@@ -245,11 +245,13 @@ class AbstractDagImporter(ABC, Generic[DefT]):
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 @@ class AbstractDagImporter(ABC, Generic[DefT]):
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 @@ class DagImporterRegistry:
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 @@ class DagImporterRegistry:
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 74d82ff0024..5a884f707d6 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 @@ class
PythonDagImporter(AbstractDagImporter[FileDagDefinition]):
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 4e517f2a799..54cfdc94df6 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 @@ import pytest
from airflow.sdk.exceptions import AirflowConfigException
from airflow.sdk.importers import (
+ DagImportError,
FileDagDefinition,
FilesystemDagDefinition,
PythonDagImporter,
@@ -208,6 +209,25 @@ class TestPythonDagImporter:
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 87e80d456a6..9ce47e01039 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 @@ class TestDagImporterRegistry:
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
