From 7e1dd805f89f5f357ca4bdcfbf7f4a1c60dbcd99 Mon Sep 17 00:00:00 2001 From: Shuchang Zheng Date: Fri, 2 May 2025 04:51:24 -0700 Subject: [PATCH] shu/remove foreign key constraints on large table actions (#2278) --- ...148-e8285b6ddcf0_drop_most_foreign_keys.py | 261 ++++++++++++++++++ skyvern/forge/sdk/db/models.py | 70 +++-- 2 files changed, 293 insertions(+), 38 deletions(-) create mode 100644 alembic/versions/2025_05_02_1148-e8285b6ddcf0_drop_most_foreign_keys.py diff --git a/alembic/versions/2025_05_02_1148-e8285b6ddcf0_drop_most_foreign_keys.py b/alembic/versions/2025_05_02_1148-e8285b6ddcf0_drop_most_foreign_keys.py new file mode 100644 index 00000000..23bedef2 --- /dev/null +++ b/alembic/versions/2025_05_02_1148-e8285b6ddcf0_drop_most_foreign_keys.py @@ -0,0 +1,261 @@ +"""drop most foreign keys + +Revision ID: e8285b6ddcf0 +Revises: 511d9da18f5d +Create Date: 2025-05-02 11:48:31.191371+00:00 + +""" + +from typing import Sequence, Union + +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "e8285b6ddcf0" +down_revision: Union[str, None] = "511d9da18f5d" +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.drop_constraint("actions_step_id_fkey", "actions", type_="foreignkey") + op.drop_constraint("actions_task_id_fkey", "actions", type_="foreignkey") + op.drop_constraint("actions_source_action_id_fkey", "actions", type_="foreignkey") + op.drop_constraint("actions_workflow_run_id_fkey", "actions", type_="foreignkey") + op.drop_constraint("actions_organization_id_fkey", "actions", type_="foreignkey") + op.drop_constraint("artifacts_task_id_fkey", "artifacts", type_="foreignkey") + op.drop_constraint("artifacts_step_id_fkey", "artifacts", type_="foreignkey") + op.drop_constraint("aws_secret_parameters_workflow_id_fkey", "aws_secret_parameters", type_="foreignkey") + op.drop_constraint( + "bitwarden_credit_card_data_parameters_workflow_id_fkey", + "bitwarden_credit_card_data_parameters", + type_="foreignkey", + ) + op.drop_constraint( + "bitwarden_login_credential_parameters_workflow_id_fkey", + "bitwarden_login_credential_parameters", + type_="foreignkey", + ) + op.drop_constraint( + "bitwarden_sensitive_information_parameters_workflow_id_fkey", + "bitwarden_sensitive_information_parameters", + type_="foreignkey", + ) + op.drop_constraint("credential_parameters_workflow_id_fkey", "credential_parameters", type_="foreignkey") + op.drop_constraint("observer_cruises_workflow_run_id_fkey", "observer_cruises", type_="foreignkey") + op.drop_constraint("observer_cruises_workflow_id_fkey", "observer_cruises", type_="foreignkey") + op.drop_constraint("observer_cruises_organization_id_fkey", "observer_cruises", type_="foreignkey") + op.drop_constraint("observer_thoughts_workflow_run_id_fkey", "observer_thoughts", type_="foreignkey") + op.drop_constraint("observer_thoughts_workflow_id_fkey", "observer_thoughts", type_="foreignkey") + op.drop_constraint("observer_thoughts_observer_cruise_id_fkey", "observer_thoughts", type_="foreignkey") + op.drop_constraint("observer_thoughts_organization_id_fkey", "observer_thoughts", type_="foreignkey") + op.drop_constraint("observer_thoughts_workflow_run_block_id_fkey", "observer_thoughts", type_="foreignkey") + op.drop_constraint("output_parameters_workflow_id_fkey", "output_parameters", type_="foreignkey") + op.drop_constraint( + "persistent_browser_sessions_organization_id_fkey", "persistent_browser_sessions", type_="foreignkey" + ) + op.drop_constraint("task_generations_organization_id_fkey", "task_generations", type_="foreignkey") + op.drop_constraint("workflow_parameters_workflow_id_fkey", "workflow_parameters", type_="foreignkey") + op.drop_constraint("workflow_run_blocks_task_id_fkey", "workflow_run_blocks", type_="foreignkey") + op.drop_constraint("workflow_run_blocks_organization_id_fkey", "workflow_run_blocks", type_="foreignkey") + op.drop_constraint( + "workflow_run_blocks_parent_workflow_run_block_id_fkey", "workflow_run_blocks", type_="foreignkey" + ) + op.drop_constraint("workflow_run_blocks_block_workflow_run_id_fkey", "workflow_run_blocks", type_="foreignkey") + op.drop_constraint("workflow_run_blocks_workflow_run_id_fkey", "workflow_run_blocks", type_="foreignkey") + op.drop_constraint( + "workflow_run_output_parameters_output_parameter_id_fkey", "workflow_run_output_parameters", type_="foreignkey" + ) + op.drop_constraint( + "workflow_run_output_parameters_workflow_run_id_fkey", "workflow_run_output_parameters", type_="foreignkey" + ) + op.drop_constraint("workflow_run_parameters_workflow_run_id_fkey", "workflow_run_parameters", type_="foreignkey") + op.drop_constraint( + "workflow_run_parameters_workflow_parameter_id_fkey", "workflow_run_parameters", type_="foreignkey" + ) + op.drop_constraint("workflow_runs_workflow_id_fkey", "workflow_runs", type_="foreignkey") + op.drop_constraint("fk_workflow_runs_organization_id", "workflow_runs", type_="foreignkey") + op.drop_constraint("workflow_runs_parent_workflow_run_id_fkey", "workflow_runs", type_="foreignkey") + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.create_foreign_key( + "workflow_runs_parent_workflow_run_id_fkey", + "workflow_runs", + "workflow_runs", + ["parent_workflow_run_id"], + ["workflow_run_id"], + ) + op.create_foreign_key( + "fk_workflow_runs_organization_id", "workflow_runs", "organizations", ["organization_id"], ["organization_id"] + ) + op.create_foreign_key( + "workflow_runs_workflow_id_fkey", "workflow_runs", "workflows", ["workflow_id"], ["workflow_id"] + ) + op.create_foreign_key( + "workflow_run_parameters_workflow_parameter_id_fkey", + "workflow_run_parameters", + "workflow_parameters", + ["workflow_parameter_id"], + ["workflow_parameter_id"], + ) + op.create_foreign_key( + "workflow_run_parameters_workflow_run_id_fkey", + "workflow_run_parameters", + "workflow_runs", + ["workflow_run_id"], + ["workflow_run_id"], + ) + op.create_foreign_key( + "workflow_run_output_parameters_workflow_run_id_fkey", + "workflow_run_output_parameters", + "workflow_runs", + ["workflow_run_id"], + ["workflow_run_id"], + ) + op.create_foreign_key( + "workflow_run_output_parameters_output_parameter_id_fkey", + "workflow_run_output_parameters", + "output_parameters", + ["output_parameter_id"], + ["output_parameter_id"], + ) + op.create_foreign_key( + "workflow_run_blocks_workflow_run_id_fkey", + "workflow_run_blocks", + "workflow_runs", + ["workflow_run_id"], + ["workflow_run_id"], + ) + op.create_foreign_key( + "workflow_run_blocks_block_workflow_run_id_fkey", + "workflow_run_blocks", + "workflow_runs", + ["block_workflow_run_id"], + ["workflow_run_id"], + ) + op.create_foreign_key( + "workflow_run_blocks_parent_workflow_run_block_id_fkey", + "workflow_run_blocks", + "workflow_run_blocks", + ["parent_workflow_run_block_id"], + ["workflow_run_block_id"], + ) + op.create_foreign_key( + "workflow_run_blocks_organization_id_fkey", + "workflow_run_blocks", + "organizations", + ["organization_id"], + ["organization_id"], + ) + op.create_foreign_key("workflow_run_blocks_task_id_fkey", "workflow_run_blocks", "tasks", ["task_id"], ["task_id"]) + op.create_foreign_key( + "workflow_parameters_workflow_id_fkey", "workflow_parameters", "workflows", ["workflow_id"], ["workflow_id"] + ) + op.create_foreign_key( + "task_generations_organization_id_fkey", + "task_generations", + "organizations", + ["organization_id"], + ["organization_id"], + ) + op.create_foreign_key( + "persistent_browser_sessions_organization_id_fkey", + "persistent_browser_sessions", + "organizations", + ["organization_id"], + ["organization_id"], + ) + op.create_foreign_key( + "output_parameters_workflow_id_fkey", "output_parameters", "workflows", ["workflow_id"], ["workflow_id"] + ) + op.create_foreign_key( + "observer_thoughts_workflow_run_block_id_fkey", + "observer_thoughts", + "workflow_run_blocks", + ["workflow_run_block_id"], + ["workflow_run_block_id"], + ) + op.create_foreign_key( + "observer_thoughts_organization_id_fkey", + "observer_thoughts", + "organizations", + ["organization_id"], + ["organization_id"], + ) + op.create_foreign_key( + "observer_thoughts_observer_cruise_id_fkey", + "observer_thoughts", + "observer_cruises", + ["observer_cruise_id"], + ["observer_cruise_id"], + ) + op.create_foreign_key( + "observer_thoughts_workflow_id_fkey", "observer_thoughts", "workflows", ["workflow_id"], ["workflow_id"] + ) + op.create_foreign_key( + "observer_thoughts_workflow_run_id_fkey", + "observer_thoughts", + "workflow_runs", + ["workflow_run_id"], + ["workflow_run_id"], + ) + op.create_foreign_key( + "observer_cruises_organization_id_fkey", + "observer_cruises", + "organizations", + ["organization_id"], + ["organization_id"], + ) + op.create_foreign_key( + "observer_cruises_workflow_id_fkey", "observer_cruises", "workflows", ["workflow_id"], ["workflow_id"] + ) + op.create_foreign_key( + "observer_cruises_workflow_run_id_fkey", + "observer_cruises", + "workflow_runs", + ["workflow_run_id"], + ["workflow_run_id"], + ) + op.create_foreign_key( + "credential_parameters_workflow_id_fkey", "credential_parameters", "workflows", ["workflow_id"], ["workflow_id"] + ) + op.create_foreign_key( + "bitwarden_sensitive_information_parameters_workflow_id_fkey", + "bitwarden_sensitive_information_parameters", + "workflows", + ["workflow_id"], + ["workflow_id"], + ) + op.create_foreign_key( + "bitwarden_login_credential_parameters_workflow_id_fkey", + "bitwarden_login_credential_parameters", + "workflows", + ["workflow_id"], + ["workflow_id"], + ) + op.create_foreign_key( + "bitwarden_credit_card_data_parameters_workflow_id_fkey", + "bitwarden_credit_card_data_parameters", + "workflows", + ["workflow_id"], + ["workflow_id"], + ) + op.create_foreign_key( + "aws_secret_parameters_workflow_id_fkey", "aws_secret_parameters", "workflows", ["workflow_id"], ["workflow_id"] + ) + op.create_foreign_key("artifacts_step_id_fkey", "artifacts", "steps", ["step_id"], ["step_id"]) + op.create_foreign_key("artifacts_task_id_fkey", "artifacts", "tasks", ["task_id"], ["task_id"]) + op.create_foreign_key( + "actions_organization_id_fkey", "actions", "organizations", ["organization_id"], ["organization_id"] + ) + op.create_foreign_key( + "actions_workflow_run_id_fkey", "actions", "workflow_runs", ["workflow_run_id"], ["workflow_run_id"] + ) + op.create_foreign_key("actions_source_action_id_fkey", "actions", "actions", ["source_action_id"], ["action_id"]) + op.create_foreign_key("actions_task_id_fkey", "actions", "tasks", ["task_id"], ["task_id"]) + op.create_foreign_key("actions_step_id_fkey", "actions", "steps", ["step_id"], ["step_id"]) + # ### end Alembic commands ### diff --git a/skyvern/forge/sdk/db/models.py b/skyvern/forge/sdk/db/models.py index 595146e0..7863bbea 100644 --- a/skyvern/forge/sdk/db/models.py +++ b/skyvern/forge/sdk/db/models.py @@ -178,8 +178,8 @@ class ArtifactModel(Base): observer_cruise_id = Column(String, index=True) observer_thought_id = Column(String, index=True) ai_suggestion_id = Column(String, index=True) - task_id = Column(String, ForeignKey("tasks.task_id")) - step_id = Column(String, ForeignKey("steps.step_id"), index=True) + task_id = Column(String) + step_id = Column(String, index=True) artifact_type = Column(String) uri = Column(String) created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False) @@ -236,11 +236,11 @@ class WorkflowRunModel(Base): __table_args__ = (Index("idx_workflow_runs_org_created", "organization_id", "created_at"),) workflow_run_id = Column(String, primary_key=True, index=True, default=generate_workflow_run_id) - workflow_id = Column(String, ForeignKey("workflows.workflow_id"), nullable=False) + workflow_id = Column(String, nullable=False) workflow_permanent_id = Column(String, nullable=False, index=True) # workfow runs with parent_workflow_run_id are nested workflow runs which won't show up in the workflow run history - parent_workflow_run_id = Column(String, ForeignKey("workflow_runs.workflow_run_id"), nullable=True, index=True) - organization_id = Column(String, ForeignKey("organizations.organization_id"), nullable=False, index=True) + parent_workflow_run_id = Column(String, nullable=True, index=True) + organization_id = Column(String, nullable=False, index=True) status = Column(String, nullable=False) failure_reason = Column(String) proxy_location = Column(String) @@ -265,7 +265,7 @@ class WorkflowParameterModel(Base): workflow_parameter_type = Column(String, nullable=False) key = Column(String, nullable=False) description = Column(String, nullable=True) - workflow_id = Column(String, ForeignKey("workflows.workflow_id"), index=True, nullable=False) + workflow_id = Column(String, index=True, nullable=False) default_value = Column(String, nullable=True) created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False) modified_at = Column( @@ -283,7 +283,7 @@ class OutputParameterModel(Base): output_parameter_id = Column(String, primary_key=True, index=True, default=generate_output_parameter_id) key = Column(String, nullable=False) description = Column(String, nullable=True) - workflow_id = Column(String, ForeignKey("workflows.workflow_id"), index=True, nullable=False) + workflow_id = Column(String, index=True, nullable=False) created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False) modified_at = Column( DateTime, @@ -298,7 +298,7 @@ class AWSSecretParameterModel(Base): __tablename__ = "aws_secret_parameters" aws_secret_parameter_id = Column(String, primary_key=True, index=True, default=generate_aws_secret_parameter_id) - workflow_id = Column(String, ForeignKey("workflows.workflow_id"), index=True, nullable=False) + workflow_id = Column(String, index=True, nullable=False) key = Column(String, nullable=False) description = Column(String, nullable=True) aws_key = Column(String, nullable=False) @@ -321,7 +321,7 @@ class BitwardenLoginCredentialParameterModel(Base): index=True, default=generate_bitwarden_login_credential_parameter_id, ) - workflow_id = Column(String, ForeignKey("workflows.workflow_id"), index=True, nullable=False) + workflow_id = Column(String, index=True, nullable=False) key = Column(String, nullable=False) description = Column(String, nullable=True) bitwarden_client_id_aws_secret_key = Column(String, nullable=False) @@ -349,7 +349,7 @@ class BitwardenSensitiveInformationParameterModel(Base): index=True, default=generate_bitwarden_sensitive_information_parameter_id, ) - workflow_id = Column(String, ForeignKey("workflows.workflow_id"), index=True, nullable=False) + workflow_id = Column(String, index=True, nullable=False) key = Column(String, nullable=False) description = Column(String, nullable=True) bitwarden_client_id_aws_secret_key = Column(String, nullable=False) @@ -378,7 +378,7 @@ class BitwardenCreditCardDataParameterModel(Base): index=True, default=generate_bitwarden_credit_card_data_parameter_id, ) - workflow_id = Column(String, ForeignKey("workflows.workflow_id"), index=True, nullable=False) + workflow_id = Column(String, index=True, nullable=False) key = Column(String, nullable=False) description = Column(String, nullable=True) bitwarden_client_id_aws_secret_key = Column(String, nullable=False) @@ -395,7 +395,7 @@ class CredentialParameterModel(Base): __tablename__ = "credential_parameters" credential_parameter_id = Column(String, primary_key=True, index=True, default=generate_credential_parameter_id) - workflow_id = Column(String, ForeignKey("workflows.workflow_id"), index=True, nullable=False) + workflow_id = Column(String, index=True, nullable=False) key = Column(String, nullable=False) description = Column(String, nullable=True) @@ -411,13 +411,11 @@ class WorkflowRunParameterModel(Base): workflow_run_id = Column( String, - ForeignKey("workflow_runs.workflow_run_id"), primary_key=True, index=True, ) workflow_parameter_id = Column( String, - ForeignKey("workflow_parameters.workflow_parameter_id"), primary_key=True, index=True, ) @@ -431,13 +429,11 @@ class WorkflowRunOutputParameterModel(Base): workflow_run_id = Column( String, - ForeignKey("workflow_runs.workflow_run_id"), primary_key=True, index=True, ) output_parameter_id = Column( String, - ForeignKey("output_parameters.output_parameter_id"), primary_key=True, index=True, ) @@ -453,7 +449,7 @@ class TaskGenerationModel(Base): __tablename__ = "task_generations" task_generation_id = Column(String, primary_key=True, default=generate_task_generation_id) - organization_id = Column(String, ForeignKey("organizations.organization_id"), nullable=False) + organization_id = Column(String, nullable=False) user_prompt = Column(String, nullable=False) user_prompt_hash = Column(String, index=True) url = Column(String) @@ -509,11 +505,11 @@ class ActionModel(Base): action_id = Column(String, primary_key=True, index=True, default=generate_action_id) action_type = Column(String, nullable=False) - source_action_id = Column(String, ForeignKey("actions.action_id"), nullable=True, index=True) - organization_id = Column(String, ForeignKey("organizations.organization_id"), nullable=True) - workflow_run_id = Column(String, ForeignKey("workflow_runs.workflow_run_id"), nullable=True) - task_id = Column(String, ForeignKey("tasks.task_id"), nullable=False, index=True) - step_id = Column(String, ForeignKey("steps.step_id"), nullable=False) + source_action_id = Column(String, nullable=True, index=True) + organization_id = Column(String, nullable=True) + workflow_run_id = Column(String, nullable=True) + task_id = Column(String, nullable=False, index=True) + step_id = Column(String, nullable=False) step_order = Column(Integer, nullable=False) action_order = Column(Integer, nullable=False) status = Column(String, nullable=False) @@ -535,15 +531,13 @@ class WorkflowRunBlockModel(Base): __table_args__ = (Index("wfrb_org_wfr_index", "organization_id", "workflow_run_id"),) workflow_run_block_id = Column(String, primary_key=True, default=generate_workflow_run_block_id) - workflow_run_id = Column(String, ForeignKey("workflow_runs.workflow_run_id"), nullable=False) + workflow_run_id = Column(String, nullable=False) # this is the inner workflow run id of the taskv2 block - block_workflow_run_id = Column(String, ForeignKey("workflow_runs.workflow_run_id"), nullable=True) - parent_workflow_run_block_id = Column( - String, ForeignKey("workflow_run_blocks.workflow_run_block_id"), nullable=True - ) - organization_id = Column(String, ForeignKey("organizations.organization_id"), nullable=True) + block_workflow_run_id = Column(String, nullable=True) + parent_workflow_run_block_id = Column(String, nullable=True) + organization_id = Column(String, nullable=True) description = Column(String, nullable=True) - task_id = Column(String, ForeignKey("tasks.task_id"), nullable=True) + task_id = Column(String, nullable=True) label = Column(String, nullable=True) block_type = Column(String, nullable=False) status = Column(String, nullable=False) @@ -579,9 +573,9 @@ class TaskV2Model(Base): # observer_cruise_id is the task_id for task v2 observer_cruise_id = Column(String, primary_key=True, default=generate_task_v2_id) status = Column(String, nullable=False, default="created") - organization_id = Column(String, ForeignKey("organizations.organization_id"), nullable=True) - workflow_run_id = Column(String, ForeignKey("workflow_runs.workflow_run_id"), nullable=True) - workflow_id = Column(String, ForeignKey("workflows.workflow_id"), nullable=True) + organization_id = Column(String, nullable=True) + workflow_run_id = Column(String, nullable=True) + workflow_id = Column(String, nullable=True) workflow_permanent_id = Column(String, nullable=True) prompt = Column(UnicodeText, nullable=True) url = Column(String, nullable=True) @@ -603,11 +597,11 @@ class ThoughtModel(Base): __table_args__ = (Index("observer_cruise_index", "organization_id", "observer_cruise_id"),) observer_thought_id = Column(String, primary_key=True, default=generate_thought_id) - organization_id = Column(String, ForeignKey("organizations.organization_id"), nullable=True) - observer_cruise_id = Column(String, ForeignKey("observer_cruises.observer_cruise_id"), nullable=False) - workflow_run_id = Column(String, ForeignKey("workflow_runs.workflow_run_id"), nullable=True) - workflow_run_block_id = Column(String, ForeignKey("workflow_run_blocks.workflow_run_block_id"), nullable=True) - workflow_id = Column(String, ForeignKey("workflows.workflow_id"), nullable=True) + organization_id = Column(String, nullable=True) + observer_cruise_id = Column(String, nullable=False) + workflow_run_id = Column(String, nullable=True) + workflow_run_block_id = Column(String, nullable=True) + workflow_id = Column(String, nullable=True) workflow_permanent_id = Column(String, nullable=True) user_input = Column(UnicodeText, nullable=True) observation = Column(String, nullable=True) @@ -631,7 +625,7 @@ class PersistentBrowserSessionModel(Base): __tablename__ = "persistent_browser_sessions" persistent_browser_session_id = Column(String, primary_key=True, default=generate_persistent_browser_session_id) - organization_id = Column(String, ForeignKey("organizations.organization_id"), nullable=False) + organization_id = Column(String, nullable=False) runnable_type = Column(String, nullable=True) runnable_id = Column(String, nullable=True) browser_id = Column(String, nullable=True)