buffer_platform/app/migration.py

"""
数据缓冲平台 - SQLite 到 MySQL 数据迁移
"""
import logging

from sqlalchemy import create_engine
from sqlalchemy.orm import make_transient, sessionmaker

from app.config import DB_PATH
from app.models import MarketData, SymbolTimestamp
from app.mysql_database import MySQLSessionLocal


def _prepare_records_for_new_session(records):
    """将 ORM 实例转换为可在新 session 中插入的 transient 状态。"""
    for record in records:
        make_transient(record)
        record.id = None

logger = logging.getLogger(__name__)


def migrate_sqlite_to_mysql():
    """从 SQLite 迁移 market_data 和 symbol_timestamps 数据到 MySQL。

    幂等设计：当 MySQL 对应表已存在数据时直接跳过，避免重复写入。
    返回 True 表示执行了迁移，False 表示跳过或失败。
    """
    if MySQLSessionLocal is None:
        logger.warning("MySQL 未初始化，跳过数据迁移")
        return False

    sqlite_engine = create_engine(
        f"sqlite:///{DB_PATH}",
        connect_args={"check_same_thread": False},
    )
    SQLiteSession = sessionmaker(bind=sqlite_engine)

    mysql_session = MySQLSessionLocal()
    try:
        market_data_count = mysql_session.query(MarketData).count()
        symbol_timestamp_count = mysql_session.query(SymbolTimestamp).count()

        if market_data_count > 0 or symbol_timestamp_count > 0:
            logger.info(
                f"MySQL 已存在数据，跳过迁移 "
                f"(market_data: {market_data_count}, symbol_timestamps: {symbol_timestamp_count})"
            )
            return False

        sqlite_session = SQLiteSession()
        try:
            market_data_records = sqlite_session.query(MarketData).all()
            symbol_timestamp_records = sqlite_session.query(SymbolTimestamp).all()
        finally:
            sqlite_session.close()

        _prepare_records_for_new_session(market_data_records)
        _prepare_records_for_new_session(symbol_timestamp_records)

        if market_data_records:
            mysql_session.add_all(market_data_records)
        if symbol_timestamp_records:
            mysql_session.add_all(symbol_timestamp_records)

        mysql_session.commit()
        logger.info(
            f"迁移完成: market_data: {len(market_data_records)} 条, "
            f"symbol_timestamps: {len(symbol_timestamp_records)} 条"
        )
        return True
    except Exception as e:
        logger.error(f"数据迁移失败: {e}")
        mysql_session.rollback()
        return False
    finally:
        mysql_session.close()