momentum-system/youtube_claims/recap/store.py

"""Postgres persistence for the recap loop (transcripts schema on the paper/research DB).
Idempotent: DDL is IF NOT EXISTS; ingest upserts the video row and rewrites its child rows in
one transaction, keyed by video_id."""

from __future__ import annotations

import json
from contextlib import contextmanager
from datetime import datetime, timezone
from pathlib import Path

import psycopg2
import psycopg2.extras

from youtube_claims.recap import settings

S = settings.PG_SCHEMA


@contextmanager
def connect():
    conn = psycopg2.connect(settings.database_url(), connect_timeout=10)
    try:
        yield conn
    finally:
        conn.close()


def apply_schema() -> None:
    ddl = (Path(__file__).parent / "schema.sql").read_text(encoding="utf-8").replace("{SCHEMA}", S)
    with connect() as conn:
        with conn.cursor() as cur:
            cur.execute(ddl)
        conn.commit()


def video_status() -> dict[str, dict]:
    with connect() as conn, conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
        cur.execute(f"SELECT video_id, status, extractor, upload_date, attempts FROM {S}.recap_videos")
        return {r["video_id"]: dict(r) for r in cur.fetchall()}


def _date(s):
    if not s:
        return None
    s = str(s)
    return f"{s[:4]}-{s[4:6]}-{s[6:]}" if len(s) == 8 and s.isdigit() else s


def _j(x):
    return json.dumps(x, default=str) if x is not None else None


def mark(video_id: str, status: str, *, error: str | None = None, meta: dict | None = None) -> None:
    """Record a video's status without touching its extraction (used for failures/skips)."""
    meta = meta or {}
    with connect() as conn, conn.cursor() as cur:
        cur.execute(
            f"INSERT INTO {S}.recap_videos (video_id, upload_date, title, duration_s, status, last_error, attempts) "
            f"VALUES (%s,%s,%s,%s,%s,%s,%s) ON CONFLICT (video_id) DO UPDATE SET status=EXCLUDED.status, "
            f"last_error=EXCLUDED.last_error, "
            f"attempts={S}.recap_videos.attempts + (CASE WHEN EXCLUDED.status='failed' THEN 1 ELSE 0 END), title=COALESCE(EXCLUDED.title, {S}.recap_videos.title), "
            f"upload_date=COALESCE(EXCLUDED.upload_date, {S}.recap_videos.upload_date), updated_at=now()",
            (video_id, _date(meta.get("upload_date")), meta.get("title"),
             int(meta["duration"]) if str(meta.get("duration") or "").isdigit() else None, status,
             (error or "")[:2000] or None, 1 if status == "failed" else 0))
        conn.commit()


