Activity support, fixes, and stronger typing (#4)

Author: cgillum

durabletask/api/state.py

@@ -2,6 +2,7 @@
 from datetime import datetime
 
 import durabletask.protos.orchestrator_service_pb2 as pb
+import durabletask.protos.helpers as helpers
 from durabletask.protos.orchestrator_service_pb2 import TaskFailureDetails
 
 
@@ -29,7 +30,7 @@ def new_orchestration_state(instance_id: str, res: pb.GetInstanceResponse) -> Or
         state.orchestrationStatus,
         state.createdTimestamp.ToDatetime(),
         state.lastUpdatedTimestamp.ToDatetime(),
-        state.input.value if state.input is not None else None,
-        state.output.value if state.output is not None else None,
-        state.customStatus.value if state.customStatus is not None else None,
+        state.input.value if not helpers.is_empty(state.input) else None,
+        state.output.value if not helpers.is_empty(state.output) else None,
+        state.customStatus.value if not helpers.is_empty(state.customStatus) else None,
         state.failureDetails if state.failureDetails.errorMessage != '' or state.failureDetails.errorType != '' else None)

durabletask/api/task_hub_client.py

@@ -1,17 +1,23 @@
 from dataclasses import dataclass
 from datetime import datetime
 import logging
-from typing import Any
+from typing import Any, TypeVar
 import uuid
 import grpc
 from durabletask.api.state import OrchestrationState, new_orchestration_state
+import durabletask.protos.helpers as helpers
 import durabletask.protos.orchestrator_service_pb2 as pb
 import durabletask.internal.shared as shared
 import simplejson as json
 
-from google.protobuf import timestamp_pb2, wrappers_pb2
+from google.protobuf import wrappers_pb2
 
+import durabletask.task.registry as registry
 from durabletask.protos.orchestrator_service_pb2_grpc import TaskHubSidecarServiceStub
+from durabletask.task.orchestration import Orchestrator
+
+TInput = TypeVar('TInput')
+TOutput = TypeVar('TOutput')
 
 
 class TaskHubGrpcClient:
@@ -24,22 +30,19 @@ def __init__(self, *,
         self._stub = TaskHubSidecarServiceStub(channel)
         self._logger = shared.get_logger(log_handler, log_formatter)
 
-    def schedule_new_orchestration(self, name: str, *,
-                                   input: Any = None,
+    def schedule_new_orchestration(self, orchestrator: Orchestrator[TInput, TOutput], *,
+                                   input: TInput | None = None,
                                    instance_id: str | None = None,
                                    start_at: datetime | None = None) -> str:
-        req = pb.CreateInstanceRequest(name=name)
-        if instance_id is None:
-            instance_id = uuid.uuid4().hex
-        req.instanceId = instance_id
-
-        if input is not None:
-            json_input = json.dumps(input)
-            req.input = wrappers_pb2.StringValue(value=json_input)
-
-        if start_at is not None:
-            req.scheduledStartTimestamp = timestamp_pb2.Timestamp()
-            req.scheduledStartTimestamp.FromDatetime(start_at)
+
+        name = registry.get_name(orchestrator)
+
+        req = pb.CreateInstanceRequest(
+            name=name,
+            instanceId=instance_id if instance_id else uuid.uuid4().hex,
+            input=wrappers_pb2.StringValue(value=json.dumps(input)) if input else None,
+            scheduledStartTimestamp=helpers.new_timestamp(start_at) if start_at else None)
+
         self._logger.info(f"Starting new '{name}' instance with ID = '{instance_id}'.")
         res: pb.CreateInstanceResponse = self._stub.StartInstance(req)
         return res.instanceId

durabletask/internal/shared.py

@@ -14,7 +14,7 @@ def get_grpc_channel(host_address: str | None) -> grpc.Channel:
     return channel
 
 
-def get_logger(log_handler: logging.Handler | None, log_formatter: logging.Formatter | None) -> logging.Logger:
+def get_logger(log_handler: logging.Handler | None = None, log_formatter: logging.Formatter | None = None) -> logging.Logger:
     logger = logging.Logger("durabletask")
 
     # Add a default log handler if none is provided

durabletask/protos/helpers.py

@@ -1,10 +1,14 @@
+import simplejson as json
 import traceback
 
 from datetime import datetime
+from typing import Any
 from google.protobuf import timestamp_pb2, wrappers_pb2
 
 from durabletask.protos.orchestrator_service_pb2 import *
 
+# TODO: The new_xxx_event methods are only used by test code and should be moved elsewhere
+
 
 def new_orchestrator_started_event(timestamp: datetime | None = None) -> HistoryEvent:
     ts = timestamp_pb2.Timestamp()
@@ -13,19 +17,15 @@ def new_orchestrator_started_event(timestamp: datetime | None = None) -> History
     return HistoryEvent(eventId=-1, timestamp=ts, orchestratorStarted=OrchestratorStartedEvent())
 
 
-def new_execution_started_event(name: str, instance_id: str, input: str | None = None) -> HistoryEvent:
-    input_: wrappers_pb2.StringValue | None = None
-    if input is not None:
-        input_ = wrappers_pb2.StringValue(value=input)
-
+def new_execution_started_event(name: str, instance_id: str, encoded_input: str | None = None) -> HistoryEvent:
     return HistoryEvent(
         eventId=-1,
         timestamp=timestamp_pb2.Timestamp(),
         executionStarted=ExecutionStartedEvent(
-            name=name, input=input_, orchestrationInstance=OrchestrationInstance(instanceId=instance_id)))
+            name=name, input=get_string_value(encoded_input), orchestrationInstance=OrchestrationInstance(instanceId=instance_id)))
 
 
-def new_timer_created_event(timer_id: int, fire_at: datetime):
+def new_timer_created_event(timer_id: int, fire_at: datetime) -> HistoryEvent:
     ts = timestamp_pb2.Timestamp()
     ts.FromDatetime(fire_at)
     return HistoryEvent(
@@ -35,7 +35,7 @@ def new_timer_created_event(timer_id: int, fire_at: datetime):
     )
 
 
