kaxil commented on code in PR #73118:
URL: https://github.com/apache/airflow/pull/73118#discussion_r4004994988
##########
task-sdk/src/airflow/sdk/importers/python_importer.py:
##########
@@ -56,12 +65,60 @@
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:
+ return self._definition.read_bytes()
+
+
+class _DefinitionBytecodeLoader(importlib.abc.Loader):
"""
- Importer for Python DAG files.
+ Execute a definition's compiled bytecode (``.pyc``) straight from its
bytes.
- This is the default importer registered with the DagImporterRegistry. It
handles
- .py files containing Python DAGs.
+ 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 get_code(self, fullname: str) -> Any:
+ data = self._definition.read_bytes()
+ if len(data) < 16 or data[:4] != importlib.util.MAGIC_NUMBER:
Review Comment:
Thinner than `SourcelessFileLoader` in two ways: `_classify_pyc` also
validates the PEP 552 flags word, and `_compile_bytecode` rejects a payload
that unmarshals to something other than a code object. Without the second, a
`.pyc` holding a marshalled string reaches the `exec()` below and runs as
source, where the old loader raises `ImportError: Non-code object`. Invalid
flags are likewise accepted here and were rejected before. `isinstance(code,
types.CodeType)` covers both, and keeps
`test_import_corrupt_pyc_captured_as_error` meaningful past the magic number.
##########
task-sdk/src/airflow/sdk/importers/zip_importer.py:
##########
@@ -166,72 +166,73 @@ def can_handle(self, definition: DagDefinition | str |
Path) -> bool:
def list_dag_definitions(
self,
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)
+ ) -> Iterator[ZipFileDagDefinition]:
+ """
+ List importable members across the bundle's zip archives.
+
+ Each member is yielded as a plain ZipFileDagDefinition;
import_definition
+ re-resolves the internal importer from the member's extension.
+ """
+ for archive in find_file_dag_definitions(bundle.path,
self.supported_extensions):
+ try:
+ with zipfile.ZipFile(archive.path) as z:
+ member_names = z.namelist()
+ except Exception as e:
+ log.warning("Skipping unreadable ZIP archive %s: %s",
archive.path, e)
+ continue
+
+ member_set = set(member_names)
+ 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,
+ archive.path,
+ )
+ continue
+
+ # Skip compiled-bytecode caches, and prefer source over a
side-by-side .pyc,
+ # so a member and its compiled form are never both imported.
+ if "__pycache__" in member_path.parts:
+ continue
+ if member_name.endswith(".pyc") and member_name[:-1] in
member_set:
+ continue
+
+ if self._get_internal_importer(member_name) is None:
+ continue
+ yield ZipFileDagDefinition(zip_path=archive.path,
file_path=member_name)
Review Comment:
`get_file_suffix()` on a `ZipFileDagDefinition` returns the member's suffix,
not `.zip`, so `ZipImporter.can_handle(member)` is False and
`get_importer(member)` routes it to `PythonDagImporter`, which imports without
`_temporary_sys_path(zip_path)`. At b5b0cb1 a cross-member DAG imports fine via
`ZipImporter.import_definition` and fails via the registry-resolved importer
with `ModuleNotFoundError: No module named 'shared_helper'`.
`registry.can_handle(member)` still returns True, so no warning either.
The new docstring promises definitions come back to the emitting importer,
but that holds only for the caller, and the one the SDK ships dispatches on
suffix. Separate from the return-type thread above: a `FileDagDefinition`
subclass would still carry `.py`.
##########
task-sdk/src/airflow/sdk/importers/base.py:
##########
@@ -234,9 +239,17 @@ def get_file_suffix(definition: DagDefinition | str |
Path) -> str | 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[FileDagDefinition]:
+ """
+ Discover file DAG definitions in a bundle by *identity* alone.
+
+ This walk decides purely from the file's name and path -- extension,
``.airflowignore``,
+ and the Python source/bytecode pairing -- and never reads a file's
contents. Deciding
+ whether a discovered file actually contains a DAG is content-based work
that belongs to
Review Comment:
Two parse-unit consequences not in the PR body.
No cheap filter is left: `safe_mode` is gone from `list_dag_definitions` and
`AbstractDagImporter` has no content pre-check, so `might_contain_dag` is
reachable only on `PythonDagImporter`. `find_dag_file_paths` drops what it
rejects (`utils/file.py#L114`) and the survivors become the file queue, one
`DagFileProcessorProcess` each (`manager.py#L1431`). On 10 example DAGs plus a
41-module `common/` package: 51 definitions here vs 10 from
`list_py_file_paths`.
That walk also yields one path per archive, so a zip is one parse unit
today; per-member definitions make it 1:N.
Should definitions map 1:1 to parse requests after the rewiring, and where
does the queue-size filter live once they do?
##########
task-sdk/src/airflow/sdk/importers/base.py:
##########
@@ -249,8 +262,14 @@ def find_file_dag_definitions(
if path.suffix.lower() not in supported_exts:
continue
- if safe_mode and not might_contain_dag(str(path), safe_mode,
conf=conf):
+ # Skip compiled-bytecode caches, and prefer source over a side-by-side
.pyc, so a
+ # module and its compiled form are never both imported (a .pyc is used
only when
+ # there is no matching .py, i.e. sourceless distribution).
+ if "__pycache__" in path.parts:
continue
+ if path.suffix.lower() == ".pyc" and path.with_suffix(".py").exists():
Review Comment:
`.exists()` is a raw stat, so this dedups against the filesystem rather than
against what the walk yielded (the zip side compares to `member_set`). With
`.airflowignore` matching `only.py` and `only.pyc` shipped beside it, discovery
returns nothing, because the ignored `.py` is still on disk. And
`PythonDagImporter(extensions=[".pyc"])` loses `sourceless.pyc` to a
`sourceless.py` sibling no importer handles. Both go away if the preference is
decided over the paths already emitted.
##########
task-sdk/src/airflow/sdk/importers/python_importer.py:
##########
@@ -56,12 +65,60 @@
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:
+ return self._definition.read_bytes()
Review Comment:
`get_data` ignores `path` and always returns the definition's own bytes,
where `SourceFileLoader.get_data` read the path it was handed. A DAG asking
`__loader__.get_data(...)` or `pkgutil.get_data` for a sibling `config.json`
gets `b"MARKER = 'i am the dag source'"` here, and the real JSON on the old
loader. Raising for a path that isn't the definition's own would at least be
loud.
##########
task-sdk/src/airflow/sdk/importers/zip_importer.py:
##########
@@ -166,72 +166,73 @@ def can_handle(self, definition: DagDefinition | str |
Path) -> bool:
def list_dag_definitions(
self,
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)
+ ) -> Iterator[ZipFileDagDefinition]:
+ """
+ List importable members across the bundle's zip archives.
+
+ Each member is yielded as a plain ZipFileDagDefinition;
import_definition
+ re-resolves the internal importer from the member's extension.
+ """
+ for archive in find_file_dag_definitions(bundle.path,
self.supported_extensions):
+ try:
+ with zipfile.ZipFile(archive.path) as z:
+ member_names = z.namelist()
+ except Exception as e:
+ log.warning("Skipping unreadable ZIP archive %s: %s",
archive.path, e)
Review Comment:
This drops the only user-visible signal for a broken archive.
`import_definition` used to raise
`DagImportError(error_type="zip_read_error")`, which reaches the import-errors
surface; `test_corrupted_zip_file` was relaxed to assert the log warning
instead. Discovery already holds the exception, so could it yield something
that fails at import, and let the user learn why the archive went missing?
##########
task-sdk/src/airflow/sdk/importers/python_importer.py:
##########
@@ -56,12 +65,60 @@
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)
Review Comment:
Not this line's fault, but it pushes the `:` separator further out: a zip
member's `__file__` becomes `archive.zip:member.py`, and
`dag.fileloc`/`relative_fileloc` already use that form via `repr()` and
`get_relative_loc`. airflow-core only understands the `os.sep` form.
`ZIP_REGEX` (`utils/file.py#L40`) yields no archive group for
`archive.zip:dag.py`, so `correct_maybe_zipped` and `open_maybe_zipped` pass it
through and the later `os.path.isfile` fails.
It came in with #72369, but this change makes per-member definitions the
only way a zip DAG gets discovered, so every zip fileloc will be the colon
form. Worth settling `/` vs `:` before the rewiring depends on it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]