vikarti.anatra/Skyvern
1
0
Fork 0
mirror of https://github.com/Skyvern-AI/skyvern.git synced 2025-09-02 10:41:04 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

add index to speed up artifact queries (#241)

Browse source 
Co-authored-by: Shuchang Zheng <wintonzheng0325@gmail.com>
This commit is contained in:
Kerem Yilmaz 2024-04-28 16:23:17 -07:00 • committed by GitHub
parent b321ac2481
commit 7301dd00b2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 32 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

28
alembic/versions/2024_04_28_2320-68d78072fdb5_add_org_task_step_index.py Normal file
View file

@ -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 ###

4
skyvern/forge/sdk/db/client.py
View file

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

3
skyvern/forge/sdk/db/models.py
View file

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