Avoid extra_dejson in ADF and Synapse async hooks (#72130)

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

Author: Vamsi-klu

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

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

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

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")