Harden claimed-file discovery errors and archive importers

- Resolve a relative discovery-error reference against the bundle, map an
  archive member reference to its archive, and drop paths outside the bundle.
  A bundle-relative reference was resolved against the working directory,
  which dropped the file or crashed the Dag processor.
- Key a single file's listing errors by the file, not ".".
- Reject a ZipImporter as a claiming importer: it imports archive members,
  whose Dags could not be persisted, so each claimed file now reports the
  configuration error.
- Warn once per registry when .py, .pyc or .zip is mapped to another importer.

Author: jason810496

airflow-core/src/airflow/dag_processing/importer_routing.py

@@ -25,19 +25,24 @@
 
 from __future__ import annotations
 
+import dataclasses
 import logging
 import os
+import weakref
 from dataclasses import dataclass
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, cast
 
+from airflow.exceptions import AirflowConfigException
 from airflow.sdk.importers import (
     AbstractDagImporter,
     DagImporterRegistry,
     DagImportError,
     DagImportResult,
     DagSourceCode,
     FilesystemDagDefinition,
+    PythonDagImporter,
+    ZipImporter,
     ZipMemberDagDefinition,
     find_file_dag_definitions,
     get_file_suffix,
@@ -54,15 +59,16 @@
 LEGACY_EXTENSIONS = frozenset({".py", ".pyc", ".zip"})
 
 _reported_registry_failures: set[str | None] = set()
+_checked_registries: weakref.WeakSet[DagImporterRegistry] = weakref.WeakSet()
 
 
 @dataclass(frozen=True)
 class BundleRef:
     """
     The part of a Dag bundle that importers read: its name and its root path.
 
-    Core passes it in place of a ``BaseDagBundle`` when it imports a single file, so an importer
-    must not rely on any other bundle attribute.
+    Core passes it in place of a ``BaseDagBundle`` when a Dag bag lists a folder or imports a file,
+    so an importer must not rely on any other bundle attribute.
     """
 
     name: str | None
@@ -88,9 +94,25 @@ def get_task_sdk_registry(bundle_name: str | None) -> DagImporterRegistry | None
             log.exception("Cannot build the Task SDK Dag importer registry for bundle %s", bundle_name)
         return None
     _reported_registry_failures.discard(bundle_name)
+    _warn_about_ignored_legacy_importers(registry)
     return registry
 
 
+def _warn_about_ignored_legacy_importers(registry: DagImporterRegistry) -> None:
+    if registry in _checked_registries:
+        return
+    _checked_registries.add(registry)
+    for ext in sorted(LEGACY_EXTENSIONS.intersection(registry.supported_extensions())):
+        try:
+            importer = registry.get_importer(f"_{ext}")
+        except Exception:
+            importer = None
+        if type(importer) not in (PythonDagImporter, ZipImporter):
+            log.warning(
+                "Ignoring the Dag importer configured for %s files: they always use the legacy importer", ext
+            )
+
+
 def _get_claimed_extensions(registry: DagImporterRegistry) -> list[str]:
     return [ext for ext in registry.supported_extensions() if ext not in LEGACY_EXTENSIONS]
 
@@ -111,11 +133,18 @@ def get_claiming_importer(
     """
     Return the Task SDK importer that claims ``path``, or ``None`` when the file is not claimed.
 
-    :raises AirflowConfigException: if the importer configured for the extension cannot be loaded.
+    :raises AirflowConfigException: if the importer configured for the extension cannot be loaded,
+        or imports archive members, which is not supported for a claimed file.
     """
     if not is_claimed(registry, path):
         return None
-    return registry.get_importer(Path(path))
+    importer = registry.get_importer(Path(path))
+    if isinstance(importer, ZipImporter):
+        raise AirflowConfigException(
+            f"{type(importer).__name__} cannot claim {get_file_suffix(Path(path))} files: an importer "
+            "that claims files must import each file as one Dag definition, not as archive members."
+        )
+    return importer
 
 
 def _get_claiming_importer_or_none(
@@ -149,24 +178,42 @@ def has_claiming_importers(registry: DagImporterRegistry) -> bool:
     return bool(_get_claimed_extensions(registry))
 
 
+def _find_bundle_file(bundle_path: Path, reference: Path) -> Path | None:
+    """
+    Return the file in the bundle that ``reference`` names, or ``None``.
+
+    A relative reference is relative to the bundle. A reference into an archive, such as
+    ``x.jar/member.py``, names the archive.
+    """
+    bundle_path = Path(os.path.normpath(bundle_path))
+    path = Path(os.path.normpath(bundle_path / reference))
+    for candidate in (path, *path.parents):
+        if not candidate.is_relative_to(bundle_path):
+            return None
+        if candidate.is_file():
+            return candidate
+    return None
+
+
 def _get_listed_path(
-    registry: DagImporterRegistry, importer: AbstractDagImporter[Any], item: object
+    registry: DagImporterRegistry, importer: AbstractDagImporter[Any], item: object, bundle_path: Path
 ) -> Path | None:
     """Return the file to parse for a listed item, or ``None`` when ``importer`` does not own one."""
     if isinstance(item, DagImportError):
         log.warning("Dag discovery error: %s", item.format_message())
         # Parsing the file lists it again and records the error as its import error.
-        path = Path(item.source_reference)
-        if not path.is_file():
-            return None
+        reference = Path(item.source_reference)
     elif isinstance(item, ZipMemberDagDefinition):
-        path = item.zip_path
+        reference = item.zip_path
     elif isinstance(item, FilesystemDagDefinition):
-        path = item.path
+        reference = item.path
     else:
         log.warning("Skipping %r: %s did not list a file", item, type(importer).__name__)
         return None
-    return path if _get_claiming_importer_or_none(registry, path) is importer else None
+    path = _find_bundle_file(bundle_path, reference)
+    if path is None or _get_claiming_importer_or_none(registry, path) is not importer:
+        return None
+    return path
 
 
 def merge_claimed_paths(
@@ -184,13 +231,13 @@ def merge_claimed_paths(
     or raises while listing, every file with its extensions is kept: parsing each one then reports
     the failure, instead of the files' Dags being treated as deleted.
     """
-    bundle_path = _cast_to_bundle(bundle).path
+    bundle_path = Path(_cast_to_bundle(bundle).path)
     claimed_paths: dict[str, None] = {}
     for importer, extensions in _group_claiming_importers(registry):
         if importer is not None:
             try:
                 for item in importer.list_dag_definitions(_cast_to_bundle(bundle), safe_mode=safe_mode):
-                    if (path := _get_listed_path(registry, importer, item)) is not None:
+                    if (path := _get_listed_path(registry, importer, item, bundle_path)) is not None:
                         claimed_paths.setdefault(os.fspath(path))
                 continue
             except Exception:
@@ -235,7 +282,9 @@ def iter_claimed_results(
     bundle = _cast_to_bundle(BundleRef(name=bundle_name, path=bundle_path))
     for item in items:
         if isinstance(item, DagImportError):
-            yield DagImportResult(errors=[item])
+            # The listing root is the file itself, so a relative reference is relative to it.
+            reference = os.fspath(file_path / item.source_reference)
+            yield DagImportResult(errors=[dataclasses.replace(item, source_reference=reference)])
             continue
         try:
             result = importer.import_definition(item, bundle)

airflow-core/tests/unit/dag_processing/fake_importers.py

@@ -21,12 +21,14 @@
 
 import contextlib
 import json
+import os
 import zipfile
 from typing import TYPE_CHECKING, Any
 
 from airflow.sdk import DAG, BaseOperator
 from airflow.sdk.importers import (
     AbstractDagImporter,
+    DagDefinition,
     DagImportError,
     DagImportResult,
     DagImportWarning,
@@ -48,6 +50,8 @@
 FAILING_SOURCE_JAR_IMPORTER = f"{__name__}.FailingSourceJarImporter"
 ERROR_LISTING_JAR_IMPORTER = f"{__name__}.ErrorListingJarImporter"
 RAISING_LISTING_JAR_IMPORTER = f"{__name__}.RaisingListingJarImporter"
+MEMBER_ERROR_LISTING_JAR_IMPORTER = f"{__name__}.MemberErrorListingJarImporter"
+NON_FILE_LISTING_JAR_IMPORTER = f"{__name__}.NonFileListingJarImporter"
 
 
 def build_dag(dag_id: str, definition: FilesystemDagDefinition, bundle: Any) -> DAG:
@@ -124,15 +128,52 @@ def get_source_code(self, definition) -> DagSourceCode:
 
 
 class ErrorListingJarImporter(JarDagImporter):
-    """Report a JAR without ``dags.txt`` as a discovery error instead of skipping it."""
+    """
+    Report a JAR without ``dags.txt`` as a discovery error instead of skipping it.
+
+    The error names the JAR relative to the listed root, as the Task SDK ``ZipImporter`` does.
+    """
+
+    member = ""
 
     def list_dag_definitions(self, bundle, *, safe_mode=True) -> Iterator[Any]:
         for definition in find_file_dag_definitions(bundle.path, self.supported_extensions):
             with zipfile.ZipFile(definition.path) as jar:
                 if "dags.txt" in jar.namelist():
                     yield definition
                 else:
-                    yield DagImportError(source_reference=repr(definition), message="no dags.txt")
+                    reference = os.path.join(definition.get_relative_loc(bundle.path), self.member)
+                    yield DagImportError(source_reference=os.path.normpath(reference), message="no dags.txt")
+
+
+class MemberErrorListingJarImporter(ErrorListingJarImporter):
+    """Report the discovery error against a member of the JAR, such as ``library.jar/Main.java``."""
+
+    member = "Main.java"
+
+
+class NonFileListingJarImporter(JarDagImporter):
+    """List a definition that is not a file next to the JARs."""
+
+    def list_dag_definitions(self, bundle, *, safe_mode=True) -> Iterator[Any]:
+        yield NotAFileDagDefinition()
+        yield from super().list_dag_definitions(bundle, safe_mode=safe_mode)
+
+
+class NotAFileDagDefinition(DagDefinition):
+    freshness_token = ""
+
+    def get_relative_loc(self, root=None) -> str:
+        return "not-a-file"
+
+    def read_bytes(self) -> bytes:
+        return b""
+
+    def as_file(self):
+        raise NotImplementedError
+
+    def __repr__(self) -> str:
+        return "not-a-file"
 
 
 class RaisingListingJarImporter(JarDagImporter):

airflow-core/tests/unit/dag_processing/test_dagbag.py

@@ -59,6 +59,7 @@
 from unit.dag_processing.fake_importers import (
     ERROR_LISTING_JAR_IMPORTER,
     FAKE_IMPORTER,
+    MEMBER_ERROR_LISTING_JAR_IMPORTER,
     task_sdk_importers,
     write_jar,
 )
@@ -1582,15 +1583,41 @@ def test_importer_that_cannot_load_is_an_import_error(self, tmp_path):
             in (dagbag.import_errors["claimed.fake"])
         )
 
-    def test_discovery_error_is_an_import_error(self, tmp_path):
+    @pytest.mark.parametrize(
+        ("importer_config", "error_key"),
+        [
+            pytest.param(ERROR_LISTING_JAR_IMPORTER, "library.jar", id="relative-to-the-file"),
+            pytest.param(MEMBER_ERROR_LISTING_JAR_IMPORTER, "library.jar/Main.java", id="archive-member"),
+        ],
+    )
+    def test_single_file_discovery_error_is_keyed_by_the_file(self, tmp_path, importer_config, error_key):
         library = write_jar(tmp_path / "library.jar")
 
-        with task_sdk_importers(ERROR_LISTING_JAR_IMPORTER):
+        with task_sdk_importers(importer_config):
             dagbag = DagBag(dag_folder=os.fspath(library), bundle_path=tmp_path, bundle_name="testing")
 
         assert dagbag.dags == {}
+        assert dagbag.import_errors == {error_key: "no dags.txt"}
+
+    def test_folder_discovery_error_is_an_import_error(self, tmp_path):
+        write_jar(tmp_path / "library.jar")
+
+        with task_sdk_importers(ERROR_LISTING_JAR_IMPORTER):
+            dagbag = DagBag(dag_folder=os.fspath(tmp_path), bundle_path=tmp_path, bundle_name="testing")
+
         assert dagbag.import_errors == {"library.jar": "no dags.txt"}
 
+    def test_archive_member_importer_cannot_claim_files(self, tmp_path):
+        with zipfile.ZipFile(tmp_path / "packaged.jar", "w") as jar:
+            jar.writestr("zipped_dag.py", PY_DAG_SOURCE.format(dag_id="zipped_dag"))
+
+        with task_sdk_importers({"classpath": "airflow.sdk.importers.ZipImporter", "extensions": [".jar"]}):
+            dagbag = DagBag(dag_folder=os.fspath(tmp_path), bundle_path=tmp_path, bundle_name="testing")
+
+        assert dagbag.dags == {}
+        assert list(dagbag.import_errors) == ["packaged.jar"]
+        assert "ZipImporter cannot claim .jar files" in dagbag.import_errors["packaged.jar"]
+
     def test_built_in_importer_claims_a_configured_extension(self, tmp_path):
         claimed = tmp_path / "python_dag.dagpy"
         claimed.write_text(PY_DAG_SOURCE.format(dag_id="python_dag"))

airflow-core/tests/unit/dag_processing/test_importer_routing.py

@@ -0,0 +1,62 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import logging
+
+import pytest
+
+from airflow.dag_processing.importer_routing import _find_bundle_file, get_task_sdk_registry
+
+from unit.dag_processing.fake_importers import FAKE_IMPORTER, task_sdk_importers, write_jar
+
+
+@pytest.mark.parametrize(
+    ("reference", "expected"),
+    [
+        pytest.param("native.jar", "native.jar", id="relative-to-the-bundle"),
+        pytest.param("native.jar/Main.java", "native.jar", id="archive-member"),
+        pytest.param("{bundle}/native.jar", "native.jar", id="absolute-inside"),
+        pytest.param("../outside.jar", None, id="relative-outside"),
+        pytest.param("{root}/outside.jar", None, id="absolute-outside"),
+        pytest.param("missing.jar", None, id="missing"),
+    ],
+)
+def test_find_bundle_file(tmp_path, reference, expected):
+    bundle_path = tmp_path / "bundle"
+    bundle_path.mkdir()
+    write_jar(bundle_path / "native.jar", "native_dag")
+    write_jar(tmp_path / "outside.jar", "outside_dag")
+
+    found = _find_bundle_file(
+        bundle_path, bundle_path.joinpath(reference.format(bundle=bundle_path, root=tmp_path))
+    )
+
+    assert found == (bundle_path / expected if expected else None)
+
+
+def test_legacy_extension_mapping_is_warned_once(caplog):
+    with task_sdk_importers({"classpath": FAKE_IMPORTER, "extensions": [".fake", ".py"]}):
+        with caplog.at_level(logging.WARNING, logger="airflow.dag_processing.importer_routing"):
+            get_task_sdk_registry("testing")
+            get_task_sdk_registry("testing")
+
+    warnings = [r.getMessage() for r in caplog.records if r.name == "airflow.dag_processing.importer_routing"]
+    assert warnings == [
+        "Ignoring the Dag importer configured for .py files: they always use the legacy importer"
+    ]