durabletask.azuremanaged 1.4.0.dev31__tar.gz → 1.4.0.dev32__tar.gz

{durabletask_azuremanaged-1.4.0.dev31 → durabletask_azuremanaged-1.4.0.dev32}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: durabletask.azuremanaged
-Version: 1.4.0.dev31
+Version: 1.4.0.dev32
 Summary: Durable Task Python SDK provider implementation for the Azure Durable Task Scheduler
 Project-URL: repository, https://github.com/microsoft/durabletask-python
 Project-URL: changelog, https://github.com/microsoft/durabletask-python/blob/main/CHANGELOG.md

{durabletask_azuremanaged-1.4.0.dev31 → durabletask_azuremanaged-1.4.0.dev32}/durabletask/azuremanaged/client.py

@@ -3,8 +3,10 @@
 
 import logging
 
-from typing import Optional
+from typing import Optional, Sequence
 
+import grpc
+import grpc.aio
 from azure.core.credentials import TokenCredential
 from azure.core.credentials_async import AsyncTokenCredential
 
@@ -13,6 +15,8 @@ from durabletask.azuremanaged.internal.durabletask_grpc_interceptor import (
     DTSDefaultClientInterceptorImpl,
 )
 from durabletask.client import AsyncTaskHubGrpcClient, TaskHubGrpcClient
+from durabletask.grpc_options import GrpcChannelOptions
+import durabletask.internal.shared as shared
 from durabletask.payload.store import PayloadStore
 
 
@@ -22,7 +26,10 @@ class DurableTaskSchedulerClient(TaskHubGrpcClient):
                  host_address: str,
                  taskhub: str,
                  token_credential: Optional[TokenCredential],
+                 channel: Optional[grpc.Channel] = None,
                  secure_channel: bool = True,
+                 interceptors: Optional[Sequence[shared.ClientInterceptor]] = None,
+                 channel_options: Optional[GrpcChannelOptions] = None,
                  default_version: Optional[str] = None,
                  payload_store: Optional[PayloadStore] = None,
                  log_handler: Optional[logging.Handler] = None,
@@ -31,17 +38,22 @@ class DurableTaskSchedulerClient(TaskHubGrpcClient):
         if not taskhub:
             raise ValueError("Taskhub value cannot be empty. Please provide a value for your taskhub")
 
