scripts/migrate_5004789_to_9149bd8.py

#!/usr/bin/env python
"""数据库迁移脚本

从 commit 5004789ff44222d41da931d4423cfb5d0e41e053
迁移到 commit 9149bd803663748f157986544eca876543b326dc

主要变更：
1. annotation_results 表：
   - dataset_id, dataset_item_id, annotation_config_id 添加外键约束
   - annotator_id 从 String 类型改为 Integer 类型，并添加外键约束
2. datasets 表：
   - creator_id 从 String 类型改为 Integer 类型，并添加外键约束
"""

import sys
from pathlib import Path

# 添加项目根目录到路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

# 注意：以下导入必须在 sys.path.insert 之后，因为需要导入项目模块
import logging  # noqa: E402
import shutil  # noqa: E402

from sqlalchemy import inspect, text  # noqa: E402
from sqlalchemy.engine import Engine  # noqa: E402
from sqlalchemy.orm import Session  # noqa: E402

from qa_annotate.database.base import DB_PATH, engine  # noqa: E402

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
logger = logging.getLogger(__name__)


def backup_database():
    """备份数据库"""
    from datetime import datetime

    backup_path = (
        DB_PATH.parent
        / f"annotations_backup_{datetime.now().strftime('%Y%m%d_%H%M%S')}.db"
    )
    shutil.copy2(DB_PATH, backup_path)
    logger.info(f"数据库已备份到: {backup_path}")
    return backup_path


def disable_foreign_keys(conn):
    """禁用外键约束"""
    conn.execute(text("PRAGMA foreign_keys = OFF"))
    logger.info("已禁用外键约束")


def enable_foreign_keys(conn):
    """启用外键约束"""
    conn.execute(text("PRAGMA foreign_keys = ON"))
    logger.info("已启用外键约束")


def restore_from_backup(backup_path: Path):
    """从备份恢复数据库"""
    logger.warning("=" * 60)
    logger.warning("开始从备份恢复数据库...")
    logger.warning(f"备份路径: {backup_path}")
    logger.warning("=" * 60)

    if not backup_path.exists():
        logger.error(f"备份文件不存在: {backup_path}")
        return False

    try:
        shutil.copy2(backup_path, DB_PATH)
        logger.info("数据库已从备份恢复")
        return True
    except Exception as e:
        logger.error(f"恢复数据库失败: {e}")
        return False


def get_table_record_counts(engine: Engine) -> dict[str, int]:
    """获取所有用户表的记录数量

    Returns:
        dict: {表名: 记录数量}
    """
    counts = {}

    with engine.connect() as conn:
        inspector = inspect(engine)
        table_names = inspector.get_table_names()

        # 排除系统表
        system_tables = {"sqlite_sequence", "sqlite_master"}
        user_tables = [t for t in table_names if t not in system_tables]

        for table_name in user_tables:
            try:
                result = conn.execute(text(f'SELECT COUNT(*) FROM "{table_name}"'))
                count = result.scalar()
                counts[table_name] = count
            except Exception as e:
                logger.warning(f"获取 {table_name} 表记录数量失败: {e}")
                counts[table_name] = -1  # 使用 -1 表示无法获取

    return counts


