auto-stock/backend/alembic/versions/0001_initial_schema_columns.py

"""initial schema columns

lifespan 에서 관리하던 ADD COLUMN / CREATE TABLE DDL 을 Alembic 으로 이전.
이 revision 이 적용된 후에는 main.py lifespan 의 raw SQL DDL 블록은 제거됨.

Revision ID: 0001
Revises:
Create Date: 2026-02-27

"""
from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql

revision: str = "0001"
down_revision: Union[str, None] = None
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # ── trade_history 컬럼 추가 ─────────────────────────────────────────────
    with op.batch_alter_table("trade_history", schema=None) as batch_op:
        for col_name, col_type in [
            ("total_amount",     sa.Numeric(15, 2)),
            ("entry_reason",     sa.Text()),
            ("exit_reason",      sa.Text()),
            ("signals",          postgresql.JSONB()),
            ("entry_trade_id",   sa.Integer()),
            ("slippage",         sa.Numeric(10, 2)),
            ("commission",       sa.Numeric(10, 2)),
            ("net_profit_loss",  sa.Numeric(15, 2)),
            ("profit_loss_rate", sa.Numeric(10, 4)),
        ]:
            try:
                batch_op.add_column(sa.Column(col_name, col_type, nullable=True))
            except (sa.exc.OperationalError, sa.exc.ProgrammingError):
                pass  # 이미 존재하는 컬럼은 무시

    # total_amount 백필
    op.execute("""
        UPDATE trade_history
        SET total_amount = executed_price * executed_quantity
        WHERE total_amount IS NULL
          AND executed_price IS NOT NULL
          AND executed_quantity IS NOT NULL
    """)

    # ── stock_analysis 컬럼 추가 ────────────────────────────────────────────
    with op.batch_alter_table("stock_analysis", schema=None) as batch_op:
        for col_name, col_type in [
            ("atr_14",        sa.Numeric(15, 2)),
            ("adx_14",        sa.Numeric(10, 2)),
            ("ml_score",      sa.Integer()),
            ("ml_confidence", sa.Numeric(5, 4)),
            ("ml_signal",     sa.String(10)),
        ]:
            try:
                batch_op.add_column(sa.Column(col_name, col_type, nullable=True))
            except Exception:
                pass

    # stock_analysis 유니크 제약 추가 (idempotent)
    op.execute("""
        DO $$ BEGIN
            IF NOT EXISTS (
                SELECT 1 FROM pg_constraint
                WHERE conname = 'uq_stock_analysis_code_date'
                  AND conrelid = 'stock_analysis'::regclass
            ) THEN
                ALTER TABLE stock_analysis
                    ADD CONSTRAINT uq_stock_analysis_code_date
                    UNIQUE (stock_code, analysis_date);
            END IF;
        END $$;
    """)

    # ── active_trades 컬럼 추가 ─────────────────────────────────────────────
    with op.batch_alter_table("active_trades", schema=None) as batch_op:
        for col_name, col_type, default in [
            ("remaining_pct",        sa.Numeric(5, 2),  "1.0"),
            ("highest_since_entry",  sa.Numeric(15, 2), None),
            ("strategy_name",        sa.String(100),    None),
        ]:
            try:
                col = sa.Column(col_name, col_type, nullable=True,
                                server_default=default if default else sa.null())
                batch_op.add_column(col)
            except Exception:
                pass

    # ── trading_sessions 테이블 ─────────────────────────────────────────────
    op.execute("""
        CREATE TABLE IF NOT EXISTS trading_sessions (
            id                SERIAL PRIMARY KEY,
            session_date      DATE UNIQUE NOT NULL,
            consecutive_losses INTEGER DEFAULT 0,
            daily_profit_loss NUMERIC(15, 2) DEFAULT 0,
            is_paused         BOOLEAN DEFAULT FALSE,
            pause_reason      TEXT,
            created_at        TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            updated_at        TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    """)

    # ── stock_bans 테이블 ───────────────────────────────────────────────────
    op.execute("""
        CREATE TABLE IF NOT EXISTS stock_bans (
            stock_code         VARCHAR(20) PRIMARY KEY,
            consecutive_losses INTEGER DEFAULT 0 NOT NULL,
            ban_until          TIMESTAMP,
            updated_at         TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
    """)


def downgrade() -> None:
    # downgrade 는 데이터 손실 위험이 있으므로 명시적으로 구현하지 않음
    # 필요 시 ALTER TABLE DROP COLUMN 을 수동으로 실행
    pass