From 800560038f264c714b8f71f276f70d920d42ea07 Mon Sep 17 00:00:00 2001 From: David Blain Date: Mon, 29 Jun 2026 16:22:59 +0200 Subject: [PATCH 1/2] refactor: Also check if session of credentials provider in cached RequestAdapter isn't closed for KiotaRequestAdapterHook, if so invalidate the cached RequestAdapter --- .../microsoft/azure/hooks/msgraph.py | 20 +++++++- .../microsoft/azure/hooks/test_msgraph.py | 47 +++++++++++++++++-- 2 files changed, 63 insertions(+), 4 deletions(-) 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 2f3bf3a40308c..3fbaae3242aae 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 @@ -54,9 +54,13 @@ if TYPE_CHECKING: from azure.core.credentials_async import AsyncTokenCredential + 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 +398,21 @@ 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 = credential._client._pipeline._transport + + 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.""" 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 dab96656d43b4..79fe568b053c6 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,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) + 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) From 933f8090bf5ae6577ea047f72ba56a4e9a4469c7 Mon Sep 17 00:00:00 2001 From: David Blain Date: Thu, 2 Jul 2026 09:55:38 +0200 Subject: [PATCH 2/2] refactor: Also check the _has_been_opened property on the transport --- .../src/airflow/providers/microsoft/azure/hooks/msgraph.py | 7 +++++-- .../azure/tests/unit/microsoft/azure/hooks/test_msgraph.py | 2 +- 2 files changed, 6 insertions(+), 3 deletions(-) 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 3fbaae3242aae..d634e91263fe6 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 @@ -54,6 +54,7 @@ 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 @@ -408,8 +409,10 @@ def _is_http_client_closed(request_adapter: RequestAdapter) -> bool: credential = cast( "ClientSecretCredential | CertificateCredential", access_token_provider._credentials ) - transport = credential._client._pipeline._transport + 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 @@ -624,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, 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 79fe568b053c6..d9ab78866ec11 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 @@ -77,7 +77,7 @@ def test_get_conn(self): @classmethod def mock_authentication_provider(self, closed: bool) -> AuthenticationProvider: - transport = Mock(spec=AioHttpTransport) + transport = Mock(spec=AioHttpTransport, _has_been_opened=not closed) transport.session = Mock(spec=ClientSession, closed=closed) pipeline = Mock(spec=AsyncPipelineClient) pipeline._transport = transport