Skip to content

Commit 849ba15

Browse files
committed
Fix template file lookup for Dags packaged in zip files
A Dag imported from a zip archive has its folder inside the archive, which Jinja's FileSystemLoader cannot read. Template files shipped next to the Dag in the archive were therefore never found: parse-time resolution silently kept the raw file path, and every task run failed with TemplateNotFound before execute(). Pointing template_searchpath into the archive failed the same way, leaving Dag authors with only a workaround outside the public Task SDK interface.
1 parent 051d5d8 commit 849ba15

5 files changed

Lines changed: 265 additions & 3 deletions

File tree

‎airflow-core/docs/core-concepts/dags.rst‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -721,6 +721,11 @@ You can either do this all inside of the Dag bundle, with a standard filesystem
721721
package1/__init__.py
722722
package1/functions.py
723723

724+
Template files used by your tasks, such as ``.sql`` or ``.sh`` files, can be packaged in the same zip file.
725+
They are found relative to the Dag file's folder inside the archive, as they would be in an unpacked Dag bundle.
726+
``template_searchpath`` entries can also point to a folder inside the archive, for example
727+
``os.path.join(os.path.dirname(__file__), "sql")``.
728+
724729
Note that packaged Dags come with some caveats:
725730

726731
* They cannot be used if you have pickling enabled for serialization

‎task-sdk/src/airflow/sdk/definitions/_internal/templater.py‎

Lines changed: 80 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,13 +20,17 @@
2020
import datetime
2121
import logging
2222
import os
23-
from collections.abc import Collection, Iterable, Iterator, Sequence
23+
import posixpath
24+
import zipfile
25+
from collections.abc import Callable, Collection, Iterable, Iterator, Sequence
2426
from dataclasses import dataclass
27+
from pathlib import Path, PurePosixPath
2528
from typing import TYPE_CHECKING, Any
2629

2730
import jinja2
2831
import jinja2.nativetypes
2932
import jinja2.sandbox
33+
from jinja2.loaders import split_template_path
3034

3135
from airflow.sdk import ObjectStoragePath
3236
from airflow.sdk.definitions._internal.mixins import ResolveMixin
@@ -374,6 +378,80 @@ def ts_nodash_with_tz_filter(value: datetime.date | datetime.time | None) -> str
374378
}
375379

376380

