open-webui/backend/open_webui/migrations/versions/240e45fa2f01_add_precharge_fields.py

73 lines
2.6 KiB
Python
Raw Normal View History

"""Add precharge fields to billing_log
Revision ID: 240e45fa2f01
Revises: f8c9d0e4a3b2
Create Date: 2025-12-06 18:00:00.000000
添加预扣费相关字段到billing_log表
- precharge_id: 预扣费事务IDUUID
- status: 记录状态precharge | settled | refunded
- estimated_tokens: 预估tokens总数
- refund_amount: 退款金额
"""
from alembic import op
import sqlalchemy as sa
revision = "240e45fa2f01"
down_revision = "f8c9d0e4a3b2"
branch_labels = None
depends_on = None
def upgrade():
"""升级数据库:添加预扣费字段"""
connection = op.get_bind()
is_sqlite = connection.dialect.name == "sqlite"
if is_sqlite:
# SQLite: 使用batch模式
with op.batch_alter_table("billing_log") as batch_op:
batch_op.add_column(sa.Column("precharge_id", sa.String(), nullable=True))
batch_op.add_column(
sa.Column("status", sa.String(20), nullable=True, server_default="settled")
)
batch_op.add_column(sa.Column("estimated_tokens", sa.Integer(), nullable=True))
batch_op.add_column(sa.Column("refund_amount", sa.Integer(), nullable=True))
batch_op.create_index("ix_billing_log_precharge_id", ["precharge_id"])
else:
# PostgreSQL: 直接操作
op.add_column("billing_log", sa.Column("precharge_id", sa.String(), nullable=True))
op.add_column(
"billing_log",
sa.Column("status", sa.String(20), nullable=True, server_default="settled"),
)
op.add_column(
"billing_log", sa.Column("estimated_tokens", sa.Integer(), nullable=True)
)
op.add_column(
"billing_log", sa.Column("refund_amount", sa.Integer(), nullable=True)
)
op.create_index("ix_billing_log_precharge_id", "billing_log", ["precharge_id"])
def downgrade():
"""降级数据库:删除预扣费字段"""
connection = op.get_bind()
is_sqlite = connection.dialect.name == "sqlite"
if is_sqlite:
with op.batch_alter_table("billing_log") as batch_op:
batch_op.drop_index("ix_billing_log_precharge_id")
batch_op.drop_column("refund_amount")
batch_op.drop_column("estimated_tokens")
batch_op.drop_column("status")
batch_op.drop_column("precharge_id")
else:
op.drop_index("ix_billing_log_precharge_id", table_name="billing_log")
op.drop_column("billing_log", "refund_amount")
op.drop_column("billing_log", "estimated_tokens")
op.drop_column("billing_log", "status")
op.drop_column("billing_log", "precharge_id")