feat(storage): add StorageManager write logic with dual-write consistency

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
 

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"
+        )