amazing-data-service/backend/app/services/realtime_service.py

"""
实时数据服务
"""
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)