[Core] Update async auth policy lock logic (#33282)

Sometimes users might be using trio and its mechanisms for
running tasks concurrently. This updates the lock initilization logic to
check if the user is in an asyncio event loop and sets the lock
accordingly.  If not, assume trio. We don't expect users to be in anything else
other than asyncio or trio event loops.

A utility function was added that tries to determine the current async
library being used.

Signed-off-by: Paul Van Eck <[email protected]>
Co-authored-by: Scott Beddall (from Dev Box) <[email protected]>

Author: pvaneck

sdk/core/azure-core/CHANGELOG.md

@@ -10,6 +10,8 @@
 
 ### Other Changes
 
+- Removed dependency on `anyio`.  #33282
+
 ## 1.29.6 (2023-12-14)
 
 ### Bugs Fixed

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

@@ -6,7 +6,6 @@
 import time
 from typing import TYPE_CHECKING, Any, Awaitable, Optional, cast, TypeVar
 
-from anyio import Lock
 from azure.core.credentials import AccessToken
 from azure.core.pipeline import PipelineRequest, PipelineResponse
 from azure.core.pipeline.policies import AsyncHTTPPolicy
@@ -15,6 +14,7 @@
 )
 from azure.core.pipeline.transport import AsyncHttpResponse as LegacyAsyncHttpResponse, HttpRequest as LegacyHttpRequest
 from azure.core.rest import AsyncHttpResponse, HttpRequest
+from azure.core.utils._utils import get_running_async_lock
 
 from .._tools_async import await_result
 
