jason810496 commented on code in PR #73118:
URL: https://github.com/apache/airflow/pull/73118#discussion_r4015565995
##########
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:
I just confirmed with the `_classify_pyc` definition in `importlib`.
```python
def _classify_pyc(data, name, exc_details):
"""Perform basic validity checking of a pyc header and return the flags
field,
which determines how the pyc should be further validated against the
source.
*data* is the contents of the pyc file. (Only the first 16 bytes are
required, though.)
*name* is the name of the module being imported. It is used for logging.
*exc_details* is a dictionary passed to ImportError if it raised for
improved debugging.
ImportError is raised when the magic number is incorrect or when the
flags
field is invalid. EOFError is raised when the data is found to be
truncated.
"""
magic = data[:4]
if magic != MAGIC_NUMBER:
message = f'bad magic number in {name!r}: {magic!r}'
_bootstrap._verbose_message('{}', message)
raise ImportError(message, **exc_details)
if len(data) < 16:
message = f'reached EOF while reading pyc header of {name!r}'
_bootstrap._verbose_message('{}', message)
raise EOFError(message)
flags = _unpack_uint32(data[4:8])
# Only the first two flags are defined.
if flags & ~0b11:
message = f'invalid flags {flags!r} in {name!r}'
raise ImportError(message, **exc_details)
return flags
```
It seems we could invoke `SourcelessFileLoader.get_code` to ensure the
correct logic.
##########
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:
Comparing with the merged logic, it seems we can't import the `.pyc` as
Kaxil pointed.
https://github.com/apache/airflow/blob/9ecc914fbd00d72617e0b8d107f04d6563d7bbda/task-sdk/src/airflow/sdk/importers/python_importer.py#L59-L85
##########
shared/module_loading/src/airflow_shared/module_loading/dag_file.py:
##########
@@ -37,6 +39,13 @@ def __call__(self, file_path: str, zip_file: zipfile.ZipFile
| None = None) -> b
class _ConfLike(Protocol):
def getimport(self, section: str, key: str, **kwargs: Any) -> Any: ...
+ class _DagDefinitionLike(Protocol):
Review Comment:
Would it be better to make this a runtime_checkable protocal? So that we can
use safely rely on the `isinstance` as the first check for the
`_DagDefinitionLike` then fallback to the str or path like handling.
##########
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:
- might_contain_dag_callable = might_contain_dag_via_default_heuristic
+ return might_contain_dag_via_default_heuristic(file_path,
zip_file=zip_file)
- return might_contain_dag_callable(file_path=file_path, zip_file=zip_file)
+ if isinstance(file_path, (str, os.PathLike)):
+ return might_contain_dag_callable(file_path=file_path,
zip_file=zip_file)
+ # Custom 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)
Review Comment:
How about having try catch when invoking the `might_contain_dag_callable`?
--
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]
