durabletask 0.1.0a1__py3-none-any.whl → 1.0.0__py3-none-any.whl

durabletask/internal/shared.py

@@ -5,30 +5,68 @@ import dataclasses
 import json
 import logging
 from types import SimpleNamespace
-from typing import Any, Dict, Union
+from typing import Any, Optional, Sequence, Union
 
 import grpc
 
+ClientInterceptor = Union[
+    grpc.UnaryUnaryClientInterceptor,
+    grpc.UnaryStreamClientInterceptor,
+    grpc.StreamUnaryClientInterceptor,
+    grpc.StreamStreamClientInterceptor
+]
+
 # Field name used to indicate that an object was automatically serialized
 # and should be deserialized as a SimpleNamespace
 AUTO_SERIALIZED = "__durabletask_autoobject__"
 
+SECURE_PROTOCOLS = ["https://", "grpcs://"]
+INSECURE_PROTOCOLS = ["http://", "grpc://"]
+
 
 def get_default_host_address() -> str:
     return "localhost:4001"
 
 
-def get_grpc_channel(host_address: Union[str, None]) -> grpc.Channel:
+def get_grpc_channel(
+        host_address: Optional[str],
+        secure_channel: bool = False,
+        interceptors: Optional[Sequence[ClientInterceptor]] = None) -> grpc.Channel:
+
     if host_address is None:
         host_address = get_default_host_address()
-    channel = grpc.insecure_channel(host_address)
+
+    for protocol in SECURE_PROTOCOLS:
+        if host_address.lower().startswith(protocol):
+            secure_channel = True
+            # remove the protocol from the host name
+            host_address = host_address[len(protocol):]
+            break
+
+    for protocol in INSECURE_PROTOCOLS:
+        if host_address.lower().startswith(protocol):
+            secure_channel = False
+            # remove the protocol from the host name
+            host_address = host_address[len(protocol):]
+            break
+
+    # Create the base channel
+    if secure_channel:
+        channel = grpc.secure_channel(host_address, grpc.ssl_channel_credentials())
+    else:
+        channel = grpc.insecure_channel(host_address)
+
+    # Apply interceptors ONLY if they exist
+    if interceptors:
+        channel = grpc.intercept_channel(channel, *interceptors)
     return channel
 
 
 def get_logger(
-        log_handler: Union[logging.Handler, None] = None,
-        log_formatter: Union[logging.Formatter, None] = None) -> logging.Logger:
-    logger = logging.Logger("durabletask")
+        name_suffix: str,
+        log_handler: Optional[logging.Handler] = None,
+        log_formatter: Optional[logging.Formatter] = None) -> logging.Logger:
+    logger = logging.Logger(f"durabletask-{name_suffix}")
 
     # Add a default log handler if none is provided
     if log_handler is None:
@@ -68,7 +106,7 @@ class InternalJSONEncoder(json.JSONEncoder):
         if dataclasses.is_dataclass(obj):
             # Dataclasses are not serializable by default, so we convert them to a dict and mark them for
             # automatic deserialization by the receiver
-            d = dataclasses.asdict(obj)
+            d = dataclasses.asdict(obj)  # type: ignore
             d[AUTO_SERIALIZED] = True
             return d
         elif isinstance(obj, SimpleNamespace):
@@ -84,7 +122,7 @@ class InternalJSONDecoder(json.JSONDecoder):
     def __init__(self, *args, **kwargs):
         super().__init__(object_hook=self.dict_to_object, *args, **kwargs)
 
-    def dict_to_object(self, d: Dict[str, Any]):
+    def dict_to_object(self, d: dict[str, Any]):
         # If the object was serialized by the InternalJSONEncoder, deserialize it as a SimpleNamespace
         if d.pop(AUTO_SERIALIZED, False):
             return SimpleNamespace(**d)

durabletask/task.py

@@ -4,10 +4,12 @@
 # See https://peps.python.org/pep-0563/
 from __future__ import annotations
 
+import math
 from abc import ABC, abstractmethod
 from datetime import datetime, timedelta