@@ -38,11 +38,17 @@ class AsyncBearerTokenCredentialPolicy(AsyncHTTPPolicy[HTTPRequestType, AsyncHTT
     def __init__(self, credential: "AsyncTokenCredential", *scopes: str, **kwargs: Any) -> None:
         super().__init__()
         self._credential = credential
-        self._lock = Lock()
         self._scopes = scopes
+        self._lock_instance = None
         self._token: Optional["AccessToken"] = None
         self._enable_cae: bool = kwargs.get("enable_cae", False)
 
+    @property
+    def _lock(self):
+        if self._lock_instance is None:
+            self._lock_instance = get_running_async_lock()
+        return self._lock_instance
+
     async def on_request(self, request: PipelineRequest[HTTPRequestType]) -> None:
         """Adds a bearer token Authorization header to request and sends request to next policy.
 

sdk/core/azure-core/azure/core/utils/_utils.py

@@ -5,8 +5,10 @@
 # license information.
 # --------------------------------------------------------------------------
 import datetime
+import sys
 from typing import (
     Any,
+    AsyncContextManager,
     Iterable,
     Iterator,
     Mapping,
@@ -161,3 +163,26 @@ def __eq__(self, other: Any) -> bool:
 
     def __repr__(self) -> str:
         return str(dict(self.items()))
+
+
+def get_running_async_lock() -> AsyncContextManager:
+    """Get a lock instance from the async library that the current context is running under.
+
+    :return: An instance of the running async library's Lock class.
+    :rtype: AsyncContextManager
+    :raises: RuntimeError if the current context is not running under an async library.
+    """
+
+    try:
+        import asyncio
+
+        # Check if we are running in an asyncio event loop.
+        asyncio.get_running_loop()
+        return asyncio.Lock()
+    except RuntimeError as err:
+        # Otherwise, assume we are running in a trio event loop if it has already been imported.
+        if "trio" in sys.modules:
+            import trio  # pylint: disable=networking-import-outside-azure-core-transport
+
+            return trio.Lock()
+        raise RuntimeError("An asyncio or trio event loop is required.") from err

sdk/core/azure-core/setup.py

@@ -69,7 +69,6 @@
     },
     python_requires=">=3.7",
     install_requires=[
-        "anyio>=3.0,<5.0",
         "requests>=2.21.0",
         "six>=1.11.0",
         "typing-extensions>=4.6.0",

sdk/core/azure-core/tests/async_tests/test_authentication_async.py

@@ -4,8 +4,9 @@
 # license information.
 # -------------------------------------------------------------------------
 import asyncio
+import sys
 import time
-from unittest.mock import Mock
+from unittest.mock import Mock, patch
 from requests import Response
 
 from azure.core.credentials import AccessToken
@@ -20,11 +21,12 @@
 )
 from azure.core.pipeline.transport import AsyncHttpTransport, HttpRequest
 import pytest
+import trio
 
-pytestmark = pytest.mark.asyncio
 from utils import HTTP_REQUESTS
 
 
+@pytest.mark.asyncio
 @pytest.mark.parametrize("http_request", HTTP_REQUESTS)
 async def test_bearer_policy_adds_header(http_request):
     """The bearer token policy should add a header containing a token from its credential"""
@@ -54,6 +56,7 @@ async def get_token(*_, **__):
     assert get_token_calls == 1
 
 
+@pytest.mark.asyncio
 @pytest.mark.parametrize("http_request", HTTP_REQUESTS)
 async def test_bearer_policy_send(http_request):
     """The bearer token policy should invoke the next policy's send method and return the result"""
@@ -72,6 +75,7 @@ async def verify_request(request):
     assert response is expected_response
 
 
+@pytest.mark.asyncio
 @pytest.mark.parametrize("http_request", HTTP_REQUESTS)
 async def test_bearer_policy_sync_send(http_request):
     """The bearer token policy should invoke the next policy's send method and return the result"""
@@ -90,6 +94,7 @@ async def verify_request(request):
     assert response is expected_response
 
 
+@pytest.mark.asyncio
 @pytest.mark.parametrize("http_request", HTTP_REQUESTS)
 async def test_bearer_policy_token_caching(http_request):
     good_for_one_hour = AccessToken("token", time.time() + 3600)
@@ -130,6 +135,7 @@ async def get_token(*_, **__):
     assert get_token_calls == 2  # token expired -> policy should call get_token
 
 
+@pytest.mark.asyncio
 @pytest.mark.parametrize("http_request", HTTP_REQUESTS)
 async def test_bearer_policy_optionally_enforces_https(http_request):
     """HTTPS enforcement should be controlled by a keyword argument, and enabled by default"""
@@ -158,6 +164,7 @@ async def assert_option_popped(request, **kwargs):
     await pipeline.run(http_request("GET", "https://secure"))
 
 
+@pytest.mark.asyncio
 @pytest.mark.parametrize("http_request", HTTP_REQUESTS)
 async def test_bearer_policy_preserves_enforce_https_opt_out(http_request):
     """The policy should use request context to preserve an opt out from https enforcement"""
@@ -175,6 +182,7 @@ def on_request(self, request):
     await pipeline.run(http_request("GET", "http://not.secure"), enforce_https=False)
 
 
+@pytest.mark.asyncio
 @pytest.mark.parametrize("http_request", HTTP_REQUESTS)
 async def test_bearer_policy_context_unmodified_by_default(http_request):
     """When no options for the policy accompany a request, the policy shouldn't add anything to the request context"""
@@ -192,6 +200,7 @@ def on_request(self, request):
     await pipeline.run(http_request("GET", "https://secure"))
 
 
+@pytest.mark.asyncio
 @pytest.mark.parametrize("http_request", HTTP_REQUESTS)
 async def test_bearer_policy_calls_sansio_methods(http_request):
     """AsyncBearerTokenCredentialPolicy should call SansIOHttpPolicy methods as does _SansIOAsyncHTTPPolicyRunner"""
@@ -440,3 +449,24 @@ async def get_token(self, *scopes, **kwargs):
 
     cred = TestTokenCredential()
     await cred.get_token("scope")
+
+
+@pytest.mark.asyncio
+async def test_async_token_credential_asyncio_lock():
+    auth_policy = AsyncBearerTokenCredentialPolicy(Mock(), "scope")
+    assert isinstance(auth_policy._lock, asyncio.Lock)
+
+
+@pytest.mark.trio
+async def test_async_token_credential_trio_lock():
+    auth_policy = AsyncBearerTokenCredentialPolicy(Mock(), "scope")
+    assert isinstance(auth_policy._lock, trio.Lock)
+
+
+def test_async_token_credential_sync():
+    """Verify that AsyncBearerTokenCredentialPolicy can be constructed in a synchronous context."""
+    auth_policy = AsyncBearerTokenCredentialPolicy(Mock(), "scope")
+    with patch.dict("sys.modules"):
+        # Ensure trio isn't in sys.modules (i.e. imported).
+        sys.modules.pop("trio", None)
+        AsyncBearerTokenCredentialPolicy(Mock(), "scope")

sdk/core/azure-core/tests/test_utils.py

@@ -2,8 +2,12 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 # ------------------------------------
+import sys
+from unittest.mock import patch
+
 import pytest
 from azure.core.utils import case_insensitive_dict
+from azure.core.utils._utils import get_running_async_lock
 
 
 @pytest.fixture()
@@ -108,3 +112,25 @@ def test_case_iter():
 
     for key in my_dict:
         assert key in keys
+
+
+@pytest.mark.asyncio
+async def test_get_running_async_module_asyncio():
+    import asyncio
+
+    assert isinstance(get_running_async_lock(), asyncio.Lock)
+
+
+@pytest.mark.trio
+async def test_get_running_async_module_trio():
+    import trio
+
+    assert isinstance(get_running_async_lock(), trio.Lock)
+
+
+def test_get_running_async_module_sync():
+    with patch.dict("sys.modules"):
+        # Ensure trio isn't in sys.modules (i.e. imported).
+        sys.modules.pop("trio", None)
+        with pytest.raises(RuntimeError):
+            get_running_async_lock()