PR feedback

Author: andystaples

durabletask-azuremanaged/durabletask/azuremanaged/client.py

@@ -6,6 +6,7 @@
 from typing import Optional
 
 from azure.core.credentials import TokenCredential
+from azure.core.credentials_async import AsyncTokenCredential
 
 from durabletask.azuremanaged.internal.durabletask_grpc_interceptor import (
     DTSAsyncDefaultClientInterceptorImpl,
@@ -81,7 +82,7 @@ class AsyncDurableTaskSchedulerClient(AsyncTaskHubGrpcClient):
     def __init__(self, *,
                  host_address: str,
                  taskhub: str,
-                 token_credential: Optional[TokenCredential],
+                 token_credential: Optional[AsyncTokenCredential],
                  secure_channel: bool = True,
                  default_version: Optional[str] = None,
                  log_handler: Optional[logging.Handler] = None,

durabletask-azuremanaged/durabletask/azuremanaged/internal/access_token_manager.py

@@ -4,6 +4,7 @@
 from typing import Optional
 
 from azure.core.credentials import AccessToken, TokenCredential
+from azure.core.credentials_async import AsyncTokenCredential
 
 import durabletask.internal.shared as shared
 
@@ -47,3 +48,40 @@ def refresh_token(self):
             # Convert UNIX timestamp to timezone-aware datetime
             self.expiry_time = datetime.fromtimestamp(self._token.expires_on, tz=timezone.utc)
             self._logger.debug(f"Token refreshed. Expires at: {self.expiry_time}")
+
+
+class AsyncAccessTokenManager:
+    """Async version of AccessTokenManager that uses AsyncTokenCredential.
+
+    This avoids blocking the event loop when acquiring or refreshing tokens."""
+
+    _token: Optional[AccessToken]
+
+    def __init__(self, token_credential: Optional[AsyncTokenCredential],
+                 refresh_interval_seconds: int = 600):
+        self._scope = "https://durabletask.io/.default"
+        self._refresh_interval_seconds = refresh_interval_seconds
+        self._logger = shared.get_logger("async_token_manager")
+
+        self._credential = token_credential
+        self._token = None
+        self.expiry_time = None
+
+    async def get_access_token(self) -> Optional[AccessToken]:
+        if self._token is None or self.is_token_expired():
+            await self.refresh_token()
+        return self._token
+
+    def is_token_expired(self) -> bool:
+        if self.expiry_time is None:
+            return True
+        return datetime.now(timezone.utc) >= (
+            self.expiry_time - timedelta(seconds=self._refresh_interval_seconds))
+
+    async def refresh_token(self):
+        if self._credential is not None:
+            self._token = await self._credential.get_token(self._scope)
+
+            # Convert UNIX timestamp to timezone-aware datetime
+            self.expiry_time = datetime.fromtimestamp(self._token.expires_on, tz=timezone.utc)
+            self._logger.debug(f"Token refreshed. Expires at: {self.expiry_time}")

durabletask-azuremanaged/durabletask/azuremanaged/internal/durabletask_grpc_interceptor.py

@@ -6,8 +6,12 @@
 
 import grpc
 from azure.core.credentials import TokenCredential
+from azure.core.credentials_async import AsyncTokenCredential
 
-from durabletask.azuremanaged.internal.access_token_manager import AccessTokenManager
+from durabletask.azuremanaged.internal.access_token_manager import (
+    AccessTokenManager,
+    AsyncAccessTokenManager,
+)
 from durabletask.internal.grpc_interceptor import (
     DefaultAsyncClientInterceptorImpl,
     DefaultClientInterceptorImpl,
@@ -34,6 +38,7 @@ def __init__(self, token_credential: Optional[TokenCredential], taskhub_name: st
             ("x-user-agent", user_agent)]  # 'user-agent' is a reserved header in grpc, so we use 'x-user-agent' instead
         super().__init__(self._metadata)
 
+        self._token_manager = None
         if token_credential is not None:
             self._token_credential = token_credential
             self._token_manager = AccessTokenManager(token_credential=self._token_credential)
@@ -45,13 +50,21 @@ def _intercept_call(
             self, client_call_details: _ClientCallDetails) -> grpc.ClientCallDetails:
         """Internal intercept_call implementation which adds metadata to grpc metadata in the RPC
             call details."""
-        # Refresh the auth token if it is present and needed
-        if self._metadata is not None:
-            for i, (key, _) in enumerate(self._metadata):
-                if key.lower() == "authorization":  # Ensure case-insensitive comparison
-                    new_token = self._token_manager.get_access_token()  # Get the new token
-                    if new_token is not None:
-                        self._metadata[i] = ("authorization", f"Bearer {new_token.token}")  # Update the token
+        # Refresh the auth token if a credential was provided. The call to
+        # get_access_token() is generally cheap, checking the expiry time and returning
+        # the cached value without a network call when still valid.
+        if self._token_manager is not None:
+            access_token = self._token_manager.get_access_token()
+            if access_token is not None:
+                # Update the existing authorization header
+                found = False
+                for i, (key, _) in enumerate(self._metadata):
+                    if key.lower() == "authorization":
+                        self._metadata[i] = ("authorization", f"Bearer {access_token.token}")
+                        found = True
+                        break
+                if not found:
+                    self._metadata.append(("authorization", f"Bearer {access_token.token}"))
 
         return super()._intercept_call(client_call_details)
 
@@ -62,7 +75,7 @@ class DTSAsyncDefaultClientInterceptorImpl(DefaultAsyncClientInterceptorImpl):
     This class implements async gRPC interceptors to add DTS-specific headers
     (task hub name, user agent, and authentication token) to all async calls."""
 
-    def __init__(self, token_credential: Optional[TokenCredential], taskhub_name: str):
+    def __init__(self, token_credential: Optional[AsyncTokenCredential], taskhub_name: str):
         try:
             # Get the version of the azuremanaged package
             sdk_version = version('durabletask-azuremanaged')
@@ -75,23 +88,34 @@ def __init__(self, token_credential: Optional[TokenCredential], taskhub_name: st
             ("x-user-agent", user_agent)]
         super().__init__(self._metadata)
 
+        # Token acquisition is deferred to the first _intercept_call invocation
+        # rather than happening in __init__, because get_token() on an
+        # AsyncTokenCredential is async and cannot be awaited in a constructor.
+        self._token_manager = None
         if token_credential is not None:
             self._token_credential = token_credential
-            self._token_manager = AccessTokenManager(token_credential=self._token_credential)
-            access_token = self._token_manager.get_access_token()
-            if access_token is not None:
-                self._metadata.append(("authorization", f"Bearer {access_token.token}"))
+            self._token_manager = AsyncAccessTokenManager(token_credential=self._token_credential)
 
-    def _intercept_call(
+    async def _intercept_call(
             self, client_call_details: _AsyncClientCallDetails) -> grpc.aio.ClientCallDetails:
         """Internal intercept_call implementation which adds metadata to grpc metadata in the RPC
             call details."""
-        # Refresh the auth token if it is present and needed
-        if self._metadata is not None:
-            for i, (key, _) in enumerate(self._metadata):
-                if key.lower() == "authorization":  # Ensure case-insensitive comparison
-                    new_token = self._token_manager.get_access_token()  # Get the new token
-                    if new_token is not None:
-                        self._metadata[i] = ("authorization", f"Bearer {new_token.token}")  # Update the token
-
-        return super()._intercept_call(client_call_details)
+        # Refresh the auth token if a credential was provided. The call to
+        # get_access_token() is generally cheap, checking the expiry time and returning
+        # the cached value without a network call when still valid.
+        if self._token_manager is not None:
+            access_token = await self._token_manager.get_access_token()
+            if access_token is not None:
+                # Update the existing authorization header, or append one if this
+                # is the first successful token acquisition (token is lazily
+                # fetched on the first call since async constructors aren't possible).
+                found = False
+                for i, (key, _) in enumerate(self._metadata):
+                    if key.lower() == "authorization":
+                        self._metadata[i] = ("authorization", f"Bearer {access_token.token}")
+                        found = True
+                        break
+                if not found:
+                    self._metadata.append(("authorization", f"Bearer {access_token.token}"))
+
+        return await super()._intercept_call(client_call_details)

durabletask/client.py

@@ -404,6 +404,12 @@ async def close(self) -> None:
         """Closes the underlying gRPC channel."""
         await self._channel.close()
 
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self.close()
+
     async def schedule_new_orchestration(self, orchestrator: Union[task.Orchestrator[TInput, TOutput], str], *,
                                          input: Optional[TInput] = None,
                                          instance_id: Optional[str] = None,

durabletask/internal/grpc_interceptor.py

@@ -93,10 +93,11 @@ class DefaultAsyncClientInterceptorImpl(
     def __init__(self, metadata: list[tuple[str, str]]):
         self._metadata = metadata
 
-    def _intercept_call(
+    async def _intercept_call(
             self, client_call_details: grpc.aio.ClientCallDetails) -> grpc.aio.ClientCallDetails:
         """Internal intercept_call implementation which adds metadata to grpc metadata in the RPC
-            call details."""
+            call details. This method is async to allow subclasses to perform async operations
+            (e.g., refreshing auth tokens) during interception."""
         new_metadata = _apply_metadata(client_call_details, self._metadata)
         if new_metadata is client_call_details.metadata:
             return client_call_details
@@ -110,17 +111,17 @@ def _intercept_call(
         )
 
     async def intercept_unary_unary(self, continuation, client_call_details, request):
-        new_client_call_details = self._intercept_call(client_call_details)
+        new_client_call_details = await self._intercept_call(client_call_details)
         return await continuation(new_client_call_details, request)
 
     async def intercept_unary_stream(self, continuation, client_call_details, request):
-        new_client_call_details = self._intercept_call(client_call_details)
+        new_client_call_details = await self._intercept_call(client_call_details)
         return await continuation(new_client_call_details, request)
 
     async def intercept_stream_unary(self, continuation, client_call_details, request_iterator):
-        new_client_call_details = self._intercept_call(client_call_details)
+        new_client_call_details = await self._intercept_call(client_call_details)
         return await continuation(new_client_call_details, request_iterator)
 
     async def intercept_stream_stream(self, continuation, client_call_details, request_iterator):
-        new_client_call_details = self._intercept_call(client_call_details)
+        new_client_call_details = await self._intercept_call(client_call_details)
         return await continuation(new_client_call_details, request_iterator)