-def new_timer_fired_event(timer_id: int, fire_at: datetime):
+def new_timer_fired_event(timer_id: int, fire_at: datetime) -> HistoryEvent:
     ts = timestamp_pb2.Timestamp()
     ts.FromDatetime(fire_at)
     return HistoryEvent(
@@ -45,6 +45,30 @@ def new_timer_fired_event(timer_id: int, fire_at: datetime):
     )
 
 
+def new_task_scheduled_event(event_id: int, name: str, encoded_input: str | None = None) -> HistoryEvent:
+    return HistoryEvent(
+        eventId=event_id,
+        timestamp=timestamp_pb2.Timestamp(),
+        taskScheduled=TaskScheduledEvent(name=name, input=get_string_value(encoded_input))
+    )
+
+
+def new_task_completed_event(event_id: int, encoded_output: str | None = None) -> HistoryEvent:
+    return HistoryEvent(
+        eventId=-1,
+        timestamp=timestamp_pb2.Timestamp(),
+        taskCompleted=TaskCompletedEvent(taskScheduledId=event_id, result=get_string_value(encoded_output))
+    )
+
+
+def new_task_failed_event(event_id: int, ex: Exception) -> HistoryEvent:
+    return HistoryEvent(
+        eventId=-1,
+        timestamp=timestamp_pb2.Timestamp(),
+        taskFailed=TaskFailedEvent(taskScheduledId=event_id, failureDetails=new_failure_details(ex))
+    )
+
+
 def new_failure_details(ex: Exception) -> TaskFailureDetails:
     return TaskFailureDetails(
         errorType=type(ex).__name__,
@@ -53,19 +77,22 @@ def new_failure_details(ex: Exception) -> TaskFailureDetails:
     )
 
 
+def get_string_value(val: str | None) -> wrappers_pb2.StringValue | None:
+    if val is None:
+        return None
+    else:
+        return wrappers_pb2.StringValue(value=val)
+
+
 def new_complete_orchestration_action(
         id: int,
         status: OrchestrationStatus,
         result: str | None = None,
         failure_details: TaskFailureDetails | None = None) -> OrchestratorAction:
 
-    result_pb: wrappers_pb2.StringValue | None = None
-    if result is not None:
-        result_pb = wrappers_pb2.StringValue(value=result)
-
     completeOrchestrationAction = CompleteOrchestrationAction(
         orchestrationStatus=status,
-        result=result_pb,
+        result=get_string_value(result),
         failureDetails=failure_details)
 
     # TODO: CarryoverEvents
@@ -77,3 +104,21 @@ def new_create_timer_action(id: int, fire_at: datetime) -> OrchestratorAction:
     timestamp = timestamp_pb2.Timestamp()
     timestamp.FromDatetime(fire_at)
     return OrchestratorAction(id=id, createTimer=CreateTimerAction(fireAt=timestamp))
+
+
+def new_schedule_task_action(id: int, name: str, input: Any) -> OrchestratorAction:
+    encoded_input = json.dumps(input) if input is not None else None
+    return OrchestratorAction(id=id, scheduleTask=ScheduleTaskAction(
+        name=name,
+        input=get_string_value(encoded_input)
+    ))
+
+
+def new_timestamp(dt: datetime) -> timestamp_pb2.Timestamp:
+    ts = timestamp_pb2.Timestamp()
+    ts.FromDatetime(dt)
+    return ts
+
+
+def is_empty(v: wrappers_pb2.StringValue):
+    return v is None or v.value == ''

durabletask/task/activities.py

@@ -0,0 +1,42 @@
+
+from typing import Callable, TypeVar
+
+
+TInput = TypeVar('TInput')
+TOutput = TypeVar('TOutput')
+
+
+class ActivityContext:
+    def __init__(self, orchestration_id: str, task_id: int):
+        self._orchestration_id = orchestration_id
+        self._task_id = task_id
+
+    @property
+    def orchestration_id(self) -> str:
+        """Get the ID of the orchestration instance that scheduled this activity.
+
+        Returns
+        -------
+        str
+            The ID of the current orchestration instance.
+        """
+        return self._orchestration_id
+
+    @property
+    def task_id(self) -> int:
+        """Get the task ID associated with this activity invocation.
+
+        The task ID is an auto-incrementing integer that is unique within
+        the scope of the orchestration instance. It can be used to distinguish
+        between multiple activity invocations that are part of the same
+        orchestration instance.
+
+        Returns
+        -------
+        str
+            The ID of the current orchestration instance.
+        """
+        return self._task_id
+
+
+Activity = Callable[[ActivityContext, TInput], TOutput]