from __future__ import annotations import asyncio import time from dataclasses import dataclass from datetime import datetime, timedelta from pathlib import Path from typing import Any, Dict, List, Tuple, Set import pyarrow.parquet as pq from app.config import settings from app.streaming.live_store import live_store, _TIMEFRAME_SECONDS @dataclass(frozen=True) class ReplayKey: symbol: str timeframe: str class HistoricalReplayProvider: """Streams OHLCV data by replaying parquet partitions as live candles.""" def __init__( self, root: str | Path | None = None, speed: float | None = None, days: int | None = None, loop: bool | None = None, ) -> None: default_root = Path(__file__).resolve().parents[3] / "data" / "parquet" / "live" resolved_root = Path(root) if root else Path(settings.HISTORICAL_REPLAY_ROOT or default_root) self._root = resolved_root self._speed = max(speed or settings.HISTORICAL_REPLAY_SPEED, 0.1) self._days_back = max(days or settings.HISTORICAL_REPLAY_DAYS, 1) self._loop = settings.HISTORICAL_REPLAY_LOOP if loop is None else loop self._tasks: Dict[ReplayKey, asyncio.Task] = {} self._pinned: Set[ReplayKey] = set() self._lock = asyncio.Lock() self._active_counts: Dict[ReplayKey, int] = {} self._bar_cache: Dict[ReplayKey, Tuple[float, List[Dict[str, Any]]]] = {} self._positions: Dict[ReplayKey, int] = {} self._last_times: Dict[ReplayKey, int] = {} def _sleep_seconds(self, timeframe: str) -> float: base = _TIMEFRAME_SECONDS.get(timeframe.lower(), 60) return max(base / self._speed, 0.5) def _ensure_history(self, key: ReplayKey) -> None: try: if not live_store.get_history(key.symbol, key.timeframe): live_store.load_historical_data(key.symbol, key.timeframe, days_back=self._days_back) except Exception: # Best-effort warmup; ignore errors so streaming can proceed pass def _resolve_partitions(self, key: ReplayKey) -> List[Tuple[datetime, Path]]: base_path = self._root / key.symbol / key.timeframe partitions: List[Tuple[datetime, Path]] = [] if not base_path.exists(): return partitions for part in base_path.glob("date=*"): if not part.is_dir(): continue _, _, date_part = part.name.partition("=") try: dt = datetime.strptime(date_part, "%Y-%m-%d") except ValueError: continue partitions.append((dt, part)) partitions.sort(key=lambda x: x[0]) return partitions def _load_bars(self, key: ReplayKey) -> List[Dict[str, Any]]: cached = self._bar_cache.get(key) now = time.time() if cached and now - cached[0] < 300: return cached[1] partitions = self._resolve_partitions(key) if not partitions: self._bar_cache[key] = (now, []) return [] cutoff_date = partitions[-1][0] - timedelta(days=self._days_back - 1) eligible = [p for p in partitions if p[0] >= cutoff_date] if not eligible: eligible = partitions[-self._days_back :] if len(partitions) >= self._days_back else partitions rows: List[Dict[str, Any]] = [] for _, part in eligible: files = sorted(part.glob("*.parquet")) for file in files: try: table = pq.read_table(file, columns=["time", "open", "high", "low", "close", "volume"]) except Exception: continue for row in table.to_pylist(): try: rows.append( { "time": int(row["time"]), "open": float(row["open"]), "high": float(row["high"]), "low": float(row["low"]), "close": float(row["close"]), "volume": float(row.get("volume") or 0.0), } ) except Exception: continue rows.sort(key=lambda r: r["time"]) self._bar_cache[key] = (now, rows) return rows async def subscribe(self, symbol: str, timeframe: str = "1m") -> Tuple[asyncio.Queue, Any]: if timeframe != "1m": raise ValueError("HistoricalReplayProvider currently supports timeframe '1m' only") key = ReplayKey(symbol=symbol.upper().replace("/", ""), timeframe=timeframe) queue: asyncio.Queue = asyncio.Queue(maxsize=200) live_store.subscribe(key.symbol, key.timeframe, queue) async with self._lock: self._active_counts[key] = self._active_counts.get(key, 0) + 1 if key not in self._tasks: self._tasks[key] = asyncio.create_task(self._run_replay(key)) async def _unsubscribe() -> None: live_store.unsubscribe(key.symbol, key.timeframe, queue) async with self._lock: self._active_counts[key] = max(0, self._active_counts.get(key, 0) - 1) if self._active_counts.get(key, 0) == 0 and key not in self._pinned: task = self._tasks.pop(key, None) if task: task.cancel() self._active_counts.pop(key, None) return queue, _unsubscribe async def ensure_stream(self, symbol: str, timeframe: str = "1m") -> None: if timeframe != "1m": raise ValueError("HistoricalReplayProvider currently supports timeframe '1m' only") key = ReplayKey(symbol=symbol.upper().replace("/", ""), timeframe=timeframe) async with self._lock: self._pinned.add(key) self._active_counts.setdefault(key, 0) if key not in self._tasks: self._tasks[key] = asyncio.create_task(self._run_replay(key)) async def release_stream(self, symbol: str, timeframe: str = "1m") -> None: key = ReplayKey(symbol=symbol.upper().replace("/", ""), timeframe=timeframe) async with self._lock: self._pinned.discard(key) if self._active_counts.get(key, 0) == 0: task = self._tasks.pop(key, None) if task: task.cancel() self._active_counts.pop(key, None) async def _run_replay(self, key: ReplayKey) -> None: self._ensure_history(key) sleep_secs = self._sleep_seconds(key.timeframe) tf_seconds = _TIMEFRAME_SECONDS.get(key.timeframe.lower(), 60) self._positions.setdefault(key, 0) self._last_times.setdefault(key, 0) while True: try: bars = self._load_bars(key) if not bars: await asyncio.sleep(5.0) continue idx = self._positions.get(key, 0) if idx >= len(bars): if not self._loop: await asyncio.sleep(sleep_secs) continue idx = 0 bar = bars[idx] self._positions[key] = idx + 1 now = int(time.time()) aligned = (now // tf_seconds) * tf_seconds last_time = self._last_times.get(key) or 0 if aligned <= last_time: aligned = last_time + tf_seconds payload = { "time": aligned, "open": bar["open"], "high": bar["high"], "low": bar["low"], "close": bar["close"], "volume": bar.get("volume", 0.0), } live_store.ingest_bar(symbol=key.symbol, timeframe=key.timeframe, bar=payload) self._last_times[key] = aligned await asyncio.sleep(sleep_secs) except asyncio.CancelledError: break except Exception: await asyncio.sleep(min(5.0, sleep_secs)) historical_replay = HistoricalReplayProvider()