Enhanced entity implementation based on .NET reference

Copilot · berndverst · Copilot · commit 539eb29286c4 · 2025-06-03T22:00:00.000Z
Co-authored-by: berndverst &lt;4535280+berndverst@users.noreply.github.com&gt;
diff --git a/durabletask/__init__.py b/durabletask/__init__.py
@@ -4,8 +4,19 @@
 """Durable Task SDK for Python"""
 
 from durabletask.worker import ConcurrencyOptions
-from durabletask.task import EntityContext, EntityState, EntityQuery, EntityQueryResult
+from durabletask.task import (
+    EntityContext, EntityState, EntityQuery, EntityQueryResult, 
+    EntityInstanceId, EntityOperationFailedException
+)
 
-__all__ = ["ConcurrencyOptions", "EntityContext", "EntityState", "EntityQuery", "EntityQueryResult"]
+__all__ = [
+    "ConcurrencyOptions", 
+    "EntityContext", 
+    "EntityState", 
+    "EntityQuery", 
+    "EntityQueryResult",
+    "EntityInstanceId",
+    "EntityOperationFailedException"
+]
 
 PACKAGE_NAME = "durabletask"
diff --git a/durabletask/client.py b/durabletask/client.py
@@ -223,15 +223,15 @@ def purge_orchestration(self, instance_id: str, recursive: bool = True):
         self._logger.info(f"Purging instance '{instance_id}'.")
         self._stub.PurgeInstances(req)
 
