117 lines
5.0 KiB
Python
117 lines
5.0 KiB
Python
"""Repurpose JobUsageStatistics for new Steps table
|
|
|
|
Revision ID: 416b9d2db10b
|
|
Revises: 25fc99e97839
|
|
Create Date: 2025-01-17 11:27:42.115755
|
|
|
|
"""
|
|
|
|
from typing import Sequence, Union
|
|
|
|
import sqlalchemy as sa
|
|
from sqlalchemy.dialects import postgresql
|
|
|
|
from alembic import op
|
|
|
|
# revision identifiers, used by Alembic.
|
|
revision: str = "416b9d2db10b"
|
|
down_revision: Union[str, None] = "25fc99e97839"
|
|
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! ###
|
|
# Rename the table
|
|
op.rename_table("job_usage_statistics", "steps")
|
|
|
|
# Rename the foreign key constraint and drop non-null constraint
|
|
op.alter_column("steps", "job_id", nullable=True)
|
|
op.drop_constraint("fk_job_usage_statistics_job_id", "steps", type_="foreignkey")
|
|
|
|
# Change id field from int to string
|
|
op.execute("ALTER TABLE steps RENAME COLUMN id TO old_id")
|
|
op.add_column("steps", sa.Column("id", sa.String(), nullable=True))
|
|
op.execute("""UPDATE steps SET id = 'step-' || gen_random_uuid()::text""")
|
|
op.drop_column("steps", "old_id")
|
|
op.alter_column("steps", "id", nullable=False)
|
|
op.create_primary_key("pk_steps_id", "steps", ["id"])
|
|
|
|
# Add new columns
|
|
op.add_column("steps", sa.Column("origin", sa.String(), nullable=True))
|
|
op.add_column("steps", sa.Column("organization_id", sa.String(), nullable=True))
|
|
op.add_column("steps", sa.Column("provider_id", sa.String(), nullable=True))
|
|
op.add_column("steps", sa.Column("provider_name", sa.String(), nullable=True))
|
|
op.add_column("steps", sa.Column("model", sa.String(), nullable=True))
|
|
op.add_column("steps", sa.Column("context_window_limit", sa.Integer(), nullable=True))
|
|
op.add_column(
|
|
"steps",
|
|
sa.Column("completion_tokens_details", postgresql.JSON(astext_type=sa.Text()), autoincrement=False, nullable=True),
|
|
)
|
|
op.add_column(
|
|
"steps",
|
|
sa.Column("tags", postgresql.JSON(astext_type=sa.Text()), autoincrement=False, nullable=True),
|
|
)
|
|
op.add_column("steps", sa.Column("tid", sa.String(), nullable=True))
|
|
|
|
# Add new foreign key constraint for provider_id
|
|
op.create_foreign_key("fk_steps_organization_id", "steps", "providers", ["provider_id"], ["id"], ondelete="RESTRICT")
|
|
|
|
# Add new foreign key constraint for provider_id
|
|
op.create_foreign_key("fk_steps_provider_id", "steps", "organizations", ["organization_id"], ["id"], ondelete="RESTRICT")
|
|
|
|
# Add new foreign key constraint for provider_id
|
|
op.create_foreign_key("fk_steps_job_id", "steps", "jobs", ["job_id"], ["id"], ondelete="SET NULL")
|
|
|
|
# Drop old step_id and step_count columns which aren't in the new model
|
|
op.drop_column("steps", "step_id")
|
|
op.drop_column("steps", "step_count")
|
|
|
|
# Add step_id to messages table
|
|
op.add_column("messages", sa.Column("step_id", sa.String(), nullable=True))
|
|
op.create_foreign_key("fk_messages_step_id", "messages", "steps", ["step_id"], ["id"], ondelete="SET NULL")
|
|
# ### end Alembic commands ###
|
|
|
|
|
|
def downgrade() -> None:
|
|
# ### commands auto generated by Alembic - please adjust! ###
|
|
# Remove step_id from messages first to avoid foreign key conflicts
|
|
op.drop_constraint("fk_messages_step_id", "messages", type_="foreignkey")
|
|
op.drop_column("messages", "step_id")
|
|
|
|
# Restore old step_count and step_id column
|
|
op.add_column("steps", sa.Column("step_count", sa.Integer(), nullable=True))
|
|
op.add_column("steps", sa.Column("step_id", sa.String(), nullable=True))
|
|
|
|
# Drop new columns and constraints
|
|
op.drop_constraint("fk_steps_provider_id", "steps", type_="foreignkey")
|
|
op.drop_constraint("fk_steps_organization_id", "steps", type_="foreignkey")
|
|
op.drop_constraint("fk_steps_job_id", "steps", type_="foreignkey")
|
|
|
|
op.drop_column("steps", "tid")
|
|
op.drop_column("steps", "tags")
|
|
op.drop_column("steps", "completion_tokens_details")
|
|
op.drop_column("steps", "context_window_limit")
|
|
op.drop_column("steps", "model")
|
|
op.drop_column("steps", "provider_name")
|
|
op.drop_column("steps", "provider_id")
|
|
op.drop_column("steps", "organization_id")
|
|
op.drop_column("steps", "origin")
|
|
|
|
# Add constraints back
|
|
op.execute("DELETE FROM steps WHERE job_id IS NULL")
|
|
op.alter_column("steps", "job_id", nullable=False)
|
|
op.create_foreign_key("fk_job_usage_statistics_job_id", "steps", "jobs", ["job_id"], ["id"], ondelete="CASCADE")
|
|
|
|
# Change id field from string back to int
|
|
op.add_column("steps", sa.Column("old_id", sa.Integer(), nullable=True))
|
|
op.execute("""UPDATE steps SET old_id = CAST(ABS(hashtext(REPLACE(id, 'step-', '')::text)) AS integer)""")
|
|
op.drop_column("steps", "id")
|
|
op.execute("ALTER TABLE steps RENAME COLUMN old_id TO id")
|
|
op.alter_column("steps", "id", nullable=False)
|
|
op.create_primary_key("pk_steps_id", "steps", ["id"])
|
|
|
|
# Rename the table
|
|
op.rename_table("steps", "job_usage_statistics")
|
|
# ### end Alembic commands ###
|