durabletask 1.3.0.dev21__tar.gz → 1.3.0.dev22__tar.gz

{durabletask-1.3.0.dev21 → durabletask-1.3.0.dev22}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: durabletask
-Version: 1.3.0.dev21
+Version: 1.3.0.dev22
 Summary: A Durable Task Client SDK for Python
 License:     MIT License
         

{durabletask-1.3.0.dev21 → durabletask-1.3.0.dev22}/durabletask/client.py

@@ -2,13 +2,13 @@
 # Licensed under the MIT License.
 
 import logging
-import uuid
 from dataclasses import dataclass
-from datetime import datetime, timezone
+from datetime import datetime
 from enum import Enum
 from typing import Any, List, Optional, Sequence, TypeVar, Union
 
 import grpc
+import grpc.aio
 
 from durabletask.entities import EntityInstanceId
 from durabletask.entities.entity_metadata import EntityMetadata
@@ -17,7 +17,19 @@ 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 import task
-from durabletask.internal.grpc_interceptor import DefaultClientInterceptorImpl
+from durabletask.internal.client_helpers import (
+    build_query_entities_req,
+    build_query_instances_req,
+    build_purge_by_filter_req,
+    build_raise_event_req,
+    build_schedule_new_orchestration_req,
+    build_signal_entity_req,
+    build_terminate_req,
+    check_continuation_token,
+    log_completion_state,
+    prepare_async_interceptors,
+    prepare_sync_interceptors,
+)
 
 TInput = TypeVar('TInput')
 TOutput = TypeVar('TOutput')
@@ -140,26 +152,22 @@ class TaskHubGrpcClient:
                  interceptors: Optional[Sequence[shared.ClientInterceptor]] = None,
                  default_version: Optional[str] = None):
 
-        # If the caller provided metadata, we need to create a new interceptor for it and
-        # add it to the list of interceptors.
-        if interceptors is not None:
-            interceptors = list(interceptors)
-            if metadata is not None:
-                interceptors.append(DefaultClientInterceptorImpl(metadata))
-        elif metadata is not None:
-            interceptors = [DefaultClientInterceptorImpl(metadata)]
-        else:
-            interceptors = None
+        interceptors = prepare_sync_interceptors(metadata, interceptors)
 
         channel = shared.get_grpc_channel(
             host_address=host_address,
             secure_channel=secure_channel,
             interceptors=interceptors
         )
+        self._channel = channel
         self._stub = stubs.TaskHubSidecarServiceStub(channel)
         self._logger = shared.get_logger("client", log_handler, log_formatter)
         self.default_version = default_version
 
+    def close(self) -> None:
+        """Closes the underlying gRPC channel."""
+        self._channel.close()
+
     def schedule_new_orchestration(self, orchestrator: Union[task.Orchestrator[TInput, TOutput], str], *,
                                    input: Optional[TInput] = None,
                                    instance_id: Optional[str] = None,
@@ -168,19 +176,12 @@ class TaskHubGrpcClient:
                                    tags: Optional[dict[str, str]] = None,
                                    version: Optional[str] = None) -> str:
 
-        name = orchestrator if isinstance(orchestrator, str) else task.get_name(orchestrator)
-
-        req = pb.CreateInstanceRequest(
-            name=name,
-            instanceId=instance_id if instance_id else uuid.uuid4().hex,
-            input=helpers.get_string_value(shared.to_json(input) if input is not None else None),
-            scheduledStartTimestamp=helpers.new_timestamp(start_at) if start_at else None,
-            version=helpers.get_string_value(version if version else self.default_version),
-            orchestrationIdReusePolicy=reuse_id_policy,
-            tags=tags
-        )
+        req = build_schedule_new_orchestration_req(
+            orchestrator, input=input, instance_id=instance_id, start_at=start_at,
+            reuse_id_policy=reuse_id_policy, tags=tags,
+            version=version if version else self.default_version)
 
-        self._logger.info(f"Starting new '{name}' instance with ID = '{req.instanceId}'.")
+        self._logger.info(f"Starting new '{req.name}' instance with ID = '{req.instanceId}'.")
         res: pb.CreateInstanceResponse = self._stub.StartInstance(req)
         return res.instanceId
 
@@ -201,24 +202,10 @@ class TaskHubGrpcClient:
         states = []
 
         while True:
-            req = pb.QueryInstancesRequest(
-                query=pb.InstanceQuery(
-                    runtimeStatus=[status.value for status in orchestration_query.runtime_status] if orchestration_query.runtime_status else None,
-                    createdTimeFrom=helpers.new_timestamp(orchestration_query.created_time_from) if orchestration_query.created_time_from else None,
-                    createdTimeTo=helpers.new_timestamp(orchestration_query.created_time_to) if orchestration_query.created_time_to else None,
-                    maxInstanceCount=orchestration_query.max_instance_count,
-                    fetchInputsAndOutputs=orchestration_query.fetch_inputs_and_outputs,
-                    continuationToken=_continuation_token
-                )
-            )
+            req = build_query_instances_req(orchestration_query, _continuation_token)
             resp: pb.QueryInstancesResponse = self._stub.QueryInstances(req)
             states += [parse_orchestration_state(res) for res in resp.orchestrationState]
