208 lines
6.4 KiB
Python
208 lines
6.4 KiB
Python
import os
|
||
import asyncio
|
||
import json
|
||
from datetime import datetime, timedelta
|
||
from typing import Dict, List, Optional
|
||
|
||
import ccxt
|
||
import pandas as pd
|
||
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Query
|
||
from fastapi.responses import JSONResponse
|
||
from fastapi.middleware.cors import CORSMiddleware
|
||
|
||
from .storage import (
|
||
ensure_storage,
|
||
read_candles,
|
||
upsert_candles,
|
||
get_last_timestamp,
|
||
)
|
||
|
||
|
||
DATA_DIR = os.environ.get("DATA_DIR", "/data")
|
||
EXCHANGE = os.environ.get("EXCHANGE", "binance")
|
||
SYMBOLS = [s.strip() for s in os.environ.get("SYMBOLS", "BTC/USDT:USDT,ETH/USDT:USDT").split(",") if s.strip()]
|
||
TIMEFRAMES = [t.strip() for t in os.environ.get("TIMEFRAMES", "1m,5m,15m,1h").split(",") if t.strip()]
|
||
START_FROM = os.environ.get("START_FROM", "2025-01-01") # 首次启动拉取起始日期（UTC）
|
||
POLL_FACTOR = float(os.environ.get("POLL_FACTOR", "0.5")) # 轮询间隔 = tf_ms * factor
|
||
|
||
ensure_storage(DATA_DIR)
|
||
|
||
app = FastAPI(title="Local Data Service", version="0.1.0")
|
||
app.add_middleware(
|
||
CORSMiddleware,
|
||
allow_origins=["*"],
|
||
allow_credentials=True,
|
||
allow_methods=["*"],
|
||
allow_headers=["*"],
|
||
)
|
||
|
||
|
||
def tf_to_ms(tf: str) -> int:
|
||
table = {
|
||
"1m": 60_000,
|
||
"3m": 3 * 60_000,
|
||
"5m": 5 * 60_000,
|
||
"15m": 15 * 60_000,
|
||
"30m": 30 * 60_000,
|
||
"1h": 60 * 60_000,
|
||
"2h": 2 * 60 * 60_000,
|
||
"4h": 4 * 60 * 60_000,
|
||
"1d": 24 * 60 * 60_000,
|
||
}
|
||
return table.get(tf, 60_000)
|
||
|
||
|
||
def parse_start_from_ms(val: str) -> int:
|
||
"""将 START_FROM 解析成毫秒级时间戳。
|
||
支持两种格式：
|
||
- YYYY-MM-DD（UTC 00:00:00）
|
||
- 整型毫秒时间戳字符串
|
||
"""
|
||
try:
|
||
return int(val)
|
||
except Exception:
|
||
pass
|
||
try:
|
||
dt = datetime.fromisoformat(val) # 允许 '2025-01-01' 或 '2025-01-01T00:00:00'
|
||
except Exception:
|
||
# 回退到固定日期
|
||
dt = datetime(2025, 1, 1)
|
||
return int(dt.timestamp() * 1000)
|
||
|
||
|
||
class Hub:
|
||
def __init__(self) -> None:
|
||
self.subscribers: Dict[str, List[WebSocket]] = {}
|
||
|
||
def topic(self, symbol: str, timeframe: str) -> str:
|
||
return f"candles::{symbol}::{timeframe}"
|
||
|
||
async def subscribe(self, ws: WebSocket, symbol: str, timeframe: str):
|
||
topic = self.topic(symbol, timeframe)
|
||
await ws.accept()
|
||
self.subscribers.setdefault(topic, []).append(ws)
|
||
|
||
def _clean(self, topic: str):
|
||
conns = self.subscribers.get(topic, [])
|
||
self.subscribers[topic] = [w for w in conns if not w.client_state.name == "DISCONNECTED"]
|
||
|
||
async def publish(self, symbol: str, timeframe: str, payload: dict):
|
||
topic = self.topic(symbol, timeframe)
|
||
conns = self.subscribers.get(topic, [])
|
||
if not conns:
|
||
return
|
||
message = json.dumps(payload, ensure_ascii=False)
|
||
dead: List[WebSocket] = []
|
||
for ws in conns:
|
||
try:
|
||
await ws.send_text(message)
|
||
except Exception:
|
||
dead.append(ws)
|
||
if dead:
|
||
self.subscribers[topic] = [w for w in conns if w not in dead]
|
||
|
||
|
||
hub = Hub()
|
||
|
||
|
||
def build_exchange():
|
||
if EXCHANGE.lower() == "binance":
|
||
return ccxt.binance({"enableRateLimit": True})
|
||
raise RuntimeError(f"Unsupported EXCHANGE: {EXCHANGE}")
|
||
|
||
|
||
async def fetch_loop(symbol: str, timeframe: str):
|
||
"""持续增量抓取并广播。"""
|
||
exchange = build_exchange()
|
||
tf_ms = tf_to_ms(timeframe)
|
||
start_since = parse_start_from_ms(START_FROM)
|
||
last_ts = get_last_timestamp(DATA_DIR, symbol, timeframe)
|
||
since = max(start_since, (last_ts + tf_ms) if last_ts else start_since)
|
||
|
||
while True:
|
||
try:
|
||
candles = exchange.fetch_ohlcv(symbol, timeframe, since=since, limit=1000)
|
||
if candles:
|
||
upsert_candles(DATA_DIR, symbol, timeframe, candles)
|
||
for row in candles[-3:]:
|
||
payload = {
|
||
"topic": f"candles.{symbol}.{timeframe}",
|
||
"type": "upsert",
|
||
"data": {
|
||
"t": row[0],
|
||
"o": row[1],
|
||
"h": row[2],
|
||
"l": row[3],
|
||
"c": row[4],
|
||
"v": row[5],
|
||
},
|
||
}
|
||
await hub.publish(symbol, timeframe, payload)
|
||
since = candles[-1][0] + tf_ms
|
||
await asyncio.sleep(max(1.0, tf_ms * POLL_FACTOR / 1000.0))
|
||
except Exception:
|
||
await asyncio.sleep(3.0)
|
||
|
||
|
||
@app.on_event("startup")
|
||
async def on_start():
|
||
ensure_storage(DATA_DIR)
|
||
for s in SYMBOLS:
|
||
for tf in TIMEFRAMES:
|
||
asyncio.create_task(fetch_loop(s, tf))
|
||
|
||
|
||
@app.get("/api/candles")
|
||
def api_candles(
|
||
symbol: str = Query(..., description="如 BTC/USDT:USDT"),
|
||
tf: str = Query("1m", description="时间周期"),
|
||
start: Optional[int] = Query(None, description="开始时间戳(ms)"),
|
||
end: Optional[int] = Query(None, description="结束时间戳(ms)"),
|
||
):
|
||
try:
|
||
df = read_candles(DATA_DIR, symbol, tf, start, end)
|
||
records = df.to_dict("records") if not df.empty else []
|
||
return JSONResponse(records)
|
||
except Exception as e:
|
||
return JSONResponse({"error": str(e)}, status_code=500)
|
||
|
||
|
||
@app.websocket("/ws")
|
||
async def ws_endpoint(websocket: WebSocket, symbol: str, tf: str, since: Optional[int] = None):
|
||
await hub.subscribe(websocket, symbol, tf)
|
||
try:
|
||
snap = read_candles(DATA_DIR, symbol, tf, since, None)
|
||
await websocket.send_text(
|
||
json.dumps(
|
||
{
|
||
"topic": f"candles.{symbol}.{tf}",
|
||
"type": "snapshot",
|
||
"data": [
|
||
{"t": int(r["timestamp"]), "o": r["open"], "h": r["high"], "l": r["low"], "c": r["close"], "v": r["volume"]}
|
||
for _, r in snap.iterrows()
|
||
],
|
||
},
|
||
ensure_ascii=False,
|
||
)
|
||
)
|
||
except Exception:
|
||
pass
|
||
try:
|
||
while True:
|
||
await asyncio.sleep(30)
|
||
await websocket.send_text(json.dumps({"type": "ping", "ts": int(datetime.utcnow().timestamp() * 1000)}))
|
||
except WebSocketDisconnect:
|
||
return
|
||
|
||
|
||
@app.get("/")
|
||
def root():
|
||
return {
|
||
"service": "Local Data Service",
|
||
"exchange": EXCHANGE,
|
||
"symbols": SYMBOLS,
|
||
"timeframes": TIMEFRAMES,
|
||
}
|
||
|
||
|