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
+66 -32
View File
@@ -13,6 +13,10 @@ _BAOSTOCK_FIELDS = "date,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]:
"""Download daily bars from akshare. Returns DataFrame with OHLCV columns."""
try:
@@ -124,50 +128,80 @@ def download_daily_batch(
start: str,
end: str,
adjust: str = "qfq",
akshare_fallback: bool = True,
akshare_fallback: bool = False,
relogin_every: int = 200,
) -> 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
speed-up when fetching thousands of symbols. Yields ``(symbol, df)`` as each
symbol completes so callers can stream results to disk; ``df`` is ``None``
when both sources fail. Each ``df`` has the same 8 columns as
:func:`download_daily`.
Logging in once (instead of per symbol) is the dominant speed-up for
thousands of symbols, but baostock drops a session after a while
(subsequent queries return ``用户未登录``). So we refresh the session every
``relogin_every`` symbols and also re-login + retry once whenever a query
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:
symbols: Internal-form symbols (``sh600000`` / ``sz000001``).
start, end: ``YYYY-MM-DD`` bounds.
adjust: ``qfq`` / ``hfq`` / ``''``.
akshare_fallback: Retry a failed symbol through akshare before yielding
``None``.
akshare_fallback: Retry a failed symbol through akshare. Off by default
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")
def _relogin() -> None:
try:
bs.logout()
except Exception:
pass
bs.login()
def _fetch(symbol: str) -> Optional[pd.DataFrame]:
"""One baostock query; returns df, or None (no data), or raises _SessionLost."""
code = f"{symbol[:2]}.{symbol[2:]}"
rs = bs.query_history_k_data_plus(
code=code, fields=_BAOSTOCK_FIELDS,
start_date=start, end_date=end, frequency="d", adjustflag=flag,
)
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 = []
while rs.next():
rows.append(rs.get_row_data())
if not rows:
return None
df = pd.DataFrame(rows, columns=["date", *_OHLCV])
# Suspended-trading days come back as empty strings; coerce to NaN
# rather than crashing the whole symbol.
df[_OHLCV] = df[_OHLCV].apply(pd.to_numeric, errors="coerce")
df["symbol"] = symbol
return df[["symbol", "date", *_OHLCV]]
bs.login()
try:
for symbol in symbols:
for i, symbol in enumerate(symbols):
if i and relogin_every and i % relogin_every == 0:
_relogin()
df: Optional[pd.DataFrame] = None
try:
code = f"{symbol[:2]}.{symbol[2:]}"
rs = bs.query_history_k_data_plus(
code=code, fields=_BAOSTOCK_FIELDS,
start_date=start, end_date=end,
frequency="d", adjustflag=flag,
)
if rs.error_code == "0":
rows = []
while rs.next():
rows.append(rs.get_row_data())
if rows:
df = pd.DataFrame(rows, columns=["date", *_OHLCV])
# Suspended-trading days come back as empty strings;
# coerce to NaN rather than crashing the whole symbol.
df[_OHLCV] = df[_OHLCV].apply(pd.to_numeric, errors="coerce")
df["symbol"] = symbol
df = df[["symbol", "date", *_OHLCV]]
else:
logger.warning("baostock error for %s: %s", symbol, rs.error_msg)
except Exception as e:
logger.warning("baostock download failed for %s: %s", symbol, e)
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:
logger.warning("baostock download failed for %s: %s", symbol, e)
break
if (df is None or df.empty) and akshare_fallback:
df = _download_akshare(symbol, start, end, adjust)