Adding accessTokenManager class for refreshing credential token

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

Author: RyanLettieri

durabletask/accessTokenManager.py

@@ -0,0 +1,29 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from azure.identity import DefaultAzureCredential
+from datetime import datetime, timedelta
+
+class AccessTokenManager:
+    def __init__(self, scope: str, refresh_buffer: int = 60):
+        self.scope = scope
+        self.refresh_buffer = refresh_buffer
+        self.credential = DefaultAzureCredential()
+        self.token = None
+        self.expiry_time = None
+
+    def get_access_token(self) -> str:
+        if self.token is None or self.is_token_expired():
+            self.refresh_token()
+        return self.token
+
+    def is_token_expired(self) -> bool:
+        if self.expiry_time is None:
+            return True
+        return datetime.utcnow() >= (self.expiry_time - timedelta(seconds=self.refresh_buffer))
+
+    def refresh_token(self):
+        new_token = self.credential.get_token(self.scope)
+        self.token = f"Bearer {new_token.token}"
+        self.expiry_time = datetime.utcnow() + timedelta(seconds=new_token.expires_on - int(datetime.utcnow().timestamp()))
+        print(f"Token refreshed. Expires at: {self.expiry_time}")

durabletask/worker.py

@@ -16,6 +16,7 @@
 import durabletask.internal.orchestrator_service_pb2 as pb
 import durabletask.internal.orchestrator_service_pb2_grpc as stubs
 import durabletask.internal.shared as shared
+from durabletask.accessTokenManager import AccessTokenManager
 from durabletask import task
 
 TInput = TypeVar('TInput')
@@ -88,14 +89,36 @@ def __init__(self, *,
                  metadata: Optional[list[tuple[str, str]]] = None,
                  log_handler=None,
                  log_formatter: Optional[logging.Formatter] = None,
-                 secure_channel: bool = False):
+                 secure_channel: bool = False,
+                 access_token_manager: AccessTokenManager = None):
         self._registry = _Registry()
         self._host_address = host_address if host_address else shared.get_default_host_address()
         self._metadata = metadata
         self._logger = shared.get_logger("worker", log_handler, log_formatter)
         self._shutdown = Event()
         self._is_running = False
         self._secure_channel = secure_channel
+        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 __enter__(self):
         return self
@@ -130,6 +153,7 @@ def run_loop():
             with concurrent.futures.ThreadPoolExecutor(max_workers=16) as executor:
                 while not self._shutdown.is_set():
                     try:
+                        self.__update_metadata_with_token()
                         # send a "Hello" message to the sidecar to ensure that it's listening
                         stub.Hello(empty_pb2.Empty())
 

examples/dts/README.md

@@ -26,6 +26,12 @@ export TASKHUB=<taskhubname>
 export ENDPOINT=<taskhubEndpoint>
 ```
 
+5. Since the samples rely on azure identity, ensure the package is installed and up-to-date
+
+```sh
+python3 -m pip install azure-identity
+```
+
 ## Running the examples
 
 With one of the sidecars running, you can simply execute any of the examples in this directory using `python3`:

examples/dts/dts_activity_sequence.py

@@ -4,7 +4,7 @@
 """End-to-end sample that demonstrates how to configure an orchestrator
 that calls an activity function in a sequence and prints the outputs."""
 from durabletask import client, task, worker
-
+from durabletask.accessTokenManager import AccessTokenManager
 
 def hello(ctx: task.ActivityContext, name: str) -> str:
     """Activity function that returns a greeting"""
@@ -47,19 +47,16 @@ def sequence(ctx: task.OrchestrationContext, _):
     exit()
 
 
-default_credential = DefaultAzureCredential()
 # Define the scope for Azure Resource Manager (ARM)
 arm_scope = "https://durabletask.io/.default"
-
-# Retrieve the access token. Note that this approach doesn't support token refresh and can't be used in production.
-access_token = "Bearer " + default_credential.get_token(arm_scope).token
+token_manager = AccessTokenManager(scope = arm_scope)
 
 metaData: list[tuple[str, str]] = [
-    ("taskhub", taskhub_name),
-    ("authorization", access_token)
+    ("taskhub", taskhub_name)
 ]
+
 # configure and start the worker
-with worker.TaskHubGrpcWorker(host_address=endpoint, metadata=metaData, secure_channel=True) as w:
+with worker.TaskHubGrpcWorker(host_address=endpoint, metadata=metaData, secure_channel=True, access_token_manager=token_manager) as w:
     w.add_orchestrator(sequence)
     w.add_activity(hello)
     w.start()