"""Backfill market_cap_proxy + exchange_proxy into existing FeatureSnapshot rows. Patches market_v1 FeatureSnapshots that are missing market_cap_proxy in their feature_json. This is needed for paper-trading backsim on events that were feature-built before the feature-builder started persisting company metadata. Usage: python -m apps.tools.backfill_market_cap python -m apps.tools.backfill_market_cap --start 2026-01-01 --end 2026-03-31 python -m apps.tools.backfill_market_cap --dry-run """ from __future__ import annotations import argparse import asyncio import datetime as dt import os from pathlib import Path _ENV_FILE = Path(__file__).parent.parent.parent / ".env" if _ENV_FILE.exists(): for _line in _ENV_FILE.read_text().splitlines(): _line = _line.strip() if _line and not _line.startswith("#") and "=" in _line: _k, _, _v = _line.partition("=") os.environ.setdefault(_k.strip(), _v.strip()) from rich.console import Console from rich.progress import track console = Console(width=120) async def run_backfill( start_date: dt.date | None, end_date: dt.date | None, dry_run: bool, ) -> None: from sqlalchemy import select, update from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine from libs.db.models import Event, FeatureSnapshot, SymbolMaster from libs.oracle_client import CompanyService, ScreenerService, make_oracle_client db_dsn = os.environ.get("POSTGRES_DSN", "") oracle_url = os.environ.get("STOCK_ORACLE_URL", "http://localhost:18001") if not db_dsn: console.print("[red]POSTGRES_DSN not set[/]") return engine = create_async_engine(db_dsn, echo=False) async_session = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False) # ── Step 1: find market_v1 snapshots without market_cap_proxy ────────── async with async_session() as session: stmt = ( select(FeatureSnapshot, Event, SymbolMaster) .join(Event, FeatureSnapshot.event_id == Event.event_id) .outerjoin(SymbolMaster, Event.symbol_id == SymbolMaster.symbol_id) .where(FeatureSnapshot.snapshot_name == "market_v1") ) if start_date: stmt = stmt.where(Event.event_date >= start_date) if end_date: stmt = stmt.where(Event.event_date <= end_date) rows = (await session.execute(stmt)).all() # Filter to those missing market_cap_proxy missing = [ (snap, event, sym) for snap, event, sym in rows if not (snap.feature_json or {}).get("market_cap_proxy") ] console.print(f"Found [bold]{len(rows)}[/] market_v1 snapshots, " f"[yellow]{len(missing)}[/] missing market_cap_proxy") if not missing or dry_run: if dry_run: console.print("[dim]Dry-run: no changes written.[/]") return # ── Step 2: collect unique tickers ───────────────────────────────────── ticker_to_snap: dict[str, list[FeatureSnapshot]] = {} for snap, event, sym in missing: if sym and sym.ticker: ticker_to_snap.setdefault(sym.ticker, []).append(snap) tickers = sorted(ticker_to_snap) console.print(f"Fetching company info for [bold]{len(tickers)}[/] unique tickers…") # ── Step 3: batch-fetch via screener first, then per-symbol CompanyService ─ ticker_mcap: dict[str, float | None] = {t: None for t in tickers} ticker_exchange: dict[str, str | None] = {t: None for t in tickers} async with make_oracle_client() as client: # Try screener first (batch — faster, one call per page) try: svc = ScreenerService(client) stocks = await svc.search_all_stocks( market_cap_min=500_000_000, exchange="NYSE,NASDAQ,AMEX", exclude_types="ETF,FUND,ADR,SPAC", ) screener_lookup = {(s.symbol or "").upper(): s for s in stocks} for t in tickers: s = screener_lookup.get(t.upper()) if s: ticker_mcap[t] = s.market_cap ticker_exchange[t] = s.exchange resolved = sum(1 for t in tickers if ticker_mcap.get(t) is not None) console.print(f" Screener resolved {resolved}/{len(tickers)}") except Exception as exc: console.print(f" [yellow]Screener failed ({exc}), falling back to CompanyService[/]") # Per-symbol CompanyService for anything still missing unresolved = [t for t in tickers if ticker_mcap.get(t) is None] if unresolved: company_svc = CompanyService(client) semaphore = asyncio.Semaphore(16) async def _fetch_one(ticker: str) -> None: async with semaphore: try: info = await company_svc.get_company(ticker) ticker_mcap[ticker] = info.market_cap ticker_exchange[ticker] = info.exchange except Exception: pass await asyncio.gather(*(_fetch_one(t) for t in unresolved)) resolved2 = sum(1 for t in unresolved if ticker_mcap.get(t) is not None) console.print(f" CompanyService resolved {resolved2}/{len(unresolved)} remaining") # ── Step 4: update feature_json ───────────────────────────────────────── updated = 0 skipped = 0 async with async_session() as session: for ticker, snaps in track(ticker_to_snap.items(), description="Updating…"): mcap = ticker_mcap.get(ticker) exch = ticker_exchange.get(ticker) if mcap is None and exch is None: skipped += len(snaps) continue for snap in snaps: fj = dict(snap.feature_json or {}) if mcap is not None: fj["market_cap_proxy"] = mcap if exch is not None: fj["exchange_proxy"] = exch await session.execute( update(FeatureSnapshot) .where(FeatureSnapshot.feature_snapshot_id == snap.feature_snapshot_id) .values(feature_json=fj) ) updated += 1 await session.commit() await engine.dispose() console.print(f"\n[bold green]Done.[/] Updated {updated} snapshots, skipped {skipped} (no data).") def main() -> None: parser = argparse.ArgumentParser(description="Backfill market_cap_proxy into FeatureSnapshots") parser.add_argument("--start", default=None, metavar="YYYY-MM-DD") parser.add_argument("--end", default=None, metavar="YYYY-MM-DD") parser.add_argument("--dry-run", action="store_true", help="Show what would be done without writing") args = parser.parse_args() start = dt.date.fromisoformat(args.start) if args.start else None end = dt.date.fromisoformat(args.end) if args.end else None asyncio.run(run_backfill(start, end, args.dry_run)) if __name__ == "__main__": main()