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 a11da87a119 Speed up zip Dag discovery and keep member file names
(#73727)
a11da87a119 is described below
commit a11da87a1195f1927f8668466df5a3ebee24051c
Author: Dilnaz Amanzholova <[email protected]>
AuthorDate: Sat Sep 26 09:40:16 2026 +0200
Speed up zip Dag discovery and keep member file names (#73727)
Safe-mode discovery reopened the archive for every member, making it
quadratic in the number of members. Members materialized to disk got a
random temporary name, which leaked into identifiers that importers derive
from the file name, such as the Dag id.
---
task-sdk/src/airflow/sdk/importers/zip_importer.py | 49 ++++++++++++----------
.../tests/task_sdk/importers/test_zip_importer.py | 27 ++++++++++++
2 files changed, 53 insertions(+), 23 deletions(-)
diff --git a/task-sdk/src/airflow/sdk/importers/zip_importer.py
b/task-sdk/src/airflow/sdk/importers/zip_importer.py
index d9de060d9cb..09b12c14f04 100644
--- a/task-sdk/src/airflow/sdk/importers/zip_importer.py
+++ b/task-sdk/src/airflow/sdk/importers/zip_importer.py
@@ -103,14 +103,11 @@ class ZipMemberDagDefinition(FileDagDefinition):
@contextlib.contextmanager
def as_file(self) -> Generator[Path, None, None]:
- with tempfile.NamedTemporaryFile(suffix=self.suffix, delete=False) as
f:
- f.write(self.read_bytes())
- temp_path = Path(f.name)
- try:
+ # Keep the member's file name: importers may derive identity from it,
such as a Dag id.
+ with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as
temp_dir:
+ temp_path = Path(temp_dir, Path(self.file_path).name)
+ temp_path.write_bytes(self.read_bytes())
yield temp_path
- finally:
- with contextlib.suppress(OSError):
- temp_path.unlink()
def __repr__(self) -> str:
return str(self.zip_path.joinpath(self.file_path))
@@ -183,8 +180,7 @@ class
ZipImporter(AbstractDagImporter[ZipMemberDagDefinition]):
"""
for archive in find_file_dag_definitions(bundle.path,
self.supported_extensions):
try:
- with zipfile.ZipFile(archive.path) as z:
- member_names = z.namelist()
+ zip_file = zipfile.ZipFile(archive.path)
except Exception as e:
log.warning("Cannot read ZIP archive %s: %s", archive.path, e)
yield DagImportError(
@@ -194,20 +190,27 @@ class
ZipImporter(AbstractDagImporter[ZipMemberDagDefinition]):
)
continue
- for member, importer in self._iter_supported_members(archive.path,
member_names):
- try:
- if safe_mode and not importer.might_contain_dag(member,
safe_mode):
- continue
- except Exception as e:
- # One unreadable member must not end discovery for the
rest of the archive.
- log.warning("Cannot read zip member %s of %s: %s",
member.file_path, archive.path, e)
- yield DagImportError(
- source_reference=member.get_relative_loc(bundle.path),
- message=f"Failed to read ZIP member: {e}",
- error_type="zip_read_error",
- )
- continue
- yield member
+ with zip_file:
+ for member, importer in
self._iter_supported_members(archive.path, zip_file.namelist()):
+ if safe_mode:
+ try:
+ # Read through the open archive: reopening it per
member re-parses the
+ # central directory and makes discovery quadratic
in the member count.
+ member._content = zip_file.read(member.file_path)
+ if not importer.might_contain_dag(member,
safe_mode):
+ continue
+ except Exception as e:
+ # One unreadable member must not end discovery for
the rest of the archive.
+ log.warning(
+ "Cannot read zip member %s of %s: %s",
member.file_path, archive.path, e
+ )
+ yield DagImportError(
+
source_reference=member.get_relative_loc(bundle.path),
+ message=f"Failed to read ZIP member: {e}",
+ error_type="zip_read_error",
+ )
+ continue
+ yield member
def _iter_supported_members(
self, zip_path: Path, member_names: list[str]
diff --git a/task-sdk/tests/task_sdk/importers/test_zip_importer.py
b/task-sdk/tests/task_sdk/importers/test_zip_importer.py
index baa5516d113..3d4a52bc696 100644
--- a/task-sdk/tests/task_sdk/importers/test_zip_importer.py
+++ b/task-sdk/tests/task_sdk/importers/test_zip_importer.py
@@ -24,6 +24,7 @@ import zipfile
from pathlib import Path
from types import SimpleNamespace
from typing import TYPE_CHECKING
+from unittest import mock
import pytest
@@ -153,6 +154,21 @@ class TestZipImporter:
assert [e.error_type for e in errors] == ["zip_read_error"]
assert errors[0].source_reference == os.path.join("bad_member.zip",
"b_bad.py")
+ @mock.patch(
+ "airflow.sdk.importers.zip_importer.zipfile.ZipFile", autospec=True,
side_effect=zipfile.ZipFile
+ )
+ def test_list_opens_archive_once(self, mock_zip_file, mock_bundle):
+ zip_path = mock_bundle.path / "many.zip"
+ with zipfile.ZipFile(zip_path, "w") as z:
+ for i in range(5):
+ z.writestr(f"dag_{i}.py", "from airflow.sdk import DAG\n")
+ mock_zip_file.reset_mock()
+
+ definitions = list(ZipImporter().list_dag_definitions(mock_bundle,
safe_mode=True))
+
+ assert len(definitions) == 5
+ mock_zip_file.assert_called_once_with(zip_path)
+
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:
@@ -256,6 +272,17 @@ class TestZipImporter:
with pytest.raises(ValueError, match="No internal importer"):
importer.get_source_code(FilesystemDagDefinition(path=zip_path))
+ def test_as_file_keeps_member_file_name(self, tmp_path):
+ zip_path = tmp_path / "bundle.zip"
+ with zipfile.ZipFile(zip_path, "w") as z:
+ z.writestr("pipelines/my_pipeline.yaml", "pipelineId:
my_pipeline\n")
+ member = ZipMemberDagDefinition(zip_path=zip_path,
file_path="pipelines/my_pipeline.yaml")
+
+ with member.as_file() as local_path:
+ assert local_path.name == "my_pipeline.yaml"
+ assert local_path.read_text() == "pipelineId: my_pipeline\n"
+ assert not local_path.exists()
+
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:
