From cc24299af3ff4b235cc4aa57fffaa98cea8b3be3 Mon Sep 17 00:00:00 2001
From: Robin Ole Heinemann <robin.ole.heinemann@gmail.com>
Date: Sun, 7 Jan 2024 21:19:46 +0100
Subject: [PATCH] backend: allow reset_task script to reset tasks with any
 state

---
 backend/scripts/reset_task.py              | 9 ++++++++-
 backend/transcribee_backend/models/task.py | 3 +++
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/backend/scripts/reset_task.py b/backend/scripts/reset_task.py
index 08528eb8..5823aa44 100644
--- a/backend/scripts/reset_task.py
+++ b/backend/scripts/reset_task.py
@@ -11,13 +11,20 @@
     parser.add_argument(
         "--uuid", required=True, type=uuid.UUID, help="Task UUID or Document UUID"
     )
+    parser.add_argument(
+        "--state",
+        type=TaskState,
+        choices=list(TaskState),
+        default=TaskState.FAILED,
+        help="State of tasks to reset",
+    )
     args = parser.parse_args()
     with SessionContextManager(path="management_command:reset_task") as session:
         task = session.execute(
             update(Task)
             .where(
                 or_(Task.id == args.uuid, Task.document_id == args.uuid),
-                Task.state == TaskState.FAILED,
+                Task.state == args.state,
             )
             .values(state=TaskState.NEW, remaining_attempts=settings.task_attempt_limit)
         )
diff --git a/backend/transcribee_backend/models/task.py b/backend/transcribee_backend/models/task.py
index 1ce3af0e..76aeefb8 100644
--- a/backend/transcribee_backend/models/task.py
+++ b/backend/transcribee_backend/models/task.py
@@ -22,6 +22,9 @@ class TaskState(enum.Enum):
     COMPLETED = "COMPLETED"
     FAILED = "FAILED"
 
+    def __str__(self):
+        return self.value
+
 
 class TaskBase(SQLModel):
     task_type: TaskType