-from typing import Any, Callable, Generator, Generic, List, TypeVar, Union
+from typing import Any, Callable, Generator, Generic, Optional, TypeVar, Union
 
+from durabletask.entities import DurableEntity, EntityInstanceId, EntityLock, EntityContext
 import durabletask.internal.helpers as pbh
 import durabletask.internal.orchestrator_service_pb2 as pb
 
@@ -34,6 +36,21 @@ class OrchestrationContext(ABC):
         """
         pass
 
+    @property
+    @abstractmethod
+    def version(self) -> Optional[str]:
+        """Get the version of the orchestration instance.
+
+        This version is set when the orchestration is scheduled and can be used
+        to determine which version of the orchestrator function is being executed.
+
+        Returns
+        -------
+        Optional[str]
+            The version of the orchestration instance, or None if not set.
+        """
+        pass
+
     @property
     @abstractmethod
     def current_utc_datetime(self) -> datetime:
@@ -69,6 +86,17 @@ class OrchestrationContext(ABC):
         """
         pass
 
+    @abstractmethod
+    def set_custom_status(self, custom_status: Any) -> None:
+        """Set the orchestration instance's custom status.
+
+        Parameters
+        ----------
+        custom_status: Any
+            A JSON-serializable custom status value to set.
+        """
+        pass
+
     @abstractmethod
     def create_timer(self, fire_at: Union[datetime, timedelta]) -> Task:
         """Create a Timer Task to fire after at the specified deadline.
@@ -87,17 +115,21 @@ class OrchestrationContext(ABC):
 
     @abstractmethod
     def call_activity(self, activity: Union[Activity[TInput, TOutput], str], *,
-                      input: Union[TInput, None] = None) -> Task[TOutput]:
+                      input: Optional[TInput] = None,
+                      retry_policy: Optional[RetryPolicy] = None,
+                      tags: Optional[dict[str, str]] = None) -> Task[TOutput]:
         """Schedule an activity for execution.
 
         Parameters
         ----------
         activity: Union[Activity[TInput, TOutput], str]
             A reference to the activity function to call.
-        input: Union[TInput, None]
+        input: Optional[TInput]
             The JSON-serializable input (or None) to pass to the activity.
-        return_type: task.Task[TOutput]
-            The JSON-serializable output type to expect from the activity result.
+        retry_policy: Optional[RetryPolicy]
+            The retry policy to use for this activity call.
+        tags: Optional[dict[str, str]]
+            Optional tags to associate with the activity invocation.
 
         Returns
         -------
@@ -107,20 +139,86 @@ class OrchestrationContext(ABC):
         pass
 
     @abstractmethod
