From a26a44e5c90ffc46d20b204c58af8c2a287de325 Mon Sep 17 00:00:00 2001 From: Lxy Date: Sat, 4 Jul 2026 10:36:14 +0800 Subject: [PATCH] feat(migration): add SQLite to MySQL data migration script --- app/migration.py | 67 ++++++++++++++++++++ tests/test_migration.py | 135 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 202 insertions(+) create mode 100644 app/migration.py create mode 100644 tests/test_migration.py diff --git a/app/migration.py b/app/migration.py new file mode 100644 index 0000000..89ae255 --- /dev/null +++ b/app/migration.py @@ -0,0 +1,67 @@ +""" +数据缓冲平台 - SQLite 到 MySQL 数据迁移 +""" +import logging + +from sqlalchemy import create_engine +from sqlalchemy.orm import sessionmaker + +from app.config import DB_PATH +from app.models import MarketData, SymbolTimestamp +from app.mysql_database import MySQLSessionLocal + +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() + + 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() diff --git a/tests/test_migration.py b/tests/test_migration.py new file mode 100644 index 0000000..d78b7ed --- /dev/null +++ b/tests/test_migration.py @@ -0,0 +1,135 @@ +""" +数据缓冲平台 - SQLite 到 MySQL 数据迁移脚本测试 +""" +import logging +from datetime import datetime +from unittest.mock import MagicMock, patch + +import pytest +from sqlalchemy import create_engine +from sqlalchemy.orm import sessionmaker + +from app.models import Base, MarketData, SymbolTimestamp + + +def _create_sqlite_source(tmp_path): + """创建包含测试数据的文件 SQLite 源数据库。""" + db_path = tmp_path / "test_source.db" + engine = create_engine(f"sqlite:///{db_path}") + MarketData.__table__.create(bind=engine) + SymbolTimestamp.__table__.create(bind=engine) + Session = sessionmaker(bind=engine) + session = Session() + fetched_at = datetime(2026, 7, 4, 10, 0, 0) + session.add_all( + [ + MarketData( + symbol="AG2606", + data_type="futures", + period="5min", + candles_json='[{"open": 100, "close": 110}]', + current_price=123.45, + fetched_at=fetched_at, + candle_count=1, + ), + MarketData( + symbol="AG2606", + data_type="futures", + period="15min", + candles_json='[{"open": 105, "close": 115}]', + current_price=123.45, + fetched_at=fetched_at, + candle_count=1, + ), + SymbolTimestamp( + symbol="AG2606", + data_type="futures", + last_refresh_at=fetched_at, + refresh_count=5, + ), + ] + ) + session.commit() + session.close() + return engine + + +def _make_mysql_session(market_data_count=0, symbol_timestamp_count=0): + """构造模拟 MySQL 会话，支持记录写入和计数查询。""" + mysql_session = MagicMock() + mysql_session.__enter__ = MagicMock(return_value=mysql_session) + mysql_session.__exit__ = MagicMock(return_value=False) + + added_market_data = [] + added_symbol_timestamps = [] + + def mock_query(model): + query = MagicMock() + if model is MarketData: + query.count.return_value = market_data_count + elif model is SymbolTimestamp: + query.count.return_value = symbol_timestamp_count + return query + + mysql_session.query.side_effect = mock_query + mysql_session.add_all.side_effect = lambda records: ( + added_market_data.extend(records) + if records and isinstance(records[0], MarketData) + else added_symbol_timestamps.extend(records) + ) + + return mysql_session, added_market_data, added_symbol_timestamps + + +class TestMigrateSqliteToMysql: + """migrate_sqlite_to_mysql 行为测试。""" + + def test_migrates_data_when_mysql_tables_are_empty(self, caplog, tmp_path): + """MySQL 表为空时，应从 SQLite 迁移全部数据。""" + sqlite_engine = _create_sqlite_source(tmp_path) + mysql_session, added_market_data, added_symbol_timestamps = _make_mysql_session( + market_data_count=0, symbol_timestamp_count=0 + ) + + with patch("app.migration.create_engine", return_value=sqlite_engine): + with patch( + "app.migration.MySQLSessionLocal", return_value=mysql_session + ): + with caplog.at_level(logging.INFO, logger="app.migration"): + from app.migration import migrate_sqlite_to_mysql + + result = migrate_sqlite_to_mysql() + + assert result is True + assert len(added_market_data) == 2 + assert len(added_symbol_timestamps) == 1 + assert added_market_data[0].symbol == "AG2606" + assert added_market_data[0].period == "5min" + assert added_symbol_timestamps[0].symbol == "AG2606" + assert "迁移完成" in caplog.text + assert "market_data: 2" in caplog.text + assert "symbol_timestamps: 1" in caplog.text + mysql_session.commit.assert_called_once() + + def test_skips_migration_when_mysql_has_data(self, caplog, tmp_path): + """MySQL 表非空时，应跳过迁移并输出提示日志。""" + sqlite_engine = _create_sqlite_source(tmp_path) + mysql_session, added_market_data, added_symbol_timestamps = _make_mysql_session( + market_data_count=1, symbol_timestamp_count=1 + ) + + with patch("app.migration.create_engine", return_value=sqlite_engine): + with patch( + "app.migration.MySQLSessionLocal", return_value=mysql_session + ): + with caplog.at_level(logging.INFO, logger="app.migration"): + from app.migration import migrate_sqlite_to_mysql + + result = migrate_sqlite_to_mysql() + + assert result is False + assert len(added_market_data) == 0 + assert len(added_symbol_timestamps) == 0 + assert "MySQL 已存在数据" in caplog.text + assert "跳过迁移" in caplog.text + mysql_session.commit.assert_not_called()