fithia2/libs/db/migrations/versions/0002_phase3_tables.py

"""Phase 3 tables: llm_call_cache, review_items, event_labels.

Revision ID: 0002
Revises: 0001
Create Date: 2026-03-12

Tables created:
- llm_call_cache: DB-backed cache for LLM API calls
- review_items: human-in-the-loop review queue
- event_labels: forward-return labels for parsed events
"""
from __future__ import annotations

from collections.abc import Sequence

import sqlalchemy as sa
from alembic import op
from sqlalchemy.dialects.postgresql import JSONB, UUID

revision: str = "0002"
down_revision: str | None = "0001"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None


def upgrade() -> None:
    op.create_table(
        "llm_call_cache",
        sa.Column("cache_key", sa.Text, primary_key=True),
        sa.Column("document_id", sa.Text, nullable=True),
        sa.Column("model_name", sa.Text, nullable=False),
        sa.Column("prompt_version", sa.Text, nullable=False),
        sa.Column("schema_version", sa.Text, nullable=False),
        sa.Column("raw_prompt", sa.Text, nullable=False),
        sa.Column("raw_response", sa.Text, nullable=False),
        sa.Column("normalized_json", JSONB, nullable=False),
        sa.Column("token_usage_json", JSONB, nullable=False),
        sa.Column("elapsed_ms", sa.Integer, nullable=False),
        sa.Column(
            "created_at_utc",
            sa.DateTime(timezone=True),
            nullable=False,
            server_default=sa.text("now()"),
        ),
    )

    op.create_table(
        "review_items",
        sa.Column("review_id", UUID(as_uuid=True), primary_key=True),
        sa.Column("entity_type", sa.Text, nullable=False),
        sa.Column("entity_id", sa.Text, nullable=False),
        sa.Column("priority", sa.Text, nullable=False),
        sa.Column("reason_codes", JSONB, nullable=False),
        sa.Column("status", sa.Text, nullable=False, server_default=sa.text("'open'")),
        sa.Column("assigned_to", sa.Text, nullable=True),
        sa.Column("snapshot_refs", JSONB, nullable=False, server_default=sa.text("'{}'")),
        sa.Column("suggested_overrides", JSONB, nullable=True),
        sa.Column("reviewer_id", sa.Text, nullable=True),
        sa.Column("resolution_type", sa.Text, nullable=True),
        sa.Column("field_overrides", JSONB, nullable=True),
        sa.Column("root_cause", sa.Text, nullable=True),
        sa.Column("notes", sa.Text, nullable=True),
        sa.Column("resolved_at", sa.DateTime(timezone=True), nullable=True),
        sa.Column(
            "created_at_utc",
            sa.DateTime(timezone=True),
            nullable=False,
            server_default=sa.text("now()"),
        ),
        sa.Column(
            "updated_at_utc",
            sa.DateTime(timezone=True),
            nullable=False,
            server_default=sa.text("now()"),
        ),
    )
    op.create_index(
        "ix_review_items_entity_open", "review_items", ["entity_type", "entity_id", "status"]
    )
    op.create_index(
        "ix_review_items_status_priority", "review_items", ["status", "priority"]
    )
    # Partial unique index: only one open review per entity
    op.create_index(
        "uix_review_items_entity_open",
        "review_items",
        ["entity_type", "entity_id"],
        unique=True,
        postgresql_where=sa.text("status = 'open'"),
    )

    op.create_table(
        "event_labels",
        sa.Column("label_id", sa.BigInteger, primary_key=True, autoincrement=True),
        sa.Column(
            "event_id",
            sa.Text,
            sa.ForeignKey("events.event_id"),
            nullable=False,
        ),
        sa.Column("entry_convention", sa.Text, nullable=False),
        sa.Column("reaction_date", sa.Date, nullable=True),
        sa.Column("entry_date", sa.Date, nullable=True),
        sa.Column("entry_price", sa.Numeric, nullable=True),
        sa.Column("fwd_return_1d", sa.Numeric, nullable=True),
        sa.Column("fwd_return_3d", sa.Numeric, nullable=True),
        sa.Column("fwd_return_5d", sa.Numeric, nullable=True),
        sa.Column("hit_pos_1r_within_3d", sa.Boolean, nullable=True),
        sa.Column("hit_neg_1r_within_3d", sa.Boolean, nullable=True),
        sa.Column("close_up_after_3d", sa.Boolean, nullable=True),
        sa.Column("close_up_after_5d", sa.Boolean, nullable=True),
        sa.Column("mfe_3d", sa.Numeric, nullable=True),
        sa.Column("mae_3d", sa.Numeric, nullable=True),
        sa.Column("mfe_5d", sa.Numeric, nullable=True),
        sa.Column("mae_5d", sa.Numeric, nullable=True),
        sa.Column("bars_to_mfe_3d", sa.Integer, nullable=True),
        sa.Column("bars_to_mae_3d", sa.Integer, nullable=True),
        sa.Column("days_to_peak_close_5d", sa.Integer, nullable=True),
        sa.Column("label_status", sa.Text, nullable=False),
        sa.Column(
            "invalid_event_for_labeling",
            sa.Boolean,
            nullable=False,
            server_default=sa.text("false"),
        ),
        sa.Column("risk_model_name", sa.Text, nullable=True),
        sa.Column("initial_stop_price", sa.Numeric, nullable=True),
        sa.Column("initial_r_value", sa.Numeric, nullable=True),
        sa.Column("label_version", sa.Text, nullable=False),
        sa.Column(
            "created_at_utc",
            sa.DateTime(timezone=True),
            nullable=False,
            server_default=sa.text("now()"),
        ),
        sa.UniqueConstraint(
            "event_id",
            "entry_convention",
            "label_version",
            name="uq_event_labels_event_convention_version",
        ),
    )
    op.create_index("ix_event_labels_event_id", "event_labels", ["event_id"])


def downgrade() -> None:
    op.drop_index("ix_event_labels_event_id", table_name="event_labels")
    op.drop_table("event_labels")
    op.drop_index("uix_review_items_entity_open", table_name="review_items")
    op.drop_index("ix_review_items_status_priority", table_name="review_items")
    op.drop_index("ix_review_items_entity_open", table_name="review_items")
    op.drop_table("review_items")
    op.drop_table("llm_call_cache")