def validate_data_integrity(
    engine: Engine, before_counts: dict[str, int] | None = None
) -> tuple[bool, list[str]]:
    """校验数据是否丢失

    检查所有表的必填字段（NOT NULL）是否有 NULL 值，以及迁移涉及的特殊字段类型是否正确。
    如果提供了迁移前的记录数量，会比较迁移前后的记录数量。

    Args:
        engine: 数据库引擎
        before_counts: 迁移前的记录数量字典 {表名: 数量}

    Returns:
        tuple: (是否通过校验, 错误信息列表)
    """
    logger.info("开始校验数据是否丢失...")
    errors = []

    with engine.connect() as conn:
        inspector = inspect(engine)
        table_names = inspector.get_table_names()

        # 排除系统表
        system_tables = {"sqlite_sequence", "sqlite_master"}
        user_tables = [t for t in table_names if t not in system_tables]

        logger.info(f"发现 {len(user_tables)} 个用户表，开始校验...")

        # 获取迁移后的记录数量
        after_counts = get_table_record_counts(engine)

        for table_name in user_tables:
            logger.info(f"校验 {table_name} 表...")

            try:
                # 获取当前表的记录数量
                current_count = after_counts.get(table_name, 0)
                logger.info(f"  {table_name} 表当前记录数: {current_count}")

                # 如果提供了迁移前的数量，比较记录数量
                if before_counts is not None and table_name in before_counts:
                    before_count = before_counts[table_name]
                    if before_count >= 0:  # 如果迁移前能获取到数量
                        if current_count < before_count:
                            errors.append(
                                f"{table_name} 表记录数量减少: 迁移前 {before_count} 条, "
                                f"迁移后 {current_count} 条, 丢失 {before_count - current_count} 条"
                            )
                        elif current_count > before_count:
                            logger.warning(
                                f"{table_name} 表记录数量增加: 迁移前 {before_count} 条, "
                                f"迁移后 {current_count} 条（可能是正常的）"
                            )
                        else:
                            logger.info(
                                f"  {table_name} 表记录数量一致: {before_count} 条"
                            )

                # 获取表的列信息
                columns = inspector.get_columns(table_name)
                not_null_columns = [
                    col["name"]
                    for col in columns
                    if not col.get("nullable", True)
                    and not col.get("primary_key", False)
                ]

                # 检查 NOT NULL 字段是否有 NULL 值
                if not_null_columns:
                    # 构建 WHERE 条件（使用双引号包裹列名以确保安全）
                    null_conditions = " OR ".join(
                        [f'"{col}" IS NULL' for col in not_null_columns]
                    )
                    # 使用参数化查询的表名（SQLite 不支持表名参数化，但表名来自元数据，相对安全）
                    result = conn.execute(
                        text(f"""
                            SELECT COUNT(*) FROM "{table_name}"
                            WHERE {null_conditions}
                        """)
                    )
                    null_count = result.scalar()
                    if null_count > 0:
                        errors.append(
                            f"{table_name} 表有 {null_count} 条记录的必填字段为 NULL（可能数据丢失）"
                        )

                # 特殊字段类型检查（迁移涉及的字段）
                if table_name == "annotation_results":
                    # 检查 annotator_id 类型（应该是 INTEGER 或 NULL）
                    if any(col["name"] == "annotator_id" for col in columns):
                        result = conn.execute(
                            text("""
                                SELECT COUNT(*) FROM annotation_results
                                WHERE annotator_id IS NOT NULL
                                  AND typeof(annotator_id) != 'integer'
                            """)
                        )
                        invalid_type_count = result.scalar()
                        if invalid_type_count > 0:
                            errors.append(
                                f"annotation_results 表有 {invalid_type_count} 条记录的 annotator_id 类型不正确"
                            )

                elif table_name == "datasets":
                    # 检查 creator_id 类型（应该是 INTEGER 或 NULL）
                    if any(col["name"] == "creator_id" for col in columns):
                        result = conn.execute(
                            text("""
                                SELECT COUNT(*) FROM datasets
                                WHERE creator_id IS NOT NULL
                                  AND typeof(creator_id) != 'integer'
                            """)
                        )
                        invalid_type_count = result.scalar()
                        if invalid_type_count > 0:
                            errors.append(
                                f"datasets 表有 {invalid_type_count} 条记录的 creator_id 类型不正确"
                            )

            except Exception as e:
                logger.warning(f"校验 {table_name} 表时发生错误: {e}")
                errors.append(f"{table_name} 表校验失败: {e}")

    # 打印记录数量统计
    logger.info("=" * 60)
    logger.info("记录数量统计:")
    for table_name in sorted(after_counts.keys()):
        count = after_counts[table_name]
        if before_counts and table_name in before_counts:
            before_count = before_counts[table_name]
            status = "✓" if count == before_count else "✗"
            logger.info(
                f"  {status} {table_name}: {count} 条"
                + (f" (迁移前: {before_count} 条)" if before_count >= 0 else "")
            )
        else:
            logger.info(f"  {table_name}: {count} 条")
    logger.info("=" * 60)

    if errors:
        logger.error("数据丢失校验失败:")
        for error in errors:
            logger.error(f"  - {error}")
        return False, errors
    else:
        logger.info("数据丢失校验通过")
        return True, []


