signal-platform/scripts/extend_snapshot_universe.py

"""Extend a *copy* of the production backtest snapshot with broad-universe OHLCV.

Research only — never writes to production Postgres.

Pipeline
--------
1. Copy ``--source`` snapshot (default ``backtest_snapshots/prod.sqlite``) to
   ``--output`` (default ``backtest_snapshots/research.sqlite``).
2. Resolve symbol pool = nasdaq_all ∪ sp500 via ``ticker_universe_service``.
3. Fetch ~5y daily bars from Alpaca for symbols missing (or short) in the copy.
4. Insert new tickers + OHLCV; mark them in side table ``research_rank_only``
   so the harness can feed signal IC without GTL/candidate replay.

Resume-friendly: re-running skips symbols that already have ≥ ``--min-bars``.

Example
-------
  python scripts/extend_snapshot_universe.py \\
      --source backtest_snapshots/prod.sqlite \\
      --output backtest_snapshots/research.sqlite \\
      --force-copy

  # smoke: first 50 missing symbols only
  python scripts/extend_snapshot_universe.py --limit 50
"""

from __future__ import annotations

import argparse
import asyncio
import shutil
import sys
import time
from datetime import date, datetime, timedelta, timezone
from pathlib import Path

from sqlalchemy import create_engine, text

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))


def _parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument(
        "--source",
        default="backtest_snapshots/prod.sqlite",
        help="Existing prod snapshot to copy (read-only after copy).",
    )
    p.add_argument(
        "--output",
        default="backtest_snapshots/research.sqlite",
        help="Research snapshot path (created/updated).",
    )
    p.add_argument(
        "--force-copy",
        action="store_true",
        help="Overwrite output by re-copying from source first.",
    )
    p.add_argument(
        "--history-days",
        type=int,
        default=1825,
        help="OHLCV lookback days (~5y). Default 1825.",
    )
    p.add_argument(
        "--min-bars",
        type=int,
        default=260,
        help="Skip re-fetch when a symbol already has this many bars.",
    )
    p.add_argument(
        "--limit",
        type=int,
        default=None,
        help="Max *new* symbols to fetch (smoke tests).",
    )
    p.add_argument(
        "--sleep",
        type=float,
        default=0.15,
        help="Seconds between Alpaca symbol requests (rate-limit cushion).",
    )
    p.add_argument(
        "--max-retries",
        type=int,
        default=5,
        help="Retries per symbol on RateLimitError.",
    )
    p.add_argument("--quiet", action="store_true")
    return p.parse_args()


def _ensure_rank_only_table(engine) -> None:
    """DDL in its own connection/transaction (don't share with ORM Session)."""
    with engine.begin() as conn:
        conn.execute(
            text(
                """
                CREATE TABLE IF NOT EXISTS research_rank_only (
                    ticker_id INTEGER PRIMARY KEY,
                    symbol TEXT NOT NULL UNIQUE
                )
                """
            )
        )


async def _resolve_pool() -> tuple[list[str], dict[str, str]]:
    """Return sorted unique symbols and source labels.

    Offline-safe: does **not** use production Postgres or SystemSetting cache
    (those require a schema). Public sources first, then FMP, then seeds.
    """
    from app.services.ticker_universe_service import (
        _SEED_UNIVERSES,
        _fetch_universe_symbols_from_fmp,
        _fetch_universe_symbols_from_public,
        _normalise_symbols,
    )

    sources: dict[str, str] = {}
    symbols: set[str] = set()

    for universe in ("nasdaq_all", "sp500"):
        cleaned: list[str] = []
        src = "none"

        public_symbols, public_failures, public_source = (
            await _fetch_universe_symbols_from_public(universe)
        )
        cleaned = _normalise_symbols(public_symbols)
        if cleaned:
            src = public_source or "public"
        else:
            if public_failures:
                print(
                    f"  WARNING: public fetch {universe}: "
                    f"{'; '.join(public_failures[:3])}"
                )
            try:
                fmp_symbols = await _fetch_universe_symbols_from_fmp(universe)
                cleaned = _normalise_symbols(fmp_symbols)
                if cleaned:
                    src = "fmp"
            except Exception as exc:
                print(f"  WARNING: FMP fetch {universe}: {exc}")

        if not cleaned:
            cleaned = _normalise_symbols(_SEED_UNIVERSES.get(universe, []))
            if cleaned:
                src = "seed"
                print(
                    f"  WARNING: {universe} fell back to seed list "
                    f"({len(cleaned)} symbols) — not full universe"
                )

        if not cleaned:
            print(f"  WARNING: universe {universe} returned no symbols")
            continue

        sources[universe] = src
        symbols.update(cleaned)
        print(f"  {universe}: {len(cleaned)} symbols (source={src})")

    return sorted(symbols), sources


