Repository navigation
Fix template file lookup for Dags packaged in zip files - #74542
sjyangkevin wants to merge 4 commits into
Conversation
594297a to
849ba15
Compare
849ba15 to
9976602
Compare
| return path | ||
|
|
||
|
|
||
| class TestCreateTemplateEnvZip: |
There was a problem hiding this comment.
One test class per subject, not per area/feature we test.
i.e. move these to be in TestTemplater.
There was a problem hiding this comment.
Thanks for pointing this out. The tests under this class have been cleanup, re-wrote, and moved under TestTemplater.
| def test_template_not_found_lists_search_paths(self, tmp_path): | ||
| archive = str(_write_zip(tmp_path / "dags.zip", {"query.sql": "SELECT 1"})) | ||
| directory = str(_write_template_dir(tmp_path / "templates", {})) | ||
| env = create_template_env(searchpath=[archive, directory]) | ||
|
|
||
| with pytest.raises(jinja2.TemplateNotFound) as ctx: | ||
| env.get_template("missing.sql") | ||
|
|
||
| assert ctx.value.name == "missing.sql" | ||
| assert ctx.value.message == f"'missing.sql' not found in search paths: {archive!r}, {directory!r}" |
There was a problem hiding this comment.
This is testing Jinja behaviour, not ours.
There was a problem hiding this comment.
After the above clean up, there are now two newly added test functions under TestTemplater.
- test_loads_template_from_zip
- test_loads_templates_from_zip_and_directory_search_paths
test_loads_template_from_zip invoke create_template_env which will go to the branch for building the zip loader in _build_template_loader. Then, it checks for 4 path patterns that the query can be rendered.
test_loads_templates_from_zip_and_directory_search_paths check two scenarios. A template in the zip is found, and a template in the folder is found, when both are specifying in searchpath.
| @pytest.mark.parametrize(("zip_first", "expected"), [(True, "from zip"), (False, "from directory")]) | ||
| def test_searches_zip_and_directory_in_order(self, tmp_path, zip_first, expected): | ||
| archive = str(_write_zip(tmp_path / "dags.zip", {"query.sql": "from zip"})) | ||
| directory = str(_write_template_dir(tmp_path / "templates", {"query.sql": "from directory"})) | ||
|
|
||
| env = create_template_env(searchpath=[archive, directory] if zip_first else [directory, archive]) | ||
|
|
||
| assert env.get_template("query.sql").render() == expected |
There was a problem hiding this comment.
the test had been cleanup.
| def test_keeps_filesystem_loader_without_zip(self, tmp_path, searchpath_parts): | ||
| (tmp_path / "templates").mkdir() | ||
| (tmp_path / "templates.zip").mkdir() | ||
| (tmp_path / "not_a_zip.zip").write_text("plain text") | ||
| searchpath = os.path.join(tmp_path, *searchpath_parts) | ||
|
|
||
| env = create_template_env(searchpath=[searchpath]) | ||
|
|
||
| assert type(env.loader) is jinja2.FileSystemLoader | ||
| assert env.loader.searchpath == [searchpath] |
There was a problem hiding this comment.
This test seems almost tautological. I can't see how this could ever fail?
There was a problem hiding this comment.
thanks for pointing this out. The test had been cleanup.
ashb
left a comment
There was a problem hiding this comment.
Some work to be done on the Tests here, Agent added too many.
developer-rpai
left a comment
There was a problem hiding this comment.
Reviewed — reading templates straight out of the archive without extracting is the right approach, and opening the zip per lookup to avoid the fork offset issue is a nice catch. Tests cover the main shapes. LGTM.
a6a3900 to
680dd24
Compare
|
Hi @ashb , thank you for the feedback. I've re-structured the test, and removed those tautological tests. I wonder if I could get your help to have another look when you have time? |
a78d265 to
d91f8a8
Compare
d91f8a8 to
c42028a
Compare
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.
A Dag in a subfolder of a zip archive that builds its template_searchpath with os.path.join(os.path.dirname(__file__), "..", "sql") still could not find its templates. Nothing resolves ".." inside an archive: the OS cannot walk into the archive file, and archive members are matched by exact name, so the path has to be normalized before it is split into the archive and the folder inside it. The tests are reduced to the cases that fail without the fix, as asked in review. The removed ones also passed on main, because they checked behaviour that comes from Jinja rather than from this change.
The template loader for zip search paths is kept instead of Jinja's ChoiceLoader because its TemplateNotFound error lists the paths that were searched, which is what a Dag author needs when a template is missing from the archive. ChoiceLoader would only report the template name, so the lookup test now also checks the error for a template found nowhere.
The error for a template found in no search path is the same as on main, so checking it tested behaviour this change does not touch, which is what the review pointed out. The test still checks that a zipped Dag loads templates from both its archive and a regular template_searchpath folder.
c42028a to
6301f7f
Compare
Why
Dags packaged as zip archives could not use template files shipped inside the archive. For a zipped Dag,
dag.folderpoints into the archive (for example.../dags.zip), and Jinja'sFileSystemLoaderonly reads real directories. A templated field such as"query": "test_sql/test.sql"therefore failed every run with:At parse time the same lookup failed silently (only logged), and pointing
template_searchpathinto the archive failed the same way.create_template_env()now reads search-path entries that are, or point inside, a zip archive directly from the archive. Normal directories keep exactly the sameFileSystemLoaderas before. Both template environment builders (DAG.get_template_env()and the operator-levelrender_template_as_native_objoverride) go through this function.closes: #59310
related: #60999
How zip entries are handled
zipfile.is_zipfile(), so detection does not depend on the.zipextension or its case.split_template_path, as inFileSystemLoader, so they cannot leave the search root. Search-path prefixes containing..are not treated as archive paths.TemplateNotFoundwith thenot found in search path(s): ...message.Testing
18 new test cases. 9 fail without this change: loading from a zip (includingDAGS.ZIP), zip-first ordering, both template environment builders with a.sqlpath nested in a dict, and parse-timeresolve_template_files()on a Dag imported by the zip importer. The other 9 pass onmainby design: they pin unchanged behaviour for non-zip search paths, the not-found message, and containment of template names. Each was checked to fail when the corresponding check in the new loader is removed.airflow dags testand the UI frombreeze start-airflow, with zipped Dags using a top-level template file, a.sqlpath nested in a dict,template_searchpathinside the archive, and{% include %}. All failed withTemplateNotFoundbefore this change and succeed after it. Unzipped copies of the same Dags succeed both before and after.BigQueryInsertJobOperator.execute(). The BigQuery call itself was not run (no GCP credentials), and this change does not touch it.mypy-task-sdk. Frombreeze verify:task-sdk-tests(3656 passed), core non-DB tests and core DB tests pass.Before/after reproduction in Breeze (
airflow dags test)repro.zipcontainsrepro_dags.py,scripts/hello.shandsql/query.sql. The same files unzipped provide thedir_*control Dags.reporter.zipis thetest.zipattached to the issue, unchanged.zip_top_levelbash_command="scripts/hello.sh"zip_nested_dict.sqlpath nested in a dict field, likeBigQueryInsertJobOperatorzip_searchpathtemplate_searchpathpointing to a folder inside the archivezip_include{% include 'scripts/hello.sh' %}test_import_path_dagdir_top_levelBefore (
main), trimmed to each Dag's result and its own task error:After (this PR):
test_import_path_dagnow gets through template rendering and fails inBigQueryInsertJobOperator.execute()withAirflowNotFoundException: The conn_id `google_cloud_default` isn't defined, since no GCP connection.repro_dags.pyscripts/hello.shisecho "hello from {{ dag.dag_id }} / {{ ti.task_id }}"andsql/query.sqlisSELECT column_a FROM test -- {{ dag.dag_id }}.Was generative AI tooling used to co-author this PR?
Generated-by: Claude Code (Opus 5.5) following the guidelines