zhifu-platform/migrations/versions/0013_token_budget.py

import sqlalchemy as sa
from alembic import op
from sqlalchemy.dialects import mysql

revision = "0013_token_budget"
down_revision = "0012_tenant_reservations"
branch_labels = None
depends_on = None


def upgrade():
    op.create_table(
        "token_policies",
        sa.Column("tenant_id", mysql.BINARY(16), primary_key=True),
        sa.Column("token_limit", mysql.BIGINT(unsigned=True), nullable=False),
        sa.Column("version", mysql.BIGINT(unsigned=True), nullable=False),
        sa.ForeignKeyConstraint(["tenant_id"], ["tenants.id"], ondelete="CASCADE"),
    )
    op.create_table(
        "token_accounts",
        sa.Column("tenant_id", mysql.BINARY(16), primary_key=True),
        sa.Column("period", mysql.DATE(), primary_key=True),
        sa.Column("token_limit", mysql.BIGINT(unsigned=True), nullable=False),
        sa.Column("reserved", mysql.BIGINT(unsigned=True), nullable=False),
        sa.Column("consumed", mysql.BIGINT(unsigned=True), nullable=False),
        sa.ForeignKeyConstraint(["tenant_id"], ["tenants.id"], ondelete="CASCADE"),
    )
    op.create_table(
        "token_entries",
        sa.Column("id", mysql.BINARY(16), primary_key=True),
        sa.Column("tenant_id", mysql.BINARY(16), nullable=False),
        sa.Column("period", mysql.DATE(), nullable=False),
        sa.Column("kind", sa.String(16), nullable=False),
        sa.Column("reference_id", mysql.BINARY(16), nullable=False),
        sa.Column("model_name", sa.String(128), nullable=False),
        sa.Column("reserved", mysql.BIGINT(unsigned=True), nullable=False),
        sa.Column("usage", sa.JSON(), nullable=True),
        sa.Column("status", sa.String(16), nullable=False),
        sa.Column("version", mysql.BIGINT(unsigned=True), nullable=False),
        sa.Column("created_at", mysql.DATETIME(fsp=6), nullable=False),
        sa.ForeignKeyConstraint(
            ["tenant_id", "period"],
            ["token_accounts.tenant_id", "token_accounts.period"],
            ondelete="CASCADE",
        ),
        sa.CheckConstraint("status IN ('PENDING','UNKNOWN','SETTLED','RELEASED')", name="status"),
    )
    op.create_index(
        "ix_token_entries_tenant_created", "token_entries", ["tenant_id", "created_at", "id"]
    )


def downgrade():
    op.drop_table("token_entries")
    op.drop_table("token_accounts")
    op.drop_table("token_policies")