add index to speed up artifact queries (#241)

Co-authored-by: Shuchang Zheng <wintonzheng0325@gmail.com>

alembic/versions/2024_04_28_2320-68d78072fdb5_add_org_task_step_index.py

@@ -0,0 +1,28 @@
+"""Add org_task_step_index
+
+Revision ID: 68d78072fdb5
+Revises: 24303f1669a7
+Create Date: 2024-04-28 23:20:28.953686+00:00
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "68d78072fdb5"
+down_revision: Union[str, None] = "24303f1669a7"
+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("org_task_step_index", "artifacts", ["organization_id", "task_id", "step_id"], unique=False)
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_index("org_task_step_index", table_name="artifacts")
+    # ### end Alembic commands ###

skyvern/forge/sdk/db/client.py

@@ -679,10 +679,10 @@ class AgentDB:
         try:
             async with self.Session() as session:
                 artifact_query = select(ArtifactModel).filter_by(task_id=task_id)
-                if step_id:
-                    artifact_query = artifact_query.filter_by(step_id=step_id)
                 if organization_id:
                     artifact_query = artifact_query.filter_by(organization_id=organization_id)
+                if step_id:
+                    artifact_query = artifact_query.filter_by(step_id=step_id)
                 if artifact_types:
                     artifact_query = artifact_query.filter(ArtifactModel.artifact_type.in_(artifact_types))
 

skyvern/forge/sdk/db/models.py

@@ -1,6 +1,6 @@
 import datetime
 
-from sqlalchemy import JSON, Boolean, Column, DateTime, Enum, ForeignKey, Integer, Numeric, String, UnicodeText
+from sqlalchemy import JSON, Boolean, Column, DateTime, Enum, ForeignKey, Index, Integer, Numeric, String, UnicodeText
 from sqlalchemy.ext.asyncio import AsyncAttrs
 from sqlalchemy.orm import DeclarativeBase
 
@@ -105,6 +105,7 @@ class OrganizationAuthTokenModel(Base):
 
 class ArtifactModel(Base):
     __tablename__ = "artifacts"
+    __table_args__ = (Index("org_task_step_index", "organization_id", "task_id", "step_id"),)
 
     artifact_id = Column(String, primary_key=True, index=True, default=generate_artifact_id)
     organization_id = Column(String, ForeignKey("organizations.organization_id"))