shu/remove foreign key constraints on large table actions (#2278)

This commit is contained in:
Shuchang Zheng
2025-05-02 04:51:24 -07:00
committed by GitHub
parent 08e8b7486e
commit 7e1dd805f8
2 changed files with 293 additions and 38 deletions

View File

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

View File

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