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

"""Live parity check: Python collectors vs the Rust ping-collector.

Runs one collection of every source with the Python code into PY_DB and with
`ping-collector` into RS_DB (both should start empty), back to back against the
real upstream APIs, then diffs all nine collector tables row by row (ids and
created/updated timestamps ignored). Exit status 1 on any difference.

    uv run python scripts/collector_parity.py \
        --py-db postgresql://postgres:ping@localhost:55432/py \
        --rs-db postgresql://postgres:ping@localhost:55432/rs \
        --collector rust/target/debug/ping-collector

Upstream snapshots can move between the two runs (a new HII hour, a revised
RID value); rerun before treating a small difference as a bug.
"""

import argparse
import os
import subprocess
import sys
from pathlib import Path

REPO = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(REPO))

TABLES = {
    "stations": (
        "id",
        "id, station_code, thai_name, english_name, latitude, longitude, geohash",
    ),
    "water_measurements": (
        "timestamp, station_id",
        "timestamp, station_id, water_level, discharge, discharge_percent, status",
    ),
    "hii_rain_stations": (
        "id",
        "id, oldcode, name_th, name_en, latitude, longitude, sub_basin_id, agency",
    ),
    "hii_rainfall": (
        "station_id, timestamp",
        "station_id, timestamp, rain_1h, rain_24h",
    ),
    "hii_wl_stations": (
        "id",
        "id, oldcode, rid_code, name_th, name_en, latitude, longitude, agency, "
        "river_name, offset_msl, ground_level_msl, min_bank_msl, "
        "critical_level_msl, critical_level_m, qmax, is_key_station",
    ),
    "hii_waterlevel": (
        "station_id, timestamp",
        "station_id, timestamp, wl_msl, wl_m, discharge, flow_rate, "
        "storage_percent, situation_level, diff_wl_bank",
    ),
    "rid_dams": (
        "dam_id",
        "dam_id, region, name_th, latitude, longitude, capacity_max_mcm, "
        "capacity_normal_mcm",
    ),
    "rid_reservoir_daily": (
        "dam_id, date",
        "dam_id, date, storage_mcm, storage_pct, inflow_mcm, outflow_mcm, level_msl",
    ),
    "openmeteo_rain": (
        "timestamp",
        "timestamp, chiang_dao, mae_taeng, mae_ngat, mae_rim, chiang_mai, "
        "catchment_mean",
    ),
}


def run_python(db_url: str, gap_days: int):
    os.environ.update(
        DB_TYPE="postgresql", POSTGRES_CONNECTION_STRING=db_url, LOG_LEVEL="WARNING"
    )
    from sqlalchemy import create_engine

    from src.config import Config
    from src.hii_collector import HiiCollector
    from src.ml import rain
    from src.rid_reservoir import RidReservoirCollector
    from src.water_scraper_v3 import EnhancedWaterMonitorScraper

    cfg = Config.get_database_config()
    scraper = EnhancedWaterMonitorScraper(cfg)
    scraper.fill_data_gaps(gap_days)
    scraper.run_scraping_cycle()
    HiiCollector(cfg).run_cycle()
    RidReservoirCollector(cfg).run_cycle()
    rain.save_to_db(rain.fetch_forecast(), create_engine(db_url), "postgresql")


def run_rust(db_url: str, collector: str, gap_days: int):
    env = {**os.environ, "POSTGRES_CONNECTION_STRING": db_url, "DB_TYPE": "postgresql"}
    for args in (["fill-gaps", str(gap_days)], ["once"]):
        subprocess.run([collector, *args], env=env, check=True)


def table_rows(db_url: str, table: str, key: str, cols: str) -> dict:
    import psycopg2

    with psycopg2.connect(db_url) as conn:
        cur = conn.cursor()
        cur.execute(f"SELECT {cols} FROM {table} ORDER BY {key}")
        nkey = len(key.split(","))
        return {row[:nkey]: row for row in cur.fetchall()}


def diff(py_db: str, rs_db: str) -> bool:
    identical = True
    for table, (key, cols) in TABLES.items():
        a = table_rows(py_db, table, key, cols)
        b = table_rows(rs_db, table, key, cols)
        only_py, only_rs = sorted(set(a) - set(b)), sorted(set(b) - set(a))
        differ = sorted(k for k in set(a) & set(b) if a[k] != b[k])
        same = not (only_py or only_rs or differ)
        identical &= same
        print(
            f"{table:<20} py {len(a):>6}  rs {len(b):>6}  only-py {len(only_py):>4}  "
            f"only-rs {len(only_rs):>4}  differing {len(differ):>4}  "
            f"{'identical' if same else 'DIFFERENT'}"
        )
        for k in only_py[:3]:
            print("   only py:", a[k])
        for k in only_rs[:3]:
            print("   only rs:", b[k])
        for k in differ[:3]:
            print("   py:", a[k], "\n   rs:", b[k])
    return identical


def main():
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--py-db", required=True)
    parser.add_argument("--rs-db", required=True)
    parser.add_argument(
        "--collector", default=str(REPO / "rust/target/debug/ping-collector")
    )
    parser.add_argument("--gap-days", type=int, default=2)
    parser.add_argument("--diff-only", action="store_true")
    args = parser.parse_args()
    if not args.diff_only:
        run_python(args.py_db, args.gap_days)
        run_rust(args.rs_db, args.collector, args.gap_days)
    ok = diff(args.py_db, args.rs_db)
    print("ALL IDENTICAL" if ok else "differences found")
    sys.exit(0 if ok else 1)


if __name__ == "__main__":
    main()