stock/backend/app/data/company.py

"""上市公司基本信息（tushare stock_company，详情页按需单查懒加载）。

不做批量同步：详情页打开才触发，ts_code= 单查一次一调用；
行即缓存 —— stock_company 表 30 天新鲜度门控（数据月更），tushare 查无此股
写墓碑行（业务字段全 NULL）做负缓存，避免无数据代码每次都穿透（0.35s 控频
+ 可能 62s 限频重试）。墓碑同样按 updated_at 参与 30 天刷新，新股上市后能自动补上。
"""
from __future__ import annotations

import asyncio
import time
from datetime import datetime, timedelta, timezone

from sqlalchemy import select
from sqlalchemy.dialects.postgresql import insert as pg_insert
from sqlalchemy.ext.asyncio import AsyncSession

from ..config import settings
from ..db import async_session
from ..models import StockCompany

_REFRESH_DAYS = 30

# 频率超限特征（等待 62s 重试一次；与 screener.market_sync / data.etf_sync._call_retry 同款语义）
_RATE_MARKS = ("频率超限", "每分钟")

# 显式列出全部字段：introduction/office/main_business/business_scope 文档标注默认不显示，
# 不传 fields 时 tushare 不返回这四列（实测 000001.SZ）
_FIELDS = (
    "ts_code,com_name,com_id,chairman,manager,secretary,reg_capital,"
    "setup_date,province,city,introduction,website,email,office,"
    "employees,main_business,business_scope"
)

_pro = None  # 惰性单例（get_pro 每次都 ts.set_token 写文件，没必要重复）


def _get_pro():
    if not settings.tushare_token:
        raise RuntimeError("未配置 TUSHARE_TOKEN，无法拉取公司简介（backend/.env）")
    global _pro
    if _pro is None:
        from .tushare_provider import get_pro

        _pro = get_pro()
    return _pro


def _call_retry(fn, *args, **kwargs):
    """同步调用 tushare 接口；「每分钟」级频率超限等 62s 重试一次。"""
    try:
        return fn(*args, **kwargs)
    except Exception as e:  # noqa: BLE001
        msg = str(e)
        if any(m in msg for m in _RATE_MARKS) and "小时" not in msg:
            time.sleep(62)
            return fn(*args, **kwargs)
        raise


def _utcnow() -> datetime:
    return datetime.now(timezone.utc).replace(tzinfo=None)


def _fresh(updated_at: datetime | None) -> bool:
    return updated_at is not None and updated_at >= _utcnow() - timedelta(days=_REFRESH_DAYS)


def _s(v) -> str | None:
    """pandas NaN / 空串 / None -> None，其余 strip。"""
    if v is None or (isinstance(v, float) and v != v):
        return None
    s = str(v).strip()
    return s or None


def _f(v) -> float | None:
    try:
        f = float(v)
    except (TypeError, ValueError):
        return None
    return None if f != f else f  # NaN -> None


def _i(v) -> int | None:
    f = _f(v)
    return None if f is None else int(f)


def fetch_company_sync(ts_code: str) -> dict | None:
    """同步拉单只公司简介（需在 to_thread 里跑）；返回行 dict，无此股返回 None。"""
    time.sleep(settings.screener_sync_interval)
    df = _call_retry(_get_pro().stock_company, ts_code=ts_code, fields=_FIELDS)
    if df is None or df.empty:
        return None
    r = df.iloc[0]
    return {
        "ts_code": ts_code,
        "com_name": _s(r.get("com_name")),
        "com_id": _s(r.get("com_id")),
        "chairman": _s(r.get("chairman")),
        "manager": _s(r.get("manager")),
        "secretary": _s(r.get("secretary")),
        "reg_capital": _f(r.get("reg_capital")),
        "setup_date": _s(r.get("setup_date")),
        "province": _s(r.get("province")),
        "city": _s(r.get("city")),
        "introduction": _s(r.get("introduction")),
        "website": _s(r.get("website")),
        "email": _s(r.get("email")),
        "office": _s(r.get("office")),
        "employees": _i(r.get("employees")),
        "main_business": _s(r.get("main_business")),
        "business_scope": _s(r.get("business_scope")),
        "updated_at": _utcnow(),
    }


_COLS = (
    "ts_code", "com_name", "com_id", "chairman", "manager", "secretary",
    "reg_capital", "setup_date", "province", "city", "introduction",
    "website", "email", "office", "employees", "main_business", "business_scope",
)


def _row_dict(row: StockCompany) -> dict:
    return {c: getattr(row, c) for c in _COLS}


async def _upsert(session: AsyncSession, row: dict) -> None:
    stmt = pg_insert(StockCompany).values(row)
    stmt = stmt.on_conflict_do_update(
        index_elements=["ts_code"],
        set_={c: stmt.excluded[c] for c in row if c != "ts_code"},
    )
    await session.execute(stmt)
    await session.commit()


# per-code 锁：同 code 首次并发 N 个请求只有 1 个打 tushare，其余等锁后双检命中；
# dict 不清理（对象数 = 全市场股票数，内存可忽略）。uvicorn 单进程场景够用，
# 多 worker 最坏情况是重复拉一次 + ON CONFLICT 幂等，无害。
_code_locks: dict[str, asyncio.Lock] = {}
_guard = asyncio.Lock()


async def get_company(session: AsyncSession, ts_code: str) -> dict | None:
    """读穿透：新鲜行直返；否则加 per-code 锁 -> 新会话双检 -> to_thread 拉 -> upsert。

    返回 None = 确认无数据（墓碑已落库）；抛异常 = tushare 拉取失败且无旧行可降级。
    """
    row = (await session.execute(
        select(StockCompany).where(StockCompany.ts_code == ts_code))).scalar_one_or_none()
    if row is not None and _fresh(row.updated_at):
        return _row_dict(row) if row.com_name is not None else None  # 墓碑 -> None
    # 释放请求会话持有的连接：后面可能隔着 1-2s 的 tushare 调用，别长占连接池。
    # 用 close() 而非 rollback()：rollback 会把会话身份映射里的实例全部 expire——
    # 包括 require_user 刚塞进 auth._session_cache 的 User，下个请求命中鉴权缓存即
    # DetachedInstanceError 500；close() 同样归还连接且已加载属性保持可访问。
    await session.close()

    async with _guard:
        lock = _code_locks.setdefault(ts_code, asyncio.Lock())
    async with lock:
        async with async_session() as s2:  # 锁内重读 + 写入走新会话
            row = (await s2.execute(
                select(StockCompany).where(StockCompany.ts_code == ts_code))).scalar_one_or_none()
            if row is not None and _fresh(row.updated_at):
                return _row_dict(row) if row.com_name is not None else None
            try:
                fetched = await asyncio.to_thread(fetch_company_sync, ts_code)
            except Exception:
                # 降级：库内有真实旧行（哪怕超 30 天）照常返回，不把「上游挂了」伪装成「无数据」
                if row is not None and row.com_name is not None:
                    return _row_dict(row)
                raise
            await _upsert(s2, fetched or {"ts_code": ts_code, "updated_at": _utcnow()})
            return fetched