diff --git a/airflow-core/docs/core-concepts/dags.rst b/airflow-core/docs/core-concepts/dags.rst index 18acf9bbdec3c..7db1c63adeebc 100644 --- a/airflow-core/docs/core-concepts/dags.rst +++ b/airflow-core/docs/core-concepts/dags.rst @@ -719,6 +719,21 @@ You can either do this all inside of the Dag bundle, with a standard filesystem package1/__init__.py package1/functions.py +Dags do not have to sit at the root of the archive. A Dag file placed inside a sub-folder is +discovered as long as that folder — and every folder above it — is an importable package, i.e. +contains an ``__init__.py``:: + + my_dag1.py + team_a/__init__.py + team_a/team_dag.py + team_a/reports/__init__.py + team_a/reports/report_dag.py + +A ``.py`` file in a sub-folder that is *not* a package (no ``__init__.py`` in that folder or one of +its parents) cannot be imported from a zip archive, so it is skipped and a warning naming the +unreachable file is logged. Add an ``__init__.py`` to every parent folder, or move the Dag to the +archive root, to make it discoverable. + Note that packaged Dags come with some caveats: * They cannot be used if you have pickling enabled for serialization diff --git a/airflow-core/newsfragments/71703.improvement.rst b/airflow-core/newsfragments/71703.improvement.rst new file mode 100644 index 0000000000000..d77f11a1af58e --- /dev/null +++ b/airflow-core/newsfragments/71703.improvement.rst @@ -0,0 +1 @@ +Dags placed in sub-folders of a packaged (zip) Dag bundle are now discovered, as long as the sub-folder and every folder above it is an importable package (contains an ``__init__.py``). Previously only Dags at the archive root were imported and sub-folder Dags were silently skipped. A Dag file in a non-package sub-folder still cannot be imported from a zip, but it is now skipped with a warning naming the unreachable file instead of being dropped silently. diff --git a/airflow-core/src/airflow/dag_processing/importers/python_importer.py b/airflow-core/src/airflow/dag_processing/importers/python_importer.py index 2eedc4dff4dec..d6c377b3a7488 100644 --- a/airflow-core/src/airflow/dag_processing/importers/python_importer.py +++ b/airflow-core/src/airflow/dag_processing/importers/python_importer.py @@ -76,6 +76,53 @@ def handle_timeout(signum, frame): signal.setitimer(signal.ITIMER_REAL, 0) +def _zip_package_dirs(zip_file: zipfile.ZipFile) -> set[str]: + """Return the archive folders that are importable packages (contain an ``__init__`` module).""" + package_dirs: set[str] = set() + for name in zip_file.namelist(): + member = Path(name) + if member.stem == "__init__" and member.suffix in (".py", ".pyc"): + package_dirs.add(member.parent.as_posix()) + return package_dirs + + +def _zip_member_module_name(zip_path: Path, package_dirs: set[str]) -> str | None: + """ + Resolve the dotted module name for a ``.py``/``.pyc`` member sitting in a sub-folder of a zip. + + ``zipimport`` can only reach a sub-folder module through its parent package, so every parent + folder must contain an ``__init__`` module. Namespace packages (folders without ``__init__``) + are not importable from a zip, so this returns ``None`` for them and the caller surfaces the + unreachable file instead of importing it. Root-level members are handled by the caller. + """ + if zip_path.stem == "__init__": + module_parts = zip_path.parts[:-1] + required_package_dirs = zip_path.parts[:-2] + else: + module_parts = (*zip_path.parts[:-1], zip_path.stem) + required_package_dirs = zip_path.parts[:-1] + + for depth in range(1, len(required_package_dirs) + 1): + if "/".join(required_package_dirs[:depth]) not in package_dirs: + return None + return ".".join(module_parts) + + +def _zip_owned_module_names( + package_dirs: set[str], candidates: list[tuple[str, zipfile.ZipInfo]] +) -> set[str]: + """Return every module and (intermediate) package name that the archive defines.""" + names: set[str] = set() + for package_dir in package_dirs: + if package_dir != ".": + names.add(package_dir.replace("/", ".")) + for mod_name, _ in candidates: + parts = mod_name.split(".") + for depth in range(1, len(parts) + 1): + names.add(".".join(parts[:depth])) + return names + + class PythonDagImporter(AbstractDagImporter): """ Importer for Python DAG files and zip archives containing Python DAGs. @@ -285,13 +332,36 @@ def _load_modules_from_zip( mods: list[ModuleType] = [] with zipfile.ZipFile(filepath) as current_zip_file: + package_dirs = _zip_package_dirs(current_zip_file) + + candidates: list[tuple[str, zipfile.ZipInfo]] = [] for zip_info in current_zip_file.infolist(): zip_path = Path(zip_info.filename) - if zip_path.suffix not in [".py", ".pyc"] or len(zip_path.parts) > 1: + if zip_path.suffix not in [".py", ".pyc"]: continue - if zip_path.stem == "__init__": - log.warning("Found %s at root of %s", zip_path.name, filepath) + if len(zip_path.parts) == 1: + if zip_path.stem == "__init__": + log.warning("Found %s at root of %s", zip_path.name, filepath) + mod_name = zip_path.stem + else: + resolved = _zip_member_module_name(zip_path, package_dirs) + if resolved is None: + # The file lives in a folder that is not an importable package, so it can + # never be loaded from the zip. Surface it when it looks like a DAG instead + # of dropping it silently, which otherwise hides the file with no diagnostic. + if might_contain_dag(zip_info.filename, safe_mode, current_zip_file): + log.warning( + "Ignoring %s in %s: it is not part of a package (no __init__.py in " + "its folder), so it cannot be imported from a zip archive. Add an " + "__init__.py to every parent folder, or move the DAG to the archive " + "root.", + zip_info.filename, + filepath, + ) + result.skipped_files.append(f"{filepath}:{zip_info.filename}") + continue + mod_name = resolved log.debug("Reading %s from %s", zip_info.filename, filepath) @@ -299,10 +369,20 @@ def _load_modules_from_zip( result.skipped_files.append(f"{filepath}:{zip_info.filename}") continue - mod_name = zip_path.stem - if mod_name in sys.modules: - del sys.modules[mod_name] + candidates.append((mod_name, zip_info)) + + # Import shallower modules first so a parent package is fully loaded before its + # sub-modules; this keeps DAG autoregistration attributed to the right module when a + # package's __init__ also defines DAGs. + candidates.sort(key=lambda candidate: candidate[0].count(".")) + + # Drop any module or package this archive defines from a previous import (e.g. a + # different zip that ships a package of the same name) so its members resolve against + # this archive rather than a stale one left in sys.modules. + for stale in _zip_owned_module_names(package_dirs, candidates): + sys.modules.pop(stale, None) + for mod_name, zip_info in candidates: DagContext.current_autoregister_module_name = mod_name try: sys.path.insert(0, filepath) diff --git a/airflow-core/tests/unit/dag_processing/test_dagbag.py b/airflow-core/tests/unit/dag_processing/test_dagbag.py index cb110ef5c7fda..ecbccc242b6d8 100644 --- a/airflow-core/tests/unit/dag_processing/test_dagbag.py +++ b/airflow-core/tests/unit/dag_processing/test_dagbag.py @@ -58,6 +58,16 @@ pytestmark = pytest.mark.db_test +_ZIP_SUBFOLDER_DAG_TEMPLATE = textwrap.dedent( + """ + from airflow.sdk import DAG + from airflow.providers.standard.operators.empty import EmptyOperator + + with DAG(dag_id={dag_id!r}, schedule=None): + EmptyOperator(task_id="noop") + """ +) + def _standard_example_dags_folder() -> Path: """ @@ -888,6 +898,57 @@ def test_get_zip_dag_registration(self, test_zip_path, expected): assert dag, f"{dag_id} was bagged" assert dag.fileloc.endswith(f"{pathlib.Path(test_zip_path).parent}/{path}") + def test_zip_dag_in_package_subfolder_is_loaded(self, tmp_path): + """A DAG placed inside a proper package (has ``__init__.py``) of a zip is discovered.""" + zipped = tmp_path / "test_zip_pkg.zip" + with zipfile.ZipFile(zipped, "w") as zf: + zf.writestr("root_dag.py", _ZIP_SUBFOLDER_DAG_TEMPLATE.format(dag_id="zip_root_dag")) + zf.writestr("pkg/__init__.py", "") + zf.writestr("pkg/nested_dag.py", _ZIP_SUBFOLDER_DAG_TEMPLATE.format(dag_id="zip_pkg_nested_dag")) + + dagbag = DagBag(dag_folder=os.fspath(zipped)) + + assert not dagbag.import_errors + assert dagbag.get_dag("zip_root_dag") is not None + nested = dagbag.get_dag("zip_pkg_nested_dag") + assert nested is not None + assert nested.fileloc.endswith(f"{zipped}/pkg/nested_dag.py") + + def test_zip_dag_in_deeply_nested_package_is_loaded(self, tmp_path): + """A DAG several package levels deep is discovered when every level is a package.""" + zipped = tmp_path / "test_zip_deep.zip" + with zipfile.ZipFile(zipped, "w") as zf: + zf.writestr("pkg/__init__.py", "") + zf.writestr("pkg/sub/__init__.py", "") + zf.writestr("pkg/sub/deep_dag.py", _ZIP_SUBFOLDER_DAG_TEMPLATE.format(dag_id="zip_deep_dag")) + + dagbag = DagBag(dag_folder=os.fspath(zipped)) + + assert not dagbag.import_errors + dag = dagbag.get_dag("zip_deep_dag") + assert dag is not None + assert dag.fileloc.endswith(f"{zipped}/pkg/sub/deep_dag.py") + + def test_zip_dag_in_non_package_subfolder_is_skipped_with_warning(self, tmp_path, caplog): + """A DAG in a sub-folder without ``__init__.py`` cannot be imported from a zip and is not + silently dropped: it is skipped and a warning names the unreachable file.""" + zipped = tmp_path / "test_zip_ns.zip" + with zipfile.ZipFile(zipped, "w") as zf: + zf.writestr("root_dag.py", _ZIP_SUBFOLDER_DAG_TEMPLATE.format(dag_id="zip_root_dag")) + zf.writestr( + "plaindir/plain_dag.py", _ZIP_SUBFOLDER_DAG_TEMPLATE.format(dag_id="zip_plaindir_dag") + ) + + dagbag = DagBag(dag_folder=os.fspath(zipped)) + + assert not dagbag.import_errors + assert dagbag.get_dag("zip_root_dag") is not None + assert dagbag.get_dag("zip_plaindir_dag") is None + assert { + "event": re.compile(r".*plaindir/plain_dag\.py.*not part of a package"), + "log_level": "warning", + } in caplog + def test_dag_registration_with_failure(self): dagbag = DagBag(dag_folder=os.devnull) found = dagbag.process_file(str(TEST_DAGS_FOLDER / "test_invalid_dup_task.py"))