-    def call_sub_orchestrator(self, orchestrator: Orchestrator[TInput, TOutput], *,
-                              input: Union[TInput, None] = None,
-                              instance_id: Union[str, None] = None) -> Task[TOutput]:
+    def call_entity(self, entity: EntityInstanceId,
+                    operation: str,
+                    input: Optional[TInput] = None) -> Task:
+        """Schedule entity function for execution.
+
+        Parameters
+        ----------
+        entity: EntityInstanceId
+            The ID of the entity instance to call.
+        operation: str
+            The name of the operation to invoke on the entity.
+        input: Optional[TInput]
+            The optional JSON-serializable input to pass to the entity function.
+
+        Returns
+        -------
+        Task
+            A Durable Task that completes when the called entity function completes or fails.
+        """
+        pass
+
+    @abstractmethod
+    def signal_entity(
+            self,
+            entity_id: EntityInstanceId,
+            operation_name: str,
+            input: Optional[TInput] = None
+    ) -> None:
+        """Signal an entity function for execution.
+
+        Parameters
+        ----------
+        entity_id: EntityInstanceId
+            The ID of the entity instance to signal.
+        operation_name: str
+            The name of the operation to invoke on the entity.
+        input: Optional[TInput]
+            The optional JSON-serializable input to pass to the entity function.
+        """
+        pass
+
+    @abstractmethod
+    def lock_entities(self, entities: list[EntityInstanceId]) -> Task[EntityLock]:
+        """Creates a Task object that locks the specified entity instances.
+
+        The locks will be acquired the next time the orchestrator yields.
+        Best practice is to immediately yield this Task and enter the returned EntityLock.
+        The lock is released when the EntityLock is exited.
+
+        Parameters
+        ----------
+        entities: list[EntityInstanceId]
+            The list of entity instance IDs to lock.
+
+        Returns
+        -------
+        EntityLock
+            A context manager object that releases the locks when exited.
+        """
+        pass
+
+    @abstractmethod
+    def call_sub_orchestrator(self, orchestrator: Union[Orchestrator[TInput, TOutput], str], *,
+                              input: Optional[TInput] = None,
+                              instance_id: Optional[str] = None,
+                              retry_policy: Optional[RetryPolicy] = None,
+                              version: Optional[str] = None) -> Task[TOutput]:
         """Schedule sub-orchestrator function for execution.
 
         Parameters
         ----------
         orchestrator: Orchestrator[TInput, TOutput]
             A reference to the orchestrator function to call.
-        input: Union[TInput, None]
+        input: Optional[TInput]
             The optional JSON-serializable input to pass to the orchestrator function.
-        instance_id: Union[str, None]
+        instance_id: Optional[str]
             A unique ID to use for the sub-orchestration instance. If not specified, a
             random UUID will be used.
+        retry_policy: Optional[RetryPolicy]
+            The retry policy to use for this sub-orchestrator call.
 
         Returns
         -------
@@ -147,9 +245,26 @@ class OrchestrationContext(ABC):
         """
         pass
 
+    @abstractmethod
+    def continue_as_new(self, new_input: Any, *, save_events: bool = False) -> None:
+        """Continue the orchestration execution as a new instance.
+
+        Parameters
+        ----------
+        new_input : Any
+            The new input to use for the new orchestration instance.
+        save_events : bool
+            A flag indicating whether to add any unprocessed external events in the new orchestration history.
+        """
+        pass
+
+    @abstractmethod
+    def _exit_critical_section(self) -> None:
+        pass
+
 
 class FailureDetails:
-    def __init__(self, message: str, error_type: str, stack_trace: Union[str, None]):
+    def __init__(self, message: str, error_type: str, stack_trace: Optional[str]):
         self._message = message
         self._error_type = error_type
         self._stack_trace = stack_trace
@@ -163,7 +278,7 @@ class FailureDetails:
         return self._error_type
 
     @property
-    def stack_trace(self) -> Union[str, None]:
+    def stack_trace(self) -> Optional[str]:
         return self._stack_trace
 
 
@@ -193,8 +308,8 @@ class OrchestrationStateError(Exception):
 class Task(ABC, Generic[T]):
     """Abstract base class for asynchronous tasks in a durable orchestration."""
     _result: T
-    _exception: Union[TaskFailedError, None]
-    _parent: Union[CompositeTask[T], None]
+    _exception: Optional[TaskFailedError]
+    _parent: Optional[CompositeTask[T]]
 
     def __init__(self) -> None:
         super().__init__()
@@ -229,9 +344,9 @@ class Task(ABC, Generic[T]):
 
 class CompositeTask(Task[T]):
     """A task that is composed of other tasks."""
-    _tasks: List[Task]
+    _tasks: list[Task]
 
-    def __init__(self, tasks: List[Task]):
+    def __init__(self, tasks: list[Task]):
         super().__init__()
         self._tasks = tasks
         self._completed_tasks = 0
@@ -241,7 +356,7 @@ class CompositeTask(Task[T]):
             if task.is_complete:
                 self.on_child_completed(task)
 
-    def get_tasks(self) -> List[Task]:
+    def get_tasks(self) -> list[Task]:
         return self._tasks
 
     @abstractmethod
@@ -249,10 +364,40 @@ class CompositeTask(Task[T]):
         pass
 
 
