durabletask 0.0.0.dev1__py3-none-any.whl

durabletask/internal/ProtoTaskHubSidecarServiceStub.py

@@ -0,0 +1,34 @@
+from typing import Any, Callable, Protocol
+
+
+class ProtoTaskHubSidecarServiceStub(Protocol):
+    """A stub class roughly matching the TaskHubSidecarServiceStub generated from the .proto file.
+    Used by Azure Functions during orchestration and entity executions to inject custom behavior,
+    as no real sidecar stub is available.
+    """
+    Hello: Callable[..., Any]
+    StartInstance: Callable[..., Any]
+    GetInstance: Callable[..., Any]
+    RewindInstance: Callable[..., Any]
+    WaitForInstanceStart: Callable[..., Any]
+    WaitForInstanceCompletion: Callable[..., Any]
+    RaiseEvent: Callable[..., Any]
+    TerminateInstance: Callable[..., Any]
+    SuspendInstance: Callable[..., Any]
+    ResumeInstance: Callable[..., Any]
+    QueryInstances: Callable[..., Any]
+    PurgeInstances: Callable[..., Any]
+    GetWorkItems: Callable[..., Any]
+    CompleteActivityTask: Callable[..., Any]
+    CompleteOrchestratorTask: Callable[..., Any]
+    CompleteEntityTask: Callable[..., Any]
+    StreamInstanceHistory: Callable[..., Any]
+    CreateTaskHub: Callable[..., Any]
+    DeleteTaskHub: Callable[..., Any]
+    SignalEntity: Callable[..., Any]
+    GetEntity: Callable[..., Any]
+    QueryEntities: Callable[..., Any]
+    CleanEntityStorage: Callable[..., Any]
+    AbandonTaskActivityWorkItem: Callable[..., Any]
+    AbandonTaskOrchestratorWorkItem: Callable[..., Any]
+    AbandonTaskEntityWorkItem: Callable[..., Any]

durabletask/internal/entity_state_shim.py

@@ -0,0 +1,66 @@
+from typing import Any, TypeVar, Union
+from typing import Optional, Type, overload
+
+import durabletask.internal.orchestrator_service_pb2 as pb
+
+TState = TypeVar("TState")
+
+
+class StateShim:
+    def __init__(self, start_state):
+        self._current_state: Any = start_state
+        self._checkpoint_state: Any = start_state
+        self._operation_actions: list[pb.OperationAction] = []
+        self._actions_checkpoint_state: int = 0
+
+    @overload
+    def get_state(self, intended_type: Type[TState], default: TState) -> TState:
+        ...
+
+    @overload
+    def get_state(self, intended_type: Type[TState]) -> Optional[TState]:
+        ...
+
+    @overload
+    def get_state(self, intended_type: None = None, default: Any = None) -> Any:
+        ...
+
+    def get_state(self, intended_type: Optional[Type[TState]] = None, default: Optional[TState] = None) -> Union[None, TState, Any]:
+        if self._current_state is None and default is not None:
+            return default
+
+        if intended_type is None:
+            return self._current_state
+
+        if isinstance(self._current_state, intended_type):
+            return self._current_state
+
+        try:
+            return intended_type(self._current_state)  # type: ignore[call-arg]
+        except Exception as ex:
+            raise TypeError(
+                f"Could not convert state of type '{type(self._current_state).__name__}' to '{intended_type.__name__}'"
+            ) from ex
+
+    def set_state(self, state):
+        self._current_state = state
+
+    def add_operation_action(self, action: pb.OperationAction):
+        self._operation_actions.append(action)
+
+    def get_operation_actions(self) -> list[pb.OperationAction]:
+        return self._operation_actions[:self._actions_checkpoint_state]
+
+    def commit(self):
+        self._checkpoint_state = self._current_state
+        self._actions_checkpoint_state = len(self._operation_actions)
+
+    def rollback(self):
+        self._current_state = self._checkpoint_state
+        self._operation_actions = self._operation_actions[:self._actions_checkpoint_state]
+
+    def reset(self):
+        self._current_state = None
+        self._checkpoint_state = None
+        self._operation_actions = []
+        self._actions_checkpoint_state = 0

durabletask/internal/exceptions.py

@@ -0,0 +1,11 @@
+import durabletask.internal.orchestrator_service_pb2 as pb
+
+
+class VersionFailureException(Exception):
+    def __init__(self, error_details: pb.TaskFailureDetails) -> None:
+        super().__init__()
+        self.error_details = error_details
+
+
+class AbandonOrchestrationError(Exception):
+    pass

