More review feedback for token passing

Signed-off-by: Ryan Lettieri <ryanLettieri@microsoft.com>

Author: RyanLettieri

durabletask-azuremanaged/durabletask/azuremanaged/client.py

@@ -12,22 +12,13 @@ class DurableTaskSchedulerClient(TaskHubGrpcClient):
     def __init__(self, *,
                  host_address: str,
                  taskhub: str,
-                 secure_channel: Optional[bool] = True,
-                 metadata: Optional[list[tuple[str, str]]] = None,
-                 token_credential: Optional[TokenCredential] = None):
+                 token_credential: TokenCredential = None,
+                 secure_channel: Optional[bool] = True):
 
         if taskhub == None:
             raise ValueError("Taskhub value cannot be empty. Please provide a value for your taskhub")
 
-        # Ensure metadata is a list
-        metadata = metadata or []
-        self._metadata = metadata.copy()  # Use a copy to avoid modifying original
-
-        # Append DurableTask-specific metadata
-        self._metadata.append(("taskhub", taskhub))
-        self._metadata.append(("dts", "True"))
-        self._metadata.append(("token_credential", token_credential))
-        self._interceptors = [DTSDefaultClientInterceptorImpl(self._metadata)]
+        self._interceptors = [DTSDefaultClientInterceptorImpl(token_credential, taskhub)]
 
         # We pass in None for the metadata so we don't construct an additional interceptor in the parent class
         # Since the parent class doesn't use anything metadata for anything else, we can set it as None

durabletask-azuremanaged/durabletask/azuremanaged/durabletask_grpc_interceptor.py

@@ -3,28 +3,23 @@
 
 from durabletask.internal.grpc_interceptor import _ClientCallDetails, DefaultClientInterceptorImpl
 from durabletask.azuremanaged.internal.access_token_manager import AccessTokenManager
-
+from azure.core.credentials import TokenCredential
 import grpc
 
 class DTSDefaultClientInterceptorImpl (DefaultClientInterceptorImpl):
     """The class implements a UnaryUnaryClientInterceptor, UnaryStreamClientInterceptor,
     StreamUnaryClientInterceptor and StreamStreamClientInterceptor from grpc to add an 
     interceptor to add additional headers to all calls as needed."""
 
-    def __init__(self, metadata: list[tuple[str, str]]):
+    def __init__(self, token_credential: TokenCredential, taskhub_name: str):
+        metadata = [("taskhub", taskhub_name)]
         super().__init__(metadata)
 
-        self._token_credential = None
-        
-        # Check what authentication we are using
-        if metadata:
-            for key, value in metadata:
-                if key.lower() == "token_credential":
-                    self._token_credential = value
-
-        self._token_manager = AccessTokenManager(token_credential=self._token_credential)
-        token = self._token_manager.get_access_token()
-        self._metadata.append(("authorization", token))
+        if token_credential is not None:
+            self._token_credential = token_credential
+            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

@@ -13,14 +13,7 @@ def __init__(self, refresh_interval_seconds: int = 600, token_credential: TokenC
         self._refresh_interval_seconds = refresh_interval_seconds
         self._logger = shared.get_logger("token_manager")
 
-        # 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._logger.debug("Using Default Azure Credentials for authentication.")
+        self._credential = token_credential
 
         self._token = self._credential.get_token(self._scope)
         self.expiry_time = None

durabletask-azuremanaged/durabletask/azuremanaged/worker.py

@@ -12,22 +12,13 @@ class DurableTaskSchedulerWorker(TaskHubGrpcWorker):
     def __init__(self, *,
                  host_address: str,
                  taskhub: str,
-                 secure_channel: Optional[bool] = True,
-                 metadata: Optional[list[tuple[str, str]]] = None,
-                 token_credential: Optional[TokenCredential] = None):
+                 token_credential: TokenCredential = None,
+                 secure_channel: Optional[bool] = True):
 
         if taskhub == None:
             raise ValueError("Taskhub value cannot be empty. Please provide a value for your taskhub")
 
