spb/hfmarketdata

1#!/usr/bin/env python32"""Convert the inline `fund_statements.coverage` JSON of existing rows to `fund_coverage_blob` references.34    cd ~/hfmarketdata && HFMD_DATA_ROOT=~/firstratedata venv/bin/python scripts/migrate_coverage.py            # all rows5    venv/bin/python scripts/migrate_coverage.py --batch 5000 --max-batches 50                                    # a slice6    venv/bin/python scripts/migrate_coverage.py --vacuum                                                         # reclaim space (offline!)78Runs while the API serves (WAL, short transactions of `--batch` rows; each batch commits, so it can be stopped9and resumed at any time: the work set is simply `WHERE coverage IS NOT NULL`). New rows are written through the10dictionary already (`ingest._replace_statements`), the read side understands both shapes, so this only shrinks the11past: 1.22 M rows × ~1.1 KB ≈ 1.3 GB of JSON → a few tens of thousands of blobs.1213Space is NOT returned to the OS by the conversion itself (free pages are reused by later inserts). `--vacuum`14runs `PRAGMA auto_vacuum=INCREMENTAL` + `VACUUM` afterwards: needs ~2× the file size of free disk, exclusive access15(stop the API and the jobs) and a few minutes for 3 GB — run it once, off-hours.1617Author: Simon-Pierre Boucher <contact@spboucher.ai>18"""19from __future__ import annotations2021import argparse22import json23import logging24import sys25import time26from pathlib import Path2728HERE = Path(__file__).resolve().parent29sys.path.insert(0, str(HERE.parent / "hfmarketdata" / "api"))3031log = logging.getLogger("migrate_coverage")323334def convert(batch: int, max_batches: int | None, sleep_s: float) -> tuple[int, int]:35    """Convert rows in id order; returns (rows converted, batches)."""36    from sqlalchemy import text3738    from core.db import engine39    from fundamentals import coverage_store as CS40    from fundamentals.models import init_db4142    init_db()   # makes sure coverage_id / fund_coverage_blob exist43    done = batches = 044    last_id = 045    t0 = time.time()46    while True:47        with engine.begin() as con:48            rows = con.execute(text("SELECT id, coverage FROM fund_statements WHERE id > :last AND coverage IS NOT NULL "49                                    "ORDER BY id LIMIT :n"), {"last": last_id, "n": batch}).all()50            if not rows:51                break52            docs = []53            for _, cov in rows:54                if isinstance(cov, (bytes, bytearray)):55                    cov = cov.decode()56                docs.append(json.loads(cov) if isinstance(cov, str) else (cov or {}))57            ids = CS.intern(docs)58            con.execute(text("UPDATE fund_statements SET coverage = NULL, coverage_id = :cid WHERE id = :id"),59                        [{"cid": cid, "id": rid} for (rid, _), cid in zip(rows, ids)])60            last_id = int(rows[-1][0])61        done += len(rows)62        batches += 163        if batches % 20 == 0:64            log.info("converted %d rows (last id %d) in %.0fs", done, last_id, time.time() - t0)65        if max_batches is not None and batches >= max_batches:66            break67        if sleep_s:68            time.sleep(sleep_s)69    return done, batches707172def main() -> int:73    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)74    ap.add_argument("--batch", type=int, default=2000, help="rows per transaction (default 2000)")75    ap.add_argument("--max-batches", type=int, help="stop after N batches (resume later)")76    ap.add_argument("--sleep", type=float, default=0.0, help="pause between batches (seconds) to stay gentle with the API")77    ap.add_argument("--vacuum", action="store_true", help="after conversion: auto_vacuum=INCREMENTAL + VACUUM (exclusive, slow)")78    ap.add_argument("-v", "--verbose", action="store_true")79    args = ap.parse_args()80    logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s")8182    from sqlalchemy import text8384    from core.config import settings85    from core.db import engine8687    log.info("database: %s", settings.state_db)88    t0 = time.time()89    done, batches = convert(args.batch, args.max_batches, args.sleep)90    with engine.connect() as con:91        left = con.execute(text("SELECT count(*) FROM fund_statements WHERE coverage IS NOT NULL")).scalar()92        blobs = con.execute(text("SELECT count(*) FROM fund_coverage_blob")).scalar()93    log.info("converted %d rows in %d batches (%.0fs); rows still inline: %d; distinct blobs: %d", done, batches,94             time.time() - t0, left, blobs)95    if args.vacuum:96        t1 = time.time()97        with engine.connect().execution_options(isolation_level="AUTOCOMMIT") as con:98            before = con.execute(text("PRAGMA page_count")).scalar()99            con.execute(text("PRAGMA auto_vacuum=INCREMENTAL"))100            con.execute(text("VACUUM"))101            after = con.execute(text("PRAGMA page_count")).scalar()102            mode = con.execute(text("PRAGMA auto_vacuum")).scalar()103        log.info("VACUUM: %s → %s pages (auto_vacuum=%s) in %.0fs", before, after, mode, time.time() - t1)104    return 0 if left == 0 or args.max_batches else 1105106107if __name__ == "__main__":108    sys.exit(main())109