spb/crypto-trade-bot

1"""WebSocket → Redis pub/sub (fan-out to API/SSE) + Postgres (latest ticker/book, trades, 1-s book samples, WS candles)."""23from __future__ import annotations45import asyncio6import json7import time8from typing import Any910import asyncpg11import redis.asyncio as aioredis1213from ktrade.ingestion.backfill import upsert_candles14from ktrade.kraken.models import PairMeta15from ktrade.kraken.ws import KrakenWS16from ktrade.logging import get_logger1718log = get_logger("ingest.stream")19CHANNEL = "ktrade:md"202122class StreamIngestor:23    def __init__(self, pool: asyncpg.Pool, redis: aioredis.Redis, ws_url: str, pairs: dict[str, PairMeta], book_depth: int, ohlc_intervals: list[int]) -> None:24        self.pool = pool25        self.redis = redis26        self.pairs = pairs27        precision = {s: (p.pair_decimals, p.lot_decimals) for s, p in pairs.items()}28        self.ws = KrakenWS(ws_url, list(pairs), self.on_event, book_depth=book_depth, ohlc_intervals=ohlc_intervals, precision=precision)29        self._last_book_sample: dict[str, int] = {}30        self._iso = None31        self.counts: dict[str, int] = {}3233    async def run(self) -> None:34        await asyncio.gather(self.ws.run(), self._health_loop())3536    async def _health_loop(self) -> None:37        while True:38            await self.redis.set("ktrade:stream:health", json.dumps({39                "connected": self.ws.connected, "last_msg_age": round(time.monotonic() - self.ws.last_msg_at, 1) if self.ws.last_msg_at else None,40                "reconnects": self.ws.reconnects, "checksum_failures": self.ws.checksum_failures, "messages": self.ws.messages, "counts": self.counts, "ts": time.time(),41            }), ex=30)42            await asyncio.sleep(5)4344    async def on_event(self, channel: str, mtype: str, data: list[dict[str, Any]], received_at: float) -> None:45        self.counts[channel] = self.counts.get(channel, 0) + 146        await self.redis.publish(CHANNEL, json.dumps({"channel": channel, "type": mtype, "data": data, "received_at": received_at}))47        try:48            if channel == "ticker":49                await self._store_ticker(data, received_at)50            elif channel == "book":51                await self._store_book(data, received_at)52            elif channel == "trade":53                await self._store_trades(data, received_at)54            elif channel == "ohlc":55                await self._store_ohlc(data)56        except Exception as e:  # noqa: BLE00157            log.warning("ingest.stream.store_error", channel=channel, error=str(e))5859    async def _store_ticker(self, data: list[dict[str, Any]], received_at: float) -> None:60        async with self.pool.acquire() as conn:61            for d in data:62                await conn.execute(63                    "INSERT INTO ticker_latest (symbol, payload, exchange_ts, received_at) VALUES ($1,$2,$3,$4) ON CONFLICT (symbol) DO UPDATE SET payload=EXCLUDED.payload, exchange_ts=EXCLUDED.exchange_ts, received_at=EXCLUDED.received_at",64                    d["symbol"], d, received_at, received_at,65                )6667    async def _store_book(self, data: list[dict[str, Any]], received_at: float) -> None:68        now = int(received_at)69        async with self.pool.acquire() as conn:70            for d in data:71                await conn.execute(72                    "INSERT INTO book_latest (symbol, payload, exchange_ts, received_at, valid) VALUES ($1,$2,$3,$4,$5) ON CONFLICT (symbol) DO UPDATE SET payload=EXCLUDED.payload, exchange_ts=EXCLUDED.exchange_ts, received_at=EXCLUDED.received_at, valid=EXCLUDED.valid",73                    d["symbol"], d, d.get("timestamp"), received_at, bool(d.get("valid", True)),74                )75                if self._last_book_sample.get(d["symbol"]) != now and d.get("best_bid"):76                    self._last_book_sample[d["symbol"]] = now77                    await conn.execute(78                        "INSERT INTO book_samples (symbol, ts, best_bid, best_ask, spread_bps, bid_depth, ask_depth, imbalance) VALUES ($1,$2,$3,$4,$5,$6,$7,$8) ON CONFLICT DO NOTHING",79                        d["symbol"], now, d["best_bid"], d["best_ask"], d.get("spread_bps"), sum(q for _, q in d["bids"]), sum(q for _, q in d["asks"]), d.get("imbalance"),80                    )8182    async def _store_trades(self, data: list[dict[str, Any]], received_at: float) -> None:83        rows = []84        for d in data:85            ts = _parse_ts(d.get("timestamp"))86            rows.append((d["symbol"], int(d["trade_id"]), ts, float(d["price"]), float(d["qty"]), d["side"], d.get("ord_type", "limit"), received_at))87        if rows:88            async with self.pool.acquire() as conn:89                await conn.executemany("INSERT INTO trades (symbol, trade_id, ts, price, volume, side, order_type, received_at) VALUES ($1,$2,$3,$4,$5,$6,$7,$8) ON CONFLICT DO NOTHING", rows)9091    async def _store_ohlc(self, data: list[dict[str, Any]]) -> None:92        now = time.time()93        async with self.pool.acquire() as conn:94            for d in data:95                iv = int(d["interval"])96                ts = int(_parse_ts(d["interval_begin"]))97                closed = ts + iv * 60 <= now98                await upsert_candles(conn, d["symbol"], iv, [(ts, float(d["open"]), float(d["high"]), float(d["low"]), float(d["close"]), float(d.get("vwap", d["close"])), float(d["volume"]), int(d.get("trades", 0)), closed)], "ws")99100101def _parse_ts(s: str | None) -> float:102    """RFC3339 with up to nanosecond precision → unix seconds (float)."""103    if not s:104        return time.time()105    import re106    from datetime import datetime107108    m = re.match(r"^(\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2})(?:\.(\d+))?(Z|[+-]\d{2}:\d{2})?$", s)109    if not m:110        return datetime.fromisoformat(s.replace("Z", "+00:00")).timestamp()111    head, frac, tz = m.group(1), m.group(2) or "", m.group(3) or "Z"112    tz = "+00:00" if tz == "Z" else tz113    return datetime.fromisoformat(f"{head}.{frac[:6].ljust(6, '0')}{tz}").timestamp()114