148 lines
4.9 KiB
Python
148 lines
4.9 KiB
Python
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 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
|
||
|
||
|