chore: Use exposed class for schedule_new_orchestration instead of using protos

Signed-off-by: Albert Callarisa <albert@diagrid.io>

Author: acroca

durabletask/aio/client.py

@@ -17,6 +17,7 @@
 from durabletask.aio.internal.grpc_interceptor import DefaultClientInterceptorImpl
 from durabletask.aio.internal.shared import ClientInterceptor, get_grpc_aio_channel
 from durabletask.client import (
+    OrchestrationIdReusePolicy,
     OrchestrationState,
     OrchestrationStatus,
     TInput,
@@ -81,7 +82,7 @@ async def schedule_new_orchestration(
         input: Optional[TInput] = None,
         instance_id: Optional[str] = None,
         start_at: Optional[datetime] = None,
-        reuse_id_policy: Optional[pb.OrchestrationIdReusePolicy] = None,
+        reuse_id_policy: Optional[OrchestrationIdReusePolicy] = None,
     ) -> str:
         name = orchestrator if isinstance(orchestrator, str) else task.get_name(orchestrator)
 
@@ -93,7 +94,7 @@ async def schedule_new_orchestration(
             else None,
             scheduledStartTimestamp=helpers.new_timestamp(start_at) if start_at else None,
             version=helpers.get_string_value(None),
-            orchestrationIdReusePolicy=reuse_id_policy,
+            orchestrationIdReusePolicy=reuse_id_policy._to_pb() if reuse_id_policy else None,
         )
 
         self._logger.info(f"Starting new '{name}' instance with ID = '{req.instanceId}'.")

durabletask/client.py

@@ -47,6 +47,28 @@ def __str__(self):
         return helpers.get_orchestration_status_str(self.value)
 
 
+class OrchestrationIdReuseAction(Enum):
+    """Action to take when scheduling an orchestration whose ID already exists."""
+
+    ERROR = pb.ERROR
+    IGNORE = pb.IGNORE
+    TERMINATE = pb.TERMINATE
+
+
+@dataclass
+class OrchestrationIdReusePolicy:
+    """Policy controlling what happens when a new orchestration is scheduled with an ID that already exists."""
+
+    action: OrchestrationIdReuseAction
+    operation_status: list[OrchestrationStatus]
+
+    def _to_pb(self) -> pb.OrchestrationIdReusePolicy:
+        return pb.OrchestrationIdReusePolicy(
+            operationStatus=[s.value for s in self.operation_status],
+            action=self.action.value,
+        )
+
+
 @dataclass
 class OrchestrationState:
     instance_id: str
@@ -166,7 +188,7 @@ def schedule_new_orchestration(
         input: Optional[TInput] = None,
         instance_id: Optional[str] = None,
         start_at: Optional[datetime] = None,
-        reuse_id_policy: Optional[pb.OrchestrationIdReusePolicy] = None,
+        reuse_id_policy: Optional[OrchestrationIdReusePolicy] = None,
     ) -> str:
         name = orchestrator if isinstance(orchestrator, str) else task.get_name(orchestrator)
 
@@ -180,7 +202,7 @@ def schedule_new_orchestration(
             input=input_pb,
             scheduledStartTimestamp=helpers.new_timestamp(start_at) if start_at else None,
             version=wrappers_pb2.StringValue(value=""),
-            orchestrationIdReusePolicy=reuse_id_policy,
+            orchestrationIdReusePolicy=reuse_id_policy._to_pb() if reuse_id_policy else None,
         )
 
         self._logger.info(f"Starting new '{name}' instance with ID = '{req.instanceId}'.")