Expose methods for closing async credential transport sessions

sdk/core/azure-core/azure/core/credentials_async.py

@@ -0,0 +1,23 @@
+# ------------------------------------
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+# ------------------------------------
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Any
+    from typing_extensions import Protocol
+    from .credentials import AccessToken
+
+    class AsyncTokenCredential(Protocol):
+        async def get_token(self, *scopes: str, **kwargs: Any) -> AccessToken:
+            pass
+
+        async def close(self) -> None:
+            pass
+
+        async def __aenter__(self):
+            pass
+
+        async def __aexit__(self, exc_type, exc_value, traceback) -> None:
+            pass

sdk/core/azure-core/azure/core/pipeline/policies/_authentication.py

@@ -29,11 +29,10 @@ class _BearerTokenCredentialPolicyBase(object):
     :param str scopes: Lets you specify the type of access needed.
     """
 
-    def __init__(self, credential, *scopes, **kwargs):  # pylint:disable=unused-argument
-        # type: (TokenCredential, *str, Mapping[str, Any]) -> None
+    def __init__(self, *scopes, **kwargs):  # pylint:disable=unused-argument
+        # type: (*str, **Any) -> None
         super(_BearerTokenCredentialPolicyBase, self).__init__()
         self._scopes = scopes
-        self._credential = credential
         self._token = None  # type: Optional[AccessToken]
 
     @staticmethod
@@ -69,6 +68,11 @@ class BearerTokenCredentialPolicy(_BearerTokenCredentialPolicyBase, SansIOHTTPPo
     :raises: :class:`~azure.core.exceptions.ServiceRequestError`
     """
 
+    def __init__(self, credential, *scopes, **kwargs):
+        # type: (TokenCredential, *str, **Any) -> None
+        self._credential = credential
+        super(BearerTokenCredentialPolicy, self).__init__(*scopes, **kwargs)
+
     def on_request(self, request):
         # type: (PipelineRequest) -> None
         """Adds a bearer token Authorization header to request and sends request to next policy.

sdk/core/azure-core/azure/core/pipeline/policies/_authentication_async.py

@@ -4,26 +4,33 @@
 # license information.
 # -------------------------------------------------------------------------
 import threading
+from typing import TYPE_CHECKING
 
-from azure.core.pipeline import PipelineRequest
 from azure.core.pipeline.policies import SansIOHTTPPolicy
 from azure.core.pipeline.policies._authentication import _BearerTokenCredentialPolicyBase
 
+if TYPE_CHECKING:
+    # pylint:disable=unused-import
+    from typing import Any
+    from azure.core.credentials_async import AsyncTokenCredential
+    from azure.core.pipeline import PipelineRequest
+
 
 class AsyncBearerTokenCredentialPolicy(_BearerTokenCredentialPolicyBase, SansIOHTTPPolicy):
     # pylint:disable=too-few-public-methods
     """Adds a bearer token Authorization header to requests.
 
     :param credential: The credential.
-    :type credential: ~azure.core.credentials.TokenCredential
+    :type credential: ~azure.core.credentials_async.AsyncTokenCredential
     :param str scopes: Lets you specify the type of access needed.
     """
 
-    def __init__(self, credential, *scopes, **kwargs):
-        super().__init__(credential, *scopes, **kwargs)
+    def __init__(self, credential: "AsyncTokenCredential", *scopes: str, **kwargs: "Any") -> None:
+        self._credential = credential
         self._lock = threading.Lock()
+        super().__init__(*scopes, **kwargs)
 
-    async def on_request(self, request: PipelineRequest):
+    async def on_request(self, request: "PipelineRequest"):
         """Adds a bearer token Authorization header to request and sends request to next policy.
 
         :param request: The pipeline request object to be modified.

sdk/identity/azure-identity/HISTORY.md

