Introduce new TaskStatuses: queued and timed_out (#170)

alembic/versions/2024_04_09_0058-8335d7fecef9_add_new_indices_to_tasks_table.py

@@ -0,0 +1,32 @@
+"""Add new indices to tasks table
+
+Revision ID: 8335d7fecef9
+Revises: ea8e24d0bc8e
+Create Date: 2024-04-09 00:58:53.060477+00:00
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "8335d7fecef9"
+down_revision: Union[str, None] = "ea8e24d0bc8e"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_index(op.f("ix_tasks_created_at"), "tasks", ["created_at"], unique=False)
+    op.create_index(op.f("ix_tasks_modified_at"), "tasks", ["modified_at"], unique=False)
+    op.create_index(op.f("ix_tasks_status"), "tasks", ["status"], unique=False)
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_index(op.f("ix_tasks_status"), table_name="tasks")
+    op.drop_index(op.f("ix_tasks_modified_at"), table_name="tasks")
+    op.drop_index(op.f("ix_tasks_created_at"), table_name="tasks")
+    # ### end Alembic commands ###

skyvern/forge/sdk/db/models.py

@@ -30,7 +30,7 @@ class TaskModel(Base):
 
     task_id = Column(String, primary_key=True, index=True, default=generate_task_id)
     organization_id = Column(String, ForeignKey("organizations.organization_id"))
-    status = Column(String)
+    status = Column(String, index=True)
     webhook_callback_url = Column(String)
     title = Column(String)
     url = Column(String)
@@ -46,8 +46,10 @@ class TaskModel(Base):
     retry = Column(Integer, nullable=True)
     error_code_mapping = Column(JSON, nullable=True)
     errors = Column(JSON, default=[], nullable=False)
-    created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False)
+    created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False, index=True)
-    modified_at = Column(DateTime, default=datetime.datetime.utcnow, onupdate=datetime.datetime.utcnow, nullable=False)
+    modified_at = Column(
+        DateTime, default=datetime.datetime.utcnow, onupdate=datetime.datetime.utcnow, nullable=False, index=True
+    )
 
 
 class StepModel(Base):

skyvern/forge/sdk/schemas/tasks.py

@@ -73,19 +73,23 @@ class TaskRequest(BaseModel):
 
 class TaskStatus(StrEnum):
     created = "created"
+    queued = "queued"
     running = "running"
+    timed_out = "timed_out"
     failed = "failed"
     terminated = "terminated"
     completed = "completed"
 
     def is_final(self) -> bool:
-        return self in {TaskStatus.failed, TaskStatus.terminated, TaskStatus.completed}
+        return self in {TaskStatus.failed, TaskStatus.terminated, TaskStatus.completed, TaskStatus.timed_out}
 
     def can_update_to(self, new_status: TaskStatus) -> bool:
         allowed_transitions: dict[TaskStatus, set[TaskStatus]] = {
-            TaskStatus.created: {TaskStatus.running},
+            TaskStatus.created: {TaskStatus.queued, TaskStatus.running, TaskStatus.timed_out},
-            TaskStatus.running: {TaskStatus.completed, TaskStatus.failed, TaskStatus.terminated},
+            TaskStatus.queued: {TaskStatus.running, TaskStatus.timed_out},
+            TaskStatus.running: {TaskStatus.completed, TaskStatus.failed, TaskStatus.terminated, TaskStatus.timed_out},
             TaskStatus.failed: set(),
+            TaskStatus.terminated: set(),
             TaskStatus.completed: set(),
         }
         return new_status in allowed_transitions[self]
@@ -97,6 +101,7 @@ class TaskStatus(StrEnum):
     def cant_have_extracted_info(self) -> bool:
         status_cant_have_extracted_information = {
             TaskStatus.created,
+            TaskStatus.queued,
             TaskStatus.running,
             TaskStatus.failed,
             TaskStatus.terminated,