"""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 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 - INTERVAL '24 hours' AND captured_at <= $2 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 market_row = await pool.fetchrow( _MARKET_SNAPSHOT_FALLBACK_SQL, ticker, 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()