Added CachedAsyncTokenCredential which keeps the session open during lifecycle of cached RequestAdapter

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,

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)
 

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)