Skip to content

Commit

Permalink
fix: disable jinja2 rednering for doc_md (#40522)
Browse files Browse the repository at this point in the history
  • Loading branch information
Lee-W authored Jul 2, 2024
1 parent 09bbc9c commit 8159f6e
Show file tree
Hide file tree
Showing 2 changed files with 13 additions and 22 deletions.
19 changes: 6 additions & 13 deletions airflow/models/dag.py
Original file line number Diff line number Diff line change
Expand Up @@ -772,20 +772,13 @@ def get_doc_md(self, doc_md: str | None) -> str | None:
if doc_md is None:
return doc_md

env = self.get_template_env(force_sandboxed=True)

if not doc_md.endswith(".md"):
template = jinja2.Template(doc_md)
else:
if doc_md.endswith(".md"):
try:
template = env.get_template(doc_md)
except jinja2.exceptions.TemplateNotFound:
return f"""
# Templating Error!
Not able to find the template file: `{doc_md}`.
"""

return template.render()
return open(doc_md).read()
except FileNotFoundError:
return doc_md

return doc_md

def _check_schedule_interval_matches_timetable(self) -> bool:
"""
Expand Down
16 changes: 7 additions & 9 deletions tests/models/test_dag.py
Original file line number Diff line number Diff line change
Expand Up @@ -3274,27 +3274,25 @@ def noop_pipeline():
assert dag.dag_id == "noop_pipeline"
assert "Regular DAG documentation" in dag.doc_md

def test_resolve_documentation_template_file_rendered(self, tmp_path):
def test_resolve_documentation_template_file_not_rendered(self, tmp_path):
"""Test that @dag uses function docs as doc_md for DAG object"""

path = tmp_path / "testfile.md"
path.write_text(
"""
raw_content = """
{% if True %}
External Markdown DAG documentation
{% endif %}
"""
)

@dag_decorator(
"test-dag", start_date=DEFAULT_DATE, template_searchpath=os.fspath(path.parent), doc_md=path.name
)
path = tmp_path / "testfile.md"
path.write_text(raw_content)

@dag_decorator("test-dag", start_date=DEFAULT_DATE, doc_md=str(path))
def markdown_docs(): ...

dag = markdown_docs()
assert isinstance(dag, DAG)
assert dag.dag_id == "test-dag"
assert dag.doc_md.strip() == "External Markdown DAG documentation"
assert dag.doc_md == raw_content

def test_fails_if_arg_not_set(self):
"""Test that @dag decorated function fails if positional argument is not set"""
Expand Down

0 comments on commit 8159f6e

Please sign in to comment.