durabletask/internal/grpc_interceptor.py

@@ -0,0 +1,65 @@
+# 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)

durabletask/internal/helpers.py

@@ -0,0 +1,288 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import traceback
+from datetime import datetime
+from typing import Optional
+
+from google.protobuf import timestamp_pb2, wrappers_pb2
+
+from durabletask.entities import EntityInstanceId
+import durabletask.internal.orchestrator_service_pb2 as pb
+
+# TODO: The new_xxx_event methods are only used by test code and should be moved elsewhere
+
+
+def new_orchestrator_started_event(timestamp: Optional[datetime] = None) -> pb.HistoryEvent:
+    ts = timestamp_pb2.Timestamp()
+    if timestamp is not None:
+        ts.FromDatetime(timestamp)
+    return pb.HistoryEvent(eventId=-1, timestamp=ts, orchestratorStarted=pb.OrchestratorStartedEvent())
+
+
+def new_execution_started_event(name: str, instance_id: str, encoded_input: Optional[str] = None,
+                                tags: Optional[dict[str, str]] = None) -> pb.HistoryEvent:
+    return pb.HistoryEvent(
+        eventId=-1,
+        timestamp=timestamp_pb2.Timestamp(),
+        executionStarted=pb.ExecutionStartedEvent(
+            name=name,
+            input=get_string_value(encoded_input),
+            orchestrationInstance=pb.OrchestrationInstance(instanceId=instance_id),
+            tags=tags))
+
+
+def new_timer_created_event(timer_id: int, fire_at: datetime) -> pb.HistoryEvent:
+    ts = timestamp_pb2.Timestamp()
+    ts.FromDatetime(fire_at)
+    return pb.HistoryEvent(
+        eventId=timer_id,
+        timestamp=timestamp_pb2.Timestamp(),
+        timerCreated=pb.TimerCreatedEvent(fireAt=ts)
+    )
+
+
+def new_timer_fired_event(timer_id: int, fire_at: datetime) -> pb.HistoryEvent:
+    ts = timestamp_pb2.Timestamp()
+    ts.FromDatetime(fire_at)
+    return pb.HistoryEvent(
+        eventId=-1,
+        timestamp=timestamp_pb2.Timestamp(),
+        timerFired=pb.TimerFiredEvent(fireAt=ts, timerId=timer_id)
+    )
+
+
+def new_task_scheduled_event(event_id: int, name: str, encoded_input: Optional[str] = None) -> pb.HistoryEvent:
+    return pb.HistoryEvent(
+        eventId=event_id,
+        timestamp=timestamp_pb2.Timestamp(),
+        taskScheduled=pb.TaskScheduledEvent(name=name, input=get_string_value(encoded_input))
+    )
+
+
+def new_task_completed_event(event_id: int, encoded_output: Optional[str] = None) -> pb.HistoryEvent:
+    return pb.HistoryEvent(
+        eventId=-1,
+        timestamp=timestamp_pb2.Timestamp(),
+        taskCompleted=pb.TaskCompletedEvent(taskScheduledId=event_id, result=get_string_value(encoded_output))
+    )
+
+
+def new_task_failed_event(event_id: int, ex: Exception) -> pb.HistoryEvent:
+    return pb.HistoryEvent(
+        eventId=-1,
+        timestamp=timestamp_pb2.Timestamp(),
+        taskFailed=pb.TaskFailedEvent(taskScheduledId=event_id, failureDetails=new_failure_details(ex))
+    )
+
+
+def new_sub_orchestration_created_event(
+        event_id: int,
+        name: str,
+        instance_id: str,
+        encoded_input: Optional[str] = None) -> pb.HistoryEvent:
+    return pb.HistoryEvent(
+        eventId=event_id,
+        timestamp=timestamp_pb2.Timestamp(),
+        subOrchestrationInstanceCreated=pb.SubOrchestrationInstanceCreatedEvent(
+            name=name,
+            input=get_string_value(encoded_input),
+            instanceId=instance_id)
+    )
+
+
+def new_sub_orchestration_completed_event(event_id: int, encoded_output: Optional[str] = None) -> pb.HistoryEvent:
+    return pb.HistoryEvent(
+        eventId=-1,
+        timestamp=timestamp_pb2.Timestamp(),
+        subOrchestrationInstanceCompleted=pb.SubOrchestrationInstanceCompletedEvent(
+            result=get_string_value(encoded_output),
+            taskScheduledId=event_id)
+    )
+
+
+def new_sub_orchestration_failed_event(event_id: int, ex: Exception) -> pb.HistoryEvent:
+    return pb.HistoryEvent(
+        eventId=-1,
+        timestamp=timestamp_pb2.Timestamp(),
+        subOrchestrationInstanceFailed=pb.SubOrchestrationInstanceFailedEvent(
+            failureDetails=new_failure_details(ex),
+            taskScheduledId=event_id)
+    )
+
+
+def new_failure_details(ex: Exception) -> pb.TaskFailureDetails:
+    return pb.TaskFailureDetails(
+        errorType=type(ex).__name__,
+        errorMessage=str(ex),
+        stackTrace=wrappers_pb2.StringValue(value=''.join(traceback.format_tb(ex.__traceback__)))
+    )
+
+
+def new_event_raised_event(name: str, encoded_input: Optional[str] = None) -> pb.HistoryEvent:
+    return pb.HistoryEvent(
+        eventId=-1,
+        timestamp=timestamp_pb2.Timestamp(),
+        eventRaised=pb.EventRaisedEvent(name=name, input=get_string_value(encoded_input))
+    )
+
+
+def new_suspend_event() -> pb.HistoryEvent:
+    return pb.HistoryEvent(
+        eventId=-1,
+        timestamp=timestamp_pb2.Timestamp(),
+        executionSuspended=pb.ExecutionSuspendedEvent()
+    )
+
+
+def new_resume_event() -> pb.HistoryEvent:
+    return pb.HistoryEvent(
+        eventId=-1,
+        timestamp=timestamp_pb2.Timestamp(),
+        executionResumed=pb.ExecutionResumedEvent()
+    )
+
+
+def new_terminated_event(*, encoded_output: Optional[str] = None) -> pb.HistoryEvent:
+    return pb.HistoryEvent(
+        eventId=-1,
+        timestamp=timestamp_pb2.Timestamp(),
+        executionTerminated=pb.ExecutionTerminatedEvent(
+            input=get_string_value(encoded_output)
+        )
+    )
+
+
+def get_string_value(val: Optional[str]) -> Optional[wrappers_pb2.StringValue]:
+    if val is None:
+        return None
+    else:
+        return wrappers_pb2.StringValue(value=val)
+
+
+def get_string_value_or_empty(val: Optional[str]) -> wrappers_pb2.StringValue:
+    if val is None:
+        return wrappers_pb2.StringValue(value="")
+    return wrappers_pb2.StringValue(value=val)
+
+
+def new_complete_orchestration_action(
+        id: int,
+        status: pb.OrchestrationStatus,
+        result: Optional[str] = None,
+        failure_details: Optional[pb.TaskFailureDetails] = None,
+        carryover_events: Optional[list[pb.HistoryEvent]] = None) -> pb.OrchestratorAction:
+    completeOrchestrationAction = pb.CompleteOrchestrationAction(
+        orchestrationStatus=status,
+        result=get_string_value(result),
+        failureDetails=failure_details,
+        carryoverEvents=carryover_events)
+
+    return pb.OrchestratorAction(id=id, completeOrchestration=completeOrchestrationAction)
+
+
+def new_create_timer_action(id: int, fire_at: datetime) -> pb.OrchestratorAction:
+    timestamp = timestamp_pb2.Timestamp()
+    timestamp.FromDatetime(fire_at)
+    return pb.OrchestratorAction(id=id, createTimer=pb.CreateTimerAction(fireAt=timestamp))
+
+
+def new_schedule_task_action(id: int, name: str, encoded_input: Optional[str],
+                             tags: Optional[dict[str, str]]) -> pb.OrchestratorAction:
+    return pb.OrchestratorAction(id=id, scheduleTask=pb.ScheduleTaskAction(
+        name=name,
+        input=get_string_value(encoded_input),
+        tags=tags
+    ))
+
+
+def new_call_entity_action(id: int,
+                           parent_instance_id: str,
+                           entity_id: EntityInstanceId,
+                           operation: str, encoded_input: Optional[str],
+                           request_id: str):
+    return pb.OrchestratorAction(id=id, sendEntityMessage=pb.SendEntityMessageAction(entityOperationCalled=pb.EntityOperationCalledEvent(
+        requestId=request_id,
+        operation=operation,
+        scheduledTime=None,
+        input=get_string_value(encoded_input),
+        parentInstanceId=get_string_value(parent_instance_id),
+        parentExecutionId=None,
+        targetInstanceId=get_string_value(str(entity_id)),
+    )))
+
+
+def new_signal_entity_action(id: int,
+                             entity_id: EntityInstanceId,
+                             operation: str,
+                             encoded_input: Optional[str],
+                             request_id: str):
+    return pb.OrchestratorAction(id=id, sendEntityMessage=pb.SendEntityMessageAction(entityOperationSignaled=pb.EntityOperationSignaledEvent(
+        requestId=request_id,
+        operation=operation,
+        scheduledTime=None,
+        input=get_string_value(encoded_input),
+        targetInstanceId=get_string_value(str(entity_id)),
+    )))
+
+
+def new_lock_entities_action(id: int, entity_message: pb.SendEntityMessageAction):
+    return pb.OrchestratorAction(id=id, sendEntityMessage=entity_message)
+
+
+def convert_to_entity_batch_request(req: pb.EntityRequest) -> tuple[pb.EntityBatchRequest, list[pb.OperationInfo]]:
+    batch_request = pb.EntityBatchRequest(entityState=req.entityState, instanceId=req.instanceId, operations=[])
+
+    operation_infos: list[pb.OperationInfo] = []
+
+    for op in req.operationRequests:
+        if op.HasField("entityOperationSignaled"):
+            batch_request.operations.append(pb.OperationRequest(requestId=op.entityOperationSignaled.requestId,
+                                                                operation=op.entityOperationSignaled.operation,
+                                                                input=op.entityOperationSignaled.input))
+            operation_infos.append(pb.OperationInfo(requestId=op.entityOperationSignaled.requestId,
+                                                    responseDestination=None))
+        elif op.HasField("entityOperationCalled"):
+            batch_request.operations.append(pb.OperationRequest(requestId=op.entityOperationCalled.requestId,
+                                                                operation=op.entityOperationCalled.operation,
+                                                                input=op.entityOperationCalled.input))
+            operation_infos.append(pb.OperationInfo(requestId=op.entityOperationCalled.requestId,
+                                                    responseDestination=pb.OrchestrationInstance(
+                                                        instanceId=op.entityOperationCalled.parentInstanceId.value,
+                                                        executionId=op.entityOperationCalled.parentExecutionId
+                                                    )))
+
+    return batch_request, operation_infos
+
+
+def new_timestamp(dt: datetime) -> timestamp_pb2.Timestamp:
+    ts = timestamp_pb2.Timestamp()
+    ts.FromDatetime(dt)
+    return ts
+
+
+def new_create_sub_orchestration_action(
+        id: int,
+        name: str,
+        instance_id: Optional[str],
+        encoded_input: Optional[str],
+        version: Optional[str]) -> pb.OrchestratorAction:
+    return pb.OrchestratorAction(id=id, createSubOrchestration=pb.CreateSubOrchestrationAction(
+        name=name,
+        instanceId=instance_id,
+        input=get_string_value(encoded_input),
+        version=get_string_value(version)
+    ))
+
+
+def is_empty(v: wrappers_pb2.StringValue):
+    return v is None or v.value == ''
+
+
+def get_orchestration_status_str(status: pb.OrchestrationStatus):
+    try:
+        const_name = pb.OrchestrationStatus.Name(status)
+        if const_name.startswith('ORCHESTRATION_STATUS_'):
+            return const_name[len('ORCHESTRATION_STATUS_'):]
+    except Exception:
+        return "UNKNOWN"

