This is an automated email from the ASF dual-hosted git repository.
o-nikolas 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 f9dca510f9f Fix global env vars shadowing team secrets in tasks and
triggers (#75998)
f9dca510f9f is described below
commit f9dca510f9f9ea8ac9652919dcce941e7286a680
Author: Niko Oliveira <[email protected]>
AuthorDate: Thu Nov 12 09:09:41 2026 -0700
Fix global env vars shadowing team secrets in tasks and triggers (#75998)
In multi-team mode, tasks and triggers looked up Variables and
Connections without passing their team to the secrets backends. The
environment variables backend is checked first, so a global
AIRFLOW_VAR_<KEY> was returned even when AIRFLOW_VAR__<TEAM>___<KEY>
was set. The docs say the team value takes precedence, and the
Execution API, which does resolve the team, was never reached.
---
.../src/airflow/jobs/triggerer_job_runner.py | 3 +-
airflow-core/tests/unit/jobs/test_triggerer_job.py | 15 ++++
.../edge3/example_dags/integration_test.py | 2 +-
.../providers/edge3/example_dags/win_test.py | 2 +-
task-sdk/src/airflow/sdk/execution_time/context.py | 80 ++++++++++++++++------
.../src/airflow/sdk/execution_time/task_runner.py | 6 ++
.../tests/task_sdk/definitions/test_connection.py | 2 +-
.../tests/task_sdk/definitions/test_variables.py | 2 +-
.../tests/task_sdk/execution_time/test_context.py | 58 ++++++++++++++++
.../task_sdk/execution_time/test_context_cache.py | 16 ++++-
.../task_sdk/execution_time/test_task_runner.py | 43 +++++++++++-
11 files changed, 201 insertions(+), 28 deletions(-)
diff --git a/airflow-core/src/airflow/jobs/triggerer_job_runner.py
b/airflow-core/src/airflow/jobs/triggerer_job_runner.py
index 649b8ca8276..afb50da0ef2 100644
--- a/airflow-core/src/airflow/jobs/triggerer_job_runner.py
+++ b/airflow-core/src/airflow/jobs/triggerer_job_runner.py
@@ -108,7 +108,7 @@ from airflow.sdk.execution_time.comms import (
_new_encoder,
_RequestFrame,
)
-from airflow.sdk.execution_time.context import AssetStateStoreAccessors
+from airflow.sdk.execution_time.context import AssetStateStoreAccessors,
set_current_team_name
from airflow.sdk.execution_time.request_handlers import (
handle_clear_asset_state_store_by_name,
handle_clear_asset_state_store_by_uri,
@@ -1611,6 +1611,7 @@ class TriggerRunner:
raise RuntimeError(f"Required first message to be a
messages.StartTriggerer, it was {msg}")
self.team_name = msg.team_name
+ set_current_team_name(msg.team_name)
await self.comms_decoder.start_reader()
diff --git a/airflow-core/tests/unit/jobs/test_triggerer_job.py
b/airflow-core/tests/unit/jobs/test_triggerer_job.py
index e38fbe0343b..2a1f6265f65 100644
--- a/airflow-core/tests/unit/jobs/test_triggerer_job.py
+++ b/airflow-core/tests/unit/jobs/test_triggerer_job.py
@@ -1546,6 +1546,21 @@ def
test_trigger_logger_fd_closed_when_upload_to_remote_raises(jobless_superviso
class TestTriggerRunner:
+ @pytest.mark.asyncio
+ @mock.patch("airflow.jobs.triggerer_job_runner.set_current_team_name",
autospec=True)
+ @mock.patch("airflow.sdk.execution_time.task_runner.SUPERVISOR_COMMS",
create=True)
+ @mock.patch("airflow.jobs.triggerer_job_runner.TriggerCommsDecoder",
autospec=True)
+ @mock.patch("asyncio.open_connection", autospec=True, return_value=(None,
None))
+ @mock.patch("airflow.jobs.triggerer_job_runner.socket", autospec=True)
+ async def test_init_comms_sets_the_current_team(
+ self, _socket, _open_connection, mock_decoder_cls, _supervisor_comms,
mock_set_current_team_name
+ ) -> None:
+ mock_decoder_cls.return_value._from_frame.return_value =
messages.StartTriggerer(team_name="team_a")
+
+ await TriggerRunner().init_comms()
+
+ mock_set_current_team_name.assert_called_once_with("team_a")
+
def test_blocked_main_thread_warning_threshold_decode(self) -> None:
with conf_vars({("triggerer",
"blocked_main_thread_warning_threshold"): "0.5"}):
trigger_runner = TriggerRunner()
diff --git
a/providers/edge3/src/airflow/providers/edge3/example_dags/integration_test.py
b/providers/edge3/src/airflow/providers/edge3/example_dags/integration_test.py
index cea09293c30..2d1bf2a4fa4 100644
---
a/providers/edge3/src/airflow/providers/edge3/example_dags/integration_test.py
+++
b/providers/edge3/src/airflow/providers/edge3/example_dags/integration_test.py
@@ -128,7 +128,7 @@ with DAG(
def connection():
try:
conn = BaseHook.get_connection("integration_test")
- print(f"Got connection {conn}")
+ print(f"Got connection {conn.conn_id}")
except AirflowNotFoundException:
print("Connection not found... but also OK.")
diff --git
a/providers/edge3/src/airflow/providers/edge3/example_dags/win_test.py
b/providers/edge3/src/airflow/providers/edge3/example_dags/win_test.py
index 43fd6f40ee5..47a4fff934a 100644
--- a/providers/edge3/src/airflow/providers/edge3/example_dags/win_test.py
+++ b/providers/edge3/src/airflow/providers/edge3/example_dags/win_test.py
@@ -311,7 +311,7 @@ with DAG(
def connection():
try:
conn = BaseHook.get_connection("integration_test")
- print(f"Got connection {conn}")
+ print(f"Got connection {conn.conn_id}")
except AirflowNotFoundException:
print("Connection 'integration_test' not found... but also
OK.")
diff --git a/task-sdk/src/airflow/sdk/execution_time/context.py
b/task-sdk/src/airflow/sdk/execution_time/context.py
index e9a36aba790..60c245bf5a7 100644
--- a/task-sdk/src/airflow/sdk/execution_time/context.py
+++ b/task-sdk/src/airflow/sdk/execution_time/context.py
@@ -22,6 +22,7 @@ import contextlib
import functools
import inspect
import json
+import os
import warnings
from collections.abc import Generator, Iterable, Iterator, Mapping, Sequence
from contextvars import ContextVar
@@ -35,6 +36,7 @@ import structlog
from pydantic import AwareDatetime, TypeAdapter
from airflow.sdk._shared.observability.metrics import stats
+from airflow.sdk._shared.secrets_backend.base import
call_secrets_backend_method
from airflow.sdk._shared.state import AssetScope
from airflow.sdk.configuration import conf
from airflow.sdk.definitions._internal.contextmanager import _CURRENT_CONTEXT
@@ -234,6 +236,22 @@ _preset_connections: ContextVar[dict[str, Connection] |
None] = ContextVar(
"_preset_connections", default=None
)
+# A plain global rather than a ContextVar, so threads started by task code see
it too.
+# Also exported to the environment, so child interpreters such as virtualenv
tasks can
+# restore it in reinit_supervisor_comms.
+_current_team_name: str | None = None
+_CURRENT_TEAM_NAME_ENV_VAR = "_AIRFLOW__CURRENT_TEAM_NAME"
+
+
+def set_current_team_name(team_name: str | None) -> None:
+ """Set the team that Connection and Variable lookups in this process and
its children belong to."""
+ global _current_team_name
+ _current_team_name = team_name
+ if team_name:
+ os.environ[_CURRENT_TEAM_NAME_ENV_VAR] = team_name
+ else:
+ os.environ.pop(_CURRENT_TEAM_NAME_ENV_VAR, None)
+
def _get_connection(conn_id: str) -> Connection:
preset = _preset_connections.get()
@@ -247,7 +265,7 @@ def _get_connection(conn_id: str) -> Connection:
# Check cache first (optional; only on dag processor)
try:
- uri = SecretCache.get_connection_uri(conn_id)
+ uri = SecretCache.get_connection_uri(conn_id,
team_name=_current_team_name)
from airflow.sdk.definitions.connection import Connection
conn = Connection.from_uri(uri, conn_id=conn_id)
@@ -261,9 +279,11 @@ def _get_connection(conn_id: str) -> Connection:
backends = ensure_secrets_backend_loaded()
for secrets_backend in backends:
try:
- conn = secrets_backend.get_connection(conn_id=conn_id) # type:
ignore[assignment]
+ conn = call_secrets_backend_method(
+ secrets_backend.get_connection, team_name=_current_team_name,
conn_id=conn_id
+ )
if conn:
- SecretCache.save_connection_uri(conn_id, conn.get_uri())
+ SecretCache.save_connection_uri(conn_id, conn.get_uri(),
team_name=_current_team_name)
_mask_connection_secrets(conn)
return conn
except AirflowSecretsBackendAccessDenied:
@@ -295,7 +315,7 @@ async def _async_get_connection(conn_id: str) -> Connection:
# Check cache first
try:
- uri = SecretCache.get_connection_uri(conn_id)
+ uri = SecretCache.get_connection_uri(conn_id,
team_name=_current_team_name)
from airflow.sdk.definitions.connection import Connection
conn = Connection.from_uri(uri, conn_id=conn_id)
@@ -314,10 +334,14 @@ async def _async_get_connection(conn_id: str) ->
Connection:
# getattr avoids triggering AsyncMock coroutine creation under
Python 3.13
async_method = getattr(secrets_backend, "aget_connection", None)
if async_method is not None:
- maybe_awaitable = async_method(conn_id)
+ maybe_awaitable = call_secrets_backend_method(
+ async_method, team_name=_current_team_name, conn_id=conn_id
+ )
conn = await maybe_awaitable if
inspect.isawaitable(maybe_awaitable) else maybe_awaitable
else:
- conn = await
sync_to_async(secrets_backend.get_connection)(conn_id) # type:
ignore[assignment]
+ conn = await sync_to_async(call_secrets_backend_method)(
+ secrets_backend.get_connection,
team_name=_current_team_name, conn_id=conn_id
+ )
if conn:
# Use aget_uri if the returned connection object supports it
(the SDK's own
@@ -329,7 +353,7 @@ async def _async_get_connection(conn_id: str) -> Connection:
uri = await aget_uri()
else:
uri = await sync_to_async(conn.get_uri)()
- SecretCache.save_connection_uri(conn_id, uri)
+ SecretCache.save_connection_uri(conn_id, uri,
team_name=_current_team_name)
await _amask_connection_secrets(conn)
return conn
except AirflowSecretsBackendAccessDenied:
@@ -389,7 +413,7 @@ def _get_variable(key: str, deserialize_json: bool) -> Any:
# Check cache first
try:
- var_val = SecretCache.get_variable(key)
+ var_val = SecretCache.get_variable(key, team_name=_current_team_name)
if var_val is not None:
return _mask_and_deserialize_variable(var_val, key,
deserialize_json)
except SecretCache.NotPresentException:
@@ -400,10 +424,12 @@ def _get_variable(key: str, deserialize_json: bool) ->
Any:
# Iterate over backends if not in cache (or expired)
for secrets_backend in backends:
try:
- var_val = secrets_backend.get_variable(key=key)
+ var_val = call_secrets_backend_method(
+ secrets_backend.get_variable, team_name=_current_team_name,
key=key
+ )
if var_val is not None:
# Save raw value before deserialization to maintain cache
consistency
- SecretCache.save_variable(key, var_val)
+ SecretCache.save_variable(key, var_val,
team_name=_current_team_name)
return _mask_and_deserialize_variable(var_val, key,
deserialize_json)
except AirflowSecretsBackendAccessDenied:
# Authoritative deny — must NOT fall through to a less-restrictive
backend.
@@ -428,7 +454,7 @@ async def _async_get_variable(key: str, deserialize_json:
bool) -> Any:
# Check cache first
try:
- var_val = SecretCache.get_variable(key)
+ var_val = SecretCache.get_variable(key, team_name=_current_team_name)
if var_val is not None:
return await _async_mask_and_deserialize_variable(var_val, key,
deserialize_json)
except SecretCache.NotPresentException:
@@ -441,12 +467,19 @@ async def _async_get_variable(key: str, deserialize_json:
bool) -> Any:
try:
async_method = getattr(secrets_backend, "aget_variable", None)
if async_method is not None:
- var_val = await async_method(key=key)
+ var_val = await call_secrets_backend_method(
+ async_method, team_name=_current_team_name, key=key
+ )
else:
- var_val = await
asyncio.to_thread(secrets_backend.get_variable, key=key)
+ var_val = await asyncio.to_thread(
+ call_secrets_backend_method,
+ secrets_backend.get_variable,
+ team_name=_current_team_name,
+ key=key,
+ )
if var_val is not None:
# Save raw value before deserialization to maintain cache
consistency
- SecretCache.save_variable(key, var_val)
+ SecretCache.save_variable(key, var_val,
team_name=_current_team_name)
return await _async_mask_and_deserialize_variable(var_val,
key, deserialize_json)
except AirflowSecretsBackendAccessDenied:
# Authoritative deny — must NOT fall through to a less-restrictive
backend.
@@ -538,7 +571,9 @@ def _set_variable(key: str, value: Any, description: str |
None = None, serializ
if isinstance(secrets_backend, ExecutionAPISecretsBackend):
continue
try:
- var_val = secrets_backend.get_variable(key=key)
+ var_val = call_secrets_backend_method(
+ secrets_backend.get_variable, team_name=_current_team_name,
key=key
+ )
if var_val is not None:
_backend_name = type(secrets_backend).__name__
log.warning(
@@ -569,7 +604,7 @@ def _set_variable(key: str, value: Any, description: str |
None = None, serializ
SUPERVISOR_COMMS.send(PutVariable(key=key, value=value,
description=description))
# Invalidate cache after setting the variable
- SecretCache.invalidate_variable(key)
+ SecretCache.invalidate_variable(key, team_name=_current_team_name)
async def _async_set_variable(
@@ -595,7 +630,12 @@ async def _async_set_variable(
if isinstance(secrets_backend, ExecutionAPISecretsBackend):
continue
try:
- var_val = await asyncio.to_thread(secrets_backend.get_variable,
key=key)
+ var_val = await asyncio.to_thread(
+ call_secrets_backend_method,
+ secrets_backend.get_variable,
+ team_name=_current_team_name,
+ key=key,
+ )
if var_val is not None:
_backend_name = type(secrets_backend).__name__
log.warning(
@@ -626,7 +666,7 @@ async def _async_set_variable(
await SUPERVISOR_COMMS.asend(PutVariable(key=key, value=value,
description=description))
# Invalidate cache after setting the variable
- SecretCache.invalidate_variable(key)
+ SecretCache.invalidate_variable(key, team_name=_current_team_name)
def _delete_variable(key: str) -> None:
@@ -643,7 +683,7 @@ def _delete_variable(key: str) -> None:
assert isinstance(msg, OKResponse)
# Invalidate cache after deleting the variable
- SecretCache.invalidate_variable(key)
+ SecretCache.invalidate_variable(key, team_name=_current_team_name)
async def _async_delete_variable(key: str) -> None:
@@ -660,7 +700,7 @@ async def _async_delete_variable(key: str) -> None:
assert isinstance(msg, OKResponse)
# Invalidate cache after deleting the variable
- SecretCache.invalidate_variable(key)
+ SecretCache.invalidate_variable(key, team_name=_current_team_name)
class ConnectionAccessor:
diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py
b/task-sdk/src/airflow/sdk/execution_time/task_runner.py
index c89cdeed176..a173c734670 100644
--- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py
+++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py
@@ -143,6 +143,7 @@ from airflow.sdk.execution_time.comms import (
ValidateInletsAndOutlets,
)
from airflow.sdk.execution_time.context import (
+ _CURRENT_TEAM_NAME_ENV_VAR,
AssetStateStoreAccessors,
ConnectionAccessor,
InletEventsAccessors,
@@ -156,6 +157,7 @@ from airflow.sdk.execution_time.context import (
context_to_airflow_vars,
get_previous_dagrun_success,
set_current_context,
+ set_current_team_name,
)
from airflow.sdk.execution_time.email_backend import (
_DEFAULT_EMAIL_BACKEND,
@@ -1305,6 +1307,8 @@ def startup(msg: StartupDetails) ->
tuple[RuntimeTaskInstance, Context, Logger]:
except Exception:
log.exception("error calling listener")
+ set_current_team_name(msg.ti_context.dag_run.team_name)
+
with _airflow_parsing_context_manager(dag_id=msg.ti.dag_id,
task_id=msg.ti.task_id):
ti = parse(msg, log)
log.debug("Dag file parsed", file=msg.dag_rel_path)
@@ -2771,6 +2775,8 @@ def reinit_supervisor_comms() -> None:
"""
import socket
+ set_current_team_name(os.environ.get(_CURRENT_TEAM_NAME_ENV_VAR))
+
if "SUPERVISOR_COMMS" not in globals():
global SUPERVISOR_COMMS
diff --git a/task-sdk/tests/task_sdk/definitions/test_connection.py
b/task-sdk/tests/task_sdk/definitions/test_connection.py
index aad4c8f46db..0f40e0c5ea6 100644
--- a/task-sdk/tests/task_sdk/definitions/test_connection.py
+++ b/task-sdk/tests/task_sdk/definitions/test_connection.py
@@ -378,7 +378,7 @@ class TestConnectionsFromSecrets:
"""Tests getting a connection from environment variable."""
mock_env_get.return_value = Connection(conn_id="something",
conn_type="some-type") # return None
Connection.get("something")
- mock_env_get.assert_called_once_with(conn_id="something")
+ mock_env_get.assert_called_once_with(conn_id="something",
team_name=None)
@conf_vars(
{
diff --git a/task-sdk/tests/task_sdk/definitions/test_variables.py
b/task-sdk/tests/task_sdk/definitions/test_variables.py
index 77f5245ee78..dd8a8a3dfbf 100644
--- a/task-sdk/tests/task_sdk/definitions/test_variables.py
+++ b/task-sdk/tests/task_sdk/definitions/test_variables.py
@@ -444,7 +444,7 @@ class TestVariableFromSecrets:
"""Tests getting a variable from environment variable."""
mock_env_get.return_value = "fake_value"
Variable.get(key="fake_var_key")
- mock_env_get.assert_called_once_with(key="fake_var_key")
+ mock_env_get.assert_called_once_with(key="fake_var_key",
team_name=None)
@conf_vars(
{
diff --git a/task-sdk/tests/task_sdk/execution_time/test_context.py
b/task-sdk/tests/task_sdk/execution_time/test_context.py
index 9dba8531117..1e711074c65 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_context.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_context.py
@@ -110,8 +110,10 @@ from airflow.sdk.execution_time.context import (
_async_get_variable_keys,
_async_set_variable,
_get_connection,
+ _get_variable,
_json_size_exceeds,
_process_connection_result_conn,
+ _set_variable,
_wrap_external_ref,
context_to_airflow_vars,
set_current_context,
@@ -1813,6 +1815,62 @@ class TestJsonSizeExceeds:
assert _json_size_exceeds(value, 10) is True
+class TestTeamSecretLookups:
+ """Lookups pass the current team, so a team's env var wins over a global
one of the same name."""
+
+ @pytest.fixture(autouse=True)
+ def _team_a_and_global_env_vars(self, monkeypatch):
+ from airflow.secrets.environment_variables import
EnvironmentVariablesBackend
+
+ monkeypatch.setenv("AIRFLOW_VAR_MY_VAR", "global_value")
+ monkeypatch.setenv("AIRFLOW_VAR__TEAM_A___MY_VAR", "team_a_value")
+ monkeypatch.setenv("AIRFLOW_CONN_MY_CONN", "http://global.example.com";)
+ monkeypatch.setenv("AIRFLOW_CONN__TEAM_A___MY_CONN",
"http://team-a.example.com";)
+
monkeypatch.setattr("airflow.sdk.execution_time.context._current_team_name",
"team_a")
+
+ backend = EnvironmentVariablesBackend()
+ backend._set_connection_class(Connection)
+ monkeypatch.setattr(
+
"airflow.sdk.execution_time.supervisor.ensure_secrets_backend_loaded", lambda:
[backend]
+ )
+
+ def test_sync_lookups(self):
+ assert _get_variable("my_var", deserialize_json=False) ==
"team_a_value"
+ assert _get_connection("my_conn").host == "team-a.example.com"
+
+ @pytest.mark.asyncio
+ async def test_async_lookups(self):
+ assert await _async_get_variable("my_var", deserialize_json=False) ==
"team_a_value"
+ assert (await _async_get_connection("my_conn")).host ==
"team-a.example.com"
+
+ @pytest.mark.asyncio
+ async def test_async_backend_methods(self, monkeypatch):
+ class AsyncBackend:
+ async def aget_variable(self, key, team_name=None):
+ return f"{team_name}_value"
+
+ async def aget_connection(self, conn_id, team_name=None):
+ return Connection(conn_id=conn_id, conn_type="http",
host=f"{team_name}.example.com")
+
+ monkeypatch.setattr(
+
"airflow.sdk.execution_time.supervisor.ensure_secrets_backend_loaded", lambda:
[AsyncBackend()]
+ )
+
+ assert await _async_get_variable("my_var", deserialize_json=False) ==
"team_a_value"
+ assert (await _async_get_connection("my_conn")).host ==
"team_a.example.com"
+
+ @pytest.mark.asyncio
+ @mock.patch("airflow.sdk.execution_time.context.log")
+ async def test_set_variable_warns_about_team_env_var(self, mock_log,
monkeypatch, mock_supervisor_comms):
+ monkeypatch.delenv("AIRFLOW_VAR_MY_VAR")
+ mock_supervisor_comms.asend.return_value = None
+
+ _set_variable("my_var", "new_value")
+ await _async_set_variable("my_var", "new_value")
+
+ assert mock_log.warning.call_count == 2
+
+
class TestTaskStateStoreAccessor:
TI_ID = UUID("01900000-0000-0000-0000-000000000001")
SCOPE = TaskScope(dag_id="dag", run_id="run", task_id="task")
diff --git a/task-sdk/tests/task_sdk/execution_time/test_context_cache.py
b/task-sdk/tests/task_sdk/execution_time/test_context_cache.py
index 58c9c51e153..924466bbaf4 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_context_cache.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_context_cache.py
@@ -79,7 +79,7 @@ class TestConnectionCacheIntegration:
result = _get_connection(conn_id)
assert result.conn_id == conn_id
assert result.conn_type == "mysql"
- mock_backend.get_connection.assert_called_once_with(conn_id=conn_id)
+ mock_backend.get_connection.assert_called_once_with(conn_id=conn_id,
team_name=None)
cached_uri = SecretCache.get_connection_uri(conn_id)
cached_conn = Connection.from_uri(cached_uri, conn_id=conn_id)
@@ -180,10 +180,22 @@ class TestVariableCacheIntegration:
result = _get_variable(key, deserialize_json=False)
assert result == value
- mock_backend.get_variable.assert_called_once_with(key=key)
+ mock_backend.get_variable.assert_called_once_with(key=key,
team_name=None)
cached_value = SecretCache.get_variable(key)
assert cached_value == value
+
@patch("airflow.sdk.execution_time.supervisor.ensure_secrets_backend_loaded")
+ def test_get_variable_cache_is_keyed_by_team(self, mock_ensure_backends,
monkeypatch):
+
monkeypatch.setattr("airflow.sdk.execution_time.context._current_team_name",
"team_a")
+ SecretCache.save_variable("my_var", "global_value")
+
+ mock_backend = MagicMock(spec=["get_variable"])
+ mock_backend.get_variable.return_value = "team_a_value"
+ mock_ensure_backends.return_value = [mock_backend]
+
+ assert _get_variable("my_var", deserialize_json=False) ==
"team_a_value"
+ assert SecretCache.get_variable("my_var", team_name="team_a") ==
"team_a_value"
+
@patch("airflow.sdk.execution_time.supervisor.ensure_secrets_backend_loaded")
def test_get_variable_from_api_saves_to_cache(self, mock_ensure_backends,
mock_supervisor_comms):
"""Test that variable from API server is saved to cache."""
diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
index 14c798def40..64beb944057 100644
--- a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
+++ b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py
@@ -111,7 +111,7 @@ from airflow.sdk.exceptions import (
ErrorType,
TaskDeferred,
)
-from airflow.sdk.execution_time import task_runner
+from airflow.sdk.execution_time import context as context_module, task_runner
from airflow.sdk.execution_time.comms import (
AssetEventResult,
AssetEventsResult,
@@ -181,6 +181,7 @@ from airflow.sdk.execution_time.context import (
TriggeringAssetEventsAccessor,
VariableAccessor,
_wrap_external_ref,
+ set_current_team_name,
)
from airflow.sdk.execution_time.lazy_sequence import LazyXComSequence
from airflow.sdk.execution_time.task_runner import (
@@ -200,6 +201,7 @@ from airflow.sdk.execution_time.task_runner import (
finalize,
get_startup_details,
parse,
+ reinit_supervisor_comms,
run,
startup,
)
@@ -2283,6 +2285,45 @@ def test_startup_and_run_dag_with_rtif(
mock_supervisor_comms.assert_has_calls(expected_calls)
[email protected]("airflow.sdk.execution_time.task_runner.set_current_team_name",
autospec=True)
+def test_startup_sets_the_current_team(mock_set_current_team_name,
mocked_parse, make_ti_context):
+ ti_context = make_ti_context()
+ ti_context.dag_run.team_name = "team_a"
+ what = StartupDetails(
+ ti=TaskInstance(
+ id=uuid7(),
+ task_id="my_task",
+ dag_id="basic_dag",
+ run_id="c",
+ try_number=1,
+ dag_version_id=uuid7(),
+ queue="default",
+ ),
+ dag_rel_path="",
+ bundle_info=FAKE_BUNDLE,
+ ti_context=ti_context,
+ start_date=timezone.utcnow(),
+ sentry_integration="",
+ )
+ mocked_parse(what, "basic_dag", CustomOperator(task_id="my_task"))
+
+ startup(what)
+
+ mock_set_current_team_name.assert_called_once_with("team_a")
+
+
[email protected](os.environ)
+def test_reinit_supervisor_comms_restores_the_current_team(monkeypatch,
mock_supervisor_comms):
+ monkeypatch.setattr(context_module, "_current_team_name", None)
+ set_current_team_name("team_a")
+ # A child interpreter starts with only the environment it inherited.
+ monkeypatch.setattr(context_module, "_current_team_name", None)
+
+ reinit_supervisor_comms()
+
+ assert context_module._current_team_name == "team_a"
+
+
@patch("os.execvp")
@patch("os.set_inheritable")
def test_task_run_with_user_impersonation(
