feat: add schema/db for new steps table (#669)

alembic/versions/416b9d2db10b_repurpose_jobusagestatistics_for_new_.py

@@ -0,0 +1,116 @@
+"""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 ###