as-event/backend/app/sources.py

"""外部数据源适配（直连现有内网库，只读）+ 可插拔 OHLC Provider。

- LegacyIndexSource     读 zs_day_data(MySQL-A 18.199) 的指数 close/percent/amount
- LegacySentimentSource 读 gp_market_sentiment(PG 16.150) 的情绪
- OHLCProvider          指数开高低收源（你后续提供）。默认 NullOHLCProvider = close_only 退化

所有源在 DSN 未配置或连接失败时**降级返回空**，绝不让 ETL 崩掉。
"""
from __future__ import annotations

import logging
from dataclasses import dataclass
from datetime import date
from typing import Dict, List, Optional, Protocol

from sqlalchemy import create_engine, text
from sqlalchemy.engine import Engine

from .config import get_settings

log = logging.getLogger("as-event.sources")
settings = get_settings()

# ---- 惰性引擎（只在配置了 DSN 时创建）----
_engines: Dict[str, Optional[Engine]] = {}


def _engine_for(dsn: Optional[str]) -> Optional[Engine]:
    if not dsn:
        return None
    if dsn not in _engines:
        try:
            _engines[dsn] = create_engine(dsn, pool_pre_ping=True, future=True)
        except Exception as e:  # noqa: BLE001
            log.warning("创建引擎失败 dsn=%s err=%s", _mask(dsn), e)
            _engines[dsn] = None
    return _engines[dsn]


def _mask(dsn: str) -> str:
    """隐藏 DSN 中的密码用于日志。"""
    try:
        head, tail = dsn.split("@", 1)
        scheme_user = head.split("//", 1)
        return f"{scheme_user[0]}//***@{tail}"
    except Exception:  # noqa: BLE001
        return "<dsn>"


def _to_float(v) -> Optional[float]:
    if v is None:
        return None
    try:
        return float(v)
    except (TypeError, ValueError):
        return None


# ============================ 数据行 ============================
@dataclass
class IndexBar:
    trade_date: date
    close: Optional[float]
    amount: Optional[float] = None
    pct_chg: Optional[float] = None
    open: Optional[float] = None
    high: Optional[float] = None
    low: Optional[float] = None
    volume: Optional[int] = None


@dataclass
class SentimentRow:
    trade_date: date
    up_down_ratio: Optional[float] = None
    median_pct_chg: Optional[float] = None
    pct_chg_gt_5_count: Optional[int] = None


# ============================ 现有库读适配 ============================
class LegacyIndexSource:
    """读 zs_day_data（MySQL-A 18.199）。字段：symbol, timestamp, close, percent, amount。"""

    def __init__(self) -> None:
        self.engine = _engine_for(settings.legacy_mysql_dsn)

    @property
    def available(self) -> bool:
        return self.engine is not None

    def fetch(self, index_code: str, start: date, end: date) -> List[IndexBar]:
        if not self.engine:
            log.info("LEGACY_MYSQL_DSN 未配置，跳过指数拉取 %s", index_code)
            return []
        sql = text(
            """
            SELECT DATE(`timestamp`) AS d, `close` AS c, `percent` AS p, `amount` AS a
            FROM zs_day_data
            WHERE symbol = :code AND DATE(`timestamp`) BETWEEN :s AND :e
            ORDER BY d ASC
            """
        )
        try:
            with self.engine.connect() as conn:
                rows = conn.execute(sql, {"code": index_code, "s": start, "e": end}).all()
        except Exception as e:  # noqa: BLE001
            log.warning("读 zs_day_data 失败 %s: %s", index_code, e)
            return []
        return [
            IndexBar(
                trade_date=r.d,
                close=_to_float(r.c),
                amount=_to_float(r.a),
                pct_chg=_to_float(r.p),
            )
            for r in rows
        ]


class LegacySentimentSource:
    """读 gp_market_sentiment（PG 16.150）。字段：trade_date, up_down_ratio, median_pct_chg, pct_chg_gt_5_count。"""

    def __init__(self) -> None:
        self.engine = _engine_for(settings.legacy_pg_dsn)

    @property
    def available(self) -> bool:
        return self.engine is not None

    def fetch(self, start: date, end: date) -> List[SentimentRow]:
        if not self.engine:
            log.info("LEGACY_PG_DSN 未配置，跳过情绪拉取")
            return []
        sql = text(
            """
            SELECT trade_date, up_down_ratio, median_pct_chg, pct_chg_gt_5_count
            FROM gp_market_sentiment
            WHERE trade_date BETWEEN :s AND :e
            ORDER BY trade_date ASC
            """
        )
        try:
            with self.engine.connect() as conn:
                rows = conn.execute(sql, {"s": start, "e": end}).all()
        except Exception as e:  # noqa: BLE001
            log.warning("读 gp_market_sentiment 失败: %s", e)
            return []
        out: List[SentimentRow] = []
        for r in rows:
            cnt = r.pct_chg_gt_5_count
            out.append(
                SentimentRow(
                    trade_date=r.trade_date,
                    up_down_ratio=_to_float(r.up_down_ratio),
                    median_pct_chg=_to_float(r.median_pct_chg),
                    pct_chg_gt_5_count=int(cnt) if cnt is not None else None,
                )
            )
        return out


# ============================ 可插拔 OHLC Provider ============================
class OHLCProvider(Protocol):
    """指数开高低收源接口。你接入自己的数据源时实现此协议。"""

    name: str

    def fetch(self, index_code: str, start: date, end: date) -> Dict[date, IndexBar]:
        """返回 {trade_date: IndexBar(带 open/high/low/close[/volume])}。"""
        ...


class NullOHLCProvider:
    """默认：无 OHLC 源。ETL 会用 close 兜底成 open=high=low=close，标 ohlc_source='close_only'。"""

    name = "none"

    def fetch(self, index_code: str, start: date, end: date) -> Dict[date, IndexBar]:
        return {}


# ------- 你后续接入 OHLC 源的示例骨架（改名注册即可）-------
# class MyOHLCProvider:
#     name = "myprovider"
#     def fetch(self, index_code, start, end):
#         # 调你自己的接口/库表，返回 {date: IndexBar(open/high/low/close/volume)}
#         return {}


_OHLC_REGISTRY = {
    NullOHLCProvider.name: NullOHLCProvider,
    # MyOHLCProvider.name: MyOHLCProvider,   # <-- 接入后取消注释
}


def build_ohlc_provider() -> OHLCProvider:
    name = settings.ohlc_provider
    cls = _OHLC_REGISTRY.get(name)
    if cls is None:
        log.warning("未知 OHLC_PROVIDER=%s，回退 none（收盘线）", name)
        cls = NullOHLCProvider
    return cls()