115 lines
4.4 KiB
Python
115 lines
4.4 KiB
Python
"""数据编排：拉取（Tushare 主 -> AKShare 兜底）+ 本地缓存。
|
||
|
||
真实行情落库到 candles 表（timeframe='1d'，**不复权底座**），回测统一从库读。
|
||
复权（qfq/hfq）在读取时按 adj_factor 表本地换算，见 api._adjust_bars。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
|
||
from sqlalchemy import func, select
|
||
from sqlalchemy.dialects.postgresql import insert as pg_insert
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
||
from ..config import settings
|
||
from ..domain import Bar
|
||
from ..models import Candle
|
||
from . import akshare_provider, tushare_provider
|
||
|
||
DEFAULT_START = settings.data_default_start or "20200101"
|
||
|
||
|
||
def _providers(source: str):
|
||
"""按优先级返回 (名称, 同步拉取函数) 列表。"""
|
||
seq = []
|
||
if source in ("auto", "tushare") and settings.tushare_token:
|
||
seq.append(("tushare", tushare_provider.fetch_daily))
|
||
if source in ("auto", "akshare"):
|
||
seq.append(("akshare", akshare_provider.fetch_daily))
|
||
return seq
|
||
|
||
|
||
async def count_cached(session: AsyncSession, symbol: str) -> int:
|
||
res = await session.execute(
|
||
select(func.count()).select_from(Candle).where(
|
||
Candle.symbol == symbol, Candle.timeframe == "1d"
|
||
)
|
||
)
|
||
return int(res.scalar() or 0)
|
||
|
||
|
||
async def is_cached(session: AsyncSession, symbol: str) -> bool:
|
||
return await count_cached(session, symbol) > 0
|
||
|
||
|
||
async def _last_cached_ts(session: AsyncSession, symbol: str):
|
||
res = await session.execute(
|
||
select(func.max(Candle.ts)).where(Candle.symbol == symbol, Candle.timeframe == "1d")
|
||
)
|
||
return res.scalar()
|
||
|
||
|
||
async def sync_symbol(
|
||
session: AsyncSession,
|
||
code: str,
|
||
start: str | None = None,
|
||
end: str | None = None,
|
||
source: str = "auto",
|
||
force: bool = False,
|
||
) -> dict:
|
||
"""增量拉取并 upsert 某标的日线（**不复权**底座）。
|
||
|
||
- 永不删除已有行：按 (symbol, timeframe, ts) 主键 upsert，
|
||
不会把 TDX 导入的 30 年历史冲掉；
|
||
- 已缓存时从最后一根的次日开始增量拉取（force 仅跳过「有缓存就返回」
|
||
的短路，用于缓存落后于最新交易日时的刷新）；
|
||
- 拉不到新行时保持原缓存不动。
|
||
"""
|
||
last_ts = await _last_cached_ts(session, code)
|
||
if last_ts is not None and not force and not start:
|
||
return {"symbol": code, "bars": await count_cached(session, code), "source": "cache"}
|
||
|
||
if last_ts is not None and not start:
|
||
# 增量：从缓存最后一根当天开始（重叠一天重新拉取，容忍数据源漏行/盘后修订）
|
||
start = last_ts.strftime("%Y%m%d")
|
||
start = start or DEFAULT_START
|
||
errors: list[str] = []
|
||
bars: list[Bar] = []
|
||
used = None
|
||
|
||
for name, fn in _providers(source):
|
||
try:
|
||
# tushare/akshare 是同步网络 IO，丢到线程池避免阻塞事件循环；
|
||
# adjust=None -> 不复权（复权在读取时按 adj_factor 换算）
|
||
bars = await asyncio.to_thread(fn, code, start, end, None)
|
||
used = name
|
||
break
|
||
except Exception as e: # noqa: BLE001
|
||
errors.append(f"{name}: {e}")
|
||
|
||
if not bars:
|
||
if last_ts is not None:
|
||
# 增量失败（如停牌/新股无新行）：保留缓存，不算错误
|
||
return {"symbol": code, "bars": await count_cached(session, code), "source": "cache"}
|
||
raise RuntimeError("所有数据源均失败 -> " + " | ".join(errors) if errors else "无可用数据源")
|
||
|
||
# upsert：不 delete，避免破坏既有底座（TDX 全量历史）
|
||
stmt = pg_insert(Candle).values([
|
||
{"symbol": code, "timeframe": "1d", "ts": b.ts, "open": b.open, "high": b.high,
|
||
"low": b.low, "close": b.close, "volume": b.volume,
|
||
"amount": b.amount, "turnover": b.turnover}
|
||
for b in bars
|
||
])
|
||
stmt = stmt.on_conflict_do_update(
|
||
index_elements=["symbol", "timeframe", "ts"],
|
||
set_={"open": stmt.excluded.open, "high": stmt.excluded.high, "low": stmt.excluded.low,
|
||
"close": stmt.excluded.close, "volume": stmt.excluded.volume,
|
||
# 增量源缺失额/换手时保留库里的旧值（如 TDX 已回补的 30 年成交额）
|
||
"amount": func.coalesce(stmt.excluded.amount, Candle.amount),
|
||
"turnover": func.coalesce(stmt.excluded.turnover, Candle.turnover)},
|
||
)
|
||
await session.execute(stmt)
|
||
await session.commit()
|
||
return {"symbol": code, "bars": len(bars), "source": used}
|