From 8f3c2b6fdb415f06c2faef0ba424f0ee2446905a Mon Sep 17 00:00:00 2001 From: Lxy Date: Sat, 4 Jul 2026 10:21:24 +0800 Subject: [PATCH] feat(storage): add StorageManager write logic with dual-write consistency --- app/storage_manager.py | 177 +++++++++++++++++++++++++ tests/test_storage_manager.py | 240 ++++++++++++++++++++++++++++++++++ 2 files changed, 417 insertions(+) diff --git a/app/storage_manager.py b/app/storage_manager.py index 1a3d596..57d783e 100644 --- a/app/storage_manager.py +++ b/app/storage_manager.py @@ -1,6 +1,7 @@ import json import logging import time +from datetime import datetime from sqlalchemy import text from sqlalchemy.orm import sessionmaker @@ -210,6 +211,182 @@ class StorageManager: logger.warning(f"MySQL 读取失败 [{cache_key}]: {e}") return None + def save_market_data(self, symbol, data): + """ + 保存行情数据：先删 Redis 缓存，再写 MySQL，MySQL 成功后回填 Redis。 + + Args: + symbol: 品种合约代码 + data: 采集脚本返回的数据，包含 timeframes, current_price, timestamp, type + + Returns: + bool: MySQL 写入是否成功 + """ + data_type = data.get("type", "futures") + timeframes = data.get("timeframes", {}) + current_price = data.get("current_price") + fetched_at = datetime.now() + + cache_keys = [f"market_data:{symbol}:{period}" for period in timeframes.keys()] + + # 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) + 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}") + + return True + + def save_symbol_timestamp(self, symbol, data_type, refresh_time): + """ + 保存合约时间戳：先删 Redis 缓存，再写 MySQL，MySQL 成功后回填 Redis。 + + Args: + symbol: 品种合约代码 + data_type: 数据类型，如 futures + refresh_time: 刷新时间（datetime） + + Returns: + bool: MySQL 写入是否成功 + """ + cache_key = f"symbol_timestamps:{symbol}" + + # 1. 删除 Redis 缓存 + if self.check_redis(): + try: + self.redis_client.delete(cache_key) + except Exception as e: + logger.warning(f"Redis 删除缓存失败 [{cache_key}]: {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: + record = ( + db.query(SymbolTimestamp) + .filter_by(symbol=symbol, data_type=data_type) + .first() + ) + if record: + record.last_refresh_at = refresh_time + record.refresh_count += 1 + else: + record = SymbolTimestamp( + symbol=symbol, + data_type=data_type, + last_refresh_at=refresh_time, + refresh_count=1, + ) + db.add(record) + 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(): + cache_value = { + "last_refresh_at": refresh_time.isoformat(), + "refresh_count": record.refresh_count, + } + 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}") + + return True + + def delete_cache(self, symbol, periods): + """ + 批量删除 Redis 行情数据缓存键。 + + Args: + symbol: 品种合约代码 + periods: 周期列表，如 ["5min", "15min"] + """ + if not self.check_redis(): + return + + cache_keys = [f"market_data:{symbol}:{period}" for period in periods] + if not cache_keys: + return + + try: + self.redis_client.delete(*cache_keys) + except Exception as e: + logger.warning(f"Redis 批量删除缓存失败 [{symbol}]: {e}") + _storage_manager = None diff --git a/tests/test_storage_manager.py b/tests/test_storage_manager.py index 7915ccb..c9df767 100644 --- a/tests/test_storage_manager.py +++ b/tests/test_storage_manager.py @@ -289,3 +289,243 @@ class TestStorageManagerReadLogic: redis_client.setex.assert_called_once_with( "symbol_timestamps:AG2606", 2592000, json.dumps(expected, ensure_ascii=False) ) + + +class TestStorageManagerWriteLogic: + """StorageManager 双写一致性测试。""" + + def test_save_market_data_deletes_cache_writes_mysql_and_updates_redis(self): + """MySQL 写入成功时，应先删缓存、再写 MySQL、最后回填 Redis。""" + redis_client = MagicMock() + mysql_engine = MagicMock() + 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}], + }, + } + + records = [] + + def mock_filter_by(symbol=None, data_type=None, period=None): + existing = next( + ( + r + for r in records + if r.symbol == symbol and r.data_type == data_type and r.period == period + ), + None, + ) + fb_query = MagicMock() + fb_query.first.return_value = existing + return fb_query + + mock_query = MagicMock() + mock_query.filter_by.side_effect = mock_filter_by + + mock_session = MagicMock() + mock_session.return_value = mock_session + mock_session.query.return_value = mock_query + mock_session.add.side_effect = records.append + mock_session.commit = MagicMock() + mock_session.rollback = MagicMock() + mock_session.__enter__ = MagicMock(return_value=mock_session) + mock_session.__exit__ = MagicMock(return_value=False) + + mock_sessionmaker = MagicMock(return_value=mock_session) + + with patch("app.storage_manager.sessionmaker", mock_sessionmaker): + with patch("app.storage_manager.datetime") as mock_datetime: + mock_datetime.now.return_value = fetched_at + result = manager.save_market_data("AG2606", data) + + assert result is True + assert len(records) == 2 + assert records[0].symbol == "AG2606" + assert records[0].data_type == "futures" + assert records[0].period == "5min" + assert records[0].current_price == 123.45 + assert records[0].candle_count == 1 + assert records[0].fetched_at == fetched_at + + redis_client.delete.assert_called_once_with( + "market_data:AG2606:5min", "market_data:AG2606:15min" + ) + redis_client.setex.assert_any_call( + "market_data:AG2606:5min", + 2592000, + json.dumps( + { + "current_price": 123.45, + "timestamp": fetched_at.isoformat(), + "candles": [{"open": 100, "close": 110}], + }, + ensure_ascii=False, + ), + ) + redis_client.setex.assert_any_call( + "market_data:AG2606:15min", + 2592000, + json.dumps( + { + "current_price": 123.45, + "timestamp": fetched_at.isoformat(), + "candles": [{"open": 105, "close": 115}], + }, + ensure_ascii=False, + ), + ) + + def test_save_market_data_returns_false_and_skips_redis_on_mysql_failure(self): + """MySQL 写入失败时应回滚事务，不更新 Redis，并返回 False。""" + redis_client = MagicMock() + mysql_engine = MagicMock() + manager = StorageManager(redis_client=redis_client, mysql_engine=mysql_engine) + + data = { + "type": "futures", + "current_price": 123.45, + "timestamp": datetime(2026, 7, 4, 10, 0, 0).isoformat(), + "timeframes": {"5min": [{"open": 100, "close": 110}]}, + } + + mock_session = MagicMock() + mock_session.return_value = mock_session + mock_session.commit.side_effect = Exception("mysql write failed") + mock_session.rollback = MagicMock() + mock_session.__enter__ = MagicMock(return_value=mock_session) + mock_session.__exit__ = MagicMock(return_value=False) + mock_sessionmaker = MagicMock(return_value=mock_session) + + with patch("app.storage_manager.sessionmaker", mock_sessionmaker): + result = manager.save_market_data("AG2606", data) + + assert result is False + redis_client.delete.assert_called_once_with("market_data:AG2606:5min") + redis_client.setex.assert_not_called() + mock_session.rollback.assert_called_once() + + def test_save_market_data_returns_true_when_redis_update_fails(self, caplog): + """MySQL 成功但 Redis 回填失败时，应返回 True 并记录警告日志。""" + redis_client = MagicMock() + redis_client.setex.side_effect = Exception("redis write failed") + mysql_engine = MagicMock() + 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}]}, + } + + mock_session = MagicMock() + mock_session.return_value = mock_session + mock_session.__enter__ = MagicMock(return_value=mock_session) + mock_session.__exit__ = MagicMock(return_value=False) + mock_sessionmaker = MagicMock(return_value=mock_session) + + with patch("app.storage_manager.sessionmaker", mock_sessionmaker): + with patch("app.storage_manager.datetime") as mock_datetime: + mock_datetime.now.return_value = fetched_at + with caplog.at_level(logging.WARNING): + result = manager.save_market_data("AG2606", data) + + assert result is True + mock_session.commit.assert_called_once() + assert "Redis 更新失败" in caplog.text + + def test_save_symbol_timestamp_deletes_cache_writes_mysql_and_updates_redis(self): + """合约时间戳双写成功时，应先删缓存、再写 MySQL、最后回填 Redis。""" + redis_client = MagicMock() + mysql_engine = MagicMock() + manager = StorageManager(redis_client=redis_client, mysql_engine=mysql_engine) + + refresh_time = datetime(2026, 7, 4, 10, 0, 0) + + records = [] + + def mock_filter_by(symbol=None, data_type=None): + existing = next( + (r for r in records if r.symbol == symbol and r.data_type == data_type), + None, + ) + fb_query = MagicMock() + fb_query.first.return_value = existing + return fb_query + + mock_query = MagicMock() + mock_query.filter_by.side_effect = mock_filter_by + + mock_session = MagicMock() + mock_session.return_value = mock_session + mock_session.query.return_value = mock_query + mock_session.add.side_effect = records.append + mock_session.commit = MagicMock() + mock_session.rollback = MagicMock() + mock_session.__enter__ = MagicMock(return_value=mock_session) + mock_session.__exit__ = MagicMock(return_value=False) + + mock_sessionmaker = MagicMock(return_value=mock_session) + + with patch("app.storage_manager.sessionmaker", mock_sessionmaker): + result = manager.save_symbol_timestamp("AG2606", "futures", refresh_time) + + assert result is True + assert len(records) == 1 + assert records[0].symbol == "AG2606" + assert records[0].data_type == "futures" + assert records[0].last_refresh_at == refresh_time + assert records[0].refresh_count == 1 + + redis_client.delete.assert_called_once_with("symbol_timestamps:AG2606") + redis_client.setex.assert_called_once_with( + "symbol_timestamps:AG2606", + 2592000, + json.dumps( + {"last_refresh_at": refresh_time.isoformat(), "refresh_count": 1}, + ensure_ascii=False, + ), + ) + + def test_save_symbol_timestamp_returns_false_on_mysql_failure(self): + """合约时间戳 MySQL 写入失败时应返回 False 且不更新 Redis。""" + redis_client = MagicMock() + mysql_engine = MagicMock() + manager = StorageManager(redis_client=redis_client, mysql_engine=mysql_engine) + + mock_session = MagicMock() + mock_session.return_value = mock_session + mock_session.commit.side_effect = Exception("mysql write failed") + mock_session.rollback = MagicMock() + mock_session.__enter__ = MagicMock(return_value=mock_session) + mock_session.__exit__ = MagicMock(return_value=False) + mock_sessionmaker = MagicMock(return_value=mock_session) + + with patch("app.storage_manager.sessionmaker", mock_sessionmaker): + result = manager.save_symbol_timestamp( + "AG2606", "futures", datetime(2026, 7, 4, 10, 0, 0) + ) + + assert result is False + redis_client.delete.assert_called_once_with("symbol_timestamps:AG2606") + redis_client.setex.assert_not_called() + mock_session.rollback.assert_called_once() + + def test_delete_cache_removes_market_data_keys(self): + """delete_cache 应批量删除指定周期的 Redis 缓存键。""" + redis_client = MagicMock() + manager = StorageManager(redis_client=redis_client, mysql_engine=None) + + manager.delete_cache("AG2606", ["5min", "15min"]) + + redis_client.delete.assert_called_once_with( + "market_data:AG2606:5min", "market_data:AG2606:15min" + )