From 67c0ae0cdfc5036cae231d3d3c25b9cddacf3a40 Mon Sep 17 00:00:00 2001 From: Ivan Afonichkin Date: Wed, 23 Aug 2023 08:15:07 +0100 Subject: [PATCH 1/4] Set context inside templates --- airflow/models/taskinstance.py | 4 +++- tests/models/test_taskinstance.py | 22 ++++++++++++++++++++++ 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py index 975e615d95414..0bd9d3d79d2a3 100644 --- a/airflow/models/taskinstance.py +++ b/airflow/models/taskinstance.py @@ -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) diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py index ac2f7f71748ef..a77db96516a65 100644 --- a/tests/models/test_taskinstance.py +++ b/tests/models/test_taskinstance.py @@ -2890,6 +2890,28 @@ def test_echo_env_variables(self, dag_maker): ti.refresh_from_db() assert ti.state == State.SUCCESS + def test_context_inside_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 + @patch.object(Stats, "incr") def test_task_stats(self, stats_mock, create_task_instance): ti = create_task_instance( From 9ccf658d681d3b1ec3b210010fec442df41f68cb Mon Sep 17 00:00:00 2001 From: Ivan Afonichkin Date: Thu, 24 Aug 2023 07:56:09 +0100 Subject: [PATCH 2/4] Fix black test --- tests/models/test_taskinstance.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py index a77db96516a65..e728b3d6792ee 100644 --- a/tests/models/test_taskinstance.py +++ b/tests/models/test_taskinstance.py @@ -2893,19 +2893,22 @@ def test_echo_env_variables(self, dag_maker): def test_context_inside_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} + 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() }}"}) + 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") From 0de85d6ff6701c26daf5bf413a6f05513234d54a Mon Sep 17 00:00:00 2001 From: Ivan Afonichkin Date: Thu, 24 Aug 2023 09:58:51 +0100 Subject: [PATCH 3/4] black fixed --- tests/models/test_taskinstance.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py index e728b3d6792ee..fcf21d475f998 100644 --- a/tests/models/test_taskinstance.py +++ b/tests/models/test_taskinstance.py @@ -2902,13 +2902,15 @@ def user_defined_macro(): 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() }}"}), + 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") From f012d6518ff9782e3734afd77fc1b66ae597e8f3 Mon Sep 17 00:00:00 2001 From: Ivan Date: Thu, 24 Aug 2023 10:48:01 +0100 Subject: [PATCH 4/4] Update tests/models/test_taskinstance.py Co-authored-by: Tzu-ping Chung --- tests/models/test_taskinstance.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/models/test_taskinstance.py b/tests/models/test_taskinstance.py index fcf21d475f998..e50917e101074 100644 --- a/tests/models/test_taskinstance.py +++ b/tests/models/test_taskinstance.py @@ -2890,7 +2890,7 @@ def test_echo_env_variables(self, dag_maker): ti.refresh_from_db() assert ti.state == State.SUCCESS - def test_context_inside_template(self, dag_maker): + def test_get_current_context_works_in_template(self, dag_maker): def user_defined_macro(): from airflow.operators.python import get_current_context