Updating dts client to refresh token

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

Author: RyanLettieri

durabletask/client.py

@@ -97,6 +97,7 @@ def __init__(self, *,
                  log_handler: Optional[logging.Handler] = None,
                  log_formatter: Optional[logging.Formatter] = None,
                  secure_channel: bool = False):
+        self._metadata = metadata
         channel = shared.get_grpc_channel(host_address, metadata, secure_channel=secure_channel)
         self._stub = stubs.TaskHubSidecarServiceStub(channel)
         self._logger = shared.get_logger("client", log_handler, log_formatter)

examples/dts/dts_activity_sequence.py

@@ -1,10 +1,9 @@
-import os
-from azure.identity import DefaultAzureCredential
-
 """End-to-end sample that demonstrates how to configure an orchestrator
 that calls an activity function in a sequence and prints the outputs."""
+import os
 from durabletask import client, task
 from externalpackages.durabletaskscheduler.durabletask_scheduler_worker import DurableTaskSchedulerWorker
+from externalpackages.durabletaskscheduler.durabletask_scheduler_client import DurableTaskSchedulerClient
 from externalpackages.durabletaskscheduler.access_token_manager import AccessTokenManager
 
 def hello(ctx: task.ActivityContext, name: str) -> str:
@@ -52,18 +51,18 @@ def sequence(ctx: task.OrchestrationContext, _):
 arm_scope = "https://durabletask.io/.default"
 token_manager = AccessTokenManager(scope = arm_scope)
 
-metaData: list[tuple[str, str]] = [
+meta_data: list[tuple[str, str]] = [
     ("taskhub", taskhub_name)
 ]
 
 # configure and start the worker
-with DurableTaskSchedulerWorker(host_address=endpoint, metadata=metaData, secure_channel=True, access_token_manager=token_manager) as w:
+with DurableTaskSchedulerWorker(host_address=endpoint, metadata=meta_data, secure_channel=True, access_token_manager=token_manager) as w:
     w.add_orchestrator(sequence)
     w.add_activity(hello)
     w.start()
 
     # Construct the client and run the orchestrations
-    c = client.TaskHubGrpcClient(host_address=endpoint, metadata=metaData, secure_channel=True)
+    c = DurableTaskSchedulerClient(host_address=endpoint, metadata=meta_data, secure_channel=True, access_token_manager=token_manager)
     instance_id = c.schedule_new_orchestration(sequence)
     state = c.wait_for_orchestration_completion(instance_id, timeout=45)
     if state and state.runtime_status == client.OrchestrationStatus.COMPLETED:

examples/dts/dts_fanout_fanin.py

@@ -0,0 +1,99 @@
+"""End-to-end sample that demonstrates how to configure an orchestrator
+that a dynamic number activity functions in parallel, waits for them all
+to complete, and prints an aggregate summary of the outputs."""
+import random
+import time
+import os
+from durabletask import client, task
+from durabletask import client, task
+from externalpackages.durabletaskscheduler.durabletask_scheduler_worker import DurableTaskSchedulerWorker
+from externalpackages.durabletaskscheduler.durabletask_scheduler_client import DurableTaskSchedulerClient
+from externalpackages.durabletaskscheduler.access_token_manager import AccessTokenManager
+
+
+def get_work_items(ctx: task.ActivityContext, _) -> list[str]:
+    """Activity function that returns a list of work items"""
+    # return a random number of work items
+    count = random.randint(2, 10)
+    print(f'generating {count} work items...')
+    return [f'work item {i}' for i in range(count)]
+
+
+def process_work_item(ctx: task.ActivityContext, item: str) -> int:
+    """Activity function that returns a result for a given work item"""
+    print(f'processing work item: {item}')
+
+    # simulate some work that takes a variable amount of time
+    time.sleep(random.random() * 5)
+
+    # return a result for the given work item, which is also a random number in this case
+    return random.randint(0, 10)
+
+
+def orchestrator(ctx: task.OrchestrationContext, _):
+    """Orchestrator function that calls the 'get_work_items' and 'process_work_item'
+    activity functions in parallel, waits for them all to complete, and prints
+    an aggregate summary of the outputs"""
+
+    work_items: list[str] = yield ctx.call_activity(get_work_items)
+
+    # execute the work-items in parallel and wait for them all to return
+    tasks = [ctx.call_activity(process_work_item, input=item) for item in work_items]
+    results: list[int] = yield task.when_all(tasks)
+
+    # return an aggregate summary of the results
+    return {
+        'work_items': work_items,
+        'results': results,
+        'total': sum(results),
+    }
+
+
+# Read the environment variable
+taskhub_name = os.getenv("TASKHUB")
+
+# Check if the variable exists
+if taskhub_name:
+    print(f"The value of TASKHUB is: {taskhub_name}")
+else:
+    print("TASKHUB is not set. Please set the TASKHUB environment variable to the name of the taskhub you wish to use")
+    print("If you are using windows powershell, run the following: $env:TASKHUB=\"<taskhubname>\"")
+    print("If you are using bash, run the following: export TASKHUB=\"<taskhubname>\"")
+    exit()
+
+# Read the environment variable
+endpoint = os.getenv("ENDPOINT")
+
+# Check if the variable exists
+if endpoint:
+    print(f"The value of ENDPOINT is: {endpoint}")
+else:
+    print("ENDPOINT is not set. Please set the ENDPOINT environment variable to the endpoint of the scheduler")
+    print("If you are using windows powershell, run the following: $env:ENDPOINT=\"<schedulerEndpoint>\"")
+    print("If you are using bash, run the following: export ENDPOINT=\"<schedulerEndpoint>\"")
+    exit()
+
+# Define the scope for Azure Resource Manager (ARM)
+arm_scope = "https://durabletask.io/.default"
+token_manager = AccessTokenManager(scope = arm_scope)
+
+meta_data: list[tuple[str, str]] = [
+    ("taskhub", taskhub_name)
+]
+
+
+# configure and start the worker
+with DurableTaskSchedulerWorker(host_address=endpoint, metadata=meta_data, secure_channel=True, access_token_manager=token_manager) 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, metadata=meta_data, secure_channel=True, access_token_manager=token_manager)
+    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:
+        print(f'Orchestration completed! Result: {state.serialized_output}')
+    elif state:
+        print(f'Orchestration failed: {state.failure_details}')