def check_and_clean_foreign_key_violations(
    engine: Engine,
) -> tuple[dict[str, int], int]:
    """检查外键约束、统计并清理违反外键约束的记录

    对于违反外键约束的记录：
    - 如果外键字段允许为 NULL，则设置为 NULL
    - 如果外键字段不允许为 NULL（NOT NULL 约束），则删除该记录

    Returns:
        tuple: (违反约束统计字典 {表名: 违反数量}, 处理的记录总数（设置为 NULL + 删除）)
    """
    logger.info("=" * 60)
    logger.info("开始检查外键约束...")
    logger.info("=" * 60)

    violations_stats = {}
    total_cleaned = 0

    with engine.begin() as conn:
        # 启用外键约束
        enable_foreign_keys(conn)

        # 检查外键约束违反
        logger.info("检查外键约束违反...")
        result = conn.execute(text("PRAGMA foreign_key_check"))
        fk_violations = result.fetchall()

        if not fk_violations:
            logger.info("未发现外键约束违反")
            return {}, 0

        # 统计违反情况
        for violation in fk_violations:
            table_name = violation[0]

            if table_name not in violations_stats:
                violations_stats[table_name] = 0
            violations_stats[table_name] += 1

        # 打印统计信息
        logger.info(f"发现 {len(fk_violations)} 条外键约束违反记录")
        logger.info("=" * 60)
        logger.info("外键约束违反统计:")
        for table_name, count in violations_stats.items():
            logger.info(f"  {table_name}: {count} 条记录")
        logger.info("=" * 60)

        # 清理违反外键约束的记录
        logger.info("开始清理违反外键约束的记录...")

        inspector = inspect(engine)

        # 按表分组处理
        violations_by_table = {}
        for violation in fk_violations:
            table_name = violation[0]
            row_id = violation[1]
            parent_table = violation[3] if len(violation) > 3 else None

            if table_name not in violations_by_table:
                violations_by_table[table_name] = []
            violations_by_table[table_name].append((row_id, parent_table))

        # 对每个表进行清理
        for table_name, violations in violations_by_table.items():
            logger.info(f"清理 {table_name} 表的违反记录...")

            try:
                # 获取表的外键信息和列信息
                fks = inspector.get_foreign_keys(table_name)
                columns = inspector.get_columns(table_name)
                column_info = {col["name"]: col for col in columns}

                # 对于每个违反的记录，找到对应的外键列并清理
                cleaned_count = 0
                deleted_count = 0
                for row_id, parent_table in violations:
                    # 找到指向 parent_table 的外键
                    target_fk = None
                    for fk in fks:
                        if fk["referred_table"] == parent_table:
                            target_fk = fk
                            break

                    if not target_fk and fks:
                        # 如果找不到，使用第一个外键（通常只有一个）
                        target_fk = fks[0]

                    if target_fk and target_fk["constrained_columns"]:
                        fk_column = target_fk["constrained_columns"][0]

                        # 检查该字段是否允许为 NULL
                        fk_column_info = column_info.get(fk_column)
                        is_nullable = (
                            fk_column_info is not None
                            and fk_column_info.get("nullable", True)
                            and not fk_column_info.get("primary_key", False)
                        )

                        try:
                            if is_nullable:
                                # 如果字段允许为 NULL，设置为 NULL
                                conn.execute(
                                    text(f"""
                                        UPDATE "{table_name}"
                                        SET "{fk_column}" = NULL
                                        WHERE rowid = :row_id
                                    """),
                                    {"row_id": row_id},
                                )
                                cleaned_count += 1
                                logger.debug(
                                    f"  已清理 {table_name} 表 rowid={row_id} 的 {fk_column} 字段（设置为 NULL）"
                                )
                            else:
                                # 如果字段不允许为 NULL，删除该记录
                                conn.execute(
                                    text(f"""
                                        DELETE FROM "{table_name}"
                                        WHERE rowid = :row_id
                                    """),
                                    {"row_id": row_id},
                                )
                                deleted_count += 1
                                logger.debug(
                                    f"  已删除 {table_name} 表 rowid={row_id} 的记录（{fk_column} 字段为 NOT NULL）"
                                )
                        except Exception as e:
                            logger.warning(
                                f"  清理 {table_name} 表 rowid={row_id} 失败: {e}"
                            )
                    else:
                        logger.warning(
                            f"  无法找到 {table_name} 表 rowid={row_id} 的外键列"
                        )

                total_cleaned += cleaned_count + deleted_count
                if cleaned_count > 0 or deleted_count > 0:
                    logger.info(
                        f"  {table_name} 表已清理 {cleaned_count} 条记录（设置为 NULL），"
                        f"删除 {deleted_count} 条记录（NOT NULL 约束）"
                    )

            except Exception as e:
                logger.error(f"清理 {table_name} 表时发生错误: {e}")

        logger.info("=" * 60)
        logger.info(f"外键约束清理完成，共处理 {total_cleaned} 条记录")
        logger.info("=" * 60)

    return violations_stats, total_cleaned