-            # Check the value for continuationToken - none or "0" indicates that there are no more results.
-            if resp.continuationToken and resp.continuationToken.value and resp.continuationToken.value != "0":
-                self._logger.info(f"Received continuation token with value {resp.continuationToken.value}, fetching next list of instances...")
-                if _continuation_token and _continuation_token.value and _continuation_token.value == resp.continuationToken.value:
-                    self._logger.warning(f"Received the same continuation token value {resp.continuationToken.value} again, stopping to avoid infinite loop.")
-                    break
+            if check_continuation_token(resp.continuationToken, _continuation_token, self._logger):
                 _continuation_token = resp.continuationToken
             else:
                 break
@@ -248,53 +235,35 @@ class TaskHubGrpcClient:
             self._logger.info(f"Waiting {timeout}s for instance '{instance_id}' to complete.")
             res: pb.GetInstanceResponse = self._stub.WaitForInstanceCompletion(req, timeout=timeout)
             state = new_orchestration_state(req.instanceId, res)
-            if not state:
-                return None
-
-            if state.runtime_status == OrchestrationStatus.FAILED and state.failure_details is not None:
-                details = state.failure_details
-                self._logger.info(f"Instance '{instance_id}' failed: [{details.error_type}] {details.message}")
-            elif state.runtime_status == OrchestrationStatus.TERMINATED:
-                self._logger.info(f"Instance '{instance_id}' was terminated.")
-            elif state.runtime_status == OrchestrationStatus.COMPLETED:
-                self._logger.info(f"Instance '{instance_id}' completed.")
-
+            log_completion_state(self._logger, instance_id, state)
             return state
         except grpc.RpcError as rpc_error:
             if rpc_error.code() == grpc.StatusCode.DEADLINE_EXCEEDED:  # type: ignore
-                # Replace gRPC error with the built-in TimeoutError
                 raise TimeoutError("Timed-out waiting for the orchestration to complete")
             else:
                 raise
 
     def raise_orchestration_event(self, instance_id: str, event_name: str, *,
-                                  data: Optional[Any] = None):
-        req = pb.RaiseEventRequest(
-            instanceId=instance_id,
-            name=event_name,
-            input=helpers.get_string_value(shared.to_json(data) if data is not None else None)
-        )
+                                  data: Optional[Any] = None) -> None:
+        req = build_raise_event_req(instance_id, event_name, data)
 
         self._logger.info(f"Raising event '{event_name}' for instance '{instance_id}'.")
         self._stub.RaiseEvent(req)
 
     def terminate_orchestration(self, instance_id: str, *,
                                 output: Optional[Any] = None,
-                                recursive: bool = True):
-        req = pb.TerminateRequest(
-            instanceId=instance_id,
-            output=helpers.get_string_value(shared.to_json(output) if output is not None else None),
-            recursive=recursive)
+                                recursive: bool = True) -> None:
+        req = build_terminate_req(instance_id, output, recursive)
 
         self._logger.info(f"Terminating instance '{instance_id}'.")
         self._stub.TerminateInstance(req)
 
-    def suspend_orchestration(self, instance_id: str):
+    def suspend_orchestration(self, instance_id: str) -> None:
         req = pb.SuspendRequest(instanceId=instance_id)
         self._logger.info(f"Suspending instance '{instance_id}'.")
         self._stub.SuspendInstance(req)
 
-    def resume_orchestration(self, instance_id: str):
+    def resume_orchestration(self, instance_id: str) -> None:
         req = pb.ResumeRequest(instanceId=instance_id)
         self._logger.info(f"Resuming instance '{instance_id}'.")
         self._stub.ResumeInstance(req)
@@ -335,29 +304,15 @@ class TaskHubGrpcClient:
                           f"created_time_to={created_time_to}, "
                           f"runtime_status={[str(status) for status in runtime_status] if runtime_status else None}, "
                           f"recursive={recursive}")
-        resp: pb.PurgeInstancesResponse = self._stub.PurgeInstances(pb.PurgeInstancesRequest(
-            purgeInstanceFilter=pb.PurgeInstanceFilter(
-                createdTimeFrom=helpers.new_timestamp(created_time_from) if created_time_from else None,
-                createdTimeTo=helpers.new_timestamp(created_time_to) if created_time_to else None,
-                runtimeStatus=[status.value for status in runtime_status] if runtime_status else None
-            ),
-            recursive=recursive
-        ))
+        req = build_purge_by_filter_req(created_time_from, created_time_to, runtime_status, recursive)
+        resp: pb.PurgeInstancesResponse = self._stub.PurgeInstances(req)
         return PurgeInstancesResult(resp.deletedInstanceCount, resp.isComplete.value)
 
     def signal_entity(self,
                       entity_instance_id: EntityInstanceId,
                       operation_name: str,
                       input: Optional[Any] = None) -> None:
-        req = pb.SignalEntityRequest(
-            instanceId=str(entity_instance_id),
-            name=operation_name,
-            input=helpers.get_string_value(shared.to_json(input) if input is not None else None),
-            requestId=str(uuid.uuid4()),
-            scheduledTime=None,
-            parentTraceContext=None,
-            requestTime=helpers.new_timestamp(datetime.now(timezone.utc))
-        )
+        req = build_signal_entity_req(entity_instance_id, operation_name, input)
         self._logger.info(f"Signaling entity '{entity_instance_id}' operation '{operation_name}'.")
         self._stub.SignalEntity(req, None)  # TODO: Cancellation timeout?
 
@@ -384,24 +339,10 @@ class TaskHubGrpcClient:
         entities = []
 
         while True:
-            query_request = pb.QueryEntitiesRequest(
-                query=pb.EntityQuery(
-                    instanceIdStartsWith=helpers.get_string_value(entity_query.instance_id_starts_with),
-                    lastModifiedFrom=helpers.new_timestamp(entity_query.last_modified_from) if entity_query.last_modified_from else None,
-                    lastModifiedTo=helpers.new_timestamp(entity_query.last_modified_to) if entity_query.last_modified_to else None,
-                    includeState=entity_query.include_state,
-                    includeTransient=entity_query.include_transient,
-                    pageSize=helpers.get_int_value(entity_query.page_size),
-                    continuationToken=_continuation_token
-                )
-            )
+            query_request = build_query_entities_req(entity_query, _continuation_token)
             resp: pb.QueryEntitiesResponse = self._stub.QueryEntities(query_request)
             entities += [EntityMetadata.from_entity_metadata(entity, query_request.query.includeState) for entity in resp.entities]
-            if resp.continuationToken and resp.continuationToken.value and resp.continuationToken.value != "0":
-                self._logger.info(f"Received continuation token with value {resp.continuationToken.value}, fetching next page of entities...")
-                if _continuation_token and _continuation_token.value and _continuation_token.value == resp.continuationToken.value:
-                    self._logger.warning(f"Received the same continuation token value {resp.continuationToken.value} again, stopping to avoid infinite loop.")
-                    break
+            if check_continuation_token(resp.continuationToken, _continuation_token, self._logger):
                 _continuation_token = resp.continuationToken
             else:
                 break
@@ -427,11 +368,249 @@ class TaskHubGrpcClient:
             empty_entities_removed += resp.emptyEntitiesRemoved
             orphaned_locks_released += resp.orphanedLocksReleased
 