381+
def _find_zip_root(path: str) -> tuple[str, str] | None:
382+
"""Return ``(archive, prefix)`` when *path* is a zip archive or a directory inside one."""
383+
candidate = Path(path)
384+
existing = next((p for p in (candidate, *candidate.parents) if os.path.exists(p)), None)
385+
if existing is None or not os.path.isfile(existing) or not zipfile.is_zipfile(existing):
386+
return None
387+
prefix = candidate.relative_to(existing).as_posix()
388+
if ".." in PurePosixPath(prefix).parts:
389+
return None
390+
return os.fspath(existing), "" if prefix == "." else prefix
391+
392+
393+
class _ZipArchiveLoader(jinja2.BaseLoader):
394+
"""Load templates from a directory inside a zip archive, reading members without extracting them."""
395+
396+
def __init__(self, archive: str, prefix: str, encoding: str = "utf-8") -> None:
397+
self.archive = archive
398+
self.prefix = prefix
399+
self.encoding = encoding
400+
401+
def get_source(
402+
self, environment: jinja2.Environment, template: str
403+
) -> tuple[str, str, Callable[[], bool]]:
404+
member = posixpath.join(self.prefix, *split_template_path(template))
405+
mtime = os.path.getmtime(self.archive)
406+
# Opened per lookup: a handle kept across the task supervisor's fork would share its file offset.
407+
with zipfile.ZipFile(self.archive) as zf:
408+
try:
409+
source = zf.read(member).decode(self.encoding)
410+
except KeyError:
411+
raise jinja2.TemplateNotFound(template) from None
412+
413+
def is_up_to_date() -> bool:
414+
try:
415+
return os.path.getmtime(self.archive) == mtime
416+
except OSError:
417+
return False
418+
419+
return source, os.path.join(self.archive, member), is_up_to_date
420+
421+
422+
class _SearchPathLoader(jinja2.BaseLoader):
423+
"""Try each search path's loader in order, reporting the search paths like ``FileSystemLoader``."""
424+
425+
def __init__(self, searchpath: list[str], loaders: list[jinja2.BaseLoader]) -> None:
426+
self.searchpath = searchpath
427+
self.loaders = loaders
428+
429+
def get_source(
430+
self, environment: jinja2.Environment, template: str
431+
) -> tuple[str, str | None, Callable[[], bool] | None]:
432+
for loader in self.loaders:
433+
try:
434+
return loader.get_source(environment, template)
435+
except jinja2.TemplateNotFound:
436+
continue
437+
plural = "path" if len(self.searchpath) == 1 else "paths"
438+
paths = ", ".join(repr(path) for path in self.searchpath)
439+
raise jinja2.TemplateNotFound(template, f"{template!r} not found in search {plural}: {paths}")
440+
441+
442+
def _build_template_loader(searchpath: list[str]) -> jinja2.BaseLoader:
443+
zip_roots = [None if os.path.isdir(path) else _find_zip_root(path) for path in searchpath]
444+
if not any(zip_roots):
445+
return jinja2.FileSystemLoader(searchpath)
446+
return _SearchPathLoader(
447+
searchpath,
448+
[
449+
_ZipArchiveLoader(*zip_root) if zip_root else jinja2.FileSystemLoader(path)
450+
for path, zip_root in zip(searchpath, zip_roots, strict=True)
451+
],
452+
)
453+
454+
377455
def create_template_env(
378456
*,
379457
native: bool = False,
@@ -391,7 +469,7 @@ def create_template_env(
391469
"cache_size": 0,
392470
}
393471
if searchpath:
394-
jinja_env_options["loader"] = jinja2.FileSystemLoader(searchpath)
472+
jinja_env_options["loader"] = _build_template_loader(searchpath)
395473
if jinja_environment_kwargs:
396474
jinja_env_options.update(jinja_environment_kwargs)
397475

‎task-sdk/tests/task_sdk/bases/test_operator.py‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,10 @@
2020
import asyncio
2121
import copy
2222
import logging
23+
import os
2324
import uuid
2425
import warnings
26+
import zipfile
2527
from datetime import UTC, date, datetime, timedelta
2628
from typing import NamedTuple
2729
from unittest import mock
@@ -120,6 +122,17 @@ def __init__(self, arg1: str = "", arg2: str = "", **kwargs):
120122
)
121123

122124

125+
class MockFileTemplateOperator(BaseOperator):
126+
"""Operator with a ``.sql`` template file nested in a dict, like ``BigQueryInsertJobOperator``."""
127+
128+
template_fields = ("configuration",)
129+
template_ext = (".sql",)
130+
131+
def __init__(self, configuration: dict, **kwargs):
132+
super().__init__(**kwargs)
133+
self.configuration = configuration
134+
135+
123136
class TestBaseOperator:
124137
def setup_method(self, method):
125138
MockOperator.start_from_trigger = False
@@ -757,6 +770,23 @@ def fn_to_template(**kwargs):
757770
task.render_template_fields({})
758771
assert task.arg2 == "foo_barbarbar"
759772

773+
@pytest.mark.parametrize("op_native", [None, True], ids=["dag-env", "operator-native-env"])
774+
def test_render_template_fields_reads_template_file_from_zipped_dag(self, tmp_path, op_native):
775+
archive = tmp_path / "dags.zip"
776+
with zipfile.ZipFile(archive, "w") as zf:
777+
zf.writestr("test_sql/test.sql", "SELECT column_a FROM {{ table }}")
778+
with DAG("zipped_dag", schedule=None, start_date=DEFAULT_DATE) as dag:
779+
task = MockFileTemplateOperator(
780+
task_id="op1",
781+
configuration={"query": {"query": "test_sql/test.sql"}},
782+
render_template_as_native_obj=op_native,
783+
)
784+
dag.fileloc = os.path.join(archive, "test.py")
785+
786+
task.render_template_fields(context={"table": "test"})
787+
788+
assert task.configuration == {"query": {"query": "SELECT column_a FROM test"}}
789+
760790
@pytest.mark.parametrize("content", [object(), uuid.uuid4()])
761791
def test_render_template_fields_no_change(self, content):
762792
"""Tests if non-templatable types remain unchanged."""