def upsert_video(rec: dict, meta: dict, *, extractor: str, transcript_source: str | None,
                 run_meta: dict | None, trades: list[dict], passed: list[dict]) -> None:
    """rec: the extraction record; trades/passed: rec items enriched with symbol/resolve/our_view
    (and entry_time_tape for trades)."""
    vid = rec["video_id"]
    sd = _date(rec.get("session_date")) or _date(rec.get("upload_date"))
    day = rec.get("day") or {}
    with connect() as conn:
        with conn.cursor() as cur:
            cur.execute(
                f"""INSERT INTO {S}.recap_videos (video_id, upload_date, session_date, title, duration_s,
                    video_kind, transcript_source, extractor, status, last_error, run_meta, record, processed_at)
                    VALUES (%s,%s,%s,%s,%s,%s,%s,%s,'done',NULL,%s,%s,%s)
                    ON CONFLICT (video_id) DO UPDATE SET upload_date=EXCLUDED.upload_date,
                    session_date=EXCLUDED.session_date, title=EXCLUDED.title, duration_s=EXCLUDED.duration_s,
                    video_kind=EXCLUDED.video_kind, transcript_source=EXCLUDED.transcript_source,
                    extractor=EXCLUDED.extractor, status='done', last_error=NULL, attempts=0, run_meta=EXCLUDED.run_meta,
                    record=EXCLUDED.record, processed_at=EXCLUDED.processed_at, updated_at=now()""",
                (vid, _date(rec.get("upload_date")), sd, meta.get("title"),
                 int(meta["duration"]) if str(meta.get("duration") or "").isdigit() else None,
                 rec.get("video_kind"), transcript_source, extractor, _j(run_meta), _j(rec),
                 datetime.now(timezone.utc)))
            for t in ("recap_day_summary", "recap_trades", "recap_passed_on", "recap_rules"):
                cur.execute(f"DELETE FROM {S}.{t} WHERE video_id=%s", (vid,))
            cur.execute(
                f"""INSERT INTO {S}.recap_day_summary (video_id, session_date, net_pnl, green, market_temperature,
                    market_quote, max_loss_hit, daily_goal_comment, sizing_comment)
                    VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s)""",
                (vid, sd, day.get("net_pnl"), day.get("green"), day.get("market_temperature"),
                 day.get("market_quote"), day.get("max_loss_hit"), day.get("daily_goal_comment"),
                 day.get("sizing_comment")))
            for i, t in enumerate(trades):
                cur.execute(
                    f"""INSERT INTO {S}.recap_trades (video_id, seq, session_date, ticker_heard, ticker_guess, symbol,
                        resolve, setup, setup_quote, entry_price, exit_price, stated_high, stated_low, shares, pnl,
                        outcome, add_on, entry_time_hint, entry_time_tape, stock_context, why_entered, tape_reads,
                        why_exited, mistake_or_lesson, our_view)
                        VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)""",
                    (vid, i, sd, t.get("ticker_heard"), t.get("ticker_guess"), t.get("symbol"), t.get("resolve"),
                     t.get("setup"), t.get("setup_quote"), t.get("entry_price"), t.get("exit_price"),
                     t.get("stated_high"), t.get("stated_low"), t.get("shares"), t.get("pnl"), t.get("outcome"),
                     bool(t.get("add_on")), t.get("entry_time_hint"), t.get("entry_time_tape"),
                     _j(t.get("stock_context")), _j(t.get("why_entered")), _j(t.get("tape_reads")),
                     _j(t.get("why_exited")), _j(t.get("mistake_or_lesson")), _j(t.get("our_view"))))
            for i, p in enumerate(passed):
                cur.execute(
                    f"""INSERT INTO {S}.recap_passed_on (video_id, seq, session_date, ticker_heard, ticker_guess,
                        symbol, resolve, reason, quote, our_view) VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)""",
                    (vid, i, sd, p.get("ticker_heard"), p.get("ticker_guess"), p.get("symbol"), p.get("resolve"),
                     p.get("reason"), p.get("quote"), _j(p.get("our_view"))))
            for i, r in enumerate(rec.get("rules_stated") or []):
                if r.get("quote"):
                    cur.execute(f"INSERT INTO {S}.recap_rules (video_id, seq, rule, quote) VALUES (%s,%s,%s,%s)",
                                (vid, i, r.get("rule"), r["quote"]))
        conn.commit()


def load_rows() -> tuple[list[dict], list[dict], list[dict]]:
    """(videos, trades, passed) for the report -- every processed video."""
    with connect() as conn, conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
        cur.execute(f"""SELECT v.video_id, v.upload_date, v.session_date, v.title, v.video_kind, v.extractor,
                        d.net_pnl, d.green, d.market_temperature, d.market_quote
                        FROM {S}.recap_videos v LEFT JOIN {S}.recap_day_summary d USING (video_id)
                        WHERE v.status='done' ORDER BY v.upload_date, v.video_id""")
        videos = [dict(r) for r in cur.fetchall()]
        cur.execute(f"SELECT * FROM {S}.recap_trades ORDER BY session_date, video_id, seq")
        trades = [dict(r) for r in cur.fetchall()]
        cur.execute(f"SELECT * FROM {S}.recap_passed_on ORDER BY session_date, video_id, seq")
        passed = [dict(r) for r in cur.fetchall()]
    return videos, trades, passed