-            if resp.continuationToken and resp.continuationToken.value and resp.continuationToken.value != "0":
-                self._logger.info(f"Received continuation token with value {resp.continuationToken.value}, cleaning next page...")
-                if _continuation_token and _continuation_token.value and _continuation_token.value == resp.continuationToken.value:
-                    self._logger.warning(f"Received the same continuation token value {resp.continuationToken.value} again, stopping to avoid infinite loop.")
-                    break
+            if check_continuation_token(resp.continuationToken, _continuation_token, self._logger):
+                _continuation_token = resp.continuationToken
+            else:
+                break
+
+        return CleanEntityStorageResult(empty_entities_removed, orphaned_locks_released)
+
+
+class AsyncTaskHubGrpcClient:
+    """Async version of TaskHubGrpcClient using grpc.aio for asyncio-based applications."""
+
+    def __init__(self, *,
+                 host_address: Optional[str] = None,
+                 metadata: Optional[list[tuple[str, str]]] = None,
+                 log_handler: Optional[logging.Handler] = None,
+                 log_formatter: Optional[logging.Formatter] = None,
+                 secure_channel: bool = False,
+                 interceptors: Optional[Sequence[shared.AsyncClientInterceptor]] = None,
+                 default_version: Optional[str] = None):
+
+        interceptors = prepare_async_interceptors(metadata, interceptors)
+
+        channel = shared.get_async_grpc_channel(
+            host_address=host_address,
+            secure_channel=secure_channel,
+            interceptors=interceptors
+        )
+        self._channel = channel
+        self._stub = stubs.TaskHubSidecarServiceStub(channel)
+        self._logger = shared.get_logger("async_client", log_handler, log_formatter)
+        self.default_version = default_version
+
+    async def close(self) -> None:
+        """Closes the underlying gRPC channel."""
+        await self._channel.close()
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self.close()
+
+    async def schedule_new_orchestration(self, orchestrator: Union[task.Orchestrator[TInput, TOutput], str], *,
+                                         input: Optional[TInput] = None,
+                                         instance_id: Optional[str] = None,
+                                         start_at: Optional[datetime] = None,
+                                         reuse_id_policy: Optional[pb.OrchestrationIdReusePolicy] = None,
+                                         tags: Optional[dict[str, str]] = None,
+                                         version: Optional[str] = None) -> str:
+
+        req = build_schedule_new_orchestration_req(
+            orchestrator, input=input, instance_id=instance_id, start_at=start_at,
+            reuse_id_policy=reuse_id_policy, tags=tags,
+            version=version if version else self.default_version)
+
+        self._logger.info(f"Starting new '{req.name}' instance with ID = '{req.instanceId}'.")
+        res: pb.CreateInstanceResponse = await self._stub.StartInstance(req)
+        return res.instanceId
+
+    async def get_orchestration_state(self, instance_id: str, *,
+                                      fetch_payloads: bool = True) -> Optional[OrchestrationState]:
+        req = pb.GetInstanceRequest(instanceId=instance_id, getInputsAndOutputs=fetch_payloads)
+        res: pb.GetInstanceResponse = await self._stub.GetInstance(req)
+        return new_orchestration_state(req.instanceId, res)
+
+    async def get_all_orchestration_states(self,
+                                           orchestration_query: Optional[OrchestrationQuery] = None
+                                           ) -> List[OrchestrationState]:
+        if orchestration_query is None:
+            orchestration_query = OrchestrationQuery()
+        _continuation_token = None
+
+        self._logger.info(f"Querying orchestration instances with query: {orchestration_query}")
+
+        states = []
+
+        while True:
+            req = build_query_instances_req(orchestration_query, _continuation_token)
+            resp: pb.QueryInstancesResponse = await self._stub.QueryInstances(req)
+            states += [parse_orchestration_state(res) for res in resp.orchestrationState]
+            if check_continuation_token(resp.continuationToken, _continuation_token, self._logger):
+                _continuation_token = resp.continuationToken
+            else:
+                break
+
+        return states
+
+    async def wait_for_orchestration_start(self, instance_id: str, *,
+                                           fetch_payloads: bool = False,
+                                           timeout: int = 60) -> Optional[OrchestrationState]:
+        req = pb.GetInstanceRequest(instanceId=instance_id, getInputsAndOutputs=fetch_payloads)
+        try:
+            self._logger.info(f"Waiting up to {timeout}s for instance '{instance_id}' to start.")
+            res: pb.GetInstanceResponse = await self._stub.WaitForInstanceStart(req, timeout=timeout)
+            return new_orchestration_state(req.instanceId, res)
+        except grpc.aio.AioRpcError as rpc_error:
+            if rpc_error.code() == grpc.StatusCode.DEADLINE_EXCEEDED:
+                raise TimeoutError("Timed-out waiting for the orchestration to start")
+            else:
+                raise
+
+    async def wait_for_orchestration_completion(self, instance_id: str, *,
+                                                fetch_payloads: bool = True,
+                                                timeout: int = 60) -> Optional[OrchestrationState]:
+        req = pb.GetInstanceRequest(instanceId=instance_id, getInputsAndOutputs=fetch_payloads)
+        try:
+            self._logger.info(f"Waiting {timeout}s for instance '{instance_id}' to complete.")
+            res: pb.GetInstanceResponse = await self._stub.WaitForInstanceCompletion(req, timeout=timeout)
+            state = new_orchestration_state(req.instanceId, res)
+            log_completion_state(self._logger, instance_id, state)
+            return state
+        except grpc.aio.AioRpcError as rpc_error:
+            if rpc_error.code() == grpc.StatusCode.DEADLINE_EXCEEDED:
+                raise TimeoutError("Timed-out waiting for the orchestration to complete")
+            else:
+                raise
+
+    async def raise_orchestration_event(self, instance_id: str, event_name: str, *,
+                                        data: Optional[Any] = None) -> None:
+        req = build_raise_event_req(instance_id, event_name, data)
+
+        self._logger.info(f"Raising event '{event_name}' for instance '{instance_id}'.")
+        await self._stub.RaiseEvent(req)
+
+    async def terminate_orchestration(self, instance_id: str, *,
+                                      output: Optional[Any] = None,
+                                      recursive: bool = True) -> None:
+        req = build_terminate_req(instance_id, output, recursive)
+
+        self._logger.info(f"Terminating instance '{instance_id}'.")
+        await self._stub.TerminateInstance(req)
+
+    async def suspend_orchestration(self, instance_id: str) -> None:
+        req = pb.SuspendRequest(instanceId=instance_id)
+        self._logger.info(f"Suspending instance '{instance_id}'.")
+        await self._stub.SuspendInstance(req)
+
+    async def resume_orchestration(self, instance_id: str) -> None:
+        req = pb.ResumeRequest(instanceId=instance_id)
+        self._logger.info(f"Resuming instance '{instance_id}'.")
+        await self._stub.ResumeInstance(req)
+
+    async def restart_orchestration(self, instance_id: str, *,
+                                    restart_with_new_instance_id: bool = False) -> str:
+        """Restarts an existing orchestration instance.
+
+        Args:
+            instance_id: The ID of the orchestration instance to restart.
+            restart_with_new_instance_id: If True, the restarted orchestration will use a new instance ID.
+                If False (default), the restarted orchestration will reuse the same instance ID.
+
+        Returns:
+            The instance ID of the restarted orchestration.
+        """
+        req = pb.RestartInstanceRequest(
+            instanceId=instance_id,
+            restartWithNewInstanceId=restart_with_new_instance_id)
+
+        self._logger.info(f"Restarting instance '{instance_id}'.")
+        res: pb.RestartInstanceResponse = await self._stub.RestartInstance(req)
+        return res.instanceId
+
+    async def purge_orchestration(self, instance_id: str, recursive: bool = True) -> PurgeInstancesResult:
+        req = pb.PurgeInstancesRequest(instanceId=instance_id, recursive=recursive)
+        self._logger.info(f"Purging instance '{instance_id}'.")
+        resp: pb.PurgeInstancesResponse = await self._stub.PurgeInstances(req)
+        return PurgeInstancesResult(resp.deletedInstanceCount, resp.isComplete.value)
+
+    async def purge_orchestrations_by(self,
+                                      created_time_from: Optional[datetime] = None,
+                                      created_time_to: Optional[datetime] = None,
+                                      runtime_status: Optional[List[OrchestrationStatus]] = None,
+                                      recursive: bool = False) -> PurgeInstancesResult:
+        self._logger.info("Purging orchestrations by filter: "
+                          f"created_time_from={created_time_from}, "
+                          f"created_time_to={created_time_to}, "
+                          f"runtime_status={[str(status) for status in runtime_status] if runtime_status else None}, "
+                          f"recursive={recursive}")
+        req = build_purge_by_filter_req(created_time_from, created_time_to, runtime_status, recursive)
+        resp: pb.PurgeInstancesResponse = await self._stub.PurgeInstances(req)
+        return PurgeInstancesResult(resp.deletedInstanceCount, resp.isComplete.value)
+
+    async def signal_entity(self,
+                            entity_instance_id: EntityInstanceId,
+                            operation_name: str,
+                            input: Optional[Any] = None) -> None:
+        req = build_signal_entity_req(entity_instance_id, operation_name, input)
+        self._logger.info(f"Signaling entity '{entity_instance_id}' operation '{operation_name}'.")
+        await self._stub.SignalEntity(req, None)
+
+    async def get_entity(self,
+                         entity_instance_id: EntityInstanceId,
+                         include_state: bool = True
+                         ) -> Optional[EntityMetadata]:
+        req = pb.GetEntityRequest(instanceId=str(entity_instance_id), includeState=include_state)
+        self._logger.info(f"Getting entity '{entity_instance_id}'.")
+        res: pb.GetEntityResponse = await self._stub.GetEntity(req)
+        if not res.exists:
+            return None
+
+        return EntityMetadata.from_entity_metadata(res.entity, include_state)
+
+    async def get_all_entities(self,
+                               entity_query: Optional[EntityQuery] = None) -> List[EntityMetadata]:
+        if entity_query is None:
+            entity_query = EntityQuery()
+        _continuation_token = None
+
+        self._logger.info(f"Retrieving entities by filter: {entity_query}")
+
+        entities = []
+
+        while True:
+            query_request = build_query_entities_req(entity_query, _continuation_token)
+            resp: pb.QueryEntitiesResponse = await self._stub.QueryEntities(query_request)
+            entities += [EntityMetadata.from_entity_metadata(entity, query_request.query.includeState) for entity in resp.entities]
+            if check_continuation_token(resp.continuationToken, _continuation_token, self._logger):
+                _continuation_token = resp.continuationToken
+            else:
+                break
+        return entities
+
+    async def clean_entity_storage(self,
+                                   remove_empty_entities: bool = True,
+                                   release_orphaned_locks: bool = True
+                                   ) -> CleanEntityStorageResult:
+        self._logger.info("Cleaning entity storage")
+
+        empty_entities_removed = 0
+        orphaned_locks_released = 0
+        _continuation_token = None
+
+        while True:
+            req = pb.CleanEntityStorageRequest(
+                removeEmptyEntities=remove_empty_entities,
+                releaseOrphanedLocks=release_orphaned_locks,
+                continuationToken=_continuation_token
+            )
+            resp: pb.CleanEntityStorageResponse = await self._stub.CleanEntityStorage(req)
+            empty_entities_removed += resp.emptyEntitiesRemoved
+            orphaned_locks_released += resp.orphanedLocksReleased
+
+            if check_continuation_token(resp.continuationToken, _continuation_token, self._logger):
                 _continuation_token = resp.continuationToken
             else:
                 break