-        # Ensure metadata is a list
-        metadata = metadata or []
-        self._metadata = metadata.copy()  # Copy to prevent modifying input
-
-        # Append DurableTask-specific metadata
-        self._metadata.append(("taskhub", taskhub))
-        self._metadata.append(("dts", "True"))
-        self._metadata.append(("token_credential", token_credential))
-        interceptors = [DTSDefaultClientInterceptorImpl(self._metadata)]
+        interceptors = [DTSDefaultClientInterceptorImpl(token_credential, taskhub)]
 
         # We pass in None for the metadata so we don't construct an additional interceptor in the parent class
         # Since the parent class doesn't use anything metadata for anything else, we can set it as None

examples/dts/dts_activity_sequence.py

@@ -4,6 +4,7 @@
 from durabletask import task
 from durabletask.azuremanaged.worker import DurableTaskSchedulerWorker
 from durabletask.azuremanaged.client import DurableTaskSchedulerClient, OrchestrationStatus
+from azure.identity import DefaultAzureCredential
 
 def hello(ctx: task.ActivityContext, name: str) -> str:
     """Activity function that returns a greeting"""
@@ -45,15 +46,18 @@ def sequence(ctx: task.OrchestrationContext, _):
     print("If you are using bash, run the following: export ENDPOINT=\"<schedulerEndpoint>\"")
     exit()
 
+credential = DefaultAzureCredential()
 
 # configure and start the worker
-with DurableTaskSchedulerWorker(host_address=endpoint, secure_channel=True, taskhub=taskhub_name) as w:
+with DurableTaskSchedulerWorker(host_address=endpoint, secure_channel=True,
+                                taskhub=taskhub_name, token_credential=credential) as w:
     w.add_orchestrator(sequence)
     w.add_activity(hello)
     w.start()
 
     # Construct the client and run the orchestrations
-    c = DurableTaskSchedulerClient(host_address=endpoint, secure_channel=True, taskhub=taskhub_name)
+    c = DurableTaskSchedulerClient(host_address=endpoint, secure_channel=True,
+                                   taskhub=taskhub_name, token_credential=credential)
     instance_id = c.schedule_new_orchestration(sequence)
     state = c.wait_for_orchestration_completion(instance_id, timeout=60)
     if state and state.runtime_status == OrchestrationStatus.COMPLETED:

examples/dts/dts_fanout_fanin.py

@@ -8,6 +8,7 @@
 from durabletask import client, task
 from durabletask.azuremanaged.worker import DurableTaskSchedulerWorker
 from durabletask.azuremanaged.client import DurableTaskSchedulerClient
+from azure.identity import DefaultAzureCredential
 
 def get_work_items(ctx: task.ActivityContext, _) -> list[str]:
     """Activity function that returns a list of work items"""
@@ -71,15 +72,19 @@ def orchestrator(ctx: task.OrchestrationContext, _):
     print("If you are using bash, run the following: export ENDPOINT=\"<schedulerEndpoint>\"")
     exit()
 
+credential = DefaultAzureCredential()
+
 # configure and start the worker
-with DurableTaskSchedulerWorker(host_address=endpoint, secure_channel=True, taskhub=taskhub_name) as w:
+with DurableTaskSchedulerWorker(host_address=endpoint, secure_channel=True,
+                                taskhub=taskhub_name, token_credential=credential) as w:
     w.add_orchestrator(orchestrator)
     w.add_activity(process_work_item)
     w.add_activity(get_work_items)
     w.start()
 
     # create a client, start an orchestration, and wait for it to finish
-    c = DurableTaskSchedulerClient(host_address=endpoint, secure_channel=True, taskhub=taskhub_name)
+    c = DurableTaskSchedulerClient(host_address=endpoint, secure_channel=True,
+                                   taskhub=taskhub_name, token_credential=credential)
     instance_id = c.schedule_new_orchestration(orchestrator)
     state = c.wait_for_orchestration_completion(instance_id, timeout=30)
     if state and state.runtime_status == client.OrchestrationStatus.COMPLETED: