""" 数据库迁移脚本:将原版两个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()