Stock-trading-programming/app/db/models.py

"""
db/models.py - SQLite 스키마 생성
기획서 v2.1 기준 4개 테이블
"""
import sqlite3
import os

DB_PATH = os.getenv("DB_PATH", "data/stockbot.db")

def _ensure_columns(cursor, table: str, columns: dict[str, str]):
    existing = {row[1] for row in cursor.execute(f"PRAGMA table_info({table})").fetchall()}
    for name, ddl in columns.items():
        if name not in existing:
            cursor.execute(f"ALTER TABLE {table} ADD COLUMN {name} {ddl}")

def init_db():
    os.makedirs(os.path.dirname(DB_PATH), exist_ok=True)
    conn = sqlite3.connect(DB_PATH)
    c = conn.cursor()

    # 체결 내역
    c.execute("""
    CREATE TABLE IF NOT EXISTS trades (
        id          INTEGER PRIMARY KEY AUTOINCREMENT,
        date        TEXT NOT NULL,
        ticker      TEXT NOT NULL,
        name        TEXT,
        entry_time  TEXT NOT NULL,
        exit_time   TEXT,
        entry_price REAL NOT NULL,
        exit_price  REAL,
        quantity    INTEGER NOT NULL,
        side        TEXT NOT NULL,
        exit_reason TEXT,
        pnl         REAL,
        fee         REAL,
        slippage    REAL,
        strategy    TEXT DEFAULT 'VB',
        ai_boosted  INTEGER DEFAULT 0
    )""")

    # 일일 요약
    c.execute("""
    CREATE TABLE IF NOT EXISTS daily_summary (
        date            TEXT PRIMARY KEY,
        total_trades    INTEGER DEFAULT 0,
        win_trades      INTEGER DEFAULT 0,
        lose_trades     INTEGER DEFAULT 0,
        gross_pnl       REAL DEFAULT 0,
        total_fee       REAL DEFAULT 0,
        net_pnl         REAL DEFAULT 0,
        max_drawdown    REAL DEFAULT 0,
        trading_stopped INTEGER DEFAULT 0
    )""")

    # 포지션 (장중 현황)
    c.execute("""
    CREATE TABLE IF NOT EXISTS positions (
        ticker       TEXT PRIMARY KEY,
        name         TEXT,
        entry_time   TEXT,
        entry_price  REAL,
        quantity     INTEGER,
        tp1_done     INTEGER DEFAULT 0,
        target_price REAL,
        stop_price   REAL,
        ai_boosted   INTEGER DEFAULT 0
    )""")

    # AI 판단 이력
    c.execute("""
    CREATE TABLE IF NOT EXISTS ai_context_log (
        date                TEXT PRIMARY KEY,
        generated_at        TEXT,
        trade_allowed       INTEGER,
        market_sentiment    TEXT,
        sentiment_score     INTEGER,
        risk_level          TEXT,
        hot_sectors         TEXT,
        avoid_sectors       TEXT,
        boosted_tickers     TEXT,
        blacklist_tickers   TEXT,
        position_size_mult  REAL,
        reason              TEXT,
        claude_tokens_used  INTEGER,
        api_call_success    INTEGER DEFAULT 1
    )""")

    c.execute("""
    CREATE TABLE IF NOT EXISTS entry_snapshots (
        id                       INTEGER PRIMARY KEY AUTOINCREMENT,
        trade_id                 INTEGER,
        date                     TEXT NOT NULL,
        ticker                   TEXT NOT NULL,
        name                     TEXT,
        entry_time               TEXT NOT NULL,
        current_price            REAL,
        entry_price              REAL,
        target_price             REAL,
        stop_price               REAL,
        today_open               REAL,
        prev_high                REAL,
        prev_low                 REAL,
        prev_amount              REAL,
        volume                   REAL,
        change_pct               REAL,
        market_sentiment         TEXT,
        sentiment_score          INTEGER,
        risk_level               TEXT,
        trade_allowed            INTEGER,
        hot_sectors              TEXT,
        avoid_sectors            TEXT,
        boosted_tickers          TEXT,
        blacklist_tickers        TEXT,
        ai_boosted               INTEGER DEFAULT 0,
        ai_win_score             REAL,
        ai_stop_loss_score       REAL,
        ai_model_version         TEXT,
        position_size_multiplier REAL,
        combined_multiplier      REAL,
        entry_reason             TEXT,
        strategy                 TEXT DEFAULT 'VB',
        created_at               TEXT NOT NULL
    )""")
    c.execute("CREATE INDEX IF NOT EXISTS idx_entry_snapshots_trade_id ON entry_snapshots(trade_id)")
    c.execute("CREATE INDEX IF NOT EXISTS idx_entry_snapshots_date_ticker ON entry_snapshots(date, ticker)")
    _ensure_columns(c, "entry_snapshots", {
        "ai_win_score": "REAL",
        "ai_stop_loss_score": "REAL",
        "ai_model_version": "TEXT",
    })

    c.execute("""
    CREATE TABLE IF NOT EXISTS post_entry_snapshots (
        id                 INTEGER PRIMARY KEY AUTOINCREMENT,
        trade_id           INTEGER,
        date               TEXT NOT NULL,
        ticker             TEXT NOT NULL,
        sample_time        TEXT NOT NULL,
        elapsed_sec        INTEGER NOT NULL,
        entry_price        REAL,
        current_price      REAL,
        return_pct         REAL,
        mfe_pct            REAL,
        mae_pct            REAL,
        volume             REAL,
        change_pct         REAL,
        position_open      INTEGER DEFAULT 1,
        created_at         TEXT NOT NULL,
        UNIQUE(trade_id, elapsed_sec)
    )""")
    c.execute("CREATE INDEX IF NOT EXISTS idx_post_entry_snapshots_trade_id ON post_entry_snapshots(trade_id)")
    c.execute("CREATE INDEX IF NOT EXISTS idx_post_entry_snapshots_date_ticker ON post_entry_snapshots(date, ticker)")

    conn.commit()
    conn.close()
    print(f"DB 초기화 완료: {DB_PATH}")

def get_conn():
    return sqlite3.connect(DB_PATH)