ephraimbuddy commented on code in PR #73118:
URL: https://github.com/apache/airflow/pull/73118#discussion_r4047596808
##########
task-sdk/src/airflow/sdk/importers/zip_importer.py:
##########
@@ -168,70 +172,78 @@ def list_dag_definitions(
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[ZipMemberDagDefinition | DagImportError]:
+ """
+ List importable members across the bundle's zip archives.
+
+ Each member is yielded as a plain ZipMemberDagDefinition;
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("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",
+ )
+ 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:
Review Comment:
With `internal_importers={".pyc": PythonDagImporter()}`, an archive
containing `dag.py` and `dag.pyc` now yields nothing: the source is
unsupported, but its presence suppresses the supported bytecode before importer
lookup. Please deduplicate over supported candidates, as filesystem discovery
does, and normalize extension case so `dag.PY` and `dag.pyc` do not both
survive.
##########
task-sdk/src/airflow/sdk/importers/zip_importer.py:
##########
@@ -168,70 +172,78 @@ def list_dag_definitions(
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[ZipMemberDagDefinition | DagImportError]:
+ """
+ List importable members across the bundle's zip archives.
+
+ Each member is yielded as a plain ZipMemberDagDefinition;
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("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",
+ )
+ 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 (importer := self._get_internal_importer(member_name)) is
None:
+ continue
+ member = ZipMemberDagDefinition(zip_path=archive.path,
file_path=member_name)
+ if safe_mode and not importer.might_contain_dag(member,
safe_mode):
Review Comment:
A zip with a valid directory but a bad member CRC raises `BadZipFile` here
and terminates the discovery generator, so later valid members are never
yielded. I reproduced this against the base, which reports the bad member as an
import error and still imports the following Dag. Please catch member-read
failures individually, yield a `DagImportError`, and continue.
##########
shared/module_loading/src/airflow_shared/module_loading/dag_file.py:
##########
@@ -102,6 +125,10 @@ def might_contain_dag(
)
if might_contain_dag_callable is None:
Review Comment:
The shipped configuration explicitly defaults to
`airflow.utils.file.might_contain_dag_via_default_heuristic`, so `getimport()`
returns a callable rather than `None`. Default-config zip discovery therefore
still reaches `as_file()` and extracts every scanned member to a temporary
file. Please recognize the configured built-in heuristic here and test with the
real default configuration.
##########
task-sdk/src/airflow/sdk/importers/python_importer.py:
##########
@@ -42,26 +45,83 @@
DagImportResult,
DagImportWarning,
DagSourceCode,
+ FileDagDefinition,
_normalize_extensions,
find_file_dag_definitions,
get_file_suffix,
)
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):
Review Comment:
Unlike `SourcelessFileLoader`, this loader has no `is_package()`, so
`spec_from_loader()` treats `__init__.pyc` as an ordinary module. A sourceless
Dag package containing `from .helper import ID` imports on the base but now
fails with `attempted relative import with no known parent package`. Please
preserve package detection and cover this case.
--
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]
