"""扩展 ck_mode_enum 约束支持 baseline 模式

Revision ID: 0017_mode_baseline
Revises: 0016_schedules
Create Date: 2026-09-21

Code Review High-1/2 修复:
  业务支持 mode='baseline'(非 LLM 统计基线),但 DB CHECK 约束只允许
  ('single', 'multi'),导致 baseline 预测写入时 CheckViolation。
  run_type='baseline' 改为 'live'(代码侧),约束无需扩展。
"""

from typing import Sequence, Union

from alembic import op

# revision identifiers, used by Alembic.
revision: str = '0017_mode_baseline'
down_revision: Union[str, None] = '0016_schedules'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    op.drop_constraint('ck_mode_enum', 'predictions', type_='check')
    op.create_check_constraint(
        'ck_mode_enum', 'predictions',
        "mode IN ('single', 'multi', 'baseline')",
    )


def downgrade() -> None:
    # 先清理 baseline 数据再恢复旧约束
    op.execute("DELETE FROM predictions WHERE mode = 'baseline'")
    op.drop_constraint('ck_mode_enum', 'predictions', type_='check')
    op.create_check_constraint(
        'ck_mode_enum', 'predictions',
        "mode IN ('single', 'multi')",
    )