-    def signal_entity(self, entity_id: str, operation_name: str, *,
+    def signal_entity(self, entity_id: Union[str, 'task.EntityInstanceId'], operation_name: str, *,
                       input: Optional[Any] = None,
                       request_id: Optional[str] = None,
                       scheduled_time: Optional[datetime] = None):
         """Signal an entity with an operation.
 
         Parameters
         ----------
-        entity_id : str
+        entity_id : Union[str, task.EntityInstanceId]
             The ID of the entity to signal.
         operation_name : str
             The name of the operation to perform.
@@ -242,22 +242,24 @@ def signal_entity(self, entity_id: str, operation_name: str, *,
         scheduled_time : Optional[datetime]
             The time to schedule the operation. If not provided, the operation is scheduled immediately.
         """
+        entity_id_str = str(entity_id) if hasattr(entity_id, '__str__') else entity_id
+        
         req = pb.SignalEntityRequest(
-            instanceId=entity_id,
+            instanceId=entity_id_str,
             name=operation_name,
             input=wrappers_pb2.StringValue(value=shared.to_json(input)) if input is not None else None,
             requestId=request_id if request_id else uuid.uuid4().hex,
             scheduledTime=helpers.new_timestamp(scheduled_time) if scheduled_time else None)
 
-        self._logger.info(f"Signaling entity '{entity_id}' with operation '{operation_name}'.")
+        self._logger.info(f"Signaling entity '{entity_id_str}' with operation '{operation_name}'.")
         self._stub.SignalEntity(req)
 
-    def get_entity(self, entity_id: str, *, include_state: bool = True) -> Optional[task.EntityState]:
+    def get_entity(self, entity_id: Union[str, 'task.EntityInstanceId'], *, include_state: bool = True) -> Optional[task.EntityState]:
         """Get the state of an entity.
 
         Parameters
         ----------
-        entity_id : str
+        entity_id : Union[str, task.EntityInstanceId]
             The ID of the entity to query.
         include_state : bool
             Whether to include the entity's state in the response.
@@ -267,7 +269,9 @@ def get_entity(self, entity_id: str, *, include_state: bool = True) -> Optional[
         Optional[EntityState]
             The entity state if it exists, None otherwise.
         """
-        req = pb.GetEntityRequest(instanceId=entity_id, includeState=include_state)
+        entity_id_str = str(entity_id) if hasattr(entity_id, '__str__') else entity_id
+        
+        req = pb.GetEntityRequest(instanceId=entity_id_str, includeState=include_state)
         res: pb.GetEntityResponse = self._stub.GetEntity(req)
 
         if not res.exists:
diff --git a/durabletask/task.py b/durabletask/task.py
@@ -5,6 +5,7 @@
 from __future__ import annotations
 
 import math
+import uuid
 from abc import ABC, abstractmethod
 from datetime import datetime, timedelta
 from typing import Any, Callable, Generator, Generic, Optional, TypeVar, Union
@@ -178,13 +179,13 @@ def continue_as_new(self, new_input: Any, *, save_events: bool = False) -> None:
         pass
 
     @abstractmethod
-    def signal_entity(self, entity_id: str, operation_name: str, *,
+    def signal_entity(self, entity_id: Union[str, 'EntityInstanceId'], operation_name: str, *,
                       input: Optional[Any] = None) -> Task:
         """Signal an entity with an operation.
 
         Parameters
         ----------
-        entity_id : str
+        entity_id : Union[str, EntityInstanceId]
             The ID of the entity to signal.
         operation_name : str
             The name of the operation to perform.
@@ -199,14 +200,14 @@ def signal_entity(self, entity_id: str, operation_name: str, *,
         pass
 
     @abstractmethod
-    def call_entity(self, entity_id: str, operation_name: str, *,
+    def call_entity(self, entity_id: Union[str, 'EntityInstanceId'], operation_name: str, *,
                     input: Optional[TInput] = None,
                     retry_policy: Optional[RetryPolicy] = None) -> Task[TOutput]:
         """Call an entity operation and wait for the result.
 
         Parameters
         ----------
-        entity_id : str
+        entity_id : Union[str, EntityInstanceId]
             The ID of the entity to call.
         operation_name : str
             The name of the operation to perform.
@@ -513,12 +514,48 @@ def task_id(self) -> int:
         return self._task_id
 
 
+@dataclass
+class EntityInstanceId:
+    """Represents the ID of a durable entity instance."""
+    name: str
+    key: str
+    
+    def __str__(self) -> str:
+        """Return the string representation in the format: name@key"""
+        return f"{self.name}@{self.key}"
+    
+    @classmethod
+    def from_string(cls, instance_id: str) -> 'EntityInstanceId':
+        """Parse an entity instance ID from string format (name@key)."""
+        if '@' not in instance_id:
+            raise ValueError(f"Invalid entity instance ID format: {instance_id}. Expected format: name@key")
+        
+        parts = instance_id.split('@', 1)
+        if len(parts) != 2 or not parts[0] or not parts[1]:
+            raise ValueError(f"Invalid entity instance ID format: {instance_id}. Expected format: name@key")
+        
+        return cls(name=parts[0], key=parts[1])
+
+
+class EntityOperationFailedException(Exception):
+    """Exception raised when an entity operation fails."""
+    
+    def __init__(self, entity_id: EntityInstanceId, operation_name: str, failure_details: FailureDetails):
+        self.entity_id = entity_id
+        self.operation_name = operation_name
+        self.failure_details = failure_details
+        super().__init__(f"Operation '{operation_name}' on entity '{entity_id}' failed: {failure_details.message}")
+
+
 class EntityContext:
+    """Context for entity operations, providing access to state and scheduling capabilities."""
+    
     def __init__(self, instance_id: str, operation_name: str, is_new_entity: bool = False):
         self._instance_id = instance_id
         self._operation_name = operation_name
         self._is_new_entity = is_new_entity
         self._state: Optional[Any] = None
+        self._entity_instance_id = EntityInstanceId.from_string(instance_id)
 
     @property
     def instance_id(self) -> str:
@@ -531,6 +568,17 @@ def instance_id(self) -> str:
         """
         return self._instance_id
 
+    @property
+    def entity_id(self) -> EntityInstanceId:
+        """Get the structured entity instance ID.
+
+        Returns
+        -------
+        EntityInstanceId
+            The structured entity instance ID.
+        """
+        return self._entity_instance_id
+
     @property
     def operation_name(self) -> str:
         """Get the name of the operation being performed on the entity.
@@ -578,6 +626,64 @@ def set_state(self, state: Any) -> None:
         """
         self._state = state
 
+    def signal_entity(self, entity_id: Union[str, EntityInstanceId], operation_name: str, *,
+                      input: Optional[Any] = None) -> None:
+        """Signal another entity with an operation (fire-and-forget).
+
+        Parameters
+        ----------
+        entity_id : Union[str, EntityInstanceId]
+            The ID of the entity to signal.
+        operation_name : str
+            The name of the operation to perform.
+        input : Optional[Any]
+            The JSON-serializable input to pass to the entity operation.
+        """
+        # Store the signal for later processing during entity execution
+        if not hasattr(self, '_signals'):
+            self._signals = []
+        
+        entity_id_str = str(entity_id) if isinstance(entity_id, EntityInstanceId) else entity_id
+        self._signals.append({
+            'entity_id': entity_id_str,
+            'operation_name': operation_name,
+            'input': input
+        })
+
+    def start_new_orchestration(self, orchestrator: Union[Orchestrator[TInput, TOutput], str], *,
+                               input: Optional[TInput] = None,
+                               instance_id: Optional[str] = None) -> str:
+        """Start a new orchestration from within an entity operation.
+
+        Parameters
+        ----------
+        orchestrator : Union[Orchestrator[TInput, TOutput], str]
+            The orchestrator function or name to start.
+        input : Optional[TInput]
+            The JSON-serializable input to pass to the orchestration.
+        instance_id : Optional[str]
+            The instance ID for the new orchestration. If not provided, a random UUID will be used.
+
+        Returns
+        -------
+        str
+            The instance ID of the new orchestration.
+        """
+        # Store the orchestration start request for later processing
+        if not hasattr(self, '_orchestrations'):
+            self._orchestrations = []
+        
+        orchestrator_name = orchestrator if isinstance(orchestrator, str) else get_name(orchestrator)
+        new_instance_id = instance_id or str(uuid.uuid4())
+        
+        self._orchestrations.append({
+            'name': orchestrator_name,
+            'input': input,
+            'instance_id': new_instance_id
+        })
+        
+        return new_instance_id
+
 
 # Orchestrators are generators that yield tasks and receive/return any type
 Orchestrator = Callable[[OrchestrationContext, TInput], Union[Generator[Task, Any, Any], TOutput]]
diff --git a/durabletask/worker.py b/durabletask/worker.py
@@ -928,12 +928,14 @@ def continue_as_new(self, new_input, *, save_events: bool = False) -> None:
 
         self.set_continued_as_new(new_input, save_events)
 
-    def signal_entity(self, entity_id: str, operation_name: str, *,
+    def signal_entity(self, entity_id: Union[str, task.EntityInstanceId], operation_name: str, *,
                       input: Optional[Any] = None) -> task.Task:
         # Create a signal entity action
+        entity_id_str = str(entity_id) if hasattr(entity_id, '__str__') else entity_id
+        
         action = pb.OrchestratorAction()
         action.sendEntitySignal.CopyFrom(pb.SendSignalAction(
-            instanceId=entity_id,
+            instanceId=entity_id_str,
             name=operation_name,
             input=ph.get_string_value(shared.to_json(input)) if input is not None else None
         ))
@@ -951,7 +953,7 @@ def signal_entity(self, entity_id: str, operation_name: str, *,
 
         return signal_task
 
-    def call_entity(self, entity_id: str, operation_name: str, *,
+    def call_entity(self, entity_id: Union[str, task.EntityInstanceId], operation_name: str, *,
                     input: Optional[Any] = None,
                     retry_policy: Optional[task.RetryPolicy] = None) -> task.Task:
         # For now, entity calls are not directly supported in orchestrations
@@ -1405,6 +1407,28 @@ def execute(self, req: pb.EntityBatchRequest) -> pb.EntityBatchResult:
                 # Update state for next operation
                 current_state = ctx.get_state()
 
+                # Process entity signals from context
+                if hasattr(ctx, '_signals'):
+                    for signal in ctx._signals:
+                        signal_action = pb.OrchestratorAction()
+                        signal_action.sendEntitySignal.CopyFrom(pb.SendSignalAction(
+                            instanceId=signal['entity_id'],
+                            name=signal['operation_name'],
+                            input=ph.get_string_value(shared.to_json(signal['input'])) if signal['input'] is not None else None
+                        ))
+                        actions.append(signal_action)
+
+                # Process orchestration starts from context
+                if hasattr(ctx, '_orchestrations'):
+                    for orch in ctx._orchestrations:
+                        orch_action = pb.OrchestratorAction()
+                        orch_action.callOrchestrator.CopyFrom(pb.CallOrchestratorAction(
+                            name=orch['name'],
+                            instanceId=orch['instance_id'],
+                            input=ph.get_string_value(shared.to_json(orch['input'])) if orch['input'] is not None else None
+                        ))
+                        actions.append(orch_action)
+
                 # Create operation result
                 result = pb.OperationResult()
                 if operation_output is not None:
diff --git a/examples/durable_entities.py b/examples/durable_entities.py
diff --git a/tests/durabletask/test_entities.py b/tests/durabletask/test_entities.py
