noaa-goes-visualization/migrate_urlcache.py

#!/usr/bin/env python
"""One-time migration: file_database.json -> the SQLite index's remote_file table.

``puller_fits.py`` recorded every URL it had fetched in an 800 MB JSON object, parsed
into memory on every run.  This moves that state into the index, where it is an
indexed lookup instead.

This state is the only record of what has already been downloaded.  Losing it means
re-fetching the entire archive from NOAA, so the migration refuses to proceed on any
inconsistency and never deletes the JSON -- it renames it to .bak only after the row
count matches, and only when asked.

Dry run by default::

    migrate_urlcache.py                 # report what would happen
    migrate_urlcache.py --apply         # write the rows
    migrate_urlcache.py --apply --backup-json
"""

import argparse
import json
import os
import sys
import time

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

from suvi import db, paths


def default_json_path():
    return os.path.join(os.path.dirname(paths.data_root()), "file_database.json")


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--json", default=None, help="source file_database.json")
    parser.add_argument("--db", default=None, help="target index (default: $SUVI_DB)")
    parser.add_argument("--apply", action="store_true", help="write rows (default: dry run)")
    parser.add_argument("--backup-json", action="store_true",
                        help="rename the JSON to .bak after a verified migration")
    parser.add_argument("--batch", type=int, default=50000)
    args = parser.parse_args(argv)

    source = args.json or default_json_path()
    if not os.path.exists(source):
        print(f"No such file: {source}")
        return 1

    size = os.path.getsize(source)
    print(f"Reading {source} ({size / 1024 / 1024:.0f} MB)...")
    started = time.time()
    with open(source, "r") as handle:
        cache = json.load(handle)
    if not isinstance(cache, dict):
        print(f"ERROR: expected a JSON object, got {type(cache).__name__}")
        return 1
    print(f"  {len(cache)} URL records in {time.time() - started:.1f}s")

    malformed = [
        url for url, mtime in cache.items()
        if not isinstance(url, str) or not isinstance(mtime, (int, float))
    ]
    if malformed:
        print(f"ERROR: {len(malformed)} records are malformed, e.g. {malformed[:3]}")
        return 1

    if not args.apply:
        print(f"\nDry run: would insert {len(cache)} rows into remote_file.")
        print("Re-run with --apply to write them.")
        return 0

    target = args.db or paths.default_db_path()
    conn = db.connect(target)
    existing = conn.execute("SELECT count(*) c FROM remote_file").fetchone()["c"]
    print(f"Writing to {target} (currently {existing} rows)...")

    now = time.time()
    written = 0
    batch = []
    for url, mtime in cache.items():
        batch.append((url, float(mtime), None, None, now))
        if len(batch) >= args.batch:
            db.record_remote_files(conn, batch)
            conn.commit()
            written += len(batch)
            batch.clear()
            print(f"  {written}/{len(cache)}")
    if batch:
        db.record_remote_files(conn, batch)
        conn.commit()
        written += len(batch)

    final = conn.execute("SELECT count(*) c FROM remote_file").fetchone()["c"]
    print(f"  wrote {written} rows; table now holds {final}")

    if final < len(cache):
        print(f"ERROR: table holds {final} rows but the JSON had {len(cache)} URLs.")
        print("       Not touching the JSON.  Investigate before re-running.")
        conn.close()
        return 1

    # Spot-check that values survived the round trip, not just the row count.
    sample = list(cache.items())[:: max(1, len(cache) // 20)][:20]
    for url, mtime in sample:
        stored = db.get_remote_mtime(conn, url)
        if stored is None or abs(stored - float(mtime)) > 1e-6:
            print(f"ERROR: round-trip mismatch for {url}: {mtime} -> {stored}")
            conn.close()
            return 1
    print(f"  spot-checked {len(sample)} URLs, all match")
    conn.close()

    if args.backup_json:
        backup = source + ".bak"
        os.rename(source, backup)
        print(f"Renamed {source} -> {backup}")
    else:
        print(f"Left {source} in place; pass --backup-json once puller_fits.py is verified.")
    return 0


if __name__ == "__main__":
    sys.exit(main())