Discover Dags in sub-folders of packaged (zip) Dag bundles

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

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.

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,24 +332,57 @@ 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)
 
                 if not might_contain_dag(zip_info.filename, safe_mode, current_zip_file, conf=conf):
                     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)

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"))