stonks-oracle/scripts/backfill_snapshot_prices.py

"""Backfill price_at_prediction for existing NULL prediction snapshots.

One-time migration script that populates price_at_prediction using the
extended fallback chain:
  1. market_snapshots within 24h of generated_at for the ticker
  2. positions table (current_price) for the ticker

Run as: .venv/bin/python scripts/backfill_snapshot_prices.py
Dry run: .venv/bin/python scripts/backfill_snapshot_prices.py --dry-run

Requires env vars: POSTGRES_USER, POSTGRES_PASSWORD, POSTGRES_HOST,
                   POSTGRES_PORT, POSTGRES_DB

Requirements: 2.3
"""

import argparse
import asyncio
import os
import sys
from datetime import timedelta

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import asyncpg  # noqa: E402

from services.shared.config import load_config  # noqa: E402

# ---------------------------------------------------------------------------
# SQL Queries
# ---------------------------------------------------------------------------

_FIND_NULL_SNAPSHOTS_SQL = """
SELECT id, ticker, generated_at
FROM prediction_snapshots
WHERE price_at_prediction IS NULL
ORDER BY generated_at DESC
"""

_MARKET_SNAPSHOT_FALLBACK_SQL = """
SELECT (data->>'c')::float AS close
FROM market_snapshots
WHERE ticker = $1
  AND snapshot_type = 'bar'
  AND data->>'c' IS NOT NULL
  AND captured_at >= $2
  AND captured_at <= $3
ORDER BY captured_at DESC
LIMIT 1
"""

_POSITIONS_FALLBACK_SQL = """
SELECT current_price
FROM positions
WHERE ticker = $1
  AND current_price IS NOT NULL
LIMIT 1
"""

_UPDATE_PRICE_SQL = """
UPDATE prediction_snapshots
SET price_at_prediction = $1
WHERE id = $2
"""


# ---------------------------------------------------------------------------
# Main backfill logic
# ---------------------------------------------------------------------------


async def backfill(dry_run: bool = False) -> None:
    config = load_config()
    dsn = config.postgres.dsn

    pool = await asyncpg.create_pool(dsn=dsn)
    assert pool is not None

    # Find all snapshots with NULL price
    rows = await pool.fetch(_FIND_NULL_SNAPSHOTS_SQL)
    total = len(rows)

    if total == 0:
        print("No prediction snapshots with NULL price_at_prediction found.")
        await pool.close()
        return

    print(f"Found {total} snapshots with NULL price_at_prediction")
    if dry_run:
        print("[DRY RUN] No updates will be performed")
    print()

    # Statistics
    found_market = 0
    found_positions = 0
    still_null = 0

    for idx, row in enumerate(rows, start=1):
        snapshot_id = row["id"]
        ticker = row["ticker"]
        generated_at = row["generated_at"]

        price: float | None = None

        # Fallback 1: market_snapshots within 24h of generated_at
        window_start = generated_at - timedelta(hours=24)
        market_row = await pool.fetchrow(
            _MARKET_SNAPSHOT_FALLBACK_SQL, ticker, window_start, generated_at
        )
        if market_row and market_row["close"] is not None:
            price = float(market_row["close"])
            found_market += 1
        else:
            # Fallback 2: positions table
            pos_row = await pool.fetchrow(_POSITIONS_FALLBACK_SQL, ticker)
            if pos_row and pos_row["current_price"] is not None:
                price = float(pos_row["current_price"])
                found_positions += 1
            else:
                still_null += 1

        # Update if we found a price
        if price is not None and not dry_run:
            await pool.execute(_UPDATE_PRICE_SQL, price, snapshot_id)

        # Progress reporting every 100 snapshots
        if idx % 100 == 0:
            action = "checked" if dry_run else "processed"
            print(
                f"  {action} {idx}/{total} snapshots "
                f"(market: {found_market}, positions: {found_positions}, "
                f"null: {still_null})"
            )

    await pool.close()

    # Final statistics
    print()
    print("=" * 60)
    print("Backfill complete" if not dry_run else "Dry run complete")
    print("=" * 60)
    print(f"  Total snapshots processed: {total}")
    print(f"  Found via market_snapshots: {found_market}")
    print(f"  Found via positions:        {found_positions}")
    print(f"  Still NULL (no data):       {still_null}")
    if dry_run:
        updated = found_market + found_positions
        print(f"\n  [DRY RUN] Would have updated {updated} snapshots")
    print()


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Backfill price_at_prediction for NULL prediction snapshots"
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Report what would be updated without making changes",
    )
    args = parser.parse_args()

    try:
        asyncio.run(backfill(dry_run=args.dry_run))
    except Exception as e:
        print(f"Backfill failed: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()