diff --git a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/msgraph.py b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/msgraph.py index 50ad3cb1e2f12..4952531dcc6fd 100644 --- a/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/msgraph.py +++ b/providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/msgraph.py @@ -27,10 +27,12 @@ from http import HTTPStatus from io import BytesIO from json import JSONDecodeError +from types import TracebackType from typing import TYPE_CHECKING, Any, cast from urllib.parse import quote, urljoin, urlparse import httpx +from azure.core.credentials_async import AsyncTokenCredential from azure.identity.aio import CertificateCredential, ClientSecretCredential from httpx import AsyncHTTPTransport, Response, Timeout from kiota_abstractions.api_error import APIError @@ -53,8 +55,7 @@ from airflow.providers.common.compat.sdk import AirflowException, AirflowNotFoundException, BaseHook, redact if TYPE_CHECKING: - from azure.core.credentials_async import AsyncTokenCredential - from azure.core.pipeline.transport._aiohttp import AioHttpTransport + from azure.core.pipeline.transport._requests_basic import RequestsTransport from kiota_abstractions.authentication import BaseBearerTokenAuthenticationProvider from kiota_abstractions.request_adapter import RequestAdapter from kiota_abstractions.response_handler import NativeResponseType @@ -91,6 +92,57 @@ def execute_callable(func: Callable, *args: Any, **kwargs: Any) -> Any: return func(*args, **filtered_kwargs) +class CachedAsyncTokenCredential(AsyncTokenCredential): # type: ignore[misc] + """ + Wraps an async Azure credential to prevent ``kiota`` from closing it after each token request. + + ``kiota_authentication_azure`` calls ``await credential.close()`` after every successful + ``get_token`` call (see ``AzureIdentityAccessTokenProvider.get_authorization_token``). That + tears down the underlying ``AioHttpTransport`` session so the next request fails with + "HTTP transport has already been closed". Suppressing ``close()`` keeps the session alive + for the lifetime of the cached ``RequestAdapter``. + """ + + def __init__(self, credential: ClientSecretCredential | CertificateCredential): + self._credential = credential + + @property + def _transport(self) -> RequestsTransport: + return self._credential._client._pipeline._transport + + @property + def closed(self) -> bool: + # _closed is set to True by AioHttpTransport.close(); check it first as it + # is authoritative even when _has_been_opened is still False. + if getattr(self._transport, "_closed", False): + return True + if not self._transport._has_been_opened and self._transport.session is None: + return False + if self._transport.session is not None: + return self._transport.session.closed + return True + + async def __aenter__(self) -> AsyncTokenCredential: + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None = None, + exc_value: BaseException | None = None, + traceback: TracebackType | None = None, + ) -> None: + await self.close() + + async def get_token(self, *args: Any, **kwargs: Any) -> Any: + return await self._credential.get_token(*args, **kwargs) + + async def get_token_info(self, *args: Any, **kwargs: Any) -> Any: + return await self._credential.get_token_info(*args, **kwargs) # type: ignore[union-attr] + + async def close(self) -> None: + """Intentionally a no-op — the credential session is closed when the adapter is evicted.""" + + class DefaultResponseHandler(ResponseHandler): """DefaultResponseHandler returns JSON payload or content in bytes or response headers.""" @@ -413,16 +465,9 @@ def _is_http_client_closed(request_adapter: RequestAdapter) -> bool: provider = cast("BaseBearerTokenAuthenticationProvider", adapter._authentication_provider) access_token_provider = cast("AzureIdentityAccessTokenProvider", provider.access_token_provider) - credential = cast( - "ClientSecretCredential | CertificateCredential", access_token_provider._credentials - ) - transport = cast("AioHttpTransport", credential._client._pipeline._transport) + credential = cast("CachedAsyncTokenCredential", access_token_provider._credentials) - if not transport._has_been_opened and transport.session is None: - return False - if transport.session is not None: - return transport.session.closed - return False + return credential.closed async def get_async_conn(self) -> RequestAdapter: """Initiate a new RequestAdapter connection asynchronously.""" @@ -487,25 +532,29 @@ def get_credentials( self.log.info("Disable instance discovery: %s", disable_instance_discovery) self.log.info("MSAL Proxies: %s", redact(msal_proxies, name="proxies")) if certificate_path or certificate_data: - return CertificateCredential( + return CachedAsyncTokenCredential( + CertificateCredential( + tenant_id=tenant_id, + client_id=login, # type: ignore + password=password, + certificate_path=certificate_path, + certificate_data=certificate_data.encode() if certificate_data else None, + authority=authority, + proxies=msal_proxies, + disable_instance_discovery=disable_instance_discovery, + connection_verify=verify, + ) + ) + return CachedAsyncTokenCredential( + ClientSecretCredential( tenant_id=tenant_id, client_id=login, # type: ignore - password=password, - certificate_path=certificate_path, - certificate_data=certificate_data.encode() if certificate_data else None, + client_secret=password, # type: ignore authority=authority, proxies=msal_proxies, disable_instance_discovery=disable_instance_discovery, connection_verify=verify, ) - return ClientSecretCredential( - tenant_id=tenant_id, - client_id=login, # type: ignore - client_secret=password, # type: ignore - authority=authority, - proxies=msal_proxies, - disable_instance_discovery=disable_instance_discovery, - connection_verify=verify, ) def test_connection(self): @@ -547,8 +596,6 @@ async def run( headers: dict[str, str] | None = None, data: dict[str, Any] | str | BytesIO | None = None, ): - self.log.info("Executing url '%s' as '%s'", url, method) - response = await self.send_request( request_info=self.request_information( url=url, @@ -624,6 +671,8 @@ async def send_request(self, request_info: RequestInformation, response_type: st conn = await self.get_async_conn() try: + self.log.info("Executing url '%s' as '%s'", request_info.url, request_info.http_method) + if response_type: return await conn.send_primitive_async( request_info=request_info, diff --git a/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_msgraph.py b/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_msgraph.py index ee9b0f87dfa55..8eb60e91eed69 100644 --- a/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_msgraph.py +++ b/providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_msgraph.py @@ -25,16 +25,9 @@ from unittest.mock import AsyncMock, Mock, patch import pytest -from aiohttp import ClientSession -from azure.core import AsyncPipelineClient -from azure.core.credentials_async import AsyncTokenCredential -from azure.core.pipeline.transport._aiohttp import AioHttpTransport -from azure.identity._internal import AadClient from httpx import AsyncClient, Response from httpx._utils import URLPattern -from kiota_abstractions.authentication import AuthenticationProvider, BaseBearerTokenAuthenticationProvider from kiota_abstractions.request_information import RequestInformation -from kiota_authentication_azure.azure_identity_access_token_provider import AzureIdentityAccessTokenProvider from kiota_http.httpx_request_adapter import HttpxRequestAdapter from kiota_serialization_json.json_parse_node import JsonParseNode from kiota_serialization_text.text_parse_node import TextParseNode @@ -44,6 +37,7 @@ from airflow.exceptions import AirflowBadRequest, AirflowConfigException, AirflowProviderDeprecationWarning from airflow.providers.common.compat.sdk import AirflowException, AirflowNotFoundException from airflow.providers.microsoft.azure.hooks.msgraph import ( + CachedAsyncTokenCredential, DefaultResponseHandler, KiotaRequestAdapterHook, execute_callable, @@ -53,14 +47,55 @@ from tests_common.test_utils.providers import get_provider_min_airflow_version from unit.microsoft.azure.test_utils import ( get_airflow_connection, + mock_authentication_provider, mock_connection, mock_json_response, mock_response, + mock_token_credentials, patch_hook, patch_hook_and_request_adapter, ) +class TestCachedAsyncTokenCredential: + @pytest.mark.parametrize( + ("closed", "expected"), + ( + pytest.param(None, False), + pytest.param(False, False), + pytest.param(True, True), + ), + ) + def test_closed(self, closed: bool | None, expected: bool): + actual = CachedAsyncTokenCredential(credential=mock_token_credentials(closed=closed)) + + assert actual.closed == expected + + @pytest.mark.asyncio + async def test_close(self): + credential = mock_token_credentials() + actual = CachedAsyncTokenCredential(credential=credential) + + await actual.close() + credential.assert_not_called() + + @pytest.mark.asyncio + async def test_get_token(self): + credential = mock_token_credentials() + actual = CachedAsyncTokenCredential(credential=credential) + + await actual.get_token() + credential.get_token.assert_called_once() + + @pytest.mark.asyncio + async def test_get_token_info(self): + credential = mock_token_credentials() + actual = CachedAsyncTokenCredential(credential=credential) + + await actual.get_token_info() + credential.get_token_info.assert_called_once() + + class TestKiotaRequestAdapterHook: def test_get_conn(self): with patch_hook(): @@ -75,35 +110,6 @@ def test_get_conn(self): assert isinstance(actual, HttpxRequestAdapter) assert actual.base_url == "https://graph.microsoft.com/v1.0/" - @classmethod - def mock_authentication_provider_never_opened(self) -> AuthenticationProvider: - """Return an auth provider whose transport has never been opened (session is None).""" - transport = Mock(spec=AioHttpTransport, _has_been_opened=False) - transport.session = None - pipeline = Mock(spec=AsyncPipelineClient) - pipeline._transport = transport - client = Mock(spec=AadClient) - client._pipeline = pipeline - credentials = Mock(spec=AsyncTokenCredential) - credentials._client = client - access_token_provider = Mock(spec=AzureIdentityAccessTokenProvider) - access_token_provider._credentials = credentials - return BaseBearerTokenAuthenticationProvider(access_token_provider=access_token_provider) - - @classmethod - def mock_authentication_provider(self, closed: bool) -> AuthenticationProvider: - transport = Mock(spec=AioHttpTransport, _has_been_opened=not closed) - transport.session = Mock(spec=ClientSession, closed=closed) - pipeline = Mock(spec=AsyncPipelineClient) - pipeline._transport = transport - client = Mock(spec=AadClient) - client._pipeline = pipeline - credentials = Mock(spec=AsyncTokenCredential) - credentials._client = client - access_token_provider = Mock(spec=AzureIdentityAccessTokenProvider) - access_token_provider._credentials = credentials - return BaseBearerTokenAuthenticationProvider(access_token_provider=access_token_provider) - @pytest.mark.asyncio async def test_get_async_conn(self): with patch_hook(): @@ -596,7 +602,7 @@ async def test_get_async_conn_rebuilds_adapter_when_credentials_session_is_close hook = KiotaRequestAdapterHook(conn_id="msgraph_api") stale_adapter = Mock(spec=HttpxRequestAdapter) stale_adapter._http_client = Mock(spec=AsyncClient, is_closed=False) - stale_adapter._authentication_provider = self.mock_authentication_provider(closed=True) + stale_adapter._authentication_provider = mock_authentication_provider(closed=True) hook.cached_request_adapters[hook.conn_id] = (hook.api_version, stale_adapter) fresh_adapter = Mock(spec=HttpxRequestAdapter) @@ -615,7 +621,7 @@ async def test_get_async_conn_does_not_rebuild_adapter_when_transport_never_open hook = KiotaRequestAdapterHook(conn_id="msgraph_api") adapter = Mock(spec=HttpxRequestAdapter) adapter._http_client = Mock(spec=AsyncClient, is_closed=False) - adapter._authentication_provider = self.mock_authentication_provider_never_opened() + adapter._authentication_provider = mock_authentication_provider() hook.cached_request_adapters[hook.conn_id] = (hook.api_version, adapter) result = await hook.get_async_conn() @@ -630,7 +636,7 @@ async def test_send_request_invalidates_cache_and_raises_on_any_error(self): adapter = Mock(spec=HttpxRequestAdapter) adapter._http_client = Mock(spec=AsyncClient, is_closed=False) - adapter._authentication_provider = self.mock_authentication_provider(closed=False) + adapter._authentication_provider = mock_authentication_provider(closed=False) adapter.send_no_response_content_async = AsyncMock(side_effect=RuntimeError("some error")) hook.cached_request_adapters[hook.conn_id] = (hook.api_version, adapter) diff --git a/providers/microsoft/azure/tests/unit/microsoft/azure/test_utils.py b/providers/microsoft/azure/tests/unit/microsoft/azure/test_utils.py index 590aea1ac45a7..bc8577ba9aa46 100644 --- a/providers/microsoft/azure/tests/unit/microsoft/azure/test_utils.py +++ b/providers/microsoft/azure/tests/unit/microsoft/azure/test_utils.py @@ -22,10 +22,17 @@ from json import JSONDecodeError from typing import Any from unittest import mock -from unittest.mock import MagicMock, patch +from unittest.mock import AsyncMock, MagicMock, Mock, patch import pytest +from aiohttp import ClientSession +from azure.core import AsyncPipelineClient +from azure.core.credentials_async import AsyncTokenCredential +from azure.core.pipeline.transport._aiohttp import AioHttpTransport +from azure.identity.aio._internal import AadClient from httpx import Headers, Response +from kiota_abstractions.authentication import AuthenticationProvider, BaseBearerTokenAuthenticationProvider +from kiota_authentication_azure.azure_identity_access_token_provider import AzureIdentityAccessTokenProvider from kiota_http.httpx_request_adapter import HttpxRequestAdapter from msgraph_core import APIVersion @@ -35,7 +42,6 @@ add_managed_identity_connection_widgets, get_async_default_azure_credential, get_field, - # _get_default_azure_credential get_sync_default_azure_credential, parse_blob_account_url, ) @@ -269,3 +275,27 @@ def patch_hook_and_request_adapter(response): mock_get_http_response.return_value = response yield [*hook_mocks, mock_get_http_response] + + +def mock_token_credentials(closed: bool | None = None): + _has_been_opened = False if closed is None else not closed + transport = Mock(spec=AioHttpTransport, _has_been_opened=_has_been_opened) + transport.session = None if closed is None else Mock(spec=ClientSession, closed=closed) + pipeline = Mock(spec=AsyncPipelineClient) + pipeline._transport = transport + client = Mock(spec=AadClient) + client._pipeline = pipeline + credentials = Mock(spec=AsyncTokenCredential) + credentials._client = client + credentials.get_token_info = AsyncMock() + return credentials + + +def mock_authentication_provider(closed: bool | None = None) -> AuthenticationProvider: + from airflow.providers.microsoft.azure.hooks.msgraph import CachedAsyncTokenCredential + + access_token_provider = Mock(spec=AzureIdentityAccessTokenProvider) + access_token_provider._credentials = CachedAsyncTokenCredential( + credential=mock_token_credentials(closed=closed) + ) + return BaseBearerTokenAuthenticationProvider(access_token_provider=access_token_provider)