vincbeck commented on code in PR #75998:
URL: https://github.com/apache/airflow/pull/75998#discussion_r4425146905
##########
task-sdk/src/airflow/sdk/execution_time/context.py:
##########
@@ -288,7 +300,9 @@ async def _async_get_connection(conn_id: str) -> Connection:
maybe_awaitable = async_method(conn_id)
Review Comment:
From my agent:
> Backends with an async method still don't get the team.
`_async_get_connection` calls `async_method(conn_id)` and `_async_get_variable`
calls `async_method(key=key)` without `team_name`. Only the sync fallback goes
through `call_secrets_backend_method`. Example: a custom backend with `async
def aget_variable(self, key, team_name=None)` returns the global value in a
trigger but the team value in a task. Fix: route `async_method` through
`call_secrets_backend_method` as well.
##########
task-sdk/src/airflow/sdk/execution_time/context.py:
##########
@@ -232,7 +242,9 @@ 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())
Review Comment:
`SecretCache` already takes `team_name`, should we pass `_current_team_name`
here and in the other cache calls? Nothing breaks today since the cache only
runs in the Dag processor, but it will leak values across teams as soon as a
team is set there.
##########
task-sdk/src/airflow/sdk/execution_time/context.py:
##########
@@ -232,7 +242,9 @@ 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(
Review Comment:
nit: this runs `inspect.signature` on every lookup for every backend, even
when multi-team is off. Maybe call the backend directly when
`_current_team_name` is `None`?
##########
task-sdk/src/airflow/sdk/execution_time/context.py:
##########
@@ -205,6 +206,15 @@ async def _amask_connection_secrets(conn: Connection) ->
None:
"_preset_connections", default=None
)
+# A plain global rather than a ContextVar, so threads started by task code see
it too.
+_current_team_name: str | None = None
+
+
+def set_current_team_name(team_name: str | None) -> None:
Review Comment:
What about the Dag processor? A top-level `Variable.get` or a Dag callback
in a team bundle still gets the global env var, right? We can do it as a
follow-up but I think it should be fixed too.
##########
providers/edge3/src/airflow/providers/edge3/example_dags/integration_test.py:
##########
@@ -128,7 +128,7 @@ def variable():
def connection():
try:
conn = BaseHook.get_connection("integration_test")
- print(f"Got connection {conn}")
+ print(f"Got connection {conn.conn_id}")
Review Comment:
Why this change (here and in `win_test.py`)? It looks unrelated to this PR.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
