stock_chart_site/backend/app/fetch/symbols_seed.py

"""KRX 전 종목 리스트를 symbols 테이블에 시드한다.

검색 UX 가 KRX 전체 종목명을 대상으로 동작해야 하므로 전 종목을 미리 적재한다.
10 개 SEED_TICKERS 는 is_seed=TRUE 로 마크.
"""
from __future__ import annotations

import logging
from dataclasses import dataclass

from sqlalchemy import text

from app.db.connection import get_engine
from app.seed.seed_tickers import SEED_CODES, SEED_TICKERS

logger = logging.getLogger(__name__)


@dataclass
class SeedReport:
    inserted: int
    updated: int
    seed_marked: int
    markets: dict[str, int]


def _fetch_market_listing(market: str) -> list[tuple[str, str]]:
    """pykrx 로 한 시장의 (code, name) 목록을 가져온다.

    pykrx 가 외부 KRX 서버에 의존하므로 호출 측에서 예외 처리한다.
    """
    from pykrx import stock as krx  # local import: heavy import

    tickers = krx.get_market_ticker_list(market=market)
    out: list[tuple[str, str]] = []
    for code in tickers:
        name = krx.get_market_ticker_name(code) or ""
        if not name:
            continue
        out.append((code, name))
    return out


def _upsert_seed_tickers() -> int:
    """SEED 10종목 강제 upsert. 네트워크 불필요 → KRX 실패와 무관하게 항상 성공.

    별도 트랜잭션이라 KRX 시드가 나중에 실패해도 살아남는다.
    """
    engine = get_engine()
    with engine.begin() as conn:
        for t in SEED_TICKERS:
            conn.execute(
                text(
                    """
                    INSERT INTO symbols (code, name, market, is_seed)
                    VALUES (:code, :name, :market, TRUE)
                    ON CONFLICT (code) DO UPDATE
                      SET name    = EXCLUDED.name,
                          market  = EXCLUDED.market,
                          is_seed = TRUE
                    """
                ),
                {"code": t.code, "name": t.name, "market": t.market},
            )
    return len(SEED_TICKERS)


def seed_symbols() -> SeedReport:
    """KOSPI + KOSDAQ 전 종목을 upsert. SEED 10 종목은 is_seed=TRUE.

    순서:
      1) SEED_TICKERS 먼저 별도 트랜잭션으로 강제 upsert (KRX 실패와 무관하게 검색 가능)
      2) KRX 리스팅 fetch (네트워크 의존) → 별도 트랜잭션으로 일괄 upsert.
         시장별 fetch 실패 시 해당 시장만 스킵하고 나머지 진행.
    """
    # 1) SEED_TICKERS — 항상 보장
    try:
        _upsert_seed_tickers()
        seed_marked = len(SEED_TICKERS)
        logger.info("seed_symbols: seed-tickers upserted (%d)", seed_marked)
    except Exception:  # noqa: BLE001
        logger.exception("seed_symbols: seed-tickers upsert failed (critical)")
        seed_marked = 0

    # 2) KRX 전 종목 — fetch 실패해도 부분 성공 허용
    market_counts: dict[str, int] = {}
    all_rows: list[tuple[str, str, str]] = []
    for market in ("KOSPI", "KOSDAQ"):
        try:
            listing = _fetch_market_listing(market)
            market_counts[market] = len(listing)
            for code, name in listing:
                all_rows.append((code, name, market))
            logger.info("seed_symbols: KRX %s fetched (%d)", market, len(listing))
        except Exception:  # noqa: BLE001
            logger.exception("seed_symbols: KRX %s fetch failed — skip market", market)
            market_counts[market] = 0

    inserted = updated = 0
    if all_rows:
        engine = get_engine()
        try:
            with engine.begin() as conn:
                for code, name, market in all_rows:
                    is_seed = code in SEED_CODES
                    res = conn.execute(
                        text(
                            """
                            INSERT INTO symbols (code, name, market, is_seed)
                            VALUES (:code, :name, :market, :is_seed)
                            ON CONFLICT (code) DO UPDATE
                              SET name    = EXCLUDED.name,
                                  market  = EXCLUDED.market,
                                  is_seed = symbols.is_seed OR EXCLUDED.is_seed
                            RETURNING (xmax = 0) AS inserted
                            """
                        ),
                        {"code": code, "name": name, "market": market, "is_seed": is_seed},
                    )
                    row = res.first()
                    if row and row[0]:
                        inserted += 1
                    else:
                        updated += 1
        except Exception:  # noqa: BLE001
            logger.exception("seed_symbols: KRX bulk upsert failed (transaction rolled back)")

    logger.info(
        "seed_symbols done: inserted=%d updated=%d seed_marked=%d markets=%s",
        inserted, updated, seed_marked, market_counts,
    )
    return SeedReport(inserted=inserted, updated=updated, seed_marked=seed_marked, markets=market_counts)