From e17a183262b0b748c6008014aef38e732508d6bd Mon Sep 17 00:00:00 2001 From: Lxy Date: Sat, 4 Jul 2026 12:16:30 +0800 Subject: [PATCH] fix(storage): address final review findings on atomic writes, migration, config and URL encoding --- app/migration.py | 38 ++++--- app/mysql_database.py | 4 +- app/services/cache.py | 7 +- app/storage_manager.py | 130 +++++++++++++++++++++++- tests/test_cache_storage_integration.py | 22 ++-- tests/test_migration.py | 22 ++++ tests/test_mysql_database.py | 23 +++++ tests/test_storage_manager.py | 89 ++++++++++++++++ 8 files changed, 311 insertions(+), 24 deletions(-) create mode 100644 tests/test_mysql_database.py diff --git a/app/migration.py b/app/migration.py index 5e45e23..9b77e64 100644 --- a/app/migration.py +++ b/app/migration.py @@ -41,12 +41,8 @@ def migrate_sqlite_to_mysql(): 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 + migrated_market_data = False + migrated_symbol_timestamps = False sqlite_session = SQLiteSession() try: @@ -55,13 +51,31 @@ def migrate_sqlite_to_mysql(): finally: sqlite_session.close() - _prepare_records_for_new_session(market_data_records) - _prepare_records_for_new_session(symbol_timestamp_records) + if market_data_count > 0: + logger.info( + f"MySQL market_data 表已存在数据，跳过迁移 " + f"(market_data: {market_data_count})" + ) + else: + _prepare_records_for_new_session(market_data_records) + if market_data_records: + mysql_session.add_all(market_data_records) + migrated_market_data = True - if market_data_records: - mysql_session.add_all(market_data_records) - if symbol_timestamp_records: - mysql_session.add_all(symbol_timestamp_records) + if symbol_timestamp_count > 0: + logger.info( + f"MySQL symbol_timestamps 表已存在数据，跳过迁移 " + f"(symbol_timestamps: {symbol_timestamp_count})" + ) + else: + _prepare_records_for_new_session(symbol_timestamp_records) + if symbol_timestamp_records: + mysql_session.add_all(symbol_timestamp_records) + migrated_symbol_timestamps = True + + if not migrated_market_data and not migrated_symbol_timestamps: + logger.info("MySQL 已存在数据，跳过迁移") + return False mysql_session.commit() logger.info( diff --git a/app/mysql_database.py b/app/mysql_database.py index f27cc63..8f66219 100644 --- a/app/mysql_database.py +++ b/app/mysql_database.py @@ -2,6 +2,7 @@ 数据缓冲平台 - MySQL 数据库连接 """ import logging +import urllib.parse from sqlalchemy import create_engine, text from sqlalchemy.orm import sessionmaker from app.config import MYSQL_HOST, MYSQL_PORT, MYSQL_USER, MYSQL_PASSWORD, MYSQL_DATABASE @@ -16,7 +17,8 @@ def init_mysql(): """初始化 MySQL 连接引擎，失败时返回 None""" global mysql_engine, MySQLSessionLocal try: - url = f"mysql+pymysql://{MYSQL_USER}:{MYSQL_PASSWORD}@{MYSQL_HOST}:{MYSQL_PORT}/{MYSQL_DATABASE}?charset=utf8mb4" + encoded_password = urllib.parse.quote_plus(MYSQL_PASSWORD) + url = f"mysql+pymysql://{MYSQL_USER}:{encoded_password}@{MYSQL_HOST}:{MYSQL_PORT}/{MYSQL_DATABASE}?charset=utf8mb4" mysql_engine = create_engine(url, pool_pre_ping=True, pool_recycle=3600) with mysql_engine.connect() as conn: conn.execute(text("SELECT 1")) diff --git a/app/services/cache.py b/app/services/cache.py index 4aa67e0..8f6b11a 100644 --- a/app/services/cache.py +++ b/app/services/cache.py @@ -294,15 +294,16 @@ def save_market_data(db: Session, symbol: str, data: Dict) -> Optional[MarketDat if storage.is_available(): try: - if storage.save_market_data(symbol, data): - storage.save_symbol_timestamp(symbol, data_type, now) + if storage.save_market_data_with_timestamp(symbol, data, now): # 注意：StorageManager 命中时返回的是内存构造的代表性 MarketData， # 不代表 SQLite 持久化记录，因此没有 SQLite 自增 id。 return _representative_market_data(symbol, data, now) else: - raise Exception("save_market_data returned False") + raise Exception("save_market_data_with_timestamp returned False") except Exception as e: logger.warning(f"StorageManager 写入失败，降级到 SQLite: {e}") + # 注意：当 StorageManager 不可用时，数据只会写入本地 SQLite 兜底， + # 这部分数据不会自动同步到主存储（Redis/MySQL），可能与主存储脱节。 return _save_to_sqlite(db, symbol, data) diff --git a/app/storage_manager.py b/app/storage_manager.py index 762e2d0..a3cf410 100644 --- a/app/storage_manager.py +++ b/app/storage_manager.py @@ -7,12 +7,13 @@ from datetime import datetime from sqlalchemy import text from sqlalchemy.orm import sessionmaker +from app.config import REDIS_TTL_SECONDS from app.models import MarketData, SymbolTimestamp logger = logging.getLogger(__name__) -# Redis 缓存 TTL：30 天 -CACHE_TTL_SECONDS = 30 * 24 * 60 * 60 +# Redis 缓存 TTL：读取自 app.config.REDIS_TTL_SECONDS（默认 30 天） +CACHE_TTL_SECONDS = REDIS_TTL_SECONDS class StorageManager: @@ -549,6 +550,131 @@ class StorageManager: return True + def save_market_data_with_timestamp(self, symbol, data, refresh_time): + """ + 在同一 MySQL 事务中原子写入行情数据与合约时间戳。 + + Redis 缓存的删除/回填在事务外进行，失败仅记录日志，不影响 MySQL 写入。 + + Args: + symbol: 品种合约代码 + data: 采集脚本返回的数据，包含 timeframes, current_price, timestamp, type + refresh_time: 刷新时间（datetime） + + Returns: + bool: MySQL 写入是否成功 + """ + data_type = data.get("type", "futures") + timeframes = data.get("timeframes", {}) + current_price = data.get("current_price") + fetched_at = self._parse_fetched_at(data.get("timestamp")) + + market_cache_keys = [ + f"market_data:{symbol}:{period}" for period in timeframes.keys() + ] + timestamp_cache_key = f"symbol_timestamps:{symbol}" + cache_keys = market_cache_keys + [timestamp_cache_key] + + # 1. 删除 Redis 缓存 + if self.check_redis() and cache_keys: + try: + self.redis_client.delete(*cache_keys) + except Exception as e: + logger.warning(f"Redis 删除缓存失败 [{symbol}]: {e}") + + # 2. 写入 MySQL（同一事务） + if not self.check_mysql(): + logger.warning(f"MySQL 不可用，无法写入行情数据与时间戳 [{symbol}]") + return False + + try: + SessionLocal = sessionmaker(bind=self.mysql_engine) + with SessionLocal() as db: + try: + for period, candles in timeframes.items(): + candles_json = json.dumps(candles, ensure_ascii=False) + record = ( + db.query(MarketData) + .filter_by(symbol=symbol, data_type=data_type, period=period) + .first() + ) + if record: + record.candles_json = candles_json + record.current_price = current_price + record.fetched_at = fetched_at + record.candle_count = len(candles) + else: + record = MarketData( + symbol=symbol, + data_type=data_type, + period=period, + candles_json=candles_json, + current_price=current_price, + fetched_at=fetched_at, + candle_count=len(candles), + ) + db.add(record) + + timestamp_record = ( + db.query(SymbolTimestamp) + .filter_by(symbol=symbol, data_type=data_type) + .first() + ) + if timestamp_record: + timestamp_record.last_refresh_at = refresh_time + timestamp_record.refresh_count += 1 + refresh_count = timestamp_record.refresh_count + else: + timestamp_record = SymbolTimestamp( + symbol=symbol, + data_type=data_type, + last_refresh_at=refresh_time, + refresh_count=1, + ) + db.add(timestamp_record) + refresh_count = 1 + + db.commit() + except Exception: + db.rollback() + raise + except Exception as e: + logger.error(f"MySQL 写入行情数据与时间戳失败 [{symbol}]: {e}") + return False + + # 3. MySQL 成功，回填 Redis（失败仅记录日志） + if self.check_redis(): + for period, candles in timeframes.items(): + cache_key = f"market_data:{symbol}:{period}" + cache_value = { + "current_price": current_price, + "timestamp": fetched_at.isoformat(), + "candles": candles, + } + try: + self.redis_client.setex( + cache_key, + CACHE_TTL_SECONDS, + json.dumps(cache_value, ensure_ascii=False), + ) + except Exception as e: + logger.warning(f"Redis 更新失败 [{cache_key}]: {e}") + + timestamp_cache_value = { + "last_refresh_at": refresh_time.isoformat(), + "refresh_count": refresh_count, + } + try: + self.redis_client.setex( + timestamp_cache_key, + CACHE_TTL_SECONDS, + json.dumps(timestamp_cache_value, ensure_ascii=False), + ) + except Exception as e: + logger.warning(f"Redis 更新失败 [{timestamp_cache_key}]: {e}") + + return True + def delete_cache(self, symbol, periods): """ 批量删除 Redis 行情数据缓存键。 diff --git a/tests/test_cache_storage_integration.py b/tests/test_cache_storage_integration.py index f354752..37adeeb 100644 --- a/tests/test_cache_storage_integration.py +++ b/tests/test_cache_storage_integration.py @@ -43,6 +43,7 @@ def storage_manager_mock(): """构造默认可用的 StorageManager mock。""" manager = MagicMock() manager.is_available.return_value = True + manager.save_market_data_with_timestamp.return_value = True manager.save_market_data.return_value = True manager.save_symbol_timestamp.return_value = True manager.get_market_data.return_value = None @@ -99,15 +100,20 @@ class TestSaveMarketData: """save_market_data 双写与降级测试。""" def test_uses_storage_manager_when_available(self, db, sample_data, storage_manager_mock): - """StorageManager 可用时优先写入存储层，并同步更新时间戳，返回代表记录。""" + """StorageManager 可用时通过原子方法写入行情数据与时间戳，并返回代表记录。""" with patch("app.services.cache.get_storage_manager", return_value=storage_manager_mock): from app.services import cache result = cache.save_market_data(db, "AG2606", sample_data) storage_manager_mock.is_available.assert_called_once() - storage_manager_mock.save_market_data.assert_called_once_with("AG2606", sample_data) - storage_manager_mock.save_symbol_timestamp.assert_called_once() + storage_manager_mock.save_market_data_with_timestamp.assert_called_once() + call_args = storage_manager_mock.save_market_data_with_timestamp.call_args[0] + assert call_args[0] == "AG2606" + assert call_args[1] == sample_data + assert isinstance(call_args[2], datetime) + storage_manager_mock.save_market_data.assert_not_called() + storage_manager_mock.save_symbol_timestamp.assert_not_called() assert isinstance(result, MarketData) assert result.symbol == "AG2606" assert result.data_type == "futures" @@ -136,21 +142,22 @@ class TestSaveMarketData: def test_falls_back_to_sqlite_when_storage_save_fails(self, db, sample_data, storage_manager_mock): """StorageManager 写入失败时降级到 SQLite 写入。""" - storage_manager_mock.save_market_data.return_value = False + storage_manager_mock.save_market_data_with_timestamp.return_value = False with patch("app.services.cache.get_storage_manager", return_value=storage_manager_mock): from app.services import cache result = cache.save_market_data(db, "AG2606", sample_data) - storage_manager_mock.save_market_data.assert_called_once() + storage_manager_mock.save_market_data_with_timestamp.assert_called_once() + storage_manager_mock.save_market_data.assert_not_called() storage_manager_mock.save_symbol_timestamp.assert_not_called() assert isinstance(result, MarketData) assert db.query(MarketData).filter_by(symbol="AG2606").count() == 2 def test_falls_back_to_sqlite_when_storage_raises(self, db, sample_data, storage_manager_mock, caplog): """StorageManager 写入异常时降级到 SQLite 写入并记录警告。""" - storage_manager_mock.save_market_data.side_effect = Exception("storage down") + storage_manager_mock.save_market_data_with_timestamp.side_effect = Exception("storage down") with patch("app.services.cache.get_storage_manager", return_value=storage_manager_mock): from app.services import cache @@ -158,6 +165,9 @@ class TestSaveMarketData: with caplog.at_level(logging.WARNING, logger="app.services.cache"): result = cache.save_market_data(db, "AG2606", sample_data) + storage_manager_mock.save_market_data_with_timestamp.assert_called_once() + storage_manager_mock.save_market_data.assert_not_called() + storage_manager_mock.save_symbol_timestamp.assert_not_called() assert isinstance(result, MarketData) assert db.query(MarketData).filter_by(symbol="AG2606").count() == 2 assert "StorageManager 写入失败" in caplog.text diff --git a/tests/test_migration.py b/tests/test_migration.py index d78b7ed..5936bef 100644 --- a/tests/test_migration.py +++ b/tests/test_migration.py @@ -133,3 +133,25 @@ class TestMigrateSqliteToMysql: assert "MySQL 已存在数据" in caplog.text assert "跳过迁移" in caplog.text mysql_session.commit.assert_not_called() + + def test_migrates_empty_table_when_other_table_has_data(self, caplog, tmp_path): + """按表分别判断：仅空表需要迁移，非空表跳过。""" + 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=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) == 0 + assert len(added_symbol_timestamps) == 1 + assert added_symbol_timestamps[0].symbol == "AG2606" + mysql_session.commit.assert_called_once() diff --git a/tests/test_mysql_database.py b/tests/test_mysql_database.py new file mode 100644 index 0000000..e79f126 --- /dev/null +++ b/tests/test_mysql_database.py @@ -0,0 +1,23 @@ +""" +数据缓冲平台 - MySQL 数据库连接测试 +""" +from unittest.mock import MagicMock, patch + +from app.mysql_database import init_mysql + + +def test_init_mysql_url_encodes_password(): + """MySQL 密码应使用 urllib.parse.quote_plus 进行 URL 编码。""" + with patch("app.mysql_database.create_engine") as mock_create_engine: + mock_engine = MagicMock() + mock_create_engine.return_value = mock_engine + with patch("app.mysql_database.MYSQL_PASSWORD", "p@ss/w+ord"): + with patch("app.mysql_database.MYSQL_USER", "user"): + with patch("app.mysql_database.MYSQL_HOST", "host"): + with patch("app.mysql_database.MYSQL_PORT", "3306"): + with patch("app.mysql_database.MYSQL_DATABASE", "buffer"): + init_mysql() + + url = mock_create_engine.call_args[0][0] + assert "p%40ss%2Fw%2Bord" in url + assert "p@ss" not in url diff --git a/tests/test_storage_manager.py b/tests/test_storage_manager.py index 49bb640..9eb77db 100644 --- a/tests/test_storage_manager.py +++ b/tests/test_storage_manager.py @@ -1,3 +1,4 @@ +import importlib import json import logging import time @@ -785,3 +786,91 @@ class TestStorageManagerWriteLogic: redis_client.delete.assert_called_once_with( "market_data:AG2606:5min", "market_data:AG2606:15min" ) + + +class TestStorageManagerConfig: + """配置读取相关测试。""" + + def test_redis_ttl_reads_config_value(self): + """StorageManager 应读取 app.config.REDIS_TTL_SECONDS 作为默认 TTL。""" + import app.config as config + import app.storage_manager as sm + + original_ttl = config.REDIS_TTL_SECONDS + original_singleton = sm._storage_manager + config.REDIS_TTL_SECONDS = 1234 + try: + importlib.reload(sm) + assert sm.CACHE_TTL_SECONDS == 1234 + finally: + config.REDIS_TTL_SECONDS = original_ttl + importlib.reload(sm) + sm._storage_manager = original_singleton + + +class TestStorageManagerAtomicWrite: + """save_market_data_with_timestamp 原子写入测试。""" + + def test_save_market_data_with_timestamp_writes_market_data_and_timestamp(self): + """同一 MySQL 事务中写入行情数据与合约时间戳。""" + from sqlalchemy import create_engine + from sqlalchemy.orm import sessionmaker + + from app.models import Base, MarketData, SymbolTimestamp + + redis_client = MagicMock() + mysql_engine = create_engine("sqlite:///:memory:") + Base.metadata.create_all(bind=mysql_engine) + manager = StorageManager(redis_client=redis_client, mysql_engine=mysql_engine) + + fetched_at = datetime(2026, 7, 4, 10, 0, 0) + data = { + "type": "futures", + "current_price": 123.45, + "timestamp": fetched_at.isoformat(), + "timeframes": { + "5min": [{"open": 100, "close": 110}], + "15min": [{"open": 105, "close": 115}], + }, + } + refresh_time = datetime(2026, 7, 4, 10, 1, 0) + + result = manager.save_market_data_with_timestamp("AG2606", data, refresh_time) + + assert result is True + SessionLocal = sessionmaker(bind=mysql_engine) + with SessionLocal() as db: + market_records = db.query(MarketData).filter_by(symbol="AG2606").all() + assert len(market_records) == 2 + periods = {r.period for r in market_records} + assert periods == {"5min", "15min"} + + ts = db.query(SymbolTimestamp).filter_by(symbol="AG2606", data_type="futures").first() + assert ts is not None + assert ts.last_refresh_at == refresh_time + assert ts.refresh_count == 1 + + redis_client.delete.assert_called_once_with( + "market_data:AG2606:5min", + "market_data:AG2606:15min", + "symbol_timestamps:AG2606", + ) + + def test_save_market_data_with_timestamp_returns_false_on_mysql_failure(self): + """MySQL 写入失败时返回 False，不同步时间戳。""" + redis_client = MagicMock() + manager = StorageManager(redis_client=redis_client, mysql_engine=None) + + data = { + "type": "futures", + "current_price": 123.45, + "timestamp": datetime(2026, 7, 4, 10, 0, 0).isoformat(), + "timeframes": {"5min": [{"open": 100, "close": 110}]}, + } + + result = manager.save_market_data_with_timestamp( + "AG2606", data, datetime(2026, 7, 4, 10, 1, 0) + ) + + assert result is False + redis_client.setex.assert_not_called()