""" 实时数据服务 """ import asyncio import logging from typing import Dict, Set, List, Callable, Optional from datetime import datetime, timedelta from sqlalchemy.orm import Session from fastapi import WebSocket from app.models.realtime import RealtimeSnapshot from app.services.base_data_service import BaseDataService from app.config import settings logger = logging.getLogger(__name__) class RealtimeManager: """实时数据管理器(单例)""" _instance = None def __new__(cls): if cls._instance is None: cls._instance = super().__new__(cls) cls._instance._initialized = False return cls._instance def __init__(self): if self._initialized: return self.subscribers: Dict[str, Set[WebSocket]] = {} self.code_callbacks: Dict[str, List[Callable]] = {} self._adapter = None self._initialized = True self._lock = asyncio.Lock() async def subscribe(self, websocket: WebSocket, codes: List[str]): """ 客户端订阅实时数据 Args: websocket: WebSocket连接 codes: 代码列表 """ await websocket.accept() async with self._lock: for code in codes: if code not in self.subscribers: self.subscribers[code] = set() # 启动SDK订阅 await self._start_sdk_subscription(code) self.subscribers[code].add(websocket) logger.info(f"WebSocket {id(websocket)} 订阅了: {codes}") async def unsubscribe(self, websocket: WebSocket, codes: List[str] = None): """ 取消订阅 Args: websocket: WebSocket连接 codes: 代码列表,None表示取消所有 """ async with self._lock: codes_to_remove = codes if codes else list(self.subscribers.keys()) for code in codes_to_remove: if code in self.subscribers: self.subscribers[code].discard(websocket) # 如果没有订阅者了,取消SDK订阅 if not self.subscribers[code]: del self.subscribers[code] await self._stop_sdk_subscription(code) logger.info(f"WebSocket {id(websocket)} 取消订阅") async def _start_sdk_subscription(self, code: str): """启动SDK订阅""" # 这里需要实现实际的SDK订阅逻辑 # 由于SDK的实时订阅是同步的回调,需要在后台线程中运行 logger.info(f"开始SDK订阅: {code}") async def _stop_sdk_subscription(self, code: str): """停止SDK订阅""" logger.info(f"停止SDK订阅: {code}") def on_sdk_data(self, code: str, data: dict): """ SDK数据回调 Args: code: 代码 data: 数据字典 """ # 保存到数据库 # self._save_snapshot(code, data) # 推送给所有订阅者 if code in self.subscribers: message = { "type": "snapshot", "code": code, "data": data, "timestamp": datetime.utcnow().isoformat() } # 异步推送 for ws in self.subscribers[code]: asyncio.create_task(self._send_to_ws(ws, message)) async def _send_to_ws(self, websocket: WebSocket, message: dict): """发送消息到WebSocket""" try: await websocket.send_json(message) except Exception as e: logger.error(f"发送WebSocket消息失败: {str(e)}") # 从订阅列表中移除 await self.unsubscribe(websocket) def _save_snapshot(self, db: Session, code: str, data: dict): """保存快照到数据库""" try: expires_at = datetime.utcnow() + timedelta(days=settings.CACHE_AUTO_CLEANUP_DAYS) snapshot = RealtimeSnapshot( code=code, security_type=data.get("security_type", "stock"), trade_time=datetime.fromisoformat(data.get("trade_time", datetime.utcnow().isoformat())), pre_close=data.get("pre_close"), last=data.get("last"), open=data.get("open"), high=data.get("high"), low=data.get("low"), close=data.get("close"), volume=data.get("volume"), amount=data.get("amount"), expires_at=expires_at ) db.add(snapshot) db.commit() except Exception as e: logger.error(f"保存快照失败: {str(e)}") # 全局实时数据管理器实例 realtime_manager = RealtimeManager() class RealtimeService: """实时数据服务""" def __init__(self, db: Session): self.db = db self.base_service = BaseDataService(db) self.manager = realtime_manager def get_latest_snapshot(self, codes: List[str]) -> Dict[str, dict]: """ 获取最新快照数据 Args: codes: 代码列表 Returns: 快照数据字典 """ result = {} for code in codes: # 查询最新的快照 snapshot = self.db.query(RealtimeSnapshot).filter( RealtimeSnapshot.code == code ).order_by(RealtimeSnapshot.trade_time.desc()).first() if snapshot: result[code] = { "code": snapshot.code, "trade_time": snapshot.trade_time.isoformat(), "pre_close": float(snapshot.pre_close) if snapshot.pre_close else None, "last": float(snapshot.last) if snapshot.last else None, "open": float(snapshot.open) if snapshot.open else None, "high": float(snapshot.high) if snapshot.high else None, "low": float(snapshot.low) if snapshot.low else None, "close": float(snapshot.close) if snapshot.close else None, "volume": int(snapshot.volume) if snapshot.volume else None, "amount": float(snapshot.amount) if snapshot.amount else None } else: result[code] = None return result async def subscribe_websocket(self, websocket: WebSocket, codes: List[str]): """订阅WebSocket""" await self.manager.subscribe(websocket, codes) async def unsubscribe_websocket(self, websocket: WebSocket, codes: List[str] = None): """取消WebSocket订阅""" await self.manager.unsubscribe(websocket, codes)