‎task-sdk/tests/task_sdk/definitions/_internal/test_templater.py‎

Lines changed: 118 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,14 +17,22 @@
1717

1818
from __future__ import annotations
1919

20+
import os
21+
import zipfile
2022
from datetime import UTC, datetime
23+
from pathlib import Path
2124
from unittest.mock import MagicMock, NonCallableMagicMock
2225

2326
import jinja2
2427
import pytest
2528

2629
from airflow.sdk import DAG, ObjectStoragePath
27-
from airflow.sdk.definitions._internal.templater import LiteralValue, SandboxedEnvironment, Templater
30+
from airflow.sdk.definitions._internal.templater import (
31+
LiteralValue,
32+
SandboxedEnvironment,
33+
Templater,
34+
create_template_env,
35+
)
2836

2937

3038
class TestTemplater:
@@ -300,6 +308,115 @@ def test_do_render_template_fields_renders_list_values(self):
300308
assert parent.items == ["first", "second"]
301309

302310

311+
def _write_zip(path: Path, members: dict[str, str]) -> Path:
312+
with zipfile.ZipFile(path, "w") as zf:
313+
for name, content in members.items():
314+
zf.writestr(name, content)
315+
return path
316+
317+
318+
def _write_template_dir(path: Path, templates: dict[str, str]) -> Path:
319+
path.mkdir()
320+
for name, content in templates.items():
321+
(path / name).write_text(content)
322+
return path
323+
324+
325+
class TestCreateTemplateEnvZip:
326+
@pytest.mark.parametrize(
327+
("archive_name", "searchpath_parts", "template_name"),
328+
[
329+
pytest.param("dags.zip", (), "sql/query.sql", id="archive"),
330+
pytest.param("dags.zip", ("sql",), "query.sql", id="directory-in-archive"),
331+
pytest.param("DAGS.ZIP", (), "sql/query.sql", id="uppercase-extension"),
332+
],
333+
)
334+
def test_loads_template_from_zip(self, tmp_path, archive_name, searchpath_parts, template_name):
335+
archive = _write_zip(tmp_path / archive_name, {"sql/query.sql": "SELECT {{ x }}"})
336+
337+
env = create_template_env(searchpath=[os.path.join(archive, *searchpath_parts)])
338+
339+
assert env.get_template(template_name).render(x=1) == "SELECT 1"
340+
341+
@pytest.mark.parametrize(
342+
"searchpath_parts",
343+
[
344+
pytest.param(("templates",), id="directory"),
345+
pytest.param(("templates.zip",), id="directory-named-like-archive"),
346+
pytest.param(("not_a_zip.zip", "sql"), id="non-zip-file"),
347+
pytest.param(("missing", "sql"), id="missing-path"),
348+
],
349+
)
350+
def test_keeps_filesystem_loader_without_zip(self, tmp_path, searchpath_parts):
351+
(tmp_path / "templates").mkdir()
352+
(tmp_path / "templates.zip").mkdir()
353+
(tmp_path / "not_a_zip.zip").write_text("plain text")
354+
searchpath = os.path.join(tmp_path, *searchpath_parts)
355+
356+
env = create_template_env(searchpath=[searchpath])
357+
358+
assert type(env.loader) is jinja2.FileSystemLoader
359+
assert env.loader.searchpath == [searchpath]
360+
361+
@pytest.mark.parametrize(("zip_first", "expected"), [(True, "from zip"), (False, "from directory")])
362+
def test_searches_zip_and_directory_in_order(self, tmp_path, zip_first, expected):
363+
archive = str(_write_zip(tmp_path / "dags.zip", {"query.sql": "from zip"}))
364+
directory = str(_write_template_dir(tmp_path / "templates", {"query.sql": "from directory"}))
365+
366+
env = create_template_env(searchpath=[archive, directory] if zip_first else [directory, archive])
367+
368+
assert env.get_template("query.sql").render() == expected
369+
370+
def test_template_not_found_lists_search_paths(self, tmp_path):
371+
archive = str(_write_zip(tmp_path / "dags.zip", {"query.sql": "SELECT 1"}))
372+
directory = str(_write_template_dir(tmp_path / "templates", {}))
373+
env = create_template_env(searchpath=[archive, directory])
374+
375+
with pytest.raises(jinja2.TemplateNotFound) as ctx:
376+
env.get_template("missing.sql")
377+
378+
assert ctx.value.name == "missing.sql"
379+
assert ctx.value.message == f"'missing.sql' not found in search paths: {archive!r}, {directory!r}"
380+
381+
@pytest.mark.parametrize(
382+
("searchpath_parts", "template_name"),
383+
[
384+
pytest.param(("sql",), "../secret.sql", id="name-leaves-directory"),
385+
pytest.param((), "../outside.sql", id="name-leaves-archive"),
386+
pytest.param(("..",), "outside.sql", id="searchpath-leaves-archive"),
387+
],
388+
)
389+
def test_does_not_escape_search_root(self, tmp_path, searchpath_parts, template_name):
390+
archive = _write_zip(
391+
tmp_path / "dags.zip",
392+
{"sql/query.sql": "SELECT 1", "secret.sql": "secret", "../outside.sql": "outside"},
393+
)
394+
env = create_template_env(searchpath=[os.path.join(archive, *searchpath_parts)])
395+
396+
with pytest.raises(jinja2.TemplateNotFound):
397+
env.get_template(template_name)
398+
399+
def test_unreadable_archive_raises_instead_of_falling_through(self, tmp_path):
400+
archive = _write_zip(tmp_path / "dags.zip", {"query.sql": "from zip"})
401+
directory = _write_template_dir(tmp_path / "templates", {"query.sql": "from directory"})
402+
env = create_template_env(searchpath=[str(archive), str(directory)])
403+
archive.write_bytes(b"no longer a zip archive")
404+
405+
with pytest.raises(zipfile.BadZipFile):
406+
env.get_template("query.sql")
407+
408+
def test_up_to_date_check_tracks_archive_mtime(self, tmp_path):
409+
archive = _write_zip(tmp_path / "dags.zip", {"query.sql": "SELECT 1"})
410+
env = create_template_env(searchpath=[str(archive)])
411+
412+
_, filename, is_up_to_date = env.loader.get_source(env, "query.sql")
413+
414+
assert filename == os.path.join(archive, "query.sql")
415+
assert is_up_to_date()
416+
os.utime(archive, (0, 0))
417+
assert not is_up_to_date()
418+
419+
303420
@pytest.fixture
304421
def env():
305422
return SandboxedEnvironment(undefined=jinja2.StrictUndefined, cache_size=0)

