From 620b5bfdeadfd7ce4192b574265ba7ed2936b549 Mon Sep 17 00:00:00 2001 From: Shuchang Zheng Date: Sat, 7 Dec 2024 17:34:18 -0800 Subject: [PATCH] Drop all the foreign keys on artifacts table except for the organization_id (#1352) --- ...drop_all_the_foreign_keys_on_artifacts_.py | 57 +++++++++++++++++++ skyvern/forge/sdk/db/models.py | 9 +-- 2 files changed, 62 insertions(+), 4 deletions(-) create mode 100644 alembic/versions/2024_12_08_0131-8069e38dc1b4_drop_all_the_foreign_keys_on_artifacts_.py diff --git a/alembic/versions/2024_12_08_0131-8069e38dc1b4_drop_all_the_foreign_keys_on_artifacts_.py b/alembic/versions/2024_12_08_0131-8069e38dc1b4_drop_all_the_foreign_keys_on_artifacts_.py new file mode 100644 index 00000000..85c94ab7 --- /dev/null +++ b/alembic/versions/2024_12_08_0131-8069e38dc1b4_drop_all_the_foreign_keys_on_artifacts_.py @@ -0,0 +1,57 @@ +"""drop all the foreign keys on artifacts table except for orgnaization_id + +Revision ID: 8069e38dc1b4 +Revises: fe49b59d836c +Create Date: 2024-12-08 01:31:56.328245+00:00 + +""" + +from typing import Sequence, Union + +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "8069e38dc1b4" +down_revision: Union[str, None] = "fe49b59d836c" +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_observer_cruise_index", "artifacts", ["organization_id", "observer_cruise_id"], unique=False) + op.drop_constraint("artifacts_workflow_run_id_fkey", "artifacts", type_="foreignkey") + op.drop_constraint("artifacts_observer_cruise_id_fkey", "artifacts", type_="foreignkey") + op.drop_constraint("artifacts_observer_thought_id_fkey", "artifacts", type_="foreignkey") + op.drop_constraint("artifacts_workflow_run_block_id_fkey", "artifacts", type_="foreignkey") + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.create_foreign_key( + "artifacts_workflow_run_block_id_fkey", + "artifacts", + "workflow_run_blocks", + ["workflow_run_block_id"], + ["workflow_run_block_id"], + ) + op.create_foreign_key( + "artifacts_observer_thought_id_fkey", + "artifacts", + "observer_thoughts", + ["observer_thought_id"], + ["observer_thought_id"], + ) + op.create_foreign_key( + "artifacts_observer_cruise_id_fkey", + "artifacts", + "observer_cruises", + ["observer_cruise_id"], + ["observer_cruise_id"], + ) + op.create_foreign_key( + "artifacts_workflow_run_id_fkey", "artifacts", "workflow_runs", ["workflow_run_id"], ["workflow_run_id"] + ) + op.drop_index("org_observer_cruise_index", table_name="artifacts") + # ### end Alembic commands ### diff --git a/skyvern/forge/sdk/db/models.py b/skyvern/forge/sdk/db/models.py index dc203dc5..6f97c330 100644 --- a/skyvern/forge/sdk/db/models.py +++ b/skyvern/forge/sdk/db/models.py @@ -162,14 +162,15 @@ class ArtifactModel(Base): __table_args__ = ( Index("org_task_step_index", "organization_id", "task_id", "step_id"), Index("org_workflow_run_index", "organization_id", "workflow_run_id"), + Index("org_observer_cruise_index", "organization_id", "observer_cruise_id"), ) artifact_id = Column(String, primary_key=True, index=True, default=generate_artifact_id) organization_id = Column(String, ForeignKey("organizations.organization_id")) - workflow_run_id = Column(String, ForeignKey("workflow_runs.workflow_run_id")) - workflow_run_block_id = Column(String, ForeignKey("workflow_run_blocks.workflow_run_block_id")) - observer_cruise_id = Column(String, ForeignKey("observer_cruises.observer_cruise_id")) - observer_thought_id = Column(String, ForeignKey("observer_thoughts.observer_thought_id")) + workflow_run_id = Column(String) + workflow_run_block_id = Column(String) + observer_cruise_id = Column(String) + observer_thought_id = Column(String) task_id = Column(String, ForeignKey("tasks.task_id")) step_id = Column(String, ForeignKey("steps.step_id"), index=True) artifact_type = Column(String)