durabletask.azuremanaged 0.0.0.dev69__py3-none-any.whl → 0.0.0.dev71__py3-none-any.whl

durabletask/azuremanaged/client.py

@@ -3,14 +3,24 @@
 
 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
 
 from durabletask.azuremanaged.internal.durabletask_grpc_interceptor import (
+    DTSAsyncDefaultClientInterceptorImpl,
     DTSDefaultClientInterceptorImpl,
 )
-from durabletask.client import TaskHubGrpcClient
+from durabletask.client import AsyncTaskHubGrpcClient, TaskHubGrpcClient
+from durabletask.grpc_options import (
+    GrpcChannelOptions,
+    GrpcClientResiliencyOptions,
+)
+import durabletask.internal.shared as shared
+from durabletask.payload.store import PayloadStore
 
 
 # Client class used for Durable Task Scheduler (DTS)
@@ -19,23 +29,113 @@ 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,
+                 resiliency_options: Optional[GrpcClientResiliencyOptions] = None,
+                 default_version: Optional[str] = None,
+                 payload_store: Optional[PayloadStore] = None,
+                 log_handler: Optional[logging.Handler] = None,
+                 log_formatter: Optional[logging.Formatter] = None):
+
+        if not taskhub:
+            raise ValueError("Taskhub value cannot be empty. Please provide a value for your 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=resolved_interceptors,
+            channel_options=channel_options,
+            resiliency_options=resiliency_options,
+            default_version=default_version,
+            payload_store=payload_store)
+
+
+# Async client class used for Durable Task Scheduler (DTS)
+class AsyncDurableTaskSchedulerClient(AsyncTaskHubGrpcClient):
+    """An async client implementation for Azure Durable Task Scheduler (DTS).
+
+    This class extends AsyncTaskHubGrpcClient to provide integration with Azure's
+    Durable Task Scheduler service using async gRPC. It handles authentication via
+    Azure credentials and configures the necessary gRPC interceptors for DTS
+    communication.
+
+    Args:
+        host_address (str): The gRPC endpoint address of the DTS service.
+        taskhub (str): The name of the task hub. Cannot be empty.
+        token_credential (Optional[TokenCredential]): Azure credential for authentication.
+            If None, anonymous authentication will be used.
+        secure_channel (bool, optional): Whether to use a secure gRPC channel (TLS).
+            Defaults to True.
+        resiliency_options (Optional[GrpcClientResiliencyOptions], optional): Client-side
+            gRPC resiliency settings forwarded to the base async client.
+        default_version (Optional[str], optional): Default version string for orchestrations.
+        payload_store (Optional[PayloadStore], optional): A payload store for
+            externalizing large payloads. If None, payloads are sent inline.
+        log_handler (Optional[logging.Handler], optional): Custom logging handler for client logs.
+        log_formatter (Optional[logging.Formatter], optional): Custom log formatter for client logs.
+
+    Raises:
+        ValueError: If taskhub is empty or None.
+
+    Example:
+        >>> from azure.identity.aio import DefaultAzureCredential
+        >>> from durabletask.azuremanaged import AsyncDurableTaskSchedulerClient
+        >>>
+        >>> credential = DefaultAzureCredential()
+        >>> async with AsyncDurableTaskSchedulerClient(
+        ...     host_address="my-dts-service.azure.com:443",
+        ...     taskhub="my-task-hub",
+        ...     token_credential=credential
+        ... ) as client:
+        ...     instance_id = await client.schedule_new_orchestration("my_orchestrator")
+    """
+
+    def __init__(self, *,
+                 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,
+                 resiliency_options: Optional[GrpcClientResiliencyOptions] = None,
                  default_version: Optional[str] = None,
+                 payload_store: Optional[PayloadStore] = None,
                  log_handler: Optional[logging.Handler] = None,
                  log_formatter: Optional[logging.Formatter] = None):
 
         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.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,
-            default_version=default_version)
+            interceptors=resolved_interceptors,
+            channel_options=channel_options,
+            resiliency_options=resiliency_options,
+            default_version=default_version,
+            payload_store=payload_store)

durabletask/azuremanaged/internal/access_token_manager.py

@@ -1,9 +1,11 @@
 # 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
+from azure.core.credentials_async import AsyncTokenCredential
 
 import durabletask.internal.shared as shared
 
@@ -19,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)
@@ -29,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.
@@ -47,3 +52,40 @@ class AccessTokenManager:
             # Convert UNIX timestamp to timezone-aware datetime
             self.expiry_time = datetime.fromtimestamp(self._token.expires_on, tz=timezone.utc)
             self._logger.debug(f"Token refreshed. Expires at: {self.expiry_time}")
