as-event/backend/app/etl.py

"""ETL：把现有内网库数据同步进自有库。

    python -m app.etl init-db
    python -m app.etl sync-index      --start 20150101 [--end 20260726] [--codes 000001.SH,399006.SZ]
    python -m app.etl sync-sentiment  --start 20150101
    python -m app.etl sync-all        --start 20150101

设计：看板只读自有库；此处是唯一「读现有库、写自有库」的地方。
内网库不可达时降级为空同步，不报错。
"""
from __future__ import annotations

import argparse
import logging
from datetime import date, datetime
from typing import Dict, List, Optional

from sqlalchemy.dialects.postgresql import insert as pg_insert

from .config import get_settings
from .db import IndexDaily, SentimentDaily, engine, init_db
from .sources import (
    IndexBar,
    LegacyIndexSource,
    LegacySentimentSource,
    build_ohlc_provider,
)

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s")
log = logging.getLogger("as-event.etl")
settings = get_settings()


def _parse_date(s: str) -> date:
    s = s.strip()
    for fmt in ("%Y%m%d", "%Y-%m-%d"):
        try:
            return datetime.strptime(s, fmt).date()
        except ValueError:
            continue
    raise ValueError(f"无法解析日期: {s}（用 YYYYMMDD 或 YYYY-MM-DD）")


# ============================ 指数同步 ============================
def sync_index(start: date, end: date, codes: Optional[List[str]] = None) -> Dict[str, int]:
    codes = codes or settings.index_code_list
    legacy = LegacyIndexSource()
    ohlc = build_ohlc_provider()
    result: Dict[str, int] = {}

    for code in codes:
        bars = legacy.fetch(code, start, end)                 # close/amount/pct
        ohlc_map = ohlc.fetch(code, start, end)               # 可能为空
        by_date: Dict[date, IndexBar] = {b.trade_date: b for b in bars}
        # 若 OHLC 源提供了 close 而现有库没有该日，也纳入
        for d, ob in ohlc_map.items():
            by_date.setdefault(d, IndexBar(trade_date=d, close=ob.close))

        rows = []
        for d, b in sorted(by_date.items()):
            ob = ohlc_map.get(d)
            close = b.close if b.close is not None else (ob.close if ob else None)
            if ob and ob.open is not None:
                o, h, l = ob.open, ob.high, ob.low
                vol = ob.volume
                src = ohlc.name
            else:
                o = h = l = close          # 收盘兜底
                vol = b.volume
                src = "close_only"
            rows.append(
                dict(
                    index_code=code, trade_date=d, open=o, high=h, low=l, close=close,
                    volume=vol, amount=b.amount, pct_chg=b.pct_chg,
                    turnover_rate=None, ohlc_source=src,
                )
            )

        _upsert_index(rows)
        result[code] = len(rows)
        log.info("指数同步 %s: %d 行 (ohlc=%s)", code, len(rows), ohlc.name)
    return result


def _upsert_index(rows: List[dict]) -> None:
    if not rows:
        return
    stmt = pg_insert(IndexDaily).values(rows)
    stmt = stmt.on_conflict_do_update(
        constraint="uq_index_daily",
        set_={
            "open": stmt.excluded.open, "high": stmt.excluded.high,
            "low": stmt.excluded.low, "close": stmt.excluded.close,
            "volume": stmt.excluded.volume, "amount": stmt.excluded.amount,
            "pct_chg": stmt.excluded.pct_chg, "turnover_rate": stmt.excluded.turnover_rate,
            "ohlc_source": stmt.excluded.ohlc_source,
        },
    )
    with engine.begin() as conn:
        conn.execute(stmt)


# ============================ 情绪同步 ============================
def sync_sentiment(start: date, end: date) -> int:
    legacy = LegacySentimentSource()
    rows = legacy.fetch(start, end)
    payload = [
        dict(
            trade_date=r.trade_date,
            up_down_ratio=r.up_down_ratio,
            median_pct_chg=r.median_pct_chg,
            pct_chg_gt_5_count=r.pct_chg_gt_5_count,
        )
        for r in rows
    ]
    _upsert_sentiment(payload)
    log.info("情绪同步: %d 行", len(payload))
    return len(payload)


def _upsert_sentiment(rows: List[dict]) -> None:
    if not rows:
        return
    stmt = pg_insert(SentimentDaily).values(rows)
    stmt = stmt.on_conflict_do_update(
        index_elements=["trade_date"],
        set_={
            "up_down_ratio": stmt.excluded.up_down_ratio,
            "median_pct_chg": stmt.excluded.median_pct_chg,
            "pct_chg_gt_5_count": stmt.excluded.pct_chg_gt_5_count,
        },
    )
    with engine.begin() as conn:
        conn.execute(stmt)


# ============================ CLI ============================
def main() -> None:
    p = argparse.ArgumentParser(description="as-event ETL")
    sub = p.add_subparsers(dest="cmd", required=True)

    sub.add_parser("init-db", help="建表（等价 ddl/own_db_init.sql）")

    for name in ("sync-index", "sync-sentiment", "sync-all"):
        sp = sub.add_parser(name)
        sp.add_argument("--start", default="20150101")
        sp.add_argument("--end", default=date.today().strftime("%Y%m%d"))
        if name in ("sync-index", "sync-all"):
            sp.add_argument("--codes", default=None, help="逗号分隔，缺省用配置")

    args = p.parse_args()
    if args.cmd == "init-db":
        init_db()
        log.info("建表完成")
        return

    start, end = _parse_date(args.start), _parse_date(args.end)
    codes = args.codes.split(",") if getattr(args, "codes", None) else None
    if args.cmd == "sync-index":
        sync_index(start, end, codes)
    elif args.cmd == "sync-sentiment":
        sync_sentiment(start, end)
    elif args.cmd == "sync-all":
        sync_index(start, end, codes)
        sync_sentiment(start, end)


if __name__ == "__main__":
    main()