backend: allow reset_task script to reset tasks with any state

bugbakery · Feb 28, 2024 · cc24299 · cc24299
1 parent c4a031d
commit cc24299
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 1 deletion.
diff --git 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
@@ -22,6 +22,9 @@ class TaskState(enum.Enum):
     COMPLETED = "COMPLETED"
     FAILED = "FAILED"
 
+    def __str__(self):
+        return self.value
+
 
 class TaskBase(SQLModel):
     task_type: TaskType