+
+
+class AsyncAccessTokenManager:
+    """Async version of AccessTokenManager that uses AsyncTokenCredential.
+
+    This avoids blocking the event loop when acquiring or refreshing tokens."""
+
+    _token: Optional[AccessToken]
+
+    def __init__(self, token_credential: Optional[AsyncTokenCredential],
+                 refresh_interval_seconds: int = 600):
+        self._scope = "https://durabletask.io/.default"
+        self._refresh_interval_seconds = refresh_interval_seconds
+        self._logger = shared.get_logger("async_token_manager")
+
+        self._credential = token_credential
+        self._token = None
+        self.expiry_time = None
+
+    async def get_access_token(self) -> Optional[AccessToken]:
+        if self._token is None or self.is_token_expired():
+            await self.refresh_token()
+        return self._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))
+
+    async def refresh_token(self):
+        if self._credential is not None:
+            self._token = await self._credential.get_token(self._scope)
+
+            # Convert UNIX timestamp to timezone-aware datetime
+            self.expiry_time = datetime.fromtimestamp(self._token.expires_on, tz=timezone.utc)
+            self._logger.debug(f"Token refreshed. Expires at: {self.expiry_time}")

durabletask/azuremanaged/internal/durabletask_grpc_interceptor.py

@@ -6,10 +6,16 @@ from typing import Optional
 
 import grpc
 from azure.core.credentials import TokenCredential
+from azure.core.credentials_async import AsyncTokenCredential
 
-from durabletask.azuremanaged.internal.access_token_manager import AccessTokenManager
+from durabletask.azuremanaged.internal.access_token_manager import (
+    AccessTokenManager,
+    AsyncAccessTokenManager,
+)
 from durabletask.internal.grpc_interceptor import (
+    DefaultAsyncClientInterceptorImpl,
     DefaultClientInterceptorImpl,
+    _AsyncClientCallDetails,
     _ClientCallDetails,
 )
 
@@ -19,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')
@@ -29,26 +39,91 @@ 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
         if token_credential is not None:
             self._token_credential = token_credential
             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:
         """Internal intercept_call implementation which adds metadata to grpc metadata in the RPC
             call details."""
-        # Refresh the auth token if it is present and needed
-        if self._metadata is not None:
-            for i, (key, _) in enumerate(self._metadata):
-                if key.lower() == "authorization":  # Ensure case-insensitive comparison
-                    new_token = self._token_manager.get_access_token()  # Get the new token
-                    if new_token is not None:
-                        self._metadata[i] = ("authorization", f"Bearer {new_token.token}")  # Update the token
+        # Refresh the auth token if a credential was provided. The call to
+        # get_access_token() is generally cheap, checking the expiry time and returning
+        # the cached value without a network call when still valid.
+        if self._token_manager is not None:
+            access_token = self._token_manager.get_access_token()
+            if access_token is not None:
+                self._upsert_authorization_header(access_token.token)
 
         return super()._intercept_call(client_call_details)
