""" K 线数据读取器 —— 从 TimescaleDB 查询历史 K 线供策略分析和回测使用。 所有方法均为只读查询,对应 data 模块 Kline 实体的字段结构。 参考:data/db/entities/kline.entity.ts """ from datetime import datetime from typing import Optional, Sequence import asyncpg from engine.data.db import get_pool from engine.data.models import KlineRecord class KlineReader: """TimescaleDB K 线只读查询器""" def __init__(self): self._pool: Optional[asyncpg.Pool] = None async def _ensure_pool(self) -> asyncpg.Pool: if self._pool is None: self._pool = await get_pool() return self._pool async def get_klines( self, symbol: str, interval: str, start_time: datetime, end_time: datetime, exchange: str = "binance", limit: int = 1000, ) -> list[KlineRecord]: """ 查询指定时间范围内的 K 线数据,支持任意周期(1m / 5m / 15m / 1h / 4h / 1d 等)。 Args: symbol: 交易对(如 BTCUSDT) interval: K 线周期(1m / 5m / 15m / 30m / 1h / 4h / 1d 等) start_time: 起始时间(含) end_time: 结束时间(含) exchange: 交易所标识 limit: 最大返回条数 """ pool = await self._ensure_pool() query = """ SELECT time, exchange, symbol, interval, open, high, low, close, volume, quote_volume, taker_buy_base_vol, taker_buy_quote_vol, trade_count, is_closed, created_at, updated_at FROM klines WHERE exchange = $1 AND symbol = $2 AND interval = $3 AND time >= $4 AND time <= $5 ORDER BY time ASC LIMIT $6 """ rows: Sequence[asyncpg.Record] = await pool.fetch( query, exchange, symbol, interval, start_time, end_time, limit, ) return [KlineRecord.from_record(row) for row in rows] async def get_latest_klines( self, symbol: str, interval: str, exchange: str = "binance", limit: int = 500, ) -> list[KlineRecord]: """ 获取最近 N 根已闭合的 K 线(策略启动时快速预热)。 仅查询 is_closed = TRUE 的 K 线,避免使用未闭合的不完整数据。 """ pool = await self._ensure_pool() query = """ SELECT time, exchange, symbol, interval, open, high, low, close, volume, quote_volume, taker_buy_base_vol, taker_buy_quote_vol, trade_count, is_closed, created_at, updated_at FROM klines WHERE exchange = $1 AND symbol = $2 AND interval = $3 AND is_closed = TRUE ORDER BY time DESC LIMIT $4 """ rows = await pool.fetch(query, exchange, symbol, interval, limit) records = [KlineRecord.from_record(row) for row in rows] records.reverse() return records async def get_klines_by_count( self, symbol: str, interval: str, count: int = 100, exchange: str = "binance", before_time: Optional[datetime] = None, ) -> list[KlineRecord]: """ 获取最新的 N 根 K 线(按时间倒序取 N 条再正序返回)。 适合策略运行时获取最近 N 根 K 线做指标计算,不关心特定时间范围。 Args: symbol: 交易对 interval: K 线周期 count: 需要获取的 K 线数量 exchange: 交易所标识 before_time: 可选,只获取该时间之前的 K 线 """ pool = await self._ensure_pool() if before_time: query = """ SELECT time, exchange, symbol, interval, open, high, low, close, volume, quote_volume, taker_buy_base_vol, taker_buy_quote_vol, trade_count, is_closed, created_at, updated_at FROM klines WHERE exchange = $1 AND symbol = $2 AND interval = $3 AND time <= $4 ORDER BY time DESC LIMIT $5 """ rows = await pool.fetch(query, exchange, symbol, interval, before_time, count) else: query = """ SELECT time, exchange, symbol, interval, open, high, low, close, volume, quote_volume, taker_buy_base_vol, taker_buy_quote_vol, trade_count, is_closed, created_at, updated_at FROM klines WHERE exchange = $1 AND symbol = $2 AND interval = $3 ORDER BY time DESC LIMIT $4 """ rows = await pool.fetch(query, exchange, symbol, interval, count) records = [KlineRecord.from_record(row) for row in rows] records.reverse() return records async def get_ohlcv_array( self, symbol: str, interval: str, exchange: str = "binance", limit: int = 200, before_time: Optional[datetime] = None, ) -> list[tuple[datetime, float, float, float, float, float]]: """ 获取 OHLCV 元组数组,直接用于 pandas DataFrame 或 TA-Lib 计算。 返回格式: [(timestamp, open, high, low, close, volume), ...] 时间升序排列。 """ pool = await self._ensure_pool() if before_time: query = """ SELECT time, open, high, low, close, volume FROM klines WHERE exchange = $1 AND symbol = $2 AND interval = $3 AND time <= $4 ORDER BY time DESC LIMIT $5 """ rows = await pool.fetch(query, exchange, symbol, interval, before_time, limit) else: query = """ SELECT time, open, high, low, close, volume FROM klines WHERE exchange = $1 AND symbol = $2 AND interval = $3 ORDER BY time DESC LIMIT $4 """ rows = await pool.fetch(query, exchange, symbol, interval, limit) return [ (row["time"], float(row["open"]), float(row["high"]), float(row["low"]), float(row["close"]), float(row["volume"])) for row in reversed(rows) ] async def get_available_symbols( self, exchange: str = "binance", ) -> list[str]: """ 查询已激活的交易对列表。 从 trading_pairs 配置表读取(仅 active=TRUE 的记录), 而非扫描 klines 时序表,避免全表扫描。 对应 data/db/entities/trading-pair.entity.ts。 """ pool = await self._ensure_pool() rows = await pool.fetch( """ SELECT tp.symbol FROM trading_pairs tp JOIN exchanges e ON tp.exchange_id = e.id WHERE e.name = $1 AND tp.active = TRUE ORDER BY tp.symbol """, exchange, ) return [row["symbol"] for row in rows] async def get_available_intervals( self, symbol: str, exchange: str = "binance", ) -> list[str]: """ 查询某交易对配置的 K 线周期列表。 从 trading_pairs.kline_intervals 读取(逗号分隔字符串), 而非扫描 klines 时序表,避免全表扫描。 对应 data/db/entities/trading-pair.entity.ts: kline_intervals 字段。 """ pool = await self._ensure_pool() row = await pool.fetchrow( """ SELECT tp.kline_intervals FROM trading_pairs tp JOIN exchanges e ON tp.exchange_id = e.id WHERE e.name = $1 AND tp.symbol = $2 AND tp.active = TRUE """, exchange, symbol, ) if row is None or not row["kline_intervals"]: return [] # kline_intervals 格式: "1m,5m,15m,1h,4h,1d" return [ interval.strip() for interval in row["kline_intervals"].split(",") if interval.strip() ]