[Identity] Allow configurable process timeouts (#28290)

AzureCliCredential, AzureDeveloperCliCredential, and AzurePowerShellCredential
now allow users to pass in a custom timeout. This addresses scenarios where
these proceses can take longer than the current default timeout values.
DefaultAzureCredential now also has an optional keyword argument to allow
users to pass in timeout values to the underlying developer credentials.

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

sdk/identity/azure-identity/CHANGELOG.md

@@ -4,7 +4,8 @@
 
 ### Features Added
 
-- Added proactive refreshing feature for managed iedentity
+- Added proactive refreshing feature for managed identity
+- Credentials that are implemented via launching a subprocess to acquire tokens now have configurable timeouts using the `process_timeout` keyword argument. This addresses scenarios where these proceses can take longer than the current default timeout values. The affected credentials are `AzureCliCredential`, `AzureDeveloperCliCredential`, and `AzurePowerShellCredential`. (Note: For `DefaultAzureCredential`, the `developer_credential_timeout` keyword argument allows users to propagate this option to `AzureCliCredential`, `AzureDeveloperCliCredential`, and `AzurePowerShellCredential` in the authentication chain.) ([#28290](https://github.com/Azure/azure-sdk-for-python/pull/28290))
 
 ### Breaking Changes
 

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

@@ -10,7 +10,7 @@
 import shutil
 import subprocess
 import sys
-from typing import Any, List, Optional
+from typing import Any, Dict, List, Optional
 import six
 
 from azure.core.credentials import AccessToken
@@ -37,12 +37,21 @@ class AzureDeveloperCliCredential:
     :keyword List[str] additionally_allowed_tenants: Specifies tenants in addition to the specified "tenant_id"
         for which the credential may acquire tokens. Add the wildcard value "*" to allow the credential to
         acquire tokens for any tenant the application can access.
+    :keyword int process_timeout: Seconds to wait for the Azure Developer CLI process to respond. Defaults
+        to 10 seconds.
     """
 
-    def __init__(self, *, tenant_id: str = "", additionally_allowed_tenants: Optional[List[str]] = None):
+    def __init__(
+        self,
+        *,
+        tenant_id: str = "",
+        additionally_allowed_tenants: Optional[List[str]] = None,
+        process_timeout: int = 10
+    ) -> None:
 
         self.tenant_id = tenant_id
         self._additionally_allowed_tenants = additionally_allowed_tenants or []
+        self._process_timeout = process_timeout
 
     def __enter__(self) -> "AzureDeveloperCliCredential":
         return self
@@ -83,7 +92,7 @@ def get_token(self, *scopes: str, **kwargs: Any) -> AccessToken:
         )
         if tenant:
             command += " --tenant-id " + tenant
-        output = _run_command(command)
+        output = _run_command(command, self._process_timeout)
 
         token = parse_token(output)
         if not token:
@@ -130,7 +139,7 @@ def sanitize_output(output):
     return re.sub(r"\"token\": \"(.*?)(\"|$)", "****", output)
 
 
-def _run_command(command):
+def _run_command(command: str, timeout: int) -> str:
     # Ensure executable exists in PATH first. This avoids a subprocess call that would fail anyway.
     if shutil.which(EXECUTABLE_NAME) is None:
         raise CredentialUnavailableError(message=CLI_NOT_FOUND)
@@ -142,12 +151,12 @@ def _run_command(command):
     try:
         working_directory = get_safe_working_dir()
 
-        kwargs = {
+        kwargs: Dict[str, Any] = {
             "stderr": subprocess.PIPE,
             "cwd": working_directory,
             "universal_newlines": True,
             "env": dict(os.environ, NO_COLOR="true"),
-            "timeout": 10,
+            "timeout": timeout,
         }
 
         return subprocess.check_output(args, **kwargs)

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

@@ -10,7 +10,7 @@
 import subprocess
 import sys
 import time
-from typing import List, Optional, Any
+from typing import List, Optional, Any, Dict
 import six
 
 from azure.core.credentials import AccessToken
@@ -36,16 +36,19 @@ class AzureCliCredential:
     :keyword List[str] additionally_allowed_tenants: Specifies tenants in addition to the specified "tenant_id"
         for which the credential may acquire tokens. Add the wildcard value "*" to allow the credential to
         acquire tokens for any tenant the application can access.
+    :keyword int process_timeout: Seconds to wait for the Azure CLI process to respond. Defaults to 10 seconds.
     """
     def __init__(
         self,
         *,
         tenant_id: str = "",
-        additionally_allowed_tenants: Optional[List[str]] = None
+        additionally_allowed_tenants: Optional[List[str]] = None,
+        process_timeout: int = 10
     ) -> None:
 
         self.tenant_id = tenant_id
         self._additionally_allowed_tenants = additionally_allowed_tenants or []
+        self._process_timeout = process_timeout
 
     def __enter__(self):
         return self
@@ -84,7 +87,7 @@ def get_token(self, *scopes: str, **kwargs: Any) -> AccessToken:
         )
         if tenant:
             command += " --tenant " + tenant
-        output = _run_command(command)
+        output = _run_command(command, self._process_timeout)
 
         token = parse_token(output)
         if not token:
@@ -135,7 +138,7 @@ def sanitize_output(output: str) -> str:
     return re.sub(r"\"accessToken\": \"(.*?)(\"|$)", "****", output)
 
 
-def _run_command(command):
+def _run_command(command: str, timeout: int) -> str:
     # Ensure executable exists in PATH first. This avoids a subprocess call that would fail anyway.
     if shutil.which(EXECUTABLE_NAME) is None:
         raise CredentialUnavailableError(message=CLI_NOT_FOUND)
@@ -147,11 +150,11 @@ def _run_command(command):
     try:
         working_directory = get_safe_working_dir()
 
-        kwargs = {
+        kwargs: Dict[str, Any] = {
             "stderr": subprocess.PIPE,
             "cwd": working_directory,
             "universal_newlines": True,
-            "timeout": 10,
+            "timeout": timeout,
             "env": dict(os.environ, AZURE_CORE_NO_COLOR="true"),
         }
         return subprocess.check_output(args, **kwargs)

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

@@ -4,7 +4,6 @@
 # ------------------------------------
 import base64
 import logging
-import platform
 import subprocess
 import sys
 from typing import List, Tuple, Optional, Any
@@ -51,16 +50,19 @@ class AzurePowerShellCredential:
     :keyword List[str] additionally_allowed_tenants: Specifies tenants in addition to the specified "tenant_id"
         for which the credential may acquire tokens. Add the wildcard value "*" to allow the credential to
         acquire tokens for any tenant the application can access.
+    :keyword int process_timeout: Seconds to wait for the Azure PowerShell process to respond. Defaults to 10 seconds.
     """
     def __init__(
         self,
         *,
         tenant_id: str = "",
-        additionally_allowed_tenants: Optional[List[str]] = None
+        additionally_allowed_tenants: Optional[List[str]] = None,
+        process_timeout: int = 10
     ) -> None:
 
         self.tenant_id = tenant_id
         self._additionally_allowed_tenants = additionally_allowed_tenants or []
+        self._process_timeout = process_timeout
 
     def __enter__(self):
         return self
@@ -96,17 +98,15 @@ def get_token(self, *scopes: str, **kwargs: Any) -> AccessToken:
             **kwargs
         )
         command_line = get_command_line(scopes, tenant_id)
-        output = run_command_line(command_line)
+        output = run_command_line(command_line, self._process_timeout)
         token = parse_token(output)
         return token
 
 
-def run_command_line(command_line: List[str]) -> str:
+def run_command_line(command_line: List[str], timeout: int) -> str:
     stdout = stderr = ""
     proc = None
-    kwargs = {}
-    if platform.python_version() >= "3.3":
-        kwargs["timeout"] = 10
+    kwargs = {"timeout": timeout}
 
     try:
         proc = start_process(command_line)

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

@@ -80,6 +80,8 @@ class DefaultAzureCredential(ChainedTokenCredential):
         :class:`~azure.identity.VisualStudioCodeCredential`. Defaults to the "Azure: Tenant" setting in VS Code's user
         settings or, when that setting has no value, the "organizations" tenant, which supports only Azure Active
         Directory work or school accounts.
+    :keyword int developer_credential_timeout: The timeout in seconds to use for developer credentials that run
+        subprocesses (e.g. AzureCliCredential, AzurePowerShellCredential). Defaults to **10** seconds.
     """
 
     def __init__(self, **kwargs: Any) -> None:  # pylint: disable=too-many-statements
@@ -116,6 +118,8 @@ def __init__(self, **kwargs: Any) -> None:  # pylint: disable=too-many-statement
             "shared_cache_tenant_id", os.environ.get(EnvironmentVariables.AZURE_TENANT_ID)
         )
 
+        developer_credential_timeout = kwargs.pop("developer_credential_timeout", 10)
+
         exclude_environment_credential = kwargs.pop("exclude_environment_credential", False)
         exclude_managed_identity_credential = kwargs.pop("exclude_managed_identity_credential", False)
         exclude_shared_token_cache_credential = kwargs.pop("exclude_shared_token_cache_credential", False)
@@ -138,7 +142,7 @@ def __init__(self, **kwargs: Any) -> None:  # pylint: disable=too-many-statement
         if not exclude_managed_identity_credential:
             credentials.append(ManagedIdentityCredential(client_id=managed_identity_client_id, **kwargs))
         if not exclude_azd_cli_credential:
-            credentials.append(AzureDeveloperCliCredential())
+            credentials.append(AzureDeveloperCliCredential(process_timeout=developer_credential_timeout))
         if not exclude_shared_token_cache_credential and SharedTokenCacheCredential.supported():
             try:
                 # username and/or tenant_id are only required when the cache contains tokens for multiple identities
@@ -151,9 +155,9 @@ def __init__(self, **kwargs: Any) -> None:  # pylint: disable=too-many-statement
         if not exclude_visual_studio_code_credential:
             credentials.append(VisualStudioCodeCredential(**vscode_args))
         if not exclude_cli_credential:
-            credentials.append(AzureCliCredential())
+            credentials.append(AzureCliCredential(process_timeout=developer_credential_timeout))
         if not exclude_powershell_credential:
-            credentials.append(AzurePowerShellCredential())
+            credentials.append(AzurePowerShellCredential(process_timeout=developer_credential_timeout))
         if not exclude_interactive_browser_credential:
             if interactive_browser_client_id:
                 credentials.append(

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

@@ -35,12 +35,21 @@ class AzureDeveloperCliCredential(AsyncContextManager):
     :keyword List[str] additionally_allowed_tenants: Specifies tenants in addition to the specified "tenant_id"
         for which the credential may acquire tokens. Add the wildcard value "*" to allow the credential to
         acquire tokens for any tenant the application can access.
+    :keyword int process_timeout: Seconds to wait for the Azure Developer CLI process to respond. Defaults
+        to 10 seconds.
     """
 
-    def __init__(self, *, tenant_id: str = "", additionally_allowed_tenants: Optional[List[str]] = None):
+    def __init__(
+        self,
+        *,
+        tenant_id: str = "",
+        additionally_allowed_tenants: Optional[List[str]] = None,
+        process_timeout: int = 10
+    ) -> None:
 
         self.tenant_id = tenant_id
         self._additionally_allowed_tenants = additionally_allowed_tenants or []
+        self._process_timeout = process_timeout
 
     @log_get_token_async
     async def get_token(self, *scopes: str, **kwargs: Any) -> AccessToken:
@@ -73,7 +82,7 @@ async def get_token(self, *scopes: str, **kwargs: Any) -> AccessToken:
 
         if tenant:
             command += " --tenant-id " + tenant
-        output = await _run_command(command)
+        output = await _run_command(command, self._process_timeout)
 
         token = parse_token(output)
         if not token:
@@ -88,7 +97,7 @@ async def close(self) -> None:
         """Calling this method is unnecessary"""
 
 
-async def _run_command(command: str) -> str:
+async def _run_command(command: str, timeout: int) -> str:
     # Ensure executable exists in PATH first. This avoids a subprocess call that would fail anyway.
     if shutil.which(EXECUTABLE_NAME) is None:
         raise CredentialUnavailableError(message=CLI_NOT_FOUND)
@@ -108,7 +117,7 @@ async def _run_command(command: str) -> str:
             cwd=working_directory,
             env=dict(os.environ, NO_COLOR="true")
         )
-        stdout_b, stderr_b = await asyncio.wait_for(proc.communicate(), 10)
+        stdout_b, stderr_b = await asyncio.wait_for(proc.communicate(), timeout)
         output = stdout_b.decode()
         stderr = stderr_b.decode()
     except OSError as ex:

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

@@ -35,16 +35,19 @@ class AzureCliCredential(AsyncContextManager):
     :keyword List[str] additionally_allowed_tenants: Specifies tenants in addition to the specified "tenant_id"
         for which the credential may acquire tokens. Add the wildcard value "*" to allow the credential to
         acquire tokens for any tenant the application can access.
+    :keyword int process_timeout: Seconds to wait for the Azure CLI process to respond. Defaults to 10 seconds.
     """
     def __init__(
         self,
         *,
         tenant_id: str = "",
-        additionally_allowed_tenants: Optional[List[str]] = None
+        additionally_allowed_tenants: Optional[List[str]] = None,
+        process_timeout: int = 10
     ) -> None:
 
         self.tenant_id = tenant_id
         self._additionally_allowed_tenants = additionally_allowed_tenants or []
+        self._process_timeout = process_timeout
 
     @log_get_token_async
     async def get_token(self, *scopes: str, **kwargs: Any) -> AccessToken:
@@ -76,7 +79,7 @@ async def get_token(self, *scopes: str, **kwargs: Any) -> AccessToken:
 
         if tenant:
             command += " --tenant " + tenant
-        output = await _run_command(command)
+        output = await _run_command(command, self._process_timeout)
 
         token = parse_token(output)
         if not token:
@@ -92,7 +95,7 @@ async def close(self) -> None:
         """Calling this method is unnecessary"""
 
 
-async def _run_command(command: str) -> str:
+async def _run_command(command: str, timeout: int) -> str:
     # Ensure executable exists in PATH first. This avoids a subprocess call that would fail anyway.
     if shutil.which(EXECUTABLE_NAME) is None:
         raise CredentialUnavailableError(message=CLI_NOT_FOUND)
@@ -112,7 +115,7 @@ async def _run_command(command: str) -> str:
             cwd=working_directory,
             env=dict(os.environ, AZURE_CORE_NO_COLOR="true")
         )
-        stdout_b, stderr_b = await asyncio.wait_for(proc.communicate(), 10)
+        stdout_b, stderr_b = await asyncio.wait_for(proc.communicate(), timeout)
         output = stdout_b.decode()
         stderr = stderr_b.decode()
     except OSError as ex:

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

@@ -29,12 +29,20 @@ class AzurePowerShellCredential(AsyncContextManager):
     :keyword List[str] additionally_allowed_tenants: Specifies tenants in addition to the specified "tenant_id"
         for which the credential may acquire tokens. Add the wildcard value "*" to allow the credential to
         acquire tokens for any tenant the application can access.
+    :keyword int process_timeout: Seconds to wait for the Azure PowerShell process to respond. Defaults to 10 seconds.
     """
 
-    def __init__(self, *, tenant_id: str = "", additionally_allowed_tenants: Optional[List[str]] = None):
+    def __init__(
+        self,
+        *,
+        tenant_id: str = "",
+        additionally_allowed_tenants: Optional[List[str]] = None,
+        process_timeout: int = 10
+    ) -> None:
 
         self.tenant_id = tenant_id
         self._additionally_allowed_tenants = additionally_allowed_tenants or []
+        self._process_timeout = process_timeout
 
     @log_get_token_async
     async def get_token(
@@ -65,23 +73,23 @@ async def get_token(
             **kwargs
         )
         command_line = get_command_line(scopes, tenant_id)
-        output = await run_command_line(command_line)
+        output = await run_command_line(command_line, self._process_timeout)
         token = parse_token(output)
         return token
 
     async def close(self) -> None:
         """Calling this method is unnecessary"""
 
 
-async def run_command_line(command_line: List[str]) -> str:
+async def run_command_line(command_line: List[str], timeout: int) -> str:
     try:
         proc = await start_process(command_line)
         stdout, stderr = await asyncio.wait_for(proc.communicate(), 10)
         if sys.platform.startswith("win") and b"' is not recognized" in stderr:
             # pwsh.exe isn't on the path; try powershell.exe
             command_line[-1] = command_line[-1].replace("pwsh", "powershell", 1)
             proc = await start_process(command_line)
-            stdout, stderr = await asyncio.wait_for(proc.communicate(), 10)
+            stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout)
 
     except OSError as ex:
         # failed to execute "cmd" or "/bin/sh"; Azure PowerShell may or may not be installed