+class WhenAllTask(CompositeTask[list[T]]):
+    """A task that completes when all of its child tasks complete."""
+
+    def __init__(self, tasks: list[Task[T]]):
+        super().__init__(tasks)
+        self._completed_tasks = 0
+        self._failed_tasks = 0
+
+    @property
+    def pending_tasks(self) -> int:
+        """Returns the number of tasks that have not yet completed."""
+        return len(self._tasks) - self._completed_tasks
+
+    def on_child_completed(self, task: Task[T]):
+        if self.is_complete:
+            raise ValueError('The task has already completed.')
+        self._completed_tasks += 1
+        if task.is_failed and self._exception is None:
+            self._exception = task.get_exception()
+            self._is_complete = True
+        if self._completed_tasks == len(self._tasks):
+            # The order of the result MUST match the order of the tasks provided to the constructor.
+            self._result = [task.get_result() for task in self._tasks]
+            self._is_complete = True
+
+    def get_completed_tasks(self) -> int:
+        return self._completed_tasks
+
+
 class CompletableTask(Task[T]):
 
     def __init__(self):
         super().__init__()
+        self._retryable_parent = None
 
     def complete(self, result: T):
         if self._is_complete:
@@ -271,39 +416,57 @@ class CompletableTask(Task[T]):
             self._parent.on_child_completed(self)
 
 
-class WhenAllTask(CompositeTask[List[T]]):
-    """A task that completes when all of its child tasks complete."""
+class RetryableTask(CompletableTask[T]):
+    """A task that can be retried according to a retry policy."""
 
-    def __init__(self, tasks: List[Task[T]]):
-        super().__init__(tasks)
-        self._completed_tasks = 0
-        self._failed_tasks = 0
+    def __init__(self, retry_policy: RetryPolicy, action: pb.OrchestratorAction,
+                 start_time: datetime, is_sub_orch: bool) -> None:
+        super().__init__()
+        self._action = action
+        self._retry_policy = retry_policy
+        self._attempt_count = 1
+        self._start_time = start_time
+        self._is_sub_orch = is_sub_orch
 
-    @property
-    def pending_tasks(self) -> int:
-        """Returns the number of tasks that have not yet completed."""
-        return len(self._tasks) - self._completed_tasks
+    def increment_attempt_count(self) -> None:
+        self._attempt_count += 1
 
-    def on_child_completed(self, task: Task[T]):
-        if self.is_complete:
-            raise ValueError('The task has already completed.')
-        self._completed_tasks += 1
-        if task.is_failed and self._exception is None:
-            self._exception = task.get_exception()
-            self._is_complete = True
-        if self._completed_tasks == len(self._tasks):
-            # The order of the result MUST match the order of the tasks provided to the constructor.
-            self._result = [task.get_result() for task in self._tasks]
-            self._is_complete = True
+    def compute_next_delay(self) -> Optional[timedelta]:
+        if self._attempt_count >= self._retry_policy.max_number_of_attempts:
+            return None
 
-    def get_completed_tasks(self) -> int:
-        return self._completed_tasks
+        retry_expiration: datetime = datetime.max
+        if self._retry_policy.retry_timeout is not None and self._retry_policy.retry_timeout != datetime.max:
+            retry_expiration = self._start_time + self._retry_policy.retry_timeout
+
+        if self._retry_policy.backoff_coefficient is None:
+            backoff_coefficient = 1.0
+        else:
+            backoff_coefficient = self._retry_policy.backoff_coefficient
+
+        if datetime.utcnow() < retry_expiration:
+            next_delay_f = math.pow(backoff_coefficient, self._attempt_count - 1) * self._retry_policy.first_retry_interval.total_seconds()
+
+            if self._retry_policy.max_retry_interval is not None:
+                next_delay_f = min(next_delay_f, self._retry_policy.max_retry_interval.total_seconds())
+                return timedelta(seconds=next_delay_f)
+
+        return None
+
+
+class TimerTask(CompletableTask[T]):
+
+    def __init__(self) -> None:
+        super().__init__()
+
+    def set_retryable_parent(self, retryable_task: RetryableTask):
+        self._retryable_parent = retryable_task
 
 
 class WhenAnyTask(CompositeTask[Task]):
     """A task that completes when any of its child tasks complete."""
 
