baodan/api/insurance/db/migrate_001.py

"""迁移脚本 001：重命名旧字段为新字段

使用方式：
    python -m insurance.db.migrate_001

或在应用启动时自动执行（需要配置 MIGRATION_ENABLED=true）
"""
from sqlalchemy import text
from insurance.db.compat import db


def migrate():
    """执行迁移：重命名旧字段为新字段。"""
    print("[migrate_001] 开始执行迁移...")

    try:
        # 检查字段是否已经迁移（防止重复执行）
        result = db.session.execute(text("""
            SELECT column_name
            FROM information_schema.columns
            WHERE table_name = 'wecom_user_mapping'
            AND column_name = 'internal_user_id'
        """))

        if result.fetchone():
            print("[migrate_001] 迁移已执行过，跳过")
            return True

        # 开始事务
        db.session.begin()

        # 1. 重命名 wecom_user_mapping 表的旧字段
        print("[migrate_001] 重命名 wecom_user_mapping.dify_user_id -> internal_user_id")

        # 添加新字段
        db.session.execute(text("""
            ALTER TABLE wecom_user_mapping
            ADD COLUMN internal_user_id VARCHAR(64)
        """))

        # 复制数据
        db.session.execute(text("""
            UPDATE wecom_user_mapping
            SET internal_user_id = dify_user_id
        """))

        # 设置 NOT NULL 约束
        db.session.execute(text("""
            ALTER TABLE wecom_user_mapping
            ALTER COLUMN internal_user_id SET NOT NULL
        """))

        # 创建新索引
        db.session.execute(text("""
            DROP INDEX IF EXISTS idx_wecom_user_mapping_dify_id
        """))
        db.session.execute(text("""
            CREATE INDEX idx_wecom_user_mapping_internal_id
            ON wecom_user_mapping(internal_user_id)
        """))

        # 删除旧字段
        db.session.execute(text("""
            ALTER TABLE wecom_user_mapping
            DROP COLUMN dify_user_id
        """))

        # 更新字段注释
        db.session.execute(text("""
            COMMENT ON COLUMN wecom_user_mapping.internal_user_id
            IS '系统内部用户 ID'
        """))

        # 2. 重命名 recommendation_records 表的旧字段
        print("[migrate_001] 重命名 recommendation_records.dify_task_id -> task_id")

        # 检查旧字段是否存在
        result = db.session.execute(text("""
            SELECT column_name
            FROM information_schema.columns
            WHERE table_name = 'recommendation_records'
            AND column_name = 'dify_task_id'
        """))

        if result.fetchone():
            # 添加新字段
            db.session.execute(text("""
                ALTER TABLE recommendation_records
                ADD COLUMN task_id VARCHAR(64)
            """))

            # 复制数据
            db.session.execute(text("""
                UPDATE recommendation_records
                SET task_id = dify_task_id
            """))

            # 删除旧字段
            db.session.execute(text("""
                ALTER TABLE recommendation_records
                DROP COLUMN dify_task_id
            """))

        # 更新表注释
        db.session.execute(text("""
            COMMENT ON TABLE wecom_user_mapping
            IS '企微用户与内部系统用户映射表'
        """))

        # 提交事务
        db.session.commit()

        print("[migrate_001] 迁移完成！")
        return True

    except Exception as e:
        db.session.rollback()
        print(f"[migrate_001] 迁移失败: {e}")
        return False


if __name__ == "__main__":
    from flask import Flask
    from configs.app_config import AppConfig

    app = Flask(__name__)
    app.config.from_object(AppConfig)

    with app.app_context():
        migrate()