durabletask.azuremanaged 1.3.0.dev20__tar.gz → 1.3.0.dev22__tar.gz

{durabletask_azuremanaged-1.3.0.dev20 → durabletask_azuremanaged-1.3.0.dev22}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: durabletask.azuremanaged
-Version: 1.3.0.dev20
+Version: 1.3.0.dev22
 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.3.0.dev22/durabletask/azuremanaged/client.py

@@ -0,0 +1,105 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import logging
+
+from typing import Optional
+
+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 AsyncTaskHubGrpcClient, TaskHubGrpcClient
+
+
+# Client class used for Durable Task Scheduler (DTS)
+class DurableTaskSchedulerClient(TaskHubGrpcClient):
+    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 = [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,
+            secure_channel=secure_channel,
+            metadata=None,
+            log_handler=log_handler,
+            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[AsyncTokenCredential],
+                 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-1.3.0.dev20 → durabletask_azuremanaged-1.3.0.dev22}/durabletask/azuremanaged/internal/access_token_manager.py

@@ -4,6 +4,7 @@ from datetime import datetime, timedelta, timezone
 from typing import Optional
 
 from azure.core.credentials import AccessToken, TokenCredential
+from azure.core.credentials_async import AsyncTokenCredential
 
 import durabletask.internal.shared as shared
 
@@ -47,3 +48,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-1.3.0.dev22/durabletask/azuremanaged/internal/durabletask_grpc_interceptor.py

@@ -0,0 +1,121 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from importlib.metadata import version
+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,
+    AsyncAccessTokenManager,
+)
+from durabletask.internal.grpc_interceptor import (
+    DefaultAsyncClientInterceptorImpl,
+    DefaultClientInterceptorImpl,
+    _AsyncClientCallDetails,
+    _ClientCallDetails,
+)
+
+
+class DTSDefaultClientInterceptorImpl (DefaultClientInterceptorImpl):
+    """The class implements a UnaryUnaryClientInterceptor, UnaryStreamClientInterceptor,
+    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):
+        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)]  # 'user-agent' is a reserved header in grpc, so we use 'x-user-agent' instead
+        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}"))
+
+    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 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:
+                # 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}"))
+
+        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)
+
+    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:
+                # 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}"))
+
+        return await super()._intercept_call(client_call_details)

{durabletask_azuremanaged-1.3.0.dev20 → durabletask_azuremanaged-1.3.0.dev22}/durabletask.azuremanaged.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: durabletask.azuremanaged
-Version: 1.3.0.dev20
+Version: 1.3.0.dev22
 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.3.0.dev20 → durabletask_azuremanaged-1.3.0.dev22}/pyproject.toml

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

durabletask_azuremanaged-1.3.0.dev20/durabletask/azuremanaged/client.py

@@ -1,41 +0,0 @@
-# Copyright (c) Microsoft Corporation.
-# Licensed under the MIT License.
-
-import logging
-
-from typing import Optional
-
-from azure.core.credentials import TokenCredential
-
-from durabletask.azuremanaged.internal.durabletask_grpc_interceptor import (
-    DTSDefaultClientInterceptorImpl,
-)
-from durabletask.client import TaskHubGrpcClient
-
-
-# Client class used for Durable Task Scheduler (DTS)
-class DurableTaskSchedulerClient(TaskHubGrpcClient):
-    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 = [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,
-            secure_channel=secure_channel,
-            metadata=None,
-            log_handler=log_handler,
-            log_formatter=log_formatter,
-            interceptors=interceptors,
-            default_version=default_version)

durabletask_azuremanaged-1.3.0.dev20/durabletask/azuremanaged/internal/durabletask_grpc_interceptor.py

@@ -1,54 +0,0 @@
-# Copyright (c) Microsoft Corporation.
-# Licensed under the MIT License.
-
-from importlib.metadata import version
-from typing import Optional
-
-import grpc
-from azure.core.credentials import TokenCredential
-
-from durabletask.azuremanaged.internal.access_token_manager import AccessTokenManager
-from durabletask.internal.grpc_interceptor import (
-    DefaultClientInterceptorImpl,
-    _ClientCallDetails,
-)
-
-
-class DTSDefaultClientInterceptorImpl (DefaultClientInterceptorImpl):
-    """The class implements a UnaryUnaryClientInterceptor, UnaryStreamClientInterceptor,
-    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):
-        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)]  # 'user-agent' is a reserved header in grpc, so we use 'x-user-agent' instead
-        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: _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
-
-        return super()._intercept_call(client_call_details)