Merge branch 'main' into andystaples/add-functions-support

Author: andystaples

durabletask/entities/entity_instance_id.py

@@ -1,6 +1,3 @@
-from typing import Optional
-
-
 class EntityInstanceId:
     def __init__(self, entity: str, key: str):
         self.entity = entity
@@ -30,8 +27,13 @@ def parse(entity_id: str) -> "EntityInstanceId":
 
         Returns
         -------
-        Optional[EntityInstanceId]
-            The parsed EntityInstanceId object, or None if the input is None.
+        EntityInstanceId
+            The parsed EntityInstanceId object.
+
+        Raises
+        ------
+        ValueError
+            If the input string is not in the correct format.
         """
         try:
             _, entity, key = entity_id.split("@", 2)

durabletask/entities/entity_metadata.py

@@ -44,8 +44,9 @@ def __init__(self,
 
     @staticmethod
     def from_entity_response(entity_response: pb.GetEntityResponse, includes_state: bool):
-        entity_id = EntityInstanceId.parse(entity_response.entity.instanceId)
-        if not entity_id:
+        try:
+            entity_id = EntityInstanceId.parse(entity_response.entity.instanceId)
+        except ValueError:
             raise ValueError("Invalid entity instance ID in entity response.")
         entity_state = None
         if includes_state:

durabletask/internal/helpers.py

@@ -20,6 +20,11 @@ def new_orchestrator_started_event(timestamp: Optional[datetime] = None) -> pb.H
     return pb.HistoryEvent(eventId=-1, timestamp=ts, orchestratorStarted=pb.OrchestratorStartedEvent())
 
 
+def new_orchestrator_completed_event() -> pb.HistoryEvent:
+    return pb.HistoryEvent(eventId=-1, timestamp=timestamp_pb2.Timestamp(),
+                           orchestratorCompleted=pb.OrchestratorCompletedEvent())
+
+
 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(
@@ -119,6 +124,18 @@ def new_failure_details(ex: Exception) -> pb.TaskFailureDetails:
     )
 
 
+def new_event_sent_event(event_id: int, instance_id: str, input: str):
+    return pb.HistoryEvent(
+        eventId=event_id,
+        timestamp=timestamp_pb2.Timestamp(),
+        eventSent=pb.EventSentEvent(
+            name="",
+            input=get_string_value(input),
+            instanceId=instance_id
+        )
+    )
+
+
 def new_event_raised_event(name: str, encoded_input: Optional[str] = None) -> pb.HistoryEvent:
     return pb.HistoryEvent(
         eventId=-1,
@@ -199,8 +216,9 @@ def new_schedule_task_action(id: int, name: str, encoded_input: Optional[str],
 def new_call_entity_action(id: int,
                            parent_instance_id: str,
                            entity_id: EntityInstanceId,
-                           operation: str, encoded_input: Optional[str],
-                           request_id: str):
+                           operation: str,
+                           encoded_input: Optional[str],
+                           request_id: str) -> pb.OrchestratorAction:
     return pb.OrchestratorAction(id=id, sendEntityMessage=pb.SendEntityMessageAction(entityOperationCalled=pb.EntityOperationCalledEvent(
         requestId=request_id,
         operation=operation,
@@ -216,7 +234,7 @@ def new_signal_entity_action(id: int,
                              entity_id: EntityInstanceId,
                              operation: str,
                              encoded_input: Optional[str],
-                             request_id: str):
+                             request_id: str) -> pb.OrchestratorAction:
     return pb.OrchestratorAction(id=id, sendEntityMessage=pb.SendEntityMessageAction(entityOperationSignaled=pb.EntityOperationSignaledEvent(
         requestId=request_id,
         operation=operation,

durabletask/internal/proto_task_hub_sidecar_service_stub.py

@@ -2,9 +2,8 @@
 
 
 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.
+    """A stub class matching the TaskHubSidecarServiceStub generated from the .proto file.
+    Allows the use of TaskHubGrpcWorker methods when a real sidecar stub is not available.
     """
     Hello: Callable[..., Any]
     StartInstance: Callable[..., Any]

durabletask/py.typed

@@ -139,7 +139,8 @@ def call_activity(self, activity: Union[Activity[TInput, TOutput], str], *,
         pass
 
     @abstractmethod
-    def call_entity(self, entity: EntityInstanceId,
+    def call_entity(self,
+                    entity: EntityInstanceId,
                     operation: str,
                     input: Optional[TInput] = None) -> Task:
         """Schedule entity function for execution.
