90 lines
2.8 KiB
Python
90 lines
2.8 KiB
Python
import asyncio
|
|
from contextlib import asynccontextmanager
|
|
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Depends
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from loguru import logger
|
|
|
|
from app.core.config import settings
|
|
from app.core.database import init_db
|
|
from app.core.redis import close_redis
|
|
from app.api.v1 import api_router
|
|
from app.websocket.manager import manager
|
|
from app.services import stock_service
|
|
from app.api.deps import get_current_user
|
|
|
|
|
|
# ── background task: push heatmap every 5s ───────────────────────────────────
|
|
|
|
async def _heatmap_pusher():
|
|
while True:
|
|
try:
|
|
data = await stock_service.get_all_stocks_spot()
|
|
await manager.broadcast_all(data)
|
|
except Exception as e:
|
|
logger.warning(f"heatmap pusher error: {e}")
|
|
await asyncio.sleep(5)
|
|
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(app: FastAPI):
|
|
logger.info("Starting up...")
|
|
await init_db()
|
|
task = asyncio.create_task(_heatmap_pusher())
|
|
yield
|
|
task.cancel()
|
|
await close_redis()
|
|
logger.info("Shut down.")
|
|
|
|
|
|
# ── app ───────────────────────────────────────────────────────────────────────
|
|
|
|
app = FastAPI(
|
|
title=settings.APP_NAME,
|
|
docs_url="/api/docs",
|
|
redoc_url="/api/redoc",
|
|
openapi_url="/api/openapi.json",
|
|
lifespan=lifespan,
|
|
)
|
|
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=settings.cors_origins,
|
|
allow_credentials=True,
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
app.include_router(api_router, prefix=settings.API_V1_PREFIX)
|
|
|
|
|
|
# ── WebSocket endpoints ───────────────────────────────────────────────────────
|
|
|
|
@app.websocket("/ws/heatmap")
|
|
async def ws_heatmap(websocket: WebSocket):
|
|
"""Subscribe to full market heatmap updates."""
|
|
await manager.connect(websocket)
|
|
try:
|
|
while True:
|
|
await websocket.receive_text()
|
|
except WebSocketDisconnect:
|
|
manager.disconnect(websocket)
|
|
|
|
|
|
@app.websocket("/ws/quote/{symbol}")
|
|
async def ws_quote(websocket: WebSocket, symbol: str):
|
|
"""Subscribe to real-time quote updates for a single stock."""
|
|
await manager.connect(websocket, symbol)
|
|
try:
|
|
while True:
|
|
await asyncio.sleep(3)
|
|
data = await stock_service.get_stock_quote(symbol)
|
|
if data:
|
|
await manager.broadcast_quote(symbol, data)
|
|
except WebSocketDisconnect:
|
|
manager.disconnect(websocket, symbol)
|
|
|
|
|
|
@app.get("/api/health")
|
|
async def health():
|
|
return {"status": "ok"}
|