Skip to content

Set context inside templates #33645

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
potiuk merged 7 commits into from
Aug 24, 2023
Merged

Set context inside templates #33645

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
67c0ae0
Set context inside templates
Aug 23, 2023
9ccf658
Fix black test
Aug 24, 2023
e790b3b
Merge branch 'main' into context-inside-templates
ivan-afonichkin Aug 24, 2023
4c227d3
Merge branch 'main' into context-inside-templates
ivan-afonichkin Aug 24, 2023
0de85d6
black fixed
Aug 24, 2023
f012d65
Update tests/models/test_taskinstance.py
ivan-afonichkin Aug 24, 2023
c9a4157
Merge branch 'main' into context-inside-templates
ivan-afonichkin Aug 24, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion airflow/models/taskinstance.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1641,7 +1641,9 @@ def signal_handler(signum, frame):
# Set the validated/merged params on the task object.
self.task.params = context["params"]

task_orig = self.render_templates(context=context)
with set_current_context(context):
task_orig = self.render_templates(context=context)

if not test_mode:
rtif = RenderedTaskInstanceFields(ti=self, render_templates=False)
RenderedTaskInstanceFields.write(rtif)
Expand Down
27 changes: 27 additions & 0 deletions tests/models/test_taskinstance.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2890,6 +2890,33 @@ def test_echo_env_variables(self, dag_maker):
ti.refresh_from_db()
assert ti.state == State.SUCCESS

def test_get_current_context_works_in_template(self, dag_maker):
def user_defined_macro():
from airflow.operators.python import get_current_context

get_current_context()

with dag_maker(
"test_context_inside_template",
start_date=DEFAULT_DATE,
end_date=DEFAULT_DATE + datetime.timedelta(days=10),
user_defined_macros={"user_defined_macro": user_defined_macro},
):

def foo(arg):
print(arg)

PythonOperator(
task_id="context_inside_template",
python_callable=foo,
op_kwargs={"arg": "{{ user_defined_macro() }}"},
),
dagrun = dag_maker.create_dagrun()
tis = dagrun.get_task_instances()
ti: TaskInstance = next(x for x in tis if x.task_id == "context_inside_template")
ti._run_raw_task()
assert ti.state == State.SUCCESS
Comment on lines +2917 to +2918

@uranusjr uranusjr Aug 24, 2023

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we need to actually run the task, or can we only test the rendering part? Less moving parts if possible.

@ivan-afonichkin ivan-afonichkin Aug 24, 2023

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's implemented as setting a context for rendering templates inside _execute_task_with_callbacks. I'm not aware perhaps, but is there a way to run _execute_task_with_callbacks without actually running a task? Or do you think there is a better place to set up this context?


@patch.object(Stats, "incr")
def test_task_stats(self, stats_mock, create_task_instance):
ti = create_task_instance(
Expand Down