Kamixitong/migrations/versions/20251212_add_security_constraints.py

"""添加安全约束

Revision ID: 20251212_add_security_constraints
Revises:
Create Date: 2025-12-12 00:00:00.000000

"""
from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql

# revision identifiers
revision = '20251212_add_security_constraints'
down_revision = None
branch_labels = None
depends_on = None


def upgrade():
    """
    添加数据库层面的安全约束
    """
    # 为License表添加约束
    # 1. 解绑次数约束
    op.create_check_constraint(
        'chk_license_unbind_count',
        'license',
        'unbind_count >= 0 AND unbind_count <= 10'  # 解绑次数不能超过10次
    )

    # 2. 有效天数约束
    op.create_check_constraint(
        'chk_license_valid_days',
        'license',
        'valid_days > 0 OR valid_days = -1'  # 有效天数必须为正数或-1（永久有效）
    )

    # 3. 状态值约束
    op.create_check_constraint(
        'chk_license_status',
        'license',
        'status IN (0, 1, 2, 3)'  # 0=未激活，1=已激活，2=已过期，3=已禁用
    )

    # 4. 绑定次数约束
    op.create_check_constraint(
        'chk_license_max_bind_times',
        'license',
        'max_bind_times > 0 AND max_bind_times <= 100'  # 最大绑定次数限制
    )

    # 为Product表添加约束
    # 1. 价格约束
    op.create_check_constraint(
        'chk_product_price',
        'product',
        'price >= 0'  # 价格不能为负数
    )

    # 2. 状态值约束
    op.create_check_constraint(
        'chk_product_status',
        'product',
        'status IN (0, 1)'  # 0=禁用，1=启用
    )

    # 为Order表添加约束
    # 1. 金额约束
    op.create_check_constraint(
        'chk_order_amount',
        'order',
        'amount > 0'  # 订单金额必须为正数
    )

    # 2. 状态值约束
    op.create_check_constraint(
        'chk_order_status',
        'order',
        'status IN (0, 1, 2, 3, 4)'  # 0=待支付，1=已支付，2=已取消，3=已退款，4=已完成
    )

    # 为Ticket表添加约束
    # 1. 优先级约束
    op.create_check_constraint(
        'chk_ticket_priority',
        'ticket',
        'priority IN (1, 2, 3, 4, 5)'  # 1=低，2=中，3=高，4=紧急，5=非常重要
    )

    # 2. 状态值约束
    op.create_check_constraint(
        'chk_ticket_status',
        'ticket',
        'status IN (0, 1, 2, 3)'  # 0=待处理，1=处理中，2=已解决，3=已关闭
    )

    # 添加索引优化查询性能
    # 1. License表的索引
    op.create_index('idx_license_product_status', 'license', ['product_id', 'status'])
    op.create_index('idx_license_type_status', 'license', ['type', 'status'])
    op.create_index('idx_license_key', 'license', ['license_key'], unique=True)

    # 2. Product表的索引
    op.create_index('idx_product_status_type', 'product', ['status', 'product_type'])

    # 3. Order表的索引
    op.create_index('idx_order_user_phone', 'order', ['user_phone'])
    op.create_index('idx_order_status_time', 'order', ['status', 'create_time'])

    # 4. Ticket表的索引
    op.create_index('idx_ticket_user_phone', 'ticket', ['user_phone'])
    op.create_index('idx_ticket_status_priority', 'ticket', ['status', 'priority'])

    # 5. Device表的索引
    op.create_index('idx_device_machine_code', 'device', ['machine_code'])
    op.create_index('idx_device_license', 'device', ['license_id'])


def downgrade():
    """
    删除安全约束
    """
    # 删除索引
    op.drop_index('idx_device_license', 'device')
    op.drop_index('idx_device_machine_code', 'device')
    op.drop_index('idx_ticket_status_priority', 'ticket')
    op.drop_index('idx_ticket_user_phone', 'ticket')
    op.drop_index('idx_order_status_time', 'order')
    op.drop_index('idx_order_user_phone', 'order')
    op.drop_index('idx_product_status_type', 'product')
    op.drop_index('idx_license_key', 'license')
    op.drop_index('idx_license_type_status', 'license')
    op.drop_index('idx_license_product_status', 'license')

    # 删除约束
    op.drop_constraint('chk_ticket_status', 'ticket', type_='check')
    op.drop_constraint('chk_ticket_priority', 'ticket', type_='check')
    op.drop_constraint('chk_order_status', 'order', type_='check')
    op.drop_constraint('chk_order_amount', 'order', type_='check')
    op.drop_constraint('chk_product_status', 'product', type_='check')
    op.drop_constraint('chk_product_price', 'product', type_='check')
    op.drop_constraint('chk_license_max_bind_times', 'license', type_='check')
    op.drop_constraint('chk_license_status', 'license', type_='check')
    op.drop_constraint('chk_license_valid_days', 'license', type_='check')
    op.drop_constraint('chk_license_unbind_count', 'license', type_='check')