-    def __init__(self, tasks: List[Task]):
+    def __init__(self, tasks: list[Task]):
         super().__init__(tasks)
 
     def on_child_completed(self, task: Task):
@@ -313,12 +476,12 @@ class WhenAnyTask(CompositeTask[Task]):
             self._result = task
 
 
-def when_all(tasks: List[Task[T]]) -> WhenAllTask[T]:
+def when_all(tasks: list[Task[T]]) -> WhenAllTask[T]:
     """Returns a task that completes when all of the provided tasks complete or when one of the tasks fail."""
     return WhenAllTask(tasks)
 
 
-def when_any(tasks: List[Task]) -> WhenAnyTask:
+def when_any(tasks: list[Task]) -> WhenAnyTask:
     """Returns a task that completes when any of the provided tasks complete or fail."""
     return WhenAnyTask(tasks)
 
@@ -362,6 +525,76 @@ Orchestrator = Callable[[OrchestrationContext, TInput], Union[Generator[Task, An
 # Activities are simple functions that can be scheduled by orchestrators
 Activity = Callable[[ActivityContext, TInput], TOutput]
 
+Entity = Union[Callable[[EntityContext, TInput], TOutput], type[DurableEntity]]
+
+
+class RetryPolicy:
+    """Represents the retry policy for an orchestration or activity function."""
+
+    def __init__(self, *,
+                 first_retry_interval: timedelta,
+                 max_number_of_attempts: int,
+                 backoff_coefficient: Optional[float] = 1.0,
+                 max_retry_interval: Optional[timedelta] = None,
+                 retry_timeout: Optional[timedelta] = None):
+        """Creates a new RetryPolicy instance.
+
+        Parameters
+        ----------
+        first_retry_interval : timedelta
+            The retry interval to use for the first retry attempt.
+        max_number_of_attempts : int
+            The maximum number of retry attempts.
+        backoff_coefficient : Optional[float]
+            The backoff coefficient to use for calculating the next retry interval.
+        max_retry_interval : Optional[timedelta]
+            The maximum retry interval to use for any retry attempt.
+        retry_timeout : Optional[timedelta]
+            The maximum amount of time to spend retrying the operation.
+        """
+        # validate inputs
+        if first_retry_interval < timedelta(seconds=0):
+            raise ValueError('first_retry_interval must be >= 0')
+        if max_number_of_attempts < 1:
+            raise ValueError('max_number_of_attempts must be >= 1')
+        if backoff_coefficient is not None and backoff_coefficient < 1:
+            raise ValueError('backoff_coefficient must be >= 1')
+        if max_retry_interval is not None and max_retry_interval < timedelta(seconds=0):
+            raise ValueError('max_retry_interval must be >= 0')
+        if retry_timeout is not None and retry_timeout < timedelta(seconds=0):
+            raise ValueError('retry_timeout must be >= 0')
+
+        self._first_retry_interval = first_retry_interval
+        self._max_number_of_attempts = max_number_of_attempts
+        self._backoff_coefficient = backoff_coefficient
+        self._max_retry_interval = max_retry_interval
+        self._retry_timeout = retry_timeout
+
+    @property
+    def first_retry_interval(self) -> timedelta:
+        """The retry interval to use for the first retry attempt."""
+        return self._first_retry_interval
+
+    @property
+    def max_number_of_attempts(self) -> int:
+        """The maximum number of retry attempts."""
+        return self._max_number_of_attempts
+
+    @property
+    def backoff_coefficient(self) -> Optional[float]:
+        """The backoff coefficient to use for calculating the next retry interval."""
+        return self._backoff_coefficient
+
+    @property
+    def max_retry_interval(self) -> Optional[timedelta]:
+        """The maximum retry interval to use for any retry attempt."""
+        return self._max_retry_interval
+
+    @property
+    def retry_timeout(self) -> Optional[timedelta]:
+        """The maximum amount of time to spend retrying the operation."""
+        return self._retry_timeout
+
 
 def get_name(fn: Callable) -> str:
     """Returns the name of the provided function"""