This is an automated email from the ASF dual-hosted git repository.
shahar1 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 3634bc89d04 Discover a bundle's Dag definitions through the importer
registry (#73728)
3634bc89d04 is described below
commit 3634bc89d04e1ac2ee8acbc261eafec6773b3012
Author: Dilnaz Amanzholova <[email protected]>
AuthorDate: Mon Sep 28 12:28:10 2026 +0200
Discover a bundle's Dag definitions through the importer registry (#73728)
* Discover a bundle's Dag definitions through the importer registry
Dag processing needs every definition in a bundle paired with the importer
that should import it: a composite importer's members, such as zip members,
cannot be routed back to it by extension. An importer whose extensions were
all overridden, for example by a bundle's own ZipImporter, kept listing its
definitions and discovered them twice.
* Materialise each configured importer spec once during discovery
---
task-sdk/src/airflow/sdk/importers/base.py | 109 +++++++++++-----
.../src/airflow/sdk/importers/python_importer.py | 10 +-
task-sdk/tests/task_sdk/importers/test_registry.py | 139 +++++++++++++++++++++
3 files changed, 222 insertions(+), 36 deletions(-)
diff --git a/task-sdk/src/airflow/sdk/importers/base.py
b/task-sdk/src/airflow/sdk/importers/base.py
index 83095bf5ad0..94916fa60eb 100644
--- a/task-sdk/src/airflow/sdk/importers/base.py
+++ b/task-sdk/src/airflow/sdk/importers/base.py
@@ -286,11 +286,14 @@ def get_file_suffix(definition: DagDefinition | str |
Path) -> str | None:
def find_file_dag_definitions(
- bundle_path: Path,
+ root: Path,
supported_extensions: Iterable[str],
) -> Iterator[FilesystemDagDefinition]:
"""
- Discover file DAG definitions in a bundle by *identity* alone.
+ Discover file Dag definitions under ``root`` by *identity* alone.
+
+ ``root`` is normally a bundle directory; it may also be a single file,
which scopes
+ discovery to that file (a zip archive is still yielded for its importer to
expand).
This walk decides purely from the file's name and path: the suffix,
``.airflowignore``, and the Python source/bytecode pairing. It never reads
@@ -298,19 +301,23 @@ def find_file_dag_definitions(
belongs to :meth:`AbstractDagImporter.import_definition`, so importers
whose
validity can only be determined by attempting the import behave the same
way.
"""
- ignore_file_syntax = conf.get_mandatory_value("core",
"DAG_IGNORE_FILE_SYNTAX", fallback="glob")
supported_exts = _normalize_extensions(supported_extensions)
+ def _is_candidate(path: Path) -> bool:
+ return path.is_file() and path.suffix.lower() in supported_exts and
"__pycache__" not in path.parts
+
def _iter_candidates() -> Iterator[Path]:
- for file_path in find_path_from_directory(bundle_path,
".airflowignore", ignore_file_syntax):
+ if root.is_file():
+ # Scoped to one artefact. ``.airflowignore`` is not re-applied:
whoever
+ # narrowed discovery to this path already walked the bundle and
honoured it.
+ if _is_candidate(root):
+ yield root
+ return
+ ignore_file_syntax = conf.get_mandatory_value("core",
"DAG_IGNORE_FILE_SYNTAX", fallback="glob")
+ for file_path in find_path_from_directory(root, ".airflowignore",
ignore_file_syntax):
path = Path(file_path)
- if not path.is_file():
- continue
- if path.suffix.lower() not in supported_exts:
- continue
- if "__pycache__" in path.parts:
- continue
- yield path
+ if _is_candidate(path):
+ yield path
candidates = list(_iter_candidates())
@@ -339,25 +346,29 @@ class _ImporterSpec:
def _parse_importer_specs(configs: Any, context: str) -> list[_ImporterSpec]:
if not isinstance(configs, list):
raise AirflowConfigException(
- f"Invalid importer configuration for {context}: expected a list of
dictionaries."
+ f"Invalid importer configuration for {context}: expected a list of
entries."
)
specs: list[_ImporterSpec] = []
for item in configs:
- if not isinstance(item, dict):
+ # A bare classpath string is shorthand for {"classpath": ...}; the
importer
+ # class then declares its own extensions and takes no kwargs.
+ entry = {"classpath": item} if isinstance(item, str) else item
+ if not isinstance(entry, dict):
raise AirflowConfigException(
- f"Invalid importer configuration for {context}: each entry
must be a dictionary."
+ f"Invalid importer configuration for {context}: "
+ "each entry must be a dictionary or a classpath string."
)
- classpath = item.get("classpath")
+ classpath = entry.get("classpath")
if not classpath:
raise AirflowConfigException(
f"Missing required 'classpath' in importer configuration for
{context}."
)
- kwargs = item.get("kwargs", {})
+ kwargs = entry.get("kwargs", {})
if not isinstance(kwargs, dict):
raise AirflowConfigException(
f"Field 'kwargs' must be a dictionary in importer
configuration for {context}."
)
- extensions = item.get("extensions")
+ extensions = entry.get("extensions")
if extensions is not None:
if not isinstance(extensions, list) or any(not isinstance(ext,
str) for ext in extensions):
raise AirflowConfigException(
@@ -423,9 +434,6 @@ class DagImporterRegistry:
Each extension can only have one importer. If an extension is already
registered,
the new importer will override it and a warning will be logged.
"""
- if importer not in self._ordered_importers:
- self._ordered_importers.append(importer)
-
if extensions is None:
extensions = _get_importer_extensions(importer)
@@ -438,6 +446,9 @@ class DagImporterRegistry:
self._warn_and_evict_extension(ext_lower,
type(importer).__name__)
self._extension_importers[ext_lower] = importer
+ if importer not in self._ordered_importers:
+ self._ordered_importers.append(importer)
+
def register_specs(self, configs: list[dict[str, Any]], context: str) ->
None:
"""Register importer specifications from configuration dictionaries."""
for spec in _parse_importer_specs(configs, context=context):
@@ -455,17 +466,8 @@ class DagImporterRegistry:
if suffix:
if suffix in self._extension_importers:
return self._extension_importers[suffix]
-
if suffix in self._extension_specs:
- spec = self._extension_specs[suffix]
- importer = self._instantiate_spec(spec)
- for e, s in list(self._extension_specs.items()):
- if s is spec:
- self._extension_importers[e] = importer
- del self._extension_specs[e]
- if importer not in self._ordered_importers:
- self._ordered_importers.append(importer)
- return importer
+ return self._materialise_spec(self._extension_specs[suffix])
for importer in reversed(self._ordered_importers):
if importer.can_handle(definition):
@@ -483,6 +485,47 @@ class DagImporterRegistry:
"""Return all registered file extensions."""
return sorted(set(self._extension_importers) |
set(self._extension_specs))
+ def list_dag_definitions(
+ self,
+ bundle: BaseDagBundle,
+ *,
+ safe_mode: bool = True,
+ ) -> Iterator[tuple[AbstractDagImporter[Any], DagDefinition |
DagImportError]]:
+ """
+ List Dag definitions in a bundle across all registered importers.
+
+ Each item is paired with the importer that listed it, which is the one
to import it
+ with: a composite importer lists definitions no extension lookup would
route back to
+ it, such as the members of an archive. A :class:`DagImportError` item
is a
+ discovery-time failure rather than a source to import.
+ """
+ # A spec registered for several extensions appears once per extension,
and
+ # materialising it drops all of them, so take one pending spec at a
time.
+ while self._extension_specs:
+ self._materialise_spec(next(iter(self._extension_specs.values())))
+
+ for importer in self._ordered_importers:
+ for item in importer.list_dag_definitions(bundle,
safe_mode=safe_mode):
+ # A file whose extension was claimed by a later registration
belongs to that
+ # importer; yielding it here as well would import it twice.
+ if (
+ isinstance(item, FilesystemDagDefinition)
+ and self._extension_importers.get(item.suffix, importer)
is not importer
+ ):
+ continue
+ yield importer, item
+
+ def _materialise_spec(self, spec: _ImporterSpec) ->
AbstractDagImporter[Any]:
+ """Instantiate a configured spec and take over every extension it was
registered for."""
+ importer = self._instantiate_spec(spec)
+ for ext, registered in list(self._extension_specs.items()):
+ if registered is spec:
+ self._extension_importers[ext] = importer
+ del self._extension_specs[ext]
+ if importer not in self._ordered_importers:
+ self._ordered_importers.append(importer)
+ return importer
+
@classmethod
def reset(cls) -> None:
"""Reset the cached importer registries (for testing)."""
@@ -529,8 +572,12 @@ class DagImporterRegistry:
existing_name,
new_name,
)
- self._extension_importers.pop(ext, None)
+ evicted = self._extension_importers.pop(ext, None)
self._extension_specs.pop(ext, None)
+ # An importer left with no extension would still list, and
duplicate, the new owner's
+ # definitions: a composite's members are not filtered by extension
ownership.
+ if evicted in self._ordered_importers and evicted not in
self._extension_importers.values():
+ self._ordered_importers.remove(evicted)
@staticmethod
def _get_bundle_importers_config(bundle_name: str) -> list[dict[str, Any]]
| None:
diff --git a/task-sdk/src/airflow/sdk/importers/python_importer.py
b/task-sdk/src/airflow/sdk/importers/python_importer.py
index f6448c12645..74d82ff0024 100644
--- a/task-sdk/src/airflow/sdk/importers/python_importer.py
+++ b/task-sdk/src/airflow/sdk/importers/python_importer.py
@@ -147,18 +147,18 @@ class
PythonDagImporter(AbstractDagImporter[FileDagDefinition]):
safe_mode: bool = True,
) -> Iterator[FileDagDefinition | DagImportError]:
"""
- List Python DAG files in a bundle matching supported extensions.
+ List Python Dag files under the bundle path matching supported
extensions.
A lightweight content sniff (``might_contain_dag``) is applied here so
files that
- clearly hold no DAG never become definitions -- keeping the discovered
set (and the
- eventual parse-process count) close to the number of real DAG files.
Zip members are
+ clearly hold no Dag never become definitions -- keeping the discovered
set (and the
+ eventual parse-process count) close to the number of real Dag files.
Zip members are
discovered by :class:`..zip_importer.ZipImporter`, not here.
"""
- if not bundle.path.is_dir():
- return
for definition in find_file_dag_definitions(bundle.path,
self.supported_extensions):
if self.might_contain_dag(definition, safe_mode):
yield definition
+ else:
+ log.debug("Skipping %r: no Airflow Dag markers found",
definition)
def import_definition(
self,
diff --git a/task-sdk/tests/task_sdk/importers/test_registry.py
b/task-sdk/tests/task_sdk/importers/test_registry.py
index 20a6b78a337..87e80d456a6 100644
--- a/task-sdk/tests/task_sdk/importers/test_registry.py
+++ b/task-sdk/tests/task_sdk/importers/test_registry.py
@@ -20,7 +20,9 @@ from __future__ import annotations
import json
import logging
+import zipfile
from pathlib import Path
+from types import SimpleNamespace
import pytest
@@ -45,6 +47,10 @@ from airflow.sdk.importers import (
from tests_common.test_utils.config import conf_vars
+def _bundle(path: Path) -> SimpleNamespace:
+ return SimpleNamespace(name="test", path=path)
+
+
class GlobalDagImporter(PythonDagImporter):
pass
@@ -219,6 +225,77 @@ class TestDagImporterRegistry:
for record in caplog.records
)
+ def test_list_dag_definitions_with_configured_extensions(self, tmp_path):
+ """Discovery only yields files matching the explicitly configured
extensions."""
+ dag_file = tmp_path / "valid_dag.custom"
+ dag_file.write_text("from airflow.sdk import DAG\n")
+ (tmp_path / "ignored.py").write_text("from airflow.sdk import DAG\n")
+ (tmp_path / "notes.txt").write_text("not a dag")
+
+ registry = DagImporterRegistry(register_defaults=False)
+ registry.register(PythonDagImporter(), extensions=[".custom"])
+
+ assert [str(d.path) for _, d in
registry.list_dag_definitions(_bundle(tmp_path))] == [str(dag_file)]
+
+ def test_list_dag_definitions_routes_to_overriding_importer(self,
tmp_path):
+ """A discovered file is routed to the importer that overrode its
extension."""
+ py_file = tmp_path / "dag.py"
+ py_file.write_text("from airflow.sdk import DAG\n")
+
+ registry = DagImporterRegistry(register_defaults=True)
+ overriding_importer = GlobalDagImporter()
+ registry.register(overriding_importer, extensions=[".py"])
+
+ assert list(registry.list_dag_definitions(_bundle(tmp_path))) == [
+ (overriding_importer, FilesystemDagDefinition(path=py_file))
+ ]
+
+ def test_list_dag_definitions_partial_override_is_listed_once(self,
tmp_path):
+ """An importer overriding only some of another importer's extensions
takes just those files."""
+ py_file = tmp_path / "dag.py"
+ pyc_file = tmp_path / "compiled.pyc"
+ py_file.write_text("from airflow.sdk import DAG\n")
+ pyc_file.write_bytes(b"compiled")
+
+ registry = DagImporterRegistry(register_defaults=False)
+ python_importer = PythonDagImporter()
+ registry.register(python_importer, extensions=[".py", ".pyc"])
+ pyc_importer = GlobalDagImporter()
+ registry.register(pyc_importer, extensions=[".pyc"])
+
+ assert sorted(
+ (type(importer).__name__, d.path.name)
+ for importer, d in
registry.list_dag_definitions(_bundle(tmp_path), safe_mode=False)
+ ) == [("GlobalDagImporter", "compiled.pyc"), ("PythonDagImporter",
"dag.py")]
+
+ def test_list_dag_definitions_yields_members_of_a_composite_importer(self,
tmp_path):
+ """A member discovered by ZipImporter is yielded even though it
imports via PythonDagImporter."""
+ archive = tmp_path / "archive.zip"
+ with zipfile.ZipFile(archive, "w") as zf:
+ zf.writestr("member_dag.py", "from airflow.sdk import DAG\n")
+
+ registry = DagImporterRegistry(register_defaults=True)
+
+ [(importer, definition)] =
registry.list_dag_definitions(_bundle(tmp_path))
+
+ assert repr(definition) == str(archive / "member_dag.py")
+ assert isinstance(importer, ZipImporter)
+
+ def
test_list_dag_definitions_materialises_a_multi_extension_spec_once(self,
tmp_path):
+ archive = tmp_path / "archive.zip"
+ with zipfile.ZipFile(archive, "w") as zf:
+ zf.writestr("member_dag.py", "from airflow.sdk import DAG\n")
+ registry = DagImporterRegistry(register_defaults=False)
+ registry.register_specs(
+ [{"classpath": "airflow.sdk.importers.ZipImporter", "extensions":
[".zip", ".zipx"]}],
+ context="test",
+ )
+
+ [(importer, definition)] =
registry.list_dag_definitions(_bundle(tmp_path))
+
+ assert repr(definition) == str(archive / "member_dag.py")
+ assert registry.get_importer("dags.zipx") is importer
+
def test_lazy_importer_instantiation(self):
"""Importer classes are not imported or instantiated until
get_importer is called."""
LazyTestImporter.instances = 0
@@ -284,6 +361,37 @@ class TestDagImporterRegistry:
assert isinstance(bundle_reg.get_importer("dag.py"),
GlobalDagImporter)
assert isinstance(bundle_reg.get_importer("dag.custom"),
BundleDagImporter)
+ @conf_vars(
+ {
+ ("dag_processor", "dag_bundle_config_list"): json.dumps(
+ [
+ {
+ "name": "test_bundle",
+ "importers": [
+ {"classpath": "airflow.sdk.importers.ZipImporter",
"extensions": [".zip"]}
+ ],
+ }
+ ]
+ )
+ }
+ )
+ def test_from_config_bundle_override_of_zip_lists_members_once(self,
tmp_path):
+ with zipfile.ZipFile(tmp_path / "dags.zip", "w") as zf:
+ zf.writestr("member.py", "from airflow.sdk import DAG\n")
+ registry = DagImporterRegistry.from_config("test_bundle")
+
+ [(importer, definition)] =
registry.list_dag_definitions(_bundle(tmp_path))
+
+ assert repr(definition) == str(tmp_path / "dags.zip" / "member.py")
+ assert importer is registry.get_importer("dags.zip")
+
+ @conf_vars({("dag_processor", "dag_importer_configs"):
json.dumps([f"{__name__}.GlobalDagImporter"])})
+ def test_from_config_accepts_a_bare_classpath_string(self):
+ """An importer that needs no extensions or kwargs can be configured as
a plain string."""
+ registry = DagImporterRegistry.from_config()
+
+ assert isinstance(registry.get_importer("dag.py"), GlobalDagImporter)
+
@pytest.mark.parametrize(
("global_cfg", "bundle_cfg", "match"),
[
@@ -461,3 +569,34 @@ class TestDagImporterRegistry:
definitions = list(find_file_dag_definitions(tmp_path, [".py",
".pyc"]))
assert {d.path.name for d in definitions} == {"workflow.PY"}
+
+ @pytest.mark.parametrize(
+ ("filename", "expected"),
+ [("workflow.py", ["workflow.py"]), ("notes.txt", [])],
+ )
+ def test_find_file_dag_definitions_bundle_pointing_at_a_single_file(self,
tmp_path, filename, expected):
+ single_file = tmp_path / filename
+ single_file.write_text("from airflow.sdk import DAG\n")
+
+ definitions = list(find_file_dag_definitions(single_file, [".py"]))
+ assert [d.path.name for d in definitions] == expected
+
+ @pytest.mark.parametrize(("safe_mode", "expected"), [(True, []), (False,
["no_markers.py"])])
+ def test_list_dag_definitions_forwards_safe_mode(self, tmp_path,
safe_mode, expected):
+ (tmp_path / "no_markers.py").write_text("print('hello')\n")
+
+ registry = DagImporterRegistry(register_defaults=False)
+ registry.register(PythonDagImporter())
+
+ definitions = registry.list_dag_definitions(_bundle(tmp_path),
safe_mode=safe_mode)
+ assert [d.path.name for _, d in definitions] == expected
+
+ def test_list_dag_definitions_passes_discovery_errors_through(self,
tmp_path):
+ (tmp_path / "corrupt.zip").write_bytes(b"not a zip")
+
+ registry = DagImporterRegistry(register_defaults=False)
+ registry.register(ZipImporter())
+
+ items = [item for _, item in
registry.list_dag_definitions(_bundle(tmp_path))]
+ assert [type(item) for item in items] == [DagImportError]
+ assert items[0].error_type == "zip_read_error"
