[Identity] Implement new protocol for all credentials (#36882)

All credentials now implement the `SupportsTokenInfo/AsyncSupportsTokenInfo` protocol,
by each having a `get_token_info` method implementation. This allows for more extensible
authentication constructs.

Signed-off-by: Paul Van Eck <[email protected]>

Author: pvaneck

sdk/identity/azure-identity/CHANGELOG.md

@@ -4,6 +4,9 @@
 
 ### Features Added
 
+- All credentials now support the `SupportsTokenInfo` protocol. Each credential now has a `get_token_info` method which returns an `AccessTokenInfo` object. The `get_token_info` method is an alternative method to `get_token` that improves support support for more complex authentication scenarios. ([#36882](https://github.com/Azure/azure-sdk-for-python/pull/36882))
+    - Information on when a token should be refreshed is now saved in `AccessTokenInfo` (if available).
+
 ### Breaking Changes
 
 ### Bugs Fixed
@@ -12,6 +15,7 @@
 
 - Added identity config validation to `ManagedIdentityCredential` to avoid non-deterministic states (e.g. both `resource_id` and `object_id` are specified). ([#36950](https://github.com/Azure/azure-sdk-for-python/pull/36950))
 - Additional validation was added for `ManagedIdentityCredential` in Azure Cloud Shell environments. ([#36438](https://github.com/Azure/azure-sdk-for-python/issues/36438))
+- Bumped minimum dependency on `azure-core` to `>=1.31.0`.
 
 ## 1.18.0b2 (2024-08-09)
 

sdk/identity/azure-identity/assets.json

@@ -2,5 +2,5 @@
   "AssetsRepo": "Azure/azure-sdk-assets",
   "AssetsRepoPrefixPath": "python",
   "TagPrefix": "python/identity/azure-identity",
-  "Tag": "python/identity/azure-identity_cb8dd6f319"
+  "Tag": "python/identity/azure-identity_61e626a4a0"
 }

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

@@ -4,7 +4,7 @@
 # ------------------------------------
 from typing import Callable
 
-from azure.core.credentials import TokenCredential
+from azure.core.credentials import TokenProvider
 from azure.core.pipeline.policies import BearerTokenCredentialPolicy
 from azure.core.pipeline import PipelineRequest, PipelineContext
 from azure.core.rest import HttpRequest
@@ -14,7 +14,7 @@ def _make_request() -> PipelineRequest[HttpRequest]:
     return PipelineRequest(HttpRequest("CredentialWrapper", "https://fakeurl"), PipelineContext(None))
 
 
-def get_bearer_token_provider(credential: TokenCredential, *scopes: str) -> Callable[[], str]:
+def get_bearer_token_provider(credential: TokenProvider, *scopes: str) -> Callable[[], str]:
     """Returns a callable that provides a bearer token.
 
     It can be used for instance to write code like:

sdk/identity/azure-identity/azure/identity/_credentials/application.py

@@ -4,9 +4,9 @@
 # ------------------------------------
 import logging
 import os
-from typing import Any, Optional
+from typing import Any, Optional, cast
 
-from azure.core.credentials import AccessToken
+from azure.core.credentials import AccessToken, AccessTokenInfo, TokenRequestOptions, SupportsTokenInfo, TokenCredential
 from .chained import ChainedTokenCredential
 from .environment import EnvironmentCredential
 from .managed_identity import ManagedIdentityCredential
@@ -83,10 +83,37 @@ def get_token(
             `message` attribute listing each authentication attempt and its error message.
         """
         if self._successful_credential:
-            token = self._successful_credential.get_token(*scopes, claims=claims, tenant_id=tenant_id, **kwargs)
+            token = cast(TokenCredential, self._successful_credential).get_token(
+                *scopes, claims=claims, tenant_id=tenant_id, **kwargs
+            )
             _LOGGER.info(
                 "%s acquired a token from %s", self.__class__.__name__, self._successful_credential.__class__.__name__
             )
             return token
 
         return super(AzureApplicationCredential, self).get_token(*scopes, claims=claims, tenant_id=tenant_id, **kwargs)
+
+    def get_token_info(self, *scopes: str, options: Optional[TokenRequestOptions] = None) -> AccessTokenInfo:
+        """Request an access token for `scopes`.
+
+        This is an alternative to `get_token` to enable certain scenarios that require additional properties
+        on the token. This method is called automatically by Azure SDK clients.
+
+        :param str scopes: desired scopes for the access token. This method requires at least one scope.
+            For more information about scopes, see https://learn.microsoft.com/entra/identity-platform/scopes-oidc.
+        :keyword options: A dictionary of options for the token request. Unknown options will be ignored. Optional.
+        :paramtype options: ~azure.core.credentials.TokenRequestOptions
+
+        :rtype: AccessTokenInfo
+        :return: An AccessTokenInfo instance containing information about the token.
+        :raises ~azure.core.exceptions.ClientAuthenticationError: authentication failed. The exception has a
+            `message` attribute listing each authentication attempt and its error message.
+        """
+        if self._successful_credential:
+            token_info = cast(SupportsTokenInfo, self._successful_credential).get_token_info(*scopes, options=options)
+            _LOGGER.info(
+                "%s acquired a token from %s", self.__class__.__name__, self._successful_credential.__class__.__name__
+            )
+            return token_info
+
+        return cast(SupportsTokenInfo, super()).get_token_info(*scopes, options=options)

sdk/identity/azure-identity/azure/identity/_credentials/authorization_code.py

@@ -4,7 +4,7 @@
 # ------------------------------------
 from typing import Optional, Any
 
-from azure.core.credentials import AccessToken
+from azure.core.credentials import AccessToken, AccessTokenInfo, TokenRequestOptions
 from azure.core.exceptions import ClientAuthenticationError
 from .._internal.aad_client import AadClient
 from .._internal.get_token_mixin import GetTokenMixin
@@ -90,10 +90,35 @@ def get_token(
             *scopes, claims=claims, tenant_id=tenant_id, client_secret=self._client_secret, **kwargs
         )
 
-    def _acquire_token_silently(self, *scopes: str, **kwargs) -> Optional[AccessToken]:
+    def get_token_info(self, *scopes: str, options: Optional[TokenRequestOptions] = None) -> AccessTokenInfo:
+        """Request an access token for `scopes`.
+
+        This is an alternative to `get_token` to enable certain scenarios that require additional properties
+        on the token. This method is called automatically by Azure SDK clients.
+
+        The first time this method is called, the credential will redeem its authorization code. On subsequent calls
+        the credential will return a cached access token or redeem a refresh token, if it acquired a refresh token upon
+        redeeming the authorization code.
+
+        :param str scopes: desired scopes for the access token. This method requires at least one scope.
+            For more information about scopes, see https://learn.microsoft.com/entra/identity-platform/scopes-oidc.
+        :keyword options: A dictionary of options for the token request. Unknown options will be ignored. Optional.
+        :paramtype options: ~azure.core.credentials.TokenRequestOptions
+
+        :rtype: AccessTokenInfo
+        :return: An AccessTokenInfo instance containing information about the token.
+        :raises ~azure.core.exceptions.ClientAuthenticationError: authentication failed. The error's ``message``
+          attribute gives a reason. Any error response from Microsoft Entra ID is available as the error's
+          ``response`` attribute.
+        """
+        return super()._get_token_base(
+            *scopes, options=options, client_secret=self._client_secret, base_method_name="get_token_info"
+        )
+
+    def _acquire_token_silently(self, *scopes: str, **kwargs) -> Optional[AccessTokenInfo]:
         return self._client.get_cached_access_token(scopes, **kwargs)
 
-    def _request_token(self, *scopes: str, **kwargs) -> AccessToken:
+    def _request_token(self, *scopes: str, **kwargs) -> AccessTokenInfo:
         if self._authorization_code:
             token = self._client.obtain_token_by_authorization_code(
                 scopes=scopes, code=self._authorization_code, redirect_uri=self._redirect_uri, **kwargs

sdk/identity/azure-identity/azure/identity/_credentials/azd_cli.py

@@ -12,7 +12,7 @@
 import sys
 from typing import Any, Dict, List, Optional
 
-from azure.core.credentials import AccessToken
+from azure.core.credentials import AccessToken, AccessTokenInfo, TokenRequestOptions
 from azure.core.exceptions import ClientAuthenticationError
 
 from .. import CredentialUnavailableError
@@ -118,10 +118,43 @@ def get_token(
         :raises ~azure.core.exceptions.ClientAuthenticationError: the credential invoked
           the Azure Developer CLI but didn't receive an access token.
         """
+        options: TokenRequestOptions = {}
+        if tenant_id:
+            options["tenant_id"] = tenant_id
+
+        token_info = self._get_token_base(*scopes, options=options, **kwargs)
+        return AccessToken(token_info.token, token_info.expires_on)
+
+    @log_get_token
+    def get_token_info(self, *scopes: str, options: Optional[TokenRequestOptions] = None) -> AccessTokenInfo:
+        """Request an access token for `scopes`.
+
+        This is an alternative to `get_token` to enable certain scenarios that require additional properties
+        on the token. This method is called automatically by Azure SDK clients. Applications calling this method
+        directly must also handle token caching because this credential doesn't cache the tokens it acquires.
+
+        :param str scopes: desired scopes for the access token. This method requires at least one scope.
+            For more information about scopes, see https://learn.microsoft.com/entra/identity-platform/scopes-oidc.
+        :keyword options: A dictionary of options for the token request. Unknown options will be ignored. Optional.
+        :paramtype options: ~azure.core.credentials.TokenRequestOptions
+
+        :rtype: AccessTokenInfo
+        :return: An AccessTokenInfo instance containing information about the token.
+
+        :raises ~azure.identity.CredentialUnavailableError: the credential was unable to invoke
+          the Azure Developer CLI.
+        :raises ~azure.core.exceptions.ClientAuthenticationError: the credential invoked
+          the Azure Developer CLI but didn't receive an access token.
+        """
+        return self._get_token_base(*scopes, options=options)
 
+    def _get_token_base(
+        self, *scopes: str, options: Optional[TokenRequestOptions] = None, **kwargs: Any
+    ) -> AccessTokenInfo:
         if not scopes:
             raise ValueError("Missing scope in request. \n")
 
+        tenant_id = options.get("tenant_id") if options else None
         if tenant_id:
             validate_tenant_id(tenant_id)
         for scope in scopes:
@@ -154,7 +187,7 @@ def get_token(
         return token
 
 
-def parse_token(output: str) -> Optional[AccessToken]:
+def parse_token(output: str) -> Optional[AccessTokenInfo]:
     """Parse to an AccessToken.
 
     In particular, convert the "expiresOn" value to epoch seconds. This value is a naive local datetime as returned by
@@ -169,7 +202,7 @@ def parse_token(output: str) -> Optional[AccessToken]:
         dt = datetime.strptime(token["expiresOn"], "%Y-%m-%dT%H:%M:%SZ")
         expires_on = dt.timestamp()
 
-        return AccessToken(token["token"], int(expires_on))
+        return AccessTokenInfo(token["token"], int(expires_on))
     except (KeyError, ValueError):
         return None
 

sdk/identity/azure-identity/azure/identity/_credentials/azure_cli.py

@@ -11,7 +11,7 @@
 import sys
 from typing import List, Optional, Any, Dict
 
-from azure.core.credentials import AccessToken
+from azure.core.credentials import AccessToken, AccessTokenInfo, TokenRequestOptions
 from azure.core.exceptions import ClientAuthenticationError
 
 from .. import CredentialUnavailableError
@@ -94,6 +94,41 @@ def get_token(
         :raises ~azure.core.exceptions.ClientAuthenticationError: the credential invoked the Azure CLI but didn't
           receive an access token.
         """
+
+        options: TokenRequestOptions = {}
+        if tenant_id:
+            options["tenant_id"] = tenant_id
+
+        token_info = self._get_token_base(*scopes, options=options, **kwargs)
+        return AccessToken(token_info.token, token_info.expires_on)
+
+    @log_get_token
+    def get_token_info(self, *scopes: str, options: Optional[TokenRequestOptions] = None) -> AccessTokenInfo:
+        """Request an access token for `scopes`.
+
+        This is an alternative to `get_token` to enable certain scenarios that require additional properties
+        on the token. This method is called automatically by Azure SDK clients. Applications calling this method
+        directly must also handle token caching because this credential doesn't cache the tokens it acquires.
+
+        :param str scopes: desired scopes for the access token. This credential allows only one scope per request.
+            For more information about scopes, see https://learn.microsoft.com/entra/identity-platform/scopes-oidc.
+        :keyword options: A dictionary of options for the token request. Unknown options will be ignored. Optional.
+        :paramtype options: ~azure.core.credentials.TokenRequestOptions
+
+        :rtype: AccessTokenInfo
+        :return: An AccessTokenInfo instance containing information about the token.
+
+        :raises ~azure.identity.CredentialUnavailableError: the credential was unable to invoke the Azure CLI.
+        :raises ~azure.core.exceptions.ClientAuthenticationError: the credential invoked the Azure CLI but didn't
+          receive an access token.
+        """
+        return self._get_token_base(*scopes, options=options)
+
+    def _get_token_base(
+        self, *scopes: str, options: Optional[TokenRequestOptions] = None, **kwargs: Any
+    ) -> AccessTokenInfo:
+
+        tenant_id = options.get("tenant_id") if options else None
         if tenant_id:
             validate_tenant_id(tenant_id)
         for scope in scopes:
@@ -126,7 +161,7 @@ def get_token(
         return token
 
 
-def parse_token(output) -> Optional[AccessToken]:
+def parse_token(output) -> Optional[AccessTokenInfo]:
     """Parse output of 'az account get-access-token' to an AccessToken.
 
     In particular, convert the "expiresOn" value to epoch seconds. This value is a naive local datetime as returned by
@@ -141,11 +176,11 @@ def parse_token(output) -> Optional[AccessToken]:
 
         # Use "expires_on" if it's present, otherwise use "expiresOn".
         if "expires_on" in token:
-            return AccessToken(token["accessToken"], int(token["expires_on"]))
+            return AccessTokenInfo(token["accessToken"], int(token["expires_on"]))
 
         dt = datetime.strptime(token["expiresOn"], "%Y-%m-%d %H:%M:%S.%f")
         expires_on = dt.timestamp()
-        return AccessToken(token["accessToken"], int(expires_on))
+        return AccessTokenInfo(token["accessToken"], int(expires_on))
     except (KeyError, ValueError):
         return None
 

sdk/identity/azure-identity/azure/identity/_credentials/azure_pipelines.py

@@ -7,7 +7,7 @@
 from typing import Any, Optional
 
 from azure.core.exceptions import ClientAuthenticationError
-from azure.core.credentials import AccessToken
+from azure.core.credentials import AccessToken, AccessTokenInfo, TokenRequestOptions
 from azure.core.rest import HttpRequest, HttpResponse
 
 from .client_assertion import ClientAssertionCredential
@@ -125,6 +125,25 @@ def get_token(
             *scopes, claims=claims, tenant_id=tenant_id, enable_cae=enable_cae, **kwargs
         )
 
+    def get_token_info(self, *scopes: str, options: Optional[TokenRequestOptions] = None) -> AccessTokenInfo:
+        """Request an access token for `scopes`.
+
+        This is an alternative to `get_token` to enable certain scenarios that require additional properties
+        on the token. This method is called automatically by Azure SDK clients.
+
+        :param str scopes: desired scope for the access token. This method requires at least one scope.
+            For more information about scopes, see https://learn.microsoft.com/entra/identity-platform/scopes-oidc.
+        :keyword options: A dictionary of options for the token request. Unknown options will be ignored. Optional.
+        :paramtype options: ~azure.core.credentials.TokenRequestOptions
+
+        :rtype: AccessTokenInfo
+        :return: An AccessTokenInfo instance containing information about the token.
+        :raises ~azure.core.exceptions.ClientAuthenticationError: authentication failed. The error's ``message``
+            attribute gives a reason.
+        """
+        validate_env_vars()
+        return self._client_assertion_credential.get_token_info(*scopes, options=options)
+
     def _get_oidc_token(self) -> str:
         request = build_oidc_request(self._service_connection_id, self._system_access_token)
         response = self._pipeline.run(request, retry_on_methods=[request.method])

sdk/identity/azure-identity/azure/identity/_credentials/azure_powershell.py

@@ -8,7 +8,7 @@
 import sys
 from typing import Any, List, Tuple, Optional
 
-from azure.core.credentials import AccessToken
+from azure.core.credentials import AccessToken, AccessTokenInfo, TokenRequestOptions
 from azure.core.exceptions import ClientAuthenticationError
 
 from .azure_cli import get_safe_working_dir
@@ -125,6 +125,42 @@ def get_token(
         :raises ~azure.core.exceptions.ClientAuthenticationError: the credential invoked Azure PowerShell but didn't
           receive an access token
         """
+
+        options: TokenRequestOptions = {}
+        if tenant_id:
+            options["tenant_id"] = tenant_id
+
+        token_info = self._get_token_base(*scopes, options=options, **kwargs)
+        return AccessToken(token_info.token, token_info.expires_on)
+
+    @log_get_token
+    def get_token_info(self, *scopes: str, options: Optional[TokenRequestOptions] = None) -> AccessTokenInfo:
+        """Request an access token for `scopes`.
+
+        This is an alternative to `get_token` to enable certain scenarios that require additional properties
+        on the token. This method is called automatically by Azure SDK clients. Applications calling this method
+        directly must also handle token caching because this credential doesn't cache the tokens it acquires.
+
+        :param str scopes: desired scopes for the access token. TThis credential allows only one scope per request.
+            For more information about scopes, see https://learn.microsoft.com/entra/identity-platform/scopes-oidc.
+        :keyword options: A dictionary of options for the token request. Unknown options will be ignored. Optional.
+        :paramtype options: ~azure.core.credentials.TokenRequestOptions
+
+        :rtype: AccessTokenInfo
+        :return: An AccessTokenInfo instance containing information about the token.
+
+        :raises ~azure.identity.CredentialUnavailableError: the credential was unable to invoke Azure PowerShell, or
+          no account is authenticated
+        :raises ~azure.core.exceptions.ClientAuthenticationError: the credential invoked Azure PowerShell but didn't
+          receive an access token
+        """
+        return self._get_token_base(*scopes, options=options)
+
+    def _get_token_base(
+        self, *scopes: str, options: Optional[TokenRequestOptions] = None, **kwargs: Any
+    ) -> AccessTokenInfo:
+
+        tenant_id = options.get("tenant_id") if options else None
         if tenant_id:
             validate_tenant_id(tenant_id)
         for scope in scopes:
@@ -185,11 +221,11 @@ def start_process(args: List[str]) -> "subprocess.Popen":
     return proc
 
 
-def parse_token(output: str) -> AccessToken:
+def parse_token(output: str) -> AccessTokenInfo:
     for line in output.split():
         if line.startswith("azsdk%"):
             _, token, expires_on = line.split("%")
-            return AccessToken(token, int(expires_on))
+            return AccessTokenInfo(token, int(expires_on))
 
     if within_dac.get():
         raise CredentialUnavailableError(message='Unexpected output from Get-AzAccessToken: "{}"'.format(output))