This is an automated email from the ASF dual-hosted git repository.
kaxil pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 0c3eb9791e9 Fix Airflow 2 session leak in mapped KubernetesPodOperator
test (#73260)
0c3eb9791e9 is described below
commit 0c3eb9791e90acfd6e462cd7695f1411e54ae595
Author: Kaxil Naik <[email protected]>
AuthorDate: Thu Sep 17 03:08:51 2026 +0100
Fix Airflow 2 session leak in mapped KubernetesPodOperator test (#73260)
Rendering a mapped operator's template fields on Airflow 2 takes its session
from a module global that get_current_task_instance_session fills in on
demand
and never clears, so rendering outside a task run leaves a session behind
and
the next TaskInstance.run anywhere in the same pytest process raises
"Session
already set for this task". In the Compat 2.11.1 provider job that surfaced
as
50 failures in common.sql and docker, none of them in the provider that
caused
it.
Wrapping the render in set_current_task_instance_session clears the global
on
exit. Airflow 3 renders without a session, so its branch is a no-op and the
test keeps asserting the rendered env var on both versions.
---
.../unit/cncf/kubernetes/operators/test_pod.py | 24 +++++++++++++++++++++-
1 file changed, 23 insertions(+), 1 deletion(-)
diff --git
a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py
b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py
index 39b3eb1b31c..5161c02b945 100644
--- a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py
+++ b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py
@@ -71,6 +71,9 @@ if AIRFLOW_V_3_0_PLUS or AIRFLOW_V_3_1_PLUS:
else:
from airflow.models.xcom import XCom # type: ignore[no-redef]
+if not AIRFLOW_V_3_0_PLUS:
+ from airflow.utils.task_instance_session import
set_current_task_instance_session
+
if TYPE_CHECKING:
from airflow.sdk import Context
@@ -118,6 +121,24 @@ def _clear_all_db_objects():
db.clear_db_dag_bundles()
+@contextmanager
+def task_instance_session():
+ """
+ Provide the session Airflow 2 renders a mapped task's template fields with.
+
+ There, ``MappedOperator.render_template_fields`` takes its session from a
module global
+ that ``get_current_task_instance_session`` fills in and never clears, so
rendering outside
+ this context manager leaves a session behind and the next
``TaskInstance.run`` anywhere in
+ the process fails with "Session already set for this task". Airflow 3
renders without a
+ session, so there is nothing to set.
+ """
+ if AIRFLOW_V_3_0_PLUS:
+ yield
+ return
+ with create_session() as session,
set_current_task_instance_session(session=session):
+ yield
+
+
def create_context(task, persist_to_db=False, map_index=None):
if task.has_dag():
dag = task.dag
@@ -461,7 +482,8 @@ class TestKubernetesPodOperator:
context = create_context(mapped, map_index=0)
context.update({"dag_run": context["ti"].dag_run, "foo":
"footemplated", "bar": "bartemplated"})
- mapped.render_template_fields(context)
+ with task_instance_session():
+ mapped.render_template_fields(context)
rendered = context["task"]
assert rendered.env_vars[0].name == "bartemplated"
