diff --git a/alembic/versions/2024_04_28_2320-68d78072fdb5_add_org_task_step_index.py b/alembic/versions/2024_04_28_2320-68d78072fdb5_add_org_task_step_index.py new file mode 100644 index 000000000..45ec8385b --- /dev/null +++ b/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 ### diff --git a/skyvern/forge/sdk/db/client.py b/skyvern/forge/sdk/db/client.py index f26c50934..f56c219c5 100644 --- a/skyvern/forge/sdk/db/client.py +++ b/skyvern/forge/sdk/db/client.py @@ -679,10 +679,10 @@ async def get_latest_n_artifacts( 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)) diff --git a/skyvern/forge/sdk/db/models.py b/skyvern/forge/sdk/db/models.py index aa681ff4a..083204833 100644 --- a/skyvern/forge/sdk/db/models.py +++ b/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"))