Async client for durabletask-azuremanged

Author: andystaples

durabletask-azuremanaged/CHANGELOG.md

@@ -5,6 +5,11 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## Unreleased
+
+- Added `AsyncDurableTaskSchedulerClient` for async/await usage with `grpc.aio`
+- Added `DTSAsyncDefaultClientInterceptorImpl` async gRPC interceptor for DTS authentication
+
 ## v1.3.0
 
 - Updates base dependency to durabletask v1.3.0

durabletask-azuremanaged/durabletask/azuremanaged/client.py

@@ -8,9 +8,10 @@
 from azure.core.credentials import TokenCredential
 
 from durabletask.azuremanaged.internal.durabletask_grpc_interceptor import (
+    DTSAsyncDefaultClientInterceptorImpl,
     DTSDefaultClientInterceptorImpl,
 )
-from durabletask.client import TaskHubGrpcClient
+from durabletask.client import AsyncTaskHubGrpcClient, TaskHubGrpcClient
 
 
 # Client class used for Durable Task Scheduler (DTS)
@@ -39,3 +40,65 @@ def __init__(self, *,
             log_formatter=log_formatter,
             interceptors=interceptors,
             default_version=default_version)
+
+
+# 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.
+        default_version (Optional[str], optional): Default version string for orchestrations.
+        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[TokenCredential],
+                 secure_channel: bool = True,
+                 default_version: Optional[str] = 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 = [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,
+            secure_channel=secure_channel,
+            metadata=None,
+            log_handler=log_handler,
+            log_formatter=log_formatter,
+            interceptors=interceptors,
+            default_version=default_version)

durabletask-azuremanaged/durabletask/azuremanaged/internal/durabletask_grpc_interceptor.py

@@ -9,7 +9,9 @@
 
 from durabletask.azuremanaged.internal.access_token_manager import AccessTokenManager
 from durabletask.internal.grpc_interceptor import (
+    DefaultAsyncClientInterceptorImpl,
     DefaultClientInterceptorImpl,
+    _AsyncClientCallDetails,
     _ClientCallDetails,
 )
 
@@ -52,3 +54,44 @@ def _intercept_call(
                         self._metadata[i] = ("authorization", f"Bearer {new_token.token}")  # Update the 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[TokenCredential], 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)
+
+        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}"))
+
+    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 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
+
+        return super()._intercept_call(client_call_details)

durabletask/internal/grpc_interceptor.py

@@ -22,7 +22,7 @@ class _ClientCallDetails(
 class _AsyncClientCallDetails(
         namedtuple(
             '_AsyncClientCallDetails',
-            ['method', 'timeout', 'metadata', 'credentials', 'wait_for_ready', 'compression']),
+            ['method', 'timeout', 'metadata', 'credentials', 'wait_for_ready']),
         grpc.aio.ClientCallDetails):
     """This is an implementation of the aio ClientCallDetails interface needed for async interceptors.
     This class takes six named values and inherits the ClientCallDetails from grpc.aio package.