125 lines
5.8 KiB
Python
125 lines
5.8 KiB
Python
"""initial schema -- users / tasks / messages / runs / usage_events
|
|
|
|
Revision ID: 0001
|
|
Revises:
|
|
Create Date: 2026-05-14
|
|
|
|
DESIGN.md section 7.4 schema. First migration.
|
|
- pgcrypto extension fallback (PG 13+ has gen_random_uuid built-in;
|
|
older versions need the extension).
|
|
- messages.payload GIN index (jsonb_path_ops).
|
|
- tasks (user_id, task_dir) and (user_id, status) composite indexes.
|
|
- users rows are INSERTed on demand by web auth entry points
|
|
(`web.auth.ensure_user_row`), NOT in this migration.
|
|
"""
|
|
from typing import Sequence, Union
|
|
|
|
from alembic import op
|
|
import sqlalchemy as sa
|
|
from sqlalchemy.dialects import postgresql
|
|
|
|
revision: str = "0001"
|
|
down_revision: Union[str, None] = None
|
|
branch_labels: Union[str, Sequence[str], None] = None
|
|
depends_on: Union[str, Sequence[str], None] = None
|
|
|
|
|
|
def upgrade() -> None:
|
|
op.execute("CREATE EXTENSION IF NOT EXISTS pgcrypto")
|
|
|
|
op.create_table(
|
|
"users",
|
|
sa.Column("user_id", postgresql.UUID(as_uuid=True), primary_key=True,
|
|
server_default=sa.text("gen_random_uuid()")),
|
|
sa.Column("email", sa.Text(), nullable=True),
|
|
sa.Column("oidc_subject", sa.Text(), nullable=True),
|
|
sa.Column("password_hash", sa.Text(), nullable=True),
|
|
sa.Column("plan", sa.Text(), nullable=True),
|
|
sa.Column("created_at", sa.DateTime(timezone=True),
|
|
server_default=sa.text("now()"), nullable=False),
|
|
)
|
|
|
|
op.create_table(
|
|
"tasks",
|
|
sa.Column("task_id", postgresql.UUID(as_uuid=True), primary_key=True,
|
|
server_default=sa.text("gen_random_uuid()")),
|
|
sa.Column("user_id", postgresql.UUID(as_uuid=True),
|
|
sa.ForeignKey("users.user_id"), nullable=False),
|
|
sa.Column("task_dir", sa.Text(), nullable=False),
|
|
sa.Column("mode", sa.Text(), nullable=False, server_default=""),
|
|
sa.Column("description", sa.Text(), nullable=False, server_default=""),
|
|
sa.Column("status", sa.Text(), nullable=False, server_default="active"),
|
|
sa.Column("model", sa.Text(), nullable=False, server_default=""),
|
|
sa.Column("model_profile", sa.Text(), nullable=False, server_default=""),
|
|
sa.Column("reasoning_effort", sa.Text(), nullable=False, server_default=""),
|
|
sa.Column("tokens_prompt", sa.Integer(), nullable=False, server_default="0"),
|
|
sa.Column("tokens_completion", sa.Integer(), nullable=False, server_default="0"),
|
|
sa.Column("cost_usd", sa.Numeric(12, 6), nullable=False, server_default="0"),
|
|
sa.Column("created_at", sa.DateTime(timezone=True),
|
|
server_default=sa.text("now()"), nullable=False),
|
|
sa.Column("updated_at", sa.DateTime(timezone=True),
|
|
server_default=sa.text("now()"), nullable=False),
|
|
)
|
|
op.create_index("ix_tasks_user_task_dir", "tasks", ["user_id", "task_dir"])
|
|
op.create_index("ix_tasks_user_status", "tasks", ["user_id", "status"])
|
|
|
|
op.create_table(
|
|
"messages",
|
|
sa.Column("message_id", postgresql.UUID(as_uuid=True), primary_key=True,
|
|
server_default=sa.text("gen_random_uuid()")),
|
|
sa.Column("task_id", postgresql.UUID(as_uuid=True),
|
|
sa.ForeignKey("tasks.task_id", ondelete="CASCADE"), nullable=False),
|
|
sa.Column("idx", sa.Integer(), nullable=False),
|
|
sa.Column("payload", postgresql.JSONB(), nullable=False),
|
|
sa.Column("tokens_in", sa.Integer(), nullable=True),
|
|
sa.Column("tokens_out", sa.Integer(), nullable=True),
|
|
sa.Column("created_at", sa.DateTime(timezone=True),
|
|
server_default=sa.text("now()"), nullable=False),
|
|
sa.UniqueConstraint("task_id", "idx", name="uq_messages_task_idx"),
|
|
)
|
|
op.create_index(
|
|
"ix_messages_payload_gin", "messages", ["payload"],
|
|
postgresql_using="gin", postgresql_ops={"payload": "jsonb_path_ops"},
|
|
)
|
|
|
|
op.create_table(
|
|
"runs",
|
|
sa.Column("run_id", postgresql.UUID(as_uuid=True), primary_key=True,
|
|
server_default=sa.text("gen_random_uuid()")),
|
|
sa.Column("task_id", postgresql.UUID(as_uuid=True),
|
|
sa.ForeignKey("tasks.task_id", ondelete="CASCADE"), nullable=False),
|
|
sa.Column("status", sa.Text(), nullable=False, server_default="pending"),
|
|
sa.Column("started_at", sa.DateTime(timezone=True), nullable=True),
|
|
sa.Column("finished_at", sa.DateTime(timezone=True), nullable=True),
|
|
sa.Column("error", sa.Text(), nullable=True),
|
|
sa.Column("tokens_p", sa.Integer(), nullable=False, server_default="0"),
|
|
sa.Column("tokens_c", sa.Integer(), nullable=False, server_default="0"),
|
|
)
|
|
op.create_index("ix_runs_task", "runs", ["task_id"])
|
|
|
|
op.create_table(
|
|
"usage_events",
|
|
sa.Column("id", sa.BigInteger(), primary_key=True, autoincrement=True),
|
|
sa.Column("user_id", postgresql.UUID(as_uuid=True), nullable=False),
|
|
sa.Column("task_id", postgresql.UUID(as_uuid=True), nullable=True),
|
|
sa.Column("run_id", postgresql.UUID(as_uuid=True), nullable=True),
|
|
sa.Column("kind", sa.Text(), nullable=False),
|
|
sa.Column("value", sa.Numeric(20, 8), nullable=False),
|
|
sa.Column("ts", sa.DateTime(timezone=True),
|
|
server_default=sa.text("now()"), nullable=False),
|
|
)
|
|
op.create_index("ix_usage_user_ts", "usage_events", ["user_id", "ts"])
|
|
|
|
|
|
def downgrade() -> None:
|
|
op.drop_index("ix_usage_user_ts", table_name="usage_events")
|
|
op.drop_table("usage_events")
|
|
op.drop_index("ix_runs_task", table_name="runs")
|
|
op.drop_table("runs")
|
|
op.drop_index("ix_messages_payload_gin", table_name="messages")
|
|
op.drop_table("messages")
|
|
op.drop_index("ix_tasks_user_status", table_name="tasks")
|
|
op.drop_index("ix_tasks_user_task_dir", table_name="tasks")
|
|
op.drop_table("tasks")
|
|
op.drop_table("users")
|