stock_chart_site/backend/app/main.py

from __future__ import annotations

import logging
import os
from contextlib import asynccontextmanager

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from sqlalchemy import text

from app.api.chart import router as chart_router
from app.api.markets import router as markets_router
from app.api.metrics import router as metrics_router
from app.api.news import router as news_router
from app.api.predict import router as predict_router
from app.api.refresh import router as refresh_router
from app.api.symbols import router as symbols_router
from app.config import settings
from app.db.connection import get_engine, ping as db_ping
from app.fetch import dart as dart_mod
from app.fetch import kis as kis_mod
from app.pipelines.scheduler import shutdown_scheduler, start_scheduler

logging.basicConfig(
    level=settings.log_level,
    format="%(asctime)s %(levelname)s %(name)s: %(message)s",
)
logger = logging.getLogger(__name__)


def _bootstrap_db() -> None:
    """첫 부팅 자동화:
      1) migrations/*.sql idempotent 적용 (timescale/pgvector 확장 + 스키마)
      2) symbols 테이블 비어있으면 pykrx 로 전 종목 시드 (SEED 10 마크 포함)

    BOOTSTRAP_DISABLED=1 이면 스킵 (테스트/CI 용). 어떤 단계든 실패해도 서버는
    뜬다 — /health/db 가 진단을 알려준다.
    """
    if os.environ.get("BOOTSTRAP_DISABLED") == "1":
        logger.info("bootstrap skipped (BOOTSTRAP_DISABLED=1)")
        return

    # 1) migrations
    try:
        from app.db.migrate import apply_all
        res = apply_all()
        logger.info("bootstrap migrate: %s", res)
    except Exception:  # noqa: BLE001
        logger.exception("bootstrap migrate failed")
        return  # 스키마 없으면 시드 불가

    # 2) symbols 시드
    #   - SEED 10종목은 매 부팅마다 무조건 upsert (10회 upsert, ms 단위, 네트워크 무관)
    #     → KRX 접근 실패한 환경에서도 최소 10종목 검색 보장
    #   - KRX 전 종목 fetch 는 symbols 가 비어있을 때만 (호출 비용 큼)
    try:
        from app.fetch.symbols_seed import _upsert_seed_tickers, seed_symbols
        n_seed = _upsert_seed_tickers()
        logger.info("bootstrap seed-tickers ensured (%d)", n_seed)

        eng = get_engine()
        with eng.connect() as conn:
            row = conn.execute(text("SELECT COUNT(*) FROM symbols")).first()
            count = int(row[0]) if row else 0
        if count <= n_seed:
            # symbols 가 SEED 만큼 또는 그 이하 → KRX 전 종목 fetch 시도
            logger.info("symbols sparse (count=%d) — running KRX listing seed", count)
            report = seed_symbols()
            logger.info("bootstrap seed_symbols: %s", report)
        else:
            logger.info("symbols already populated (count=%d) — skip KRX listing seed", count)
    except Exception:  # noqa: BLE001
        logger.exception("bootstrap seed_symbols failed")


@asynccontextmanager
async def lifespan(_: FastAPI):
    _bootstrap_db()
    # 스케줄러는 옵션. CI/테스트에서 disable 하고 싶으면 SCHEDULER_DISABLED 같은 env 추가 가능.
    if os.environ.get("SCHEDULER_DISABLED") == "1":
        logger.info("scheduler skipped (SCHEDULER_DISABLED=1)")
    else:
        try:
            start_scheduler()
        except Exception:  # noqa: BLE001
            logger.exception("scheduler start failed")
    yield
    shutdown_scheduler()


app = FastAPI(title="stock_chart_site", version="0.1.0", lifespan=lifespan)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

app.include_router(refresh_router)
app.include_router(symbols_router)
app.include_router(chart_router)
app.include_router(predict_router)
app.include_router(metrics_router)
app.include_router(news_router)
app.include_router(markets_router)


def _resolved_device() -> str:
    if settings.model_device != "auto":
        return settings.model_device
    try:
        import torch  # noqa: WPS433
        return "cuda" if torch.cuda.is_available() else "cpu"
    except Exception:  # noqa: BLE001
        return "cpu"


@app.get("/health")
def health() -> dict[str, object]:
    return {"ok": True, "device": _resolved_device(), "version": "0.1.0"}


@app.get("/health/db")
def health_db() -> dict[str, object]:
    return {"ok": True, **db_ping()}


@app.get("/health/keys")
def health_keys() -> dict[str, object]:
    """등록된 외부 키들 ping (key 값은 노출하지 않음)."""
    return {
        "kis": kis_mod.ping(),
        "dart": dart_mod.ping(),
        # huggingface 는 모델 다운로드 시점에 확인 (별도 ping 호출 비용 회피)
    }


@app.get("/health/models")
def health_models() -> dict[str, object]:
    """Chronos / LGBM 가용성 진단.

    Chronos: lazy 로드 첫 호출이라 30초~수 분 걸릴 수 있음 (HuggingFace 다운로드).
    LGBM: 체크포인트 디렉토리 스캔 — retrain 안 돈 cold start 에선 비어있음.
    """
    from pathlib import Path

    from app.models import chronos as chronos_mod

    lgbm_dir = Path(os.environ.get("LGBM_MODEL_DIR", "/app/data/models"))
    lgbm_files: list[str] = []
    if lgbm_dir.exists():
        lgbm_files = sorted(p.name for p in lgbm_dir.glob("*.pkl"))

    return {
        "chronos": chronos_mod.ping(),
        "lgbm": {
            "model_dir": str(lgbm_dir),
            "checkpoint_count": len(lgbm_files),
            "samples": lgbm_files[:8],  # 너무 많으면 잘라서.
            "status": "ok" if lgbm_files else "no_checkpoints (cold start, run retrain_weekly)",
        },
    }