Add bearer token provider (#32655)

* Add bearer token provider

* Only creates the policy once

* Bump azure-core for typing

* black

* Revert "black"

This reverts commit 6454f84.

* black

* Feedback

---------

Co-authored-by: xiangyan99 <xiangsjtu@gmail.com>

Author: lmazuel

sdk/identity/azure-identity/azure/identity/__init__.py

@@ -28,6 +28,7 @@
     WorkloadIdentityCredential,
 )
 from ._persistent_cache import TokenCachePersistenceOptions
+from ._bearer_token_provider import get_bearer_token_provider
 
 
 __all__ = [
@@ -55,6 +56,7 @@
     "UsernamePasswordCredential",
     "VisualStudioCodeCredential",
     "WorkloadIdentityCredential",
+    "get_bearer_token_provider",
 ]
 
 from ._version import VERSION

sdk/identity/azure-identity/azure/identity/_bearer_token_provider.py

@@ -0,0 +1,46 @@
+# ------------------------------------
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+# ------------------------------------
+from typing import Callable
+
+from azure.core.credentials import TokenCredential
+from azure.core.pipeline.policies import BearerTokenCredentialPolicy
+from azure.core.pipeline import PipelineRequest, PipelineContext
+from azure.core.rest import HttpRequest
+
+
+def _make_request() -> PipelineRequest[HttpRequest]:
+    return PipelineRequest(HttpRequest("CredentialWrapper", "https://fakeurl"), PipelineContext(None))
+
+
+def get_bearer_token_provider(credential: TokenCredential, *scopes: str) -> Callable[[], str]:
+    """Returns a callable that provides a bearer token.
+
+    It can be used for instance to write code like:
+
+    .. code-block:: python
+
+        from azure.identity import DefaultAzureCredential, get_bearer_token_provider
+
+        credential = DefaultAzureCredential()
+        bearer_token_provider = get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
+
+        # Usage
+        request.headers["Authorization"] = "Bearer " + bearer_token_provider()
+
+    :param credential: The credential used to authenticate the request.
+    :type credential: ~azure.core.credentials.TokenCredential
+    :param str scopes: The scopes required for the bearer token.
+    :rtype: callable
+    :return: A callable that returns a bearer token.
+    """
+
+    policy = BearerTokenCredentialPolicy(credential, *scopes)
+
+    def wrapper() -> str:
+        request = _make_request()
+        policy.on_request(request)
+        return request.http_request.headers["Authorization"][len("Bearer ") :]
+
+    return wrapper

sdk/identity/azure-identity/azure/identity/aio/__init__.py

@@ -21,6 +21,7 @@
     ClientAssertionCredential,
     WorkloadIdentityCredential,
 )
+from ._bearer_token_provider import get_bearer_token_provider
 
 
 __all__ = [
@@ -39,4 +40,5 @@
     "VisualStudioCodeCredential",
     "ClientAssertionCredential",
     "WorkloadIdentityCredential",
+    "get_bearer_token_provider",
 ]

sdk/identity/azure-identity/azure/identity/aio/_bearer_token_provider.py

@@ -0,0 +1,47 @@
+# ------------------------------------
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+# ------------------------------------
+from typing import Callable, Coroutine, Any
+
+from azure.core.credentials_async import AsyncTokenCredential
+from azure.core.pipeline.policies import AsyncBearerTokenCredentialPolicy
+from azure.core.pipeline import PipelineRequest, PipelineContext
+from azure.core.rest import HttpRequest
+
+
+def _make_request() -> PipelineRequest[HttpRequest]:
+    return PipelineRequest(HttpRequest("CredentialWrapper", "https://fakeurl"), PipelineContext(None))
+
+
+def get_bearer_token_provider(credential: AsyncTokenCredential, *scopes: str) -> Callable[[], Coroutine[Any, Any, str]]:
+    """Returns a callable that provides a bearer token.
+
+    It can be used for instance to write code like:
+
+    .. code-block:: python
+
+        from azure.identity.aio import DefaultAzureCredential, get_bearer_token_provider
+
+        credential = DefaultAzureCredential()
+        bearer_token_provider = get_bearer_token_provider(credential, "https://cognitiveservices.azure.com/.default")
+
+
+        # Usage
+        request.headers["Authorization"] = "Bearer " + await bearer_token_provider()
+
+    :param credential: The credential used to authenticate the request.
+    :type credential: ~azure.core.credentials_async.AsyncTokenCredential
+    :param str scopes: The scopes required for the bearer token.
+    :rtype: coroutine
+    :return: A coroutine that returns a bearer token.
+    """
+
+    policy = AsyncBearerTokenCredentialPolicy(credential, *scopes)
+
+    async def wrapper() -> str:
+        request = _make_request()
+        await policy.on_request(request)
+        return request.http_request.headers["Authorization"][len("Bearer ") :]
+
+    return wrapper

sdk/identity/azure-identity/setup.py

@@ -47,6 +47,7 @@
         "Programming Language :: Python :: 3.9",
         "Programming Language :: Python :: 3.10",
         "Programming Language :: Python :: 3.11",
+        "Programming Language :: Python :: 3.12",
         "License :: OSI Approved :: MIT License",
     ],
     zip_safe=False,
@@ -59,7 +60,7 @@
     ),
     python_requires=">=3.7",
     install_requires=[
-        "azure-core<2.0.0,>=1.11.0",
+        "azure-core<2.0.0,>=1.23.0",
         "cryptography>=2.5",
         "msal<2.0.0,>=1.24.0",
         "msal-extensions<2.0.0,>=0.3.0",

sdk/identity/azure-identity/tests/test_bearer_token_provider.py

@@ -0,0 +1,20 @@
+# ------------------------------------
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+# ------------------------------------
+
+from azure.core.credentials import AccessToken
+from azure.identity import get_bearer_token_provider
+
+
+class MockCredential:
+    def get_token(self, *scopes, **kwargs):
+        assert len(scopes) == 1
+        assert scopes[0] == "scope"
+        return AccessToken("mock_token", 42)
+
+
+def test_get_bearer_token_provider():
+
+    func = get_bearer_token_provider(MockCredential(), "scope")
+    assert func() == "mock_token"

sdk/identity/azure-identity/tests/test_bearer_token_provider_async.py

@@ -0,0 +1,23 @@
+# ------------------------------------
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+# ------------------------------------
+
+from azure.core.credentials import AccessToken
+from azure.identity.aio import get_bearer_token_provider
+
+import pytest
+
+
+class MockCredential:
+    async def get_token(self, *scopes, **kwargs):
+        assert len(scopes) == 1
+        assert scopes[0] == "scope"
+        return AccessToken("mock_token", 42)
+
+
+@pytest.mark.asyncio
+async def test_get_bearer_token_provider():
+
+    func = get_bearer_token_provider(MockCredential(), "scope")
+    assert await func() == "mock_token"