def get_user_id_by_username(db: Session, username: str) -> int | None:
    """根据用户名获取用户ID"""
    if not username:
        return None

    result = db.execute(
        text("SELECT id FROM users WHERE username = :username"), {"username": username}
    ).first()

    if result:
        return result[0]
    return None


def migrate_annotation_results(engine: Engine):
    """迁移 annotation_results 表"""
    logger.info("开始迁移 annotation_results 表...")

    with engine.begin() as conn:
        # 禁用外键约束以避免迁移过程中的数据丢失
        disable_foreign_keys(conn)
        # 检查表是否存在
        inspector = inspect(engine)
        if "annotation_results" not in inspector.get_table_names():
            logger.warning("annotation_results 表不存在，跳过迁移")
            enable_foreign_keys(conn)
            return

        # 检查 annotator_id 列的类型
        columns = {
            col["name"]: col for col in inspector.get_columns("annotation_results")
        }

        if "annotator_id" not in columns:
            logger.warning("annotation_results 表没有 annotator_id 列，跳过迁移")
            enable_foreign_keys(conn)
            return

        annotator_id_type = str(columns["annotator_id"]["type"])
        is_string_type = (
            "VARCHAR" in annotator_id_type
            or "TEXT" in annotator_id_type
            or "CHAR" in annotator_id_type
        )

        if not is_string_type:
            logger.info("annotator_id 已经是 Integer 类型，检查外键约束...")
            # 检查是否已有外键约束
            fks = inspector.get_foreign_keys("annotation_results")
            has_fk = any(fk["constrained_columns"] == ["annotator_id"] for fk in fks)
            if has_fk:
                logger.info("annotation_results 表已经迁移完成，跳过")
                enable_foreign_keys(conn)
                return

        # 清理可能存在的旧表（如果之前的迁移失败）
        logger.info("清理可能存在的旧表...")
        conn.execute(text("DROP TABLE IF EXISTS annotation_results_new"))

        # 清理可能存在的旧索引
        for index_name in [
            "ix_annotation_results_dataset_id",
            "ix_annotation_results_dataset_item_id",
            "ix_annotation_results_annotation_config_id",
            "ix_annotation_results_annotator_id",
        ]:
            conn.execute(text(f"DROP INDEX IF EXISTS {index_name}"))

        # 创建新表
        logger.info("创建新的 annotation_results 表...")
        conn.execute(
            text("""
            CREATE TABLE annotation_results_new (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                dataset_id INTEGER NOT NULL,
                dataset_item_id INTEGER NOT NULL,
                annotation_config_id INTEGER NOT NULL,
                value_json TEXT NOT NULL,
                annotator_id INTEGER,
                annotator_name TEXT,
                created_at DATETIME NOT NULL,
                updated_at DATETIME NOT NULL,
                duration_seconds REAL,
                confidence REAL,
                notes TEXT,
                custom_fields_json TEXT,
                FOREIGN KEY (dataset_id) REFERENCES datasets(id) ON DELETE CASCADE,
                FOREIGN KEY (dataset_item_id) REFERENCES qa_pairs(id) ON DELETE CASCADE,
                FOREIGN KEY (annotation_config_id) REFERENCES annotation_configs(id) ON DELETE CASCADE,
                FOREIGN KEY (annotator_id) REFERENCES users(id) ON DELETE SET NULL
            )
        """)
        )

        # 创建索引
        logger.info("创建索引...")
        conn.execute(
            text(
                "CREATE INDEX ix_annotation_results_dataset_id ON annotation_results_new(dataset_id)"
            )
        )
        conn.execute(
            text(
                "CREATE INDEX ix_annotation_results_dataset_item_id ON annotation_results_new(dataset_item_id)"
            )
        )
        conn.execute(
            text(
                "CREATE INDEX ix_annotation_results_annotation_config_id ON annotation_results_new(annotation_config_id)"
            )
        )
        conn.execute(
            text(
                "CREATE INDEX ix_annotation_results_annotator_id ON annotation_results_new(annotator_id)"
            )
        )

        # 迁移数据
        logger.info("迁移数据...")
        with Session(engine) as session:
            # 获取所有记录
            result = conn.execute(text("SELECT * FROM annotation_results"))
            rows = result.fetchall()
            column_names = result.keys()

            migrated_count = 0
            skipped_count = 0

            for row in rows:
                row_dict = dict(zip(column_names, row))

                # 处理 annotator_id：如果是字符串，尝试转换为用户ID
                annotator_id = row_dict.get("annotator_id")
                if annotator_id and isinstance(annotator_id, str):
                    # 尝试将字符串转换为整数（如果已经是数字字符串）
                    try:
                        annotator_id = int(annotator_id)
                    except ValueError:
                        # 如果不是数字，尝试通过用户名查找
                        user_id = get_user_id_by_username(session, annotator_id)
                        if user_id:
                            annotator_id = user_id
                        else:
                            logger.warning(
                                f"无法找到用户: {annotator_id}，将 annotator_id 设置为 NULL"
                            )
                            annotator_id = None

                # 插入到新表
                try:
                    conn.execute(
                        text("""
                            INSERT INTO annotation_results_new (
                                id, dataset_id, dataset_item_id, annotation_config_id,
                                value_json, annotator_id, annotator_name,
                                created_at, updated_at, duration_seconds,
                                confidence, notes, custom_fields_json
                            ) VALUES (
                                :id, :dataset_id, :dataset_item_id, :annotation_config_id,
                                :value_json, :annotator_id, :annotator_name,
                                :created_at, :updated_at, :duration_seconds,
                                :confidence, :notes, :custom_fields_json
                            )
                        """),
                        {
                            "id": row_dict["id"],
                            "dataset_id": row_dict["dataset_id"],
                            "dataset_item_id": row_dict["dataset_item_id"],
                            "annotation_config_id": row_dict["annotation_config_id"],
                            "value_json": row_dict.get("value_json", ""),
                            "annotator_id": annotator_id,
                            "annotator_name": row_dict.get("annotator_name"),
                            "created_at": row_dict.get("created_at"),
                            "updated_at": row_dict.get("updated_at"),
                            "duration_seconds": row_dict.get("duration_seconds"),
                            "confidence": row_dict.get("confidence"),
                            "notes": row_dict.get("notes"),
                            "custom_fields_json": row_dict.get("custom_fields_json"),
                        },
                    )
                    migrated_count += 1
                except Exception as e:
                    logger.error(f"迁移记录失败 (id={row_dict['id']}): {e}")
                    skipped_count += 1

            logger.info(
                f"数据迁移完成: {migrated_count} 条成功, {skipped_count} 条失败"
            )

        # 删除旧表
        logger.info("删除旧表...")
        conn.execute(text("DROP TABLE annotation_results"))

        # 重命名新表
        logger.info("重命名新表...")
        conn.execute(
            text("ALTER TABLE annotation_results_new RENAME TO annotation_results")
        )

        # 启用外键约束
        enable_foreign_keys(conn)

        logger.info("annotation_results 表迁移完成")