+
+
+class DTSAsyncDefaultClientInterceptorImpl(DefaultAsyncClientInterceptorImpl):
+    """Async version of DTSDefaultClientInterceptorImpl for use with grpc.aio channels.
+
+    This class implements async gRPC interceptors to add DTS-specific headers
+    (task hub name, user agent, and authentication token) to all async calls."""
+
+    def __init__(self, token_credential: Optional[AsyncTokenCredential], taskhub_name: str):
+        try:
+            # Get the version of the azuremanaged package
+            sdk_version = version('durabletask-azuremanaged')
+        except Exception:
+            # Fallback if version cannot be determined
+            sdk_version = "unknown"
+        user_agent = f"durabletask-python/{sdk_version}"
+        self._metadata = [
+            ("taskhub", taskhub_name),
+            ("x-user-agent", user_agent)]
+        super().__init__(self._metadata)
+
+        # Token acquisition is deferred to the first _intercept_call invocation
+        # rather than happening in __init__, because get_token() on an
+        # AsyncTokenCredential is async and cannot be awaited in a constructor.
+        self._token_manager = None
+        if token_credential is not None:
+            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
+            call details."""
+        # Refresh the auth token if a credential was provided. The call to
+        # get_access_token() is generally cheap, checking the expiry time and returning
+        # the cached value without a network call when still valid.
+        if self._token_manager is not None:
+            access_token = await self._token_manager.get_access_token()
+            if access_token is not None:
+                self._upsert_authorization_header(access_token.token)
+
+        return await super()._intercept_call(client_call_details)

durabletask/azuremanaged/worker.py

@@ -2,13 +2,23 @@
 # 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,
+    GrpcWorkerResiliencyOptions,
+)
+import durabletask.internal.shared as shared
+from durabletask.payload.store import PayloadStore
 from durabletask.worker import ConcurrencyOptions, TaskHubGrpcWorker
 
 
@@ -27,9 +37,13 @@ class DurableTaskSchedulerWorker(TaskHubGrpcWorker):
             If None, anonymous authentication will be used.
         secure_channel (bool, optional): Whether to use a secure gRPC channel (TLS).
             Defaults to True.
+        resiliency_options (Optional[GrpcWorkerResiliencyOptions], optional): Worker-side
+            gRPC resiliency settings forwarded to the base worker.
         concurrency_options (Optional[ConcurrencyOptions], optional): Configuration
             for controlling worker concurrency limits. If None, default concurrency
             settings will be used.
+        payload_store (Optional[PayloadStore], optional): A payload store for
+            externalizing large payloads. If None, payloads are sent inline.
         log_handler (Optional[logging.Handler], optional): Custom logging handler for worker logs.
         log_formatter (Optional[logging.Formatter], optional): Custom log formatter for worker logs.
 
@@ -61,23 +75,41 @@ 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,
+                 resiliency_options: Optional[GrpcWorkerResiliencyOptions] = None,
                  concurrency_options: Optional[ConcurrencyOptions] = None,
+                 payload_store: Optional[PayloadStore] = None,
                  log_handler: Optional[logging.Handler] = None,
                  log_formatter: Optional[logging.Formatter] = None):
 
         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,
-            concurrency_options=concurrency_options)
+            interceptors=resolved_interceptors,
+            channel_options=channel_options,
+            resiliency_options=resiliency_options,
+            concurrency_options=concurrency_options,
+            # DTS natively supports long timers so chunking is unnecessary
+            maximum_timer_interval=None,
+            payload_store=payload_store
+        )

{durabletask_azuremanaged-0.0.0.dev69.dist-info → durabletask_azuremanaged-0.0.0.dev71.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: durabletask.azuremanaged
-Version: 0.0.0.dev69
+Version: 0.0.0.dev71
 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
@@ -12,3 +12,5 @@ Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 Requires-Dist: durabletask>=0.0.0dev1
 Requires-Dist: azure-identity>=1.19.0
+Provides-Extra: azure-blob-payloads
+Requires-Dist: durabletask[azure-blob-payloads]>=1.4.0; extra == "azure-blob-payloads"

durabletask_azuremanaged-0.0.0.dev71.dist-info/RECORD

@@ -0,0 +1,11 @@
+durabletask/azuremanaged/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+durabletask/azuremanaged/client.py,sha256=NBsXXQVwNkxl1bYnSGwxzsvuKXf2M7dZgJX8OQLF6_A,6680
+durabletask/azuremanaged/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+durabletask/azuremanaged/worker.py,sha256=jWxDSjMcUiAg6qZn8uFYb_tp8lWbqnDwlUYpn24z6bs,5243
+durabletask/azuremanaged/internal/access_token_manager.py,sha256=4czHfvPhZZX1MksQP02dG2ARPhZiYF1FUD57253XmVg,3901
+durabletask/azuremanaged/internal/durabletask_grpc_interceptor.py,sha256=XasSBmCJtnuAnJqmo_laR1O2t0PgLySwtX7qWilA468,5751
+durabletask/azuremanaged/internal/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+durabletask_azuremanaged-0.0.0.dev71.dist-info/METADATA,sha256=4z-TY2mj3n8MQyUSw_FuGbIyxfmQ51Jse7CtImSxhLM,792
+durabletask_azuremanaged-0.0.0.dev71.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+durabletask_azuremanaged-0.0.0.dev71.dist-info/top_level.txt,sha256=EBVyuKWnjOwq8bJI1Uvb9U3c4fzQxACWj9p83he6fik,12
+durabletask_azuremanaged-0.0.0.dev71.dist-info/RECORD,,

{durabletask_azuremanaged-0.0.0.dev69.dist-info → durabletask_azuremanaged-0.0.0.dev71.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (82.0.0)
+Generator: setuptools (82.0.1)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

durabletask_azuremanaged-0.0.0.dev69.dist-info/RECORD

@@ -1,11 +0,0 @@
-durabletask/azuremanaged/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-durabletask/azuremanaged/client.py,sha256=5ehYArdhV4EWSkllxkMwj2vk5hKYpywEqd1Rm86LoF4,1601
-durabletask/azuremanaged/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-durabletask/azuremanaged/worker.py,sha256=NRrymZKCvkspj3qzBbghuIMLV9rJ_QYf2EC4udlkmlM,3715
-durabletask/azuremanaged/internal/access_token_manager.py,sha256=Q8k9A2O8mNMyS-5jGRWoH-ZWFtzQ5Nt4anA_Bida5Kk,2167
-durabletask/azuremanaged/internal/durabletask_grpc_interceptor.py,sha256=MUbqlFLnppCVjUyuK4ISF8os1WHQ20GLdYto6P-WS2A,2527
-durabletask/azuremanaged/internal/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-durabletask_azuremanaged-0.0.0.dev69.dist-info/METADATA,sha256=kepuVCYiQ1CB60BL-gB07SOciixIW0K-g6LxvACpn78,669
-durabletask_azuremanaged-0.0.0.dev69.dist-info/WHEEL,sha256=YCfwYGOYMi5Jhw2fU4yNgwErybb2IX5PEwBKV4ZbdBo,91
-durabletask_azuremanaged-0.0.0.dev69.dist-info/top_level.txt,sha256=EBVyuKWnjOwq8bJI1Uvb9U3c4fzQxACWj9p83he6fik,12
-durabletask_azuremanaged-0.0.0.dev69.dist-info/RECORD,,