externalpackages/durabletaskscheduler/durabletask_scheduler_client.py

@@ -1,6 +1,64 @@
-from durabletask import TaskHubGrpcClient
+from durabletask.client import TaskHubGrpcClient
+from externalpackages.durabletaskscheduler.access_token_manager import AccessTokenManager
 
 class DurableTaskSchedulerClient(TaskHubGrpcClient):
-    def __init__(self, *args, **kwargs):
+    def __init__(self, *args, access_token_manager: AccessTokenManager, **kwargs):
         # Initialize the base class
-        super().__init__(*args, **kwargs)
+        super().__init__(*args, **kwargs)
+        self._access_token_manager = access_token_manager
+        self.__update_metadata_with_token()
+
+    def __update_metadata_with_token(self):
+        """
+        Add or update the `authorization` key in the metadata with the current access token.
+        """
+        if self._access_token_manager is not None:
+            token = self._access_token_manager.get_access_token()
+            
+            # Check if "authorization" already exists in the metadata
+            updated = False
+            for i, (key, _) in enumerate(self._metadata):
+                if key == "authorization":
+                    self._metadata[i] = ("authorization", token)
+                    updated = True
+                    break
+            
+            # If not updated, add a new entry
+            if not updated:
+                self._metadata.append(("authorization", token))
+
+    def schedule_new_orchestration(self, *args, **kwargs) -> str:
+        self.__update_metadata_with_token()
+        return super().schedule_new_orchestration(*args, **kwargs)
+
+    def get_orchestration_state(self, *args, **kwargs):
+        self.__update_metadata_with_token()
+        super().get_orchestration_state(*args, **kwargs)
+
+    def wait_for_orchestration_start(self, *args, **kwargs):
+        self.__update_metadata_with_token()
+        super().wait_for_orchestration_start(*args, **kwargs)
+
+    def wait_for_orchestration_completion(self, *args, **kwargs):
+        self.__update_metadata_with_token()
+        super().wait_for_orchestration_completion(*args, **kwargs)
+
+    def raise_orchestration_event(self, *args, **kwargs):
+        self.__update_metadata_with_token()
+        super().raise_orchestration_event(*args, **kwargs)
+
+    def terminate_orchestration(self, *args, **kwargs):
+        self.__update_metadata_with_token()
+        super().terminate_orchestration(*args, **kwargs)
+
+    def suspend_orchestration(self, *args, **kwargs):
+        self.__update_metadata_with_token()
+        super().suspend_orchestration(*args, **kwargs)
+
+    def resume_orchestration(self, *args, **kwargs):
+        self.__update_metadata_with_token()
+        super().resume_orchestration(*args, **kwargs)
+
+    def purge_orchestration(self, *args, **kwargs):
+        self.__update_metadata_with_token()
+        super().purge_orchestration(*args, **kwargs)

externalpackages/durabletaskscheduler/durabletask_scheduler_worker.py

@@ -10,7 +10,7 @@
 from externalpackages.durabletaskscheduler.access_token_manager import AccessTokenManager
 
 class DurableTaskSchedulerWorker(TaskHubGrpcWorker):
-    def __init__(self, *args, access_token_manager: AccessTokenManager = None, **kwargs):
+    def __init__(self, *args, access_token_manager: AccessTokenManager, **kwargs):
         # Initialize the base class
         super().__init__(*args, **kwargs)
         self._access_token_manager = access_token_manager