durabletask-1.3.0.dev22/durabletask/internal/client_helpers.py

@@ -0,0 +1,199 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from __future__ import annotations
+
+import logging
+import uuid
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Any, List, Optional, Sequence, TypeVar, Union
+
+import durabletask.internal.helpers as helpers
+import durabletask.internal.orchestrator_service_pb2 as pb
+import durabletask.internal.shared as shared
+from durabletask import task
+from durabletask.internal.grpc_interceptor import (
+    DefaultAsyncClientInterceptorImpl,
+    DefaultClientInterceptorImpl,
+)
+
+if TYPE_CHECKING:
+    from durabletask.client import (
+        EntityQuery,
+        OrchestrationQuery,
+        OrchestrationState,
+        OrchestrationStatus,
+    )
+    from durabletask.entities import EntityInstanceId
+
+TInput = TypeVar('TInput')
+TOutput = TypeVar('TOutput')
+
+
+def prepare_sync_interceptors(
+        metadata: Optional[list[tuple[str, str]]],
+        interceptors: Optional[Sequence[shared.ClientInterceptor]]
+) -> Optional[list[shared.ClientInterceptor]]:
+    """Prepare the list of sync gRPC interceptors, adding a metadata interceptor if needed."""
+    result: Optional[list[shared.ClientInterceptor]] = None
+    if interceptors is not None:
+        result = list(interceptors)
+        if metadata is not None:
+            result.append(DefaultClientInterceptorImpl(metadata))
+    elif metadata is not None:
+        result = [DefaultClientInterceptorImpl(metadata)]
+    return result
+
+
+def prepare_async_interceptors(
+        metadata: Optional[list[tuple[str, str]]],
+        interceptors: Optional[Sequence[shared.AsyncClientInterceptor]]
+) -> Optional[list[shared.AsyncClientInterceptor]]:
+    """Prepare the list of async gRPC interceptors, adding a metadata interceptor if needed."""
+    result: Optional[list[shared.AsyncClientInterceptor]] = None
+    if interceptors is not None:
+        result = list(interceptors)
+        if metadata is not None:
+            result.append(DefaultAsyncClientInterceptorImpl(metadata))
+    elif metadata is not None:
+        result = [DefaultAsyncClientInterceptorImpl(metadata)]
+    return result
+
+
+def build_schedule_new_orchestration_req(
+        orchestrator: Union[task.Orchestrator[TInput, TOutput], str], *,
+        input: Optional[TInput],
+        instance_id: Optional[str],
+        start_at: Optional[datetime],
+        reuse_id_policy: Optional[pb.OrchestrationIdReusePolicy],
+        tags: Optional[dict[str, str]],
+        version: Optional[str]) -> pb.CreateInstanceRequest:
+    """Build a CreateInstanceRequest for scheduling a new orchestration."""
+    name = orchestrator if isinstance(orchestrator, str) else task.get_name(orchestrator)
+    return pb.CreateInstanceRequest(
+        name=name,
+        instanceId=instance_id if instance_id else uuid.uuid4().hex,
+        input=helpers.get_string_value(shared.to_json(input) if input is not None else None),
+        scheduledStartTimestamp=helpers.new_timestamp(start_at) if start_at else None,
+        version=helpers.get_string_value(version),
+        orchestrationIdReusePolicy=reuse_id_policy,
+        tags=tags
+    )
+
+
+def build_query_instances_req(
+        orchestration_query: OrchestrationQuery,
+        continuation_token) -> pb.QueryInstancesRequest:
+    """Build a QueryInstancesRequest from an OrchestrationQuery."""
+    return pb.QueryInstancesRequest(
+        query=pb.InstanceQuery(
+            runtimeStatus=[status.value for status in orchestration_query.runtime_status] if orchestration_query.runtime_status else None,
+            createdTimeFrom=helpers.new_timestamp(orchestration_query.created_time_from) if orchestration_query.created_time_from else None,
+            createdTimeTo=helpers.new_timestamp(orchestration_query.created_time_to) if orchestration_query.created_time_to else None,
+            maxInstanceCount=orchestration_query.max_instance_count,
+            fetchInputsAndOutputs=orchestration_query.fetch_inputs_and_outputs,
+            continuationToken=continuation_token
+        )
+    )
+
+
+def build_purge_by_filter_req(
+        created_time_from: Optional[datetime],
+        created_time_to: Optional[datetime],
+        runtime_status: Optional[List[OrchestrationStatus]],
+        recursive: bool) -> pb.PurgeInstancesRequest:
+    """Build a PurgeInstancesRequest for purging orchestrations by filter."""
+    return pb.PurgeInstancesRequest(
+        purgeInstanceFilter=pb.PurgeInstanceFilter(
+            createdTimeFrom=helpers.new_timestamp(created_time_from) if created_time_from else None,
+            createdTimeTo=helpers.new_timestamp(created_time_to) if created_time_to else None,
+            runtimeStatus=[status.value for status in runtime_status] if runtime_status else None
+        ),
+        recursive=recursive
+    )
+
+
+def build_query_entities_req(
+        entity_query: EntityQuery,
+        continuation_token) -> pb.QueryEntitiesRequest:
+    """Build a QueryEntitiesRequest from an EntityQuery."""
+    return pb.QueryEntitiesRequest(
+        query=pb.EntityQuery(
+            instanceIdStartsWith=helpers.get_string_value(entity_query.instance_id_starts_with),
+            lastModifiedFrom=helpers.new_timestamp(entity_query.last_modified_from) if entity_query.last_modified_from else None,
+            lastModifiedTo=helpers.new_timestamp(entity_query.last_modified_to) if entity_query.last_modified_to else None,
+            includeState=entity_query.include_state,
+            includeTransient=entity_query.include_transient,
+            pageSize=helpers.get_int_value(entity_query.page_size),
+            continuationToken=continuation_token
+        )
+    )
+
+
+def check_continuation_token(resp_token, prev_token, logger: logging.Logger) -> bool:
+    """Check if a continuation token indicates more pages. Returns True to continue, False to stop."""
+    if resp_token and resp_token.value and resp_token.value != "0":
+        logger.info(f"Received continuation token with value {resp_token.value}, fetching next page...")
+        if prev_token and prev_token.value and prev_token.value == resp_token.value:
+            logger.warning(f"Received the same continuation token value {resp_token.value} again, stopping to avoid infinite loop.")
+            return False
+        return True
+    return False
+
+
+def log_completion_state(
+        logger: logging.Logger,
+        instance_id: str,
+        state: Optional[OrchestrationState]):
+    """Log the final state of a completed orchestration."""
+    if not state:
+        return
+    # Compare against proto constants to avoid circular imports with client.py
+    status_val = state.runtime_status.value
+    if status_val == pb.ORCHESTRATION_STATUS_FAILED and state.failure_details is not None:
+        details = state.failure_details
+        logger.info(f"Instance '{instance_id}' failed: [{details.error_type}] {details.message}")
+    elif status_val == pb.ORCHESTRATION_STATUS_TERMINATED:
+        logger.info(f"Instance '{instance_id}' was terminated.")
+    elif status_val == pb.ORCHESTRATION_STATUS_COMPLETED:
+        logger.info(f"Instance '{instance_id}' completed.")
+
+
+def build_raise_event_req(
+        instance_id: str,
+        event_name: str,
+        data: Optional[Any] = None) -> pb.RaiseEventRequest:
+    """Build a RaiseEventRequest for raising an orchestration event."""
+    return pb.RaiseEventRequest(
+        instanceId=instance_id,
+        name=event_name,
+        input=helpers.get_string_value(shared.to_json(data) if data is not None else None)
+    )
+
+
+def build_terminate_req(
+        instance_id: str,
+        output: Optional[Any] = None,
+        recursive: bool = True) -> pb.TerminateRequest:
+    """Build a TerminateRequest for terminating an orchestration."""
+    return pb.TerminateRequest(
+        instanceId=instance_id,
+        output=helpers.get_string_value(shared.to_json(output) if output is not None else None),
+        recursive=recursive
+    )
+
+
+def build_signal_entity_req(
+        entity_instance_id: EntityInstanceId,
+        operation_name: str,
+        input: Optional[Any] = None) -> pb.SignalEntityRequest:
+    """Build a SignalEntityRequest for signaling an entity."""
+    return pb.SignalEntityRequest(
+        instanceId=str(entity_instance_id),
+        name=operation_name,
+        input=helpers.get_string_value(shared.to_json(input) if input is not None else None),
+        requestId=str(uuid.uuid4()),
+        scheduledTime=None,
+        parentTraceContext=None,
+        requestTime=helpers.new_timestamp(datetime.now(timezone.utc))
+    )

