Set context inside templates (#33645)

* Set context inside templates --------- Co-authored-by: Ivan Afonichkin <[email protected]> Co-authored-by: Tzu-ping Chung <[email protected]> (cherry picked from commit 9fa782f)
apache · Aug 28, 2023 · b0f9afb · b0f9afb
1 parent e6f6dea
commit b0f9afb
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 1 deletion.
diff --git 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
@@ -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
+
     @patch.object(Stats, "incr")
     def test_task_stats(self, stats_mock, create_task_instance):
         ti = create_task_instance(