buffer_platform/buffer_platform_v2/scripts/migrate_data.py

"""
数据库迁移脚本：将原版两个SQLite数据库合并到V2统一数据库

原版数据库：
  - buffer.db（5张表：market_data, symbol_timestamps, scheduled_tasks, users, sessions）
  - futures_analysis.db（14张表：futures_analysis, watched_symbols, ai_model_configs, ...）

V2数据库：
  - buffer_v2.db（19张表，合并上述两个库的所有表）

使用方法：
  python scripts/migrate_data.py [--source-dir SOURCE_DIR] [--target-db TARGET_DB] [--dry-run]

示例：
  # 预览迁移（不实际执行）
  python scripts/migrate_data.py --dry-run

  # 执行迁移（默认路径）
  python scripts/migrate_data.py

  # 指定源目录
  python scripts/migrate_data.py --source-dir ../data
"""

import argparse
import os
import shutil
import sqlite3
import sys
from datetime import datetime
from pathlib import Path

# 默认路径
DEFAULT_SOURCE_DIR = Path(__file__).parent.parent.parent / "data"
DEFAULT_TARGET_DB = Path(__file__).parent.parent / "data" / "buffer_v2.db"

# 从 buffer.db 迁移的表
BUFFER_DB_TABLES = [
    "market_data",
    "symbol_timestamps",
    "scheduled_tasks",
    "users",
    "sessions",
]

# 从 futures_analysis.db 迁移的表
ANALYSIS_DB_TABLES = [
    "futures_analysis",
    "watched_symbols",
    "ai_model_configs",
    "analysis_settings",
    "ai_analysis_cache",
    "review_dates",
    "symbol_rankings",
    "trading_plans",
    "symbol_scores_v2",
    "trading_plans_v2",
    "sector_heat",
    "review_plans_v2",
    "trade_records",
    "trade_import_batches",
]


def get_table_row_count(conn: sqlite3.Connection, table: str) -> int:
    """获取表的行数"""
    try:
        cursor = conn.execute(f"SELECT COUNT(*) FROM {table}")
        return cursor.fetchone()[0]
    except sqlite3.OperationalError:
        return -1  # 表不存在


def get_table_columns(conn: sqlite3.Connection, table: str) -> list[str]:
    """获取表的列名列表"""
    cursor = conn.execute(f"PRAGMA table_info({table})")
    return [row[1] for row in cursor.fetchall()]


def check_table_exists(conn: sqlite3.Connection, table: str) -> bool:
    """检查表是否存在"""
    cursor = conn.execute(
        "SELECT name FROM sqlite_master WHERE type='table' AND name=?", (table,)
    )
    return cursor.fetchone() is not None


def backup_database(db_path: Path) -> Path | None:
    """备份数据库文件"""
    if not db_path.exists():
        return None
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    backup_path = db_path.with_suffix(f".db.backup_{timestamp}")
    shutil.copy2(db_path, backup_path)
    return backup_path


