Add more indexes (#1858)

alembic/versions/2025_03_01_0428-a21b9f4f51d2_add_index_for_workflow_run_and_workflow_.py

@@ -0,0 +1,35 @@
+"""add index for workflow run and workflow tables
+
+Revision ID: a21b9f4f51d2
+Revises: b4bb0b98912a
+Create Date: 2025-03-01 04:28:57.760711+00:00
+
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "a21b9f4f51d2"
+down_revision: Union[str, None] = "b4bb0b98912a"
+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_steps_task_id"), "steps", ["task_id"], unique=False)
+    op.create_index("idx_tasks_org_created", "tasks", ["organization_id", "created_at"], unique=False)
+    op.create_index("idx_workflow_runs_org_created", "workflow_runs", ["organization_id", "created_at"], unique=False)
+    op.create_index(op.f("ix_workflow_runs_modified_at"), "workflow_runs", ["modified_at"], unique=False)
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_index(op.f("ix_workflow_runs_modified_at"), table_name="workflow_runs")
+    op.drop_index("idx_workflow_runs_org_created", table_name="workflow_runs")
+    op.drop_index("idx_tasks_org_created", table_name="tasks")
+    op.drop_index(op.f("ix_steps_task_id"), table_name="steps")
+    # ### end Alembic commands ###

skyvern/forge/sdk/db/models.py

@@ -56,6 +56,7 @@ class Base(AsyncAttrs, DeclarativeBase):
 
 class TaskModel(Base):
     __tablename__ = "tasks"
+    __table_args__ = (Index("idx_tasks_org_created", "organization_id", "created_at"),)
 
     task_id = Column(String, primary_key=True, index=True, default=generate_task_id)
     organization_id = Column(String, ForeignKey("organizations.organization_id"))
@@ -101,7 +102,7 @@ class StepModel(Base):
 
     step_id = Column(String, primary_key=True, index=True, default=generate_step_id)
     organization_id = Column(String, ForeignKey("organizations.organization_id"))
-    task_id = Column(String, ForeignKey("tasks.task_id"))
+    task_id = Column(String, ForeignKey("tasks.task_id"), index=True)
     status = Column(String)
     output = Column(JSON)
     order = Column(Integer)
@@ -230,6 +231,7 @@ class WorkflowModel(Base):
 
 class WorkflowRunModel(Base):
     __tablename__ = "workflow_runs"
+    __table_args__ = (Index("idx_workflow_runs_org_created", "organization_id", "created_at"),)
 
     workflow_run_id = Column(String, primary_key=True, index=True, default=generate_workflow_run_id)
     workflow_id = Column(String, ForeignKey("workflows.workflow_id"), nullable=False)
@@ -250,6 +252,7 @@ class WorkflowRunModel(Base):
         default=datetime.datetime.utcnow,
         onupdate=datetime.datetime.utcnow,
         nullable=False,
+        index=True,
     )