This is an automated email from the ASF dual-hosted git repository.
jason810496 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 b23ce7e9f5f Identity-only DAG discovery and byte-importing (#73118)
b23ce7e9f5f is described below
commit b23ce7e9f5f12fc39bb9429e515f3a12b53a216e
Author: Tzu-ping Chung <[email protected]>
AuthorDate: Wed Sep 23 21:35:03 2026 +0800
Identity-only DAG discovery and byte-importing (#73118)
* Identity-only DAG discovery and byte-importing
Discovery now filters purely by a source's identity rather than its
contents. find_file_dag_definitions and list_dag_definitions match on
extension and .airflowignore, skip __pycache__, and prefer a .py over a
side-by-side .pyc, without reading any file or taking a safe_mode flag.
The might_contain_dag heuristic instead runs a single time, in
import_definition; plain files were previously scanned both at discovery
and again at import, so this drops a redundant scan while still reporting
non-DAG sources through DagImportResult.skipped_definitions.
ZipImporter now processes one member at a time instead of a whole
archive. It enumerates individual members during discovery, applying the
same __pycache__ and .py/.pyc rules and the ZipSlip defence there, and
imports a single member per call, rather than listing whole .zip files
and walking their members inside import_definition.
Modules are loaded from a definition's bytes. Source and bytecode are
read through the definition with custom loaders in place of the
file-backed SourceFileLoader and SourcelessFileLoader, so archive members
are no longer extracted to a temporary file to be scanned or imported and
.py and .pyc members load in place. Accordingly, might_contain_dag now
takes a definition rather than a path, and PythonDagImporter is
parameterized over the file and zip-member definition types it handles.
Tests cover identity-only discovery, per-member zip import, import-time
skipping for both files and members, .py/.pyc dedup, .pyc import for a
file and a member, and corrupt-bytecode handling.
* Rework Task SDK DAG importers around a file-source abstraction
Discovery now filters only by identity (extension, .airflowignore,
__pycache__, .py-over-.pyc); DAG detection runs once at import for both
files and zip members, surfaced through skipped_definitions. ZipImporter
processes one member at a time, and modules load from a definition's
bytes instead of an extracted file.
File sources are modelled as an abstract FileDagDefinition with
FilesystemDagDefinition and ZipMemberDagDefinition backings; importers
bind to the interface, so PythonDagImporter handles archive members
without referencing zip. The shared might_contain_dag also accepts a
definition, reading its bytes in place and falling back to as_file() only
for a custom might_contain_dag_callable; existing path/zip callers are
unchanged.
* Rename private function
* Dedup .py/.pyc over discovered files
find_file_dag_definitions preferred a source over a side-by-side .pyc by
checking path.with_suffix(".py").exists(), a raw filesystem stat.
This dropped a .pyc whose .py sibling was on disk but never discovered,
excluded by .airflowignore, or not in the importer's supported extensions,
leaving nothing to import.
Decide the preference over the candidates the walk actually yields instead,
keyed by (directory, stem) so the pairing also matches case-insensitively on
the extension. Candidates keep the walk's order.
* Reject .pyc payloads that aren't code objects
_DefinitionBytecodeLoader.get_code passed whatever marshal.loads returned
straight to exec(), so a .pyc with a valid magic header but a payload that
unmarshals to a non-code object (e.g. a marshalled string) was run as source
instead of rejected. Validate isinstance(code, types.CodeType) and raise
ImportError otherwise, keeping the corrupt-.pyc handling meaningful past the
magic-number check.
* Reject sibling-resource reads from source loader
_DefinitionSourceLoader.get_data ignored its path argument and always
returned
the definition's own bytes, so a pkgutil.get_data or __loader__.get_data
call
for a sibling resource silently got the DAG source instead.
Raise FileNotFoundError when the requested path isn't the module's own
get_filename() value; the normal import path, which only asks for that
value,
is unchanged.
* Report unreadable archives as import errors
list_dag_definitions is widened to yield DagDefinition | DagImportError so a
ZipImporter can be reported for the caller to forward into DagImportResult.
* Have definitions control the import environment
A zip member is imported correctly only with its archive on sys.path, but
that
setup lived in ZipImporter.import_definition. A member routed through the
registry by its .py suffix went to PythonDagImporter, which imported it
without
the archive on sys.path and failed cross-member imports.
Move the setup onto the definition: FileDagDefinition.import_context is
no-op
by default; ZipMemberDagDefinition overrides it to put the archive on
sys.path,
and PythonDagImporter wraps the load in it.
* Use os.sep form for zip member filelocs
ZIP_REGEX in airflow-core splits on os.sep, so we need to match. Join with
Path.joinpath instead (i.e. archive.zip/member.py) instead.
* Move content sniffing back to discovery
Identity-only discovery surfaced every extension match as a definition, so a
bundle of 10 DAGs beside a 40-module helper package became 50 definitions,
and,
once wired, 50 parse processes. Content detection at import shrank each
process's work but not their count.
Move the lightweight might_contain_dag sniff back into discovery:
list_dag_definitions applies it (gated by a safe_mode parameter) so a source
that clearly holds no DAG never becomes a definition. might_contain_dag is a
method on AbstractDagImporter that returns True by default, overridden by
PythonDagImporter with the real heuristic; ZipImporter calls it on each
member's internal importer, so importers with no cheap check (non-Python
ones)
keep every member. The import path no longer re-sniffs.
safe_mode moves from import_definition (where it was a redundant second
check)
to list_dag_definitions, where the filter decision actually lives; it is
dropped from import_definition and the ABC.
* Harden Dag discovery for zip members and sourceless packages
Review of the identity-only discovery rewrite surfaced four cases where
the new byte-based paths behave differently from the file-backed ones
they replace.
The shipped configuration names the built-in heuristic in
might_contain_dag_callable, so the callable is never absent and every
scanned archive member was still written to a temporary file, which is
the cost this rewrite set out to remove. A member that failed to read
ended discovery for its whole archive, losing later valid Dags that the
file-backed path reported individually. The source-over-bytecode
preference was decided before importer lookup, so an unsupported source
could hide the bytecode beside it, and a member whose extension differed
only in case was never paired at all. A bytes-backed __init__.pyc was
treated as an ordinary module, so a sourceless package's relative
imports failed where the file-backed loader resolved them.
---------
Co-authored-by: LIU ZHE YOU <[email protected]>
---
.../src/airflow_shared/module_loading/__init__.py | 1 +
.../src/airflow_shared/module_loading/dag_file.py | 89 ++++++---
.../tests/module_loading/test_dag_file.py | 101 ++++++++++
task-sdk/src/airflow/sdk/importers/__init__.py | 6 +-
task-sdk/src/airflow/sdk/importers/base.py | 144 +++++++++++----
.../src/airflow/sdk/importers/python_importer.py | 204 ++++++++++++---------
task-sdk/src/airflow/sdk/importers/zip_importer.py | 192 +++++++++++--------
.../task_sdk/importers/test_python_importer.py | 193 ++++++++++++++++---
task-sdk/tests/task_sdk/importers/test_registry.py | 65 +++++--
.../tests/task_sdk/importers/test_zip_importer.py | 181 ++++++++++++++----
10 files changed, 870 insertions(+), 306 deletions(-)
diff --git
a/shared/module_loading/src/airflow_shared/module_loading/__init__.py
b/shared/module_loading/src/airflow_shared/module_loading/__init__.py
index 3bf8eacab8e..60bead945e3 100644
--- a/shared/module_loading/src/airflow_shared/module_loading/__init__.py
+++ b/shared/module_loading/src/airflow_shared/module_loading/__init__.py
@@ -30,6 +30,7 @@ from typing import TYPE_CHECKING
from .dag_file import (
MODIFIED_DAG_MODULE_NAME as MODIFIED_DAG_MODULE_NAME,
UNUSUAL_MODULE_PREFIX as UNUSUAL_MODULE_PREFIX,
+ accepts_dag_definition as accepts_dag_definition,
get_unique_dag_module_name as get_unique_dag_module_name,
might_contain_dag as might_contain_dag,
might_contain_dag_via_default_heuristic as
might_contain_dag_via_default_heuristic,
diff --git
a/shared/module_loading/src/airflow_shared/module_loading/dag_file.py
b/shared/module_loading/src/airflow_shared/module_loading/dag_file.py
index c4d95db8ba5..6b026e9096e 100644
--- a/shared/module_loading/src/airflow_shared/module_loading/dag_file.py
+++ b/shared/module_loading/src/airflow_shared/module_loading/dag_file.py
@@ -20,6 +20,7 @@
from __future__ import annotations
import hashlib
+import os
import re
import zipfile
from pathlib import Path
@@ -29,14 +30,26 @@ UNUSUAL_MODULE_PREFIX = "unusual_prefix_"
MODIFIED_DAG_MODULE_NAME =
f"{UNUSUAL_MODULE_PREFIX}{{path_hash}}_{{module_name}}"
if TYPE_CHECKING:
+ from contextlib import AbstractContextManager
from typing import Protocol
class _MightContainDagCallable(Protocol):
- def __call__(self, file_path: str, zip_file: zipfile.ZipFile | None =
None) -> bool: ...
+ def __call__(
+ self,
+ file_path: str | _DagDefinitionLike,
+ zip_file: zipfile.ZipFile | None = None,
+ ) -> bool: ...
class _ConfLike(Protocol):
def getimport(self, section: str, key: str, **kwargs: Any) -> Any: ...
+ class _DagDefinitionLike(Protocol):
+ """Structural view of a DagDefinition: read its bytes, or materialize
it as a file."""
+
+ def read_bytes(self) -> bytes: ...
+
+ def as_file(self) -> AbstractContextManager[Path]: ...
+
def get_unique_dag_module_name(file_path: str) -> str:
"""Return a unique module name in the format unusual_prefix_{sha1 of
module's file path}_{original module name}."""
@@ -47,41 +60,67 @@ def get_unique_dag_module_name(file_path: str) -> str:
raise ValueError("file_path should be a string to generate unique module
name")
-def might_contain_dag_via_default_heuristic(file_path: str, zip_file:
zipfile.ZipFile | None = None) -> bool:
+def accepts_dag_definition(func: _MightContainDagCallable) ->
_MightContainDagCallable:
+ """
+ Mark a ``might_contain_dag_callable`` as accepting a Dag definition, not
only a path.
+
+ A marked callable is handed the definition itself, so an archive member or
any other
+ non-filesystem source is checked without being written to a temporary file
first.
+ """
+ func.accepts_dag_definition = True # type: ignore[attr-defined]
+ return func
+
+
+@accepts_dag_definition
+def might_contain_dag_via_default_heuristic(
+ file_path: str | _DagDefinitionLike,
+ zip_file: zipfile.ZipFile | None = None,
+) -> bool:
"""
Heuristic that guesses whether a Python file contains an Airflow DAG
definition.
- :param file_path: Path to the file to be checked.
- :param zip_file: if passed, checks the archive. Otherwise, check local
filesystem.
+ :param file_path: path to the file to check, or a DagDefinition-like
object whose bytes
+ are read directly (nothing is read from disk).
+ :param zip_file: if passed, checks the named member inside the archive.
Otherwise, check
+ the local filesystem.
:return: True, if file might contain DAGs.
"""
- if zip_file:
+ if not isinstance(file_path, (str, os.PathLike)):
+ data = file_path.read_bytes()
+ elif zip_file:
with zip_file.open(file_path) as current_file:
- content = current_file.read()
+ data = current_file.read()
+ elif zipfile.is_zipfile(file_path):
+ return True
else:
- if zipfile.is_zipfile(file_path):
- return True
with open(file_path, "rb") as dag_file:
- content = dag_file.read()
- content = content.lower()
- if b"airflow" not in content:
+ data = dag_file.read()
+ data = data.lower()
+ if b"airflow" not in data:
return False
- return any(s in content for s in (b"dag", b"asset"))
+ return any(s in data for s in (b"dag", b"asset"))
def might_contain_dag(
- file_path: str,
- safe_mode: bool,
+ file_path: str | _DagDefinitionLike,
+ safe_mode: bool = True,
zip_file: zipfile.ZipFile | None = None,
*,
conf: _ConfLike,
) -> bool:
"""
- Check whether a Python file contains Airflow DAGs.
-
- When safe_mode is off (with False value), this function always returns
True.
-
- If might_contain_dag_callable isn't specified, it uses airflow default
heuristic.
+ Check whether a source might contain Airflow DAGs.
+
+ ``file_path`` may be a filesystem path (optionally with a ``zip_file``
archive whose
+ member it names) or a DagDefinition-like object exposing ``read_bytes()``
and
+ ``as_file()``. Passing a definition lets the check run against an
in-memory or
+ archive-backed source without materializing a file. When safe_mode is off
(with False
+ value), this function always returns True.
+
+ A callable marked with :func:`accepts_dag_definition`, including the
default heuristic,
+ is handed the definition and reads its bytes directly. Any other callable
only
+ understands the legacy ``(file_path, zip_file)`` signature, so a
definition is
+ materialized through its own ``as_file()`` and passed by path for
compatibility.
"""
if not safe_mode:
return True
@@ -102,6 +141,12 @@ def might_contain_dag(
)
if might_contain_dag_callable is None:
- might_contain_dag_callable = might_contain_dag_via_default_heuristic
-
- return might_contain_dag_callable(file_path=file_path, zip_file=zip_file)
+ return might_contain_dag_via_default_heuristic(file_path,
zip_file=zip_file)
+
+ if isinstance(file_path, (str, os.PathLike)) or getattr(
+ might_contain_dag_callable, "accepts_dag_definition", False
+ ):
+ return might_contain_dag_callable(file_path=file_path,
zip_file=zip_file)
+ # Legacy callables only accept (file_path, zip_file); let the definition
materialize itself.
+ with file_path.as_file() as materialized:
+ return might_contain_dag_callable(file_path=str(materialized),
zip_file=None)
diff --git a/shared/module_loading/tests/module_loading/test_dag_file.py
b/shared/module_loading/tests/module_loading/test_dag_file.py
index f802362f43c..7f97e4351a8 100644
--- a/shared/module_loading/tests/module_loading/test_dag_file.py
+++ b/shared/module_loading/tests/module_loading/test_dag_file.py
@@ -17,14 +17,19 @@
# under the License.
from __future__ import annotations
+import contextlib
import logging
+import tempfile
+from pathlib import Path
from unittest import mock
from airflow_shared.module_loading import (
MODIFIED_DAG_MODULE_NAME,
UNUSUAL_MODULE_PREFIX,
+ accepts_dag_definition,
get_unique_dag_module_name,
might_contain_dag,
+ might_contain_dag_via_default_heuristic,
)
@@ -74,3 +79,99 @@ def
test_might_contain_dag_logs_warning_on_broken_config(tmp_path, caplog) -> No
assert result is True
assert "Failed to load might_contain_dag_callable from config" in
caplog.text
+
+
+class _FakeDefinition:
+ """Structural DagDefinition: read its bytes, or materialize it as a
file."""
+
+ def __init__(self, data: bytes) -> None:
+ self._data = data
+
+ def read_bytes(self) -> bytes:
+ return self._data
+
+ @contextlib.contextmanager
+ def as_file(self):
+ with tempfile.NamedTemporaryFile(suffix=".py", delete=False) as f:
+ f.write(self._data)
+ path = Path(f.name)
+ try:
+ yield path
+ finally:
+ path.unlink(missing_ok=True)
+
+
+def test_might_contain_dag_from_definition() -> None:
+ # The default heuristic reads a definition's bytes directly, with nothing
on disk.
+ mock_conf = mock.MagicMock()
+ mock_conf.getimport.return_value = None
+
+ assert (
+ might_contain_dag(
+ _FakeDefinition(b"from airflow import DAG\ndag = DAG('x')"),
safe_mode=True, conf=mock_conf
+ )
+ is True
+ )
+ assert (
+ might_contain_dag(_FakeDefinition(b"def add(x, y): return x + y"),
safe_mode=True, conf=mock_conf)
+ is False
+ )
+ assert might_contain_dag(_FakeDefinition(b"anything"), safe_mode=False,
conf=mock_conf) is True
+
+
+def test_default_heuristic_accepts_definition() -> None:
+ # The default heuristic reads a definition's bytes directly, same as it
reads a path.
+ assert might_contain_dag_via_default_heuristic(_FakeDefinition(b"from
airflow import DAG")) is True
+ assert might_contain_dag_via_default_heuristic(_FakeDefinition(b"x = 1"))
is False
+
+
+class _NoFileDefinition(_FakeDefinition):
+ """A definition that refuses to materialize, proving the bytes path was
taken."""
+
+ def as_file(self):
+ raise AssertionError("definition should not be materialized")
+
+
+def test_configured_default_heuristic_reads_bytes() -> None:
+ # The shipped config points might_contain_dag_callable at the default
heuristic, so a
+ # definition must still be passed straight through rather than written to
a temp file.
+ mock_conf = mock.MagicMock()
+ mock_conf.getimport.return_value = might_contain_dag_via_default_heuristic
+
+ definition = _NoFileDefinition(b"from airflow import DAG")
+ assert might_contain_dag(definition, safe_mode=True, conf=mock_conf) is
True
+
+
+def test_marked_custom_callable_receives_definition() -> None:
+ # Any callable can opt into taking definitions by marking itself.
+ @accepts_dag_definition
+ def custom(file_path, zip_file=None):
+ return b"airflow" in file_path.read_bytes()
+
+ mock_conf = mock.MagicMock()
+ mock_conf.getimport.return_value = custom
+
+ definition = _NoFileDefinition(b"from airflow import DAG")
+ assert might_contain_dag(definition, safe_mode=True, conf=mock_conf) is
True
+
+
+def test_might_contain_dag_definition_materialized_for_custom_callable() ->
None:
+ # A custom callable only understands (file_path, zip_file); a definition
is materialized
+ # through its own as_file() so the callable gets a real, readable path.
+ seen: dict[str, object] = {}
+
+ def custom(file_path, zip_file=None):
+ seen["path"] = file_path
+ seen["zip_file"] = zip_file
+ seen["data"] = Path(file_path).read_bytes()
+ return b"airflow" in seen["data"]
+
+ mock_conf = mock.MagicMock()
+ mock_conf.getimport.return_value = custom
+
+ assert (
+ might_contain_dag(_FakeDefinition(b"from airflow import DAG"),
safe_mode=True, conf=mock_conf) is True
+ )
+ assert seen["data"] == b"from airflow import DAG"
+ assert seen["zip_file"] is None
+ assert not Path(str(seen["path"])).exists()
diff --git a/task-sdk/src/airflow/sdk/importers/__init__.py
b/task-sdk/src/airflow/sdk/importers/__init__.py
index f424fbf56ee..70c97ebb7e6 100644
--- a/task-sdk/src/airflow/sdk/importers/__init__.py
+++ b/task-sdk/src/airflow/sdk/importers/__init__.py
@@ -27,13 +27,14 @@ from airflow.sdk.importers.base import (
DagImportWarning,
DagSourceCode,
FileDagDefinition,
+ FilesystemDagDefinition,
find_file_dag_definitions,
get_file_suffix,
get_importer_registry,
reset_importer_registry,
)
from airflow.sdk.importers.python_importer import PythonDagImporter
-from airflow.sdk.importers.zip_importer import ZipFileDagDefinition,
ZipImporter
+from airflow.sdk.importers.zip_importer import ZipImporter,
ZipMemberDagDefinition
__all__ = [
"AbstractDagImporter",
@@ -44,9 +45,10 @@ __all__ = [
"DagImporterRegistry",
"DagSourceCode",
"FileDagDefinition",
+ "FilesystemDagDefinition",
"PythonDagImporter",
- "ZipFileDagDefinition",
"ZipImporter",
+ "ZipMemberDagDefinition",
"find_file_dag_definitions",
"get_file_suffix",
"get_importer_registry",
diff --git a/task-sdk/src/airflow/sdk/importers/base.py
b/task-sdk/src/airflow/sdk/importers/base.py
index 9c07892a526..83095bf5ad0 100644
--- a/task-sdk/src/airflow/sdk/importers/base.py
+++ b/task-sdk/src/airflow/sdk/importers/base.py
@@ -21,12 +21,12 @@ from __future__ import annotations
import contextlib
import functools
import logging
+import os
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from pathlib import Path
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Generic, TypeVar
-from airflow.sdk._shared.module_loading.dag_file import might_contain_dag
from airflow.sdk._shared.module_loading.file_discovery import
find_path_from_directory
from airflow.sdk.configuration import conf
from airflow.sdk.exceptions import AirflowConfigException
@@ -41,6 +41,8 @@ if TYPE_CHECKING:
log = logging.getLogger(__name__)
+DefT = TypeVar("DefT", bound="DagDefinition")
+
class DagDefinition(ABC):
"""Abstract base class for a DAG source definition."""
@@ -76,8 +78,34 @@ class DagDefinition(ABC):
"""Return string representation used by import error and warning
objects."""
-@dataclass
class FileDagDefinition(DagDefinition):
+ """
+ A DAG definition backed by a named, file-like resource.
+
+ Adds a filename ``suffix`` to the base content interface -- the trait that
makes a
+ source routable by file extension -- so extension-based importers can work
against
+ any file-like definition (a real file, an archive member, or a future
variant)
+ without recognizing concrete types.
+ """
+
+ @property
+ @abstractmethod
+ def suffix(self) -> str:
+ """Lowercased file extension including the leading dot (for example
``.py``)."""
+
+ def import_context(self) -> contextlib.AbstractContextManager[None]:
+ """
+ Prepare the environment needed to import this definition, then restore
it.
+
+ The default does nothing. A definition that needs setup -- e.g. an
archive member
+ placing its archive on ``sys.path`` for cross-member imports --
overrides this, so any
+ importer can load it without knowing where the definition came from.
+ """
+ return contextlib.nullcontext()
+
+
+@dataclass
+class FilesystemDagDefinition(FileDagDefinition):
"""A DAG definition backed by a file on the local filesystem."""
path: Path
@@ -98,6 +126,10 @@ class FileDagDefinition(DagDefinition):
except ValueError:
return str(self.path)
+ @property
+ def suffix(self) -> str:
+ return self.path.suffix.lower()
+
def read_bytes(self) -> bytes:
return self.path.read_bytes()
@@ -190,8 +222,14 @@ def _get_importer_extensions(importer:
AbstractDagImporter) -> list[str]:
return []
-class AbstractDagImporter(ABC):
- """Abstract base class for DAG importers."""
+class AbstractDagImporter(ABC, Generic[DefT]):
+ """
+ Abstract base for DAG importers, generic over the definition type it emits.
+
+ :meth:`.list_dag_definitions` yields definitions of :class:`DagDefinition`
+ subtypes, and those same objects are fed back to :meth:`import_definition`,
+ so a concrete importer only ever deals with its own definition type.
+ """
@abstractmethod
def can_handle(self, definition: DagDefinition | str | Path) -> bool:
@@ -203,16 +241,20 @@ class AbstractDagImporter(ABC):
bundle: BaseDagBundle,
*,
safe_mode: bool = True,
- ) -> Iterator[DagDefinition]:
- """List DAG definitions in a bundle that this importer can handle."""
+ ) -> Iterator[DefT | DagImportError]:
+ """
+ List DAG definitions in a bundle that this importer can handle
(identity-only discovery).
+
+ 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.
+ """
@abstractmethod
def import_definition(
self,
- definition: DagDefinition,
+ definition: DefT,
bundle: BaseDagBundle,
- *,
- safe_mode: bool = True,
) -> DagImportResult:
"""Import DAGs from a DAG definition."""
@@ -220,38 +262,68 @@ class AbstractDagImporter(ABC):
def get_source_code(self, definition: DagDefinition) -> DagSourceCode:
"""Retrieve the raw source code and its language identifier for the
specified DAG definition."""
+ def might_contain_dag(self, definition: DagDefinition, safe_mode: bool) ->
bool:
+ """
+ Cheap, optional pre-check for whether a discovered 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.
+ """
+ return True
+
def get_file_suffix(definition: DagDefinition | str | Path) -> str | None:
- """Extract lowercase file suffix from a definition, path, or filename."""
- path = (
- definition
- if isinstance(definition, (str, Path))
- else getattr(definition, "path", getattr(definition, "file_path",
None))
- )
- return Path(path).suffix.lower() if path else None
+ """Extract the lowercase file suffix from a file-like definition, path, or
filename."""
+ match definition:
+ case Path():
+ return definition.suffix.lower()
+ case str():
+ return os.path.splitext(definition)[-1].lower()
+ case FileDagDefinition():
+ return definition.suffix
+ return None
def find_file_dag_definitions(
bundle_path: Path,
supported_extensions: Iterable[str],
- safe_mode: bool = True,
-) -> Iterator[DagDefinition]:
- """Find file DAG definitions in a bundle matching given extensions and
respecting .airflowignore."""
+) -> Iterator[FilesystemDagDefinition]:
+ """
+ Discover file DAG definitions in a bundle by *identity* alone.
+
+ This walk decides purely from the file's name and path: the suffix,
+ ``.airflowignore``, and the Python source/bytecode pairing. It never reads
+ any file content. Deciding whether a discovered file actually contains a
DAG
+ belongs to :meth:`AbstractDagImporter.import_definition`, so importers
whose
+ validity can only be determined by attempting the import behave the same
way.
+ """
ignore_file_syntax = conf.get_mandatory_value("core",
"DAG_IGNORE_FILE_SYNTAX", fallback="glob")
supported_exts = _normalize_extensions(supported_extensions)
- for file_path in find_path_from_directory(bundle_path, ".airflowignore",
ignore_file_syntax):
- path = Path(file_path)
-
- if not path.is_file():
- continue
-
- if path.suffix.lower() not in supported_exts:
- continue
-
- if safe_mode and not might_contain_dag(str(path), safe_mode,
conf=conf):
+ def _iter_candidates() -> Iterator[Path]:
+ for file_path in find_path_from_directory(bundle_path,
".airflowignore", ignore_file_syntax):
+ path = Path(file_path)
+ if not path.is_file():
+ continue
+ if path.suffix.lower() not in supported_exts:
+ continue
+ if "__pycache__" in path.parts:
+ continue
+ yield path
+
+ candidates = list(_iter_candidates())
+
+ # A .pyc is discovered only when no matching .py survived the same
filtering
+ # (a genuinely sourceless module), so an ignored or unsupported .py sibling
+ # never suppresses it. Sources are keyed by (directory, stem) so the
.py/.pyc
+ # pairing matches case-insensitively on the extension, and candidates keep
the
+ # walk's order.
+ source_keys = {(path.parent, path.stem) for path in candidates if
path.suffix.lower() == ".py"}
+ for path in candidates:
+ if path.suffix.lower() == ".pyc" and (path.parent, path.stem) in
source_keys:
continue
- yield FileDagDefinition(path=path)
+ yield FilesystemDagDefinition(path=path)
@dataclass(frozen=True)
@@ -312,9 +384,9 @@ class DagImporterRegistry:
is logged. The built-in PythonDagImporter handles .py and ZipImporter
handles .zip files.
"""
- _extension_importers: dict[str, AbstractDagImporter]
+ _extension_importers: dict[str, AbstractDagImporter[Any]]
_extension_specs: dict[str, _ImporterSpec]
- _ordered_importers: list[AbstractDagImporter]
+ _ordered_importers: list[AbstractDagImporter[Any]]
def __init__(self, register_defaults: bool = True) -> None:
self._extension_importers = {}
@@ -344,7 +416,7 @@ class DagImporterRegistry:
return registry
- def register(self, importer: AbstractDagImporter, extensions: list[str] |
None = None) -> None:
+ def register(self, importer: AbstractDagImporter[Any], extensions:
list[str] | None = None) -> None:
"""
Register an importer.
@@ -377,7 +449,7 @@ class DagImporterRegistry:
self._warn_and_evict_extension(ext_lower, spec.classpath)
self._extension_specs[ext_lower] = spec
- def get_importer(self, definition: DagDefinition | str | Path) ->
AbstractDagImporter | None:
+ def get_importer(self, definition: DagDefinition | str | Path) ->
AbstractDagImporter[Any] | None:
"""Get the appropriate importer for a definition or file, or None if
unsupported."""
suffix = get_file_suffix(definition)
if suffix:
@@ -424,7 +496,7 @@ class DagImporterRegistry:
self.register(ZipImporter())
@staticmethod
- def _instantiate_spec(spec: _ImporterSpec) -> AbstractDagImporter:
+ def _instantiate_spec(spec: _ImporterSpec) -> AbstractDagImporter[Any]:
from airflow.sdk._shared.module_loading import import_string
try:
diff --git a/task-sdk/src/airflow/sdk/importers/python_importer.py
b/task-sdk/src/airflow/sdk/importers/python_importer.py
index 264cbd6b7a3..f6448c12645 100644
--- a/task-sdk/src/airflow/sdk/importers/python_importer.py
+++ b/task-sdk/src/airflow/sdk/importers/python_importer.py
@@ -14,17 +14,20 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
+
"""Python DAG importer - imports DAGs from Python files."""
from __future__ import annotations
import functools
-import importlib.machinery
+import importlib.abc
import importlib.util
import logging
-import os
+import marshal
+import signal
import sys
import traceback
+import types
import warnings
from pathlib import Path
from typing import TYPE_CHECKING, Any
@@ -42,6 +45,7 @@ from airflow.sdk.importers.base import (
DagImportResult,
DagImportWarning,
DagSourceCode,
+ FileDagDefinition,
_normalize_extensions,
find_file_dag_definitions,
get_file_suffix,
@@ -49,19 +53,80 @@ from airflow.sdk.importers.base import (
if TYPE_CHECKING:
from collections.abc import Iterator
- from types import ModuleType
from airflow.dag_processing.bundles.base import BaseDagBundle # noqa:
SDK002
log = logging.getLogger(__name__)
-class PythonDagImporter(AbstractDagImporter):
+class _DefinitionSourceLoader(importlib.abc.SourceLoader):
+ """
+ A SourceLoader that executes a DagDefinition straight from its bytes.
+
+ It needs no file on disk: :meth:`.get_data`` returns the definition's
+ source, and :meth:`get_filename` reports the definition's repr, so
+ ``__file__`` and tracebacks stay meaningful.
+
+ Bytecode caching is left disabled (the inherited ``path_stats`` raises
+ ``OSError``) since it's not particularly useful in dag processors.
+ """
+
+ def __init__(self, definition: DagDefinition) -> None:
+ self._definition = definition
+
+ def get_filename(self, fullname: str) -> str:
+ return repr(self._definition)
+
+ def get_data(self, path: str) -> bytes:
+ # The machinery only asks for get_filename(), i.e. the module's own
+ # source. Any other path is a sibling-resource request this
bytes-backed
+ # loader can't serve, so fail loud instead of returning the DAG source.
+ if path != self.get_filename(path):
+ raise FileNotFoundError(path)
+ return self._definition.read_bytes()
+
+
+class _DefinitionBytecodeLoader(importlib.abc.Loader):
+ """
+ Execute a definition's compiled bytecode (``.pyc``) straight from its
bytes.
+
+ The bytes-based counterpart of
:class:`importlib.machinery.SourcelessFileLoader`
+ (which is file-backed); reading through the definition keeps archive
members from
+ being extracted just to run them.
+ """
+
+ def __init__(self, definition: DagDefinition) -> None:
+ self._definition = definition
+
+ def get_filename(self, fullname: str) -> str:
+ return repr(self._definition)
+
+ def is_package(self, fullname: str) -> bool:
+ # importlib's file loaders report an ``__init__`` module as a package.
Without this a
+ # sourceless Dag package loses ``__path__`` and its relative imports
fail.
+ return Path(self.get_filename(fullname)).stem == "__init__"
+
+ def get_code(self, fullname: str) -> types.CodeType:
+ data = self._definition.read_bytes()
+ if len(data) < 16 or data[:4] != importlib.util.MAGIC_NUMBER:
+ raise ImportError(f"Incompatible or corrupt bytecode for
{self._definition!r}")
+ if int.from_bytes(data[4:8], "little") & ~0b11: # Reject undefined
PEP 552 flag bits.
+ raise ImportError(f"Invalid bytecode flags for
{self._definition!r}")
+ if not isinstance(code := marshal.loads(data[16:]), types.CodeType):
+ raise ImportError(f"Bytecode for {self._definition!r} does not
contain a code object")
+ return code
+
+ def exec_module(self, module: types.ModuleType) -> None:
+ exec(self.get_code(module.__name__), module.__dict__)
+
+
+class PythonDagImporter(AbstractDagImporter[FileDagDefinition]):
"""
- Importer for Python DAG files.
+ Importer for Python DAG sources.
- This is the default importer registered with the DagImporterRegistry. It
handles
- .py files containing Python DAGs.
+ This is the default importer registered with the DagImporterRegistry for
+ ``.py`` and ``.pyc`` files. The importer can import both from plain files
+ from the local filesystem, or members inside a zip archive.
"""
supported_extensions = [".py", ".pyc"]
@@ -80,39 +145,34 @@ class PythonDagImporter(AbstractDagImporter):
bundle: BaseDagBundle,
*,
safe_mode: bool = True,
- ) -> Iterator[DagDefinition]:
- """List Python DAG definitions in a bundle matching supported
extensions."""
- yield from find_file_dag_definitions(bundle.path,
self.supported_extensions, safe_mode=safe_mode)
+ ) -> Iterator[FileDagDefinition | DagImportError]:
+ """
+ List Python DAG files in a bundle matching supported extensions.
+
+ 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.
+ """
+ if not bundle.path.is_dir():
+ return
+ for definition in find_file_dag_definitions(bundle.path,
self.supported_extensions):
+ if self.might_contain_dag(definition, safe_mode):
+ yield definition
def import_definition(
self,
- definition: DagDefinition,
+ definition: FileDagDefinition,
bundle: BaseDagBundle,
- *,
- safe_mode: bool = True,
) -> DagImportResult:
- """
- Import DAGs from a Python DAG definition.
-
- :param definition: The definition to import from.
- :param bundle: The DAG bundle containing the definition.
- :param safe_mode: If True, skip files that don't appear to contain
DAGs.
- :return: DagImportResult with imported DAGs and any errors.
- """
+ """Import DAGs from a Python DAG definition."""
result = DagImportResult(definition=definition)
DagContext.autoregistered_dags.clear()
captured_warnings: list[warnings.WarningMessage] = []
try:
with warnings.catch_warnings(record=True) as captured_warnings:
- with definition.as_file() as local_path:
- filepath = os.fspath(local_path)
- modules = self._load_modules_from_file(
- filepath,
- safe_mode,
- result,
- bundle=bundle,
- )
+ modules = self._load_modules(definition, result, bundle=bundle)
except AirflowConfigException:
# Configuration errors (e.g., invalid timeout type) should
propagate
raise
@@ -140,12 +200,7 @@ class PythonDagImporter(AbstractDagImporter):
)
)
- self._process_modules(
- modules,
- result,
- bundle=bundle,
- )
-
+ self._process_modules(modules, result, bundle=bundle)
return result
def get_source_code(self, definition: DagDefinition) -> DagSourceCode:
@@ -155,77 +210,58 @@ class PythonDagImporter(AbstractDagImporter):
source_code="# Sourceless bytecode (.pyc) — source code not
available\n",
language="python",
)
- return DagSourceCode(
- source_code=definition.read_text(encoding="utf-8"),
- language="python",
- )
+ return
DagSourceCode(source_code=definition.read_text(encoding="utf-8"),
language="python")
- def might_contain_dag(self, file_path: str | Path, safe_mode: bool = True)
-> bool:
- """Check whether a file might contain Airflow DAGs according to safe
mode heuristics."""
- if not safe_mode:
- return True
- return might_contain_dag(str(file_path), safe_mode, conf=conf)
+ def might_contain_dag(self, definition: DagDefinition, safe_mode: bool) ->
bool:
+ """Sniff a Python DAG source's bytes for the Airflow/DAG markers."""
+ return might_contain_dag(definition, safe_mode=safe_mode, conf=conf)
- def _load_modules_from_file(
+ def _load_modules(
self,
- filepath: str,
- safe_mode: bool,
+ definition: FileDagDefinition,
result: DagImportResult,
bundle: BaseDagBundle,
- ) -> list[ModuleType]:
- definition = result.definition
-
- import signal
-
- def sigsegv_handler(signum, frame):
- msg = f"Received SIGSEGV signal while processing {filepath}."
+ ) -> list[types.ModuleType]:
+ def _handle_sigsegv(signum, frame):
+ msg = f"Received SIGSEGV signal while processing {definition!r}."
log.error(msg)
result.errors.append(
- DagImportError(
- source_reference=repr(definition),
- message=msg,
- error_type="segfault",
- )
+ DagImportError(source_reference=repr(definition), message=msg,
error_type="segfault")
)
try:
- signal.signal(signal.SIGSEGV, sigsegv_handler)
+ signal.signal(signal.SIGSEGV, _handle_sigsegv)
except (ValueError, AttributeError):
log.warning("SIGSEGV signal handler registration failed. Not in
the main thread")
- if not self.might_contain_dag(filepath, safe_mode):
- log.debug("File %s assumed to contain no DAGs. Skipping.",
filepath)
- if definition is not None:
- result.skipped_definitions.append(definition)
- return []
-
- log.debug("Importing %s (bundle: %s)", filepath, bundle.name)
- mod_name = get_unique_dag_module_name(filepath)
+ log.debug("Importing %r (bundle: %s)", definition, bundle.name)
+ mod_name = get_unique_dag_module_name(repr(definition))
if mod_name in sys.modules:
del sys.modules[mod_name]
DagContext.current_autoregister_module_name = mod_name
- def parse(mod_name: str, filepath: str) -> list[ModuleType]:
+ def parse() -> list[types.ModuleType]:
try:
- loader: importlib.machinery.SourceFileLoader |
importlib.machinery.SourcelessFileLoader
- if Path(filepath).suffix.lower() == ".pyc":
- loader =
importlib.machinery.SourcelessFileLoader(mod_name, filepath)
- else:
- loader = importlib.machinery.SourceFileLoader(mod_name,
filepath)
- spec = importlib.util.spec_from_loader(mod_name, loader)
- new_module = importlib.util.module_from_spec(spec) # type:
ignore[arg-type]
- sys.modules[spec.name] = new_module # type: ignore[union-attr]
- loader.exec_module(new_module)
- return [new_module]
+ with definition.import_context():
+ loader: importlib.abc.Loader
+ if get_file_suffix(definition) == ".pyc":
+ loader = _DefinitionBytecodeLoader(definition)
+ else:
+ loader = _DefinitionSourceLoader(definition)
+ spec = importlib.util.spec_from_loader(mod_name, loader)
+ new_module = importlib.util.module_from_spec(spec) #
type: ignore[arg-type]
+ sys.modules[mod_name] = new_module
+ loader.exec_module(new_module)
+ return [new_module]
except KeyboardInterrupt:
sys.modules.pop(mod_name, None)
raise
except BaseException as e:
sys.modules.pop(mod_name, None)
DagContext.autoregistered_dags.clear()
- log.exception("Failed to import: %s", filepath)
+ log.exception("Failed to import: %r", definition)
if self._dagbag_import_error_tracebacks:
stacktrace =
traceback.format_exc(limit=-self._dagbag_import_error_traceback_depth)
else:
@@ -244,7 +280,7 @@ class PythonDagImporter(AbstractDagImporter):
try:
from airflow import settings # noqa: SDK002
- dagbag_import_timeout =
settings.get_dagbag_import_timeout(filepath)
+ dagbag_import_timeout =
settings.get_dagbag_import_timeout(repr(definition))
except (ImportError, AttributeError):
dagbag_import_timeout = 30.0
@@ -254,16 +290,16 @@ class PythonDagImporter(AbstractDagImporter):
)
if dagbag_import_timeout <= 0:
- return parse(mod_name, filepath)
+ return parse()
timeout_msg = (
- f"DagBag import timeout for {filepath} after
{dagbag_import_timeout}s.\n"
+ f"DagBag import timeout for {definition!r} after
{dagbag_import_timeout}s.\n"
"Please take a look at these docs to improve your DAG import
time:\n"
"*
https://airflow.apache.org/docs/apache-airflow/stable/best-practices.html#top-level-python-code\n";
"*
https://airflow.apache.org/docs/apache-airflow/stable/best-practices.html#reducing-dag-complexity";
)
with timeout(seconds=dagbag_import_timeout, error_message=timeout_msg):
- return parse(mod_name, filepath)
+ return parse()
def _process_modules(
self,
diff --git a/task-sdk/src/airflow/sdk/importers/zip_importer.py
b/task-sdk/src/airflow/sdk/importers/zip_importer.py
index bc55eda0d94..d9de060d9cb 100644
--- a/task-sdk/src/airflow/sdk/importers/zip_importer.py
+++ b/task-sdk/src/airflow/sdk/importers/zip_importer.py
@@ -20,6 +20,7 @@ from __future__ import annotations
import contextlib
import logging
+import os
import sys
import tempfile
import threading
@@ -36,6 +37,7 @@ from airflow.sdk.importers.base import (
DagImportError,
DagImportResult,
DagSourceCode,
+ FileDagDefinition,
_get_importer_extensions,
_normalize_extensions,
_parse_importer_specs,
@@ -53,23 +55,8 @@ log = logging.getLogger(__name__)
_sys_path_lock = threading.RLock()
[email protected]
-def _temporary_sys_path(path: str) -> Generator[None, None, None]:
- """Safely prepend a path to sys.path with synchronization and
restoration."""
- with _sys_path_lock:
- already_present = path in sys.path
- if not already_present:
- sys.path.insert(0, path)
- try:
- yield
- finally:
- if not already_present:
- with contextlib.suppress(ValueError):
- sys.path.remove(path)
-
-
@dataclass
-class ZipFileDagDefinition(DagDefinition):
+class ZipMemberDagDefinition(FileDagDefinition):
"""A DAG definition backed by a file inside a ZIP archive."""
zip_path: Path
@@ -87,8 +74,26 @@ class ZipFileDagDefinition(DagDefinition):
def get_relative_loc(self, root: Path | None = None) -> str:
if root is not None:
with contextlib.suppress(ValueError):
- return f"{self.zip_path.relative_to(root)}:{self.file_path}"
- return f"{self.zip_path}:{self.file_path}"
+ return
str(self.zip_path.relative_to(root).joinpath(self.file_path))
+ return str(self.zip_path.joinpath(self.file_path))
+
+ @property
+ def suffix(self) -> str:
+ return os.path.splitext(self.file_path)[-1].lower()
+
+ @contextlib.contextmanager
+ def import_context(self) -> Generator[None, None, None]:
+ path = str(self.zip_path)
+ with _sys_path_lock:
+ already_present = path in sys.path
+ if not already_present:
+ sys.path.insert(0, path)
+ try:
+ yield
+ finally:
+ if not already_present:
+ with contextlib.suppress(ValueError):
+ sys.path.remove(path)
def read_bytes(self) -> bytes:
if self._content is None:
@@ -98,8 +103,7 @@ class ZipFileDagDefinition(DagDefinition):
@contextlib.contextmanager
def as_file(self) -> Generator[Path, None, None]:
- suffix = Path(self.file_path).suffix
- with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as f:
+ with tempfile.NamedTemporaryFile(suffix=self.suffix, delete=False) as
f:
f.write(self.read_bytes())
temp_path = Path(f.name)
try:
@@ -109,25 +113,25 @@ class ZipFileDagDefinition(DagDefinition):
temp_path.unlink()
def __repr__(self) -> str:
- return f"{self.zip_path}:{self.file_path}"
+ return str(self.zip_path.joinpath(self.file_path))
-class ZipImporter(AbstractDagImporter):
+class ZipImporter(AbstractDagImporter[ZipMemberDagDefinition]):
"""Composite importer responsible for routing archive members to internal
importers."""
supported_extensions = [".zip"]
def __init__(
self,
- internal_importers: dict[str, AbstractDagImporter | dict[str, Any]]
+ internal_importers: dict[str, AbstractDagImporter[Any] | dict[str,
Any]]
| list[dict[str, Any]]
| None = None,
extensions: list[str] | None = None,
) -> None:
if extensions is not None:
self.supported_extensions = _normalize_extensions(extensions)
- self._internal_extension_importers: dict[str, AbstractDagImporter] = {}
- self._ordered_internal_importers: list[AbstractDagImporter] = []
+ self._internal_extension_importers: dict[str,
AbstractDagImporter[Any]] = {}
+ self._ordered_internal_importers: list[AbstractDagImporter[Any]] = []
if internal_importers is None:
from airflow.sdk.importers.python_importer import PythonDagImporter
@@ -168,70 +172,98 @@ class ZipImporter(AbstractDagImporter):
bundle: BaseDagBundle,
*,
safe_mode: bool = True,
- ) -> Iterator[DagDefinition]:
- """List zip archive DAG definitions in a bundle matching supported
extensions."""
- yield from find_file_dag_definitions(bundle.path,
self.supported_extensions, safe_mode=safe_mode)
-
- def import_definition(
- self,
- definition: DagDefinition,
- bundle: BaseDagBundle,
- *,
- safe_mode: bool = True,
- ) -> DagImportResult:
+ ) -> Iterator[ZipMemberDagDefinition | DagImportError]:
"""
- Import DAGs from a ZIP archive by routing its members to internal
importers.
+ List importable members across the bundle's zip archives.
- The archive itself is placed on ``sys.path`` so Python imports between
- members resolve via ``zipimport``. A real file is materialized on
demand
- with :meth:`.as_file()` for internal importers.
+ Each member is yielded as a plain ZipMemberDagDefinition;
import_definition
+ re-resolves the internal importer from the member's extension. An
archive or a
+ member that cannot be read is reported as a DagImportError and the
remaining
+ members are still yielded.
"""
- result = DagImportResult(definition=definition)
-
- with definition.as_file() as local_zip_path:
+ for archive in find_file_dag_definitions(bundle.path,
self.supported_extensions):
try:
- with zipfile.ZipFile(local_zip_path) as z:
+ with zipfile.ZipFile(archive.path) as z:
member_names = z.namelist()
except Exception as e:
- result.errors.append(
- DagImportError(
-
source_reference=definition.get_relative_loc(bundle.path),
- message=f"Failed to read ZIP archive: {e}",
- error_type="zip_read_error",
- )
+ log.warning("Cannot read ZIP archive %s: %s", archive.path, e)
+ yield DagImportError(
+ source_reference=archive.get_relative_loc(bundle.path),
+ message=f"Failed to read ZIP archive: {e}",
+ error_type="zip_read_error",
)
- return result
-
- with _temporary_sys_path(str(local_zip_path)):
- for member_name in member_names:
- if member_name.endswith("/") or
member_name.startswith("__MACOSX/"):
- continue
- # ZipSlip defence: reject traversal or absolute member
names.
- member_path = Path(member_name)
- if member_path.is_absolute() or ".." in member_path.parts:
- log.warning(
- "Skipping zip member %r in %s: directory traversal
patterns detected",
- member_name,
- definition,
- )
- continue
-
- importer = self._get_internal_importer(member_name)
- if importer is None:
- continue
+ continue
- nested_def = ZipFileDagDefinition(zip_path=local_zip_path,
file_path=member_name)
- if not importer.can_handle(nested_def):
+ for member, importer in self._iter_supported_members(archive.path,
member_names):
+ try:
+ if safe_mode and not importer.might_contain_dag(member,
safe_mode):
continue
+ except Exception as e:
+ # One unreadable member must not end discovery for the
rest of the archive.
+ log.warning("Cannot read zip member %s of %s: %s",
member.file_path, archive.path, e)
+ yield DagImportError(
+ source_reference=member.get_relative_loc(bundle.path),
+ message=f"Failed to read ZIP member: {e}",
+ error_type="zip_read_error",
+ )
+ continue
+ yield member
+
+ def _iter_supported_members(
+ self, zip_path: Path, member_names: list[str]
+ ) -> Iterator[tuple[ZipMemberDagDefinition, AbstractDagImporter[Any]]]:
+ """Yield the archive's importable members, source preferred over
bytecode."""
+ candidates: list[tuple[str, AbstractDagImporter[Any]]] = []
+ for member_name in member_names:
+ if member_name.endswith("/") or
member_name.startswith("__MACOSX/"):
+ continue
+ # ZipSlip defence: reject traversal or absolute member names.
+ member_path = Path(member_name)
+ if member_path.is_absolute() or ".." in member_path.parts:
+ log.warning(
+ "Skipping zip member %r in %s: directory traversal
patterns detected",
+ member_name,
+ zip_path,
+ )
+ continue
+ if "__pycache__" in member_path.parts:
+ continue
+ if (importer := self._get_internal_importer(member_name)) is not
None:
+ candidates.append((member_name, importer))
+
+ # A .pyc is kept only when no supported .py sibling survived the same
filtering, so an
+ # unsupported source never hides a genuinely sourceless module.
Suffixes are compared
+ # lowercased, so `dag.PY` and `dag.pyc` still pair up.
+ source_stems = {os.path.splitext(name)[0] for name, _ in candidates if
get_file_suffix(name) == ".py"}
+ for member_name, importer in candidates:
+ if get_file_suffix(member_name) == ".pyc" and
os.path.splitext(member_name)[0] in source_stems:
+ continue
+ yield ZipMemberDagDefinition(zip_path=zip_path,
file_path=member_name), importer
- member_result = importer.import_definition(nested_def,
bundle, safe_mode=safe_mode)
- result.dags.extend(member_result.dags)
- result.errors.extend(member_result.errors)
- result.warnings.extend(member_result.warnings)
-
result.skipped_definitions.extend(member_result.skipped_definitions)
- result.dependencies.extend(member_result.dependencies)
+ def import_definition(
+ self,
+ definition: ZipMemberDagDefinition,
+ bundle: BaseDagBundle,
+ ) -> DagImportResult:
+ """
+ Import a single archive member.
- return result
+ The internal importer is resolved from the member's extension; the
member's
+ ``import_context`` places the archive on ``sys.path`` so imports
between members
+ resolve via ``zipimport``.
+ """
+ importer = self._get_internal_importer(definition.file_path)
+ if importer is None:
+ result = DagImportResult(definition=definition)
+ result.errors.append(
+ DagImportError(
+ source_reference=definition.get_relative_loc(bundle.path),
+ message=f"No internal importer registered for zip member
{definition.file_path}",
+ error_type="import",
+ )
+ )
+ return result
+ return importer.import_definition(definition, bundle)
def get_source_code(self, definition: DagDefinition) -> DagSourceCode:
"""
@@ -248,7 +280,9 @@ class ZipImporter(AbstractDagImporter):
raise ValueError(f"No internal importer to read source for
{definition!r}")
return importer.get_source_code(definition)
- def _register_internal(self, importer: AbstractDagImporter, extensions:
list[str] | None = None) -> None:
+ def _register_internal(
+ self, importer: AbstractDagImporter[Any], extensions: list[str] | None
= None
+ ) -> None:
if importer not in self._ordered_internal_importers:
self._ordered_internal_importers.append(importer)
exts = extensions if extensions is not None else
_get_importer_extensions(importer)
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 46bd283de8c..4e517f2a799 100644
--- a/task-sdk/tests/task_sdk/importers/test_python_importer.py
+++ b/task-sdk/tests/task_sdk/importers/test_python_importer.py
@@ -18,10 +18,16 @@
from __future__ import annotations
+import contextlib
+import importlib.util
import logging
+import marshal
import py_compile
import signal
import sys
+import tempfile
+import zipfile
+from pathlib import Path
from types import SimpleNamespace
from unittest import mock
@@ -30,8 +36,11 @@ import pytest
from airflow.sdk.exceptions import AirflowConfigException
from airflow.sdk.importers import (
FileDagDefinition,
+ FilesystemDagDefinition,
PythonDagImporter,
+ ZipMemberDagDefinition,
)
+from airflow.sdk.importers.python_importer import _DefinitionSourceLoader
@pytest.fixture
@@ -41,6 +50,41 @@ def mock_bundle(tmp_path):
return SimpleNamespace(name="test_bundle", path=bundle_dir)
+class _InMemoryDagDefinition(FileDagDefinition):
+ """A file-like definition backed purely by in-memory bytes (neither a file
nor a zip member)."""
+
+ def __init__(self, name: str, source: bytes) -> None:
+ self._name = name
+ self._source = source
+
+ @property
+ def suffix(self) -> str:
+ return Path(self._name).suffix.lower()
+
+ @property
+ def freshness_token(self) -> str:
+ return str(len(self._source))
+
+ def get_relative_loc(self, root: Path | None = None) -> str:
+ return self._name
+
+ def read_bytes(self) -> bytes:
+ return self._source
+
+ @contextlib.contextmanager
+ def as_file(self):
+ with tempfile.NamedTemporaryFile(suffix=self.suffix, delete=False) as
f:
+ f.write(self._source)
+ tmp = Path(f.name)
+ try:
+ yield tmp
+ finally:
+ tmp.unlink(missing_ok=True)
+
+ def __repr__(self) -> str:
+ return f"<memory:{self._name}>"
+
+
class TestPythonDagImporter:
"""Test the PythonDagImporter implementation."""
@@ -49,7 +93,7 @@ class TestPythonDagImporter:
dag_file.write_text("from airflow.sdk import DAG\ndag =
DAG('test_dag_1')\n")
importer = PythonDagImporter()
- definition = FileDagDefinition(path=dag_file)
+ definition = FilesystemDagDefinition(path=dag_file)
result = importer.import_definition(definition, bundle=mock_bundle)
assert len(result.dags) == 1
@@ -63,40 +107,117 @@ class TestPythonDagImporter:
dag_file.write_text("from airflow.sdk import DAG\ndef broken(\n")
importer = PythonDagImporter()
- result = importer.import_definition(FileDagDefinition(path=dag_file),
bundle=mock_bundle)
+ result =
importer.import_definition(FilesystemDagDefinition(path=dag_file),
bundle=mock_bundle)
assert len(result.errors) == 1
assert result.errors[0].error_type == "import"
assert not any("bad_dag" in m for m in sys.modules)
- def test_skip_non_dag_file_in_safe_mode(self, mock_bundle):
+ def test_import_non_dag_file_yields_no_dags(self, mock_bundle):
+ # Import does not re-sniff -- discovery already filters non-DAG files.
Importing one
+ # directly just yields no DAGs (and no error); it is not reported as
skipped.
helper_file = mock_bundle.path / "helper.py"
helper_file.write_text("def util(): return 42\n")
importer = PythonDagImporter()
- definition = FileDagDefinition(path=helper_file)
- result = importer.import_definition(definition, bundle=mock_bundle,
safe_mode=True)
+ definition = FilesystemDagDefinition(path=helper_file)
+ result = importer.import_definition(definition, bundle=mock_bundle)
+
+ assert result.dags == []
+ assert result.errors == []
+ assert result.skipped_definitions == []
+
+ def test_import_non_dag_zip_member_yields_no_dags(self, mock_bundle):
+ # Same for a zip member imported directly: no sniff at import, so no
DAGs and no skip.
+ zip_path = mock_bundle.path / "helpers.zip"
+ with zipfile.ZipFile(zip_path, "w") as z:
+ z.writestr("helper.py", "def util():\n return 42\n")
+
+ importer = PythonDagImporter()
+ definition = ZipMemberDagDefinition(zip_path=zip_path,
file_path="helper.py")
+ result = importer.import_definition(definition, bundle=mock_bundle)
+
+ assert result.dags == []
+ assert result.errors == []
+ assert result.skipped_definitions == []
+
+ def test_import_corrupt_pyc_captured_as_error(self, mock_bundle):
+ # A .pyc whose header is not valid CPython bytecode must surface as an
import
+ # error (from _DefinitionBytecodeLoader's magic check), not crash the
importer.
+ bad_pyc = mock_bundle.path / "broken.pyc"
+ bad_pyc.write_bytes(b"this is not valid python bytecode at all!!")
+
+ importer = PythonDagImporter()
+ definition = FilesystemDagDefinition(path=bad_pyc)
+ result = importer.import_definition(definition, bundle=mock_bundle)
assert len(result.dags) == 0
- assert len(result.errors) == 0
- assert result.skipped_definitions == [definition]
+ assert len(result.errors) == 1
+ assert result.errors[0].error_type == "import"
- @pytest.mark.parametrize(
- ("safe_mode", "expected_files"),
- [
- (True, {"sample_dag.py"}),
- (False, {"sample_dag.py", "helper.py"}),
- ],
- )
- def test_list_dag_definitions(self, mock_bundle, safe_mode,
expected_files):
+ def test_source_loader_get_data_rejects_foreign_path(self):
+ # get_data serves the module's own source, but a sibling-resource
request must fail
+ # loud rather than return the DAG source for it.
+ class _Def:
+ def read_bytes(self) -> bytes:
+ return b"SOURCE = 1"
+
+ def __repr__(self) -> str:
+ return "mydef"
+
+ loader = _DefinitionSourceLoader(_Def())
+ assert loader.get_data(loader.get_filename("mydef")) == b"SOURCE = 1"
+ with pytest.raises(FileNotFoundError):
+ loader.get_data("some/other/config.json")
+
+ def test_import_pyc_with_non_code_payload_captured_as_error(self,
mock_bundle):
+ # A .pyc with a valid magic header but a payload that unmarshals to a
non-code
+ # object must be rejected, not exec()'d as source.
+ header = importlib.util.MAGIC_NUMBER + b"\x00" * 12
+ bad_pyc = mock_bundle.path / "not_code.pyc"
+ bad_pyc.write_bytes(header + marshal.dumps("i am a string, not a code
object"))
+
+ importer = PythonDagImporter()
+ definition = FilesystemDagDefinition(path=bad_pyc)
+ result = importer.import_definition(definition, bundle=mock_bundle)
+
+ assert len(result.dags) == 0
+ assert len(result.errors) == 1
+ assert result.errors[0].error_type == "import"
+
+ def test_imports_arbitrary_file_like_definition(self, mock_bundle):
+ # The importer must handle any FileDagDefinition through the
interface, with no
+ # knowledge of File/Zip concrete types -- so a third, in-memory
backing works too.
+ source = b"from airflow.sdk import DAG\ndag = DAG('in_memory_dag')\n"
+ definition = _InMemoryDagDefinition("in_memory_dag.py", source)
+
+ result = PythonDagImporter().import_definition(definition,
bundle=mock_bundle)
+
+ assert [d.dag_id for d in result.dags] == ["in_memory_dag"]
+ assert result.errors == []
+
+ def test_list_dag_definitions(self, mock_bundle):
+ # Discovery applies the lightweight sniff, so a .py with no DAG
markers (helper.py) is
+ # filtered out and never becomes a definition; only the real DAG file
is returned.
dag_file = mock_bundle.path / "sample_dag.py"
dag_file.write_text("from airflow.sdk import DAG\ndag =
DAG('test_dag_1')\n")
(mock_bundle.path / "helper.py").write_text("def helper():\n return
42\n")
(mock_bundle.path / "notes.txt").write_text("hello")
importer = PythonDagImporter()
- defs = list(importer.list_dag_definitions(mock_bundle,
safe_mode=safe_mode))
- assert {d.path.name for d in defs} == expected_files
+ defs = list(importer.list_dag_definitions(mock_bundle))
+ assert {d.path.name for d in defs} == {"sample_dag.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
+ (mock_bundle.path / "bar.pyc").write_bytes(b"airflow dag") #
sourceless (has markers) -> kept
+ cache = mock_bundle.path / "__pycache__"
+ cache.mkdir()
+ (cache / "foo.cpython-311.pyc").write_bytes(b"compiled") # cache ->
skipped
+
+ defs = list(PythonDagImporter().list_dag_definitions(mock_bundle))
+ assert sorted(d.path.name for d in defs) == ["bar.pyc", "foo.py"]
@pytest.mark.parametrize(
("filename", "is_bytecode", "expected_content"),
@@ -112,7 +233,7 @@ class TestPythonDagImporter:
else:
dag_file.write_text(expected_content)
- src =
PythonDagImporter().get_source_code(FileDagDefinition(path=dag_file))
+ src =
PythonDagImporter().get_source_code(FilesystemDagDefinition(path=dag_file))
assert src.language == "python"
assert src.source_code == expected_content
@@ -123,17 +244,36 @@ class TestPythonDagImporter:
py_compile.compile(str(source_file), cfile=str(pyc_file))
importer = PythonDagImporter()
- result = importer.import_definition(FileDagDefinition(path=pyc_file),
bundle=mock_bundle)
+ result =
importer.import_definition(FilesystemDagDefinition(path=pyc_file),
bundle=mock_bundle)
assert len(result.dags) == 1
assert result.dags[0].dag_id == "compiled_dag"
assert len(result.errors) == 0
+ def test_import_sourceless_package_resolves_relative_import(self,
mock_bundle, tmp_path):
+ # An __init__.pyc must load as a package, otherwise its relative
imports fail with
+ # "attempted relative import with no known parent package".
+ pkg = mock_bundle.path / "sourceless_pkg"
+ pkg.mkdir()
+ (tmp_path / "helper.py").write_text("ID = 5\n")
+ (tmp_path / "__init__.py").write_text(
+ "from airflow.sdk import DAG\nfrom .helper import ID\ndag =
DAG(f'sourceless_pkg_{ID}')\n"
+ )
+ for name in ("helper", "__init__"):
+ py_compile.compile(str(tmp_path / f"{name}.py"), cfile=str(pkg /
f"{name}.pyc"))
+
+ result = PythonDagImporter().import_definition(
+ FilesystemDagDefinition(path=pkg / "__init__.pyc"),
bundle=mock_bundle
+ )
+
+ assert result.errors == []
+ assert [d.dag_id for d in result.dags] == ["sourceless_pkg_5"]
+
def test_file_dag_definition_freshness_token(self, tmp_path):
dag_file = tmp_path / "fresh_dag.py"
dag_file.write_text("from airflow.sdk import DAG\n")
stat = dag_file.stat()
- assert FileDagDefinition(path=dag_file).freshness_token ==
f"{stat.st_mtime_ns}-{stat.st_size}"
+ assert FilesystemDagDefinition(path=dag_file).freshness_token ==
f"{stat.st_mtime_ns}-{stat.st_size}"
def test_python_importer_custom_extensions(self, mock_bundle):
importer = PythonDagImporter(extensions=[".custom_py"])
@@ -165,7 +305,7 @@ class TestPythonDagImporter:
dag_file.write_text("from airflow.sdk import DAG\ndef broken(\n")
importer = PythonDagImporter()
- result = importer.import_definition(FileDagDefinition(path=dag_file),
bundle=mock_bundle)
+ result =
importer.import_definition(FilesystemDagDefinition(path=dag_file),
bundle=mock_bundle)
assert len(result.errors) == 1
assert (result.errors[0].stacktrace is not None) == expect_traceback
@@ -183,16 +323,15 @@ class TestPythonDagImporter:
),
):
importer.import_definition(
- FileDagDefinition(path=mock_bundle.path / "dag.py"),
+ FilesystemDagDefinition(path=mock_bundle.path / "dag.py"),
bundle=mock_bundle,
- safe_mode=False,
)
- @mock.patch.object(PythonDagImporter, "_load_modules_from_file",
side_effect=TypeError("unexpected None"))
+ @mock.patch.object(PythonDagImporter, "_load_modules",
side_effect=TypeError("unexpected None"))
def test_unexpected_type_error_captured_in_result_errors(self, mock_load,
mock_bundle):
importer = PythonDagImporter()
result = importer.import_definition(
- FileDagDefinition(path=mock_bundle.path / "dag.py"),
+ FilesystemDagDefinition(path=mock_bundle.path / "dag.py"),
bundle=mock_bundle,
)
@@ -213,7 +352,7 @@ class TestPythonDagImporter:
registered_handler = handler
with mock.patch("signal.signal", side_effect=mock_signal_func):
- result =
importer.import_definition(FileDagDefinition(path=dag_file), bundle=mock_bundle)
+ result =
importer.import_definition(FilesystemDagDefinition(path=dag_file),
bundle=mock_bundle)
assert callable(registered_handler)
registered_handler(signal.SIGSEGV, None)
@@ -230,7 +369,7 @@ class TestPythonDagImporter:
mock.patch("signal.signal", side_effect=ValueError("signal only
works in main thread")),
caplog.at_level(logging.WARNING),
):
- result =
importer.import_definition(FileDagDefinition(path=dag_file), bundle=mock_bundle)
+ result =
importer.import_definition(FilesystemDagDefinition(path=dag_file),
bundle=mock_bundle)
assert len(result.dags) == 1
assert "SIGSEGV signal handler registration failed. Not in the main
thread" in caplog.text
diff --git a/task-sdk/tests/task_sdk/importers/test_registry.py
b/task-sdk/tests/task_sdk/importers/test_registry.py
index 44dc5929aae..20a6b78a337 100644
--- a/task-sdk/tests/task_sdk/importers/test_registry.py
+++ b/task-sdk/tests/task_sdk/importers/test_registry.py
@@ -32,10 +32,10 @@ from airflow.sdk.importers import (
DagImportError,
DagImportResult,
DagSourceCode,
- FileDagDefinition,
+ FilesystemDagDefinition,
PythonDagImporter,
- ZipFileDagDefinition,
ZipImporter,
+ ZipMemberDagDefinition,
find_file_dag_definitions,
get_file_suffix,
get_importer_registry,
@@ -167,10 +167,27 @@ class TestDagImporterRegistry:
assert registry.get_importer("custom://dags/sample") is importer
assert not registry.can_handle("other://dags/sample")
- def
test_abstract_dag_importer_has_no_extension_attributes_or_methods(self):
- """AbstractDagImporter must not define file extension attributes or
methods."""
+ def test_abstract_dag_importer_has_no_extension_attributes(self):
+ """AbstractDagImporter must not define file-extension attributes."""
assert not hasattr(AbstractDagImporter, "supported_extensions")
- assert not hasattr(AbstractDagImporter, "might_contain_dag")
+
+ def test_might_contain_dag_defaults_to_true(self):
+ """An importer with no cheap heuristic inherits the base default and
keeps everything."""
+
+ class _Bare(AbstractDagImporter):
+ def can_handle(self, definition):
+ return True
+
+ def list_dag_definitions(self, bundle, *, safe_mode=True):
+ return iter(())
+
+ def import_definition(self, definition, bundle):
+ raise NotImplementedError
+
+ def get_source_code(self, definition):
+ raise NotImplementedError
+
+ assert _Bare().might_contain_dag(object(), safe_mode=True) is True
def test_custom_importer_file_pattern_can_handle(self):
"""Registry resolves file definitions via can_handle when importer has
no supported_extensions."""
@@ -380,8 +397,8 @@ class TestDagImporterRegistry:
("path/to/FOO.PY", ".py"),
(Path("archive.ZIP"), ".zip"),
("no_extension", ""),
- (FileDagDefinition(path=Path("my_dag.py")), ".py"),
- (ZipFileDagDefinition(zip_path=Path("a.zip"),
file_path="nested/workflow.py"), ".py"),
+ (FilesystemDagDefinition(path=Path("my_dag.py")), ".py"),
+ (ZipMemberDagDefinition(zip_path=Path("a.zip"),
file_path="nested/workflow.py"), ".py"),
(None, None),
],
)
@@ -416,17 +433,31 @@ class TestDagImporterRegistry:
def test_dag_import_error_format_message(self, error, expected):
assert error.format_message() == expected
- @pytest.mark.parametrize(
- ("safe_mode", "expected_files"),
- [
- (True, {"workflow.py"}),
- (False, {"workflow.py", "script.py"}),
- ],
- )
- def test_find_file_dag_definitions_safe_mode(self, tmp_path, safe_mode,
expected_files):
+ def test_find_file_dag_definitions_is_identity_only(self, tmp_path):
+ # Discovery matches on extension alone and never reads file contents,
so a .py with no
+ # DAG (script.py) is returned just like a real DAG file; only the .csv
is filtered out
+ # (wrong extension). Content-based filtering happens later at
import_definition.
(tmp_path / "workflow.py").write_text("from airflow.sdk import DAG\n")
(tmp_path / "script.py").write_text("print('hello')\n")
(tmp_path / "data.csv").write_text("a,b,c\n")
- definitions = list(find_file_dag_definitions(tmp_path, [".py"],
safe_mode=safe_mode))
- assert {d.path.name for d in definitions} == expected_files
+ definitions = list(find_file_dag_definitions(tmp_path, [".py"]))
+ assert {d.path.name for d in definitions} == {"workflow.py",
"script.py"}
+
+ def test_find_file_dag_definitions_pyc_dedup_over_discovered_set(self,
tmp_path):
+ # The .py-over-.pyc preference is decided over what discovery yields,
not a raw
+ # filesystem stat. Here the .py is not a supported extension, so it is
never a
+ # discovery candidate and must not suppress the sourceless .pyc beside
it.
+ (tmp_path / "sourceless.py").write_text("from airflow.sdk import
DAG\n")
+ (tmp_path / "sourceless.pyc").write_bytes(b"compiled")
+
+ definitions = list(find_file_dag_definitions(tmp_path, [".pyc"]))
+ assert {d.path.name for d in definitions} == {"sourceless.pyc"}
+
+ def
test_find_file_dag_definitions_pyc_dedup_case_insensitive_extension(self,
tmp_path):
+ # The source file's extension may be any case; the .pyc beside it is
still deduped.
+ (tmp_path / "workflow.PY").write_text("from airflow.sdk import DAG\n")
+ (tmp_path / "workflow.pyc").write_bytes(b"compiled")
+
+ definitions = list(find_file_dag_definitions(tmp_path, [".py",
".pyc"]))
+ assert {d.path.name for d in definitions} == {"workflow.PY"}
diff --git a/task-sdk/tests/task_sdk/importers/test_zip_importer.py
b/task-sdk/tests/task_sdk/importers/test_zip_importer.py
index 023db51cb16..baa5516d113 100644
--- a/task-sdk/tests/task_sdk/importers/test_zip_importer.py
+++ b/task-sdk/tests/task_sdk/importers/test_zip_importer.py
@@ -18,10 +18,12 @@
from __future__ import annotations
+import os
import py_compile
import zipfile
from pathlib import Path
from types import SimpleNamespace
+from typing import TYPE_CHECKING
import pytest
@@ -29,13 +31,17 @@ from airflow.sdk.exceptions import AirflowConfigException
from airflow.sdk.importers import (
AbstractDagImporter,
DagDefinition,
+ DagImportError,
DagImportResult,
DagSourceCode,
- FileDagDefinition,
- PythonDagImporter,
- ZipFileDagDefinition,
+ FilesystemDagDefinition,
ZipImporter,
+ ZipMemberDagDefinition,
)
+from airflow.sdk.importers.python_importer import PythonDagImporter
+
+if TYPE_CHECKING:
+ from airflow.sdk import DAG
class CustomInternalNonExtensionImporter(AbstractDagImporter):
@@ -57,6 +63,25 @@ class
CustomInternalNonExtensionImporter(AbstractDagImporter):
return DagSourceCode(source_code="", language="text")
+def _import_all(
+ importer: AbstractDagImporter[DagDefinition],
+ bundle,
+) -> tuple[list[DAG], list[DagImportError]]:
+ """Enumerate an importer's definitions and import each, aggregating
dags/errors."""
+ dags, errors = [], []
+ for item in importer.list_dag_definitions(bundle):
+ match item:
+ case DagImportError():
+ errors.append(item)
+ case DagDefinition():
+ result = importer.import_definition(item, bundle=bundle)
+ dags.extend(result.dags)
+ errors.extend(result.errors)
+ case _:
+ raise ValueError(f"unrecognized dag definition {item!r}")
+ return dags, errors
+
+
class TestZipImporter:
"""Test the ZipImporter composite implementation."""
@@ -64,22 +89,69 @@ class TestZipImporter:
def mock_bundle(self, tmp_path):
return SimpleNamespace(name="test_bundle", path=tmp_path)
- @pytest.mark.parametrize(
- ("safe_mode", "expected_count"),
- [
- (True, 1),
- (False, 2),
- ],
- )
- def test_list_dag_definitions(self, mock_bundle, safe_mode,
expected_count):
+ def test_list_dag_definitions(self, mock_bundle):
zip_path = mock_bundle.path / "sample.zip"
with zipfile.ZipFile(zip_path, "w") as z:
z.writestr("dag.py", "from airflow.sdk import DAG\n")
(mock_bundle.path / "corrupt.zip").write_bytes(b"not a valid zip and
no dag markers")
- definitions = list(ZipImporter().list_dag_definitions(mock_bundle,
safe_mode=safe_mode))
- assert len(definitions) == expected_count
- assert any(d.path == zip_path for d in definitions)
+ # The valid archive yields its member; the unreadable one is surfaced
in-band as a
+ # DagImportError rather than dropped.
+ items = list(ZipImporter().list_dag_definitions(mock_bundle))
+ members = [i for i in items if not isinstance(i, DagImportError)]
+ errors = [i for i in items if isinstance(i, DagImportError)]
+ assert [(d.zip_path, d.file_path) for d in members] == [(zip_path,
"dag.py")]
+ assert [e.error_type for e in errors] == ["zip_read_error"]
+
+ def test_list_prefers_source_over_pyc_and_skips_pycache(self, mock_bundle):
+ zip_path = mock_bundle.path / "compiled.zip"
+ with zipfile.ZipFile(zip_path, "w") as z:
+ z.writestr("foo.py", "from airflow.sdk import DAG\n")
+ z.writestr("foo.pyc", b"compiled") # side-by-side -> skipped
+ z.writestr("bar.pyc", b"airflow dag") # sourceless (has markers)
-> kept
+ z.writestr("__pycache__/foo.cpython-311.pyc", b"compiled") #
cache -> skipped
+
+ definitions = list(ZipImporter().list_dag_definitions(mock_bundle))
+ assert sorted(d.file_path for d in definitions) == ["bar.pyc",
"foo.py"]
+
+ def test_list_dedups_over_supported_candidates_only(self, mock_bundle):
+ # An unsupported .py must not suppress the supported .pyc beside it.
+ zip_path = mock_bundle.path / "bytecode_only.zip"
+ with zipfile.ZipFile(zip_path, "w") as z:
+ z.writestr("dag.py", "from airflow.sdk import DAG\n")
+ z.writestr("dag.pyc", b"airflow dag")
+
+ importer = ZipImporter(internal_importers={".pyc":
PythonDagImporter()})
+ definitions = list(importer.list_dag_definitions(mock_bundle))
+ assert [d.file_path for d in definitions] == ["dag.pyc"]
+
+ def test_list_pairs_source_and_bytecode_case_insensitively(self,
mock_bundle):
+ # `dag.PY` and `dag.pyc` are the same module, so only the source
survives.
+ zip_path = mock_bundle.path / "mixed_case.zip"
+ with zipfile.ZipFile(zip_path, "w") as z:
+ z.writestr("dag.PY", "from airflow.sdk import DAG\n")
+ z.writestr("dag.pyc", b"airflow dag")
+
+ definitions = list(ZipImporter().list_dag_definitions(mock_bundle))
+ assert [d.file_path for d in definitions] == ["dag.PY"]
+
+ def test_list_continues_past_unreadable_member(self, mock_bundle):
+ # A single bad member is reported and the rest of the archive is still
discovered.
+ zip_path = mock_bundle.path / "bad_member.zip"
+ payload = b"# airflow dag BBBB\n"
+ with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_STORED) as z:
+ z.writestr("a_dag.py", "from airflow.sdk import DAG\n")
+ z.writestr("b_bad.py", payload)
+ z.writestr("c_dag.py", "from airflow.sdk import DAG\n")
+ # Rewrite the stored bytes in place so the recorded CRC no longer
matches.
+ zip_path.write_bytes(zip_path.read_bytes().replace(payload, b"#
airflow dag CCCC\n"))
+
+ items = list(ZipImporter().list_dag_definitions(mock_bundle))
+ members = [i for i in items if not isinstance(i, DagImportError)]
+ errors = [i for i in items if isinstance(i, DagImportError)]
+ assert [d.file_path for d in members] == ["a_dag.py", "c_dag.py"]
+ assert [e.error_type for e in errors] == ["zip_read_error"]
+ assert errors[0].source_reference == os.path.join("bad_member.zip",
"b_bad.py")
def test_import_zip_archive_with_dags(self, mock_bundle):
zip_path = mock_bundle.path / "sample_dags.zip"
@@ -87,11 +159,11 @@ class TestZipImporter:
z.writestr("dag_a.py", "from airflow.sdk import DAG\ndag =
DAG('zip_dag_a')\n")
z.writestr("helper.py", "def util(): return 1\n")
- result =
ZipImporter().import_definition(FileDagDefinition(path=zip_path),
bundle=mock_bundle)
- assert len(result.dags) == 1
- assert result.dags[0].dag_id == "zip_dag_a"
- assert result.dags[0].bundle_name == "test_bundle"
- assert len(result.errors) == 0
+ dags, errors = _import_all(ZipImporter(), mock_bundle)
+ assert len(dags) == 1
+ assert dags[0].dag_id == "zip_dag_a"
+ assert dags[0].bundle_name == "test_bundle"
+ assert len(errors) == 0
def test_import_zip_archive_with_pyc_dag(self, mock_bundle, tmp_path):
source_file = tmp_path / "compiled_dag.py"
@@ -104,12 +176,14 @@ class TestZipImporter:
z.write(pyc_file, arcname="compiled_dag.pyc")
importer = ZipImporter()
- result = importer.import_definition(FileDagDefinition(path=zip_path),
bundle=mock_bundle)
- assert len(result.dags) == 1
- assert result.dags[0].dag_id == "zip_pyc_dag"
- assert len(result.errors) == 0
+ dags, errors = _import_all(importer, mock_bundle)
+ assert len(dags) == 1
+ assert dags[0].dag_id == "zip_pyc_dag"
+ assert len(errors) == 0
- src = importer.get_source_code(ZipFileDagDefinition(zip_path=zip_path,
file_path="compiled_dag.pyc"))
+ src = importer.get_source_code(
+ ZipMemberDagDefinition(zip_path=zip_path,
file_path="compiled_dag.pyc")
+ )
assert src.language == "python"
assert "Sourceless bytecode" in src.source_code
@@ -121,18 +195,45 @@ class TestZipImporter:
z.writestr("../evil_dag.py", "from airflow.sdk import DAG\ndag =
DAG('evil_dag')\n")
z.writestr("valid_dag.py", "from airflow.sdk import DAG\ndag =
DAG('valid_dag')\n")
- result =
ZipImporter().import_definition(FileDagDefinition(path=zip_path),
bundle=mock_bundle)
- assert len(result.dags) == 1
- assert result.dags[0].dag_id == "valid_dag"
+ dags, _ = _import_all(ZipImporter(), mock_bundle)
+ assert len(dags) == 1
+ assert dags[0].dag_id == "valid_dag"
assert not (mock_bundle.path.parent / "evil_dag.py").exists()
def test_corrupted_zip_file(self, mock_bundle):
bad_zip = mock_bundle.path / "corrupted.zip"
bad_zip.write_bytes(b"not a real zip")
- result =
ZipImporter().import_definition(FileDagDefinition(path=bad_zip),
bundle=mock_bundle)
- assert len(result.errors) == 1
- assert result.errors[0].error_type == "zip_read_error"
+ # An unreadable archive is surfaced in-band as a DagImportError, not
silently dropped.
+ items = list(ZipImporter().list_dag_definitions(mock_bundle))
+ assert len(items) == 1
+ assert isinstance(items[0], DagImportError)
+ assert items[0].error_type == "zip_read_error"
+
+ def test_zip_member_cross_import_via_python_importer(self, mock_bundle):
+ # A member imported directly through PythonDagImporter (as the
registry routes it by
+ # suffix, not via ZipImporter) still resolves sibling-member imports,
because the
+ # definition's import_context puts its archive on sys.path.
+ zip_path = mock_bundle.path / "cross.zip"
+ with zipfile.ZipFile(zip_path, "w") as z:
+ z.writestr("helper_mod.py", "VALUE = 7\n")
+ z.writestr(
+ "main_dag.py",
+ "from airflow.sdk import DAG\nimport helper_mod\ndag =
DAG(f'cross_{helper_mod.VALUE}')\n",
+ )
+
+ member = ZipMemberDagDefinition(zip_path=zip_path,
file_path="main_dag.py")
+ result = PythonDagImporter().import_definition(member,
bundle=mock_bundle)
+
+ assert result.errors == []
+ assert [d.dag_id for d in result.dags] == ["cross_7"]
+
+ def test_zip_member_fileloc_uses_os_sep(self, tmp_path):
+ # The fileloc/relative loc join with os.sep (archive.zip/member.py) --
the form
+ # airflow-core's ZIP_REGEX and open_maybe_zipped understand -- not a
colon.
+ member = ZipMemberDagDefinition(zip_path=tmp_path / "a.zip",
file_path="sub/dag.py")
+ assert repr(member) == os.path.join(str(tmp_path / "a.zip"),
"sub/dag.py")
+ assert member.get_relative_loc(tmp_path) == os.path.join("a.zip",
"sub/dag.py")
def test_get_source_code_reads_member_not_archive(self, tmp_path):
zip_path = tmp_path / "source_dags.zip"
@@ -145,20 +246,22 @@ class TestZipImporter:
# A zip is a directory of DAG files: each member is its own source
unit,
# rendered through its file-type importer (same single-member
semantics as
# the legacy code view's open_maybe_zipped).
- src_member =
importer.get_source_code(ZipFileDagDefinition(zip_path=zip_path,
file_path="my_dag.py"))
+ src_member = importer.get_source_code(
+ ZipMemberDagDefinition(zip_path=zip_path, file_path="my_dag.py")
+ )
assert src_member.language == "python"
assert src_member.source_code == dag_content
# The archive as a whole has no source, the same way a directory does
not.
with pytest.raises(ValueError, match="No internal importer"):
- importer.get_source_code(FileDagDefinition(path=zip_path))
+ importer.get_source_code(FilesystemDagDefinition(path=zip_path))
def test_zip_dag_definition_freshness_token(self, tmp_path):
zip_path = tmp_path / "fresh_bundle.zip"
with zipfile.ZipFile(zip_path, "w") as z:
z.writestr("dag.py", "from airflow.sdk import DAG\n")
- member_def = ZipFileDagDefinition(zip_path=zip_path,
file_path="dag.py")
+ member_def = ZipMemberDagDefinition(zip_path=zip_path,
file_path="dag.py")
stat = zip_path.stat()
assert member_def.freshness_token ==
f"{stat.st_mtime_ns}-{stat.st_size}-dag.py"
@@ -178,9 +281,9 @@ class TestZipImporter:
with zipfile.ZipFile(zip_path, "w") as z:
z.writestr("dag.custom_py", "from airflow.sdk import DAG\ndag =
DAG('custom_zip_dag')\n")
- res = importer.import_definition(FileDagDefinition(path=zip_path),
bundle=mock_bundle)
- assert len(res.dags) == 1
- assert res.dags[0].dag_id == "custom_zip_dag"
+ dags, _ = _import_all(importer, mock_bundle)
+ assert len(dags) == 1
+ assert dags[0].dag_id == "custom_zip_dag"
def test_zip_importer_internal_importers_from_dict(self):
importer = ZipImporter(
@@ -232,6 +335,6 @@ class TestZipImporter:
with zipfile.ZipFile(zip_path, "w") as z:
z.writestr("my_workflow_file", "steps:\n - run: echo hello\n")
- res = importer.import_definition(FileDagDefinition(path=zip_path),
bundle=mock_bundle)
- assert len(res.dags) == 1
- assert res.dags[0].dag_id == "workflow_dag"
+ dags, _ = _import_all(importer, mock_bundle)
+ assert len(dags) == 1
+ assert dags[0].dag_id == "workflow_dag"
