kaxil commented on code in PR #71890:
URL: https://github.com/apache/airflow/pull/71890#discussion_r3978612040
##########
task-sdk/src/airflow/sdk/definitions/connection.py:
##########
@@ -306,6 +322,26 @@ def extra_dejson(self) -> dict:
return extra
+ async def aextra_dejson(self) -> dict:
Review Comment:
Following up on Amogh's second point, here are the concrete callers that
still go through the sync property right after resolving a connection on the
event loop: `SSHHookAsync._get_conn` -> `_parse_extras` (ssh.py:606),
`SlackHook.get_async_client` -> `_get_conn_params` (slack.py:175), and
`HttpAsyncHook.config` -> `_process_extra_options_from_connection`
(http.py:66). msgraph already works around it with
`json.loads(connection.extra)` and a TODO pointing at #54350, which skips
masking altogether. Nothing calls `aextra_dejson` yet, so is a follow-up
planned to move those over?
##########
task-sdk/src/airflow/sdk/execution_time/context.py:
##########
@@ -297,7 +297,7 @@ async def _async_get_connection(conn_id: str) -> Connection:
conn = await
sync_to_async(secrets_backend.get_connection)(conn_id) # type:
ignore[assignment]
if conn:
- SecretCache.save_connection_uri(conn_id, conn.get_uri())
+ SecretCache.save_connection_uri(conn_id, await conn.aget_uri())
Review Comment:
`aget_uri()` only exists on the Task SDK `Connection`, and this line sits
inside the `try` whose `except Exception` logs at debug and moves on to the
next backend. So a backend returning any other connection-shaped object
surfaces as `AirflowNotFoundException("The conn_id ... isn't defined")` rather
than as an attribute error: `MetastoreBackend` hands back
`airflow.models.Connection`, which has `get_uri` but no `aget_uri`, and custom
backends that build their own objects are in the same position. The SSH test in
this diff is the shape of it, a `MagicMock` connection no longer survives the
loop. Nothing in-tree hits this today since every client-context backend gets
the SDK class injected via `_set_connection_class`, but would a `getattr(conn,
"aget_uri", None)` fallback to `get_uri()` be worth it here, mirroring the
`hasattr(hook, "aget_connection")` check in
`common.compat.get_async_connection`?
##########
task-sdk/tests/task_sdk/execution_time/test_context.py:
##########
@@ -1272,6 +1272,44 @@ def get_connection(self, conn_id: str) -> Connection |
None:
mock_supervisor_comms.send.assert_not_called()
mock_supervisor_comms.asend.assert_not_called()
+ @pytest.mark.asyncio
+ async def test_async_get_connection_uses_aget_uri_not_get_uri(self,
mock_supervisor_comms):
+ """_async_get_connection must call aget_uri() when caching, never the
sync get_uri().
+
+ get_uri() accesses extra_dejson which calls mask_secret() ->
comms.send()
+ from the event-loop thread, triggering DeadlockImminentError in
Airflow 3.3.1.
+ aget_uri() uses amask_secret() -> asend() and is safe in async
contexts.
+ """
+ from airflow.sdk.execution_time.cache import SecretCache
+
+ sample_connection = Connection(
+ conn_id="test_conn",
+ conn_type="postgres",
+ host="localhost",
+ port=5432,
+ extra='{"sslmode": "require"}',
+ )
+
+ class MockSecretsBackend:
+ def get_connection(self, conn_id: str) -> Connection | None:
+ return sample_connection if conn_id == "test_conn" else None
+
+ with (
+ patch(
+
"airflow.sdk.execution_time.supervisor.ensure_secrets_backend_loaded",
autospec=True
+ ) as mock_load,
+ mock.patch.object(sample_connection, "aget_uri") as mock_aget_uri,
Review Comment:
Mocking `aget_uri` means the test never reaches `aextra_dejson` ->
`amask_secret` -> `asend`, so it pins the call shape rather than the thing that
was broken. Dropping these two `patch.object` mocks and asserting
`mock_supervisor_comms.send.assert_not_called()`, the way the test just above
does, would fail on the old code for the actual reason.
--
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]
