From 419114b87bacb93f6fec3dfa9ff8ea75887f08b3 Mon Sep 17 00:00:00 2001 From: Yuxuan Yan Date: Tue, 9 Jun 2026 14:10:51 +0800 Subject: [PATCH] fix: keep baostock session alive across bulk downloads MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- data/downloader.py | 98 +++++++++++++++++++++++++++++++--------------- 1 file changed, 66 insertions(+), 32 deletions(-) diff --git a/data/downloader.py b/data/downloader.py index 5b3d759..2c15734 100644 --- a/data/downloader.py +++ b/data/downloader.py @@ -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)