From 7f6ecbd2d476dabf00062a0bf639b4e49c8682eb Mon Sep 17 00:00:00 2001
From: Fabian Martinez <46371672+famarting@users.noreply.github.com>
Date: Tue, 3 Dec 2024 12:17:38 +0100
Subject: [PATCH] support set custom status

Signed-off-by: Fabian Martinez <46371672+famarting@users.noreply.github.com>
---
 examples/workflow/task_chaining.py                          | 1 +
 .../dapr/ext/workflow/dapr_workflow_context.py              | 4 ++++
 ext/dapr-ext-workflow/dapr/ext/workflow/workflow_context.py | 6 ++++++
 3 files changed, 11 insertions(+)

diff --git a/examples/workflow/task_chaining.py b/examples/workflow/task_chaining.py
index c24e340c..c67308d5 100644
--- a/examples/workflow/task_chaining.py
+++ b/examples/workflow/task_chaining.py
@@ -27,6 +27,7 @@ def task_chain_workflow(ctx: wf.DaprWorkflowContext, wf_input: int):
     except Exception as e:
         yield ctx.call_activity(error_handler, input=str(e))
         raise
+    # TODO update to set custom status
     return [result1, result2, result3]
 
 
diff --git a/ext/dapr-ext-workflow/dapr/ext/workflow/dapr_workflow_context.py b/ext/dapr-ext-workflow/dapr/ext/workflow/dapr_workflow_context.py
index dbcccd64..2dee46fe 100644
--- a/ext/dapr-ext-workflow/dapr/ext/workflow/dapr_workflow_context.py
+++ b/ext/dapr-ext-workflow/dapr/ext/workflow/dapr_workflow_context.py
@@ -53,6 +53,10 @@ def current_utc_datetime(self) -> datetime:
     def is_replaying(self) -> bool:
         return self.__obj.is_replaying
 
+    def set_custom_status(self, custom_status: str) -> None:
+        self._logger.debug(f'{self.instance_id}: Setting custom status to {custom_status}')
+        self.__obj.set_custom_status(custom_status)
+
     def create_timer(self, fire_at: Union[datetime, timedelta]) -> task.Task:
         self._logger.debug(f'{self.instance_id}: Creating timer to fire at {fire_at} time')
         return self.__obj.create_timer(fire_at)
diff --git a/ext/dapr-ext-workflow/dapr/ext/workflow/workflow_context.py b/ext/dapr-ext-workflow/dapr/ext/workflow/workflow_context.py
index e0e3c736..36674236 100644
--- a/ext/dapr-ext-workflow/dapr/ext/workflow/workflow_context.py
+++ b/ext/dapr-ext-workflow/dapr/ext/workflow/workflow_context.py
@@ -84,6 +84,12 @@ def is_replaying(self) -> bool:
         """
         pass
 
+    @abstractmethod
+    def set_custom_status(self, custom_status: str) -> None:
+        """Set the custom status.
+        """
+        pass
+
     @abstractmethod
     def create_timer(self, fire_at: Union[datetime, timedelta]) -> task.Task:
         """Create a Timer Task to fire after at the specified deadline.