jason810496 commented on code in PR #73118:
URL: https://github.com/apache/airflow/pull/73118#discussion_r4080922419
##########
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:
Addressed in
https://github.com/apache/airflow/pull/73118/commits/d2f744b0971e999d15a7c53e5282665accf2f51d
##########
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:
Addressed in
https://github.com/apache/airflow/pull/73118/commits/d2f744b0971e999d15a7c53e5282665accf2f51d
##########
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:
Addressed in
https://github.com/apache/airflow/pull/73118/commits/d2f744b0971e999d15a7c53e5282665accf2f51d
##########
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:
Addressed in
https://github.com/apache/airflow/pull/73118/commits/d2f744b0971e999d15a7c53e5282665accf2f51d
--
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]