durabletask/internal/orchestration_entity_context.py

@@ -0,0 +1,115 @@
+from datetime import datetime
+from typing import Generator, List, Optional, Tuple, Union
+
+from durabletask.internal.helpers import get_string_value
+import durabletask.internal.orchestrator_service_pb2 as pb
+from durabletask.entities import EntityInstanceId
+
+
+class OrchestrationEntityContext:
+    def __init__(self, instance_id: str):
+        self.instance_id = instance_id
+
+        self.lock_acquisition_pending = False
+
+        self.critical_section_id = None
+        self.critical_section_locks: list[EntityInstanceId] = []
+        self.available_locks: list[EntityInstanceId] = []
+
+    @property
+    def is_inside_critical_section(self) -> bool:
+        return self.critical_section_id is not None
+
+    def get_available_entities(self) -> Generator[EntityInstanceId, None, None]:
+        if self.is_inside_critical_section:
+            for available_lock in self.available_locks:
+                yield available_lock
+
+    def validate_suborchestration_transition(self) -> Tuple[bool, str]:
+        if self.is_inside_critical_section:
+            return False, "While holding locks, cannot call suborchestrators."
+        return True, ""
+
+    def validate_operation_transition(self, target_instance_id: EntityInstanceId, one_way: bool) -> Tuple[bool, str]:
+        if self.is_inside_critical_section:
+            lock_to_use = target_instance_id
+            if one_way:
+                if target_instance_id in self.critical_section_locks:
+                    return False, "Must not signal a locked entity from a critical section."
+            else:
+                try:
+                    self.available_locks.remove(lock_to_use)
+                except ValueError:
+                    if self.lock_acquisition_pending:
+                        return False, "Must await the completion of the lock request prior to calling any entity."
+                    if lock_to_use in self.critical_section_locks:
+                        return False, "Must not call an entity from a critical section while a prior call to the same entity is still pending."
+                    else:
+                        return False, "Must not call an entity from a critical section if it is not one of the locked entities."
+        return True, ""
+
+    def validate_acquire_transition(self) -> Tuple[bool, str]:
+        if self.is_inside_critical_section:
+            return False, "Must not enter another critical section from within a critical section."
+        return True, ""
+
+    def recover_lock_after_call(self, target_instance_id: EntityInstanceId):
+        if self.is_inside_critical_section:
+            self.available_locks.append(target_instance_id)
+
+    def emit_lock_release_messages(self):
+        if self.is_inside_critical_section:
+            for entity_id in self.critical_section_locks:
+                unlock_event = pb.SendEntityMessageAction(entityUnlockSent=pb.EntityUnlockSentEvent(
+                    criticalSectionId=self.critical_section_id,
+                    targetInstanceId=get_string_value(str(entity_id)),
+                    parentInstanceId=get_string_value(self.instance_id)
+                ))
+                yield unlock_event
+
+            self.critical_section_locks = []
+            self.available_locks = []
+            self.critical_section_id = None
+
+    def emit_request_message(self, target, operation_name: str, one_way: bool, operation_id: str,
+                             scheduled_time_utc: datetime, input: Optional[str],
+                             request_time: Optional[datetime] = None, create_trace: bool = False):
+        raise NotImplementedError()
+
+    def emit_acquire_message(self, critical_section_id: str, entities: List[EntityInstanceId]) -> Union[Tuple[None, None], Tuple[pb.SendEntityMessageAction, pb.OrchestrationInstance]]:
+        if not entities:
+            return None, None
+
+        # Acquire the locks in a globally fixed order to avoid deadlocks
+        # Also remove duplicates - this can be optimized for perf if necessary
+        entity_ids = sorted(entities)
+        entity_ids_dedup = []
+        for i, entity_id in enumerate(entity_ids):
+            if entity_id != entity_ids[i - 1] if i > 0 else True:
+                entity_ids_dedup.append(entity_id)
+
+        target = pb.OrchestrationInstance(instanceId=str(entity_ids_dedup[0]))
+        request = pb.SendEntityMessageAction(entityLockRequested=pb.EntityLockRequestedEvent(
+            criticalSectionId=critical_section_id,
+            parentInstanceId=get_string_value(self.instance_id),
+            lockSet=[str(eid) for eid in entity_ids_dedup],
+            position=0,
+        ))
+
+        self.critical_section_id = critical_section_id
+        self.critical_section_locks = entity_ids_dedup
+        self.lock_acquisition_pending = True
+
+        return request, target
+
+    def complete_acquire(self, critical_section_id):
+        if self.critical_section_id != critical_section_id:
+            raise RuntimeError(f"Unexpected lock acquire for critical section ID '{critical_section_id}' (expected '{self.critical_section_id}')")
+        self.available_locks = self.critical_section_locks
+        self.lock_acquisition_pending = False
+
+    def adjust_outgoing_message(self, instance_id: str, request_message, capped_time: datetime) -> str:
+        raise NotImplementedError()
+
+    def deserialize_entity_response_event(self, event_content: str):
+        raise NotImplementedError()