fix: keep baostock session alive across bulk downloads

baostock drops a session after a few hundred queries; every later query
then returned 用户未登录 and the symbol failed. download_daily_batch now
refreshes the session every relogin_every symbols and re-logs in + retries
once on a detected session loss. akshare fallback now defaults off in the
batch path since it is slow/unreliable on this network and the re-login
keeps baostock as the fast primary.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
Yuxuan Yan
2026-06-09 14:10:51 +08:00
parent 1caa63faeb
commit 419114b87b
+56 -22
View File
@@ -13,6 +13,10 @@ _BAOSTOCK_FIELDS = "date,open,high,low,close,volume,amount"
_OHLCV = ["open", "high", "low", "close", "volume", "amount"] _OHLCV = ["open", "high", "low", "close", "volume", "amount"]
class _SessionLost(Exception):
"""baostock reported the session was dropped (``用户未登录``)."""
def _download_akshare(symbol: str, start: str, end: str, adjust: str = "qfq") -> Optional[pd.DataFrame]: def _download_akshare(symbol: str, start: str, end: str, adjust: str = "qfq") -> Optional[pd.DataFrame]:
"""Download daily bars from akshare. Returns DataFrame with OHLCV columns.""" """Download daily bars from akshare. Returns DataFrame with OHLCV columns."""
try: try:
@@ -124,50 +128,80 @@ def download_daily_batch(
start: str, start: str,
end: str, end: str,
adjust: str = "qfq", adjust: str = "qfq",
akshare_fallback: bool = True, akshare_fallback: bool = False,
relogin_every: int = 200,
) -> Iterator[Tuple[str, Optional[pd.DataFrame]]]: ) -> Iterator[Tuple[str, Optional[pd.DataFrame]]]:
"""Download many symbols under a single baostock session. """Download many symbols, keeping a baostock session alive across the run.
Logging into baostock once per call (instead of per symbol) is the dominant Logging in once (instead of per symbol) is the dominant speed-up for
speed-up when fetching thousands of symbols. Yields ``(symbol, df)`` as each thousands of symbols, but baostock drops a session after a while
symbol completes so callers can stream results to disk; ``df`` is ``None`` (subsequent queries return ``用户未登录``). So we refresh the session every
when both sources fail. Each ``df`` has the same 8 columns as ``relogin_every`` symbols and also re-login + retry once whenever a query
:func:`download_daily`. reports the session is gone. Yields ``(symbol, df)`` as each symbol
completes; ``df`` is ``None`` when no data is available. Each ``df`` has the
same 8 columns as :func:`download_daily`.
Args: Args:
symbols: Internal-form symbols (``sh600000`` / ``sz000001``). symbols: Internal-form symbols (``sh600000`` / ``sz000001``).
start, end: ``YYYY-MM-DD`` bounds. start, end: ``YYYY-MM-DD`` bounds.
adjust: ``qfq`` / ``hfq`` / ``''``. adjust: ``qfq`` / ``hfq`` / ``''``.
akshare_fallback: Retry a failed symbol through akshare before yielding akshare_fallback: Retry a failed symbol through akshare. Off by default
``None``. because akshare is unreliable on the deployment network and each
failed attempt is slow; baostock + re-login is the fast path.
relogin_every: Proactively refresh the baostock session every N symbols.
""" """
flag = _BAOSTOCK_ADJUST.get(adjust, "2") flag = _BAOSTOCK_ADJUST.get(adjust, "2")
def _relogin() -> None:
try:
bs.logout()
except Exception:
pass
bs.login() bs.login()
try:
for symbol in symbols: def _fetch(symbol: str) -> Optional[pd.DataFrame]:
df: Optional[pd.DataFrame] = None """One baostock query; returns df, or None (no data), or raises _SessionLost."""
try:
code = f"{symbol[:2]}.{symbol[2:]}" code = f"{symbol[:2]}.{symbol[2:]}"
rs = bs.query_history_k_data_plus( rs = bs.query_history_k_data_plus(
code=code, fields=_BAOSTOCK_FIELDS, code=code, fields=_BAOSTOCK_FIELDS,
start_date=start, end_date=end, start_date=start, end_date=end, frequency="d", adjustflag=flag,
frequency="d", adjustflag=flag,
) )
if rs.error_code == "0": if rs.error_code != "0":
if "未登录" in (rs.error_msg or ""):
raise _SessionLost(rs.error_msg)
logger.warning("baostock error for %s: %s", symbol, rs.error_msg)
return None
rows = [] rows = []
while rs.next(): while rs.next():
rows.append(rs.get_row_data()) rows.append(rs.get_row_data())
if rows: if not rows:
return None
df = pd.DataFrame(rows, columns=["date", *_OHLCV]) df = pd.DataFrame(rows, columns=["date", *_OHLCV])
# Suspended-trading days come back as empty strings; # Suspended-trading days come back as empty strings; coerce to NaN
# coerce to NaN rather than crashing the whole symbol. # rather than crashing the whole symbol.
df[_OHLCV] = df[_OHLCV].apply(pd.to_numeric, errors="coerce") df[_OHLCV] = df[_OHLCV].apply(pd.to_numeric, errors="coerce")
df["symbol"] = symbol df["symbol"] = symbol
df = df[["symbol", "date", *_OHLCV]] return df[["symbol", "date", *_OHLCV]]
else:
logger.warning("baostock error for %s: %s", symbol, rs.error_msg) bs.login()
try:
for i, symbol in enumerate(symbols):
if i and relogin_every and i % relogin_every == 0:
_relogin()
df: Optional[pd.DataFrame] = None
for attempt in (1, 2):
try:
df = _fetch(symbol)
break
except _SessionLost:
if attempt == 1:
_relogin() # session dropped — refresh and retry once
continue
logger.warning("baostock session lost for %s after relogin", symbol)
except Exception as e: except Exception as e:
logger.warning("baostock download failed for %s: %s", symbol, e) logger.warning("baostock download failed for %s: %s", symbol, e)
break
if (df is None or df.empty) and akshare_fallback: if (df is None or df.empty) and akshare_fallback:
df = _download_akshare(symbol, start, end, adjust) df = _download_akshare(symbol, start, end, adjust)