初始化工程
This commit is contained in:
@@ -0,0 +1 @@
|
||||
"""数据源适配层 — 能力探测 / 调度 / Repository。"""
|
||||
@@ -0,0 +1,76 @@
|
||||
"""Capability 定义(§5.1)。
|
||||
|
||||
业务代码只依赖 CapabilitySet,不读 tiers.yaml,不感知"档位"。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from enum import StrEnum
|
||||
|
||||
|
||||
class Cap(StrEnum):
|
||||
"""所有 capability 的命名常量。新增能力时只在这里加一行。"""
|
||||
|
||||
QUOTE_BY_SYMBOL = "quote.by_symbol"
|
||||
QUOTE_BATCH = "quote.batch"
|
||||
QUOTE_POOL = "quote.pool"
|
||||
KLINE_DAILY_BY_SYMBOL = "kline.daily.by_symbol"
|
||||
KLINE_DAILY_BATCH = "kline.daily.batch"
|
||||
KLINE_MINUTE_BY_SYMBOL = "kline.minute.by_symbol"
|
||||
KLINE_MINUTE_BATCH = "kline.minute.batch"
|
||||
INTRADAY = "intraday"
|
||||
INTRADAY_BATCH = "intraday.batch"
|
||||
DEPTH5 = "depth5"
|
||||
DEPTH5_BATCH = "depth5.batch"
|
||||
WEBSOCKET = "websocket"
|
||||
FINANCIAL = "financial"
|
||||
ADJ_FACTOR = "adj_factor"
|
||||
|
||||
|
||||
@dataclass(slots=True, frozen=True)
|
||||
class CapabilityLimits:
|
||||
"""单个 capability 的运行时限制。"""
|
||||
rpm: int | None = None # 次/分钟,None 表示未知或不限
|
||||
batch: int | None = None # 标的/次
|
||||
subscribe: int | None = None # WS 订阅上限
|
||||
|
||||
|
||||
class CapabilitySet:
|
||||
"""探测得到的"用户当前可用能力"。业务代码的唯一真理源。"""
|
||||
|
||||
def __init__(self, caps: dict[Cap, CapabilityLimits] | None = None) -> None:
|
||||
self._caps: dict[Cap, CapabilityLimits] = dict(caps or {})
|
||||
|
||||
def has(self, cap: Cap) -> bool:
|
||||
return cap in self._caps
|
||||
|
||||
def limits(self, cap: Cap) -> CapabilityLimits | None:
|
||||
return self._caps.get(cap)
|
||||
|
||||
def require(self, cap: Cap) -> CapabilityLimits:
|
||||
"""断言可用,否则抛 CapabilityDenied。"""
|
||||
if cap not in self._caps:
|
||||
raise CapabilityDenied(cap)
|
||||
return self._caps[cap]
|
||||
|
||||
def all(self) -> dict[Cap, CapabilityLimits]:
|
||||
return dict(self._caps)
|
||||
|
||||
def to_dict(self) -> dict[str, dict]:
|
||||
return {
|
||||
str(cap): {
|
||||
"rpm": lim.rpm,
|
||||
"batch": lim.batch,
|
||||
"subscribe": lim.subscribe,
|
||||
}
|
||||
for cap, lim in self._caps.items()
|
||||
}
|
||||
|
||||
|
||||
class CapabilityDenied(Exception):
|
||||
"""请求的 capability 当前不可用。"""
|
||||
|
||||
def __init__(self, cap: Cap, suggestion: str | None = None) -> None:
|
||||
self.cap = cap
|
||||
self.suggestion = suggestion or f"加购『{cap}』能力可解锁"
|
||||
super().__init__(f"capability not available: {cap}; {self.suggestion}")
|
||||
@@ -0,0 +1,109 @@
|
||||
"""数据源 SDK 封装(§5)。
|
||||
|
||||
进程内单例;Key 来源(优先级):secrets.json > .env。
|
||||
用户改 Key 后需要 `reset_clients()`,然后 `get_client()` 会拿新的。
|
||||
|
||||
5 档体系下服务器归属:
|
||||
- none 档(无 key / 无效 key) → 数据源 SDK .free()(free-api 服务器)
|
||||
- free 档(免费有效 key) → 数据源 SDK .free()(key 被 SDK 忽略,运行时走 free-api)
|
||||
- starter/pro/expert(付费 key) → 数据源 SDK (api_key=key, base_url)
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
from tickflow import AsyncTickFlow, TickFlow
|
||||
|
||||
from app import secrets_store
|
||||
|
||||
_sync_client: TickFlow | None = None
|
||||
_async_client: AsyncTickFlow | None = None
|
||||
|
||||
|
||||
# ===== 服务器归属判定 =====
|
||||
|
||||
# free-api 服务器默认节点(SDK 默认值),none/free 档运行时走这里。
|
||||
FREE_ENDPOINT = "https://free-api.tickflow.org"
|
||||
# 付费端点默认节点(starter+ 运行时走这里,也是端点切换的默认值)。
|
||||
PAID_ENDPOINT = "https://api.tickflow.org"
|
||||
|
||||
|
||||
def _should_use_free_server() -> bool:
|
||||
"""是否应走 free-api 服务器。
|
||||
|
||||
判定依据:无 key,或当前档位为 none/free。
|
||||
付费档(starter+)走付费端点。
|
||||
"""
|
||||
if not secrets_store.get_tickflow_key():
|
||||
return True
|
||||
# 有 key 时按探测出的档位判定(避免读 capabilities.json 在首次启动前未生成的边界)
|
||||
from app.tickflow.policy import base_tier_name
|
||||
return base_tier_name() in ("none", "free")
|
||||
|
||||
|
||||
def _base_url() -> str | None:
|
||||
"""从 secrets.json 读取用户自定义端点,没有则返回 None(用 SDK 默认)。"""
|
||||
return secrets_store.load().get("tickflow_base_url") or None
|
||||
|
||||
|
||||
def get_client() -> TickFlow:
|
||||
"""同步客户端。能力探测、盘后管道用。"""
|
||||
global _sync_client
|
||||
if _sync_client is None:
|
||||
key = secrets_store.get_tickflow_key()
|
||||
if _should_use_free_server():
|
||||
# none/free 档:走 free-api 服务器(无 key 或免费 key 被 SDK 忽略)
|
||||
_sync_client = TickFlow.free()
|
||||
else:
|
||||
_sync_client = TickFlow(api_key=key, base_url=_base_url())
|
||||
return _sync_client
|
||||
|
||||
|
||||
def get_async_client() -> AsyncTickFlow:
|
||||
"""异步客户端。FastAPI 请求路径上用。"""
|
||||
global _async_client
|
||||
if _async_client is None:
|
||||
key = secrets_store.get_tickflow_key()
|
||||
if _should_use_free_server():
|
||||
_async_client = AsyncTickFlow.free()
|
||||
else:
|
||||
_async_client = AsyncTickFlow(api_key=key, base_url=_base_url())
|
||||
return _async_client
|
||||
|
||||
|
||||
def reset_clients() -> None:
|
||||
"""Key 变化后调用 — 让下一次 get_client() 拿新实例。"""
|
||||
global _sync_client, _async_client
|
||||
_sync_client = None
|
||||
_async_client = None
|
||||
|
||||
|
||||
def current_mode() -> str:
|
||||
"""供 UI 显示当前模式。三态:
|
||||
|
||||
- "none" : 无 key / 无效 key(走 free-api,仅历史日K)
|
||||
- "free" : 免费有效 key(走 free-api,仅历史日K)
|
||||
- "api_key" : 付费 key(starter+,走付费端点,有实时行情)
|
||||
"""
|
||||
if not secrets_store.get_tickflow_key():
|
||||
return "none"
|
||||
from app.tickflow.policy import base_tier_name
|
||||
tier = base_tier_name()
|
||||
if tier in ("none", "free"):
|
||||
return "free" if tier == "free" else "none"
|
||||
return "api_key"
|
||||
|
||||
|
||||
def current_endpoint() -> str:
|
||||
"""返回当前显示用的端点 URL(对应 endpoints.json 列表项)。
|
||||
|
||||
- none/free 档:显示 free-api 服务器节点
|
||||
- 付费档:显示用户自定义端点(测速切换后)或默认付费节点 api.tickflow.org
|
||||
"""
|
||||
if _should_use_free_server():
|
||||
return FREE_ENDPOINT
|
||||
# 自定义端点(付费模式测速切换后):优先返回
|
||||
base = _base_url()
|
||||
if base:
|
||||
return base.rstrip("/")
|
||||
return PAID_ENDPOINT
|
||||
@@ -0,0 +1,550 @@
|
||||
"""能力探测 + CapabilitySet 持久化(§5.3)。
|
||||
|
||||
探测策略:逐 capability 用最小代价请求试探。
|
||||
- 成功 → 记录可用,优先取响应头 X-RateLimit-* 否则用 tiers.yaml 默认
|
||||
- 抛权限错 → 不可用
|
||||
- 抛其他错 → 不可用(谨慎,保留日志)
|
||||
|
||||
Tier Label 算法见 §5.3:基线档 + 补丁能力。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import yaml
|
||||
|
||||
from app.config import settings
|
||||
from app import secrets_store
|
||||
|
||||
from .capabilities import Cap, CapabilityLimits, CapabilitySet
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_CAPSET_CACHE_FILE = "capabilities.json"
|
||||
|
||||
# 缓存 schema 版本。capabilities 模型有结构性变更时 bump(如新增/拆分 Cap),
|
||||
# 旧缓存(无此字段或版本更低)会被判定过期,触发重新探测。
|
||||
# v2: 拆分 depth5 → depth5(单只) + depth5.batch(批量)
|
||||
# v3: 探测补全 quote.batch(此前 tiers.yaml 声明了但 _probe_real 漏探测)
|
||||
# v4: 5 档重构 —— 新增 none 档(无key/无效key),free 档重定义(走 free-api 服务器,
|
||||
# 仅历史日K)。判定改为复权因子分水岭:_classify_tier 接管档位判定。
|
||||
_CACHE_SCHEMA_VERSION = 4
|
||||
|
||||
# 探测用最小代价请求:挑流通性最好的 1 只标的试
|
||||
_PROBE_SYMBOL = "600000.SH" # 浦发银行,长期不会退市
|
||||
|
||||
|
||||
def _load_tiers_yaml() -> dict[str, dict[str, dict[str, Any]]]:
|
||||
for path in [settings.tiers_yaml, Path("/app/tiers.yaml"), Path("../tiers.yaml")]:
|
||||
if path.exists():
|
||||
with path.open(encoding="utf-8") as f:
|
||||
return yaml.safe_load(f)
|
||||
raise FileNotFoundError("tiers.yaml not found")
|
||||
|
||||
|
||||
def _tier_to_capset(tier_def: dict[str, dict[str, Any]]) -> CapabilitySet:
|
||||
caps: dict[Cap, CapabilityLimits] = {}
|
||||
for cap_name, limits_dict in tier_def.items():
|
||||
try:
|
||||
cap = Cap(cap_name)
|
||||
except ValueError:
|
||||
logger.warning("unknown cap in tiers.yaml: %s", cap_name)
|
||||
continue
|
||||
caps[cap] = CapabilityLimits(
|
||||
rpm=limits_dict.get("rpm"),
|
||||
batch=limits_dict.get("batch"),
|
||||
subscribe=limits_dict.get("subscribe"),
|
||||
)
|
||||
return CapabilitySet(caps)
|
||||
|
||||
|
||||
def _is_transient(e: Exception) -> bool:
|
||||
"""是否为"可重试的瞬时错误"——网络抖动 / 限流 / 服务端 5xx。
|
||||
|
||||
与权限/参数错误(403/401/400/404)区分:后者重试也无用,不重试。
|
||||
用类名匹配而非 import SDK 异常,避免探测期对 SDK 内部耦合。
|
||||
"""
|
||||
cls = e.__class__.__name__
|
||||
if cls in {
|
||||
"RateLimitError", "InternalServerError", "APIError",
|
||||
"ConnectionError", "TimeoutError", "ConnectError",
|
||||
"ConnectTimeout", "ReadTimeout", "RemoteProtocolError",
|
||||
"httpx.ConnectError", "httpx.TimeoutException",
|
||||
}:
|
||||
return True
|
||||
# APIError 体系下,status_code 5xx/429 视为瞬时
|
||||
status = getattr(e, "status_code", None)
|
||||
if isinstance(status, int) and (status == 429 or status >= 500):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _call_with_retry(fn, attempts: int = 3, backoff: float = 0.6) -> None:
|
||||
"""调用 fn();对瞬时错误退避重试,权限/参数错误立即抛出。
|
||||
|
||||
attempts=总尝试次数(含首次)。返回 None,异常由调用方分类。
|
||||
"""
|
||||
last_exc: Exception | None = None
|
||||
for i in range(attempts):
|
||||
try:
|
||||
fn()
|
||||
return
|
||||
except Exception as e: # noqa: BLE001
|
||||
last_exc = e
|
||||
# 权限/参数类错误:重试无意义,立即抛出交给 try_call 归类
|
||||
if not _is_transient(e):
|
||||
raise
|
||||
# 瞬时错误:最后一轮不再 sleep
|
||||
if i < attempts - 1:
|
||||
time.sleep(backoff * (i + 1))
|
||||
# 重试耗尽,抛出最后一次异常
|
||||
assert last_exc is not None
|
||||
raise last_exc
|
||||
|
||||
|
||||
def _probe_real(tiers: dict) -> tuple[CapabilitySet, list[str]]:
|
||||
"""逐 capability 试探。需要 API key。
|
||||
|
||||
**关键**:探测始终在付费端点上进行,用 key 鉴权验证有效性。
|
||||
绝不能读旧 capabilities 缓存的档位来选服务器 —— 否则首次保存 key 时,
|
||||
旧缓存是 none 档 → get_client() 返回 free 服务器 → free 服务器忽略 key →
|
||||
乱填 key 也能拿到日K → 误判成 free 档(鸡生蛋蛋生鸡的循环依赖 bug)。
|
||||
|
||||
返回 (capset, probe_log)。
|
||||
"""
|
||||
from tickflow import TickFlow
|
||||
from .client import _base_url, PAID_ENDPOINT
|
||||
|
||||
key = secrets_store.get_tickflow_key()
|
||||
# 探测专用客户端:强制走付费端点验证 key。
|
||||
# base_url 用用户自定义端点(若已配置测速切换),否则默认付费端点。
|
||||
probe_base = _base_url() or PAID_ENDPOINT
|
||||
tf = TickFlow(api_key=key, base_url=probe_base)
|
||||
available: dict[Cap, CapabilityLimits] = {}
|
||||
log: list[str] = []
|
||||
|
||||
def try_call(cap: Cap, fn, default_limits: dict[str, Any]) -> None:
|
||||
try:
|
||||
_call_with_retry(fn)
|
||||
available[cap] = CapabilityLimits(
|
||||
rpm=default_limits.get("rpm"),
|
||||
batch=default_limits.get("batch"),
|
||||
subscribe=default_limits.get("subscribe"),
|
||||
)
|
||||
log.append(f"✓ {cap}")
|
||||
except Exception as e: # noqa: BLE001
|
||||
msg = str(e).lower()
|
||||
cls = e.__class__.__name__
|
||||
# PermissionError 类名 / HTTP 403 / 中英文权限关键词都算"明确无权限"
|
||||
is_perm_denied = (
|
||||
cls in {"PermissionError", "AuthorizationError"}
|
||||
or "permission" in msg or "unauthorized" in msg
|
||||
or "403" in msg or "forbidden" in msg
|
||||
or "套餐" in msg or "权限" in msg or "需要" in msg
|
||||
)
|
||||
if is_perm_denied:
|
||||
log.append(f"✗ {cap}(无权限)")
|
||||
else:
|
||||
# 重试耗尽仍失败的瞬时错误 — 标记为疑似,而非直接判定"无此能力"
|
||||
log.append(f"? {cap} ({cls}: {e})")
|
||||
|
||||
# 用各档默认上限作为占位(无 X-RateLimit-* 头时)
|
||||
# 取所有档的并集,逐 cap 试探
|
||||
all_caps_defaults: dict[str, dict[str, Any]] = {}
|
||||
for tier in ("free", "starter", "pro", "expert"):
|
||||
for cap_name, lim in tiers.get(tier, {}).items():
|
||||
all_caps_defaults.setdefault(cap_name, lim)
|
||||
|
||||
def defaults(cap: Cap) -> dict[str, Any]:
|
||||
return all_caps_defaults.get(str(cap), {})
|
||||
|
||||
# 全部用 keyword-only 形式调用,符合 SDK 真实签名
|
||||
# quote.by_symbol
|
||||
try_call(Cap.QUOTE_BY_SYMBOL,
|
||||
lambda: tf.quotes.get(symbols=[_PROBE_SYMBOL], as_dataframe=False),
|
||||
defaults(Cap.QUOTE_BY_SYMBOL))
|
||||
|
||||
# quote.batch — 批量行情(POST /v1/quotes)。用 get_by_symbols 试探。
|
||||
try_call(Cap.QUOTE_BATCH,
|
||||
lambda: tf.quotes.get_by_symbols([_PROBE_SYMBOL], as_dataframe=False),
|
||||
defaults(Cap.QUOTE_BATCH))
|
||||
|
||||
# quote.pool — 用一个真实存在的 universe id 试探。
|
||||
# universes.list() 在 Free 也开放,先拿任意一个 universe id 再用 get_by_universes 试。
|
||||
def _probe_pool():
|
||||
unis = tf.universes.list()
|
||||
if not unis:
|
||||
raise RuntimeError("no universes available")
|
||||
first_id = unis[0]["id"] if isinstance(unis[0], dict) else getattr(unis[0], "id")
|
||||
return tf.quotes.get_by_universes([first_id], as_dataframe=False)
|
||||
|
||||
try_call(Cap.QUOTE_POOL, _probe_pool, defaults(Cap.QUOTE_POOL))
|
||||
|
||||
# kline.daily.by_symbol — Free 也有
|
||||
try_call(Cap.KLINE_DAILY_BY_SYMBOL,
|
||||
lambda: tf.klines.get(_PROBE_SYMBOL, period="1d", count=1, as_dataframe=False),
|
||||
defaults(Cap.KLINE_DAILY_BY_SYMBOL))
|
||||
|
||||
# kline.daily.batch
|
||||
try_call(Cap.KLINE_DAILY_BATCH,
|
||||
lambda: tf.klines.batch([_PROBE_SYMBOL], period="1d", count=1, as_dataframe=False),
|
||||
defaults(Cap.KLINE_DAILY_BATCH))
|
||||
|
||||
# kline.minute.by_symbol
|
||||
try_call(Cap.KLINE_MINUTE_BY_SYMBOL,
|
||||
lambda: tf.klines.get(_PROBE_SYMBOL, period="1m", count=1, as_dataframe=False),
|
||||
defaults(Cap.KLINE_MINUTE_BY_SYMBOL))
|
||||
|
||||
# kline.minute.batch
|
||||
try_call(Cap.KLINE_MINUTE_BATCH,
|
||||
lambda: tf.klines.batch([_PROBE_SYMBOL], period="1m", count=1, as_dataframe=False),
|
||||
defaults(Cap.KLINE_MINUTE_BATCH))
|
||||
|
||||
# intraday
|
||||
try_call(Cap.INTRADAY,
|
||||
lambda: tf.klines.intraday(_PROBE_SYMBOL, count=1, as_dataframe=False),
|
||||
defaults(Cap.INTRADAY))
|
||||
|
||||
# intraday.batch
|
||||
try_call(Cap.INTRADAY_BATCH,
|
||||
lambda: tf.klines.intraday_batch([_PROBE_SYMBOL], count=1, as_dataframe=False),
|
||||
defaults(Cap.INTRADAY_BATCH))
|
||||
|
||||
# depth5 — 按标的查(单只)
|
||||
try_call(Cap.DEPTH5,
|
||||
lambda: tf.depth.get(_PROBE_SYMBOL),
|
||||
defaults(Cap.DEPTH5))
|
||||
|
||||
# depth5.batch — 批量查(SDK 0.1.23+ 提供 depth.batch,对应官方 /v1/depth/batch 端点)
|
||||
try_call(Cap.DEPTH5_BATCH,
|
||||
lambda: tf.depth.batch([_PROBE_SYMBOL]),
|
||||
defaults(Cap.DEPTH5_BATCH))
|
||||
|
||||
# financial — SDK 提供 income / balance_sheet / cash_flow / metrics / shares
|
||||
# 用 metrics 探测(单据最小)
|
||||
try_call(Cap.FINANCIAL,
|
||||
lambda: tf.financials.metrics([_PROBE_SYMBOL], latest=True, as_dataframe=False),
|
||||
defaults(Cap.FINANCIAL))
|
||||
|
||||
# adj_factor — 实际在 klines.ex_factors
|
||||
try_call(Cap.ADJ_FACTOR,
|
||||
lambda: tf.klines.ex_factors([_PROBE_SYMBOL], as_dataframe=False),
|
||||
defaults(Cap.ADJ_FACTOR))
|
||||
|
||||
# websocket 不在探测期试连接(成本太高且阻塞),按档位默认推断
|
||||
# 若 expert 的其他 cap 都通,则推断 websocket 也可用
|
||||
if (Cap.FINANCIAL in available and Cap.INTRADAY_BATCH in available):
|
||||
available[Cap.WEBSOCKET] = CapabilityLimits(
|
||||
subscribe=defaults(Cap.WEBSOCKET).get("subscribe", 100),
|
||||
)
|
||||
log.append("✓ websocket (inferred from expert tier)")
|
||||
|
||||
return CapabilitySet(available), log
|
||||
|
||||
|
||||
def detect_capabilities(force: bool = False) -> CapabilitySet:
|
||||
"""探测当前 API Key 的能力集。"""
|
||||
cache_path = settings.data_dir / _CAPSET_CACHE_FILE
|
||||
if not force and cache_path.exists():
|
||||
with cache_path.open(encoding="utf-8") as f:
|
||||
cached = json.load(f)
|
||||
# schema 版本校验:旧缓存或缺版本号 → 过期,丢弃后重新探测
|
||||
if cached.get("schema_version") == _CACHE_SCHEMA_VERSION:
|
||||
return _capset_from_json(cached)
|
||||
logger.info("capabilities 缓存 schema 版本过期(缓存=%s, 当前=%d), 重新探测",
|
||||
cached.get("schema_version"), _CACHE_SCHEMA_VERSION)
|
||||
|
||||
tiers = _load_tiers_yaml()
|
||||
if settings.use_free_mode:
|
||||
# 无 key —— 归 none 档(走 free-api 服务器,仅历史日K)
|
||||
capset = _tier_to_capset(tiers["none"])
|
||||
_persist(capset, "None", log=["无 API Key(无档 · free-api 服务器)"], missing=[], extras=[])
|
||||
return capset
|
||||
|
||||
# 有 API key — 真实探测
|
||||
try:
|
||||
capset, probe_log = _probe_real(tiers)
|
||||
# 判定档位:无效 key → none,免费 key → free,付费 → starter/pro/expert
|
||||
classified = _classify_tier(capset, tiers)
|
||||
if classified.is_invalid:
|
||||
# 无效 key(连单只日K都拿不到):归 none 档,标记要求清除 key
|
||||
capset = _tier_to_capset(tiers["none"])
|
||||
probe_log.append("⚠ Key 无效(单只日K也无法获取),判定为无档")
|
||||
_persist(capset, "None", log=probe_log, missing=[], extras=[], invalid_key=True)
|
||||
return capset
|
||||
if classified.is_free:
|
||||
# 免费有效 key:能力按 free 档(= none 档能力,走 free-api 服务器)
|
||||
capset = _tier_to_capset(tiers["free"])
|
||||
_persist(capset, "Free", log=probe_log + ["✓ 免费有效 key(运行时走 free-api 服务器)"], missing=[], extras=[])
|
||||
return capset
|
||||
# 付费档(starter+) — 探测出的能力即为真实可用
|
||||
label, missing, extras = _compute_label_and_missing(capset, tiers)
|
||||
capset = _override_limits_with_detected_tier(capset, label, tiers)
|
||||
_persist(capset, label, log=probe_log, missing=missing, extras=extras)
|
||||
return capset
|
||||
except Exception as e:
|
||||
logger.exception("detect_capabilities failed; using none baseline: %s", e)
|
||||
capset = _tier_to_capset(tiers["none"])
|
||||
_persist(capset, "None(探测失败)", log=[f"探测失败:{e}"], missing=[], extras=[])
|
||||
return capset
|
||||
|
||||
|
||||
# ===== Tier 代表性 capability(signature caps)=====
|
||||
# 拥有**任意一个**即认作该档及以上。自上而下匹配。
|
||||
# 这套设计的好处:单个 capability 探测的 transient 失败不会把整体档位"误降"。
|
||||
TIER_SIGNATURES: dict[str, set[Cap]] = {
|
||||
"expert": {Cap.FINANCIAL, Cap.INTRADAY_BATCH, Cap.WEBSOCKET},
|
||||
"pro": {Cap.KLINE_MINUTE_BATCH, Cap.KLINE_MINUTE_BY_SYMBOL,
|
||||
Cap.INTRADAY, Cap.DEPTH5, Cap.DEPTH5_BATCH},
|
||||
"starter": {Cap.QUOTE_BATCH, Cap.KLINE_DAILY_BATCH,
|
||||
Cap.ADJ_FACTOR, Cap.QUOTE_POOL},
|
||||
# free / none 不需 signature — 由 _classify_tier 的分水岭逻辑判定
|
||||
}
|
||||
|
||||
|
||||
@dataclass(slots=True, frozen=True)
|
||||
class TierClassification:
|
||||
"""档位判定结果。
|
||||
|
||||
判定依据是"复权因子分水岭":
|
||||
- 连单只日K都没有 → 无效 key(is_invalid),归 none 档
|
||||
- 有单只日K、无复权因子 → 免费 key(is_free)
|
||||
- 有复权因子 → 付费档(starter+),具体档位由 signature 决定
|
||||
"""
|
||||
|
||||
tier: str # "none" / "free" / "starter" / "pro" / "expert"
|
||||
is_invalid: bool # 无效 key(连单只日K都拿不到)
|
||||
is_free: bool # 免费有效 key(有日K、无复权因子)
|
||||
|
||||
|
||||
def _classify_tier(capset: CapabilitySet, tiers: dict) -> TierClassification:
|
||||
"""根据探测出的能力集判定档位。
|
||||
|
||||
分水岭是 KLINE_DAILY_BY_SYMBOL(单只日K)与 ADJ_FACTOR(复权因子):
|
||||
- 无单只日K → none(无效 key)
|
||||
- 有日K无复权 → free(免费 key)
|
||||
- 有复权因子 → 走 signature 判定 starter/pro/expert
|
||||
"""
|
||||
held = set(capset.all().keys())
|
||||
|
||||
# 1) 连单只日K都没有 → 无效 key
|
||||
if Cap.KLINE_DAILY_BY_SYMBOL not in held:
|
||||
return TierClassification(tier="none", is_invalid=True, is_free=False)
|
||||
|
||||
# 2) 有日K但无复权因子 → 免费 key
|
||||
if Cap.ADJ_FACTOR not in held:
|
||||
return TierClassification(tier="free", is_invalid=False, is_free=True)
|
||||
|
||||
# 3) 有复权因子 → 付费档,按 signature 自上而下判定
|
||||
if held & TIER_SIGNATURES["expert"]:
|
||||
base = "expert"
|
||||
elif held & TIER_SIGNATURES["pro"]:
|
||||
base = "pro"
|
||||
elif held & TIER_SIGNATURES["starter"]:
|
||||
base = "starter"
|
||||
else:
|
||||
# 有复权因子但无任何代表能力 — 兜底为 starter(复权本身是 starter 特征)
|
||||
base = "starter"
|
||||
return TierClassification(tier=base, is_invalid=False, is_free=False)
|
||||
|
||||
# 补丁友好命名(label 后缀用)
|
||||
_CAP_ALIASES: dict[Cap, str] = {
|
||||
Cap.KLINE_MINUTE_BATCH: "分钟K",
|
||||
Cap.KLINE_MINUTE_BY_SYMBOL: "分钟K",
|
||||
Cap.INTRADAY: "分时",
|
||||
Cap.INTRADAY_BATCH: "批量分时",
|
||||
Cap.DEPTH5: "五档",
|
||||
Cap.DEPTH5_BATCH: "批量五档",
|
||||
Cap.WEBSOCKET: "WS",
|
||||
Cap.FINANCIAL: "财务",
|
||||
Cap.ADJ_FACTOR: "复权",
|
||||
Cap.QUOTE_BATCH: "批量行情",
|
||||
Cap.QUOTE_POOL: "标的池",
|
||||
Cap.KLINE_DAILY_BATCH: "日K批量",
|
||||
}
|
||||
|
||||
|
||||
def _override_limits_with_detected_tier(
|
||||
capset: CapabilitySet, label: str, tiers: dict,
|
||||
) -> CapabilitySet:
|
||||
"""探测完成后,用判档对应的 limits 覆盖每个 cap 的速率/批量。
|
||||
|
||||
判档前每个 cap 用的是"所有档默认值的并集"(为了不漏数据),
|
||||
判档后才知道用户真实档位,limits 用该档的实际值更准。
|
||||
label 可能是 "Pro" / "Pro + 分钟K" / "Pro+" 等组合形式 — 取第一个词当作基线档名。
|
||||
"""
|
||||
base_name = label.split()[0].split("+")[0].strip().lower() # "Pro + 分钟K" → "pro"
|
||||
tier_limits = tiers.get(base_name, {})
|
||||
new_caps: dict[Cap, CapabilityLimits] = {}
|
||||
for cap, _old_lim in capset.all().items():
|
||||
spec = tier_limits.get(cap.value)
|
||||
if spec:
|
||||
new_caps[cap] = CapabilityLimits(
|
||||
rpm=spec.get("rpm"),
|
||||
batch=spec.get("batch"),
|
||||
subscribe=spec.get("subscribe"),
|
||||
)
|
||||
else:
|
||||
# 不在该档定义里(extras),用 expert 档兜底(最宽松)
|
||||
expert_spec = tiers.get("expert", {}).get(cap.value, {})
|
||||
new_caps[cap] = CapabilityLimits(
|
||||
rpm=expert_spec.get("rpm"),
|
||||
batch=expert_spec.get("batch"),
|
||||
subscribe=expert_spec.get("subscribe"),
|
||||
)
|
||||
return CapabilitySet(new_caps)
|
||||
|
||||
|
||||
def _tier_caps_set(tiers: dict, tier_name: str) -> set[Cap]:
|
||||
"""读 tiers.yaml 的某档定义,转为 Cap 集合。"""
|
||||
return {Cap(c) for c in tiers.get(tier_name, {}).keys() if c in {x.value for x in Cap}}
|
||||
|
||||
|
||||
def _compute_label_and_missing(
|
||||
capset: CapabilitySet, tiers: dict,
|
||||
) -> tuple[str, list[str], list[str]]:
|
||||
"""返回 (label, missing_caps, extra_caps)。
|
||||
|
||||
label:档位标签。
|
||||
missing_caps:本档**应有但未探测到**的 capability(用于诊断:可能是探测 bug 或权限丢失)。
|
||||
extra_caps:超出本档的额外 capability(自定义组合)。
|
||||
"""
|
||||
held = set(capset.all().keys())
|
||||
|
||||
# 1) 完全匹配 — 干净命中某档
|
||||
for tier_name in ["free", "starter", "pro", "expert"]:
|
||||
if held == _tier_caps_set(tiers, tier_name):
|
||||
return tier_name.capitalize(), [], []
|
||||
|
||||
# 2) 按 signature 自上而下判档
|
||||
if held & TIER_SIGNATURES["expert"]:
|
||||
base = "expert"
|
||||
elif held & TIER_SIGNATURES["pro"]:
|
||||
base = "pro"
|
||||
elif held & TIER_SIGNATURES["starter"]:
|
||||
base = "starter"
|
||||
else:
|
||||
base = "free"
|
||||
|
||||
base_caps = _tier_caps_set(tiers, base)
|
||||
missing = sorted(c.value for c in (base_caps - held))
|
||||
extras = base_caps and (held - base_caps) or set() # extras 是超出该档的部分
|
||||
|
||||
# 实际超出 = held 中"既不属于本档、也不属于本档下方任何档"的 cap
|
||||
# 简化:extras = held - base_caps
|
||||
extras_set = held - base_caps
|
||||
|
||||
# 3) 拼 label
|
||||
if not extras_set:
|
||||
# 完全在本档内(可能缺一两项 — 由 missing 反映)
|
||||
return base.capitalize(), missing, []
|
||||
|
||||
# 补丁过多 → 用 "≈" 形式
|
||||
if len(extras_set) > 3:
|
||||
return f"{base.capitalize()}+", missing, sorted(c.value for c in extras_set)
|
||||
|
||||
suffix = sorted({_CAP_ALIASES.get(e, str(e)) for e in extras_set})
|
||||
return f"{base.capitalize()} + " + " + ".join(suffix), missing, sorted(c.value for c in extras_set)
|
||||
|
||||
|
||||
def _compute_label(capset: CapabilitySet, tiers: dict) -> str:
|
||||
"""对外简化签名 — 只要 label。"""
|
||||
label, _missing, _extras = _compute_label_and_missing(capset, tiers)
|
||||
return label
|
||||
|
||||
|
||||
def _persist(
|
||||
capset: CapabilitySet,
|
||||
label: str,
|
||||
log: list[str] | None = None,
|
||||
missing: list[str] | None = None,
|
||||
extras: list[str] | None = None,
|
||||
invalid_key: bool = False,
|
||||
) -> None:
|
||||
settings.data_dir.mkdir(parents=True, exist_ok=True)
|
||||
cache_path = settings.data_dir / _CAPSET_CACHE_FILE
|
||||
payload = {
|
||||
"schema_version": _CACHE_SCHEMA_VERSION,
|
||||
"label": label,
|
||||
"capabilities": capset.to_dict(),
|
||||
"probe_log": log or [],
|
||||
"missing_caps": missing or [], # 本档应有但未探测到
|
||||
"extras_caps": extras or [], # 超出本档的额外能力
|
||||
"invalid_key": invalid_key, # 探测出的 key 无效(连单只日K都拿不到)
|
||||
}
|
||||
with cache_path.open("w", encoding="utf-8") as f:
|
||||
json.dump(payload, f, ensure_ascii=False, indent=2)
|
||||
|
||||
|
||||
def _capset_from_json(data: dict[str, Any]) -> CapabilitySet:
|
||||
caps: dict[Cap, CapabilityLimits] = {}
|
||||
for cap_name, lim in data.get("capabilities", {}).items():
|
||||
try:
|
||||
cap = Cap(cap_name)
|
||||
except ValueError:
|
||||
continue
|
||||
caps[cap] = CapabilityLimits(
|
||||
rpm=lim.get("rpm"),
|
||||
batch=lim.get("batch"),
|
||||
subscribe=lim.get("subscribe"),
|
||||
)
|
||||
return CapabilitySet(caps)
|
||||
|
||||
|
||||
def tier_label() -> str:
|
||||
cache_path = settings.data_dir / _CAPSET_CACHE_FILE
|
||||
if cache_path.exists():
|
||||
with cache_path.open(encoding="utf-8") as f:
|
||||
return json.load(f).get("label", "Unknown")
|
||||
return "Unknown"
|
||||
|
||||
|
||||
def probe_log() -> list[str]:
|
||||
cache_path = settings.data_dir / _CAPSET_CACHE_FILE
|
||||
if cache_path.exists():
|
||||
with cache_path.open(encoding="utf-8") as f:
|
||||
return json.load(f).get("probe_log", [])
|
||||
return []
|
||||
|
||||
|
||||
def missing_caps() -> list[str]:
|
||||
"""本档应有但未探测到的 capability — 通常意味着探测有 bug 或权限边界。"""
|
||||
cache_path = settings.data_dir / _CAPSET_CACHE_FILE
|
||||
if cache_path.exists():
|
||||
with cache_path.open(encoding="utf-8") as f:
|
||||
return json.load(f).get("missing_caps", [])
|
||||
return []
|
||||
|
||||
|
||||
def extras_caps() -> list[str]:
|
||||
cache_path = settings.data_dir / _CAPSET_CACHE_FILE
|
||||
if cache_path.exists():
|
||||
with cache_path.open(encoding="utf-8") as f:
|
||||
return json.load(f).get("extras_caps", [])
|
||||
return []
|
||||
|
||||
|
||||
def is_invalid_key() -> bool:
|
||||
"""最近一次探测是否判定 key 无效(连单只日K都拿不到)。
|
||||
|
||||
settings 层据此清除已存的 key,避免乱填的 key 被持久化。
|
||||
"""
|
||||
cache_path = settings.data_dir / _CAPSET_CACHE_FILE
|
||||
if cache_path.exists():
|
||||
with cache_path.open(encoding="utf-8") as f:
|
||||
return bool(json.load(f).get("invalid_key", False))
|
||||
return False
|
||||
|
||||
|
||||
def base_tier_name() -> str:
|
||||
"""当前档位的基础名(小写): none / free / starter / pro / expert。
|
||||
|
||||
供 client 层判断"是否走 free-api 服务器"(none/free → free 服务器)。
|
||||
"""
|
||||
label = tier_label()
|
||||
return label.split()[0].split("+")[0].strip().lower()
|
||||
@@ -0,0 +1,158 @@
|
||||
"""标的池(Universe)定义(§6.3)。
|
||||
|
||||
Phase 1 实现:
|
||||
- 常用指数成份(沪深 300 / 中证 500 / 上证 50)用数据源 `quote.pool` 端点拉取并缓存
|
||||
- 全 A 通过 instruments.batch 获取
|
||||
- 自选池 = 用户的 watchlist
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from datetime import date
|
||||
from pathlib import Path
|
||||
from typing import Literal
|
||||
|
||||
import polars as pl
|
||||
|
||||
from app.config import settings
|
||||
from app.tickflow.client import get_client
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
PoolId = Literal["CSI300", "CSI500", "SSE50", "CN_Equity_A", "CN_Index", "watchlist"]
|
||||
|
||||
# 数据源 universe id 是内部命名(见 tf.universes.list())。
|
||||
# 没有官方对照表,启动时按名称模糊匹配从 universes.list() 里找。
|
||||
# 常见名:沪深300 / 中证500 / 上证50 / 全 A
|
||||
_POOL_NAME_HINTS = {
|
||||
"CSI300": ["沪深300", "HS300", "CSI300"],
|
||||
"CSI500": ["中证500", "ZZ500", "CSI500"],
|
||||
"SSE50": ["上证50", "SH50", "SSE50"],
|
||||
}
|
||||
|
||||
|
||||
def _find_universe_id(hints: list[str]) -> str | None:
|
||||
"""从 universes.list() 里按 name/id 子串匹配找一个 universe id。"""
|
||||
try:
|
||||
tf = get_client()
|
||||
unis = tf.universes.list()
|
||||
except Exception as e: # noqa: BLE001
|
||||
logger.warning("universes.list failed: %s", e)
|
||||
return None
|
||||
for u in unis or []:
|
||||
item = u if isinstance(u, dict) else {"id": getattr(u, "id", ""), "name": getattr(u, "name", "")}
|
||||
haystack = (item.get("id", "") + " " + item.get("name", "")).lower()
|
||||
for h in hints:
|
||||
if h.lower() in haystack:
|
||||
return item["id"]
|
||||
return None
|
||||
|
||||
|
||||
def _pool_cache_path(pool_id: str) -> Path:
|
||||
return settings.data_dir / "pools" / f"{pool_id}.parquet"
|
||||
|
||||
|
||||
def get_pool(pool_id: PoolId, refresh: bool = False) -> list[str]:
|
||||
"""返回标的池里的 symbol 列表。"""
|
||||
if pool_id == "watchlist":
|
||||
return _load_watchlist()
|
||||
|
||||
cache = _pool_cache_path(pool_id)
|
||||
if cache.exists() and not refresh:
|
||||
df = pl.read_parquet(cache)
|
||||
return df["symbol"].to_list()
|
||||
|
||||
symbols = _fetch_pool(pool_id)
|
||||
if symbols:
|
||||
cache.parent.mkdir(parents=True, exist_ok=True)
|
||||
pl.DataFrame({"symbol": symbols, "as_of": [date.today()] * len(symbols)}).write_parquet(cache)
|
||||
return symbols
|
||||
|
||||
|
||||
def _fetch_pool(pool_id: PoolId) -> list[str]:
|
||||
"""从数据源拉取池成份。
|
||||
|
||||
实现:先用 universes.list 找到 universe id,再 quotes.get_by_universes 拉成份。
|
||||
"""
|
||||
tf = get_client()
|
||||
|
||||
if pool_id in _POOL_NAME_HINTS:
|
||||
uid = _find_universe_id(_POOL_NAME_HINTS[pool_id])
|
||||
if not uid:
|
||||
logger.warning("无法在数据源 universes 列表里匹配到 %s", pool_id)
|
||||
return []
|
||||
try:
|
||||
df = tf.quotes.get_by_universes([uid], as_dataframe=True)
|
||||
if df is not None and len(df) > 0 and "symbol" in df.columns:
|
||||
return df["symbol"].astype(str).tolist()
|
||||
except Exception as e: # noqa: BLE001
|
||||
logger.warning("fetch pool %s via universe %s failed: %s", pool_id, uid, e)
|
||||
|
||||
if pool_id == "CN_Equity_A":
|
||||
# 全 A — 优先直接用 CN_Equity_A universe (包含沪深京三市)
|
||||
uid = _find_universe_id(["CN_Equity_A", "沪深京A股", "全A"])
|
||||
if uid:
|
||||
try:
|
||||
df = tf.quotes.get_by_universes([uid], as_dataframe=True)
|
||||
if df is not None and len(df) > 0 and "symbol" in df.columns:
|
||||
return sorted(set(df["symbol"].astype(str).tolist()))
|
||||
except Exception as e: # noqa: BLE001
|
||||
logger.warning("fetch CN_Equity_A via universe %s failed: %s", uid, e)
|
||||
|
||||
# fallback: 聚合申万一级行业 (覆盖度较低, 缺北交所/新股)
|
||||
try:
|
||||
unis = tf.universes.list()
|
||||
except Exception as e: # noqa: BLE001
|
||||
logger.warning("universes.list failed: %s", e)
|
||||
unis = []
|
||||
sw1_ids = []
|
||||
for u in unis or []:
|
||||
item = u if isinstance(u, dict) else {"id": getattr(u, "id", "")}
|
||||
uid = item.get("id", "")
|
||||
if "SW1_" in uid:
|
||||
sw1_ids.append(uid)
|
||||
if sw1_ids:
|
||||
try:
|
||||
df = tf.quotes.get_by_universes(sw1_ids, as_dataframe=True)
|
||||
if df is not None and "symbol" in df.columns:
|
||||
return sorted(set(df["symbol"].astype(str).tolist()))
|
||||
except Exception as e: # noqa: BLE001
|
||||
logger.warning("aggregate SW1 fetch failed: %s", e)
|
||||
|
||||
if pool_id == "CN_Index":
|
||||
uid = _find_universe_id(["CN_Index", "沪深指数", "指数"])
|
||||
ids = [uid] if uid else ["CN_Index"]
|
||||
try:
|
||||
df = tf.quotes.get_by_universes(ids, as_dataframe=True)
|
||||
if df is not None and len(df) > 0 and "symbol" in df.columns:
|
||||
return sorted(set(df["symbol"].astype(str).tolist()))
|
||||
except Exception as e: # noqa: BLE001
|
||||
logger.warning("fetch CN_Index via universe %s failed: %s", ids, e)
|
||||
|
||||
return []
|
||||
|
||||
|
||||
def _load_watchlist() -> list[str]:
|
||||
"""读取用户自选(由 watchlist service 维护)。"""
|
||||
path = settings.data_dir / "user_data" / "watchlist.parquet"
|
||||
if not path.exists():
|
||||
return []
|
||||
df = pl.read_parquet(path)
|
||||
if df.is_empty() or "symbol" not in df.columns:
|
||||
return []
|
||||
return df["symbol"].to_list()
|
||||
|
||||
|
||||
# 兜底:Free 用户/无 API 时给一个小型可用集合,让 UI 不至于空白
|
||||
DEMO_SYMBOLS = [
|
||||
"600000.SH", # 浦发银行
|
||||
"600036.SH", # 招商银行
|
||||
"600519.SH", # 贵州茅台
|
||||
"601318.SH", # 中国平安
|
||||
"601398.SH", # 工商银行
|
||||
"000001.SZ", # 平安银行
|
||||
"000333.SZ", # 美的集团
|
||||
"000651.SZ", # 格力电器
|
||||
"000858.SZ", # 五粮液
|
||||
"002594.SZ", # 比亚迪
|
||||
]
|
||||
@@ -0,0 +1,988 @@
|
||||
"""Repository 层(§7.4)。
|
||||
|
||||
数据分层:
|
||||
- DuckDB 视图: 冷查询(统计、元数据、用户自定义SQL)
|
||||
- Polars 缓存: 热路径(enriched 最新日 ~5500行 + instruments ~5500行)
|
||||
- Polars scan_parquet: 分钟K/历史日K (predicate pushdown)
|
||||
|
||||
缓存生命周期:
|
||||
- startup 时不加载(数据可能为空)
|
||||
- pipeline 完成后调用 refresh_cache()
|
||||
- 服务层通过 get_enriched_latest() / get_instruments() 获取缓存
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import threading
|
||||
from datetime import date
|
||||
from pathlib import Path
|
||||
|
||||
import duckdb
|
||||
import polars as pl
|
||||
|
||||
from app.config import settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class DataStore:
|
||||
"""唯一的存储入口 — 进程启动时创建。"""
|
||||
|
||||
def __init__(self, data_dir: Path | None = None) -> None:
|
||||
self.data_dir = Path(data_dir or settings.data_dir)
|
||||
self.data_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# 关键子目录(§7.2)
|
||||
for sub in (
|
||||
"kline_daily",
|
||||
"kline_daily_enriched",
|
||||
"kline_index_daily",
|
||||
"kline_index_enriched",
|
||||
"kline_minute",
|
||||
"adj_factor",
|
||||
"financials",
|
||||
"instruments",
|
||||
"instruments_index",
|
||||
"instruments_ext",
|
||||
"kline_ext",
|
||||
"pools",
|
||||
"backtest_results",
|
||||
"screener_results",
|
||||
"ai_cache",
|
||||
"user_data",
|
||||
"depth5",
|
||||
):
|
||||
(self.data_dir / sub).mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# 财务数据子目录
|
||||
for sub in ("metrics", "income", "balance_sheet", "cash_flow"):
|
||||
(self.data_dir / "financials" / sub).mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# DuckDB 内存模式 — 不建 .db 文件(§7.1)
|
||||
self.db = duckdb.connect(database=":memory:")
|
||||
self._register_views()
|
||||
|
||||
def _register_views(self) -> None:
|
||||
"""把 Parquet 目录挂载为 DuckDB 视图(§7.3)。"""
|
||||
d = self.data_dir.as_posix()
|
||||
statements = [
|
||||
f"""CREATE OR REPLACE VIEW kline_daily AS
|
||||
SELECT * FROM read_parquet('{d}/kline_daily/**/*.parquet', union_by_name=true)""",
|
||||
f"""CREATE OR REPLACE VIEW kline_enriched AS
|
||||
SELECT * FROM read_parquet('{d}/kline_daily_enriched/**/*.parquet', union_by_name=true)""",
|
||||
f"""CREATE OR REPLACE VIEW kline_index_daily AS
|
||||
SELECT * FROM read_parquet('{d}/kline_index_daily/**/*.parquet', union_by_name=true)""",
|
||||
f"""CREATE OR REPLACE VIEW kline_index_enriched AS
|
||||
SELECT * FROM read_parquet('{d}/kline_index_enriched/**/*.parquet', union_by_name=true)""",
|
||||
f"""CREATE OR REPLACE VIEW kline_minute AS
|
||||
SELECT * FROM read_parquet('{d}/kline_minute/**/*.parquet', union_by_name=true)""",
|
||||
f"""CREATE OR REPLACE VIEW adj_factor AS
|
||||
SELECT * FROM read_parquet('{d}/adj_factor/**/*.parquet', union_by_name=true)""",
|
||||
f"""CREATE OR REPLACE VIEW instruments AS
|
||||
SELECT * FROM read_parquet('{d}/instruments/**/*.parquet', union_by_name=true)""",
|
||||
f"""CREATE OR REPLACE VIEW instruments_index AS
|
||||
SELECT * FROM read_parquet('{d}/instruments_index/**/*.parquet', union_by_name=true)""",
|
||||
f"""CREATE OR REPLACE VIEW instruments_ext AS
|
||||
SELECT * FROM read_parquet('{d}/instruments_ext/**/*.parquet', union_by_name=true)""",
|
||||
f"""CREATE OR REPLACE VIEW kline_ext AS
|
||||
SELECT * FROM read_parquet('{d}/kline_ext/**/*.parquet', union_by_name=true)""",
|
||||
# 财务数据视图
|
||||
f"""CREATE OR REPLACE VIEW financials_metrics AS
|
||||
SELECT * FROM read_parquet('{d}/financials/metrics/*.parquet', union_by_name=true)""",
|
||||
f"""CREATE OR REPLACE VIEW financials_income AS
|
||||
SELECT * FROM read_parquet('{d}/financials/income/*.parquet', union_by_name=true)""",
|
||||
f"""CREATE OR REPLACE VIEW financials_balance_sheet AS
|
||||
SELECT * FROM read_parquet('{d}/financials/balance_sheet/*.parquet', union_by_name=true)""",
|
||||
f"""CREATE OR REPLACE VIEW financials_cash_flow AS
|
||||
SELECT * FROM read_parquet('{d}/financials/cash_flow/*.parquet', union_by_name=true)""",
|
||||
# 五档盘口 sealed 真假涨停(独立旁路存储,不进 enriched)
|
||||
f"""CREATE OR REPLACE VIEW depth5 AS
|
||||
SELECT * FROM read_parquet('{d}/depth5/**/*.parquet', union_by_name=true)""",
|
||||
]
|
||||
for sql in statements:
|
||||
try:
|
||||
self.db.execute(sql)
|
||||
except duckdb.IOException:
|
||||
logger.debug("view registration skipped (no parquet yet): %s", sql[:60])
|
||||
|
||||
|
||||
class KlineRepository:
|
||||
"""日 K / 分钟 K 的读写入口。"""
|
||||
|
||||
def __init__(self, store: DataStore) -> None:
|
||||
self.store = store
|
||||
self.db = store.db
|
||||
self._lock = threading.Lock()
|
||||
|
||||
# ---- Polars 缓存 ----
|
||||
self._enriched_cache: pl.DataFrame | None = None # 最新一天 (~5500行)
|
||||
self._enriched_cache_date: date | None = None
|
||||
self._live_agg_cache: pl.DataFrame | None = None # 预计算聚合表 (~5500行)
|
||||
self._live_agg_cache_date: date | None = None
|
||||
self._instruments_cache: pl.DataFrame | None = None
|
||||
# 完整 enriched 历史 (含所有指标, 供 filter_history 策略使用)
|
||||
self._enriched_history_cache: pl.DataFrame | None = None # ~100万行
|
||||
self._enriched_history_start: date | None = None
|
||||
self._index_instruments_cache: pl.DataFrame | None = None
|
||||
|
||||
# parquet glob 路径
|
||||
self._enriched_glob = str(store.data_dir / "kline_daily_enriched" / "**" / "*.parquet")
|
||||
self._index_enriched_glob = str(store.data_dir / "kline_index_enriched" / "**" / "*.parquet")
|
||||
self._minute_glob = str(store.data_dir / "kline_minute" / "**" / "*.parquet")
|
||||
self._inst_glob = str(store.data_dir / "instruments" / "**" / "*.parquet")
|
||||
self._index_inst_glob = str(store.data_dir / "instruments_index" / "**" / "*.parquet")
|
||||
|
||||
def execute_all(self, sql: str, params: list | None = None) -> list[tuple]:
|
||||
"""线程安全的 SELECT → fetchall。DuckDB 单 connection 非线程安全,所有读路径须走此方法。"""
|
||||
with self._lock:
|
||||
return self.db.execute(sql, params or []).fetchall()
|
||||
|
||||
def execute_one(self, sql: str, params: list | None = None) -> tuple | None:
|
||||
"""线程安全的 SELECT → fetchone。"""
|
||||
with self._lock:
|
||||
return self.db.execute(sql, params or []).fetchone()
|
||||
|
||||
# ================================================================
|
||||
# Polars 缓存管理
|
||||
# ================================================================
|
||||
|
||||
def refresh_cache(self) -> None:
|
||||
"""刷新 Polars 缓存。在 pipeline 完成后、服务启动时调用。"""
|
||||
self._refresh_instruments()
|
||||
self._refresh_index_instruments()
|
||||
self._refresh_enriched()
|
||||
|
||||
def clear_cache(self) -> None:
|
||||
"""清空所有 Polars 内存缓存。
|
||||
|
||||
与 refresh_cache 的区别: refresh_cache 在磁盘无数据时会提前 return,
|
||||
导致内存里的旧缓存残留 (clear 数据后看板仍显示旧数据的根因)。
|
||||
本方法无条件清空, 供清除数据/重置场景调用。
|
||||
"""
|
||||
self._enriched_cache = None
|
||||
self._enriched_cache_date = None
|
||||
self._enriched_history_cache = None
|
||||
self._enriched_history_start = None
|
||||
self._live_agg_cache = None
|
||||
self._live_agg_cache_date = None
|
||||
self._instruments_cache = None
|
||||
self._index_instruments_cache = None
|
||||
|
||||
def _refresh_enriched(self) -> None:
|
||||
"""从 parquet 加载 enriched 最新日到内存 + 构建聚合表。
|
||||
|
||||
enriched parquet 仅存 14 列基础数据。启动时读入历史数据并即时计算完整指标,
|
||||
将结果缓存在内存中供各服务使用。
|
||||
|
||||
优化: 扩大历史读取范围, 同时缓存完整历史 (含指标), 供 filter_history 策略直接复用。
|
||||
"""
|
||||
try:
|
||||
latest = self._latest_enriched_date_duckdb()
|
||||
if not latest:
|
||||
# 磁盘已无数据: 必须清空内存缓存, 否则旧数据会残留
|
||||
# (清数据后看板仍显示旧数据的根因)
|
||||
self.clear_cache()
|
||||
return
|
||||
|
||||
# Step 1: 直接读最新日期的分区文件 (仅 14 列)
|
||||
enriched_dir = self.store.data_dir / "kline_daily_enriched"
|
||||
ds = latest.isoformat() if hasattr(latest, "isoformat") else str(latest)
|
||||
target_parquet = enriched_dir / f"date={ds}" / "part.parquet"
|
||||
|
||||
if not target_parquet.exists():
|
||||
return
|
||||
|
||||
df_latest = pl.read_parquet(target_parquet)
|
||||
if df_latest.is_empty():
|
||||
return
|
||||
|
||||
# Step 2: 读近 300 天 14 列数据 → compute → filter(latest) → 缓存
|
||||
# 300 日历天 ≈ 210 交易日, 覆盖 filter_history 最大 lookback(90) + warmup(60)
|
||||
try:
|
||||
from datetime import timedelta
|
||||
from app.indicators.pipeline import compute_indicators, compute_signals, compute_limit_signals
|
||||
start_full = latest - timedelta(days=300)
|
||||
read_cols = [c for c in ["symbol", "date", "open", "high", "low", "close",
|
||||
"volume", "amount", "raw_close", "raw_high", "raw_low"]
|
||||
if c in df_latest.columns]
|
||||
lf = (
|
||||
pl.scan_parquet(self._enriched_glob)
|
||||
.filter(pl.col("date") >= start_full)
|
||||
.sort(["symbol", "date"])
|
||||
)
|
||||
df_hist = lf.select(read_cols).collect()
|
||||
if not df_hist.is_empty():
|
||||
instruments = self._instruments_cache if self._instruments_cache is not None else pl.DataFrame()
|
||||
df_full = compute_indicators(df_hist)
|
||||
df_full = compute_signals(df_full)
|
||||
if instruments is not None and not instruments.is_empty():
|
||||
df_full = compute_limit_signals(df_full, instruments)
|
||||
|
||||
# JOIN instruments 到完整历史 (filter_history/basic_filter 需要 name/股本等列)
|
||||
if instruments is not None and not instruments.is_empty():
|
||||
inst_cols = [c for c in ["name", "total_shares", "float_shares"]
|
||||
if c in instruments.columns and c not in df_full.columns]
|
||||
if inst_cols:
|
||||
df_full = df_full.join(
|
||||
instruments.select(["symbol", *inst_cols]).unique(subset=["symbol"]),
|
||||
on="symbol",
|
||||
how="left",
|
||||
)
|
||||
|
||||
# 缓存完整历史 (含指标+必要基础信息) 供 filter_history/backtest 直接复用
|
||||
self._enriched_history_cache = df_full
|
||||
self._enriched_history_start = df_full["date"].min()
|
||||
logger.info("enriched 历史缓存: %d rows, %s ~ %s",
|
||||
len(df_full), self._enriched_history_start, latest)
|
||||
|
||||
# 只取最新一天作为 enriched_cache
|
||||
df_today = df_full.filter(pl.col("date") == latest)
|
||||
if not df_today.is_empty():
|
||||
self._enriched_cache = df_today
|
||||
self._enriched_cache_date = latest
|
||||
# 构建盘中递推基准: 若最新分区是今天的实时盘中数据,
|
||||
# 递推状态必须停在上一交易日, 不能把今天作为“昨日”。
|
||||
self._build_live_agg(self._live_agg_baseline_date(latest))
|
||||
logger.info("enriched 缓存已计算: %d 只, 日期 %s (即时计算)", len(df_today), latest)
|
||||
return
|
||||
except Exception as e: # noqa: BLE001
|
||||
logger.warning("enriched 即时计算失败, 使用原始 14 列缓存: %s", e)
|
||||
|
||||
# 降级: 直接使用 14 列数据 + 构建 live_agg
|
||||
self._enriched_cache = df_latest
|
||||
self._enriched_cache_date = latest
|
||||
self._build_live_agg(self._live_agg_baseline_date(latest))
|
||||
|
||||
logger.info("enriched 缓存已加载: %d 只, 日期 %s", len(df_latest), latest)
|
||||
except Exception as e: # noqa: BLE001
|
||||
logger.warning("enriched 缓存刷新失败: %s", e)
|
||||
|
||||
def _build_live_agg(self, latest: date) -> None:
|
||||
"""从 OHLCV 即时计算递推状态 + 窗口聚合, 构建盘中实时聚合表。
|
||||
|
||||
优化: 优先使用 _enriched_history_cache (启动时已计算), 避免重复 compute_indicators。
|
||||
"""
|
||||
from datetime import timedelta
|
||||
from app.indicators.pipeline import _ema_alpha
|
||||
|
||||
start_60d = latest - timedelta(days=90) # 日历90天 ≈ 60个交易日
|
||||
|
||||
# 优先使用已有的历史缓存 (避免重复 scan_parquet + compute_indicators)
|
||||
if self._enriched_history_cache is not None and not self._enriched_history_cache.is_empty():
|
||||
hist_all = self._enriched_history_cache
|
||||
if "date" in hist_all.columns and hist_all["date"].min() <= start_60d:
|
||||
# 从历史缓存中提取所需列 (历史缓存已有指标列)
|
||||
base_cols = ["symbol", "date", "open", "high", "low", "close", "volume",
|
||||
"raw_close", "raw_high", "raw_low"]
|
||||
needed = [c for c in base_cols if c in hist_all.columns]
|
||||
df_hist = hist_all.filter(
|
||||
(pl.col("date") >= start_60d) & (pl.col("date") <= latest)
|
||||
).select(needed).sort(["symbol", "date"])
|
||||
|
||||
# 用历史缓存的指标列提取最新日状态 (无需再次 compute_indicators)
|
||||
state_source = hist_all.filter(pl.col("date") == latest)
|
||||
|
||||
state_cols = [
|
||||
"symbol",
|
||||
"ema5", "ema10", "ema20", "ema30", "ema60",
|
||||
"macd_dea",
|
||||
"kdj_k", "kdj_d",
|
||||
"atr_14",
|
||||
"close", "high", "low",
|
||||
"annual_vol_20d",
|
||||
]
|
||||
existing_state = [c for c in state_cols if c in state_source.columns]
|
||||
agg_a = state_source.select(existing_state)
|
||||
else:
|
||||
df_hist = pl.DataFrame()
|
||||
agg_a = pl.DataFrame()
|
||||
else:
|
||||
# 降级: 读 parquet + compute_indicators
|
||||
df_hist, agg_a = self._build_live_agg_from_parquet(latest, start_60d)
|
||||
|
||||
if df_hist.is_empty():
|
||||
self._live_agg_cache = pl.DataFrame()
|
||||
self._live_agg_cache_date = None
|
||||
return
|
||||
|
||||
if agg_a.is_empty():
|
||||
self._live_agg_cache = pl.DataFrame()
|
||||
self._live_agg_cache_date = None
|
||||
return
|
||||
|
||||
# 单独计算 _ema12 / _ema26 (compute_indicators 内部会 drop 掉)
|
||||
df_ema = df_hist.sort(["symbol", "date"]).with_columns([
|
||||
pl.col("close").ewm_mean(alpha=_ema_alpha(12), adjust=False).over("symbol").alias("_ema12"),
|
||||
pl.col("close").ewm_mean(alpha=_ema_alpha(26), adjust=False).over("symbol").alias("_ema26"),
|
||||
]).filter(pl.col("date") == latest).select("symbol", "_ema12", "_ema26")
|
||||
|
||||
agg_a = agg_a.join(df_ema, on="symbol", how="inner")
|
||||
|
||||
# 单独计算 RSI 状态列 (compute_indicators 内部会 drop 掉)
|
||||
df_rsi_base = df_hist.sort(["symbol", "date"]).with_columns(
|
||||
pl.col("close").diff().over("symbol").alias("_daily_delta")
|
||||
)
|
||||
gain = pl.when(pl.col("_daily_delta") > 0).then(pl.col("_daily_delta")).otherwise(0.0)
|
||||
loss = pl.when(pl.col("_daily_delta") < 0).then(-pl.col("_daily_delta")).otherwise(0.0)
|
||||
rsi_exprs = []
|
||||
for n in (6, 14, 24):
|
||||
a = 1.0 / n
|
||||
rsi_exprs.append(gain.ewm_mean(alpha=a, adjust=False).over("symbol").alias(f"_rsi_avg_gain_{n}"))
|
||||
rsi_exprs.append(loss.ewm_mean(alpha=a, adjust=False).over("symbol").alias(f"_rsi_avg_loss_{n}"))
|
||||
df_rsi = (
|
||||
df_rsi_base
|
||||
.with_columns(rsi_exprs)
|
||||
.filter(pl.col("date") == latest)
|
||||
.select("symbol", *[f"_rsi_avg_gain_{n}" for n in (6, 14, 24)],
|
||||
*[f"_rsi_avg_loss_{n}" for n in (6, 14, 24)])
|
||||
)
|
||||
agg_a = agg_a.join(df_rsi, on="symbol", how="inner")
|
||||
|
||||
# 前复权因子: adj_factor = close(复权) / raw_close(原始)
|
||||
if "raw_close" in df_hist.columns:
|
||||
adj_factor_df = (
|
||||
df_hist.filter(pl.col("date") == latest)
|
||||
.select("symbol", (pl.col("close") / pl.col("raw_close")).alias("_adj_factor"))
|
||||
)
|
||||
agg_a = agg_a.join(adj_factor_df, on="symbol", how="left")
|
||||
if "_adj_factor" in agg_a.columns:
|
||||
agg_a = agg_a.with_columns(pl.col("_adj_factor").fill_null(1.0))
|
||||
|
||||
# annual_vol_20d 递推状态: 最近 19 天日收益率的部分和 / 平方和
|
||||
df_daily_pct = (
|
||||
df_hist.sort(["symbol", "date"])
|
||||
.with_columns(
|
||||
pl.col("close").pct_change().over("symbol").alias("_daily_pct")
|
||||
)
|
||||
)
|
||||
df_vol = df_daily_pct.group_by("symbol").agg([
|
||||
pl.col("_daily_pct").tail(19).sum().alias("_vol_19d_pct_sum"),
|
||||
(pl.col("_daily_pct") ** 2).tail(19).sum().alias("_vol_19d_pct_sq_sum"),
|
||||
])
|
||||
agg_a = agg_a.join(df_vol, on="symbol", how="left")
|
||||
|
||||
# 昨日连板数: 从 enriched parquet 取 (用于增量计算同向 +1)
|
||||
lf = pl.scan_parquet(self._enriched_glob).filter(pl.col("date") == latest)
|
||||
consec_cols = [c for c in ["symbol", "consecutive_limit_ups", "consecutive_limit_downs"]
|
||||
if c in lf.collect_schema().names()]
|
||||
if len(consec_cols) == 3:
|
||||
consec_df = lf.select(consec_cols).collect()
|
||||
if not consec_df.is_empty():
|
||||
consec = consec_df.select(
|
||||
"symbol",
|
||||
pl.col("consecutive_limit_ups").alias("_prev_consec_up"),
|
||||
pl.col("consecutive_limit_downs").alias("_prev_consec_down"),
|
||||
)
|
||||
agg_a = agg_a.join(consec, on="symbol", how="left")
|
||||
|
||||
# B类: 按 symbol 分组聚合 — 窗口统计
|
||||
agg_b = (
|
||||
df_hist.sort(["symbol", "date"])
|
||||
.group_by("symbol")
|
||||
.agg([
|
||||
pl.col("close").tail(4).sum().alias("_ma5_partial_sum"),
|
||||
pl.col("close").tail(9).sum().alias("_ma10_partial_sum"),
|
||||
pl.col("close").tail(19).sum().alias("_ma20_partial_sum"),
|
||||
pl.col("close").tail(29).sum().alias("_ma30_partial_sum"),
|
||||
pl.col("close").tail(59).sum().alias("_ma60_partial_sum"),
|
||||
|
||||
pl.col("close").tail(19).sum().alias("_boll_partial_sum"),
|
||||
(pl.col("close").tail(19) ** 2).sum().alias("_boll_partial_sq_sum"),
|
||||
|
||||
pl.col("high").tail(59).max().alias("_high_59d"),
|
||||
pl.col("low").tail(59).min().alias("_low_59d"),
|
||||
|
||||
pl.col("close").tail(5).first().alias("_close_5d_ago"),
|
||||
pl.col("close").tail(10).first().alias("_close_10d_ago"),
|
||||
pl.col("close").tail(20).first().alias("_close_20d_ago"),
|
||||
pl.col("close").tail(30).first().alias("_close_30d_ago"),
|
||||
pl.col("close").tail(60).first().alias("_close_60d_ago"),
|
||||
|
||||
pl.col("volume").tail(4).sum().alias("_vol_ma5_partial_sum"),
|
||||
pl.col("volume").tail(9).sum().alias("_vol_ma10_partial_sum"),
|
||||
|
||||
pl.col("low").tail(8).min().alias("_kdj_8d_low"),
|
||||
pl.col("high").tail(8).max().alias("_kdj_8d_high"),
|
||||
|
||||
pl.col("close").tail(59).len().alias("_window_len"),
|
||||
])
|
||||
)
|
||||
|
||||
self._live_agg_cache = agg_a.join(agg_b, on="symbol", how="inner")
|
||||
self._live_agg_cache_date = latest
|
||||
|
||||
def _live_agg_baseline_date(self, latest: date) -> date:
|
||||
"""盘中递推基准日期。当天实时分区存在时使用上一可用交易日。"""
|
||||
if latest != date.today():
|
||||
return latest
|
||||
try:
|
||||
row = self.execute_one(
|
||||
"SELECT max(date) FROM kline_enriched WHERE date < ?",
|
||||
[latest],
|
||||
)
|
||||
if row and row[0]:
|
||||
d = row[0]
|
||||
return d if isinstance(d, date) else date.fromisoformat(str(d))
|
||||
except Exception: # noqa: BLE001
|
||||
pass
|
||||
return latest
|
||||
|
||||
def _build_live_agg_from_parquet(self, latest: date, start_60d: date) -> tuple[pl.DataFrame, pl.DataFrame]:
|
||||
"""降级路径: 从 parquet 读取数据并计算指标 (当 _enriched_history_cache 不可用时)。"""
|
||||
from app.indicators.pipeline import compute_indicators
|
||||
|
||||
lf = (
|
||||
pl.scan_parquet(self._enriched_glob)
|
||||
.filter(pl.col("date") >= start_60d)
|
||||
.filter(pl.col("date") <= latest)
|
||||
.sort(["symbol", "date"])
|
||||
)
|
||||
|
||||
read_cols = [c for c in ["symbol", "date", "open", "high", "low", "close", "volume",
|
||||
"raw_close", "raw_high", "raw_low"]
|
||||
if c in lf.collect_schema().names()]
|
||||
df_hist = lf.select(read_cols).collect()
|
||||
|
||||
if df_hist.is_empty():
|
||||
return df_hist, pl.DataFrame()
|
||||
|
||||
df_with_indicators = compute_indicators(df_hist)
|
||||
|
||||
state_cols = [
|
||||
"symbol",
|
||||
"ema5", "ema10", "ema20", "ema30", "ema60",
|
||||
"macd_dea",
|
||||
"kdj_k", "kdj_d",
|
||||
"atr_14",
|
||||
"close", "high", "low",
|
||||
"annual_vol_20d",
|
||||
]
|
||||
existing_state = [c for c in state_cols if c in df_with_indicators.columns]
|
||||
agg_a = df_with_indicators.filter(pl.col("date") == latest).select(existing_state)
|
||||
|
||||
return df_hist, agg_a
|
||||
|
||||
def _refresh_instruments(self) -> None:
|
||||
"""加载 instruments 到内存。"""
|
||||
try:
|
||||
df = pl.scan_parquet(self._inst_glob).collect()
|
||||
if not df.is_empty():
|
||||
self._instruments_cache = df
|
||||
logger.info("instruments 缓存已加载: %d 只", len(df))
|
||||
except Exception as e: # noqa: BLE001
|
||||
logger.warning("instruments 缓存刷新失败: %s", e)
|
||||
|
||||
def _refresh_index_instruments(self) -> None:
|
||||
"""加载指数 instruments 到内存。"""
|
||||
try:
|
||||
df = pl.scan_parquet(self._index_inst_glob).collect()
|
||||
if not df.is_empty():
|
||||
self._index_instruments_cache = df
|
||||
logger.info("index instruments 缓存已加载: %d 只", len(df))
|
||||
except Exception as e: # noqa: BLE001
|
||||
logger.debug("index instruments 缓存刷新跳过: %s", e)
|
||||
|
||||
def get_enriched_latest(self) -> tuple[pl.DataFrame, date | None]:
|
||||
"""返回缓存的 enriched 最新日 DataFrame + 日期。如无缓存则懒加载。"""
|
||||
if self._enriched_cache is None:
|
||||
self._refresh_enriched()
|
||||
if self._enriched_cache is None:
|
||||
return pl.DataFrame(), self._enriched_cache_date
|
||||
return self._enriched_cache, self._enriched_cache_date
|
||||
|
||||
def get_enriched_history(self, target_date: date, lookback_days: int) -> pl.DataFrame | None:
|
||||
"""返回预计算的 enriched 历史数据 (仅 lookback 范围, 不含 warmup)。
|
||||
|
||||
warmup 部分在 _refresh_enriched 计算指标时已使用, 策略只需要最终的 lookback 窗口。
|
||||
返回 ~33万行 (90日历天) 而非 ~107万行, filter_history 策略的 group_by 快 20x+。
|
||||
"""
|
||||
cache = self._enriched_history_cache
|
||||
if cache is None or cache.is_empty():
|
||||
return None
|
||||
if "date" not in cache.columns:
|
||||
return None
|
||||
cache_max = cache["date"].max()
|
||||
cache_min = cache["date"].min()
|
||||
from datetime import timedelta
|
||||
# 验证缓存覆盖完整范围 (含 warmup)
|
||||
warmup_start = target_date - timedelta(days=(lookback_days + 60) * 2)
|
||||
if cache_min > warmup_start or cache_max < target_date:
|
||||
return None
|
||||
# 只返回 lookback 范围 (日历天数 ≈ 2/3 交易日, 足够覆盖)
|
||||
lookback_start = target_date - timedelta(days=lookback_days)
|
||||
return cache.filter((pl.col("date") >= lookback_start) & (pl.col("date") <= target_date))
|
||||
|
||||
def get_enriched_range(
|
||||
self,
|
||||
start: date,
|
||||
end: date,
|
||||
symbols: list[str] | None = None,
|
||||
columns: list[str] | None = None,
|
||||
) -> pl.DataFrame | None:
|
||||
"""从预计算 enriched 历史缓存返回完整区间;缓存不覆盖时返回 None。"""
|
||||
if self._enriched_history_cache is None:
|
||||
self._refresh_enriched()
|
||||
cache = self._enriched_history_cache
|
||||
if cache is None or cache.is_empty() or "date" not in cache.columns:
|
||||
return None
|
||||
|
||||
cache_min = cache["date"].min()
|
||||
cache_max = cache["date"].max()
|
||||
if cache_min > start or cache_max < end:
|
||||
return None
|
||||
|
||||
df = cache.filter((pl.col("date") >= start) & (pl.col("date") <= end))
|
||||
if symbols is not None:
|
||||
df = df.filter(pl.col("symbol").is_in(symbols))
|
||||
if columns and not df.is_empty():
|
||||
existing = [c for c in columns if c in df.columns]
|
||||
if "symbol" not in existing and "symbol" in df.columns:
|
||||
existing.insert(0, "symbol")
|
||||
if "date" not in existing and "date" in df.columns:
|
||||
existing.insert(1, "date")
|
||||
df = df.select(existing)
|
||||
return df.sort(["symbol", "date"])
|
||||
|
||||
def get_live_agg(self) -> pl.DataFrame:
|
||||
"""返回盘中实时指标预计算聚合表。如无缓存则懒加载。"""
|
||||
if self._live_agg_cache is None:
|
||||
self._refresh_enriched()
|
||||
if self._live_agg_cache is None:
|
||||
return pl.DataFrame()
|
||||
return self._live_agg_cache
|
||||
|
||||
def get_instruments(self) -> pl.DataFrame:
|
||||
"""返回缓存的 instruments DataFrame。如无缓存则懒加载。"""
|
||||
if self._instruments_cache is None:
|
||||
self._refresh_instruments()
|
||||
if self._instruments_cache is None:
|
||||
return pl.DataFrame()
|
||||
return self._instruments_cache
|
||||
|
||||
def get_index_instruments(self) -> pl.DataFrame:
|
||||
"""返回缓存的指数 instruments DataFrame。如无缓存则懒加载。"""
|
||||
if self._index_instruments_cache is None:
|
||||
self._refresh_index_instruments()
|
||||
if self._index_instruments_cache is None:
|
||||
return pl.DataFrame()
|
||||
return self._index_instruments_cache
|
||||
|
||||
def get_index_symbol_set(self) -> set[str]:
|
||||
"""返回已缓存指数 symbol 集合。"""
|
||||
df = self.get_index_instruments()
|
||||
if df.is_empty() or "symbol" not in df.columns:
|
||||
return set()
|
||||
return set(df["symbol"].cast(pl.Utf8).to_list())
|
||||
|
||||
def enriched_latest_date(self) -> date | None:
|
||||
"""返回缓存中的 enriched 最新日期。"""
|
||||
return self._enriched_cache_date
|
||||
|
||||
# ================================================================
|
||||
# 热路径: Polars 查询 (Chart / Screener / Signals / Intraday)
|
||||
# ================================================================
|
||||
|
||||
def get_daily(
|
||||
self,
|
||||
symbol: str,
|
||||
start: date,
|
||||
end: date,
|
||||
columns: list[str] | None = None,
|
||||
) -> pl.DataFrame:
|
||||
"""单股日K查询 — 从14列parquet读取后即时计算指标。"""
|
||||
from datetime import timedelta
|
||||
|
||||
# 扩展范围用于指标预热 (MA60 需要 ~60 交易日 ≈ 120 日历日)
|
||||
warmup_start = start - timedelta(days=150)
|
||||
|
||||
# 扫描14列 parquet
|
||||
df = self._scan_daily_symbol(symbol, warmup_start, end, None)
|
||||
if not df.is_empty():
|
||||
df = self._compute_enriched_range(df)
|
||||
|
||||
# 尝试用缓存数据覆盖最新日 (盘中更准确)
|
||||
cached, cache_date = self.get_enriched_latest()
|
||||
if not df.is_empty() and cached is not None and not cached.is_empty() and cache_date:
|
||||
if start <= cache_date <= end:
|
||||
cached_part = self._filter_cached(cached, symbol, None)
|
||||
if not cached_part.is_empty():
|
||||
df = df.filter(pl.col("date") != cache_date)
|
||||
common_cols = [c for c in df.columns if c in cached_part.columns]
|
||||
df = pl.concat([df.select(common_cols), cached_part.select(common_cols)])
|
||||
|
||||
# 裁剪到请求范围
|
||||
if not df.is_empty():
|
||||
df = df.filter((pl.col("date") >= start) & (pl.col("date") <= end))
|
||||
|
||||
if columns and not df.is_empty():
|
||||
existing = [c for c in columns if c in df.columns]
|
||||
df = df.select(existing)
|
||||
|
||||
return df
|
||||
|
||||
def get_daily_batch(
|
||||
self,
|
||||
symbols: list[str],
|
||||
start: date,
|
||||
end: date,
|
||||
columns: list[str] | None = None,
|
||||
) -> pl.DataFrame:
|
||||
"""批量日K查询。"""
|
||||
cached, cache_date = self.get_enriched_latest()
|
||||
if cached is not None and not cached.is_empty() and cache_date:
|
||||
if start >= cache_date:
|
||||
return self._filter_cached_batch(cached, symbols, columns)
|
||||
|
||||
# 回退 scan_parquet
|
||||
return self._scan_daily_batch(symbols, start, end, columns)
|
||||
|
||||
def get_index_daily(
|
||||
self,
|
||||
symbol: str,
|
||||
start: date,
|
||||
end: date,
|
||||
columns: list[str] | None = None,
|
||||
) -> pl.DataFrame:
|
||||
"""指数日K查询 — 从独立指数 enriched parquet 读取后即时计算通用指标。"""
|
||||
from datetime import timedelta
|
||||
|
||||
warmup_start = start - timedelta(days=150)
|
||||
df = self._scan_index_daily_symbol(symbol, warmup_start, end, None)
|
||||
if not df.is_empty():
|
||||
df = self._compute_index_enriched_range(df)
|
||||
df = df.filter((pl.col("date") >= start) & (pl.col("date") <= end))
|
||||
if columns and not df.is_empty():
|
||||
existing = [c for c in columns if c in df.columns]
|
||||
df = df.select(existing)
|
||||
return df
|
||||
|
||||
def get_minute(
|
||||
self,
|
||||
symbol: str,
|
||||
trade_date: date,
|
||||
) -> pl.DataFrame:
|
||||
"""分钟K查询 — Polars scan_parquet + predicate pushdown。"""
|
||||
try:
|
||||
return pl.scan_parquet(self._minute_glob).filter(
|
||||
(pl.col("symbol") == symbol)
|
||||
& (pl.col("datetime").dt.date() == trade_date)
|
||||
).sort("datetime").collect()
|
||||
except Exception as e: # noqa: BLE001
|
||||
logger.warning("分钟K查询失败: %s", e)
|
||||
return pl.DataFrame()
|
||||
|
||||
# ================================================================
|
||||
# Polars 查询内部方法
|
||||
# ================================================================
|
||||
|
||||
def _compute_enriched_range(self, df: pl.DataFrame) -> pl.DataFrame:
|
||||
"""对14列enriched数据即时计算完整指标+信号。输入应含足够预热行数。"""
|
||||
from app.indicators.pipeline import compute_indicators, compute_signals, compute_limit_signals, filter_halt_days
|
||||
if df.is_empty() or df.height < 2:
|
||||
return df
|
||||
# 兜底过滤历史脏数据中的停牌日 (close 可能被填充为前收盘价)
|
||||
df = filter_halt_days(df)
|
||||
if df.is_empty() or df.height < 2:
|
||||
return df
|
||||
try:
|
||||
df = compute_indicators(df)
|
||||
df = compute_signals(df)
|
||||
instruments = self.get_instruments()
|
||||
df = compute_limit_signals(df, instruments)
|
||||
except Exception as e: # noqa: BLE001
|
||||
logger.warning("on-demand compute failed: %s", e)
|
||||
return df
|
||||
|
||||
def _compute_index_enriched_range(self, df: pl.DataFrame) -> pl.DataFrame:
|
||||
"""指数只计算通用技术指标和通用信号,跳过涨跌停/股本/市值逻辑。"""
|
||||
from app.indicators.pipeline import compute_indicators, compute_signals
|
||||
if df.is_empty() or df.height < 2:
|
||||
return df
|
||||
try:
|
||||
df = compute_indicators(df)
|
||||
df = compute_signals(df)
|
||||
except Exception as e: # noqa: BLE001
|
||||
logger.warning("index on-demand compute failed: %s", e)
|
||||
return df
|
||||
|
||||
def _filter_cached(self, cached: pl.DataFrame, symbol: str, columns: list[str] | None) -> pl.DataFrame:
|
||||
df = cached.filter(pl.col("symbol") == symbol)
|
||||
if columns and not df.is_empty():
|
||||
existing = [c for c in columns if c in df.columns]
|
||||
df = df.select(existing)
|
||||
return df
|
||||
|
||||
def _filter_cached_batch(self, cached: pl.DataFrame, symbols: list[str], columns: list[str] | None) -> pl.DataFrame:
|
||||
df = cached.filter(pl.col("symbol").is_in(symbols))
|
||||
if columns and not df.is_empty():
|
||||
existing = [c for c in columns if c in df.columns]
|
||||
df = df.select(existing)
|
||||
return df.sort(["symbol", "date"])
|
||||
|
||||
def _scan_daily_symbol(self, symbol: str, start: date, end: date, columns: list[str] | None) -> pl.DataFrame:
|
||||
try:
|
||||
lf = pl.scan_parquet(self._enriched_glob,
|
||||
cast_options=pl.ScanCastOptions(integer_cast="allow-float")).filter(
|
||||
(pl.col("symbol") == symbol)
|
||||
& (pl.col("date") >= start)
|
||||
& (pl.col("date") <= end)
|
||||
).sort("date")
|
||||
if columns:
|
||||
schema_names = lf.collect_schema().names()
|
||||
existing = [c for c in columns if c in schema_names]
|
||||
lf = lf.select(existing)
|
||||
return lf.collect()
|
||||
except Exception as e: # noqa: BLE001
|
||||
logger.warning("日K查询失败: %s", e)
|
||||
return pl.DataFrame()
|
||||
|
||||
def _scan_daily_batch(self, symbols: list[str], start: date, end: date, columns: list[str] | None) -> pl.DataFrame:
|
||||
try:
|
||||
lf = pl.scan_parquet(self._enriched_glob,
|
||||
cast_options=pl.ScanCastOptions(integer_cast="allow-float")).filter(
|
||||
(pl.col("symbol").is_in(symbols))
|
||||
& (pl.col("date") >= start)
|
||||
& (pl.col("date") <= end)
|
||||
).sort(["symbol", "date"])
|
||||
if columns:
|
||||
schema_names = lf.collect_schema().names()
|
||||
existing = [c for c in columns if c in schema_names]
|
||||
lf = lf.select(existing)
|
||||
return lf.collect()
|
||||
except Exception as e: # noqa: BLE001
|
||||
logger.warning("日K批量查询失败: %s", e)
|
||||
return pl.DataFrame()
|
||||
|
||||
def _scan_index_daily_symbol(self, symbol: str, start: date, end: date, columns: list[str] | None) -> pl.DataFrame:
|
||||
try:
|
||||
lf = pl.scan_parquet(self._index_enriched_glob,
|
||||
cast_options=pl.ScanCastOptions(integer_cast="allow-float")).filter(
|
||||
(pl.col("symbol") == symbol)
|
||||
& (pl.col("date") >= start)
|
||||
& (pl.col("date") <= end)
|
||||
).sort("date")
|
||||
if columns:
|
||||
schema_names = lf.collect_schema().names()
|
||||
existing = [c for c in columns if c in schema_names]
|
||||
lf = lf.select(existing)
|
||||
return lf.collect()
|
||||
except Exception as e: # noqa: BLE001
|
||||
logger.warning("指数日K查询失败: %s", e)
|
||||
return pl.DataFrame()
|
||||
|
||||
def _merge_cached_and_scan(
|
||||
self,
|
||||
cached: pl.DataFrame,
|
||||
cache_date: date,
|
||||
symbol: str,
|
||||
start: date,
|
||||
end: date,
|
||||
columns: list[str] | None,
|
||||
) -> pl.DataFrame:
|
||||
"""合并缓存部分 + scan 历史部分。
|
||||
|
||||
历史部分用 strict < cache_date, 避免与缓存重复。
|
||||
两部分 schema 可能不一致 (增量 vs 全量), concat 前对齐列。
|
||||
"""
|
||||
hist = self._scan_daily_symbol(symbol, start, cache_date, columns)
|
||||
cached_part = self._filter_cached(cached, symbol, columns)
|
||||
if hist.is_empty():
|
||||
return cached_part
|
||||
if cached_part.is_empty():
|
||||
return hist
|
||||
# 去重: 历史部分可能包含 cache_date, 去掉后再合并
|
||||
hist = hist.filter(pl.col("date") < cache_date)
|
||||
# 对齐列: 取交集, 统一类型
|
||||
common_cols = [c for c in hist.columns if c in cached_part.columns]
|
||||
hist = hist.select(common_cols)
|
||||
cached_part = cached_part.select(common_cols)
|
||||
# 统一类型: 历史可能是 Float64, 缓存可能是 Int64, 统一为 cast
|
||||
for c in common_cols:
|
||||
if hist[c].dtype != cached_part[c].dtype:
|
||||
# 统一到更宽的类型
|
||||
target = hist[c].dtype if hist.height > cached_part.height else cached_part[c].dtype
|
||||
hist = hist.with_columns(pl.col(c).cast(target))
|
||||
cached_part = cached_part.with_columns(pl.col(c).cast(target))
|
||||
return pl.concat([hist, cached_part])
|
||||
|
||||
# ================================================================
|
||||
# DuckDB 查询 (冷路径: 统计/元数据/自定义SQL)
|
||||
# ================================================================
|
||||
|
||||
def latest_minute_date(self, symbol: str) -> date | None:
|
||||
try:
|
||||
with self._lock:
|
||||
row = self.db.execute(
|
||||
"SELECT max(CAST(datetime AS DATE)) FROM kline_minute WHERE symbol = ?",
|
||||
[symbol],
|
||||
).fetchone()
|
||||
if row and row[0]:
|
||||
return row[0] if isinstance(row[0], date) else date.fromisoformat(str(row[0]))
|
||||
except duckdb.CatalogException:
|
||||
pass
|
||||
return None
|
||||
|
||||
def earliest_daily_date(self) -> date | None:
|
||||
"""本地日K数据的最早日期。"""
|
||||
try:
|
||||
with self._lock:
|
||||
res = self.db.execute(
|
||||
"SELECT min(date) FROM kline_daily",
|
||||
).fetchone()
|
||||
if res and res[0]:
|
||||
d = res[0]
|
||||
return d if isinstance(d, date) else date.fromisoformat(str(d))
|
||||
except Exception:
|
||||
return None
|
||||
return None
|
||||
|
||||
def earliest_minute_date(self) -> date | None:
|
||||
"""本地分钟K数据的最早日期。"""
|
||||
try:
|
||||
with self._lock:
|
||||
res = self.db.execute(
|
||||
"SELECT min(CAST(datetime AS DATE)) FROM kline_minute",
|
||||
).fetchone()
|
||||
if res and res[0]:
|
||||
d = res[0]
|
||||
return d if isinstance(d, date) else date.fromisoformat(str(d))
|
||||
except Exception:
|
||||
return None
|
||||
return None
|
||||
|
||||
def latest_daily_date(self) -> date | None:
|
||||
"""本地日K数据的最新日期。"""
|
||||
try:
|
||||
with self._lock:
|
||||
res = self.db.execute(
|
||||
"SELECT max(date) FROM kline_daily",
|
||||
).fetchone()
|
||||
if res and res[0]:
|
||||
d = res[0]
|
||||
return d if isinstance(d, date) else date.fromisoformat(str(d))
|
||||
except Exception:
|
||||
return None
|
||||
return None
|
||||
|
||||
def _latest_enriched_date_duckdb(self) -> date | None:
|
||||
try:
|
||||
with self._lock:
|
||||
res = self.db.execute(
|
||||
"SELECT max(date) FROM kline_enriched",
|
||||
).fetchone()
|
||||
if res and res[0]:
|
||||
d = res[0]
|
||||
return d if isinstance(d, date) else date.fromisoformat(str(d))
|
||||
except Exception: # noqa: BLE001
|
||||
return None
|
||||
return None
|
||||
|
||||
# ================================================================
|
||||
# 写入 (Pipeline / Sync)
|
||||
# ================================================================
|
||||
|
||||
def append_daily(self, df: pl.DataFrame) -> None:
|
||||
"""按日分区写入日K数据 (merge-upsert)。"""
|
||||
if df.is_empty():
|
||||
return
|
||||
self._write_daily_partition(df, "kline_daily")
|
||||
|
||||
def append_enriched(self, df: pl.DataFrame) -> None:
|
||||
"""按日分区写入 enriched 数据 (merge-upsert)。磁盘仅写入 14 列存储列。"""
|
||||
if df.is_empty():
|
||||
return
|
||||
from app.indicators.pipeline import ENRICHED_STORAGE_COLS
|
||||
storage_cols = [c for c in ENRICHED_STORAGE_COLS if c in df.columns]
|
||||
df_storage = df.select(storage_cols)
|
||||
self._write_daily_partition(df_storage, "kline_daily_enriched")
|
||||
|
||||
def append_index_daily(self, df: pl.DataFrame) -> None:
|
||||
"""按日分区写入指数日K数据 (merge-upsert)。"""
|
||||
if df.is_empty():
|
||||
return
|
||||
self._write_daily_partition(df, "kline_index_daily")
|
||||
|
||||
def append_index_enriched(self, df: pl.DataFrame) -> None:
|
||||
"""按日分区写入指数 enriched 数据。磁盘仅写入通用基础行情窄表。"""
|
||||
if df.is_empty():
|
||||
return
|
||||
from app.indicators.pipeline import ENRICHED_STORAGE_COLS
|
||||
storage_cols = [c for c in ENRICHED_STORAGE_COLS if c in df.columns]
|
||||
df_storage = df.select(storage_cols)
|
||||
self._write_daily_partition(df_storage, "kline_index_enriched")
|
||||
|
||||
def save_index_instruments(self, df: pl.DataFrame) -> None:
|
||||
"""保存指数标的维表。"""
|
||||
if df.is_empty() or "symbol" not in df.columns:
|
||||
return
|
||||
out = self.store.data_dir / "instruments_index" / "instruments_index.parquet"
|
||||
out.parent.mkdir(parents=True, exist_ok=True)
|
||||
df.unique(subset=["symbol"], keep="last").sort("symbol").write_parquet(out)
|
||||
self._index_instruments_cache = None
|
||||
self._refresh_index_instruments()
|
||||
|
||||
def refresh_index_views(self) -> None:
|
||||
"""刷新指数相关 DuckDB 视图。"""
|
||||
d = self.store.data_dir.as_posix()
|
||||
statements = [
|
||||
f"""CREATE OR REPLACE VIEW kline_index_daily AS
|
||||
SELECT * FROM read_parquet('{d}/kline_index_daily/**/*.parquet', union_by_name=true)""",
|
||||
f"""CREATE OR REPLACE VIEW kline_index_enriched AS
|
||||
SELECT * FROM read_parquet('{d}/kline_index_enriched/**/*.parquet', union_by_name=true)""",
|
||||
f"""CREATE OR REPLACE VIEW instruments_index AS
|
||||
SELECT * FROM read_parquet('{d}/instruments_index/**/*.parquet', union_by_name=true)""",
|
||||
]
|
||||
for sql in statements:
|
||||
try:
|
||||
with self._lock:
|
||||
self.db.execute(sql)
|
||||
except Exception as e: # noqa: BLE001
|
||||
logger.debug("index view refresh skipped: %s", e)
|
||||
|
||||
def _write_daily_partition(self, df: pl.DataFrame, table: str) -> None:
|
||||
"""按 date 分区写入 parquet,每个日期一个文件,支持 merge-upsert。"""
|
||||
base = self.store.data_dir / table
|
||||
for date_df in df.partition_by("date"):
|
||||
dt = date_df["date"][0]
|
||||
ds = dt.isoformat() if hasattr(dt, "isoformat") else str(dt)
|
||||
out = base / f"date={ds}" / "part.parquet"
|
||||
out.parent.mkdir(parents=True, exist_ok=True)
|
||||
if out.exists():
|
||||
existing = pl.read_parquet(out)
|
||||
date_df = pl.concat([existing, date_df], how="diagonal_relaxed").unique(
|
||||
subset=["symbol", "date"], keep="last"
|
||||
)
|
||||
date_df = date_df.sort(["symbol", "date"])
|
||||
date_df.write_parquet(out)
|
||||
|
||||
def flush_live_daily(self, df: pl.DataFrame) -> None:
|
||||
"""覆写当天 kline_daily 分区 (实时行情落盘, 非merge)。"""
|
||||
if df.is_empty() or "date" not in df.columns:
|
||||
return
|
||||
base = self.store.data_dir / "kline_daily"
|
||||
dt = df["date"][0]
|
||||
ds = dt.isoformat() if hasattr(dt, "isoformat") else str(dt)
|
||||
out = base / f"date={ds}" / "part.parquet"
|
||||
out.parent.mkdir(parents=True, exist_ok=True)
|
||||
df.sort(["symbol", "date"]).write_parquet(out)
|
||||
|
||||
def flush_live_enriched(self, df: pl.DataFrame) -> None:
|
||||
"""覆写当天 kline_daily_enriched 分区 (实时 enriched 落盘, 非merge)。
|
||||
|
||||
内存缓存保留完整指标列供各服务使用,磁盘仅写入 14 列存储列。
|
||||
"""
|
||||
if df.is_empty() or "date" not in df.columns:
|
||||
return
|
||||
# 内存缓存: 保留完整 66 列
|
||||
self._enriched_cache = df.sort(["symbol"])
|
||||
dt = df["date"][0]
|
||||
self._enriched_cache_date = dt
|
||||
# 磁盘写入: 仅 14 列存储列
|
||||
from app.indicators.pipeline import ENRICHED_STORAGE_COLS
|
||||
storage_cols = [c for c in ENRICHED_STORAGE_COLS if c in df.columns]
|
||||
df_storage = df.select(storage_cols).sort(["symbol"])
|
||||
base = self.store.data_dir / "kline_daily_enriched"
|
||||
ds = dt.isoformat() if hasattr(dt, "isoformat") else str(dt)
|
||||
out = base / f"date={ds}" / "part.parquet"
|
||||
out.parent.mkdir(parents=True, exist_ok=True)
|
||||
df_storage.write_parquet(out)
|
||||
@@ -0,0 +1,66 @@
|
||||
"""请求调度器(§5.6)。
|
||||
|
||||
按 capability 分别维护令牌桶。Phase 0:基础实现;Phase 1 接入批量合并、优先级队列。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
|
||||
from .capabilities import Cap, CapabilitySet
|
||||
|
||||
|
||||
@dataclass
|
||||
class _Bucket:
|
||||
capacity: int # 每分钟令牌数
|
||||
tokens: float
|
||||
last_refill: float # 单位:秒
|
||||
|
||||
def consume(self, n: int = 1) -> float:
|
||||
"""尝试消费 n 个令牌,返回需要等待的秒数(0 表示无需等待)。"""
|
||||
now = time.monotonic()
|
||||
elapsed = now - self.last_refill
|
||||
# 60s 内补满 capacity,匀速补
|
||||
refill = elapsed * (self.capacity / 60.0)
|
||||
self.tokens = min(self.capacity, self.tokens + refill)
|
||||
self.last_refill = now
|
||||
|
||||
if self.tokens >= n:
|
||||
self.tokens -= n
|
||||
return 0.0
|
||||
deficit = n - self.tokens
|
||||
# 还需多少秒才能补齐
|
||||
wait = deficit / (self.capacity / 60.0)
|
||||
# 不预扣,留给下一次再竞争(避免饿死优先级高的请求)
|
||||
return wait
|
||||
|
||||
|
||||
class Scheduler:
|
||||
"""每个 capability 一个桶。"""
|
||||
|
||||
def __init__(self, capset: CapabilitySet) -> None:
|
||||
self._capset = capset
|
||||
self._buckets: dict[Cap, _Bucket] = {}
|
||||
self._locks: dict[Cap, asyncio.Lock] = {}
|
||||
for cap, lim in capset.all().items():
|
||||
if lim.rpm:
|
||||
self._buckets[cap] = _Bucket(
|
||||
capacity=lim.rpm,
|
||||
tokens=lim.rpm,
|
||||
last_refill=time.monotonic(),
|
||||
)
|
||||
self._locks[cap] = asyncio.Lock()
|
||||
|
||||
async def acquire(self, cap: Cap, n: int = 1) -> None:
|
||||
"""阻塞直到拿到 n 个令牌。无桶 = 不限流(由调用方保证)。"""
|
||||
bucket = self._buckets.get(cap)
|
||||
if bucket is None:
|
||||
return
|
||||
lock = self._locks[cap]
|
||||
async with lock:
|
||||
while True:
|
||||
wait = bucket.consume(n)
|
||||
if wait == 0:
|
||||
return
|
||||
await asyncio.sleep(wait)
|
||||
Reference in New Issue
Block a user