diff --git a/alembic/versions/2025_03_01_0428-a21b9f4f51d2_add_index_for_workflow_run_and_workflow_.py b/alembic/versions/2025_03_01_0428-a21b9f4f51d2_add_index_for_workflow_run_and_workflow_.py new file mode 100644 index 00000000..6f07b631 --- /dev/null +++ b/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 ### diff --git a/skyvern/forge/sdk/db/models.py b/skyvern/forge/sdk/db/models.py index 1db65fa8..cf8fbc86 100644 --- a/skyvern/forge/sdk/db/models.py +++ b/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, )