Add entity types, context, and client methods for durable entities

Co-authored-by: berndverst <[email protected]>

Author: Copilot

durabletask/__init__.py

@@ -4,7 +4,8 @@
 """Durable Task SDK for Python"""
 
 from durabletask.worker import ConcurrencyOptions
+from durabletask.task import EntityContext, EntityState, EntityQuery, EntityQueryResult
 
-__all__ = ["ConcurrencyOptions"]
+__all__ = ["ConcurrencyOptions", "EntityContext", "EntityState", "EntityQuery", "EntityQueryResult"]
 
 PACKAGE_NAME = "durabletask"

durabletask/client.py

@@ -222,3 +222,141 @@ def purge_orchestration(self, instance_id: str, recursive: bool = True):
         req = pb.PurgeInstancesRequest(instanceId=instance_id, recursive=recursive)
         self._logger.info(f"Purging instance '{instance_id}'.")
         self._stub.PurgeInstances(req)
+
+    def signal_entity(self, entity_id: str, 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
+            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.
+        request_id : Optional[str]
+            A unique request ID for the operation. If not provided, a random UUID will be used.
+        scheduled_time : Optional[datetime]
+            The time to schedule the operation. If not provided, the operation is scheduled immediately.
+        """
+        req = pb.SignalEntityRequest(
+            instanceId=entity_id,
+            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._stub.SignalEntity(req)
+
+    def get_entity(self, entity_id: str, *, include_state: bool = True) -> Optional[task.EntityState]:
+        """Get the state of an entity.
+
+        Parameters
+        ----------
+        entity_id : str
+            The ID of the entity to query.
+        include_state : bool
+            Whether to include the entity's state in the response.
+
+        Returns
+        -------
+        Optional[EntityState]
+            The entity state if it exists, None otherwise.
+        """
+        req = pb.GetEntityRequest(instanceId=entity_id, includeState=include_state)
+        res: pb.GetEntityResponse = self._stub.GetEntity(req)
+        
+        if not res.exists:
+            return None
+
+        entity_metadata = res.entity
+        return task.EntityState(
+            instance_id=entity_metadata.instanceId,
+            last_modified_time=entity_metadata.lastModifiedTime.ToDatetime(),
+            backlog_queue_size=entity_metadata.backlogQueueSize,
+            locked_by=entity_metadata.lockedBy.value if not helpers.is_empty(entity_metadata.lockedBy) else None,
+            serialized_state=entity_metadata.serializedState.value if not helpers.is_empty(entity_metadata.serializedState) else None)
+
+    def query_entities(self, query: task.EntityQuery) -> task.EntityQueryResult:
+        """Query entities based on the provided criteria.
+
+        Parameters
+        ----------
+        query : EntityQuery
+            The query criteria for entities.
+
+        Returns
+        -------
+        EntityQueryResult
+            The query result containing matching entities and continuation token.
+        """
+        # Build the protobuf query
+        pb_query = pb.EntityQuery(
+            includeState=query.include_state,
+            includeTransient=query.include_transient)
+
+        if query.instance_id_starts_with is not None:
+            pb_query.instanceIdStartsWith = wrappers_pb2.StringValue(value=query.instance_id_starts_with)
+        if query.last_modified_from is not None:
+            pb_query.lastModifiedFrom = helpers.new_timestamp(query.last_modified_from)
+        if query.last_modified_to is not None:
+            pb_query.lastModifiedTo = helpers.new_timestamp(query.last_modified_to)
+        if query.page_size is not None:
+            pb_query.pageSize = wrappers_pb2.Int32Value(value=query.page_size)
+        if query.continuation_token is not None:
+            pb_query.continuationToken = wrappers_pb2.StringValue(value=query.continuation_token)
+
+        req = pb.QueryEntitiesRequest(query=pb_query)
+        res: pb.QueryEntitiesResponse = self._stub.QueryEntities(req)
+
+        # Convert response to Python objects
+        entities = []
+        for entity_metadata in res.entities:
+            entities.append(task.EntityState(
+                instance_id=entity_metadata.instanceId,
+                last_modified_time=entity_metadata.lastModifiedTime.ToDatetime(),
+                backlog_queue_size=entity_metadata.backlogQueueSize,
+                locked_by=entity_metadata.lockedBy.value if not helpers.is_empty(entity_metadata.lockedBy) else None,
+                serialized_state=entity_metadata.serializedState.value if not helpers.is_empty(entity_metadata.serializedState) else None))
+
+        return task.EntityQueryResult(
+            entities=entities,
+            continuation_token=res.continuationToken.value if not helpers.is_empty(res.continuationToken) else None)
+
+    def clean_entity_storage(self, *,
+                            remove_empty_entities: bool = True,
+                            release_orphaned_locks: bool = True,
+                            continuation_token: Optional[str] = None) -> tuple[int, int, Optional[str]]:
+        """Clean up entity storage by removing empty entities and releasing orphaned locks.
+
+        Parameters
+        ----------
+        remove_empty_entities : bool
+            Whether to remove entities that have no state.
+        release_orphaned_locks : bool
+            Whether to release locks that are no longer held by active orchestrations.
+        continuation_token : Optional[str]
+            A continuation token from a previous cleanup operation.
+
+        Returns
+        -------
+        tuple[int, int, Optional[str]]
+            A tuple containing (empty_entities_removed, orphaned_locks_released, continuation_token).
+        """
+        req = pb.CleanEntityStorageRequest(
+            removeEmptyEntities=remove_empty_entities,
+            releaseOrphanedLocks=release_orphaned_locks)
+
+        if continuation_token is not None:
+            req.continuationToken = wrappers_pb2.StringValue(value=continuation_token)
+
+        self._logger.info("Cleaning entity storage.")
+        res: pb.CleanEntityStorageResponse = self._stub.CleanEntityStorage(req)
+
+        return (res.emptyEntitiesRemoved, 
+                res.orphanedLocksReleased,
+                res.continuationToken.value if not helpers.is_empty(res.continuationToken) else None)

durabletask/task.py

@@ -8,6 +8,7 @@
 from abc import ABC, abstractmethod
 from datetime import datetime, timedelta
 from typing import Any, Callable, Generator, Generic, Optional, TypeVar, Union
+from dataclasses import dataclass
 
 import durabletask.internal.helpers as pbh
 import durabletask.internal.orchestrator_service_pb2 as pb
@@ -176,6 +177,51 @@ 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, *,
+                      input: Optional[Any] = None) -> Task:
+        """Signal an entity with an operation.
+
+        Parameters
+        ----------
+        entity_id : str
+            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.
+
+        Returns
+        -------
+        Task
+            A Durable Task that completes when the entity operation is scheduled.
+        """
+        pass
+
+    @abstractmethod
+    def call_entity(self, entity_id: str, 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
+            The ID of the entity to call.
+        operation_name : str
+            The name of the operation to perform.
+        input : Optional[TInput]
+            The JSON-serializable input to pass to the entity operation.
+        retry_policy : Optional[RetryPolicy]
+            The retry policy to use for this entity call.
+
+        Returns
+        -------
+        Task[TOutput]
+            A Durable Task that completes when the entity operation completes or fails.
+        """
+        pass
+
 
 class FailureDetails:
     def __init__(self, message: str, error_type: str, stack_trace: Optional[str]):
@@ -219,6 +265,40 @@ class OrchestrationStateError(Exception):
     pass
 
 
+@dataclass
+class EntityState:
+    """Represents the state of a durable entity."""
+    instance_id: str
+    last_modified_time: datetime
+    backlog_queue_size: int
+    locked_by: Optional[str]
+    serialized_state: Optional[str]
+
+    @property
+    def exists(self) -> bool:
+        """Returns True if the entity exists (has been created), False otherwise."""
+        return self.serialized_state is not None
+
+
+@dataclass
+class EntityQuery:
+    """Represents a query for durable entities."""
+    instance_id_starts_with: Optional[str] = None
+    last_modified_from: Optional[datetime] = None
+    last_modified_to: Optional[datetime] = None
+    include_state: bool = False
+    include_transient: bool = False
+    page_size: Optional[int] = None
+    continuation_token: Optional[str] = None
+
+
+@dataclass
+class EntityQueryResult:
+    """Represents the result of an entity query."""
+    entities: list[EntityState]
+    continuation_token: Optional[str] = None
+
+
 class Task(ABC, Generic[T]):
     """Abstract base class for asynchronous tasks in a durable orchestration."""
     _result: T
@@ -433,12 +513,81 @@ def task_id(self) -> int:
         return self._task_id
 
 
+class EntityContext:
+    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
+
+    @property
+    def instance_id(self) -> str:
+        """Get the ID of the entity instance.
+
+        Returns
+        -------
+        str
+            The ID of the current entity instance.
+        """
+        return self._instance_id
+
+    @property
+    def operation_name(self) -> str:
+        """Get the name of the operation being performed on the entity.
+
+        Returns
+        -------
+        str
+            The name of the operation.
+        """
+        return self._operation_name
+
+    @property
+    def is_new_entity(self) -> bool:
+        """Get a value indicating whether this is a newly created entity.
+
+        Returns
+        -------
+        bool
+            True if this is the first operation on this entity, False otherwise.
+        """
+        return self._is_new_entity
+
+    def get_state(self, state_type: type[T] = None) -> Optional[T]:
+        """Get the current state of the entity.
+
+        Parameters
+        ----------
+        state_type : type[T], optional
+            The type to deserialize the state to. If not provided, returns the raw state.
+
+        Returns
+        -------
+        Optional[T]
+            The current state of the entity, or None if the entity has no state.
+        """
+        return self._state
+
+    def set_state(self, state: Any) -> None:
+        """Set the current state of the entity.
+
+        Parameters
+        ----------
+        state : Any
+            The new state for the entity. Must be JSON-serializable.
+        """
+        self._state = state
+
+
 # Orchestrators are generators that yield tasks and receive/return any type
 Orchestrator = Callable[[OrchestrationContext, TInput], Union[Generator[Task, Any, Any], TOutput]]
 
 # Activities are simple functions that can be scheduled by orchestrators
 Activity = Callable[[ActivityContext, TInput], TOutput]
 
+# Entities are stateful objects that can receive operations and maintain state
+Entity = Callable[['EntityContext', TInput], TOutput]
+
 
 class RetryPolicy:
     """Represents the retry policy for an orchestration or activity function."""