From ab60f51e032f26b2111759ea64849d4191249f64 Mon Sep 17 00:00:00 2001 From: Vamsi-klu Date: Thu, 27 Aug 2026 06:48:20 +0000 Subject: [PATCH] Avoid extra_dejson in ADF and Synapse async hooks extra_dejson can mask secrets via a sync send on the triggerer event loop, which raises AsyncToSync. Parse extras with json.loads instead, matching the MSGraph workaround. closes: #55728 --- .../microsoft/azure/hooks/data_factory.py | 8 ++- .../microsoft/azure/hooks/synapse.py | 7 +- .../azure/hooks/test_data_factory.py | 68 +++++++++++++++++-- .../azure/hooks/test_synapse_pipeline.py | 43 +++++++++++- 4 files changed, 117 insertions(+), 9 deletions(-) diff --git a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/data_factory.py b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/data_factory.py index 61ae3c8848b9c..7df86c7be620a 100644 --- a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/data_factory.py +++ b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/data_factory.py @@ -34,6 +34,7 @@ from __future__ import annotations import inspect +import json import time from collections.abc import Callable from functools import wraps @@ -1124,7 +1125,8 @@ async def bind_argument(arg: Any, default_key: str) -> None: if arg not in bound_args.arguments or bound_args.arguments[arg] is None: self = args[0] conn = await get_async_connection(self.conn_id) - extras = conn.extra_dejson + # extra_dejson can call mask_secret -> sync send on the triggerer loop. + extras = json.loads(conn.extra) if conn.extra else {} default_value = extras.get(default_key) or extras.get( f"extra__azure_data_factory__{default_key}" ) @@ -1175,7 +1177,8 @@ async def get_async_conn(self) -> AsyncDataFactoryManagementClient: return self._async_conn conn = await get_async_connection(self.conn_id) - extras = conn.extra_dejson + # extra_dejson can call mask_secret -> sync send on the triggerer loop. + extras = json.loads(conn.extra) if conn.extra else {} tenant = get_field(extras, "tenantId") try: @@ -1208,6 +1211,7 @@ async def get_async_conn(self) -> AsyncDataFactoryManagementClient: async def refresh_conn(self) -> AsyncDataFactoryManagementClient: # type: ignore[override] self._conn = None + await self.close() return await self.get_async_conn() @provide_targeted_factory_async diff --git a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/synapse.py b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/synapse.py index e37c6bca9f213..c54f90e6bf790 100644 --- a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/synapse.py +++ b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/synapse.py @@ -16,6 +16,7 @@ # under the License. from __future__ import annotations +import json import time from typing import TYPE_CHECKING, Any @@ -29,6 +30,7 @@ from azure.synapse.artifacts.aio import ArtifactsClient as AsyncArtifactsClient from azure.synapse.spark import SparkClient +from airflow.providers.common.compat.connection import get_async_connection from airflow.providers.common.compat.sdk import AirflowException, AirflowTaskTimeout, BaseHook from airflow.providers.microsoft.azure.utils import ( add_managed_identity_connection_widgets, @@ -489,8 +491,9 @@ async def get_async_conn(self) -> AsyncArtifactsClient: if self._async_conn is not None: return self._async_conn - conn = self.get_connection(self.conn_id) - extras = conn.extra_dejson + conn = await get_async_connection(self.conn_id) + # extra_dejson can call mask_secret -> sync send on the triggerer loop. + extras = json.loads(conn.extra) if conn.extra else {} tenant = self._get_field(extras, "tenantId") credential: AsyncCredentials diff --git a/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_data_factory.py b/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_data_factory.py index 1b64d29c3ce64..16e6eb237e880 100644 --- a/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_data_factory.py +++ b/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_data_factory.py @@ -16,6 +16,7 @@ # under the License. from __future__ import annotations +import json import os from unittest import mock from unittest.mock import MagicMock, PropertyMock, patch @@ -716,20 +717,31 @@ async def test_get_adf_pipeline_run_status_cancelled(self, mock_get_pipeline_run assert response == mock_status @pytest.mark.asyncio - @mock.patch(f"{MODULE}.AzureDataFactoryAsyncHook.get_connection") + @mock.patch(f"{MODULE}.get_async_connection", new_callable=mock.AsyncMock) @mock.patch(f"{MODULE}.AzureDataFactoryAsyncHook.get_async_conn") - async def test_get_pipeline_run_exception_without_resource(self, mock_conn, mock_get_connection): + async def test_get_pipeline_run_exception_without_resource(self, mock_conn, mock_get_async_connection): """ Test get_pipeline_run function without passing the resource name to check the decorator function and raise exception """ mock_connection = Connection(extra={"factory_name": DATAFACTORY_NAME}) - mock_get_connection.return_value = mock_connection + mock_get_async_connection.return_value = mock_connection mock_conn.return_value.pipeline_runs.get.return_value = MagicMock() hook = AzureDataFactoryAsyncHook(AZURE_DATA_FACTORY_CONN_ID) with pytest.raises(AirflowException): await hook.get_pipeline_run(RUN_ID, None, DATAFACTORY_NAME) + @staticmethod + def _conn_with_raising_extra_dejson(extra: dict, login="clientId", password="clientSecret"): + conn = mock.Mock() + conn.login = login + conn.password = password + conn.extra = json.dumps(extra) + type(conn).extra_dejson = PropertyMock( + side_effect=RuntimeError("You cannot use AsyncToSync in the same thread as an async event loop") + ) + return conn + @pytest.mark.asyncio @pytest.mark.parametrize( "mocked_connection", @@ -786,6 +798,50 @@ async def test_get_async_conn(self, mocked_connection): response = await hook.get_async_conn() assert isinstance(response, DataFactoryManagementClient) + @pytest.mark.asyncio + async def test_get_async_conn_does_not_touch_extra_dejson(self): + conn = self._conn_with_raising_extra_dejson( + {"tenantId": "tenantId", "subscriptionId": "subscriptionId"} + ) + hook = AzureDataFactoryAsyncHook(AZURE_DATA_FACTORY_CONN_ID) + with ( + mock.patch(f"{MODULE}.get_async_connection", new=mock.AsyncMock(return_value=conn)), + mock.patch(f"{MODULE}.AsyncClientSecretCredential"), + mock.patch(f"{MODULE}.AsyncDataFactoryManagementClient") as mock_client, + ): + response = await hook.get_async_conn() + assert response is mock_client.return_value + + @pytest.mark.asyncio + async def test_get_async_conn_uses_get_async_connection(self): + conn = self._conn_with_raising_extra_dejson( + {"tenantId": "tenantId", "subscriptionId": "subscriptionId"} + ) + hook = AzureDataFactoryAsyncHook(AZURE_DATA_FACTORY_CONN_ID) + with ( + mock.patch( + f"{MODULE}.get_async_connection", new=mock.AsyncMock(return_value=conn) + ) as mock_get_async_connection, + mock.patch(f"{MODULE}.AsyncClientSecretCredential"), + mock.patch(f"{MODULE}.AsyncDataFactoryManagementClient"), + ): + await hook.get_async_conn() + mock_get_async_connection.assert_awaited_once_with(AZURE_DATA_FACTORY_CONN_ID) + + @pytest.mark.asyncio + @mock.patch(f"{MODULE}.AzureDataFactoryAsyncHook.get_async_conn") + async def test_provide_targeted_factory_async_does_not_touch_extra_dejson(self, mock_get_async_conn): + conn = self._conn_with_raising_extra_dejson( + {"resource_group_name": RESOURCE_GROUP_NAME, "factory_name": DATAFACTORY_NAME} + ) + mock_get_async_conn.return_value.pipeline_runs.get = mock.AsyncMock(return_value=MagicMock()) + hook = AzureDataFactoryAsyncHook(AZURE_DATA_FACTORY_CONN_ID) + with mock.patch(f"{MODULE}.get_async_connection", new=mock.AsyncMock(return_value=conn)): + await hook.get_pipeline_run(RUN_ID, None, None) + mock_get_async_conn.return_value.pipeline_runs.get.assert_awaited_once_with( + RESOURCE_GROUP_NAME, DATAFACTORY_NAME, RUN_ID + ) + @pytest.mark.asyncio @pytest.mark.parametrize( "mocked_connection", @@ -904,10 +960,14 @@ def test_get_field_non_prefixed_extras(self): @pytest.mark.asyncio @mock.patch(f"{MODULE}.AzureDataFactoryAsyncHook.get_async_conn") async def test_refresh_conn(self, mock_get_async_conn): - """Test refresh_conn method _conn is reset and get_async_conn is called""" + """Test refresh_conn closes the async client before recreating it.""" hook = AzureDataFactoryAsyncHook(AZURE_DATA_FACTORY_CONN_ID) + mock_async_conn = mock.AsyncMock() + hook._async_conn = mock_async_conn await hook.refresh_conn() assert not hook._conn + mock_async_conn.close.assert_awaited_once() + assert hook._async_conn is None assert mock_get_async_conn.called @pytest.mark.asyncio diff --git a/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_synapse_pipeline.py b/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_synapse_pipeline.py index b35e589d02831..436d2aa4474f6 100644 --- a/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_synapse_pipeline.py +++ b/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_synapse_pipeline.py @@ -16,7 +16,8 @@ # under the License. from __future__ import annotations -from unittest.mock import AsyncMock, MagicMock, patch +import json +from unittest.mock import AsyncMock, MagicMock, PropertyMock, patch import pytest from azure.synapse.artifacts import ArtifactsClient @@ -247,6 +248,46 @@ async def test_get_async_conn_client_secret(self, mock_credential, mock_client): credential=mock_credential.return_value, ) + @pytest.mark.asyncio + @patch(f"{MODULE}.AsyncArtifactsClient") + @patch(f"{MODULE}.AsyncClientSecretCredential") + async def test_get_async_conn_does_not_touch_extra_dejson(self, mock_credential, mock_client): + conn = MagicMock() + conn.login = "clientId" + conn.password = "clientSecret" + conn.extra = json.dumps({"tenantId": "tenantId"}) + type(conn).extra_dejson = PropertyMock( + side_effect=RuntimeError("You cannot use AsyncToSync in the same thread as an async event loop") + ) + hook = AzureSynapsePipelineAsyncHook( + azure_synapse_conn_id=DEFAULT_CONNECTION_CLIENT_SECRET, + azure_synapse_workspace_dev_endpoint=AZURE_SYNAPSE_WORKSPACE_DEV_ENDPOINT, + ) + with patch(f"{MODULE}.get_async_connection", new=AsyncMock(return_value=conn)): + result = await hook.get_async_conn() + assert result is mock_client.return_value + mock_credential.assert_called_with( + client_id="clientId", + client_secret="clientSecret", + tenant_id="tenantId", + ) + + @pytest.mark.asyncio + @patch(f"{MODULE}.AsyncArtifactsClient") + @patch(f"{MODULE}.AsyncClientSecretCredential") + async def test_get_async_conn_uses_get_async_connection(self, mock_credential, mock_client): + conn = MagicMock() + conn.login = "clientId" + conn.password = "clientSecret" + conn.extra = json.dumps({"tenantId": "tenantId"}) + hook = AzureSynapsePipelineAsyncHook( + azure_synapse_conn_id=DEFAULT_CONNECTION_CLIENT_SECRET, + azure_synapse_workspace_dev_endpoint=AZURE_SYNAPSE_WORKSPACE_DEV_ENDPOINT, + ) + with patch(f"{MODULE}.get_async_connection", new=AsyncMock(return_value=conn)) as mock_get: + await hook.get_async_conn() + mock_get.assert_awaited_once_with(DEFAULT_CONNECTION_CLIENT_SECRET) + @pytest.mark.asyncio @patch(f"{MODULE}.AsyncArtifactsClient") @patch(f"{MODULE}.AsyncDefaultAzureCredential")