def migrate_datasets(engine: Engine):
    """迁移 datasets 表"""
    logger.info("开始迁移 datasets 表...")

    with engine.begin() as conn:
        # 禁用外键约束以避免迁移过程中的数据丢失
        disable_foreign_keys(conn)
        # 检查表是否存在
        inspector = inspect(engine)
        if "datasets" not in inspector.get_table_names():
            logger.warning("datasets 表不存在，跳过迁移")
            enable_foreign_keys(conn)
            return

        # 检查 creator_id 列的类型
        columns = {col["name"]: col for col in inspector.get_columns("datasets")}

        if "creator_id" not in columns:
            logger.warning("datasets 表没有 creator_id 列，跳过迁移")
            enable_foreign_keys(conn)
            return

        creator_id_type = str(columns["creator_id"]["type"])
        is_string_type = (
            "VARCHAR" in creator_id_type
            or "TEXT" in creator_id_type
            or "CHAR" in creator_id_type
        )

        if not is_string_type:
            logger.info("creator_id 已经是 Integer 类型，检查外键约束...")
            # 检查是否已有外键约束
            fks = inspector.get_foreign_keys("datasets")
            has_fk = any(fk["constrained_columns"] == ["creator_id"] for fk in fks)
            if has_fk:
                logger.info("datasets 表已经迁移完成，跳过")
                enable_foreign_keys(conn)
                return

        # 清理可能存在的旧表（如果之前的迁移失败）
        logger.info("清理可能存在的旧表...")
        conn.execute(text("DROP TABLE IF EXISTS datasets_new"))

        # 获取所有列信息，构建新表结构
        logger.info("创建新的 datasets 表...")

        # 构建列定义
        column_defs = []
        for col_name, col_info in columns.items():
            col_type = str(col_info["type"])
            nullable = "NULL" if col_info.get("nullable", True) else "NOT NULL"
            default = ""
            if col_info.get("default") is not None:
                default = f" DEFAULT {col_info['default']}"

            # 特殊处理 creator_id：改为 INTEGER
            if col_name == "creator_id":
                col_type = "INTEGER"

            # 处理主键
            if col_info.get("primary_key"):
                column_defs.append(f"{col_name} {col_type} PRIMARY KEY AUTOINCREMENT")
            else:
                column_defs.append(f"{col_name} {col_type} {nullable}{default}")

        # 构建 CREATE TABLE 语句
        create_sql = f"CREATE TABLE datasets_new (\n    {',\n    '.join(column_defs)}"

        # 检查是否已有外键约束
        fks = inspector.get_foreign_keys("datasets")
        existing_fk_columns = set()
        for fk in fks:
            existing_fk_columns.update(fk["constrained_columns"])

        # 添加 creator_id 的外键约束（如果还没有）
        if "creator_id" not in existing_fk_columns:
            create_sql += ",\n    FOREIGN KEY (creator_id) REFERENCES users(id) ON DELETE SET NULL"

        create_sql += "\n)"

        # 创建新表
        conn.execute(text(create_sql))

        # 复制索引
        result = conn.execute(
            text("""
            SELECT sql FROM sqlite_master
            WHERE type='index' AND tbl_name='datasets' AND sql IS NOT NULL
        """)
        )
        for row in result:
            index_sql = row[0]
            if index_sql:
                new_index_sql = index_sql.replace("datasets", "datasets_new")
                conn.execute(text(new_index_sql))

        # 迁移数据
        logger.info("迁移数据...")
        with Session(engine) as session:
            # 获取所有记录
            result = conn.execute(text("SELECT * FROM datasets"))
            rows = result.fetchall()
            column_names = result.keys()

            migrated_count = 0
            skipped_count = 0

            for row in rows:
                row_dict = dict(zip(column_names, row))

                # 处理 creator_id：如果是字符串，尝试转换为用户ID
                creator_id = row_dict.get("creator_id")
                if creator_id and isinstance(creator_id, str):
                    # 尝试将字符串转换为整数（如果已经是数字字符串）
                    try:
                        creator_id = int(creator_id)
                    except ValueError:
                        # 如果不是数字，尝试通过用户名查找
                        user_id = get_user_id_by_username(session, creator_id)
                        if user_id:
                            creator_id = user_id
                        else:
                            logger.warning(
                                f"无法找到用户: {creator_id}，将 creator_id 设置为 NULL"
                            )
                            creator_id = None

                # 构建插入语句（动态获取所有列）
                columns_str = ", ".join(column_names)
                placeholders = ", ".join([f":{col}" for col in column_names])

                # 更新 creator_id
                row_dict["creator_id"] = creator_id

                # 插入到新表
                try:
                    conn.execute(
                        text(
                            f"INSERT INTO datasets_new ({columns_str}) VALUES ({placeholders})"
                        ),
                        row_dict,
                    )
                    migrated_count += 1
                except Exception as e:
                    logger.error(
                        f"迁移记录失败 (id={row_dict.get('id', 'unknown')}): {e}"
                    )
                    skipped_count += 1

            logger.info(
                f"数据迁移完成: {migrated_count} 条成功, {skipped_count} 条失败"
            )

        # 删除旧表
        logger.info("删除旧表...")
        conn.execute(text("DROP TABLE datasets"))

        # 重命名新表
        logger.info("重命名新表...")
        conn.execute(text("ALTER TABLE datasets_new RENAME TO datasets"))

        # 启用外键约束
        enable_foreign_keys(conn)

        logger.info("datasets 表迁移完成")


