durabletask 0.1.0a5__py3-none-any.whl → 0.2.0__py3-none-any.whl

durabletask/internal/shared.py

@@ -5,39 +5,67 @@ import dataclasses
 import json
 import logging
 from types import SimpleNamespace
-from typing import Any, Dict, List, Tuple, Union
+from typing import Any, Optional, Sequence, Union
 
 import grpc
 
-from durabletask.internal.grpc_interceptor import DefaultClientInterceptorImpl
+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], metadata: Union[List[Tuple[str, str]], None], secure_channel: bool = False) -> 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()
 
+    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)
 
-    if metadata is not None and len(metadata) > 0:
-        interceptors = [DefaultClientInterceptorImpl(metadata)]
+    # Apply interceptors ONLY if they exist
+    if interceptors:
         channel = grpc.intercept_channel(channel, *interceptors)
     return channel
 
+
 def get_logger(
         name_suffix: str,
-        log_handler: Union[logging.Handler, None] = None,
-        log_formatter: Union[logging.Formatter, None] = None) -> logging.Logger:
+        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
@@ -78,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):
@@ -94,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,9 +4,10 @@
 # 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
 
 import durabletask.internal.helpers as pbh
 import durabletask.internal.orchestrator_service_pb2 as pb
@@ -69,6 +70,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 +99,18 @@ 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) -> 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.
 
         Returns
         -------
@@ -108,19 +121,22 @@ class OrchestrationContext(ABC):
 
     @abstractmethod
     def call_sub_orchestrator(self, orchestrator: Orchestrator[TInput, TOutput], *,
-                              input: Union[TInput, None] = None,
-                              instance_id: Union[str, None] = None) -> Task[TOutput]:
+                              input: Optional[TInput] = None,
+                              instance_id: Optional[str] = None,
+                              retry_policy: Optional[RetryPolicy] = 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
         -------
@@ -162,7 +178,7 @@ class OrchestrationContext(ABC):
 
 
 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
@@ -176,7 +192,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
 
 
@@ -206,8 +222,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__()
@@ -242,9 +258,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
@@ -254,7 +270,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
@@ -262,10 +278,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:
@@ -284,39 +330,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):
@@ -326,12 +390,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)
 
@@ -376,6 +440,74 @@ Orchestrator = Callable[[OrchestrationContext, TInput], Union[Generator[Task, An
 Activity = Callable[[ActivityContext, TInput], TOutput]
 
 
+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"""
     name = fn.__name__