@@ -264,8 +265,8 @@ def new_uuid(self) -> str:
 
         The default implementation of this method creates a name-based UUID
         using the algorithm from RFC 4122 §4.3. The name input used to generate
-        this value is a combination of the orchestration instance ID and an
-        internally managed sequence number.
+        this value is a combination of the orchestration instance ID, the current UTC datetime,
+        and an internally managed counter.
 
         Returns
         -------

durabletask/task.py

@@ -12,7 +12,7 @@
 from threading import Event, Thread
 from types import GeneratorType
 from enum import Enum
-from typing import Any, Generator, Optional, Sequence, TypeVar, Union
+from typing import Any, Generator, Optional, Sequence, Tuple, TypeVar, Union
 import uuid
 from packaging.version import InvalidVersion, parse
 
@@ -25,6 +25,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
@@ -680,7 +681,7 @@ def _execute_orchestrator(
     def _cancel_orchestrator(
             self,
             req: pb.OrchestratorRequest,
-            stub: stubs.TaskHubSidecarServiceStub,
+            stub: Union[stubs.TaskHubSidecarServiceStub, ProtoTaskHubSidecarServiceStub],
             completionToken,
     ):
         stub.AbandonTaskOrchestratorWorkItem(
@@ -693,7 +694,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
@@ -726,7 +727,7 @@ def _execute_activity(
     def _cancel_activity(
             self,
             req: pb.ActivityRequest,
-            stub: stubs.TaskHubSidecarServiceStub,
+            stub: Union[stubs.TaskHubSidecarServiceStub, ProtoTaskHubSidecarServiceStub],
             completionToken,
     ):
         stub.AbandonTaskActivityWorkItem(
@@ -754,9 +755,10 @@ def _execute_entity_batch(
         for operation in req.operations:
             start_time = datetime.now(timezone.utc)
             executor = _EntityExecutor(self._registry, self._logger)
-            entity_instance_id = EntityInstanceId.parse(instance_id)
-            if not entity_instance_id:
-                raise RuntimeError(f"Invalid entity instance ID '{operation.requestId}' in entity operation request.")
+            try:
+                entity_instance_id = EntityInstanceId.parse(instance_id)
+            except ValueError:
+                raise RuntimeError(f"Invalid entity instance ID '{instance_id}' in entity operation request.")
 
             operation_result = None
 
@@ -808,7 +810,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(
@@ -831,9 +833,8 @@ def __init__(self, instance_id: str, registry: _Registry):
         self._pending_actions: dict[int, pb.OrchestratorAction] = {}
         self._pending_tasks: dict[int, task.CompletableTask] = {}
         # Maps entity ID to task ID
-        self._entity_task_id_map: dict[str, tuple[EntityInstanceId, int, Optional[str]]] = {}
-        # Maps criticalSectionId to task ID
-        self._entity_lock_id_map: dict[str, int] = {}
+        self._entity_task_id_map: dict[str, tuple[EntityInstanceId, int]] = {}
+        self._entity_lock_task_id_map: dict[str, tuple[EntityInstanceId, int]] = {}
         self._sequence_number = 0
         self._new_uuid_counter = 0
         self._current_utc_datetime = datetime(1000, 1, 1)
@@ -1170,12 +1171,7 @@ def call_entity_function_helper(
             raise RuntimeError(error_message)
 
         encoded_input = shared.to_json(input) if input is not None else None
-        action = ph.new_call_entity_action(id,
-                                           self.instance_id,
-                                           entity_id,
-                                           operation,
-                                           encoded_input,
-                                           self.new_uuid())
+        action = ph.new_call_entity_action(id, self.instance_id, entity_id, operation, encoded_input, self.new_uuid())
         self._pending_actions[id] = action
 
         fn_task = task.CompletableTask()
@@ -1262,14 +1258,14 @@ def continue_as_new(self, new_input, *, save_events: bool = False) -> None:
         self.set_continued_as_new(new_input, save_events)
 
     def new_uuid(self) -> str:
-        URL_NAMESPACE: str = "9e952958-5e33-4daf-827f-2fa12937b875"
+        NAMESPACE_UUID: str = "9e952958-5e33-4daf-827f-2fa12937b875"
 
         uuid_name_value = \
             f"{self._instance_id}" \
             f"_{self.current_utc_datetime.strftime(DATETIME_STRING_FORMAT)}" \
             f"_{self._new_uuid_counter}"
         self._new_uuid_counter += 1
-        namespace_uuid = uuid.uuid5(uuid.NAMESPACE_OID, URL_NAMESPACE)
+        namespace_uuid = uuid.uuid5(uuid.NAMESPACE_OID, NAMESPACE_UUID)
         return str(uuid.uuid5(namespace_uuid, uuid_name_value))
 
 
@@ -1612,32 +1608,11 @@ def process_event(
                     raise TypeError("Unexpected sub-orchestration task type")
             elif event.HasField("eventRaised"):
                 if event.eventRaised.name in ctx._entity_task_id_map:
-                    # This eventRaised represents the result of an entity operation after being translated to the old
-                    # entity protocol by the Durable WebJobs extension
-                    entity_id, task_id, action_type = ctx._entity_task_id_map.get(event.eventRaised.name, (None, None, None))
-                    if entity_id is None:
-                        raise RuntimeError(f"Could not retrieve entity ID for entity-related eventRaised with ID '{event.eventId}'")
-                    if task_id is None:
-                        raise RuntimeError(f"Could not retrieve task ID for entity-related eventRaised with ID '{event.eventId}'")
-                    entity_task = ctx._pending_tasks.pop(task_id, None)
-                    if not entity_task:
-                        raise RuntimeError(f"Could not retrieve entity task for entity-related eventRaised with ID '{event.eventId}'")
-                    result = None
-                    if not ph.is_empty(event.eventRaised.input):
-                        # TODO: Investigate why the event result is wrapped in a dict with "result" key
-                        result = shared.from_json(event.eventRaised.input.value)["result"]
-                    if action_type == "entityOperationCalled":
-                        ctx._entity_context.recover_lock_after_call(entity_id)
-                        entity_task.complete(result)
-                        ctx.resume()
-                    elif action_type == "entityLockRequested":
-                        ctx._entity_context.complete_acquire(event.eventRaised.name)
-                        entity_task.complete(EntityLock(ctx))
-                        ctx.resume()
-                    else:
-                        raise RuntimeError(f"Unknown action type '{action_type}' for entity-related eventRaised "
-                                           f"with ID '{event.eventId}'")
-
+                    entity_id, task_id = ctx._entity_task_id_map.get(event.eventRaised.name, (None, None))
+                    self._handle_entity_event_raised(ctx, event, entity_id, task_id, False)
+                elif event.eventRaised.name in ctx._entity_lock_task_id_map:
+                    entity_id, task_id = ctx._entity_lock_task_id_map.get(event.eventRaised.name, (None, None))
+                    self._handle_entity_event_raised(ctx, event, entity_id, task_id, True)
                 else:
                     # event names are case-insensitive
                     event_name = event.eventRaised.name.casefold()
@@ -1705,8 +1680,9 @@ def process_event(
                     raise _get_wrong_action_type_error(
                         entity_call_id, expected_method_name, action
                     )
-                entity_id = EntityInstanceId.parse(event.entityOperationCalled.targetInstanceId.value)
-                if not entity_id:
+                try:
+                    entity_id = EntityInstanceId.parse(event.entityOperationCalled.targetInstanceId.value)
+                except ValueError:
                     raise RuntimeError(f"Could not parse entity ID from targetInstanceId '{event.entityOperationCalled.targetInstanceId.value}'")
                 ctx._entity_task_id_map[event.entityOperationCalled.requestId] = (entity_id, entity_call_id, None)
             elif event.HasField("entityOperationSignaled"):
@@ -1802,15 +1778,11 @@ def process_event(
                 action = ctx._pending_actions.pop(event.eventId, None)
                 if action and action.HasField("sendEntityMessage"):
                     if action.sendEntityMessage.HasField("entityOperationCalled"):
-                        action_type = "entityOperationCalled"
+                        entity_id, event_id = self._parse_entity_event_sent_input(event)
+                        ctx._entity_task_id_map[event_id] = (entity_id, event.eventId)
                     elif action.sendEntityMessage.HasField("entityLockRequested"):
-                        action_type = "entityLockRequested"
-                    else:
-                        return
-
-                    entity_id = EntityInstanceId.parse(event.eventSent.instanceId)
-                    event_id = json.loads(event.eventSent.input.value)["id"]
-                    ctx._entity_task_id_map[event_id] = (entity_id, event.eventId, action_type)
+                        entity_id, event_id = self._parse_entity_event_sent_input(event)
+                        ctx._entity_lock_task_id_map[event_id] = (entity_id, event.eventId)
             else:
                 eventType = event.WhichOneof("eventType")
                 raise task.OrchestrationStateError(
@@ -1820,6 +1792,44 @@ def process_event(
             # The orchestrator generator function completed
             ctx.set_complete(generatorStopped.value, pb.ORCHESTRATION_STATUS_COMPLETED)
 
+    def _parse_entity_event_sent_input(self, event: pb.HistoryEvent) -> Tuple[EntityInstanceId, str]:
+        try:
+            entity_id = EntityInstanceId.parse(event.eventSent.instanceId)
+        except ValueError:
+            raise RuntimeError(f"Could not parse entity ID from instanceId '{event.eventSent.instanceId}'")
+        try:
+            event_id = json.loads(event.eventSent.input.value)["id"]
+        except (json.JSONDecodeError, KeyError, TypeError) as ex:
+            raise RuntimeError(f"Could not parse event ID from eventSent input '{event.eventSent.input.value}'") from ex
+        return entity_id, event_id
+
+    def _handle_entity_event_raised(self,
+                                    ctx: _RuntimeOrchestrationContext,
+                                    event: pb.HistoryEvent,
+                                    entity_id: Optional[EntityInstanceId],
+                                    task_id: Optional[int],
+                                    is_lock_event: bool):
+        # This eventRaised represents the result of an entity operation after being translated to the old
+        # entity protocol by the Durable WebJobs extension
+        if entity_id is None:
+            raise RuntimeError(f"Could not retrieve entity ID for entity-related eventRaised with ID '{event.eventId}'")
+        if task_id is None:
+            raise RuntimeError(f"Could not retrieve task ID for entity-related eventRaised with ID '{event.eventId}'")
+        entity_task = ctx._pending_tasks.pop(task_id, None)
+        if not entity_task:
+            raise RuntimeError(f"Could not retrieve entity task for entity-related eventRaised with ID '{event.eventId}'")
+        result = None
+        if not ph.is_empty(event.eventRaised.input):
+            # TODO: Investigate why the event result is wrapped in a dict with "result" key
+            result = shared.from_json(event.eventRaised.input.value)["result"]
+        if is_lock_event:
+            ctx._entity_context.complete_acquire(event.eventRaised.name)
+            entity_task.complete(EntityLock(ctx))
+        else:
+            ctx._entity_context.recover_lock_after_call(entity_id)
+            entity_task.complete(result)
+        ctx.resume()
+
     def evaluate_orchestration_versioning(self, versioning: Optional[VersioningOptions], orchestration_version: Optional[str]) -> Optional[pb.TaskFailureDetails]:
         if versioning is None:
             return None

durabletask/worker.py

@@ -5,6 +5,7 @@
 import os
 import threading
 from datetime import timedelta
+import uuid
 
 import pytest
 
@@ -532,3 +533,39 @@ def empty_orchestrator(ctx: task.OrchestrationContext, _):
     assert state.serialized_input is None
     assert state.serialized_output is None
     assert state.serialized_custom_status == "\"foobaz\""
+
+
+def test_new_uuid():
+    def noop(_: task.ActivityContext, _1):
+        pass
+
+    def empty_orchestrator(ctx: task.OrchestrationContext, _):
+        # Assert that two new_uuid calls return different values
+        results = [ctx.new_uuid(), ctx.new_uuid()]
+        yield ctx.call_activity("noop")
+        # Assert that new_uuid still returns a unique value after replay
+        results.append(ctx.new_uuid())
+        return results
+
+    # Start a worker, which will connect to the sidecar in a background thread
+    with DurableTaskSchedulerWorker(host_address=endpoint, secure_channel=True,
+                                    taskhub=taskhub_name, token_credential=None) as w:
+        w.add_orchestrator(empty_orchestrator)
+        w.add_activity(noop)
+        w.start()
+
+        c = DurableTaskSchedulerClient(host_address=endpoint, secure_channel=True,
+                                       taskhub=taskhub_name, token_credential=None)
+        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
+    results = json.loads(state.serialized_output or "\"\"")
+    assert isinstance(results, list) and len(results) == 3
+    assert uuid.UUID(results[0]) != uuid.UUID(results[1])
+    assert uuid.UUID(results[0]) != uuid.UUID(results[2])
+    assert uuid.UUID(results[1]) != uuid.UUID(results[2])