Northern-Thailand-Ping-River-Monitor/scripts/backfill_rid_reservoir.py

#!/usr/bin/env python3
"""Backfill rid_reservoir_daily with RID large-dam history (Mae Ngat et al.).

Two paths, both idempotent and both skipping what is already stored, so a
rerun repairs holes left by transient failures and is safe alongside the
hourly live collector:

  --dam-id (default: Mae Ngat)  one dam, whole range, via api/dam — a handful
                                of requests for the entire 2009-today archive
  --all-dams                    all ~35 dams, one request per calendar day via
                                api/dams — thousands of requests, ~25 minutes

Usage:
    uv run scripts/backfill_rid_reservoir.py                     # Mae Ngat since 2018-08-01
    uv run scripts/backfill_rid_reservoir.py --start 2009-01-01  # full archive
    uv run scripts/backfill_rid_reservoir.py --refresh           # rewrite stored days too
    uv run scripts/backfill_rid_reservoir.py --all-dams --start 2015-01-01
    uv run scripts/backfill_rid_reservoir.py --db-url postgresql://...
"""

import argparse
import datetime
import logging
import os
import sys

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

from src.config import Config
from src.rid_reservoir import (
    MAE_NGAT_DAM_ID,
    RidReservoirStore,
    backfill,
    backfill_dam,
)

DEFAULT_START = datetime.date(2018, 8, 1)  # start of the water_measurements grid


def main(argv=None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--start", type=datetime.date.fromisoformat, default=DEFAULT_START
    )
    parser.add_argument("--end", type=datetime.date.fromisoformat, default=None)
    parser.add_argument("--db-url", default=None)
    parser.add_argument("--throttle", type=float, default=0.4)
    parser.add_argument(
        "--dam-id",
        default=MAE_NGAT_DAM_ID,
        help="dam to backfill via the fast range endpoint (default Mae Ngat)",
    )
    parser.add_argument(
        "--all-dams",
        action="store_true",
        help="every dam, one request per calendar day (slow full-fleet path)",
    )
    parser.add_argument(
        "--chunk-days",
        type=int,
        default=1830,
        help="days per range request; the endpoint imposes no limit of its own",
    )
    parser.add_argument(
        "--refresh",
        action="store_true",
        help="re-fetch days already stored (adds level_msl to api/dams rows)",
    )
    args = parser.parse_args(argv)

    logging.basicConfig(
        level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s"
    )
    if args.db_url:
        # postgresql+psycopg2://... -> postgresql (driver suffix is not a dialect)
        connection_string = args.db_url
        db_type = args.db_url.split(":", 1)[0].split("+", 1)[0]
    else:
        cfg = Config.get_database_config()
        if cfg["type"] not in ("sqlite", "postgresql", "mysql"):
            print(f"requires a SQL database, got {cfg['type']}", file=sys.stderr)
            return 1
        connection_string, db_type = cfg["connection_string"], cfg["type"]

    store = RidReservoirStore(connection_string, db_type)
    if not store.connect():
        print(
            "database connection failed — check the connection string",
            file=sys.stderr,
        )
        return 1
    end = args.end or datetime.date.today()
    span_days = (end - args.start).days + 1
    dam_id = None if args.all_dams else args.dam_id
    missing = span_days - len(store.present_dates(args.start, end, dam_id=dam_id))
    if args.all_dams:
        saved = backfill(store, args.start, end, throttle_seconds=args.throttle)
        print(f"backfilled {saved} dam-day rows ({missing} days were missing)")
        return 0 if saved or missing == 0 else 1

    stats = {}
    saved = backfill_dam(
        store,
        dam_id=args.dam_id,
        start=args.start,
        end=end,
        chunk_days=args.chunk_days,
        throttle_seconds=max(args.throttle, 1.0),
        skip_present=not args.refresh,
        stats=stats,
    )
    still_missing = span_days - len(
        store.present_dates(args.start, end, dam_id=args.dam_id)
    )
    print(
        f"backfilled {saved} dam-day rows for dam {args.dam_id} "
        f"(requests: {stats.get('requests', 0)}, {missing} days were missing, "
        f"{still_missing} never published by the source)"
    )
    # A rerun saves nothing once the archive is complete — only a real
    # transport/database failure is an error here.
    return 1 if stats.get("aborted") or stats.get("failures") else 0


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