penti/backend/app/tasks/crawlers/prices.py

"""Stock price data crawler (주가 데이터 수집)."""
import time
from datetime import date, datetime, timedelta
from io import BytesIO
from typing import List, Optional

import pandas as pd
import requests as rq
from tqdm import tqdm
from sqlalchemy.orm import Session
from sqlalchemy import func

from app.models.asset import Asset
from app.models.price import PriceData


def get_price_data_from_naver(
    ticker: str,
    start_date: str,
    end_date: str
) -> Optional[pd.DataFrame]:
    """
    Naver에서 주가 데이터 다운로드.

    Args:
        ticker: 종목코드
        start_date: 시작일 (YYYYMMDD)
        end_date: 종료일 (YYYYMMDD)

    Returns:
        주가 DataFrame (실패 시 None)
    """
    try:
        url = f'''https://fchart.stock.naver.com/siseJson.nhn?symbol={ticker}&requestType=1&startTime={start_date}&endTime={end_date}&timeframe=day'''

        # 데이터 다운로드
        data = rq.get(url, timeout=30).content
        data_price = pd.read_csv(BytesIO(data))

        # 데이터 클렌징
        price = data_price.iloc[:, 0:6]
        price.columns = ['날짜', '시가', '고가', '저가', '종가', '거래량']
        price = price.dropna()
        price['날짜'] = price['날짜'].str.extract("(\d+)")
        price['날짜'] = pd.to_datetime(price['날짜'])
        price['종목코드'] = ticker

        return price

    except Exception as e:
        print(f"종목 {ticker} 가격 데이터 다운로드 오류: {e}")
        return None


def process_price_data(
    db_session: Session,
    tickers: Optional[List[str]] = None,
    start_date: Optional[str] = None,
    sleep_time: float = 0.5
) -> dict:
    """
    주가 데이터 수집 및 저장.

    Args:
        db_session: 데이터베이스 세션
        tickers: 종목코드 리스트 (None이면 전체 종목)
        start_date: 시작일 (YYYYMMDD, None이면 최근 저장 날짜 다음날)
        sleep_time: 요청 간격 (초)

    Returns:
        {'success': 성공 종목 수, 'failed': 실패 종목 리스트}
    """
    # 종목 리스트 조회
    if tickers is None:
        assets = db_session.query(Asset).filter(
            Asset.is_active == True,
            Asset.stock_type == '보통주'  # 보통주만 조회
        ).all()
        tickers = [asset.ticker for asset in assets]
        print(f"전체 {len(tickers)}개 종목 주가 수집 시작")
    else:
        print(f"{len(tickers)}개 종목 주가 수집 시작")

    # 종료일 (오늘)
    end_date = date.today().strftime("%Y%m%d")

    # 결과 추적
    success_count = 0
    error_list = []

    # 전종목 주가 다운로드 및 저장
    for ticker in tqdm(tickers):
        try:
            # 최근 저장 날짜 조회
            latest_record = db_session.query(
                func.max(PriceData.timestamp)
            ).filter(
                PriceData.ticker == ticker
            ).scalar()

            if latest_record and start_date is None:
                # 최근 날짜 다음날부터
                from_date = (latest_record.date() + timedelta(days=1)).strftime("%Y%m%d")
            elif start_date:
                from_date = start_date
            else:
                # 기본값: 1년 전부터
                from_date = (date.today() - timedelta(days=365)).strftime("%Y%m%d")

            # 이미 최신 상태면 스킵
            if from_date >= end_date:
                continue

            # Naver에서 데이터 다운로드
            price_df = get_price_data_from_naver(ticker, from_date, end_date)

            if price_df is None or price_df.empty:
                continue

            # 데이터베이스 저장
            save_price_to_db(price_df, db_session)
            success_count += 1

        except Exception as e:
            print(f"종목 {ticker} 처리 오류: {e}")
            error_list.append(ticker)

        # 요청 간격
        time.sleep(sleep_time)

    print(f"\n주가 수집 완료: 성공 {success_count}개, 실패 {len(error_list)}개")
    if error_list:
        print(f"실패 종목: {error_list[:10]}...")  # 처음 10개만 출력

    return {
        'success': success_count,
        'failed': error_list
    }


def save_price_to_db(price_df: pd.DataFrame, db_session: Session):
    """
    주가 데이터를 PostgreSQL에 저장 (UPSERT).

    Args:
        price_df: 주가 DataFrame
        db_session: 데이터베이스 세션
    """
    for _, row in price_df.iterrows():
        # 기존 레코드 조회
        existing = db_session.query(PriceData).filter(
            PriceData.ticker == row['종목코드'],
            PriceData.timestamp == row['날짜']
        ).first()

        if existing:
            # 업데이트
            existing.open = row['시가'] if row['시가'] else None
            existing.high = row['고가'] if row['고가'] else None
            existing.low = row['저가'] if row['저가'] else None
            existing.close = row['종가']
            existing.volume = int(row['거래량']) if row['거래량'] else None
        else:
            # 신규 삽입
            price_data = PriceData(
                ticker=row['종목코드'],
                timestamp=row['날짜'],
                open=row['시가'] if row['시가'] else None,
                high=row['고가'] if row['고가'] else None,
                low=row['저가'] if row['저가'] else None,
                close=row['종가'],
                volume=int(row['거래량']) if row['거래량'] else None
            )
            db_session.add(price_data)

    db_session.commit()


def update_recent_prices(
    db_session: Session,
    days: int = 30,
    sleep_time: float = 0.5
) -> dict:
    """
    최근 N일 주가 데이터 업데이트.

    Args:
        db_session: 데이터베이스 세션
        days: 최근 N일
        sleep_time: 요청 간격 (초)

    Returns:
        {'success': 성공 종목 수, 'failed': 실패 종목 리스트}
    """
    start_date = (date.today() - timedelta(days=days)).strftime("%Y%m%d")
    return process_price_data(db_session, start_date=start_date, sleep_time=sleep_time)