This is an automated email from the ASF dual-hosted git repository.
jason810496 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 1f5d7691663 [AIP-85] Add Dag importers to Task SDK in preparation for
rewiring (#72369)
1f5d7691663 is described below
commit 1f5d7691663016c898718768991a2dae522b4475
Author: Dilnaz Amanzholova <[email protected]>
AuthorDate: Mon Sep 14 03:45:40 2026 +0200
[AIP-85] Add Dag importers to Task SDK in preparation for rewiring (#72369)
* feat: Restructure the AbstractDagImporter/DagImporterRegistry and move
the /importers base classes into the task SDK
* fix: run the checks
* Decouple Dag importers from Dag bundles in Task SDK
Pass bundle name and path as primitive parameters to importer methods
instead of
requiring a BaseDagBundle instance. This effort saves us from passing whole
bundle object into the DagBag for the import_definition() method.
* Fix core Airflow import violations in shared libraries and Task SDK Dag
importers
* Harden ZipImporter against path traversal. Remove core dependencies from
Task SDK DAG importers and standardize exceptions. Clean up test boilerplate
and redundant checks.
* Pass Dag bundles directly to Task SDK Dag importers
* Reimplement Zip import def to extract on-demand
* Simplify ZipImporter's get_source_code
The whole-zip-file concat logic is removed since we are not supposed to
treat
the entire zip file as one dag source. Airflow's mental model is to treat a
zip
more like a directory, and only one member in the zip is the source file of
a
given dag. This tradition should be kept. Viewing an entire zip should
instead
be handled as a feature on the dag bundle level.
* Improve Dag importer resilience, discovery filtering, and bytecode support
* Pass conf explicitly to might_contain_dag instead of frame introspection
Resolving which config module to read (airflow.configuration vs
airflow.sdk.configuration) by walking the call stack only worked by
coincidence, since every current caller's immediate frame already sits
in an airflow-prefixed module. It silently breaks for anything else in
the chain (tests, wrapped calls) and re-walks the stack on every file
scanned in safe mode. Every caller already has the right conf object
in scope, so require it explicitly instead.
---------
Co-authored-by: Tzu-ping Chung <[email protected]>
Co-authored-by: ZHE YOU LIU <[email protected]>
---
.../src/airflow/dag_processing/importers/base.py | 2 +-
.../dag_processing/importers/python_importer.py | 6 +-
airflow-core/src/airflow/dag_processing/manager.py | 2 +-
airflow-core/src/airflow/utils/file.py | 59 +--
airflow-core/tests/unit/utils/test_file.py | 7 +-
.../src/airflow_shared/module_loading/__init__.py | 3 +
.../src/airflow_shared/module_loading/dag_file.py | 84 ++++
.../tests/module_loading/test_dag_file.py | 52 ++-
task-sdk/src/airflow/sdk/importers/__init__.py | 54 +++
task-sdk/src/airflow/sdk/importers/base.py | 482 +++++++++++++++++++++
.../src/airflow/sdk/importers/python_importer.py | 297 +++++++++++++
task-sdk/src/airflow/sdk/importers/zip_importer.py | 270 ++++++++++++
task-sdk/tests/task_sdk/docs/test_public_api.py | 1 +
.../tests/task_sdk/importers/__init__.py | 7 -
.../task_sdk/importers/test_python_importer.py | 236 ++++++++++
task-sdk/tests/task_sdk/importers/test_registry.py | 432 ++++++++++++++++++
.../tests/task_sdk/importers/test_zip_importer.py | 237 ++++++++++
17 files changed, 2162 insertions(+), 69 deletions(-)
diff --git a/airflow-core/src/airflow/dag_processing/importers/base.py
b/airflow-core/src/airflow/dag_processing/importers/base.py
index 643c16611f9..036dfab98d7 100644
--- a/airflow-core/src/airflow/dag_processing/importers/base.py
+++ b/airflow-core/src/airflow/dag_processing/importers/base.py
@@ -155,7 +155,7 @@ class AbstractDagImporter(ABC):
continue
# Apply safe_mode heuristic if enabled
- if safe_mode and not might_contain_dag(file_path, safe_mode):
+ if safe_mode and not might_contain_dag(file_path, safe_mode,
conf=conf):
continue
yield file_path
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 47f77a86f73..2eedc4dff4d 100644
--- a/airflow-core/src/airflow/dag_processing/importers/python_importer.py
+++ b/airflow-core/src/airflow/dag_processing/importers/python_importer.py
@@ -112,7 +112,7 @@ class PythonDagImporter(AbstractDagImporter):
path = Path(file_path)
try:
if path.is_file() and (path.suffix.lower() == ".py" or
zipfile.is_zipfile(path)):
- if might_contain_dag(file_path, safe_mode):
+ if might_contain_dag(file_path, safe_mode, conf=conf):
yield file_path
except Exception:
log.exception("Error while examining %s", file_path)
@@ -216,7 +216,7 @@ class PythonDagImporter(AbstractDagImporter):
except ValueError:
log.warning("SIGSEGV signal handler registration failed. Not in
the main thread")
- if not might_contain_dag(filepath, safe_mode):
+ if not might_contain_dag(filepath, safe_mode, conf=conf):
log.debug("File %s assumed to contain no DAGs. Skipping.",
filepath)
result.skipped_files.append(filepath)
return []
@@ -295,7 +295,7 @@ class PythonDagImporter(AbstractDagImporter):
log.debug("Reading %s from %s", zip_info.filename, filepath)
- if not might_contain_dag(zip_info.filename, safe_mode,
current_zip_file):
+ 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
diff --git a/airflow-core/src/airflow/dag_processing/manager.py
b/airflow-core/src/airflow/dag_processing/manager.py
index 80bbf40d441..ee71d86a9a8 100644
--- a/airflow-core/src/airflow/dag_processing/manager.py
+++ b/airflow-core/src/airflow/dag_processing/manager.py
@@ -988,7 +988,7 @@ class DagFileProcessorManager(LoggingMixin):
with zipfile.ZipFile(abs_path) as z:
for info in z.infolist():
# Use the configured discovery safe mode
- if might_contain_dag(info.filename,
self.dag_discovery_safe_mode, z):
+ if might_contain_dag(info.filename,
self.dag_discovery_safe_mode, z, conf=conf):
yield os.path.join(abs_path, info.filename)
except zipfile.BadZipFile:
self.log.exception("There was an error accessing ZIP file %s",
abs_path)
diff --git a/airflow-core/src/airflow/utils/file.py
b/airflow-core/src/airflow/utils/file.py
index feeaa5239c3..c762fffaf9f 100644
--- a/airflow-core/src/airflow/utils/file.py
+++ b/airflow-core/src/airflow/utils/file.py
@@ -18,7 +18,6 @@
from __future__ import annotations
import ast
-import hashlib
import logging
import os
import re
@@ -28,7 +27,11 @@ from io import TextIOWrapper
from pathlib import Path
from typing import overload
-from airflow._shared.module_loading import MODIFIED_DAG_MODULE_NAME
+from airflow._shared.module_loading import (
+ get_unique_dag_module_name as get_unique_dag_module_name,
+ might_contain_dag as might_contain_dag,
+ might_contain_dag_via_default_heuristic as
might_contain_dag_via_default_heuristic,
+)
from airflow.configuration import conf
log = logging.getLogger(__name__)
@@ -108,7 +111,7 @@ def find_dag_file_paths(directory: str | os.PathLike[str],
safe_mode: bool) -> l
path = Path(file_path)
try:
if path.is_file() and (path.suffix == ".py" or
zipfile.is_zipfile(path)):
- if might_contain_dag(file_path, safe_mode):
+ if might_contain_dag(file_path, safe_mode, conf=conf):
file_paths.append(file_path)
except Exception:
log.exception("Error while examining %s", file_path)
@@ -119,47 +122,6 @@ def find_dag_file_paths(directory: str | os.PathLike[str],
safe_mode: bool) -> l
COMMENT_PATTERN = re.compile(r"\s*#.*")
-def might_contain_dag(file_path: str, safe_mode: bool, zip_file:
zipfile.ZipFile | None = None) -> bool:
- """
- Check whether a Python file contains Airflow DAGs.
-
- When safe_mode is off (with False value), this function always returns
True.
-
- If might_contain_dag_callable isn't specified, it uses airflow default
heuristic
- """
- if not safe_mode:
- return True
-
- might_contain_dag_callable = conf.getimport(
- "core",
- "might_contain_dag_callable",
- fallback="airflow.utils.file.might_contain_dag_via_default_heuristic",
- )
- return might_contain_dag_callable(file_path=file_path, zip_file=zip_file)
-
-
-def might_contain_dag_via_default_heuristic(file_path: str, zip_file:
zipfile.ZipFile | None = None) -> bool:
- """
- Heuristic that guesses whether a Python file contains an Airflow DAG
definition.
-
- :param file_path: Path to the file to be checked.
- :param zip_file: if passed, checks the archive. Otherwise, check local
filesystem.
- :return: True, if file might contain DAGs.
- """
- if zip_file:
- with zip_file.open(file_path) as current_file:
- content = current_file.read()
- else:
- if zipfile.is_zipfile(file_path):
- return True
- with open(file_path, "rb") as dag_file:
- content = dag_file.read()
- content = content.lower()
- if b"airflow" not in content:
- return False
- return any(s in content for s in (b"dag", b"asset"))
-
-
def _find_imported_modules(module: ast.Module) -> Generator[str, None, None]:
for st in module.body:
if isinstance(st, ast.Import):
@@ -180,15 +142,6 @@ def iter_airflow_imports(file_path: str) -> Generator[str,
None, None]:
yield m
-def get_unique_dag_module_name(file_path: str) -> str:
- """Return a unique module name in the format unusual_prefix_{sha1 of
module's file path}_{original module name}."""
- if isinstance(file_path, str):
- path_hash = hashlib.sha1(file_path.encode("utf-8"),
usedforsecurity=False).hexdigest()
- org_mod_name = re.sub(r"[.-]", "_", Path(file_path).stem)
- return MODIFIED_DAG_MODULE_NAME.format(path_hash=path_hash,
module_name=org_mod_name)
- raise ValueError("file_path should be a string to generate unique module
name")
-
-
def __getattr__(name: str):
if name == "find_path_from_directory":
import warnings
diff --git a/airflow-core/tests/unit/utils/test_file.py
b/airflow-core/tests/unit/utils/test_file.py
index cc55c1ac063..11502777128 100644
--- a/airflow-core/tests/unit/utils/test_file.py
+++ b/airflow-core/tests/unit/utils/test_file.py
@@ -25,6 +25,7 @@ from unittest import mock
import pytest
from airflow._shared.module_loading import find_path_from_directory
+from airflow.configuration import conf
from airflow.utils import file as file_utils
from airflow.utils.file import (
correct_maybe_zipped,
@@ -139,7 +140,7 @@ class TestListPyFilesPath:
def test_might_contain_dag_with_default_callable(self):
file_path_with_dag = os.path.join(TEST_DAGS_FOLDER,
"test_scheduler_dags.py")
- assert file_utils.might_contain_dag(file_path=file_path_with_dag,
safe_mode=True)
+ assert file_utils.might_contain_dag(file_path=file_path_with_dag,
safe_mode=True, conf=conf)
@conf_vars({("core", "might_contain_dag_callable"):
"unit.utils.test_file.might_contain_dag"})
def test_might_contain_dag(self):
@@ -149,10 +150,10 @@ class TestListPyFilesPath:
# There is a DAG defined in the file_path_with_dag, however, the
might_contain_dag_callable
# returns False no matter what, which is used to test
might_contain_dag_callable actually
# overrides the default function
- assert not file_utils.might_contain_dag(file_path=file_path_with_dag,
safe_mode=True)
+ assert not file_utils.might_contain_dag(file_path=file_path_with_dag,
safe_mode=True, conf=conf)
# With safe_mode is False, the user defined callable won't be invoked
- assert file_utils.might_contain_dag(file_path=file_path_with_dag,
safe_mode=False)
+ assert file_utils.might_contain_dag(file_path=file_path_with_dag,
safe_mode=False, conf=conf)
def test_get_modules(self):
file_path = os.path.join(TEST_DAGS_FOLDER, "test_imports.py")
diff --git
a/shared/module_loading/src/airflow_shared/module_loading/__init__.py
b/shared/module_loading/src/airflow_shared/module_loading/__init__.py
index 238506ae52d..3bf8eacab8e 100644
--- a/shared/module_loading/src/airflow_shared/module_loading/__init__.py
+++ b/shared/module_loading/src/airflow_shared/module_loading/__init__.py
@@ -30,6 +30,9 @@ from typing import TYPE_CHECKING
from .dag_file import (
MODIFIED_DAG_MODULE_NAME as MODIFIED_DAG_MODULE_NAME,
UNUSUAL_MODULE_PREFIX as UNUSUAL_MODULE_PREFIX,
+ get_unique_dag_module_name as get_unique_dag_module_name,
+ might_contain_dag as might_contain_dag,
+ might_contain_dag_via_default_heuristic as
might_contain_dag_via_default_heuristic,
)
from .file_discovery import (
find_path_from_directory as find_path_from_directory,
diff --git
a/shared/module_loading/src/airflow_shared/module_loading/dag_file.py
b/shared/module_loading/src/airflow_shared/module_loading/dag_file.py
index d4fdc80c737..c4d95db8ba5 100644
--- a/shared/module_loading/src/airflow_shared/module_loading/dag_file.py
+++ b/shared/module_loading/src/airflow_shared/module_loading/dag_file.py
@@ -19,5 +19,89 @@
from __future__ import annotations
+import hashlib
+import re
+import zipfile
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
UNUSUAL_MODULE_PREFIX = "unusual_prefix_"
MODIFIED_DAG_MODULE_NAME =
f"{UNUSUAL_MODULE_PREFIX}{{path_hash}}_{{module_name}}"
+
+if TYPE_CHECKING:
+ from typing import Protocol
+
+ class _MightContainDagCallable(Protocol):
+ def __call__(self, file_path: str, zip_file: zipfile.ZipFile | None =
None) -> bool: ...
+
+ class _ConfLike(Protocol):
+ def getimport(self, section: str, key: str, **kwargs: Any) -> Any: ...
+
+
+def get_unique_dag_module_name(file_path: str) -> str:
+ """Return a unique module name in the format unusual_prefix_{sha1 of
module's file path}_{original module name}."""
+ if isinstance(file_path, str):
+ path_hash = hashlib.sha1(file_path.encode("utf-8"),
usedforsecurity=False).hexdigest()
+ org_mod_name = re.sub(r"[.-]", "_", Path(file_path).stem)
+ return MODIFIED_DAG_MODULE_NAME.format(path_hash=path_hash,
module_name=org_mod_name)
+ raise ValueError("file_path should be a string to generate unique module
name")
+
+
+def might_contain_dag_via_default_heuristic(file_path: str, zip_file:
zipfile.ZipFile | None = None) -> bool:
+ """
+ Heuristic that guesses whether a Python file contains an Airflow DAG
definition.
+
+ :param file_path: Path to the file to be checked.
+ :param zip_file: if passed, checks the archive. Otherwise, check local
filesystem.
+ :return: True, if file might contain DAGs.
+ """
+ if zip_file:
+ with zip_file.open(file_path) as current_file:
+ content = current_file.read()
+ else:
+ if zipfile.is_zipfile(file_path):
+ return True
+ with open(file_path, "rb") as dag_file:
+ content = dag_file.read()
+ content = content.lower()
+ if b"airflow" not in content:
+ return False
+ return any(s in content for s in (b"dag", b"asset"))
+
+
+def might_contain_dag(
+ file_path: str,
+ safe_mode: bool,
+ zip_file: zipfile.ZipFile | None = None,
+ *,
+ conf: _ConfLike,
+) -> bool:
+ """
+ Check whether a Python file contains Airflow DAGs.
+
+ When safe_mode is off (with False value), this function always returns
True.
+
+ If might_contain_dag_callable isn't specified, it uses airflow default
heuristic.
+ """
+ if not safe_mode:
+ return True
+
+ might_contain_dag_callable: _MightContainDagCallable | None = None
+ try:
+ might_contain_dag_callable = conf.getimport(
+ "core",
+ "might_contain_dag_callable",
+ fallback=None,
+ )
+ except Exception as e:
+ import logging
+
+ logging.getLogger(__name__).warning(
+ "Failed to load might_contain_dag_callable from config, falling
back to default heuristic: %s",
+ e,
+ )
+
+ if might_contain_dag_callable is None:
+ might_contain_dag_callable = might_contain_dag_via_default_heuristic
+
+ return might_contain_dag_callable(file_path=file_path, zip_file=zip_file)
diff --git a/shared/module_loading/tests/module_loading/test_dag_file.py
b/shared/module_loading/tests/module_loading/test_dag_file.py
index c7319384a93..f802362f43c 100644
--- a/shared/module_loading/tests/module_loading/test_dag_file.py
+++ b/shared/module_loading/tests/module_loading/test_dag_file.py
@@ -17,10 +17,60 @@
# under the License.
from __future__ import annotations
-from airflow_shared.module_loading import MODIFIED_DAG_MODULE_NAME,
UNUSUAL_MODULE_PREFIX
+import logging
+from unittest import mock
+
+from airflow_shared.module_loading import (
+ MODIFIED_DAG_MODULE_NAME,
+ UNUSUAL_MODULE_PREFIX,
+ get_unique_dag_module_name,
+ might_contain_dag,
+)
def test_constants() -> None:
"""Test that the constants are as expected."""
assert UNUSUAL_MODULE_PREFIX == "unusual_prefix_"
assert MODIFIED_DAG_MODULE_NAME ==
"unusual_prefix_{path_hash}_{module_name}"
+
+
+def test_get_unique_dag_module_name() -> None:
+ mod_name = get_unique_dag_module_name("/path/to/my_dag.py")
+ assert mod_name.startswith("unusual_prefix_")
+ assert mod_name.endswith("_my_dag")
+
+
+def test_might_contain_dag(tmp_path) -> None:
+ mock_conf = mock.MagicMock()
+ mock_conf.getimport.return_value = None
+
+ dag_file = tmp_path / "test_dag.py"
+ dag_file.write_text("from airflow import DAG\ndag = DAG('test')")
+ assert might_contain_dag(str(dag_file), safe_mode=True, conf=mock_conf) is
True
+
+ non_dag_file = tmp_path / "helper.py"
+ non_dag_file.write_text("def add(x, y): return x + y")
+ assert might_contain_dag(str(non_dag_file), safe_mode=True,
conf=mock_conf) is False
+ assert might_contain_dag(str(non_dag_file), safe_mode=False,
conf=mock_conf) is True
+
+
+def test_might_contain_dag_with_explicit_conf() -> None:
+ mock_conf = mock.MagicMock()
+ mock_conf.getimport.return_value = lambda file_path, zip_file=None: False
+
+ assert might_contain_dag("sample.py", safe_mode=True, conf=mock_conf) is
False
+ mock_conf.getimport.assert_called_once_with("core",
"might_contain_dag_callable", fallback=None)
+
+
+def test_might_contain_dag_logs_warning_on_broken_config(tmp_path, caplog) ->
None:
+ dag_file = tmp_path / "test_dag.py"
+ dag_file.write_text("from airflow import DAG\ndag = DAG('test')")
+
+ mock_conf = mock.MagicMock()
+ mock_conf.getimport.side_effect = ImportError("No module named
'broken_module'")
+
+ with caplog.at_level(logging.WARNING):
+ result = might_contain_dag(str(dag_file), safe_mode=True,
conf=mock_conf)
+
+ assert result is True
+ assert "Failed to load might_contain_dag_callable from config" in
caplog.text
diff --git a/task-sdk/src/airflow/sdk/importers/__init__.py
b/task-sdk/src/airflow/sdk/importers/__init__.py
new file mode 100644
index 00000000000..f424fbf56ee
--- /dev/null
+++ b/task-sdk/src/airflow/sdk/importers/__init__.py
@@ -0,0 +1,54 @@
+# 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.
+"""DAG Importer - pluggable mechanism for importing DAGs from different file
formats."""
+
+from __future__ import annotations
+
+from airflow.sdk.importers.base import (
+ AbstractDagImporter,
+ DagDefinition,
+ DagImporterRegistry,
+ DagImportError,
+ DagImportResult,
+ DagImportWarning,
+ DagSourceCode,
+ FileDagDefinition,
+ find_file_dag_definitions,
+ get_file_suffix,
+ get_importer_registry,
+ reset_importer_registry,
+)
+from airflow.sdk.importers.python_importer import PythonDagImporter
+from airflow.sdk.importers.zip_importer import ZipFileDagDefinition,
ZipImporter
+
+__all__ = [
+ "AbstractDagImporter",
+ "DagDefinition",
+ "DagImportError",
+ "DagImportResult",
+ "DagImportWarning",
+ "DagImporterRegistry",
+ "DagSourceCode",
+ "FileDagDefinition",
+ "PythonDagImporter",
+ "ZipFileDagDefinition",
+ "ZipImporter",
+ "find_file_dag_definitions",
+ "get_file_suffix",
+ "get_importer_registry",
+ "reset_importer_registry",
+]
diff --git a/task-sdk/src/airflow/sdk/importers/base.py
b/task-sdk/src/airflow/sdk/importers/base.py
new file mode 100644
index 00000000000..9c07892a526
--- /dev/null
+++ b/task-sdk/src/airflow/sdk/importers/base.py
@@ -0,0 +1,482 @@
+# 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.
+"""Abstract base class for DAG importers."""
+
+from __future__ import annotations
+
+import contextlib
+import functools
+import logging
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from airflow.sdk._shared.module_loading.dag_file import might_contain_dag
+from airflow.sdk._shared.module_loading.file_discovery import
find_path_from_directory
+from airflow.sdk.configuration import conf
+from airflow.sdk.exceptions import AirflowConfigException
+
+if TYPE_CHECKING:
+ from collections.abc import Generator, Iterable, Iterator
+
+ from typing_extensions import Self
+
+ from airflow.dag_processing.bundles.base import BaseDagBundle # noqa:
SDK002
+ from airflow.sdk import DAG
+
+log = logging.getLogger(__name__)
+
+
+class DagDefinition(ABC):
+ """Abstract base class for a DAG source definition."""
+
+ @property
+ @abstractmethod
+ def freshness_token(self) -> str:
+ """Opaque, generalized token representing the current state of the
source."""
+
+ @abstractmethod
+ def get_relative_loc(self, root: Path | None = None) -> str:
+ """Get relative location of the definition to a root directory."""
+
+ @abstractmethod
+ def read_bytes(self) -> bytes:
+ """Read and return the content of the resource as bytes."""
+
+ def read_text(self, encoding: str = "utf-8") -> str:
+ """Read and return the content of the resource as a string."""
+ return self.read_bytes().decode(encoding)
+
+ @abstractmethod
+ def as_file(self) -> contextlib.AbstractContextManager[Path]:
+ """
+ Return a context manager yielding a Path pointing to a local file.
+
+ For file-backed resources, this is the actual file path.
+ For others, a temp file is created and cleaned up.
+ """
+
+ @abstractmethod
+ def __repr__(self) -> str:
+ """Return string representation used by import error and warning
objects."""
+
+
+@dataclass
+class FileDagDefinition(DagDefinition):
+ """A DAG definition backed by a file on the local filesystem."""
+
+ path: Path
+
+ @property
+ def freshness_token(self) -> str:
+ try:
+ stat = self.path.stat()
+ return f"{stat.st_mtime_ns}-{stat.st_size}"
+ except OSError:
+ return ""
+
+ def get_relative_loc(self, root: Path | None = None) -> str:
+ if root is None:
+ return str(self.path)
+ try:
+ return str(self.path.relative_to(root))
+ except ValueError:
+ return str(self.path)
+
+ def read_bytes(self) -> bytes:
+ return self.path.read_bytes()
+
+ @contextlib.contextmanager
+ def as_file(self) -> Generator[Path, None, None]:
+ yield self.path
+
+ def __repr__(self) -> str:
+ return str(self.path)
+
+
+@dataclass
+class DagImportError:
+ """Structured error information for DAG import failures."""
+
+ source_reference: str
+ message: str
+ error_type: str = "import"
+ line_number: int | None = None
+ column_number: int | None = None
+ context: str | None = None
+ suggestion: str | None = None
+ stacktrace: str | None = None
+
+ def format_message(self) -> str:
+ """Format the error as a human-readable single-line string."""
+ loc_parts = []
+ if self.line_number is not None:
+ loc_parts.append(f"line {self.line_number}")
+ if self.column_number is not None:
+ loc_parts.append(f"column {self.column_number}")
+ loc_str = f" ({', '.join(loc_parts)})" if loc_parts else ""
+
+ parts = [f"Error in {self.source_reference}{loc_str}
[{self.error_type}]: {self.message.strip()}"]
+ if self.context:
+ parts.append(f"Context: {' '.join(self.context.split())}")
+ if self.suggestion:
+ parts.append(f"Suggestion: {self.suggestion.strip()}")
+ return "; ".join(parts)
+
+
+@dataclass
+class DagImportWarning:
+ """Warning information for non-fatal issues during DAG import."""
+
+ source_reference: str
+ message: str
+ warning_type: str = "import"
+ line_number: int | None = None
+ context: dict[str, Any] | None = None
+
+
+@dataclass
+class DagImportResult:
+ """Result of importing DAGs from a definition."""
+
+ definition: DagDefinition | None = None
+ dags: list[DAG] = field(default_factory=list)
+ errors: list[DagImportError] = field(default_factory=list)
+ skipped_definitions: list[DagDefinition] = field(default_factory=list)
+ warnings: list[DagImportWarning] = field(default_factory=list)
+ dependencies: list[DagDefinition] = field(default_factory=list)
+
+ @property
+ def success(self) -> bool:
+ """Return True if no fatal errors occurred."""
+ return not self.errors
+
+
+@dataclass
+class DagSourceCode:
+ """Raw source code and its language identifier for a DAG definition."""
+
+ source_code: str
+ language: str
+
+
+def _normalize_extensions(extensions: Iterable[str]) -> list[str]:
+ """Normalize file extensions to lowercase with leading dot."""
+ return [ext.lower() if ext.startswith(".") else f".{ext.lower()}" for ext
in extensions]
+
+
+def _get_importer_extensions(importer: AbstractDagImporter) -> list[str]:
+ """Extract supported extensions from an importer via duck typing."""
+ exts = getattr(importer, "supported_extensions", None)
+ if callable(exts):
+ return _normalize_extensions(exts())
+ if exts is not None:
+ return _normalize_extensions(exts)
+ return []
+
+
+class AbstractDagImporter(ABC):
+ """Abstract base class for DAG importers."""
+
+ @abstractmethod
+ def can_handle(self, definition: DagDefinition | str | Path) -> bool:
+ """Check if this importer can handle the given definition."""
+
+ @abstractmethod
+ def list_dag_definitions(
+ self,
+ bundle: BaseDagBundle,
+ *,
+ safe_mode: bool = True,
+ ) -> Iterator[DagDefinition]:
+ """List DAG definitions in a bundle that this importer can handle."""
+
+ @abstractmethod
+ def import_definition(
+ self,
+ definition: DagDefinition,
+ bundle: BaseDagBundle,
+ *,
+ safe_mode: bool = True,
+ ) -> DagImportResult:
+ """Import DAGs from a DAG definition."""
+
+ @abstractmethod
+ def get_source_code(self, definition: DagDefinition) -> DagSourceCode:
+ """Retrieve the raw source code and its language identifier for the
specified DAG definition."""
+
+
+def get_file_suffix(definition: DagDefinition | str | Path) -> str | None:
+ """Extract lowercase file suffix from a definition, path, or filename."""
+ path = (
+ definition
+ if isinstance(definition, (str, Path))
+ else getattr(definition, "path", getattr(definition, "file_path",
None))
+ )
+ return Path(path).suffix.lower() if path else None
+
+
+def find_file_dag_definitions(
+ bundle_path: Path,
+ supported_extensions: Iterable[str],
+ safe_mode: bool = True,
+) -> Iterator[DagDefinition]:
+ """Find file DAG definitions in a bundle matching given extensions and
respecting .airflowignore."""
+ ignore_file_syntax = conf.get_mandatory_value("core",
"DAG_IGNORE_FILE_SYNTAX", fallback="glob")
+ supported_exts = _normalize_extensions(supported_extensions)
+
+ for file_path in find_path_from_directory(bundle_path, ".airflowignore",
ignore_file_syntax):
+ path = Path(file_path)
+
+ if not path.is_file():
+ continue
+
+ if path.suffix.lower() not in supported_exts:
+ continue
+
+ if safe_mode and not might_contain_dag(str(path), safe_mode,
conf=conf):
+ continue
+ yield FileDagDefinition(path=path)
+
+
+@dataclass(frozen=True)
+class _ImporterSpec:
+ """Declarative specification for a DAG importer."""
+
+ classpath: str
+ kwargs: dict[str, Any] = field(default_factory=dict)
+ extensions: list[str] | None = None
+ context: str = "importer configuration"
+
+
+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."
+ )
+ specs: list[_ImporterSpec] = []
+ for item in configs:
+ if not isinstance(item, dict):
+ raise AirflowConfigException(
+ f"Invalid importer configuration for {context}: each entry
must be a dictionary."
+ )
+ classpath = item.get("classpath")
+ if not classpath:
+ raise AirflowConfigException(
+ f"Missing required 'classpath' in importer configuration for
{context}."
+ )
+ kwargs = item.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")
+ if extensions is not None:
+ if not isinstance(extensions, list) or any(not isinstance(ext,
str) for ext in extensions):
+ raise AirflowConfigException(
+ f"Field 'extensions' must be a list of strings in importer
configuration for {context}."
+ )
+ extensions = _normalize_extensions(extensions)
+ specs.append(
+ _ImporterSpec(
+ classpath=classpath,
+ kwargs=kwargs,
+ extensions=extensions,
+ context=context,
+ )
+ )
+ return specs
+
+
+class DagImporterRegistry:
+ """
+ Registry for DAG importers. Manages importers by file extension and
generic definition.
+
+ Each file extension can only be handled by one importer at a time. If
multiple
+ importers claim the same extension, the last registered one wins and a
warning
+ is logged. The built-in PythonDagImporter handles .py and ZipImporter
handles .zip files.
+ """
+
+ _extension_importers: dict[str, AbstractDagImporter]
+ _extension_specs: dict[str, _ImporterSpec]
+ _ordered_importers: list[AbstractDagImporter]
+
+ def __init__(self, register_defaults: bool = True) -> None:
+ self._extension_importers = {}
+ self._extension_specs = {}
+ self._ordered_importers = []
+ if register_defaults:
+ self._register_default_importers()
+
+ @classmethod
+ def from_config(cls, bundle_name: str | None = None) -> Self:
+ """Create and configure a DagImporterRegistry with 3-tier
precedence."""
+ registry = cls(register_defaults=True)
+
+ global_importers = conf.getjson("dag_processor",
"dag_importer_configs", fallback=None)
+ if global_importers:
+ if not isinstance(global_importers, list):
+ raise AirflowConfigException(
+ "Section `dag_processor` key `dag_importer_configs` must
be a list "
+ f"but got {global_importers.__class__.__name__}"
+ )
+ registry.register_specs(global_importers, context="global
configuration")
+
+ if bundle_name:
+ bundle_importers = cls._get_bundle_importers_config(bundle_name)
+ if bundle_importers:
+ registry.register_specs(bundle_importers, context=f"bundle
'{bundle_name}'")
+
+ return registry
+
+ def register(self, importer: AbstractDagImporter, extensions: list[str] |
None = None) -> None:
+ """
+ Register an importer.
+
+ 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)
+
+ if extensions:
+ normalized_extensions = _normalize_extensions(extensions)
+ if hasattr(importer, "supported_extensions"):
+ with contextlib.suppress(AttributeError, TypeError):
+ importer.supported_extensions = normalized_extensions
+ for ext_lower in normalized_extensions:
+ self._warn_and_evict_extension(ext_lower,
type(importer).__name__)
+ self._extension_importers[ext_lower] = 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):
+ if spec.extensions is None:
+ self.register(self._instantiate_spec(spec))
+ continue
+
+ for ext_lower in spec.extensions:
+ self._warn_and_evict_extension(ext_lower, spec.classpath)
+ self._extension_specs[ext_lower] = spec
+
+ def get_importer(self, definition: DagDefinition | str | Path) ->
AbstractDagImporter | None:
+ """Get the appropriate importer for a definition or file, or None if
unsupported."""
+ suffix = get_file_suffix(definition)
+ 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
+
+ for importer in reversed(self._ordered_importers):
+ if importer.can_handle(definition):
+ return importer
+ return None
+
+ def can_handle(self, definition: DagDefinition | str | Path) -> bool:
+ """Check if any registered importer can handle this definition/file."""
+ suffix = get_file_suffix(definition)
+ if suffix and (suffix in self._extension_importers or suffix in
self._extension_specs):
+ return True
+ return any(importer.can_handle(definition) for importer in
reversed(self._ordered_importers))
+
+ def supported_extensions(self) -> list[str]:
+ """Return all registered file extensions."""
+ return sorted(set(self._extension_importers) |
set(self._extension_specs))
+
+ @classmethod
+ def reset(cls) -> None:
+ """Reset the cached importer registries (for testing)."""
+ reset_importer_registry()
+
+ def _register_default_importers(self) -> None:
+ from airflow.sdk.importers.python_importer import PythonDagImporter
+ from airflow.sdk.importers.zip_importer import ZipImporter
+
+ self.register(PythonDagImporter())
+ self.register(ZipImporter())
+
+ @staticmethod
+ def _instantiate_spec(spec: _ImporterSpec) -> AbstractDagImporter:
+ from airflow.sdk._shared.module_loading import import_string
+
+ try:
+ importer_class = import_string(spec.classpath)
+ importer = importer_class(**spec.kwargs)
+ except Exception as err:
+ raise AirflowConfigException(
+ f"Failed to load DAG importer '{spec.classpath}' for
{spec.context}: {err}"
+ ) from err
+
+ if not isinstance(importer, AbstractDagImporter):
+ raise AirflowConfigException(
+ f"Configured DAG importer {type(importer).__module__}."
+ f"{type(importer).__qualname__} for {spec.context} must
inherit "
+ "from AbstractDagImporter."
+ )
+
+ if spec.extensions is not None and hasattr(importer,
"supported_extensions"):
+ with contextlib.suppress(AttributeError, TypeError):
+ importer.supported_extensions = spec.extensions
+ return importer
+
+ def _warn_and_evict_extension(self, ext: str, new_name: str) -> None:
+ if ext in self._extension_importers or ext in self._extension_specs:
+ existing = self._extension_importers.get(ext)
+ existing_name = type(existing).__name__ if existing else
self._extension_specs[ext].classpath
+ log.warning(
+ "Extension '%s' already registered by %s, overriding with %s",
+ ext,
+ existing_name,
+ new_name,
+ )
+ self._extension_importers.pop(ext, None)
+ self._extension_specs.pop(ext, None)
+
+ @staticmethod
+ def _get_bundle_importers_config(bundle_name: str) -> list[dict[str, Any]]
| None:
+ """Retrieve importer configs for a specific bundle from
configuration."""
+ bundle_config_list = conf.getjson("dag_processor",
"dag_bundle_config_list", fallback=None)
+ if isinstance(bundle_config_list, list):
+ for item in bundle_config_list:
+ if isinstance(item, dict) and item.get("name") == bundle_name:
+ return item.get("importers")
+ return None
+
+
[email protected]
+def get_importer_registry(bundle_name: str | None = None) ->
DagImporterRegistry:
+ """Get the cached DagImporterRegistry instance for global or bundle
scope."""
+ return DagImporterRegistry.from_config(bundle_name=bundle_name)
+
+
+def reset_importer_registry() -> None:
+ """Reset cached importer registries."""
+ get_importer_registry.cache_clear()
diff --git a/task-sdk/src/airflow/sdk/importers/python_importer.py
b/task-sdk/src/airflow/sdk/importers/python_importer.py
new file mode 100644
index 00000000000..264cbd6b7a3
--- /dev/null
+++ b/task-sdk/src/airflow/sdk/importers/python_importer.py
@@ -0,0 +1,297 @@
+# 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.
+"""Python DAG importer - imports DAGs from Python files."""
+
+from __future__ import annotations
+
+import functools
+import importlib.machinery
+import importlib.util
+import logging
+import os
+import sys
+import traceback
+import warnings
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from airflow.sdk._shared.module_loading.dag_file import
get_unique_dag_module_name, might_contain_dag
+from airflow.sdk.configuration import conf
+from airflow.sdk.definitions._internal.contextmanager import DagContext
+from airflow.sdk.definitions.dag import DAG
+from airflow.sdk.exceptions import AirflowConfigException
+from airflow.sdk.execution_time.timeout import timeout
+from airflow.sdk.importers.base import (
+ AbstractDagImporter,
+ DagDefinition,
+ DagImportError,
+ DagImportResult,
+ DagImportWarning,
+ DagSourceCode,
+ _normalize_extensions,
+ find_file_dag_definitions,
+ get_file_suffix,
+)
+
+if TYPE_CHECKING:
+ from collections.abc import Iterator
+ from types import ModuleType
+
+ from airflow.dag_processing.bundles.base import BaseDagBundle # noqa:
SDK002
+
+log = logging.getLogger(__name__)
+
+
+class PythonDagImporter(AbstractDagImporter):
+ """
+ Importer for Python DAG files.
+
+ This is the default importer registered with the DagImporterRegistry. It
handles
+ .py files containing Python DAGs.
+ """
+
+ supported_extensions = [".py", ".pyc"]
+
+ def __init__(self, extensions: list[str] | None = None) -> None:
+ if extensions is not None:
+ self.supported_extensions = _normalize_extensions(extensions)
+
+ def can_handle(self, definition: DagDefinition | str | Path) -> bool:
+ """Check if this importer can handle the given definition based on
file extension."""
+ suffix = get_file_suffix(definition)
+ return suffix in self.supported_extensions if suffix else False
+
+ def list_dag_definitions(
+ self,
+ bundle: BaseDagBundle,
+ *,
+ safe_mode: bool = True,
+ ) -> Iterator[DagDefinition]:
+ """List Python DAG definitions in a bundle matching supported
extensions."""
+ yield from find_file_dag_definitions(bundle.path,
self.supported_extensions, safe_mode=safe_mode)
+
+ def import_definition(
+ self,
+ definition: DagDefinition,
+ bundle: BaseDagBundle,
+ *,
+ safe_mode: bool = True,
+ ) -> DagImportResult:
+ """
+ Import DAGs from a Python DAG definition.
+
+ :param definition: The definition to import from.
+ :param bundle: The DAG bundle containing the definition.
+ :param safe_mode: If True, skip files that don't appear to contain
DAGs.
+ :return: DagImportResult with imported DAGs and any errors.
+ """
+ result = DagImportResult(definition=definition)
+ DagContext.autoregistered_dags.clear()
+ captured_warnings: list[warnings.WarningMessage] = []
+
+ try:
+ with warnings.catch_warnings(record=True) as captured_warnings:
+ with definition.as_file() as local_path:
+ filepath = os.fspath(local_path)
+ modules = self._load_modules_from_file(
+ filepath,
+ safe_mode,
+ result,
+ bundle=bundle,
+ )
+ except AirflowConfigException:
+ # Configuration errors (e.g., invalid timeout type) should
propagate
+ raise
+ except Exception as e:
+ result.errors.append(
+ DagImportError(
+ source_reference=repr(definition),
+ message=str(e),
+ error_type="import",
+ stacktrace=traceback.format_exc(),
+ )
+ )
+ return result
+
+ for warn_msg in captured_warnings:
+ category = warn_msg.category.__name__
+ if (module := warn_msg.category.__module__) != "builtins":
+ category = f"{module}.{category}"
+ result.warnings.append(
+ DagImportWarning(
+ source_reference=repr(definition),
+ message=str(warn_msg.message),
+ warning_type=category,
+ line_number=warn_msg.lineno,
+ )
+ )
+
+ self._process_modules(
+ modules,
+ result,
+ bundle=bundle,
+ )
+
+ return result
+
+ def get_source_code(self, definition: DagDefinition) -> DagSourceCode:
+ """Retrieve the raw source code for the Python definition."""
+ if get_file_suffix(definition) == ".pyc":
+ return DagSourceCode(
+ source_code="# Sourceless bytecode (.pyc) — source code not
available\n",
+ language="python",
+ )
+ return DagSourceCode(
+ source_code=definition.read_text(encoding="utf-8"),
+ language="python",
+ )
+
+ def might_contain_dag(self, file_path: str | Path, safe_mode: bool = True)
-> bool:
+ """Check whether a file might contain Airflow DAGs according to safe
mode heuristics."""
+ if not safe_mode:
+ return True
+ return might_contain_dag(str(file_path), safe_mode, conf=conf)
+
+ def _load_modules_from_file(
+ self,
+ filepath: str,
+ safe_mode: bool,
+ result: DagImportResult,
+ bundle: BaseDagBundle,
+ ) -> list[ModuleType]:
+ definition = result.definition
+
+ import signal
+
+ def sigsegv_handler(signum, frame):
+ msg = f"Received SIGSEGV signal while processing {filepath}."
+ log.error(msg)
+ result.errors.append(
+ DagImportError(
+ source_reference=repr(definition),
+ message=msg,
+ error_type="segfault",
+ )
+ )
+
+ try:
+ signal.signal(signal.SIGSEGV, sigsegv_handler)
+ except (ValueError, AttributeError):
+ log.warning("SIGSEGV signal handler registration failed. Not in
the main thread")
+
+ if not self.might_contain_dag(filepath, safe_mode):
+ log.debug("File %s assumed to contain no DAGs. Skipping.",
filepath)
+ if definition is not None:
+ result.skipped_definitions.append(definition)
+ return []
+
+ log.debug("Importing %s (bundle: %s)", filepath, bundle.name)
+ mod_name = get_unique_dag_module_name(filepath)
+
+ if mod_name in sys.modules:
+ del sys.modules[mod_name]
+
+ DagContext.current_autoregister_module_name = mod_name
+
+ def parse(mod_name: str, filepath: str) -> list[ModuleType]:
+ try:
+ loader: importlib.machinery.SourceFileLoader |
importlib.machinery.SourcelessFileLoader
+ if Path(filepath).suffix.lower() == ".pyc":
+ loader =
importlib.machinery.SourcelessFileLoader(mod_name, filepath)
+ else:
+ loader = importlib.machinery.SourceFileLoader(mod_name,
filepath)
+ spec = importlib.util.spec_from_loader(mod_name, loader)
+ new_module = importlib.util.module_from_spec(spec) # type:
ignore[arg-type]
+ sys.modules[spec.name] = new_module # type: ignore[union-attr]
+ loader.exec_module(new_module)
+ return [new_module]
+ except KeyboardInterrupt:
+ sys.modules.pop(mod_name, None)
+ raise
+ except BaseException as e:
+ sys.modules.pop(mod_name, None)
+ DagContext.autoregistered_dags.clear()
+ log.exception("Failed to import: %s", filepath)
+ if self._dagbag_import_error_tracebacks:
+ stacktrace =
traceback.format_exc(limit=-self._dagbag_import_error_traceback_depth)
+ else:
+ stacktrace = None
+ result.errors.append(
+ DagImportError(
+ source_reference=repr(definition),
+ message=str(e),
+ error_type="import",
+ stacktrace=stacktrace,
+ )
+ )
+ return []
+
+ dagbag_import_timeout: float
+ try:
+ from airflow import settings # noqa: SDK002
+
+ dagbag_import_timeout =
settings.get_dagbag_import_timeout(filepath)
+ except (ImportError, AttributeError):
+ dagbag_import_timeout = 30.0
+
+ if not isinstance(dagbag_import_timeout, (int, float)):
+ raise AirflowConfigException(
+ f"Value ({dagbag_import_timeout}) from
get_dagbag_import_timeout must be int or float"
+ )
+
+ if dagbag_import_timeout <= 0:
+ return parse(mod_name, filepath)
+
+ timeout_msg = (
+ f"DagBag import timeout for {filepath} after
{dagbag_import_timeout}s.\n"
+ "Please take a look at these docs to improve your DAG import
time:\n"
+ "*
https://airflow.apache.org/docs/apache-airflow/stable/best-practices.html#top-level-python-code\n";
+ "*
https://airflow.apache.org/docs/apache-airflow/stable/best-practices.html#reducing-dag-complexity";
+ )
+ with timeout(seconds=dagbag_import_timeout, error_message=timeout_msg):
+ return parse(mod_name, filepath)
+
+ def _process_modules(
+ self,
+ mods: list[Any],
+ result: DagImportResult,
+ bundle: BaseDagBundle,
+ ) -> None:
+ """Extract DAG objects from modules. Validation happens in
bag_dag()."""
+ top_level_dags: set[tuple[DAG, Any]] = {
+ (o, m) for m in mods for o in m.__dict__.values() if isinstance(o,
DAG)
+ }
+ top_level_dags.update(DagContext.autoregistered_dags)
+
+ DagContext.current_autoregister_module_name = None
+ DagContext.autoregistered_dags.clear()
+
+ for dag, _mod in top_level_dags:
+ dag.bundle_name = bundle.name
+ dag.fileloc = repr(result.definition)
+ if result.definition is not None:
+ dag.relative_fileloc =
result.definition.get_relative_loc(bundle.path)
+ result.dags.append(dag)
+ log.debug("Found DAG %s", dag.dag_id)
+
+ @functools.cached_property
+ def _dagbag_import_error_tracebacks(self) -> bool:
+ return conf.getboolean("core", "dagbag_import_error_tracebacks")
+
+ @functools.cached_property
+ def _dagbag_import_error_traceback_depth(self) -> int:
+ return conf.getint("core", "dagbag_import_error_traceback_depth")
diff --git a/task-sdk/src/airflow/sdk/importers/zip_importer.py
b/task-sdk/src/airflow/sdk/importers/zip_importer.py
new file mode 100644
index 00000000000..bc55eda0d94
--- /dev/null
+++ b/task-sdk/src/airflow/sdk/importers/zip_importer.py
@@ -0,0 +1,270 @@
+# 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.
+"""Zip archive DAG importer."""
+
+from __future__ import annotations
+
+import contextlib
+import logging
+import sys
+import tempfile
+import threading
+import zipfile
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from airflow.sdk.exceptions import AirflowConfigException
+from airflow.sdk.importers.base import (
+ AbstractDagImporter,
+ DagDefinition,
+ DagImporterRegistry,
+ DagImportError,
+ DagImportResult,
+ DagSourceCode,
+ _get_importer_extensions,
+ _normalize_extensions,
+ _parse_importer_specs,
+ find_file_dag_definitions,
+ get_file_suffix,
+)
+
+if TYPE_CHECKING:
+ from collections.abc import Generator, Iterator
+
+ from airflow.dag_processing.bundles.base import BaseDagBundle # noqa:
SDK002
+
+log = logging.getLogger(__name__)
+
+_sys_path_lock = threading.RLock()
+
+
[email protected]
+def _temporary_sys_path(path: str) -> Generator[None, None, None]:
+ """Safely prepend a path to sys.path with synchronization and
restoration."""
+ with _sys_path_lock:
+ already_present = path in sys.path
+ if not already_present:
+ sys.path.insert(0, path)
+ try:
+ yield
+ finally:
+ if not already_present:
+ with contextlib.suppress(ValueError):
+ sys.path.remove(path)
+
+
+@dataclass
+class ZipFileDagDefinition(DagDefinition):
+ """A DAG definition backed by a file inside a ZIP archive."""
+
+ zip_path: Path
+ file_path: str
+ _content: bytes | None = field(default=None, repr=False, compare=False)
+
+ @property
+ def freshness_token(self) -> str:
+ try:
+ stat = self.zip_path.stat()
+ except OSError:
+ return ""
+ return f"{stat.st_mtime_ns}-{stat.st_size}-{self.file_path}"
+
+ def get_relative_loc(self, root: Path | None = None) -> str:
+ if root is not None:
+ with contextlib.suppress(ValueError):
+ return f"{self.zip_path.relative_to(root)}:{self.file_path}"
+ return f"{self.zip_path}:{self.file_path}"
+
+ def read_bytes(self) -> bytes:
+ if self._content is None:
+ with zipfile.ZipFile(self.zip_path) as z:
+ self._content = z.read(self.file_path)
+ return self._content
+
+ @contextlib.contextmanager
+ def as_file(self) -> Generator[Path, None, None]:
+ suffix = Path(self.file_path).suffix
+ with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as f:
+ f.write(self.read_bytes())
+ temp_path = Path(f.name)
+ try:
+ yield temp_path
+ finally:
+ with contextlib.suppress(OSError):
+ temp_path.unlink()
+
+ def __repr__(self) -> str:
+ return f"{self.zip_path}:{self.file_path}"
+
+
+class ZipImporter(AbstractDagImporter):
+ """Composite importer responsible for routing archive members to internal
importers."""
+
+ supported_extensions = [".zip"]
+
+ def __init__(
+ self,
+ internal_importers: dict[str, AbstractDagImporter | dict[str, Any]]
+ | list[dict[str, Any]]
+ | None = None,
+ extensions: list[str] | None = None,
+ ) -> None:
+ if extensions is not None:
+ self.supported_extensions = _normalize_extensions(extensions)
+ self._internal_extension_importers: dict[str, AbstractDagImporter] = {}
+ self._ordered_internal_importers: list[AbstractDagImporter] = []
+
+ if internal_importers is None:
+ from airflow.sdk.importers.python_importer import PythonDagImporter
+
+ self._register_internal(PythonDagImporter())
+ elif isinstance(internal_importers, list):
+ specs = _parse_importer_specs(internal_importers,
context="internal_importers of ZipImporter")
+ for spec in specs:
+ importer = DagImporterRegistry._instantiate_spec(spec)
+ self._register_internal(importer, extensions=spec.extensions)
+ elif isinstance(internal_importers, dict):
+ for ext, cfg in internal_importers.items():
+ if isinstance(cfg, AbstractDagImporter):
+ self._register_internal(cfg, extensions=[ext])
+ elif isinstance(cfg, dict):
+ specs = _parse_importer_specs(
+ [cfg], context=f"internal_importers configuration for
extension '{ext}'"
+ )
+ importer = DagImporterRegistry._instantiate_spec(specs[0])
+ self._register_internal(importer,
extensions=specs[0].extensions or [ext])
+ else:
+ raise AirflowConfigException(
+ f"Invalid internal importer configuration for
extension '{ext}': "
+ f"expected AbstractDagImporter or dictionary, got
{type(cfg).__name__}."
+ )
+ else:
+ raise AirflowConfigException(
+ f"Field 'internal_importers' must be a list or dictionary, got
{type(internal_importers).__name__}."
+ )
+
+ def can_handle(self, definition: DagDefinition | str | Path) -> bool:
+ """Check if this importer can handle the given definition based on
file extension."""
+ suffix = get_file_suffix(definition)
+ return suffix in self.supported_extensions if suffix else False
+
+ def list_dag_definitions(
+ self,
+ bundle: BaseDagBundle,
+ *,
+ safe_mode: bool = True,
+ ) -> Iterator[DagDefinition]:
+ """List zip archive DAG definitions in a bundle matching supported
extensions."""
+ yield from find_file_dag_definitions(bundle.path,
self.supported_extensions, safe_mode=safe_mode)
+
+ def import_definition(
+ self,
+ definition: DagDefinition,
+ bundle: BaseDagBundle,
+ *,
+ safe_mode: bool = True,
+ ) -> DagImportResult:
+ """
+ Import DAGs from a ZIP archive by routing its members to internal
importers.
+
+ The archive itself is placed on ``sys.path`` so Python imports between
+ members resolve via ``zipimport``. A real file is materialized on
demand
+ with :meth:`.as_file()` for internal importers.
+ """
+ result = DagImportResult(definition=definition)
+
+ with definition.as_file() as local_zip_path:
+ try:
+ with zipfile.ZipFile(local_zip_path) as z:
+ member_names = z.namelist()
+ except Exception as e:
+ result.errors.append(
+ DagImportError(
+
source_reference=definition.get_relative_loc(bundle.path),
+ message=f"Failed to read ZIP archive: {e}",
+ error_type="zip_read_error",
+ )
+ )
+ return result
+
+ with _temporary_sys_path(str(local_zip_path)):
+ for member_name in member_names:
+ if member_name.endswith("/") or
member_name.startswith("__MACOSX/"):
+ continue
+ # ZipSlip defence: reject traversal or absolute member
names.
+ member_path = Path(member_name)
+ if member_path.is_absolute() or ".." in member_path.parts:
+ log.warning(
+ "Skipping zip member %r in %s: directory traversal
patterns detected",
+ member_name,
+ definition,
+ )
+ continue
+
+ importer = self._get_internal_importer(member_name)
+ if importer is None:
+ continue
+
+ nested_def = ZipFileDagDefinition(zip_path=local_zip_path,
file_path=member_name)
+ if not importer.can_handle(nested_def):
+ continue
+
+ member_result = importer.import_definition(nested_def,
bundle, safe_mode=safe_mode)
+ result.dags.extend(member_result.dags)
+ result.errors.extend(member_result.errors)
+ result.warnings.extend(member_result.warnings)
+
result.skipped_definitions.extend(member_result.skipped_definitions)
+ result.dependencies.extend(member_result.dependencies)
+
+ return result
+
+ def get_source_code(self, definition: DagDefinition) -> DagSourceCode:
+ """
+ Return the source of a single archive member.
+
+ A zip is treated as a directory of DAG files: each member is its own
+ source unit, so this renders exactly the member named by ``definition``
+ through its file-type internal importer. The archive as a whole has no
+ source, the same way a directory does not; passing an archive-level
+ definition here is a category error and raises.
+ """
+ importer = self._get_internal_importer(definition)
+ if importer is None:
+ raise ValueError(f"No internal importer to read source for
{definition!r}")
+ return importer.get_source_code(definition)
+
+ def _register_internal(self, importer: AbstractDagImporter, extensions:
list[str] | None = None) -> None:
+ if importer not in self._ordered_internal_importers:
+ self._ordered_internal_importers.append(importer)
+ exts = extensions if extensions is not None else
_get_importer_extensions(importer)
+ if exts:
+ normalized = _normalize_extensions(exts)
+ if hasattr(importer, "supported_extensions"):
+ with contextlib.suppress(AttributeError, TypeError):
+ importer.supported_extensions = normalized
+ for ext in normalized:
+ self._internal_extension_importers[ext] = importer
+
+ def _get_internal_importer(self, member: DagDefinition | str) ->
AbstractDagImporter | None:
+ suffix = get_file_suffix(member)
+ if suffix and suffix in self._internal_extension_importers:
+ return self._internal_extension_importers[suffix]
+ for importer in reversed(self._ordered_internal_importers):
+ if importer.can_handle(member):
+ return importer
+ return None
diff --git a/task-sdk/tests/task_sdk/docs/test_public_api.py
b/task-sdk/tests/task_sdk/docs/test_public_api.py
index 09614581317..eea7740cc4b 100644
--- a/task-sdk/tests/task_sdk/docs/test_public_api.py
+++ b/task-sdk/tests/task_sdk/docs/test_public_api.py
@@ -53,6 +53,7 @@ def test_airflow_sdk_no_unexpected_exports():
"definitions",
"exceptions",
"execution_time",
+ "importers",
"io",
"lineage",
"listener",
diff --git
a/shared/module_loading/src/airflow_shared/module_loading/dag_file.py
b/task-sdk/tests/task_sdk/importers/__init__.py
similarity index 77%
copy from shared/module_loading/src/airflow_shared/module_loading/dag_file.py
copy to task-sdk/tests/task_sdk/importers/__init__.py
index d4fdc80c737..13a83393a91 100644
--- a/shared/module_loading/src/airflow_shared/module_loading/dag_file.py
+++ b/task-sdk/tests/task_sdk/importers/__init__.py
@@ -1,4 +1,3 @@
-#
# 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
@@ -15,9 +14,3 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""Dag file utilities for finding and loading Dag files."""
-
-from __future__ import annotations
-
-UNUSUAL_MODULE_PREFIX = "unusual_prefix_"
-MODIFIED_DAG_MODULE_NAME =
f"{UNUSUAL_MODULE_PREFIX}{{path_hash}}_{{module_name}}"
diff --git a/task-sdk/tests/task_sdk/importers/test_python_importer.py
b/task-sdk/tests/task_sdk/importers/test_python_importer.py
new file mode 100644
index 00000000000..46bd283de8c
--- /dev/null
+++ b/task-sdk/tests/task_sdk/importers/test_python_importer.py
@@ -0,0 +1,236 @@
+# 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.
+"""Tests for PythonDagImporter."""
+
+from __future__ import annotations
+
+import logging
+import py_compile
+import signal
+import sys
+from types import SimpleNamespace
+from unittest import mock
+
+import pytest
+
+from airflow.sdk.exceptions import AirflowConfigException
+from airflow.sdk.importers import (
+ FileDagDefinition,
+ PythonDagImporter,
+)
+
+
[email protected]
+def mock_bundle(tmp_path):
+ bundle_dir = tmp_path / "bundle"
+ bundle_dir.mkdir(parents=True, exist_ok=True)
+ return SimpleNamespace(name="test_bundle", path=bundle_dir)
+
+
+class TestPythonDagImporter:
+ """Test the PythonDagImporter implementation."""
+
+ def test_import_successful_dag(self, mock_bundle):
+ dag_file = mock_bundle.path / "sample_dag.py"
+ dag_file.write_text("from airflow.sdk import DAG\ndag =
DAG('test_dag_1')\n")
+
+ importer = PythonDagImporter()
+ definition = FileDagDefinition(path=dag_file)
+ result = importer.import_definition(definition, bundle=mock_bundle)
+
+ assert len(result.dags) == 1
+ assert result.dags[0].dag_id == "test_dag_1"
+ assert result.dags[0].bundle_name == "test_bundle"
+ assert result.dags[0].relative_fileloc == "sample_dag.py"
+ assert len(result.errors) == 0
+
+ def test_import_syntax_error_cleans_sys_modules(self, mock_bundle):
+ dag_file = mock_bundle.path / "bad_dag.py"
+ dag_file.write_text("from airflow.sdk import DAG\ndef broken(\n")
+
+ importer = PythonDagImporter()
+ result = importer.import_definition(FileDagDefinition(path=dag_file),
bundle=mock_bundle)
+
+ assert len(result.errors) == 1
+ assert result.errors[0].error_type == "import"
+ assert not any("bad_dag" in m for m in sys.modules)
+
+ def test_skip_non_dag_file_in_safe_mode(self, mock_bundle):
+ helper_file = mock_bundle.path / "helper.py"
+ helper_file.write_text("def util(): return 42\n")
+
+ importer = PythonDagImporter()
+ definition = FileDagDefinition(path=helper_file)
+ result = importer.import_definition(definition, bundle=mock_bundle,
safe_mode=True)
+
+ assert len(result.dags) == 0
+ assert len(result.errors) == 0
+ assert result.skipped_definitions == [definition]
+
+ @pytest.mark.parametrize(
+ ("safe_mode", "expected_files"),
+ [
+ (True, {"sample_dag.py"}),
+ (False, {"sample_dag.py", "helper.py"}),
+ ],
+ )
+ def test_list_dag_definitions(self, mock_bundle, safe_mode,
expected_files):
+ dag_file = mock_bundle.path / "sample_dag.py"
+ dag_file.write_text("from airflow.sdk import DAG\ndag =
DAG('test_dag_1')\n")
+ (mock_bundle.path / "helper.py").write_text("def helper():\n return
42\n")
+ (mock_bundle.path / "notes.txt").write_text("hello")
+
+ importer = PythonDagImporter()
+ defs = list(importer.list_dag_definitions(mock_bundle,
safe_mode=safe_mode))
+ assert {d.path.name for d in defs} == expected_files
+
+ @pytest.mark.parametrize(
+ ("filename", "is_bytecode", "expected_content"),
+ [
+ ("source_dag.py", False, "# My DAG\nfrom airflow.sdk import
DAG\n"),
+ ("source_dag.pyc", True, "# Sourceless bytecode (.pyc) — source
code not available\n"),
+ ],
+ )
+ def test_get_source_code(self, tmp_path, filename, is_bytecode,
expected_content):
+ dag_file = tmp_path / filename
+ if is_bytecode:
+ dag_file.write_bytes(b"\x00\x00\x00\x00bytecode")
+ else:
+ dag_file.write_text(expected_content)
+
+ src =
PythonDagImporter().get_source_code(FileDagDefinition(path=dag_file))
+ assert src.language == "python"
+ assert src.source_code == expected_content
+
+ def test_import_pyc_file(self, mock_bundle, tmp_path):
+ source_file = tmp_path / "compiled_dag.py"
+ source_file.write_text("from airflow.sdk import DAG\ndag =
DAG('compiled_dag')\n")
+ pyc_file = mock_bundle.path / "compiled_dag.pyc"
+ py_compile.compile(str(source_file), cfile=str(pyc_file))
+
+ importer = PythonDagImporter()
+ result = importer.import_definition(FileDagDefinition(path=pyc_file),
bundle=mock_bundle)
+
+ assert len(result.dags) == 1
+ assert result.dags[0].dag_id == "compiled_dag"
+ assert len(result.errors) == 0
+
+ def test_file_dag_definition_freshness_token(self, tmp_path):
+ dag_file = tmp_path / "fresh_dag.py"
+ dag_file.write_text("from airflow.sdk import DAG\n")
+ stat = dag_file.stat()
+ assert FileDagDefinition(path=dag_file).freshness_token ==
f"{stat.st_mtime_ns}-{stat.st_size}"
+
+ def test_python_importer_custom_extensions(self, mock_bundle):
+ importer = PythonDagImporter(extensions=[".custom_py"])
+ assert importer.can_handle("dag.custom_py")
+ assert not importer.can_handle("dag.py")
+ assert importer.supported_extensions == [".custom_py"]
+
+ dag_file = mock_bundle.path / "sample_dag.custom_py"
+ dag_file.write_text("from airflow.sdk import DAG\ndag =
DAG('custom_py_dag')\n")
+ defs = list(importer.list_dag_definitions(mock_bundle))
+ assert len(defs) == 1
+ assert defs[0].path == dag_file
+
+ @pytest.mark.parametrize(
+ ("enable_traceback", "expect_traceback"),
+ [
+ (True, True),
+ (False, False),
+ ],
+ )
+ @mock.patch("airflow.sdk.importers.python_importer.conf")
+ def test_import_error_traceback_configuration(
+ self, mock_conf, enable_traceback, expect_traceback, mock_bundle
+ ):
+ mock_conf.getboolean.return_value = enable_traceback
+ mock_conf.getint.return_value = 2
+
+ dag_file = mock_bundle.path / "bad.py"
+ dag_file.write_text("from airflow.sdk import DAG\ndef broken(\n")
+
+ importer = PythonDagImporter()
+ result = importer.import_definition(FileDagDefinition(path=dag_file),
bundle=mock_bundle)
+
+ assert len(result.errors) == 1
+ assert (result.errors[0].stacktrace is not None) == expect_traceback
+
+ def test_invalid_dagbag_import_timeout_raises_custom_exception(self,
mock_bundle):
+ mock_settings = mock.MagicMock()
+ mock_settings.get_dagbag_import_timeout.return_value =
"invalid_timeout_str"
+
+ importer = PythonDagImporter()
+ with (
+ mock.patch.dict("sys.modules", {"airflow":
mock.MagicMock(settings=mock_settings)}),
+ pytest.raises(
+ AirflowConfigException,
+ match=r"Value \(invalid_timeout_str\) from
get_dagbag_import_timeout must be int or float",
+ ),
+ ):
+ importer.import_definition(
+ FileDagDefinition(path=mock_bundle.path / "dag.py"),
+ bundle=mock_bundle,
+ safe_mode=False,
+ )
+
+ @mock.patch.object(PythonDagImporter, "_load_modules_from_file",
side_effect=TypeError("unexpected None"))
+ def test_unexpected_type_error_captured_in_result_errors(self, mock_load,
mock_bundle):
+ importer = PythonDagImporter()
+ result = importer.import_definition(
+ FileDagDefinition(path=mock_bundle.path / "dag.py"),
+ bundle=mock_bundle,
+ )
+
+ assert len(result.errors) == 1
+ assert result.errors[0].error_type == "import"
+ assert "unexpected None" in result.errors[0].message
+
+ def test_sigsegv_handler_registration_and_execution(self, mock_bundle):
+ dag_file = mock_bundle.path / "sample_dag.py"
+ dag_file.write_text("from airflow.sdk import DAG\ndag =
DAG('test_dag')\n")
+
+ importer = PythonDagImporter()
+ registered_handler = None
+
+ def mock_signal_func(signum, handler):
+ nonlocal registered_handler
+ if signum == signal.SIGSEGV:
+ registered_handler = handler
+
+ with mock.patch("signal.signal", side_effect=mock_signal_func):
+ result =
importer.import_definition(FileDagDefinition(path=dag_file), bundle=mock_bundle)
+ assert callable(registered_handler)
+
+ registered_handler(signal.SIGSEGV, None)
+ assert len(result.errors) == 1
+ assert result.errors[0].error_type == "segfault"
+ assert "Received SIGSEGV signal while processing" in
result.errors[0].message
+
+ def test_sigsegv_handler_registration_failure_logged(self, mock_bundle,
caplog):
+ dag_file = mock_bundle.path / "sample_dag.py"
+ dag_file.write_text("from airflow.sdk import DAG\ndag =
DAG('test_dag')\n")
+
+ importer = PythonDagImporter()
+ with (
+ mock.patch("signal.signal", side_effect=ValueError("signal only
works in main thread")),
+ caplog.at_level(logging.WARNING),
+ ):
+ result =
importer.import_definition(FileDagDefinition(path=dag_file), bundle=mock_bundle)
+
+ assert len(result.dags) == 1
+ assert "SIGSEGV signal handler registration failed. Not in the main
thread" in caplog.text
diff --git a/task-sdk/tests/task_sdk/importers/test_registry.py
b/task-sdk/tests/task_sdk/importers/test_registry.py
new file mode 100644
index 00000000000..44dc5929aae
--- /dev/null
+++ b/task-sdk/tests/task_sdk/importers/test_registry.py
@@ -0,0 +1,432 @@
+# 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.
+"""Tests for the DagImporterRegistry."""
+
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+
+import pytest
+
+from airflow.sdk.exceptions import AirflowConfigException
+from airflow.sdk.importers import (
+ AbstractDagImporter,
+ DagDefinition,
+ DagImporterRegistry,
+ DagImportError,
+ DagImportResult,
+ DagSourceCode,
+ FileDagDefinition,
+ PythonDagImporter,
+ ZipFileDagDefinition,
+ ZipImporter,
+ find_file_dag_definitions,
+ get_file_suffix,
+ get_importer_registry,
+ reset_importer_registry,
+)
+
+from tests_common.test_utils.config import conf_vars
+
+
+class GlobalDagImporter(PythonDagImporter):
+ pass
+
+
+class BundleDagImporter(PythonDagImporter):
+ pass
+
+
+class CustomBundleNonExtensionImporter(AbstractDagImporter):
+ """A non-extension custom importer that routes by URI prefix."""
+
+ def can_handle(self, definition: DagDefinition | str | Path) -> bool:
+ return isinstance(definition, str) and
definition.startswith("custom://")
+
+ def list_dag_definitions(self, bundle, **kwargs):
+ return iter([])
+
+ def import_definition(self, definition, bundle=None, **kwargs):
+ return DagImportResult()
+
+ def get_source_code(self, definition):
+ return DagSourceCode(source_code="", language="text")
+
+
+class LazyTestImporter(PythonDagImporter):
+ instances = 0
+
+ def __init__(self, **kwargs):
+ super().__init__()
+ self.kwargs = kwargs
+ LazyTestImporter.instances += 1
+
+
+class TestDagImporterRegistry:
+ """Test the DagImporterRegistry."""
+
+ @pytest.fixture(autouse=True)
+ def _clean_registry(self):
+ """Reset the registry before and after each test."""
+ reset_importer_registry()
+ yield
+ reset_importer_registry()
+
+ def test_singleton_pattern(self):
+ """Registry should return the same instance."""
+ registry1 = get_importer_registry()
+ registry2 = get_importer_registry()
+ assert registry1 is registry2
+
+ def test_default_importers_registered(self):
+ """Registry should have Python importer by default."""
+ registry = get_importer_registry()
+ extensions = registry.supported_extensions()
+
+ assert ".py" in extensions
+ assert ".zip" in extensions
+
+ def test_get_importer_for_python(self):
+ """Should return PythonDagImporter for .py files."""
+ registry = get_importer_registry()
+ importer = registry.get_importer("test.py")
+
+ assert importer is not None
+ assert isinstance(importer, PythonDagImporter)
+
+ def test_get_importer_for_zip(self):
+ """Should return ZipImporter for .zip files."""
+ registry = get_importer_registry()
+ importer = registry.get_importer("test.zip")
+
+ assert importer is not None
+ assert isinstance(importer, ZipImporter)
+
+ def test_get_importer_for_unknown(self):
+ """Should return None for unknown file types."""
+ registry = get_importer_registry()
+ importer = registry.get_importer("test.txt")
+
+ assert importer is None
+
+ def test_can_handle_supported_files(self):
+ """can_handle should return True for supported file types."""
+ registry = get_importer_registry()
+
+ assert registry.can_handle("dag.py")
+ assert registry.can_handle(Path("subdir/dag.py"))
+
+ def test_can_handle_unsupported_files(self):
+ """can_handle should return False for unsupported file types."""
+ registry = get_importer_registry()
+
+ assert not registry.can_handle("readme.txt")
+ assert not registry.can_handle("config.json")
+ assert not registry.can_handle("script.sh")
+
+ def test_case_insensitive_extension_matching(self):
+ """Extension matching should be case-insensitive."""
+ registry = get_importer_registry()
+
+ # All these should be handled
+ assert registry.can_handle("dag.PY")
+ assert registry.can_handle("dag.Py")
+
+ def test_reset_clears_singleton(self):
+ """reset() should clear the singleton instance."""
+ registry1 = get_importer_registry()
+ DagImporterRegistry.reset()
+ registry2 = get_importer_registry()
+
+ # Should be different instances after reset
+ assert registry1 is not registry2
+
+ def test_custom_importer_without_extensions(self):
+ """Registry should resolve non-file-based importers via can_handle."""
+ registry = get_importer_registry()
+ importer = CustomBundleNonExtensionImporter()
+ registry.register(importer)
+
+ assert registry.can_handle("custom://dags/sample")
+ assert registry.get_importer("custom://dags/sample") is importer
+ assert not registry.can_handle("other://dags/sample")
+
+ def
test_abstract_dag_importer_has_no_extension_attributes_or_methods(self):
+ """AbstractDagImporter must not define file extension attributes or
methods."""
+ assert not hasattr(AbstractDagImporter, "supported_extensions")
+ assert not hasattr(AbstractDagImporter, "might_contain_dag")
+
+ def test_custom_importer_file_pattern_can_handle(self):
+ """Registry resolves file definitions via can_handle when importer has
no supported_extensions."""
+
+ class PatternDagImporter(PythonDagImporter):
+ def can_handle(self, definition: DagDefinition | str | Path) ->
bool:
+ return "special_dag_" in str(definition) and
str(definition).endswith(".json")
+
+ registry = get_importer_registry()
+ importer = PatternDagImporter()
+ registry.register(importer)
+
+ assert registry.can_handle("special_dag_1.json")
+ assert registry.get_importer("special_dag_1.json") is importer
+ assert not registry.can_handle("normal_dag.json")
+ assert registry.get_importer("normal_dag.json") is None
+
+ def test_override_importer_for_extension(self, caplog):
+ """Registering an importer for an existing extension logs a warning
and evicts the previous importer."""
+ registry = DagImporterRegistry(register_defaults=True)
+ custom_importer = GlobalDagImporter()
+
+ with caplog.at_level(logging.WARNING):
+ registry.register(custom_importer, extensions=[".py"])
+
+ assert registry.get_importer("test.py") is custom_importer
+ assert any(
+ record.levelno == logging.WARNING and "already registered" in
record.message
+ for record in caplog.records
+ )
+
+ def test_lazy_importer_instantiation(self):
+ """Importer classes are not imported or instantiated until
get_importer is called."""
+ LazyTestImporter.instances = 0
+ reg = DagImporterRegistry(register_defaults=False)
+ reg.register_specs(
+ [
+ {
+ "classpath": f"{__name__}.LazyTestImporter",
+ "extensions": [".lazy", ".lazy2"],
+ "kwargs": {"param": "value"},
+ }
+ ],
+ context="test",
+ )
+
+ assert LazyTestImporter.instances == 0
+ assert reg.can_handle("file.lazy")
+ assert reg.can_handle("file.lazy2")
+ assert LazyTestImporter.instances == 0
+ assert set(reg.supported_extensions()) == {".lazy", ".lazy2"}
+ assert LazyTestImporter.instances == 0
+
+ importer1 = reg.get_importer("file.lazy")
+ assert LazyTestImporter.instances == 1
+ assert isinstance(importer1, LazyTestImporter)
+ assert importer1.kwargs == {"param": "value"}
+
+ importer2 = reg.get_importer("file.lazy2")
+ assert importer2 is importer1
+ assert LazyTestImporter.instances == 1
+
+ def test_from_config_three_tier_precedence(self):
+ """Bundle config overrides global config, which overrides defaults."""
+ global_config = [
+ {
+ "classpath": f"{__name__}.GlobalDagImporter",
+ "extensions": [".py", ".custom"],
+ }
+ ]
+ bundle_config_list = [
+ {
+ "name": "test_bundle",
+ "importers": [
+ {
+ "classpath": f"{__name__}.BundleDagImporter",
+ "extensions": [".custom"],
+ }
+ ],
+ }
+ ]
+
+ with conf_vars(
+ {
+ ("dag_processor", "dag_importer_configs"):
json.dumps(global_config),
+ ("dag_processor", "dag_bundle_config_list"):
json.dumps(bundle_config_list),
+ }
+ ):
+ global_reg = DagImporterRegistry.from_config()
+ assert isinstance(global_reg.get_importer("dag.py"),
GlobalDagImporter)
+ assert isinstance(global_reg.get_importer("dag.custom"),
GlobalDagImporter)
+
+ bundle_reg = DagImporterRegistry.from_config("test_bundle")
+ assert isinstance(bundle_reg.get_importer("dag.py"),
GlobalDagImporter)
+ assert isinstance(bundle_reg.get_importer("dag.custom"),
BundleDagImporter)
+
+ @pytest.mark.parametrize(
+ ("global_cfg", "bundle_cfg", "match"),
+ [
+ (json.dumps({"invalid": "object"}), None, "key
`dag_importer_configs` must be a list"),
+ (None, [{"extensions": [".py"]}], "Missing required 'classpath'"),
+ (None, [{"classpath": "invalid.path"}], "Failed to load DAG
importer"),
+ (
+ None,
+ [{"classpath": "builtins.dict"}],
+ r"Configured DAG importer builtins\.dict for bundle
'test_bundle' must inherit from AbstractDagImporter\.",
+ ),
+ (
+ json.dumps([{"classpath": "builtins.dict"}]),
+ None,
+ r"Configured DAG importer builtins\.dict for global
configuration must inherit from AbstractDagImporter\.",
+ ),
+ ],
+ )
+ def test_from_config_invalid_configs(self, global_cfg, bundle_cfg, match):
+ """Invalid configurations raise AirflowConfigException."""
+ overrides = {}
+ if global_cfg:
+ overrides[("dag_processor", "dag_importer_configs")] = global_cfg
+ if bundle_cfg:
+ overrides[("dag_processor", "dag_bundle_config_list")] =
json.dumps(
+ [
+ {
+ "name": "test_bundle",
+ "importers": bundle_cfg,
+ }
+ ]
+ )
+
+ with conf_vars(overrides), pytest.raises(AirflowConfigException,
match=match):
+ DagImporterRegistry.from_config("test_bundle" if bundle_cfg else
None)
+
+ def test_get_importer_registry_caching_and_isolation(self):
+ """get_importer_registry caches instances per bundle name and clears
on reset."""
+ reg1 = get_importer_registry()
+ reg2 = get_importer_registry()
+ assert reg1 is reg2
+
+ bundle_a = get_importer_registry("bundle_a")
+ bundle_a2 = get_importer_registry("bundle_a")
+ bundle_b = get_importer_registry("bundle_b")
+
+ assert bundle_a is bundle_a2
+ assert bundle_a is not bundle_b
+ assert bundle_a is not reg1
+
+ reset_importer_registry()
+ new_reg = get_importer_registry()
+ new_bundle_a = get_importer_registry("bundle_a")
+ assert new_reg is not reg1
+ assert new_bundle_a is not bundle_a
+
+ def test_bundle_config_with_extensions_and_non_extension_importers(self):
+ """Verify bundle config with extensions, non-extension importers, and
archive importer."""
+ bundle_config_list = [
+ {
+ "name": "dags-folder",
+ "classpath":
"airflow.dag_processing.bundles.local.LocalDagBundle",
+ "kwargs": {},
+ "importers": [
+ {
+ "classpath":
"airflow.sdk.importers.python_importer.PythonDagImporter",
+ "extensions": [".py"],
+ },
+ {
+ "classpath":
f"{__name__}.CustomBundleNonExtensionImporter",
+ },
+ {
+ "classpath":
"airflow.sdk.importers.zip_importer.ZipImporter",
+ "extensions": [".zip"],
+ "kwargs": {
+ "internal_importers": [
+ {
+ "classpath":
"airflow.sdk.importers.python_importer.PythonDagImporter",
+ "extensions": [".py"],
+ }
+ ]
+ },
+ },
+ ],
+ }
+ ]
+
+ with conf_vars({("dag_processor", "dag_bundle_config_list"):
json.dumps(bundle_config_list)}):
+ registry = DagImporterRegistry.from_config("dags-folder")
+
+ # 1. Importer with extensions (.py)
+ assert registry.can_handle("sample_dag.py")
+ py_importer = registry.get_importer("sample_dag.py")
+ assert isinstance(py_importer, PythonDagImporter)
+
+ # 2. Non-extension importer (custom://)
+ assert registry.can_handle("custom://my-pipeline-1")
+ custom_importer = registry.get_importer("custom://my-pipeline-1")
+ assert isinstance(custom_importer,
CustomBundleNonExtensionImporter)
+
+ # 3. Archive importer (.zip) with internal_importers
+ assert registry.can_handle("bundle.zip")
+ zip_importer = registry.get_importer("bundle.zip")
+ assert isinstance(zip_importer, ZipImporter)
+ assert ".py" in zip_importer._internal_extension_importers
+
+ @pytest.mark.parametrize(
+ ("input_val", "expected"),
+ [
+ ("foo.py", ".py"),
+ ("path/to/FOO.PY", ".py"),
+ (Path("archive.ZIP"), ".zip"),
+ ("no_extension", ""),
+ (FileDagDefinition(path=Path("my_dag.py")), ".py"),
+ (ZipFileDagDefinition(zip_path=Path("a.zip"),
file_path="nested/workflow.py"), ".py"),
+ (None, None),
+ ],
+ )
+ def test_get_file_suffix(self, input_val, expected):
+ assert get_file_suffix(input_val) == expected
+
+ @pytest.mark.parametrize(
+ ("error", "expected"),
+ [
+ (
+ DagImportError(
+ source_reference="dag.py",
+ message="syntax error",
+ error_type="syntax",
+ ),
+ "Error in dag.py [syntax]: syntax error",
+ ),
+ (
+ DagImportError(
+ source_reference="dag.py",
+ message="unexpected token",
+ error_type="syntax",
+ line_number=10,
+ column_number=5,
+ context="def broken(\n",
+ suggestion="close the parenthesis",
+ ),
+ "Error in dag.py (line 10, column 5) [syntax]: unexpected
token; Context: def broken(; Suggestion: close the parenthesis",
+ ),
+ ],
+ )
+ def test_dag_import_error_format_message(self, error, expected):
+ assert error.format_message() == expected
+
+ @pytest.mark.parametrize(
+ ("safe_mode", "expected_files"),
+ [
+ (True, {"workflow.py"}),
+ (False, {"workflow.py", "script.py"}),
+ ],
+ )
+ def test_find_file_dag_definitions_safe_mode(self, tmp_path, safe_mode,
expected_files):
+ (tmp_path / "workflow.py").write_text("from airflow.sdk import DAG\n")
+ (tmp_path / "script.py").write_text("print('hello')\n")
+ (tmp_path / "data.csv").write_text("a,b,c\n")
+
+ definitions = list(find_file_dag_definitions(tmp_path, [".py"],
safe_mode=safe_mode))
+ assert {d.path.name for d in definitions} == expected_files
diff --git a/task-sdk/tests/task_sdk/importers/test_zip_importer.py
b/task-sdk/tests/task_sdk/importers/test_zip_importer.py
new file mode 100644
index 00000000000..023db51cb16
--- /dev/null
+++ b/task-sdk/tests/task_sdk/importers/test_zip_importer.py
@@ -0,0 +1,237 @@
+# 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.
+"""Tests for ZipImporter."""
+
+from __future__ import annotations
+
+import py_compile
+import zipfile
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+
+from airflow.sdk.exceptions import AirflowConfigException
+from airflow.sdk.importers import (
+ AbstractDagImporter,
+ DagDefinition,
+ DagImportResult,
+ DagSourceCode,
+ FileDagDefinition,
+ PythonDagImporter,
+ ZipFileDagDefinition,
+ ZipImporter,
+)
+
+
+class CustomInternalNonExtensionImporter(AbstractDagImporter):
+ """An internal importer that routes archive members by filename pattern."""
+
+ def can_handle(self, definition: DagDefinition | str | Path) -> bool:
+ path_str = str(getattr(definition, "file_path", definition))
+ return "workflow" in path_str
+
+ def list_dag_definitions(self, bundle, **kwargs):
+ return iter([])
+
+ def import_definition(self, definition, bundle=None, **kwargs):
+ from airflow.sdk import DAG
+
+ return DagImportResult(dags=[DAG("workflow_dag")])
+
+ def get_source_code(self, definition):
+ return DagSourceCode(source_code="", language="text")
+
+
+class TestZipImporter:
+ """Test the ZipImporter composite implementation."""
+
+ @pytest.fixture
+ def mock_bundle(self, tmp_path):
+ return SimpleNamespace(name="test_bundle", path=tmp_path)
+
+ @pytest.mark.parametrize(
+ ("safe_mode", "expected_count"),
+ [
+ (True, 1),
+ (False, 2),
+ ],
+ )
+ def test_list_dag_definitions(self, mock_bundle, safe_mode,
expected_count):
+ zip_path = mock_bundle.path / "sample.zip"
+ with zipfile.ZipFile(zip_path, "w") as z:
+ z.writestr("dag.py", "from airflow.sdk import DAG\n")
+ (mock_bundle.path / "corrupt.zip").write_bytes(b"not a valid zip and
no dag markers")
+
+ definitions = list(ZipImporter().list_dag_definitions(mock_bundle,
safe_mode=safe_mode))
+ assert len(definitions) == expected_count
+ assert any(d.path == zip_path for d in definitions)
+
+ def test_import_zip_archive_with_dags(self, mock_bundle):
+ zip_path = mock_bundle.path / "sample_dags.zip"
+ with zipfile.ZipFile(zip_path, "w") as z:
+ z.writestr("dag_a.py", "from airflow.sdk import DAG\ndag =
DAG('zip_dag_a')\n")
+ z.writestr("helper.py", "def util(): return 1\n")
+
+ result =
ZipImporter().import_definition(FileDagDefinition(path=zip_path),
bundle=mock_bundle)
+ assert len(result.dags) == 1
+ assert result.dags[0].dag_id == "zip_dag_a"
+ assert result.dags[0].bundle_name == "test_bundle"
+ assert len(result.errors) == 0
+
+ def test_import_zip_archive_with_pyc_dag(self, mock_bundle, tmp_path):
+ source_file = tmp_path / "compiled_dag.py"
+ source_file.write_text("from airflow.sdk import DAG\ndag =
DAG('zip_pyc_dag')\n")
+ pyc_file = tmp_path / "compiled_dag.pyc"
+ py_compile.compile(str(source_file), cfile=str(pyc_file))
+
+ zip_path = mock_bundle.path / "sample_pyc_dags.zip"
+ with zipfile.ZipFile(zip_path, "w") as z:
+ z.write(pyc_file, arcname="compiled_dag.pyc")
+
+ importer = ZipImporter()
+ result = importer.import_definition(FileDagDefinition(path=zip_path),
bundle=mock_bundle)
+ assert len(result.dags) == 1
+ assert result.dags[0].dag_id == "zip_pyc_dag"
+ assert len(result.errors) == 0
+
+ src = importer.get_source_code(ZipFileDagDefinition(zip_path=zip_path,
file_path="compiled_dag.pyc"))
+ assert src.language == "python"
+ assert "Sourceless bytecode" in src.source_code
+
+ def test_zipslip_traversal_and_metadata_skipped(self, mock_bundle):
+ zip_path = mock_bundle.path / "malicious.zip"
+ with zipfile.ZipFile(zip_path, "w") as z:
+ z.writestr("subfolder/", "")
+ z.writestr("__MACOSX/._dag.py", "apple double metadata")
+ z.writestr("../evil_dag.py", "from airflow.sdk import DAG\ndag =
DAG('evil_dag')\n")
+ z.writestr("valid_dag.py", "from airflow.sdk import DAG\ndag =
DAG('valid_dag')\n")
+
+ result =
ZipImporter().import_definition(FileDagDefinition(path=zip_path),
bundle=mock_bundle)
+ assert len(result.dags) == 1
+ assert result.dags[0].dag_id == "valid_dag"
+ assert not (mock_bundle.path.parent / "evil_dag.py").exists()
+
+ def test_corrupted_zip_file(self, mock_bundle):
+ bad_zip = mock_bundle.path / "corrupted.zip"
+ bad_zip.write_bytes(b"not a real zip")
+
+ result =
ZipImporter().import_definition(FileDagDefinition(path=bad_zip),
bundle=mock_bundle)
+ assert len(result.errors) == 1
+ assert result.errors[0].error_type == "zip_read_error"
+
+ def test_get_source_code_reads_member_not_archive(self, tmp_path):
+ zip_path = tmp_path / "source_dags.zip"
+ dag_content = "from airflow.sdk import DAG\ndag = DAG('src_dag')\n"
+ with zipfile.ZipFile(zip_path, "w") as z:
+ z.writestr("my_dag.py", dag_content)
+
+ importer = ZipImporter()
+
+ # A zip is a directory of DAG files: each member is its own source
unit,
+ # rendered through its file-type importer (same single-member
semantics as
+ # the legacy code view's open_maybe_zipped).
+ src_member =
importer.get_source_code(ZipFileDagDefinition(zip_path=zip_path,
file_path="my_dag.py"))
+ assert src_member.language == "python"
+ assert src_member.source_code == dag_content
+
+ # The archive as a whole has no source, the same way a directory does
not.
+ with pytest.raises(ValueError, match="No internal importer"):
+ importer.get_source_code(FileDagDefinition(path=zip_path))
+
+ def test_zip_dag_definition_freshness_token(self, tmp_path):
+ zip_path = tmp_path / "fresh_bundle.zip"
+ with zipfile.ZipFile(zip_path, "w") as z:
+ z.writestr("dag.py", "from airflow.sdk import DAG\n")
+
+ member_def = ZipFileDagDefinition(zip_path=zip_path,
file_path="dag.py")
+ stat = zip_path.stat()
+ assert member_def.freshness_token ==
f"{stat.st_mtime_ns}-{stat.st_size}-dag.py"
+
+ def test_zip_importer_internal_importers_from_list_of_specs(self,
tmp_path, mock_bundle):
+ importer = ZipImporter(
+ internal_importers=[
+ {
+ "classpath":
"airflow.sdk.importers.python_importer.PythonDagImporter",
+ "extensions": [".custom_py"],
+ }
+ ]
+ )
+ assert ".custom_py" in importer._internal_extension_importers
+ assert ".py" not in importer._internal_extension_importers
+
+ zip_path = mock_bundle.path / "custom_dags.zip"
+ with zipfile.ZipFile(zip_path, "w") as z:
+ z.writestr("dag.custom_py", "from airflow.sdk import DAG\ndag =
DAG('custom_zip_dag')\n")
+
+ res = importer.import_definition(FileDagDefinition(path=zip_path),
bundle=mock_bundle)
+ assert len(res.dags) == 1
+ assert res.dags[0].dag_id == "custom_zip_dag"
+
+ def test_zip_importer_internal_importers_from_dict(self):
+ importer = ZipImporter(
+ internal_importers={
+ ".py": PythonDagImporter(),
+ ".alt": {
+ "classpath":
"airflow.sdk.importers.python_importer.PythonDagImporter",
+ "extensions": [".alt"],
+ },
+ }
+ )
+ assert ".py" in importer._internal_extension_importers
+ assert ".alt" in importer._internal_extension_importers
+
+ @pytest.mark.parametrize(
+ ("config", "match"),
+ [
+ ("not_list_or_dict", "must be a list or dictionary"),
+ ([{"extensions": [".py"]}], "Missing required 'classpath'"),
+ ([{"classpath": "invalid.path"}], "Failed to load DAG importer"),
+ (
+ [{"classpath": "builtins.dict"}],
+ r"must inherit from AbstractDagImporter",
+ ),
+ ({".py": {"kwargs": {}}}, "Missing required 'classpath'"),
+ ({".py": "invalid"}, "expected AbstractDagImporter or dictionary"),
+ ],
+ )
+ def test_zip_importer_invalid_configurations(self, config, match):
+ with pytest.raises(AirflowConfigException, match=match):
+ ZipImporter(internal_importers=config)
+
+ def test_zip_importer_custom_extensions(self):
+ importer = ZipImporter(extensions=[".bundle", ".zip"])
+ assert importer.can_handle("test.bundle")
+ assert importer.can_handle("test.zip")
+ assert set(importer.supported_extensions) == {".bundle", ".zip"}
+
+ def test_zip_importer_internal_importers_non_extension(self, mock_bundle):
+ """ZipImporter can route archive members to non-extension internal
importers."""
+ importer = ZipImporter(
+ internal_importers=[
+ {
+ "classpath":
f"{__name__}.CustomInternalNonExtensionImporter",
+ }
+ ]
+ )
+ zip_path = mock_bundle.path / "workflow_archive.zip"
+ with zipfile.ZipFile(zip_path, "w") as z:
+ z.writestr("my_workflow_file", "steps:\n - run: echo hello\n")
+
+ res = importer.import_definition(FileDagDefinition(path=zip_path),
bundle=mock_bundle)
+ assert len(res.dags) == 1
+ assert res.dags[0].dag_id == "workflow_dag"
