Allow passing custom stub to execute operations

Author: andystaples

durabletask/internal/proto_task_hub_sidecar_service_stub.py

@@ -0,0 +1,38 @@
+from typing import Callable
+
+
+class ProtoTaskHubSidecarServiceStub(object):
+    """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.
+    """
+
+    def __init__(self):
+        """Constructor.
+        """
+        self.Hello: Callable[..., None]
+        self.StartInstance: Callable[..., None]
+        self.GetInstance: Callable[..., None]
+        self.RewindInstance: Callable[..., None]
+        self.WaitForInstanceStart: Callable[..., None]
+        self.WaitForInstanceCompletion: Callable[..., None]
+        self.RaiseEvent: Callable[..., None]
+        self.TerminateInstance: Callable[..., None]
+        self.SuspendInstance: Callable[..., None]
+        self.ResumeInstance: Callable[..., None]
+        self.QueryInstances: Callable[..., None]
+        self.PurgeInstances: Callable[..., None]
+        self.GetWorkItems: Callable[..., None]
+        self.CompleteActivityTask: Callable[..., None]
+        self.CompleteOrchestratorTask: Callable[..., None]
+        self.CompleteEntityTask: Callable[..., None]
+        self.StreamInstanceHistory: Callable[..., None]
+        self.CreateTaskHub: Callable[..., None]
+        self.DeleteTaskHub: Callable[..., None]
+        self.SignalEntity: Callable[..., None]
+        self.GetEntity: Callable[..., None]
+        self.QueryEntities: Callable[..., None]
+        self.CleanEntityStorage: Callable[..., None]
+        self.AbandonTaskActivityWorkItem: Callable[..., None]
+        self.AbandonTaskOrchestratorWorkItem: Callable[..., None]
+        self.AbandonTaskEntityWorkItem: Callable[..., None]

durabletask/worker.py

@@ -23,6 +23,7 @@
 from durabletask.internal.helpers import new_timestamp
 from durabletask.entities import DurableEntity, EntityLock, EntityInstanceId, EntityContext
 from durabletask.internal.orchestration_entity_context import OrchestrationEntityContext
+from durabletask.internal.proto_task_hub_sidecar_service_stub import ProtoTaskHubSidecarServiceStub
 import durabletask.internal.helpers as ph
 import durabletask.internal.exceptions as pe
 import durabletask.internal.orchestrator_service_pb2 as pb
@@ -629,7 +630,7 @@ def stop(self):
     def _execute_orchestrator(
             self,
             req: pb.OrchestratorRequest,
-            stub: stubs.TaskHubSidecarServiceStub,
+            stub: Union[stubs.TaskHubSidecarServiceStub, ProtoTaskHubSidecarServiceStub],
             completionToken,
     ):
         try:
@@ -677,7 +678,7 @@ def _execute_orchestrator(
     def _cancel_orchestrator(
             self,
             req: pb.OrchestratorRequest,
-            stub: stubs.TaskHubSidecarServiceStub,
+            stub: Union[stubs.TaskHubSidecarServiceStub, ProtoTaskHubSidecarServiceStub],
             completionToken,
     ):
         stub.AbandonTaskOrchestratorWorkItem(
@@ -690,7 +691,7 @@ def _cancel_orchestrator(
     def _execute_activity(
             self,
             req: pb.ActivityRequest,
-            stub: stubs.TaskHubSidecarServiceStub,
+            stub: Union[stubs.TaskHubSidecarServiceStub, ProtoTaskHubSidecarServiceStub],
             completionToken,
     ):
         instance_id = req.orchestrationInstance.instanceId
@@ -723,7 +724,7 @@ def _execute_activity(
     def _cancel_activity(
             self,
             req: pb.ActivityRequest,
-            stub: stubs.TaskHubSidecarServiceStub,
+            stub: Union[stubs.TaskHubSidecarServiceStub, ProtoTaskHubSidecarServiceStub],
             completionToken,
     ):
         stub.AbandonTaskActivityWorkItem(
@@ -736,7 +737,7 @@ def _cancel_activity(
     def _execute_entity_batch(
             self,
             req: Union[pb.EntityBatchRequest, pb.EntityRequest],
-            stub: stubs.TaskHubSidecarServiceStub,
+            stub: Union[stubs.TaskHubSidecarServiceStub, ProtoTaskHubSidecarServiceStub],
             completionToken,
     ):
         if isinstance(req, pb.EntityRequest):
@@ -804,7 +805,7 @@ def _execute_entity_batch(
     def _cancel_entity_batch(
             self,
             req: Union[pb.EntityBatchRequest, pb.EntityRequest],
-            stub: stubs.TaskHubSidecarServiceStub,
+            stub: Union[stubs.TaskHubSidecarServiceStub, ProtoTaskHubSidecarServiceStub],
             completionToken,
     ):
         stub.AbandonTaskEntityWorkItem(