|
17 | 17 |
|
18 | 18 | from __future__ import annotations |
19 | 19 |
|
| 20 | +import os |
| 21 | +import zipfile |
20 | 22 | from datetime import UTC, datetime |
| 23 | +from pathlib import Path |
21 | 24 | from unittest.mock import MagicMock, NonCallableMagicMock |
22 | 25 |
|
23 | 26 | import jinja2 |
24 | 27 | import pytest |
25 | 28 |
|
26 | 29 | 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 | +) |
28 | 36 |
|
29 | 37 |
|
30 | 38 | class TestTemplater: |
@@ -300,6 +308,115 @@ def test_do_render_template_fields_renders_list_values(self): |
300 | 308 | assert parent.items == ["first", "second"] |
301 | 309 |
|
302 | 310 |
|
| 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 | + |
303 | 420 | @pytest.fixture |
304 | 421 | def env(): |
305 | 422 | return SandboxedEnvironment(undefined=jinja2.StrictUndefined, cache_size=0) |
|
0 commit comments