-        interceptors = [DTSDefaultClientInterceptorImpl(token_credential, taskhub)]
+        resolved_interceptors: list[shared.ClientInterceptor] = (
+            list(interceptors) if interceptors is not None else []
+        )
+        resolved_interceptors.append(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
         super().__init__(
             host_address=host_address,
+            channel=channel,
             secure_channel=secure_channel,
             metadata=None,
             log_handler=log_handler,
             log_formatter=log_formatter,
-            interceptors=interceptors,
+            interceptors=resolved_interceptors,
+            channel_options=channel_options,
             default_version=default_version,
             payload_store=payload_store)
 
@@ -88,7 +100,10 @@ class AsyncDurableTaskSchedulerClient(AsyncTaskHubGrpcClient):
                  host_address: str,
                  taskhub: str,
                  token_credential: Optional[AsyncTokenCredential],
+                 channel: Optional[grpc.aio.Channel] = None,
                  secure_channel: bool = True,
+                 interceptors: Optional[Sequence[shared.AsyncClientInterceptor]] = None,
+                 channel_options: Optional[GrpcChannelOptions] = None,
                  default_version: Optional[str] = None,
                  payload_store: Optional[PayloadStore] = None,
                  log_handler: Optional[logging.Handler] = None,
@@ -97,16 +112,21 @@ class AsyncDurableTaskSchedulerClient(AsyncTaskHubGrpcClient):
         if not taskhub:
             raise ValueError("Taskhub value cannot be empty. Please provide a value for your taskhub")
 
-        interceptors = [DTSAsyncDefaultClientInterceptorImpl(token_credential, taskhub)]
+        resolved_interceptors: list[shared.AsyncClientInterceptor] = (
+            list(interceptors) if interceptors is not None else []
+        )
+        resolved_interceptors.append(DTSAsyncDefaultClientInterceptorImpl(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
         super().__init__(
             host_address=host_address,
+            channel=channel,
             secure_channel=secure_channel,
             metadata=None,
             log_handler=log_handler,
             log_formatter=log_formatter,
-            interceptors=interceptors,
+            interceptors=resolved_interceptors,
+            channel_options=channel_options,
             default_version=default_version,
             payload_store=payload_store)

{durabletask_azuremanaged-1.4.0.dev31 → durabletask_azuremanaged-1.4.0.dev32}/durabletask/azuremanaged/internal/access_token_manager.py

@@ -1,6 +1,7 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 from datetime import datetime, timedelta, timezone
+from threading import Lock
 from typing import Optional
 
 from azure.core.credentials import AccessToken, TokenCredential
@@ -20,6 +21,7 @@ class AccessTokenManager:
         self._logger = shared.get_logger("token_manager")
 
         self._credential = token_credential
+        self._refresh_lock = Lock()
 
         if self._credential is not None:
             self._token = self._credential.get_token(self._scope)
@@ -30,7 +32,9 @@ class AccessTokenManager:
 
     def get_access_token(self) -> Optional[AccessToken]:
         if self._token is None or self.is_token_expired():
-            self.refresh_token()
+            with self._refresh_lock:
+                if self._token is None or self.is_token_expired():
+                    self.refresh_token()
         return self._token
 
     # Checks if the token is expired, or if it will expire in the next "refresh_interval_seconds" seconds.

{durabletask_azuremanaged-1.4.0.dev31 → durabletask_azuremanaged-1.4.0.dev32}/durabletask/azuremanaged/internal/durabletask_grpc_interceptor.py

@@ -25,7 +25,11 @@ class DTSDefaultClientInterceptorImpl (DefaultClientInterceptorImpl):
     StreamUnaryClientInterceptor and StreamStreamClientInterceptor from grpc to add an
     interceptor to add additional headers to all calls as needed."""
 
-    def __init__(self, token_credential: Optional[TokenCredential], taskhub_name: str):
+    def __init__(
+            self,
+            token_credential: Optional[TokenCredential],
+            taskhub_name: str,
+            worker_id: Optional[str] = None):
         try:
             # Get the version of the azuremanaged package
             sdk_version = version('durabletask-azuremanaged')
@@ -35,7 +39,9 @@ class DTSDefaultClientInterceptorImpl (DefaultClientInterceptorImpl):
         user_agent = f"durabletask-python/{sdk_version}"
         self._metadata = [
             ("taskhub", taskhub_name),
-            ("x-user-agent", user_agent)]  # 'user-agent' is a reserved header in grpc, so we use 'x-user-agent' instead
+            ("x-user-agent", user_agent)]  # 'user-agent' is a reserved header; use 'x-user-agent'
+        if worker_id is not None:
+            self._metadata.append(("workerid", worker_id))
         super().__init__(self._metadata)
 
         self._token_manager = None
@@ -44,7 +50,17 @@ class DTSDefaultClientInterceptorImpl (DefaultClientInterceptorImpl):
             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._upsert_authorization_header(access_token.token)
+
+    def _upsert_authorization_header(self, token: str) -> None:
+        found = False
+        for i, (key, _) in enumerate(self._metadata):
+            if key.lower() == "authorization":
+                self._metadata[i] = ("authorization", f"Bearer {token}")
+                found = True
+                break
+        if not found:
+            self._metadata.append(("authorization", f"Bearer {token}"))
 
     def _intercept_call(
             self, client_call_details: _ClientCallDetails) -> grpc.ClientCallDetails:
@@ -56,15 +72,7 @@ class DTSDefaultClientInterceptorImpl (DefaultClientInterceptorImpl):
         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}"))
+                self._upsert_authorization_header(access_token.token)
 
         return super()._intercept_call(client_call_details)
 
@@ -96,6 +104,16 @@ class DTSAsyncDefaultClientInterceptorImpl(DefaultAsyncClientInterceptorImpl):
             self._token_credential = token_credential
             self._token_manager = AsyncAccessTokenManager(token_credential=self._token_credential)
 
+    def _upsert_authorization_header(self, token: str) -> None:
+        found = False
+        for i, (key, _) in enumerate(self._metadata):
+            if key.lower() == "authorization":
+                self._metadata[i] = ("authorization", f"Bearer {token}")
+                found = True
+                break
+        if not found:
+            self._metadata.append(("authorization", f"Bearer {token}"))
+
     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
@@ -106,16 +124,6 @@ class DTSAsyncDefaultClientInterceptorImpl(DefaultAsyncClientInterceptorImpl):
         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}"))
+                self._upsert_authorization_header(access_token.token)
 
         return await super()._intercept_call(client_call_details)

{durabletask_azuremanaged-1.4.0.dev31 → durabletask_azuremanaged-1.4.0.dev32}/durabletask/azuremanaged/worker.py

@@ -2,13 +2,19 @@
 # Licensed under the MIT License.
 
 import logging
+import os
+import socket
+import uuid
 
-from typing import Optional
+from typing import Optional, Sequence
 
+import grpc
 from azure.core.credentials import TokenCredential
 
 from durabletask.azuremanaged.internal.durabletask_grpc_interceptor import \
     DTSDefaultClientInterceptorImpl
+from durabletask.grpc_options import GrpcChannelOptions
+import durabletask.internal.shared as shared
 from durabletask.payload.store import PayloadStore
 from durabletask.worker import ConcurrencyOptions, TaskHubGrpcWorker
 
@@ -64,7 +70,10 @@ class DurableTaskSchedulerWorker(TaskHubGrpcWorker):
                  host_address: str,
                  taskhub: str,
                  token_credential: Optional[TokenCredential],
+                 channel: Optional[grpc.Channel] = None,
                  secure_channel: bool = True,
+                 interceptors: Optional[Sequence[shared.ClientInterceptor]] = None,
+                 channel_options: Optional[GrpcChannelOptions] = None,
                  concurrency_options: Optional[ConcurrencyOptions] = None,
                  payload_store: Optional[PayloadStore] = None,
                  log_handler: Optional[logging.Handler] = None,
@@ -73,17 +82,25 @@ class DurableTaskSchedulerWorker(TaskHubGrpcWorker):
         if not taskhub:
             raise ValueError("The taskhub value cannot be empty.")
 
-        interceptors = [DTSDefaultClientInterceptorImpl(token_credential, taskhub)]
+        worker_id = f"{socket.gethostname()}:{os.getpid()}:{uuid.uuid4()}"
+        resolved_interceptors: list[shared.ClientInterceptor] = (
+            list(interceptors) if interceptors is not None else []
+        )
+        resolved_interceptors.append(
+            DTSDefaultClientInterceptorImpl(token_credential, taskhub, worker_id=worker_id)
+        )
 
         # 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
         super().__init__(
             host_address=host_address,
+            channel=channel,
             secure_channel=secure_channel,
             metadata=None,
             log_handler=log_handler,
             log_formatter=log_formatter,
-            interceptors=interceptors,
+            interceptors=resolved_interceptors,
+            channel_options=channel_options,
             concurrency_options=concurrency_options,
             # DTS natively supports long timers so chunking is unnecessary
             maximum_timer_interval=None,

{durabletask_azuremanaged-1.4.0.dev31 → durabletask_azuremanaged-1.4.0.dev32}/durabletask.azuremanaged.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: durabletask.azuremanaged
-Version: 1.4.0.dev31
+Version: 1.4.0.dev32
 Summary: Durable Task Python SDK provider implementation for the Azure Durable Task Scheduler
 Project-URL: repository, https://github.com/microsoft/durabletask-python
 Project-URL: changelog, https://github.com/microsoft/durabletask-python/blob/main/CHANGELOG.md

{durabletask_azuremanaged-1.4.0.dev31 → durabletask_azuremanaged-1.4.0.dev32}/pyproject.toml

@@ -9,7 +9,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "durabletask.azuremanaged"
-version = "1.4.0.dev31"
+version = "1.4.0.dev32"
 description = "Durable Task Python SDK provider implementation for the Azure Durable Task Scheduler"
 keywords = [
     "durable",