async def _fetch_symbol_bars(
    provider,
    symbol: str,
    start: date,
    end: date,
    *,
    max_retries: int,
    sleep_s: float,
) -> list:
    from app.exceptions import ProviderError, RateLimitError

    for attempt in range(max_retries):
        try:
            bars = await provider.fetch_ohlcv(symbol, start, end)
            if sleep_s > 0:
                await asyncio.sleep(sleep_s)
            return bars
        except RateLimitError:
            wait = min(60.0, 2.0 ** attempt)
            print(f"  rate limited on {symbol}; sleep {wait:.0f}s")
            await asyncio.sleep(wait)
        except ProviderError as exc:
            if attempt + 1 >= max_retries:
                raise
            await asyncio.sleep(1.0)
            _ = exc
    return []


async def _main() -> None:
    args = _parse_args()
    source = Path(args.source)
    output = Path(args.output)
    if not source.exists():
        raise SystemExit(f"Source snapshot not found: {source}")

    if args.force_copy or not output.exists():
        output.parent.mkdir(parents=True, exist_ok=True)
        if output.exists():
            output.unlink()
        print(f"Copying {source} → {output}")
        shutil.copy2(source, output)
    else:
        print(f"Updating existing research snapshot: {output}")

    from app.config import settings
    from app.providers.alpaca import AlpacaOHLCVProvider

    if not settings.alpaca_api_key or not settings.alpaca_api_secret:
        raise SystemExit("ALPACA_API_KEY / ALPACA_API_SECRET required in .env")

    provider = AlpacaOHLCVProvider(settings.alpaca_api_key, settings.alpaca_api_secret)
    end = date.today()
    start = end - timedelta(days=int(args.history_days))

    print("Resolving universe pool (nasdaq_all ∪ sp500)…")
    pool, sources = await _resolve_pool()
    print(f"Pool size: {len(pool)} (sources={sources})")

    # Sync sqlite via raw SQL — one short transaction per symbol so a failed
    # write never leaves the session in "transaction is inactive".
    engine = create_engine(
        f"sqlite:///{output.resolve().as_posix()}",
        future=True,
    )
    _ensure_rank_only_table(engine)

    with engine.connect() as conn:
        existing_rows = conn.execute(
            text("SELECT id, symbol FROM tickers")
        ).fetchall()
        existing_ids = {str(sym): int(tid) for tid, sym in existing_rows}
        prod_symbols = set(existing_ids)

        bar_counts: dict[str, int] = {}
        for sym, tid in existing_ids.items():
            n = conn.execute(
                text("SELECT COUNT(*) FROM ohlcv_records WHERE ticker_id = :tid"),
                {"tid": tid},
            ).scalar_one()
            bar_counts[sym] = int(n)

        to_fetch: list[str] = []
        for sym in pool:
            if sym in existing_ids and bar_counts.get(sym, 0) >= args.min_bars:
                continue
            to_fetch.append(sym)

        if args.limit is not None:
            to_fetch = to_fetch[: max(0, int(args.limit))]

        print(f"Symbols to fetch/extend: {len(to_fetch)}")
        ok = 0
        fail = 0
        t0 = time.monotonic()

        insert_ohlcv = text(
            """
            INSERT INTO ohlcv_records
                (ticker_id, date, open, high, low, close, volume, created_at)
            VALUES
                (:ticker_id, :date, :open, :high, :low, :close, :volume, :created_at)
            """
        )

        for index, sym in enumerate(to_fetch, 1):
            try:
                bars = await _fetch_symbol_bars(
                    provider,
                    sym,
                    start,
                    end,
                    max_retries=args.max_retries,
                    sleep_s=args.sleep,
                )
            except Exception as exc:
                fail += 1
                if not args.quiet:
                    print(f"  [{index}/{len(to_fetch)}] {sym} FAIL {exc}")
                continue

            if not bars:
                fail += 1
                if not args.quiet:
                    print(f"  [{index}/{len(to_fetch)}] {sym} empty")
                continue

            try:
                with engine.begin() as write:
                    ticker_id = existing_ids.get(sym)
                    is_new = ticker_id is None
                    if is_new:
                        write.execute(
                            text(
                                "INSERT INTO tickers (symbol, name, created_at) "
                                "VALUES (:sym, NULL, :created)"
                            ),
                            {
                                "sym": sym,
                                "created": datetime.now(timezone.utc).isoformat(),
                            },
                        )
                        ticker_id = int(
                            write.execute(
                                text("SELECT id FROM tickers WHERE symbol = :sym"),
                                {"sym": sym},
                            ).scalar_one()
                        )
                        existing_ids[sym] = ticker_id

                    write.execute(
                        text(
                            "DELETE FROM ohlcv_records WHERE ticker_id = :tid "
                            "AND date >= :start AND date <= :end"
                        ),
                        {
                            "tid": ticker_id,
                            "start": start.isoformat(),
                            "end": end.isoformat(),
                        },
                    )
                    now = datetime.now(timezone.utc).replace(tzinfo=None)
                    write.execute(
                        insert_ohlcv,
                        [
                            {
                                "ticker_id": ticker_id,
                                "date": b.date.isoformat(),
                                "open": float(b.open),
                                "high": float(b.high),
                                "low": float(b.low),
                                "close": float(b.close),
                                "volume": int(b.volume),
                                "created_at": now.isoformat(),
                            }
                            for b in bars
                        ],
                    )
                    if is_new:
                        write.execute(
                            text(
                                "INSERT OR REPLACE INTO research_rank_only "
                                "(ticker_id, symbol) VALUES (:tid, :sym)"
                            ),
                            {"tid": ticker_id, "sym": sym},
                        )
            except Exception as exc:
                fail += 1
                if not args.quiet:
                    print(f"  [{index}/{len(to_fetch)}] {sym} WRITE FAIL {exc}")
                continue

            ok += 1
            if not args.quiet and (index % 25 == 0 or index == len(to_fetch)):
                elapsed = time.monotonic() - t0
                print(
                    f"  progress {index}/{len(to_fetch)} ok={ok} fail={fail} "
                    f"elapsed={elapsed/60:.1f}m last={sym} bars={len(bars)}"
                )

        rank_only_n = conn.execute(
            text("SELECT COUNT(*) FROM research_rank_only")
        ).scalar_one()
        ticker_n = conn.execute(text("SELECT COUNT(*) FROM tickers")).scalar_one()
        ohlcv_n = conn.execute(
            text("SELECT COUNT(*) FROM ohlcv_records")
        ).scalar_one()

    print("Done.")
    print(f"  output: {output}")
    print(f"  tickers: {ticker_n}")
    print(f"  ohlcv rows: {ohlcv_n}")
    print(f"  research_rank_only: {rank_only_n}")
    print(f"  fetched ok/fail: {ok}/{fail}")


if __name__ == "__main__":
    asyncio.run(_main())