Passing token credential as an argument rather than 2 strings

Signed-off-by: Ryan Lettieri <[email protected]>

Author: RyanLettieri

durabletask-azuremanaged/durabletask/azuremanaged/client.py

@@ -5,7 +5,7 @@
 from durabletask.client import TaskHubGrpcClient, OrchestrationStatus
 from durabletask.azuremanaged.internal.access_token_manager import AccessTokenManager
 from durabletask.azuremanaged.durabletask_grpc_interceptor import DTSDefaultClientInterceptorImpl
-from azure.identity import DefaultAzureCredential
+from azure.core.credentials import TokenCredential
 
 # Client class used for Durable Task Scheduler (DTS)
 class DurableTaskSchedulerClient(TaskHubGrpcClient):
@@ -14,8 +14,7 @@ def __init__(self, *,
                  taskhub: str,
                  secure_channel: Optional[bool] = True,
                  metadata: Optional[list[tuple[str, str]]] = None,
-                 use_managed_identity: Optional[bool] = False,
-                 client_id: Optional[str] = None):
+                 token_credential: Optional[TokenCredential] = None):
 
         if taskhub == None:
             raise ValueError("Taskhub value cannot be empty. Please provide a value for your taskhub")
@@ -27,14 +26,7 @@ def __init__(self, *,
         # Append DurableTask-specific metadata
         self._metadata.append(("taskhub", taskhub))
         self._metadata.append(("dts", "True"))
-        self._metadata.append(("use_managed_identity", str(use_managed_identity)))
-        self._metadata.append(("client_id", str(client_id or "None")))
-
-        self._access_token_manager = AccessTokenManager(use_managed_identity=use_managed_identity,
-                                                        client_id=client_id)
-        token = self._access_token_manager.get_access_token()
-        self._metadata.append(("authorization", token))
-
+        self._metadata.append(("token_credential", token_credential))
         self._interceptors = [DTSDefaultClientInterceptorImpl(self._metadata)]
 
         # We pass in None for the metadata so we don't construct an additional interceptor in the parent class

durabletask-azuremanaged/durabletask/azuremanaged/durabletask_grpc_interceptor.py

@@ -13,20 +13,18 @@ class DTSDefaultClientInterceptorImpl (DefaultClientInterceptorImpl):
 
     def __init__(self, metadata: list[tuple[str, str]]):
         super().__init__(metadata)
-
-        use_managed_identity = False
-        client_id = None
+        
+        self._token_credential = None
 
         # Check what authentication we are using
         if metadata:
             for key, value in metadata:
-                if key.lower() == "use_managed_identity":
-                    self.use_managed_identity = value.strip().lower() == "true"  # Convert to boolean
-                elif key.lower() == "client_id":
-                    self.client_id = value
+                if key.lower() == "token_credential":
+                    self._token_credential = value
 
-        self._token_manager = AccessTokenManager(use_managed_identity=use_managed_identity,
-                                                 client_id=client_id)
+        self._token_manager = AccessTokenManager(token_credential=self._token_credential)
+        token = self._token_manager.get_access_token()
+        self._metadata.append(("authorization", token))
 
     def _intercept_call(
                 self, client_call_details: _ClientCallDetails) -> grpc.ClientCallDetails:

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

@@ -4,47 +4,43 @@
 from datetime import datetime, timedelta, timezone
 from typing import Optional
 import durabletask.internal.shared as shared
+from azure.core.credentials import TokenCredential
 
 # By default, when there's 10minutes left before the token expires, refresh the token
 class AccessTokenManager:
-    def __init__(self, refresh_interval_seconds: int = 600, use_managed_identity: bool = False, client_id: str = None):
-        self.scope = "https://durabletask.io/.default"
-        self.refresh_interval_seconds = refresh_interval_seconds
-        self._use_managed_identity = use_managed_identity
-        self._client_id = client_id
+    def __init__(self, refresh_interval_seconds: int = 600, token_credential: TokenCredential = None):
+        self._scope = "https://durabletask.io/.default"
+        self._refresh_interval_seconds = refresh_interval_seconds
         self._logger = shared.get_logger("token_manager")
 
-        # Choose the appropriate credential based on use_managed_identity
-        if self._use_managed_identity:
-            if not self._client_id:
-                self._logger.debug("Using System Assigned Managed Identity for authentication.")
-                self.credential = ManagedIdentityCredential()
-            else:
-                self._logger.debug("Using User Assigned Managed Identity for authentication.")
-                self.credential = ManagedIdentityCredential(client_id=self._client_id)
+        # Choose the appropriate credential. 
+        # Both TokenCredential and DefaultAzureCredential get_token methods return an AccessToken
+        if token_credential:
+            self._logger.debug("Using user provided token credentials.")
+            self._credential = token_credential
         else:
-            self.credential = DefaultAzureCredential()
+            self._credential = DefaultAzureCredential()
             self._logger.debug("Using Default Azure Credentials for authentication.")
 
-        self.token = None
+        self._token = self._credential.get_token(self._scope)
         self.expiry_time = None
 
     def get_access_token(self) -> str:
-        if self.token is None or self.is_token_expired():
+        if self._token is None or self.is_token_expired():
             self.refresh_token()
-        return self.token
+        return self._token
 
     # Checks if the token is expired, or if it will expire in the next "refresh_interval_seconds" seconds.
     # For example, if the token is created to have a lifespan of 2 hours, and the refresh buffer is set to 30 minutes,
     # We will grab a new token when there're 30minutes left on the lifespan of the 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))
+        return datetime.now(timezone.utc) >= (self.expiry_time - timedelta(seconds=self._refresh_interval_seconds))
 
     def refresh_token(self):
-        new_token = self.credential.get_token(self.scope)
-        self.token = f"Bearer {new_token.token}"
+        new_token = self._credential.get_token(self._scope)
+        self._token = f"Bearer {new_token.token}"
 
         # Convert UNIX timestamp to timezone-aware datetime
         self.expiry_time = datetime.fromtimestamp(new_token.expires_on, tz=timezone.utc)

durabletask-azuremanaged/durabletask/azuremanaged/worker.py

@@ -5,16 +5,16 @@
 from durabletask.worker import TaskHubGrpcWorker
 from durabletask.azuremanaged.internal.access_token_manager import AccessTokenManager
 from durabletask.azuremanaged.durabletask_grpc_interceptor import DTSDefaultClientInterceptorImpl
+from azure.core.credentials import TokenCredential
 
 # Worker class used for Durable Task Scheduler (DTS)
 class DurableTaskSchedulerWorker(TaskHubGrpcWorker):
     def __init__(self, *,
                  host_address: str,
                  taskhub: str,
-                 secure_channel: bool,
+                 secure_channel: Optional[bool] = True,
                  metadata: Optional[list[tuple[str, str]]] = None,
-                 use_managed_identity: Optional[bool] = False,
-                 client_id: Optional[str] = None):
+                 token_credential: Optional[TokenCredential] = None):
 
         if taskhub == None:
             raise ValueError("Taskhub value cannot be empty. Please provide a value for your taskhub")