def add_foreign_keys_to_annotation_results(engine: Engine):
    """为 annotation_results 表添加外键约束（如果还没有）"""
    logger.info("检查 annotation_results 表的外键约束...")

    inspector = inspect(engine)

    if "annotation_results" not in inspector.get_table_names():
        logger.warning("annotation_results 表不存在，跳过")
        return

    fks = inspector.get_foreign_keys("annotation_results")
    fk_columns = set()
    for fk in fks:
        fk_columns.update(fk["constrained_columns"])

    # 检查是否需要添加外键
    needs_fk = {
        "dataset_id": "datasets",
        "dataset_item_id": "qa_pairs",
        "annotation_config_id": "annotation_configs",
    }

    for col, ref_table in needs_fk.items():
        if col not in fk_columns:
            logger.info(f"为 {col} 添加外键约束...")
            # SQLite 不支持直接添加外键，需要重建表
            # 这里假设表已经通过 migrate_annotation_results 迁移过了
            logger.warning(
                f"{col} 的外键约束应该在表创建时添加，如果缺失可能需要重新迁移"
            )


def main():
    """主函数"""
    logger.info("=" * 60)
    logger.info("开始数据库迁移")
    logger.info(f"数据库路径: {DB_PATH}")
    logger.info("=" * 60)

    # 检查数据库文件是否存在
    if not DB_PATH.exists():
        logger.error(f"数据库文件不存在: {DB_PATH}")
        sys.exit(1)

    # 备份数据库
    logger.info("备份数据库...")
    backup_path = backup_database()
    logger.info(f"备份完成: {backup_path}")

    # 记录迁移前的记录数量
    logger.info("记录迁移前的数据统计...")
    before_counts = get_table_record_counts(engine)
    logger.info("迁移前记录数量统计:")
    for table_name in sorted(before_counts.keys()):
        count = before_counts[table_name]
        if count >= 0:
            logger.info(f"  {table_name}: {count} 条")
        else:
            logger.warning(f"  {table_name}: 无法获取数量")
    logger.info("=" * 60)

    migration_success = False
    try:
        # 迁移 annotation_results 表
        migrate_annotation_results(engine)

        # 迁移 datasets 表
        migrate_datasets(engine)

        # 检查外键约束
        add_foreign_keys_to_annotation_results(engine)

        # 数据丢失校验
        logger.info("=" * 60)
        logger.info("开始数据丢失校验...")
        logger.info("=" * 60)

        is_valid, errors = validate_data_integrity(engine, before_counts)

        if not is_valid:
            logger.error("=" * 60)
            logger.error("数据丢失校验失败！")
            logger.error("=" * 60)
            logger.error("错误详情:")
            for error in errors:
                logger.error(f"  - {error}")
            logger.error("=" * 60)
            logger.error("开始回退到备份...")
            logger.error("=" * 60)

            if restore_from_backup(backup_path):
                logger.error("已成功回退到备份")
            else:
                logger.error("回退失败，请手动恢复数据库")
            sys.exit(1)

        # 检查外键约束、统计并清理违反外键约束的记录
        violations_stats, processed_count = check_and_clean_foreign_key_violations(
            engine
        )

        migration_success = True

        logger.info("=" * 60)
        logger.info("数据库迁移完成！")
        logger.info("数据丢失校验通过！")
        if violations_stats:
            logger.info(f"外键约束违反统计: {sum(violations_stats.values())} 条记录")
            logger.info(
                f"已处理 {processed_count} 条违反外键约束的记录（设置为 NULL 或删除）"
            )
        else:
            logger.info("未发现外键约束违反")
        logger.info("=" * 60)

        # 用户确认
        logger.info("=" * 60)
        logger.info("迁移和校验已完成，请确认是否接受此次迁移")
        logger.info(f"备份文件位置: {backup_path}")
        logger.info("=" * 60)

        try:
            user_input = input("确认接受迁移？(y/n): ").strip().lower()
            if user_input in ("y", "yes"):
                logger.info("=" * 60)
                logger.info("用户已确认，迁移完成！")
                logger.info("=" * 60)
            else:
                logger.warning("=" * 60)
                logger.warning("用户已取消，开始回退到备份...")
                logger.warning("=" * 60)

                if restore_from_backup(backup_path):
                    logger.warning("已成功回退到备份")
                    logger.warning("迁移已取消")
                else:
                    logger.error("回退失败，请手动恢复数据库")
                sys.exit(0)

        except (KeyboardInterrupt, EOFError):
            logger.warning("")
            logger.warning("=" * 60)
            logger.warning("用户中断操作，开始回退到备份...")
            logger.warning("=" * 60)

            if restore_from_backup(backup_path):
                logger.warning("已成功回退到备份")
                logger.warning("迁移已取消")
            else:
                logger.error("回退失败，请手动恢复数据库")
            sys.exit(0)

    except Exception as e:
        logger.error("=" * 60)
        logger.error(f"迁移过程中发生异常: {e}", exc_info=True)
        logger.error("=" * 60)

        if not migration_success:
            logger.error("开始回退到备份...")
            if restore_from_backup(backup_path):
                logger.error("已成功回退到备份")
            else:
                logger.error("回退失败，请手动恢复数据库")
        sys.exit(1)


if __name__ == "__main__":
    main()