import logging import os import shutil import threading from datetime import datetime from typing import List, Optional import pandas as pd import pyarrow.dataset as ds logger = logging.getLogger("datasvc") _lock = threading.Lock() def ensure_storage(base_dir: str): os.makedirs(base_dir, exist_ok=True) def _path(base_dir: str, symbol: str, timeframe: str) -> str: safe_symbol = symbol.replace("/", "_").replace(":", "_") d = os.path.join(base_dir, timeframe) os.makedirs(d, exist_ok=True) return os.path.join(d, f"{safe_symbol}.parquet") def read_candles(base_dir: str, symbol: str, timeframe: str, start: Optional[int], end: Optional[int]) -> pd.DataFrame: p = _path(base_dir, symbol, timeframe) if not os.path.exists(p): return pd.DataFrame(columns=["timestamp", "open", "high", "low", "close", "volume"]) # empty try: df = pd.read_parquet(p) except Exception as exc: with _lock: backup = _backup_corrupted_file(p) extra = f",已备份至 {backup}" if backup else "" logger.warning( "读取缓存失败,将视为空数据 [%s %s]%s:%s", symbol, timeframe, extra, exc, ) return pd.DataFrame(columns=["timestamp", "open", "high", "low", "close", "volume"]) if start is not None: df = df[df["timestamp"] >= int(start)] if end is not None: df = df[df["timestamp"] <= int(end)] df = df.sort_values("timestamp") return df def upsert_candles(base_dir: str, symbol: str, timeframe: str, candles: List[List[float]]): p = _path(base_dir, symbol, timeframe) new_df = pd.DataFrame(candles, columns=["timestamp", "open", "high", "low", "close", "volume"]) with _lock: if os.path.exists(p): try: old = pd.read_parquet(p) except Exception as exc: backup = _backup_corrupted_file(p) extra = f",已备份至 {backup}" if backup else "" logger.warning( "读取缓存失败,准备重建文件 [%s %s]%s:%s", symbol, timeframe, extra, exc, ) old = pd.DataFrame(columns=["timestamp", "open", "high", "low", "close", "volume"]) merged = pd.concat([old, new_df], ignore_index=True) merged = merged.drop_duplicates(subset=["timestamp"], keep="last").sort_values("timestamp") else: merged = new_df.sort_values("timestamp") temp_path = f"{p}.tmp" try: merged.to_parquet(temp_path, index=False) os.replace(temp_path, p) finally: if os.path.exists(temp_path): try: os.remove(temp_path) except OSError: pass def write_candles_snapshot(base_dir: str, symbol: str, timeframe: str, df: pd.DataFrame): columns = ["timestamp", "open", "high", "low", "close", "volume"] if df.empty: safe_df = pd.DataFrame(columns=columns) else: safe_df = df[columns].copy() safe_df = safe_df.drop_duplicates(subset=["timestamp"], keep="last").sort_values("timestamp").reset_index(drop=True) p = _path(base_dir, symbol, timeframe) with _lock: temp_path = f"{p}.tmp" try: safe_df.to_parquet(temp_path, index=False) os.replace(temp_path, p) finally: if os.path.exists(temp_path): try: os.remove(temp_path) except OSError: pass def candle_path(base_dir: str, symbol: str, timeframe: str) -> str: return _path(base_dir, symbol, timeframe) def get_last_timestamp(base_dir: str, symbol: str, timeframe: str) -> Optional[int]: p = _path(base_dir, symbol, timeframe) if not os.path.exists(p): return None try: df = pd.read_parquet(p) except Exception as exc: with _lock: backup = _backup_corrupted_file(p) extra = f",已备份至 {backup}" if backup else "" logger.warning( "获取最后时间戳失败 [%s %s]%s:%s", symbol, timeframe, extra, exc, ) return None if df.empty: return None return int(df["timestamp"].iloc[-1]) def read_candle_exact(base_dir: str, symbol: str, timeframe: str, timestamp: int) -> pd.DataFrame: p = _path(base_dir, symbol, timeframe) if not os.path.exists(p): return pd.DataFrame(columns=["timestamp", "open", "high", "low", "close", "volume"]) try: dataset = ds.dataset(p, format="parquet") table = dataset.to_table(filter=ds.field("timestamp") == int(timestamp)) except Exception as exc: with _lock: backup = _backup_corrupted_file(p) extra = f",已备份至 {backup}" if backup else "" logger.warning( "读取指定时间 K 线失败 [%s %s]%s:%s", symbol, timeframe, extra, exc, ) return pd.DataFrame(columns=["timestamp", "open", "high", "low", "close", "volume"]) if table.num_rows == 0: return pd.DataFrame(columns=["timestamp", "open", "high", "low", "close", "volume"]) return table.to_pandas() def _backup_corrupted_file(path: str) -> Optional[str]: try: if not os.path.exists(path): return None timestamp = datetime.utcnow().strftime("%Y%m%d%H%M%S") backup_path = f"{path}.corrupted.{timestamp}" shutil.move(path, backup_path) return backup_path except Exception as exc: logger.warning("备份损坏文件失败 (%s):%s", path, exc) return None