@@ -24,15 +24,9 @@ def __init__(self, *,
         self._metadata = metadata.copy()  # Copy to prevent modifying input
 
         # Append DurableTask-specific metadata
-        self._metadata.append(("taskhub", taskhub or "default-taskhub"))
+        self._metadata.append(("taskhub", taskhub))
         self._metadata.append(("dts", "True"))
-        self._metadata.append(("use_managed_identity", str(use_managed_identity)))
-        self._metadata.append(("client_id", str(client_id or "None")))
-
-        self._access_token_manager = AccessTokenManager(use_managed_identity=use_managed_identity,
-                                                        client_id=client_id)
-        token = self._access_token_manager.get_access_token()
-        self._metadata.append(("authorization", token))
+        self._metadata.append(("token_credential", token_credential))
         interceptors = [DTSDefaultClientInterceptorImpl(self._metadata)]
 
         # We pass in None for the metadata so we don't construct an additional interceptor in the parent class

examples/dts/dts_activity_sequence.py

@@ -47,7 +47,7 @@ def sequence(ctx: task.OrchestrationContext, _):
 
 
 # configure and start the worker
-with DurableTaskSchedulerWorker(host_address=endpoint, secure_channel=True, use_managed_identity=False, client_id="", taskhub=taskhub_name) as w:
+with DurableTaskSchedulerWorker(host_address=endpoint, secure_channel=True, taskhub=taskhub_name) as w:
     w.add_orchestrator(sequence)
     w.add_activity(hello)
     w.start()