Chan/datasvc/app/storage.py

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