‎task-sdk/tests/task_sdk/importers/test_zip_importer.py‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020

2121
import os
2222
import py_compile
23+
import textwrap
2324
import zipfile
2425
from pathlib import Path
2526
from types import SimpleNamespace
@@ -181,6 +182,37 @@ def test_import_zip_archive_with_dags(self, mock_bundle):
181182
assert dags[0].bundle_name == "test_bundle"
182183
assert len(errors) == 0
183184

185+
def test_imported_dag_resolves_template_file_from_archive(self, mock_bundle):
186+
zip_path = mock_bundle.path / "templated.zip"
187+
with zipfile.ZipFile(zip_path, "w") as z:
188+
z.writestr(
189+
"pkg/templated_dag.py",
190+
textwrap.dedent(
191+
"""\
192+
from airflow.sdk import DAG, BaseOperator
193+
194+
class SqlOperator(BaseOperator):
195+
template_fields = ("sql",)
196+
template_ext = (".sql",)
197+
198+
def __init__(self, sql, **kwargs):
199+
super().__init__(**kwargs)
200+
self.sql = sql
201+
202+
with DAG("zip_templated_dag"):
203+
SqlOperator(task_id="t", sql="sql/query.sql")
204+
"""
205+
),
206+
)
207+
z.writestr("pkg/sql/query.sql", "SELECT 1")
208+
209+
dags, errors = _import_all(ZipImporter(), mock_bundle)
210+
assert errors == []
211+
(dag,) = dags
212+
dag.resolve_template_files()
213+
214+
assert dag.get_task("t").sql == "SELECT 1"
215+
184216
def test_import_zip_archive_with_pyc_dag(self, mock_bundle, tmp_path):
185217
source_file = tmp_path / "compiled_dag.py"
186218
source_file.write_text("from airflow.sdk import DAG\ndag = DAG('zip_pyc_dag')\n")

0 commit comments

Comments
 (0)