def migrate_database(
    source_dir: Path,
    target_db: Path,
    dry_run: bool = False,
) -> dict:
    """执行数据迁移

    Args:
        source_dir: 原始数据库所在目录
        target_db: 目标V2数据库路径
        dry_run: 仅预览不执行

    Returns:
        迁移统计字典
    """
    buffer_db_path = source_dir / "buffer.db"
    analysis_db_path = source_dir / "futures_analysis.db"

    stats = {
        "buffer_db_tables": {},
        "analysis_db_tables": {},
        "errors": [],
        "dry_run": dry_run,
    }

    # -----------------------------------------------------------------------
    # 1. 检查源数据库
    # -----------------------------------------------------------------------
    print("=" * 60)
    print("数据库迁移：buffer.db + futures_analysis.db → buffer_v2.db")
    print("=" * 60)

    if not buffer_db_path.exists():
        stats["errors"].append(f"源数据库不存在: {buffer_db_path}")
        print(f"[ERROR] 源数据库不存在: {buffer_db_path}")
        return stats

    if not analysis_db_path.exists():
        stats["errors"].append(f"源数据库不存在: {analysis_db_path}")
        print(f"[ERROR] 源数据库不存在: {analysis_db_path}")
        return stats

    print(f"\n[OK] 源数据库已找到:")
    print(f"  - {buffer_db_path}")
    print(f"  - {analysis_db_path}")

    # -----------------------------------------------------------------------
    # 2. 统计源数据库数据量
    # -----------------------------------------------------------------------
    print("\n" + "-" * 60)
    print("源数据库统计:")
    print("-" * 60)

    # 统计 buffer.db
    buffer_conn = sqlite3.connect(str(buffer_db_path))
    print(f"\n  buffer.db ({buffer_db_path}):")
    for table in BUFFER_DB_TABLES:
        count = get_table_row_count(buffer_conn, table)
        stats["buffer_db_tables"][table] = count
        status = f"{count} 行" if count >= 0 else "表不存在"
        print(f"    {table:30s} {status}")
    buffer_conn.close()

    # 统计 futures_analysis.db
    analysis_conn = sqlite3.connect(str(analysis_db_path))
    print(f"\n  futures_analysis.db ({analysis_db_path}):")
    for table in ANALYSIS_DB_TABLES:
        count = get_table_row_count(analysis_conn, table)
        stats["analysis_db_tables"][table] = count
        status = f"{count} 行" if count >= 0 else "表不存在"
        print(f"    {table:30s} {status}")
    analysis_conn.close()

    total_rows = sum(v for v in stats["buffer_db_tables"].values() if v > 0)
    total_rows += sum(v for v in stats["analysis_db_tables"].values() if v > 0)
    print(f"\n  总计: {total_rows} 行数据待迁移")

    # -----------------------------------------------------------------------
    # 3. 如果 dry_run，到此结束
    # -----------------------------------------------------------------------
    if dry_run:
        print("\n[DRY RUN] 预览完成，未执行实际迁移")
        return stats

    # -----------------------------------------------------------------------
    # 4. 备份目标数据库（如果已存在）
    # -----------------------------------------------------------------------
    if target_db.exists():
        backup_path = backup_database(target_db)
        print(f"\n[BACKUP] 已备份目标数据库: {backup_path}")

    # 确保目标目录存在
    target_db.parent.mkdir(parents=True, exist_ok=True)

    # -----------------------------------------------------------------------
    # 5. 创建目标数据库并建表
    # -----------------------------------------------------------------------
    print(f"\n[TARGET] 创建目标数据库: {target_db}")

    # 使用V2项目的模型创建表
    # 导入V2模型并创建表
    sys.path.insert(0, str(Path(__file__).parent.parent))
    from app.database import Base, engine
    from app.models import (  # noqa: F401
        MarketData,
        SymbolTimestamp,
        ScheduledTask,
        User,
        Session,
        FuturesAnalysis,
        WatchedSymbol,
        AIModelConfig,
        AnalysisSettings,
        AIAnalysisCache,
        ReviewDate,
        SymbolRanking,
        TradingPlan,
        SymbolScoreV2,
        TradingPlanV2,
        SectorHeat,
        ReviewPlanV2,
        TradeRecord,
        TradeImportBatch,
    )

    import asyncio

    async def create_tables():
        async with engine.begin() as conn:
            await conn.run_sync(Base.metadata.create_all)

    asyncio.run(create_tables())
    print("[OK] 目标数据库表已创建")

    # -----------------------------------------------------------------------
    # 6. 迁移 buffer.db 数据
    # -----------------------------------------------------------------------
    print("\n" + "-" * 60)
    print("迁移 buffer.db 数据...")
    print("-" * 60)

    # 使用 SQLite 的 ATTACH 方式直接迁移（最高效）
    target_conn = sqlite3.connect(str(target_db))
    target_conn.execute("PRAGMA journal_mode=WAL")
    target_conn.execute("PRAGMA synchronous=NORMAL")

    buffer_conn = sqlite3.connect(str(buffer_db_path))

    for table in BUFFER_DB_TABLES:
        if not check_table_exists(buffer_conn, table):
            print(f"  [SKIP] {table}: 源表不存在")
            continue

        source_count = get_table_row_count(buffer_conn, table)
        if source_count == 0:
            print(f"  [SKIP] {table}: 源表为空")
            continue

        # 获取源表和目标表的列名，取交集
        source_cols = get_table_columns(buffer_conn, table)
        target_cols = get_table_columns(target_conn, table)
        common_cols = [c for c in source_cols if c in target_cols]

        if not common_cols:
            print(f"  [ERROR] {table}: 无匹配列")
            stats["errors"].append(f"{table}: 无匹配列")
            continue

        cols_str = ", ".join(common_cols)

        try:
            # 从源数据库读取数据
            cursor = buffer_conn.execute(f"SELECT {cols_str} FROM {table}")
            rows = cursor.fetchall()

            if not rows:
                print(f"  [SKIP] {table}: 无数据")
                continue

            # 插入到目标数据库
            placeholders = ", ".join(["?"] * len(common_cols))
            insert_sql = f"INSERT OR REPLACE INTO {table} ({cols_str}) VALUES ({placeholders})"

            target_conn.executemany(insert_sql, rows)
            target_conn.commit()

            migrated_count = len(rows)
            print(f"  [OK] {table}: 迁移 {migrated_count} 行")

        except sqlite3.Error as e:
            print(f"  [ERROR] {table}: {e}")
            stats["errors"].append(f"{table}: {e}")

    buffer_conn.close()

    # -----------------------------------------------------------------------
    # 7. 迁移 futures_analysis.db 数据
    # -----------------------------------------------------------------------
    print("\n" + "-" * 60)
    print("迁移 futures_analysis.db 数据...")
    print("-" * 60)

    analysis_conn = sqlite3.connect(str(analysis_db_path))

    for table in ANALYSIS_DB_TABLES:
        if not check_table_exists(analysis_conn, table):
            print(f"  [SKIP] {table}: 源表不存在")
            continue

        source_count = get_table_row_count(analysis_conn, table)
        if source_count == 0:
            print(f"  [SKIP] {table}: 源表为空")
            continue

        # 获取源表和目标表的列名，取交集
        source_cols = get_table_columns(analysis_conn, table)
        target_cols = get_table_columns(target_conn, table)
        common_cols = [c for c in source_cols if c in target_cols]

        if not common_cols:
            print(f"  [ERROR] {table}: 无匹配列")
            stats["errors"].append(f"{table}: 无匹配列")
            continue

        cols_str = ", ".join(common_cols)

        try:
            # 从源数据库读取数据
            cursor = analysis_conn.execute(f"SELECT {cols_str} FROM {table}")
            rows = cursor.fetchall()

            if not rows:
                print(f"  [SKIP] {table}: 无数据")
                continue

            # 插入到目标数据库
            placeholders = ", ".join(["?"] * len(common_cols))
            insert_sql = f"INSERT OR REPLACE INTO {table} ({cols_str}) VALUES ({placeholders})"

            target_conn.executemany(insert_sql, rows)
            target_conn.commit()

            migrated_count = len(rows)
            print(f"  [OK] {table}: 迁移 {migrated_count} 行")

        except sqlite3.Error as e:
            print(f"  [ERROR] {table}: {e}")
            stats["errors"].append(f"{table}: {e}")

    analysis_conn.close()

    # -----------------------------------------------------------------------
    # 8. 修复自增ID序列
    # -----------------------------------------------------------------------
    print("\n" + "-" * 60)
    print("修复自增ID序列...")
    print("-" * 60)

    all_tables = BUFFER_DB_TABLES + ANALYSIS_DB_TABLES
    for table in all_tables:
        if not check_table_exists(target_conn, table):
            continue
        try:
            # 获取当前最大ID
            cursor = target_conn.execute(f"SELECT MAX(id) FROM {table}")
            max_id = cursor.fetchone()[0]
            if max_id is not None:
                # 更新 sqlite_sequence 表
                target_conn.execute(
                    "INSERT OR REPLACE INTO sqlite_sequence (name, seq) VALUES (?, ?)",
                    (table, max_id),
                )
                print(f"  [OK] {table}: max_id = {max_id}")
        except sqlite3.Error:
            pass  # 某些表可能没有id列

    target_conn.commit()

    # -----------------------------------------------------------------------
    # 9. 验证迁移结果
    # -----------------------------------------------------------------------
    print("\n" + "-" * 60)
    print("验证迁移结果...")
    print("-" * 60)

    migrated_total = 0
    for table in all_tables:
        if not check_table_exists(target_conn, table):
            print(f"  [WARN] {table}: 目标表不存在")
            continue
        count = get_table_row_count(target_conn, table)
        if count > 0:
            print(f"  {table:30s} {count} 行")
            migrated_total += count

    target_conn.close()

    print(f"\n  迁移总计: {migrated_total} 行")

    # -----------------------------------------------------------------------
    # 10. 完成
    # -----------------------------------------------------------------------
    print("\n" + "=" * 60)
    if stats["errors"]:
        print(f"迁移完成，但有 {len(stats['errors'])} 个错误:")
        for err in stats["errors"]:
            print(f"  - {err}")
    else:
        print("迁移完成，无错误！")
    print("=" * 60)

    return stats


def main():
    parser = argparse.ArgumentParser(description="数据库迁移脚本")
    parser.add_argument(
        "--source-dir",
        type=Path,
        default=DEFAULT_SOURCE_DIR,
        help=f"原始数据库所在目录 (默认: {DEFAULT_SOURCE_DIR})",
    )
    parser.add_argument(
        "--target-db",
        type=Path,
        default=DEFAULT_TARGET_DB,
        help=f"目标V2数据库路径 (默认: {DEFAULT_TARGET_DB})",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="仅预览迁移，不实际执行",
    )
    args = parser.parse_args()

    print(f"\n配置:")
    print(f"  源目录: {args.source_dir}")
    print(f"  目标数据库: {args.target_db}")
    print(f"  模式: {'预览' if args.dry_run else '执行'}")
    print()

    stats = migrate_database(
        source_dir=args.source_dir,
        target_db=args.target_db,
        dry_run=args.dry_run,
    )

    # 退出码
    if stats["errors"]:
        sys.exit(1)
    sys.exit(0)


if __name__ == "__main__":
    main()