Invalidate cached RequestAdapter of KiotaRequestAdapterHook if session of AuthenticationProvider is closed

providers/microsoft/azure/src/airflow/providers/microsoft/azure/hooks/msgraph.py

@@ -54,9 +54,14 @@
 
 if TYPE_CHECKING:
     from azure.core.credentials_async import AsyncTokenCredential
+    from azure.core.pipeline.transport._aiohttp import AioHttpTransport
+    from kiota_abstractions.authentication import BaseBearerTokenAuthenticationProvider
     from kiota_abstractions.request_adapter import RequestAdapter
     from kiota_abstractions.response_handler import NativeResponseType
     from kiota_abstractions.serialization import ParsableFactory
+    from kiota_authentication_azure.azure_identity_access_token_provider import (
+        AzureIdentityAccessTokenProvider,
+    )
 
     from airflow.providers.common.compat.sdk import Connection
 
@@ -394,7 +399,23 @@ def get_conn(self) -> RequestAdapter:
     @staticmethod
     def _is_http_client_closed(request_adapter: RequestAdapter) -> bool:
         """Return True when the underlying httpx AsyncClient has been closed."""
-        return cast("HttpxRequestAdapter", request_adapter)._http_client.is_closed
+        adapter = cast("HttpxRequestAdapter", request_adapter)
+
+        if adapter._http_client.is_closed:
+            return True
+
+        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)
+
+        if not transport._has_been_opened and transport.session is None:
+            return True
+        if transport.session is not None:
+            return transport.session.closed
+        return False
 
     async def get_async_conn(self) -> RequestAdapter:
         """Initiate a new RequestAdapter connection asynchronously."""
@@ -606,7 +627,7 @@ async def send_request(self, request_info: RequestInformation, response_type: st
                 request_info=request_info,
                 error_map=self.error_mapping(),
             )
-        except Exception as e:
+        except (RuntimeError, ValueError) as e:
             self.log.warning(
                 "Request failed for conn_id '%s': %s. Invalidating cached request adapter.",
                 self.conn_id,

providers/microsoft/azure/tests/unit/microsoft/azure/hooks/test_msgraph.py

@@ -25,9 +25,16 @@
 from unittest.mock import AsyncMock, Mock, patch
 
 import pytest
-from httpx import Response
+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
@@ -68,6 +75,20 @@ def test_get_conn(self):
             assert isinstance(actual, HttpxRequestAdapter)
             assert actual.base_url == "https://graph.microsoft.com/v1.0/"
 
+    @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():
@@ -541,7 +562,26 @@ async def test_get_async_conn_rebuilds_adapter_when_http_client_is_closed(self):
             hook = KiotaRequestAdapterHook(conn_id="msgraph_api")
 
             stale_adapter = Mock(spec=HttpxRequestAdapter)
-            stale_adapter._http_client = Mock(is_closed=True)
+            stale_adapter._http_client = Mock(spec=AsyncClient, is_closed=True)
+            hook.cached_request_adapters[hook.conn_id] = (hook.api_version, stale_adapter)
+
+            fresh_adapter = Mock(spec=HttpxRequestAdapter)
+            fresh_adapter._http_client = Mock(is_closed=False)
+
+            with patch.object(hook, "_build_request_adapter", return_value=("v1.0", fresh_adapter)):
+                result = await hook.get_async_conn()
+
+            assert result is fresh_adapter
+            assert hook.cached_request_adapters[hook.conn_id] == ("v1.0", fresh_adapter)
+
+    @pytest.mark.asyncio
+    async def test_get_async_conn_rebuilds_adapter_when_credentials_session_is_closed(self):
+        """get_async_conn evicts and rebuilds the adapter when the cached HTTP client is already closed."""
+        with patch_hook():
+            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)
             hook.cached_request_adapters[hook.conn_id] = (hook.api_version, stale_adapter)
 
             fresh_adapter = Mock(spec=HttpxRequestAdapter)
@@ -560,7 +600,8 @@ async def test_send_request_invalidates_cache_and_raises_on_any_error(self):
             hook = KiotaRequestAdapterHook(conn_id="msgraph_api")
 
             adapter = Mock(spec=HttpxRequestAdapter)
-            adapter._http_client = Mock(is_closed=False)
+            adapter._http_client = Mock(spec=AsyncClient, is_closed=False)
+            adapter._authentication_provider = self.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)