add index to speed up artifact queries (#241)

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

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"))