@@ -4,6 +4,8 @@
 
 - All credential pipelines include `ProxyPolicy`
 ([#8945](https://github.com/Azure/azure-sdk-for-python/pull/8945))
+- Async credentials are async context managers and have an async `close` method
+([#9090](https://github.com/Azure/azure-sdk-for-python/pull/9090))
 
 
 ## 1.1.0 (2019-11-27)

sdk/identity/azure-identity/README.md

@@ -213,6 +213,24 @@ async transport, such as [aiohttp](https://pypi.org/project/aiohttp/). See
 [azure-core documentation](../../core/azure-core/README.md#transport)
 for more information.
 
+Async credentials should be closed when they're no longer needed. Each async
+credential is an async context manager and defines an async `close` method. For
+example:
+
+```py
+from azure.identity.aio import DefaultAzureCredential
+
+# call close when the credential is no longer needed
+credential = DefaultAzureCredential()
+...
+await credential.close()
+
+# alternatively, use the credential as an async context manager
+credential = DefaultAzureCredential()
+async with credential:
+  ...
+```
+
 This example demonstrates authenticating the asynchronous `SecretClient` from
 [azure-keyvault-secrets][azure_keyvault_secrets] with an asynchronous
 credential.

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

@@ -54,6 +54,16 @@ def __init__(
         self._pipeline = AsyncPipeline(transport=transport, policies=policies)
         super().__init__(**kwargs)
 
+    async def __aenter__(self):
+        await self._pipeline.__aenter__()
+        return self
+
+    async def __aexit__(self, *args):
+        await self.close()
+
+    async def close(self) -> None:
+        await self._pipeline.__aexit__()
+
     async def request_token(
         self,
         scopes: "Iterable[str]",

sdk/identity/azure-identity/azure/identity/aio/_credentials/base.py

@@ -0,0 +1,21 @@
+# ------------------------------------
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+# ------------------------------------
+import abc
+
+
+class AsyncCredentialBase(abc.ABC):
+    @abc.abstractmethod
+    async def close(self):
+        pass
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, *args):
+        await self.close()
+
+    @abc.abstractmethod
+    async def get_token(self, *scopes, **kwargs):
+        pass

sdk/identity/azure-identity/azure/identity/aio/_credentials/chained.py

@@ -2,17 +2,19 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 # ------------------------------------
+import asyncio
 from typing import TYPE_CHECKING
 
 from azure.core.exceptions import ClientAuthenticationError
 from ... import ChainedTokenCredential as SyncChainedTokenCredential
+from .base import AsyncCredentialBase
 
 if TYPE_CHECKING:
     from typing import Any
     from azure.core.credentials import AccessToken
 
 
-class ChainedTokenCredential(SyncChainedTokenCredential):
+class ChainedTokenCredential(SyncChainedTokenCredential, AsyncCredentialBase):
     """A sequence of credentials that is itself a credential.
 
     Its :func:`get_token` method calls ``get_token`` on each credential in the sequence, in order, returning the first
@@ -22,6 +24,11 @@ class ChainedTokenCredential(SyncChainedTokenCredential):
     :type credentials: :class:`azure.core.credentials.TokenCredential`
     """
 
+    async def close(self):
+        """Close the transport sessions of all credentials in the chain."""
+
+        await asyncio.gather(*(credential.close() for credential in self.credentials))
+
     async def get_token(self, *scopes: str, **kwargs: "Any") -> "AccessToken":  # pylint:disable=unused-argument
         """Asynchronously request a token from each credential, in order, returning the first token received.
 

sdk/identity/azure-identity/azure/identity/aio/_credentials/client_credential.py

@@ -4,6 +4,7 @@
 # ------------------------------------
 from typing import TYPE_CHECKING
 
+from .base import AsyncCredentialBase
 from .._authn_client import AsyncAuthnClient
 from ..._base import ClientSecretCredentialBase, CertificateCredentialBase
 
@@ -12,7 +13,7 @@
     from azure.core.credentials import AccessToken
 
 
-class ClientSecretCredential(ClientSecretCredentialBase):
+class ClientSecretCredential(ClientSecretCredentialBase, AsyncCredentialBase):
     """Authenticates as a service principal using a client ID and client secret.
 
     :param str tenant_id: ID of the service principal's tenant. Also called its 'directory' ID.
@@ -28,6 +29,15 @@ def __init__(self, tenant_id: str, client_id: str, client_secret: str, **kwargs:
         super(ClientSecretCredential, self).__init__(tenant_id, client_id, client_secret, **kwargs)
         self._client = AsyncAuthnClient(tenant=tenant_id, **kwargs)
 
+    async def __aenter__(self):
+        await self._client.__aenter__()
+        return self
+
+    async def close(self):
+        """Close the credential's transport session."""
+
+        await self._client.__aexit__()
+
     async def get_token(self, *scopes: str, **kwargs: "Any") -> "AccessToken":  # pylint:disable=unused-argument
         """Asynchronously request an access token for `scopes`.
 
@@ -44,7 +54,7 @@ async def get_token(self, *scopes: str, **kwargs: "Any") -> "AccessToken":  # py
         return token  # type: ignore
 
 
-class CertificateCredential(CertificateCredentialBase):
+class CertificateCredential(CertificateCredentialBase, AsyncCredentialBase):
     """Authenticates as a service principal using a certificate.
 
     :param str tenant_id: ID of the service principal's tenant. Also called its 'directory' ID.
@@ -57,6 +67,15 @@ class CertificateCredential(CertificateCredentialBase):
           defines authorities for other clouds.
     """
 
+    async def __aenter__(self):
+        await self._client.__aenter__()
+        return self
+
+    async def close(self):
+        """Close the credential's transport session."""
+
+        await self._client.__aexit__()
+
     async def get_token(self, *scopes: str, **kwargs: "Any") -> "AccessToken":  # pylint:disable=unused-argument
         """Asynchronously request an access token for `scopes`.
 

sdk/identity/azure-identity/azure/identity/aio/_credentials/environment.py

@@ -8,13 +8,14 @@
 from azure.core.exceptions import ClientAuthenticationError
 from ..._constants import EnvironmentVariables
 from .client_credential import CertificateCredential, ClientSecretCredential
+from .base import AsyncCredentialBase
 
 if TYPE_CHECKING:
     from typing import Any, Optional, Union
     from azure.core.credentials import AccessToken
 
 
-class EnvironmentCredential:
+class EnvironmentCredential(AsyncCredentialBase):
     """A credential configured by environment variables.
 
     This credential is capable of authenticating as a service principal using a client secret or a certificate, or as
@@ -50,6 +51,17 @@ def __init__(self, **kwargs: "Any") -> None:
                 **kwargs
             )
 
+    async def __aenter__(self):
+        if self._credential:
+            await self._credential.__aenter__()
+        return self
+
+    async def close(self):
+        """Close the credential's transport session."""
+
+        if self._credential:
+            await self._credential.__aexit__()
+
     async def get_token(self, *scopes: str, **kwargs: "Any") -> "AccessToken":
         """Asynchronously request an access token for `scopes`.
 

sdk/identity/azure-identity/azure/identity/aio/_credentials/managed_identity.py

@@ -2,6 +2,7 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 # ------------------------------------
+import abc
 import os
 from typing import TYPE_CHECKING
 
@@ -10,6 +11,7 @@
 from azure.core.pipeline.policies import AsyncRetryPolicy
 
 from azure.identity._credentials.managed_identity import _ManagedIdentityBase
+from .base import AsyncCredentialBase
 from .._authn_client import AsyncAuthnClient
 from ..._constants import Endpoints, EnvironmentVariables
 
@@ -37,6 +39,15 @@ def __new__(cls, *args, **kwargs):
     def __init__(self, **kwargs: "Any") -> None:
         pass
 
+    async def __aenter__(self):
+        pass
+
+    async def __aexit__(self, *args):
+        pass
+
+    async def close(self):
+        """Close the credential's transport session."""
+
     async def get_token(self, *scopes: str, **kwargs: "Any") -> "AccessToken":  # pylint:disable=unused-argument
         """Asynchronously request an access token for `scopes`.
 
@@ -49,10 +60,23 @@ async def get_token(self, *scopes: str, **kwargs: "Any") -> "AccessToken":  # py
         return AccessToken()
 
 
-class _AsyncManagedIdentityBase(_ManagedIdentityBase):
+class _AsyncManagedIdentityBase(_ManagedIdentityBase, AsyncCredentialBase):
     def __init__(self, endpoint: str, **kwargs: "Any") -> None:
         super().__init__(endpoint=endpoint, client_cls=AsyncAuthnClient, **kwargs)
 
+    async def __aenter__(self):
+        await self._client.__aenter__()
+        return self
+
+    async def close(self):
+        """Close the credential's transport session."""
+
+        await self._client.__aexit__()
+
+    @abc.abstractmethod
+    async def get_token(self, *scopes, **kwargs):
+        pass
+
     @staticmethod
     def _create_config(**kwargs: "Any") -> "Configuration":
         """Build a default configuration for the credential's HTTP pipeline."""

sdk/identity/azure-identity/azure/identity/aio/_credentials/shared_cache.py

@@ -9,21 +9,31 @@
 from ..._internal.shared_token_cache import NO_TOKEN, SharedTokenCacheBase
 from .._internal.aad_client import AadClient
 from .._internal.exception_wrapper import wrap_exceptions
+from .base import AsyncCredentialBase
 
 if TYPE_CHECKING:
     from typing import Any
     from azure.core.credentials import AccessToken
     from ..._internal.aad_client import AadClientBase
 
 
-class SharedTokenCacheCredential(SharedTokenCacheBase):
+class SharedTokenCacheCredential(SharedTokenCacheBase, AsyncCredentialBase):
     """Authenticates using tokens in the local cache shared between Microsoft applications.
 
     :param str username:
         Username (typically an email address) of the user to authenticate as. This is required because the local cache
         may contain tokens for multiple identities.
     """
 
+    async def __aenter__(self):
+        await self._client.__aenter__()
+        return self
+
+    async def close(self):
+        """Close the credential's transport session."""
+
+        await self._client.__aexit__()
+
     @wrap_exceptions
     async def get_token(self, *scopes: str, **kwargs: "Any") -> "AccessToken":  # pylint:disable=unused-argument
         """Get an access token for `scopes` from the shared cache.

sdk/identity/azure-identity/azure/identity/aio/_internal/aad_client.py

@@ -19,8 +19,19 @@
 
 
 class AadClient(AadClientBase):
-    # pylint:disable=arguments-differ
+    async def __aenter__(self):
+        await self._client.session.__aenter__()
+        return self
+
+    async def __aexit__(self, *args):
+        await self.close()
+
+    async def close(self) -> None:
+        """Close the client's transport session."""
 
+        await self._client.session.__aexit__()
+
+    # pylint:disable=arguments-differ
     def obtain_token_by_authorization_code(
         self, *args: "Any", loop: "asyncio.AbstractEventLoop" = None, **kwargs: "Any"
     ) -> "AccessToken":