durabletask-1.3.0.dev22/durabletask/internal/grpc_interceptor.py

@@ -0,0 +1,127 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from collections import namedtuple
+
+import grpc
+import grpc.aio
+
+
+class _ClientCallDetails(
+        namedtuple(
+            '_ClientCallDetails',
+            ['method', 'timeout', 'metadata', 'credentials', 'wait_for_ready', 'compression']),
+        grpc.ClientCallDetails):
+    """This is an implementation of the ClientCallDetails interface needed for interceptors.
+    This class takes six named values and inherits the ClientCallDetails from grpc package.
+    This class encloses the values that describe a RPC to be invoked.
+    """
+    pass
+
+
+class _AsyncClientCallDetails(
+        namedtuple(
+            '_AsyncClientCallDetails',
+            ['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 five named values and inherits the ClientCallDetails from grpc.aio package.
+    This class encloses the values that describe a RPC to be invoked.
+    """
+    pass
+
+
+def _apply_metadata(client_call_details, metadata):
+    """Shared logic for applying metadata to call details. Returns the updated metadata list."""
+    if metadata is None:
+        return client_call_details.metadata
+
+    if client_call_details.metadata is not None:
+        new_metadata = list(client_call_details.metadata)
+    else:
+        new_metadata = []
+
+    new_metadata.extend(metadata)
+    return new_metadata
+
+
+class DefaultClientInterceptorImpl (
+        grpc.UnaryUnaryClientInterceptor, grpc.UnaryStreamClientInterceptor,
+        grpc.StreamUnaryClientInterceptor, grpc.StreamStreamClientInterceptor):
+    """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, metadata: list[tuple[str, str]]):
+        super().__init__()
+        self._metadata = metadata
+
+    def _intercept_call(
+            self, client_call_details: grpc.ClientCallDetails) -> grpc.ClientCallDetails:
+        """Internal intercept_call implementation which adds metadata to grpc metadata in the RPC
+            call details."""
+        new_metadata = _apply_metadata(client_call_details, self._metadata)
+        if new_metadata is client_call_details.metadata:
+            return client_call_details
+
+        return _ClientCallDetails(
+            client_call_details.method, client_call_details.timeout, new_metadata,
+            client_call_details.credentials, client_call_details.wait_for_ready, client_call_details.compression)
+
+    def intercept_unary_unary(self, continuation, client_call_details, request):
+        new_client_call_details = self._intercept_call(client_call_details)
+        return continuation(new_client_call_details, request)
+
+    def intercept_unary_stream(self, continuation, client_call_details, request):
+        new_client_call_details = self._intercept_call(client_call_details)
+        return continuation(new_client_call_details, request)
+
+    def intercept_stream_unary(self, continuation, client_call_details, request):
+        new_client_call_details = self._intercept_call(client_call_details)
+        return continuation(new_client_call_details, request)
+
+    def intercept_stream_stream(self, continuation, client_call_details, request):
+        new_client_call_details = self._intercept_call(client_call_details)
+        return continuation(new_client_call_details, request)
+
+
+class DefaultAsyncClientInterceptorImpl(
+        grpc.aio.UnaryUnaryClientInterceptor, grpc.aio.UnaryStreamClientInterceptor,
+        grpc.aio.StreamUnaryClientInterceptor, grpc.aio.StreamStreamClientInterceptor):
+    """Async gRPC interceptor that adds metadata headers to all calls."""
+
+    def __init__(self, metadata: list[tuple[str, str]]):
+        self._metadata = metadata
+
+    async def _intercept_call(
+            self, client_call_details: grpc.aio.ClientCallDetails) -> grpc.aio.ClientCallDetails:
+        """Internal intercept_call implementation which adds metadata to grpc metadata in the RPC
+            call details. This method is async to allow subclasses to perform async operations
+            (e.g., refreshing auth tokens) during interception."""
+        new_metadata = _apply_metadata(client_call_details, self._metadata)
+        if new_metadata is client_call_details.metadata:
+            return client_call_details
+
+        return _AsyncClientCallDetails(
+            client_call_details.method,
+            client_call_details.timeout,
+            new_metadata,
+            client_call_details.credentials,
+            client_call_details.wait_for_ready,
+        )
+
+    async def intercept_unary_unary(self, continuation, client_call_details, request):
+        new_client_call_details = await self._intercept_call(client_call_details)
+        return await continuation(new_client_call_details, request)
+
+    async def intercept_unary_stream(self, continuation, client_call_details, request):
+        new_client_call_details = await self._intercept_call(client_call_details)
+        return await continuation(new_client_call_details, request)
+
+    async def intercept_stream_unary(self, continuation, client_call_details, request_iterator):
+        new_client_call_details = await self._intercept_call(client_call_details)
+        return await continuation(new_client_call_details, request_iterator)
+
+    async def intercept_stream_stream(self, continuation, client_call_details, request_iterator):
+        new_client_call_details = await self._intercept_call(client_call_details)
+        return await continuation(new_client_call_details, request_iterator)

{durabletask-1.3.0.dev21 → durabletask-1.3.0.dev22}/durabletask/internal/shared.py

@@ -8,6 +8,7 @@ from types import SimpleNamespace
 from typing import Any, Optional, Sequence, Union
 
 import grpc
+import grpc.aio
 
 ClientInterceptor = Union[
     grpc.UnaryUnaryClientInterceptor,
@@ -16,6 +17,13 @@ ClientInterceptor = Union[
     grpc.StreamStreamClientInterceptor
 ]
 
+AsyncClientInterceptor = Union[
+    grpc.aio.UnaryUnaryClientInterceptor,
+    grpc.aio.UnaryStreamClientInterceptor,
+    grpc.aio.StreamUnaryClientInterceptor,
+    grpc.aio.StreamStreamClientInterceptor
+]
+
 # Field name used to indicate that an object was automatically serialized
 # and should be deserialized as a SimpleNamespace
 AUTO_SERIALIZED = "__durabletask_autoobject__"
@@ -62,6 +70,38 @@ def get_grpc_channel(
     return channel
 
 
+def get_async_grpc_channel(
+        host_address: Optional[str],
+        secure_channel: bool = False,
+        interceptors: Optional[Sequence[AsyncClientInterceptor]] = None) -> grpc.aio.Channel:
+
+    if host_address is None:
+        host_address = get_default_host_address()
+
+    for protocol in SECURE_PROTOCOLS:
+        if host_address.lower().startswith(protocol):
+            secure_channel = True
+            host_address = host_address[len(protocol):]
+            break
+
+    for protocol in INSECURE_PROTOCOLS:
+        if host_address.lower().startswith(protocol):
+            secure_channel = False
+            host_address = host_address[len(protocol):]
+            break
+
+    if secure_channel:
+        channel = grpc.aio.secure_channel(
+            host_address, grpc.ssl_channel_credentials(),
+            interceptors=interceptors)
+    else:
+        channel = grpc.aio.insecure_channel(
+            host_address,
+            interceptors=interceptors)
+
+    return channel
+
+
 def get_logger(
         name_suffix: str,
         log_handler: Optional[logging.Handler] = None,

{durabletask-1.3.0.dev21 → durabletask-1.3.0.dev22}/durabletask.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: durabletask
-Version: 1.3.0.dev21
+Version: 1.3.0.dev22
 Summary: A Durable Task Client SDK for Python
 License:     MIT License
         

{durabletask-1.3.0.dev21 → durabletask-1.3.0.dev22}/durabletask.egg-info/SOURCES.txt

@@ -18,6 +18,7 @@ durabletask/entities/entity_instance_id.py
 durabletask/entities/entity_lock.py
 durabletask/entities/entity_metadata.py
 durabletask/entities/entity_operation_failed_exception.py
+durabletask/internal/client_helpers.py
 durabletask/internal/entity_state_shim.py
 durabletask/internal/exceptions.py
 durabletask/internal/grpc_interceptor.py

{durabletask-1.3.0.dev21 → durabletask-1.3.0.dev22}/pyproject.toml

@@ -9,7 +9,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "durabletask"
-version = "1.3.0.dev21"
+version = "1.3.0.dev22"
 description = "A Durable Task Client SDK for Python"
 keywords = [
     "durable",
@@ -41,3 +41,4 @@ include = ["durabletask", "durabletask.*"]
 [tool.pytest.ini_options]
 minversion = "6.0"
 testpaths = ["tests"]
+asyncio_mode = "auto"

durabletask-1.3.0.dev21/durabletask/internal/grpc_interceptor.py

@@ -1,65 +0,0 @@
-# Copyright (c) Microsoft Corporation.
-# Licensed under the MIT License.
-
-from collections import namedtuple
-
-import grpc
-
-
-class _ClientCallDetails(
-        namedtuple(
-            '_ClientCallDetails',
-            ['method', 'timeout', 'metadata', 'credentials', 'wait_for_ready', 'compression']),
-        grpc.ClientCallDetails):
-    """This is an implementation of the ClientCallDetails interface needed for interceptors.
-    This class takes six named values and inherits the ClientCallDetails from grpc package.
-    This class encloses the values that describe a RPC to be invoked.
-    """
-    pass
-
-
-class DefaultClientInterceptorImpl (
-        grpc.UnaryUnaryClientInterceptor, grpc.UnaryStreamClientInterceptor,
-        grpc.StreamUnaryClientInterceptor, grpc.StreamStreamClientInterceptor):
-    """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, metadata: list[tuple[str, str]]):
-        super().__init__()
-        self._metadata = metadata
-
-    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."""
-        if self._metadata is None:
-            return client_call_details
-
-        if client_call_details.metadata is not None:
-            metadata = list(client_call_details.metadata)
-        else:
-            metadata = []
-
-        metadata.extend(self._metadata)
-        client_call_details = _ClientCallDetails(
-            client_call_details.method, client_call_details.timeout, metadata,
-            client_call_details.credentials, client_call_details.wait_for_ready, client_call_details.compression)
-
-        return client_call_details
-
-    def intercept_unary_unary(self, continuation, client_call_details, request):
-        new_client_call_details = self._intercept_call(client_call_details)
-        return continuation(new_client_call_details, request)
-
-    def intercept_unary_stream(self, continuation, client_call_details, request):
-        new_client_call_details = self._intercept_call(client_call_details)
-        return continuation(new_client_call_details, request)
-
-    def intercept_stream_unary(self, continuation, client_call_details, request):
-        new_client_call_details = self._intercept_call(client_call_details)
-        return continuation(new_client_call_details, request)
-
-    def intercept_stream_stream(self, continuation, client_call_details, request):
-        new_client_call_details = self._intercept_call(client_call_details)
-        return continuation(new_client_call_details, request)