Update durabletask protos, set custom status (#31)

Signed-off-by: Fabian Martinez <[email protected]>

CHANGELOG.md

@@ -5,6 +5,16 @@ 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).
 
+## v0.2.0 (Unreleased)
+
+### New
+
+- Support for orchestration custom status ([#31](https://github.com/microsoft/durabletask-python/pull/31)) - contributed by [@famarting](https://github.com/famarting)
+
+### Updates
+
+- Updated `durabletask-protobuf` submodule reference to latest
+
 ## v0.1.1a1
 
 ### New

durabletask/client.py

@@ -104,7 +104,8 @@ def __init__(self, *,
     def schedule_new_orchestration(self, orchestrator: Union[task.Orchestrator[TInput, TOutput], str], *,
                                    input: Union[TInput, None] = None,
                                    instance_id: Union[str, None] = None,
-                                   start_at: Union[datetime, None] = None) -> str:
+                                   start_at: Union[datetime, None] = None,
+                                   reuse_id_policy: Union[pb.OrchestrationIdReusePolicy, None] = None) -> str:
 
         name = orchestrator if isinstance(orchestrator, str) else task.get_name(orchestrator)
 
@@ -113,7 +114,9 @@ def schedule_new_orchestration(self, orchestrator: Union[task.Orchestrator[TInpu
             instanceId=instance_id if instance_id else uuid.uuid4().hex,
             input=wrappers_pb2.StringValue(value=shared.to_json(input)) if input is not None else None,
             scheduledStartTimestamp=helpers.new_timestamp(start_at) if start_at else None,
-            version=wrappers_pb2.StringValue(value=""))
+            version=wrappers_pb2.StringValue(value=""),
+            orchestrationIdReusePolicy=reuse_id_policy,
+            )
 
         self._logger.info(f"Starting new '{name}' instance with ID = '{req.instanceId}'.")
         res: pb.CreateInstanceResponse = self._stub.StartInstance(req)

durabletask/task.py

@@ -71,6 +71,12 @@ def is_replaying(self) -> bool:
         """
         pass
 
+    @abstractmethod
+    def set_custom_status(self, custom_status: str) -> None:
+        """Set the custom status.
+        """
+        pass
+
     @abstractmethod
     def create_timer(self, fire_at: Union[datetime, timedelta]) -> Task:
         """Create a Timer Task to fire after at the specified deadline.

durabletask/worker.py

@@ -10,7 +10,7 @@
                     TypeVar, Union)
 
 import grpc
-from google.protobuf import empty_pb2
+from google.protobuf import empty_pb2, wrappers_pb2
 
 import durabletask.internal.helpers as ph
 import durabletask.internal.helpers as pbh
@@ -188,8 +188,8 @@ def stop(self):
     def _execute_orchestrator(self, req: pb.OrchestratorRequest, stub: stubs.TaskHubSidecarServiceStub):
         try:
             executor = _OrchestrationExecutor(self._registry, self._logger)
-            actions = executor.execute(req.instanceId, req.pastEvents, req.newEvents)
-            res = pb.OrchestratorResponse(instanceId=req.instanceId, actions=actions)
+            result = executor.execute(req.instanceId, req.pastEvents, req.newEvents)
+            res = pb.OrchestratorResponse(instanceId=req.instanceId, actions=result.actions, customStatus=wrappers_pb2.StringValue(value=result.custom_status))
         except Exception as ex:
             self._logger.exception(f"An error occurred while trying to execute instance '{req.instanceId}': {ex}")
             failure_details = pbh.new_failure_details(ex)
@@ -242,6 +242,7 @@ def __init__(self, instance_id: str):
         self._pending_events: Dict[str, List[task.CompletableTask]] = {}
         self._new_input: Optional[Any] = None
         self._save_events = False
+        self._custom_status: str = ""
 
     def run(self, generator: Generator[task.Task, Any, Any]):
         self._generator = generator
@@ -355,6 +356,9 @@ def is_replaying(self) -> bool:
     def current_utc_datetime(self, value: datetime):
         self._current_utc_datetime = value
 
+    def set_custom_status(self, custom_status: str) -> None:
+        self._custom_status = custom_status
+
     def create_timer(self, fire_at: Union[datetime, timedelta]) -> task.Task:
         return self.create_timer_internal(fire_at)
 
@@ -457,6 +461,14 @@ def continue_as_new(self, new_input, *, save_events: bool = False) -> None:
         self.set_continued_as_new(new_input, save_events)
 
 
+class ExecutionResults:
+    actions: List[pb.OrchestratorAction]
+    custom_status: str
+
+    def __init__(self, actions: List[pb.OrchestratorAction], custom_status: str):
+        self.actions = actions
+        self.custom_status = custom_status
+
 class _OrchestrationExecutor:
     _generator: Optional[task.Orchestrator] = None
 
@@ -466,7 +478,7 @@ def __init__(self, registry: _Registry, logger: logging.Logger):
         self._is_suspended = False
         self._suspended_events: List[pb.HistoryEvent] = []
 
-    def execute(self, instance_id: str, old_events: Sequence[pb.HistoryEvent], new_events: Sequence[pb.HistoryEvent]) -> List[pb.OrchestratorAction]:
+    def execute(self, instance_id: str, old_events: Sequence[pb.HistoryEvent], new_events: Sequence[pb.HistoryEvent]) -> ExecutionResults:
         if not new_events:
             raise task.OrchestrationStateError("The new history event list must have at least one event in it.")
 
@@ -501,7 +513,7 @@ def execute(self, instance_id: str, old_events: Sequence[pb.HistoryEvent], new_e
         actions = ctx.get_actions()
         if self._logger.level <= logging.DEBUG:
             self._logger.debug(f"{instance_id}: Returning {len(actions)} action(s): {_get_action_summary(actions)}")
-        return actions
+        return ExecutionResults(actions=actions, custom_status=ctx._custom_status)
 
     def process_event(self, ctx: _RuntimeOrchestrationContext, event: pb.HistoryEvent) -> None:
         if self._is_suspended and _is_suspendable(event):

tests/test_orchestration_e2e.py

@@ -441,3 +441,26 @@ def throw_activity(ctx: task.ActivityContext, _):
         assert state.failure_details.message.endswith("Activity task #1 failed: Kah-BOOOOM!!!")
         assert state.failure_details.stack_trace is not None
         assert throw_activity_counter == 4
+
+def test_custom_status():
+
+    def empty_orchestrator(ctx: task.OrchestrationContext, _):
+        ctx.set_custom_status("foobaz")
+
+    # Start a worker, which will connect to the sidecar in a background thread
+    with worker.TaskHubGrpcWorker() as w:
+        w.add_orchestrator(empty_orchestrator)
+        w.start()
+
+        c = client.TaskHubGrpcClient()
+        id = c.schedule_new_orchestration(empty_orchestrator)
+        state = c.wait_for_orchestration_completion(id, timeout=30)
+
+    assert state is not None
+    assert state.name == task.get_name(empty_orchestrator)
+    assert state.instance_id == id
+    assert state.failure_details is None
+    assert state.runtime_status == client.OrchestrationStatus.COMPLETED
+    assert state.serialized_input is None
+    assert state.serialized_output is None
+    assert state.serialized_custom_status is "\"foobaz\""