This is an automated email from the ASF dual-hosted git repository. vatsrahul1001 pushed a commit to branch backport-d0577848b6-v3-3-test in repository https://gitbox.apache.org/repos/asf/airflow.git
commit eea6c055da30b1f21966a321bed67051d50955cf Author: David Blain <[email protected]> AuthorDate: Thu Sep 10 17:53:27 2026 +0200 Fix DeadlockImminentError when a connection is resolved inside an async task (#71890) * Add aget_uri and aextra_dejson to Connection to fix DeadlockImminentError in async tasks get_uri() accesses extra_dejson, which calls the synchronous mask_secret() → comms.send() from within the event-loop thread. Any async hook or task that calls aget_hook() / aget_connection() triggers this path, and Airflow 3.3.1's DeadlockImminentError detection surfaces the bug. Add aextra_dejson() — an async method that awaits amask_secret() instead of the blocking mask_secret() — and aget_uri(), which delegates URI assembly to a new shared _build_uri() helper and calls await self.aextra_dejson(). This keeps the entire connection-serialisation path safely on the async stack. The sync get_uri() / extra_dejson are unchanged; _build_uri() is the single source of truth for the URI format, shared by both paths. Co-authored-by: Copilot <[email protected]> (cherry picked from commit d0577848b68b37ad01418bb1c61dce605a435bcd) --- .../ssh/tests/unit/ssh/hooks/test_ssh_async.py | 8 ++-- task-sdk/docs/deferred-vs-async-operators.rst | 2 +- task-sdk/src/airflow/sdk/definitions/connection.py | 43 ++++++++++++++++-- task-sdk/src/airflow/sdk/execution_time/context.py | 11 ++++- .../tests/task_sdk/definitions/test_connection.py | 52 ++++++++++++++++++++++ .../tests/task_sdk/execution_time/test_context.py | 37 +++++++++++++++ 6 files changed, 144 insertions(+), 9 deletions(-) diff --git a/providers/ssh/tests/unit/ssh/hooks/test_ssh_async.py b/providers/ssh/tests/unit/ssh/hooks/test_ssh_async.py index 1e56fb4c372..4b47900ca9b 100644 --- a/providers/ssh/tests/unit/ssh/hooks/test_ssh_async.py +++ b/providers/ssh/tests/unit/ssh/hooks/test_ssh_async.py @@ -125,9 +125,11 @@ class TestSSHHookAsync: mock_ssh_client = mock.AsyncMock() - with mock.patch("asgiref.sync.sync_to_async") as mock_sync: - mock_sync.return_value = mock.AsyncMock(return_value=mock_conn_obj) - + with mock.patch( + "airflow.providers.ssh.hooks.ssh.get_async_connection", + new_callable=mock.AsyncMock, + return_value=mock_conn_obj, + ): with mock.patch("asyncssh.connect", new_callable=mock.AsyncMock) as mock_connect: mock_connect.return_value = mock_ssh_client result = await hook._get_conn() diff --git a/task-sdk/docs/deferred-vs-async-operators.rst b/task-sdk/docs/deferred-vs-async-operators.rst index 466060ab3a4..a2a0df16054 100644 --- a/task-sdk/docs/deferred-vs-async-operators.rst +++ b/task-sdk/docs/deferred-vs-async-operators.rst @@ -234,7 +234,7 @@ This allows multiple paginated requests to be performed efficiently within a sin @task async def get_users(): - hook = KiotaRequestAdapterHook.get_hook(conn_id="msgraph_default") + hook = await KiotaRequestAdapterHook.aget_hook(conn_id="msgraph_default") return await hook.paginated_run(url="users") diff --git a/task-sdk/src/airflow/sdk/definitions/connection.py b/task-sdk/src/airflow/sdk/definitions/connection.py index 06a95a3b868..8e9ab232a7c 100644 --- a/task-sdk/src/airflow/sdk/definitions/connection.py +++ b/task-sdk/src/airflow/sdk/definitions/connection.py @@ -153,10 +153,13 @@ class Connection: else: self.__dict__.update(attrs.asdict(self.from_uri(uri, conn_id=conn_id), recurse=False)) - def get_uri(self) -> str: - """Generate and return connection in URI format.""" - from urllib.parse import parse_qsl + def _build_uri(self, extra_dejson: dict) -> str: + """ + Build the connection URI given a pre-resolved extra_dejson dict. + Shared by ``get_uri`` (sync) and ``aget_uri`` (async) so the + URI-assembly logic lives in exactly one place. + """ if self.conn_type: uri = f"{self.conn_type.lower().replace('_', '-')}://" else: @@ -202,7 +205,6 @@ class Connection: if self.extra: try: - extra_dejson = self.extra_dejson query: str | None = urlencode(extra_dejson) except TypeError: query = None @@ -213,6 +215,20 @@ class Connection: return uri + def get_uri(self) -> str: + """Generate and return connection in URI format.""" + return self._build_uri(self.extra_dejson) + + async def aget_uri(self) -> str: + """ + Async version of ``get_uri``, safe for use inside an async task. + + Calls ``aextra_dejson`` so that secret masking uses ``asend()`` + instead of the synchronous ``send()``, preventing + ``DeadlockImminentError`` when invoked from within an async context. + """ + return self._build_uri(await self.aextra_dejson()) + def get_hook(self, *, hook_params=None): """Return hook based on conn_type.""" from airflow.sdk._shared.module_loading import import_string @@ -306,6 +322,25 @@ class Connection: return extra + async def aextra_dejson(self) -> dict: + """ + Async version of ``extra_dejson``, safe for use inside an async task. + + Uses ``amask_secret`` instead of the synchronous ``mask_secret``, so calling + this from within an async context does not trigger ``DeadlockImminentError``. + """ + from airflow.sdk.log import amask_secret + + extra: dict = {} + if self.extra: + try: + extra = json.loads(self.extra) + except JSONDecodeError: + log.exception("Failed to deserialize extra property `extra`, returning empty dictionary") + else: + await amask_secret(extra) + return extra + def get_extra_dejson(self) -> dict: """Deserialize extra property to JSON.""" import warnings diff --git a/task-sdk/src/airflow/sdk/execution_time/context.py b/task-sdk/src/airflow/sdk/execution_time/context.py index 58c58e5945e..a860af4c6e5 100644 --- a/task-sdk/src/airflow/sdk/execution_time/context.py +++ b/task-sdk/src/airflow/sdk/execution_time/context.py @@ -269,7 +269,16 @@ 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()) + # Use aget_uri if the returned connection object supports it (the SDK's own + # Connection class does); otherwise fall back to the sync get_uri, since backends + # can hand back other connection-shaped objects (e.g. MetastoreBackend returns + # airflow.models.Connection, which has no aget_uri). + aget_uri = getattr(conn, "aget_uri", None) + if aget_uri is not None: + uri = await aget_uri() + else: + uri = await sync_to_async(conn.get_uri)() + SecretCache.save_connection_uri(conn_id, uri) await _amask_connection_secrets(conn) return conn except AirflowSecretsBackendAccessDenied: diff --git a/task-sdk/tests/task_sdk/definitions/test_connection.py b/task-sdk/tests/task_sdk/definitions/test_connection.py index c973fc58c91..d471cccfcf6 100644 --- a/task-sdk/tests/task_sdk/definitions/test_connection.py +++ b/task-sdk/tests/task_sdk/definitions/test_connection.py @@ -234,6 +234,58 @@ class TestConnections: connection.extra = '{"auth": {"type": "oauth"}, "headers": {"User-Agent": "Airflow"}}' assert connection.extra_dejson == {"auth": {"type": "oauth"}, "headers": {"User-Agent": "Airflow"}} + @pytest.mark.asyncio + @mock.patch("airflow.sdk.definitions.connection.Connection.aextra_dejson") + async def test_aget_uri(self, mock_aextra_dejson): + """aget_uri must produce the same URI as get_uri and use aextra_dejson.""" + extra = {"charset": "utf8", "timeout": "30"} + mock_aextra_dejson.return_value = extra + + conn = Connection( + conn_id="test_conn", + conn_type="mysql", + host="localhost", + login="user", + password="password", + schema="test_schema", + port=3306, + extra='{"charset": "utf8", "timeout": "30"}', + ) + + uri = await conn.aget_uri() + assert uri == conn.get_uri() + mock_aextra_dejson.assert_awaited_once() + + @pytest.mark.asyncio + async def test_aextra_dejson_calls_amask_secret(self): + """aextra_dejson must use amask_secret (async), never the sync mask_secret.""" + connection = Connection( + conn_id="test_conn", + conn_type="http", + extra='{"api_key": "secret"}', + ) + + with ( + mock.patch("airflow.sdk.log.amask_secret") as mock_amask, + mock.patch("airflow.sdk.log.mask_secret") as mock_mask, + ): + result = await connection.aextra_dejson() + + assert result == {"api_key": "secret"} + mock_amask.assert_awaited_once_with({"api_key": "secret"}) + mock_mask.assert_not_called() + + @pytest.mark.asyncio + async def test_aextra_dejson_no_extra(self): + """aextra_dejson must return an empty dict without calling amask_secret when extra is None.""" + connection = Connection(conn_id="test_conn", conn_type="http") + + with mock.patch("airflow.sdk.log.amask_secret") as mock_amask: + result = await connection.aextra_dejson() + + assert result == {} + mock_amask.assert_not_called() + class TestConnectionsFromSecrets: def test_get_connection_secrets_backend(self, mock_supervisor_comms, tmp_path): 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 23e132432b2..a7a79d65185 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_context.py +++ b/task-sdk/tests/task_sdk/execution_time/test_context.py @@ -1227,6 +1227,43 @@ class TestAsyncGetConnection: 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(SecretCache, "save_connection_uri"), + ): + mock_load.return_value = [MockSecretsBackend()] + + await _async_get_connection("test_conn") + + # get_uri() would reach the sync mask_secret() -> comms.send(), which deadlocks + # on the event-loop thread; aget_uri() must go through amask_secret() -> comms.asend(). + mock_supervisor_comms.send.assert_not_called() + mock_supervisor_comms.asend.assert_awaited() + class TestSecretsBackend: """Test that connection resolution uses the backend chain correctly."""
