This is an automated email from the ASF dual-hosted git repository.
dabla pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 6fd1994ce28 refactor: Introduced CachedAsyncTokenCredential which
keeps the session open during lifecycle of cached RequestAdapter (#69365)
6fd1994ce28 is described below
commit 6fd1994ce2878c7317aa8a2f9b334beda0041961
Author: David Blain <[email protected]>
AuthorDate: Tue Jul 7 15:27:27 2026 +0200
refactor: Introduced CachedAsyncTokenCredential which keeps the session
open during lifecycle of cached RequestAdapter (#69365)
---
.../providers/microsoft/azure/hooks/msgraph.py | 99 ++++++++++++++++------
.../unit/microsoft/azure/hooks/test_msgraph.py | 84 +++++++++---------
.../azure/tests/unit/microsoft/azure/test_utils.py | 34 +++++++-
3 files changed, 151 insertions(+), 66 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 50ad3cb1e2f..4952531dcc6 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 contextlib import suppress
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.connection import
get_async_connection
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 @@ class KiotaRequestAdapterHook(BaseHook):
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 @@ class KiotaRequestAdapterHook(BaseHook):
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 @@ class KiotaRequestAdapterHook(BaseHook):
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 @@ class KiotaRequestAdapterHook(BaseHook):
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 ee9b0f87dfa..8eb60e91eed 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 typing import cast
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 opentelemetry.trace import Span
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.file_loading import
load_file_from_resources, load_
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 @@ class TestKiotaRequestAdapterHook:
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 @@ class TestKiotaRequestAdapterHook:
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 @@ class TestKiotaRequestAdapterHook:
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 @@ class TestKiotaRequestAdapterHook:
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 590aea1ac45..bc8577ba9aa 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 contextlib import ExitStack, contextmanager
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 @@ from airflow.providers.microsoft.azure.utils import (
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)
