diff --git a/docs/DATA_SOURCES.md b/docs/DATA_SOURCES.md index e01fee2..062544d 100644 --- a/docs/DATA_SOURCES.md +++ b/docs/DATA_SOURCES.md @@ -183,8 +183,8 @@ Oct 2024 flood). Mae Kuang Udom Thara is the second upstream reservoir. | Source | What | Access | |---|---|---| -| `https://lsim.rid.go.th/ForeCast?reservoirid=22` | Mae Ngat daily status/forecast (RID) | Open, scrape | -| `https://app.rid.go.th/reservoir/` | RID reservoir DB, per-reservoir daily detail with date-range URLs | Open, scrape (confirm URL pattern before hard-coding) | +| `https://app.rid.go.th/reservoir/api/dams` | **INGESTED** — daily snapshot of all ~35 large dams (storage/inflow/outflow MCM, % of usable). `POST` with form field `date=YYYY-MM-DD` (empty = today); GET returns 404 "Unknown method." Archive ≥ 2009; `level_msl` (`DMD_Q`) populated in older years only. Mae Ngat = `DAM_ID 200103` — hit 113% usable capacity, ~19 MCM/day inflow, in Oct 2024. Collected daily by `src/rid_reservoir.py` into `rid_dams` + `rid_reservoir_daily`; backfill via `scripts/backfill_rid_reservoir.py` | Open, no auth | +| `https://lsim.rid.go.th/ForeCast?reservoirid=22` | Mae Ngat daily status/forecast (RID) | Open, scrape — timed out from outside RID network when probed 2026-08-13 | | `https://water.egat.co.th` | EGAT dams (Bhumibol/Sirikit) hourly+daily inflow/outflow/level | Endpoint catalog not public; contact EGAT (0-2436-8186). Only relevant downstream of Bhumibol | | ThaiWater `/v2/large-dam/*`, `dam_rulecurve/graph` | All large/medium dams incl. hourly | Requires HII API key (§2.2) | @@ -239,7 +239,7 @@ Oct 2024 flood). Mae Kuang Udom Thara is the second upstream reservoir. | HII `rain_24h` | ✅ Ingested hourly via `src/hii_collector.py` → `hii_rain_stations` + `hii_rainfall` (Ping-filtered; ~300 stations) | — | | HII `waterlevel_load` | ✅ Ingested hourly via `src/hii_collector.py` → `hii_wl_stations` + `hii_waterlevel` (125 Ping stations, m MSL; `rid_code` column maps mirrors like `ridhydro_P.1` → `P.1`, `offset_msl` converts MSL → gauge datum) | — | | HII `waterlevel_graph` | ✅ Backfill via `scripts/backfill_hii_waterlevel.py` → `hii_waterlevel` (hourly MSL + discharge, archive ≥2019; full-year windows per request; upserts never overwrite live-snapshot columns) | Run once on the box: `python scripts/backfill_hii_waterlevel.py` (defaults: 2019-01-01 → today, RID-mirror + key stations; `--stations P.1,P.67`, `--all` for every Ping station) | -| Mae Ngat reservoir (lsim.rid.go.th) | ❌ Not ingested | Add daily storage/release scrape | +| Mae Ngat reservoir (app.rid.go.th) | ✅ Ingested (2026-08-13) | Daily storage/inflow/outflow for all large dams → `rid_reservoir_daily`; candidate model features for next retrain | | Satellite QPE (GSMaP/IMERG) | ❌ | Backfill training rainfall (GEE) | | Open-Meteo forecasts | ❌ | Add forecast features (live + 2021 archive for training) | | HII API key (dams, forecasts) | ❌ | Contact HII for sanctioned access | diff --git a/scripts/backfill_rain_db.py b/scripts/backfill_rain_db.py index 709310c..e1f2c7a 100644 --- a/scripts/backfill_rain_db.py +++ b/scripts/backfill_rain_db.py @@ -26,7 +26,9 @@ def main(argv=None) -> int: level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s" ) if args.db_url: - connection_string, db_type = args.db_url, args.db_url.split(":", 1)[0] + # 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"): diff --git a/scripts/backfill_rid_reservoir.py b/scripts/backfill_rid_reservoir.py new file mode 100644 index 0000000..a088871 --- /dev/null +++ b/scripts/backfill_rid_reservoir.py @@ -0,0 +1,69 @@ +#!/usr/bin/env python3 +"""Backfill rid_reservoir_daily with RID large-dam history (Mae Ngat et al.). + +One request per day against app.rid.go.th/reservoir/api/dams (archive reaches +back to at least 2009). Days already stored are skipped, so reruns only fetch +what is missing — safe alongside the hourly live collector, and a rerun +repairs holes left by transient failures. + +Usage: + uv run scripts/backfill_rid_reservoir.py # missing days since 2018-08-01 + uv run scripts/backfill_rid_reservoir.py --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 RidReservoirStore, backfill + +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) + 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 + missing = span_days - len(store.present_dates(args.start, end)) + 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 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/config.py b/src/config.py index 05b8f30..de24488 100644 --- a/src/config.py +++ b/src/config.py @@ -86,6 +86,11 @@ class Config: "yes", ) HII_BASIN_CODE = int(os.getenv("HII_BASIN_CODE", "6")) # 6 = Ping Basin + + # RID large-dam daily status (app.rid.go.th/reservoir) — Mae Ngat et al. + ENABLE_RESERVOIR_COLLECTION = os.getenv( + "ENABLE_RESERVOIR_COLLECTION", "true" + ).lower() in ("1", "true", "yes") # TTL for the /api/hii/*/latest response cache; source data changes hourly HII_CACHE_TTL_SECONDS = int(os.getenv("HII_CACHE_TTL_SECONDS", "120")) # TTL for the /measurements/latest response cache (hottest endpoint) diff --git a/src/rid_reservoir.py b/src/rid_reservoir.py new file mode 100644 index 0000000..a7bd584 --- /dev/null +++ b/src/rid_reservoir.py @@ -0,0 +1,351 @@ +"""Collector for RID large-dam daily status (app.rid.go.th/reservoir). + +The Royal Irrigation Department reservoir app exposes an unauthenticated +JSON API: ``POST https://app.rid.go.th/reservoir/api/dams`` with form field +``date=YYYY-MM-DD`` (empty = today) returns a daily snapshot of every large +dam in Thailand — storage, inflow and outflow in MCM — with archive depth +back to at least 2009. GET returns 404 ("Unknown method."); the POST body +may be empty but must carry a Content-Length. + +The reservoir that matters for P.1 flood forecasting is Mae Ngat Somboon +Chon (DAM_ID 200103), the only large dam upstream of Chiang Mai: during the +Oct 2024 record flood it reached 113% of usable capacity with inflow spikes +of ~19 MCM/day. All dams in the payload are stored (same request cost); +filtering happens at feature-build time. + +See docs/DATA_SOURCES.md for the endpoint catalog. +""" + +import datetime +import logging +import time +from typing import Any, Dict, List, Optional + +import requests + +logger = logging.getLogger(__name__) + +RID_DAMS_URL = "https://app.rid.go.th/reservoir/api/dams" +MAE_NGAT_DAM_ID = "200103" + + +def _to_float(value: Any) -> Optional[float]: + """API numerics arrive as strings ('222.01'), ' - ' placeholders, or None.""" + if value is None: + return None + if isinstance(value, str): + value = value.replace(",", "").strip() + if value in ("", "-"): + return None + try: + return float(value) + except (TypeError, ValueError): + return None + + +def parse_dam_records(payload: Dict) -> List[Dict]: + """Flatten the regions/dams payload into per-dam daily rows.""" + records = [] + for region in payload.get("regions") or []: + for dam in region.get("dams") or []: + dam_id = dam.get("DAM_ID") + try: + date = datetime.date.fromisoformat(dam.get("DMD_Date") or "") + except ValueError: + continue + if not dam_id: + continue + records.append( + { + "dam_id": dam_id, + "region": region.get("region_name"), + "name_th": dam.get("DAM_Name"), + "latitude": _to_float(dam.get("DAM_Lat")), + "longitude": _to_float(dam.get("DAM_Lon")), + "capacity_max_mcm": _to_float(dam.get("DAM_QMax")), + "capacity_normal_mcm": _to_float(dam.get("DAM_QStore")), + "date": date, + "storage_mcm": _to_float(dam.get("DMD_QUse")), + "storage_pct": _to_float(dam.get("PERCENT_DMD_QUse")), + "inflow_mcm": _to_float(dam.get("DMD_Inflow")), + "outflow_mcm": _to_float(dam.get("DMD_Outflow")), + "level_msl": _to_float(dam.get("DMD_Q")), + } + ) + return records + + +class RidReservoirClient: + """HTTP client for the RID reservoir daily-status API.""" + + def __init__( + self, + url: str = RID_DAMS_URL, + session: Optional[requests.Session] = None, + timeout: int = 60, + ): + self.url = url + self.session = session or requests.Session() + self.timeout = timeout + + def fetch_day(self, date: Optional[datetime.date] = None) -> List[Dict]: + """Daily rows for every large dam on `date` (None = today).""" + data = {"date": date.isoformat()} if date else {"date": ""} + response = self.session.post( + self.url, + data=data, + headers={"X-Requested-With": "XMLHttpRequest"}, + timeout=self.timeout, + ) + response.raise_for_status() + return parse_dam_records(response.json()) + + +class RidReservoirStore: + """SQL persistence for dam metadata + daily measurements. + + Shares the app's relational database; writes rid_dams (metadata) and + rid_reservoir_daily keyed (dam_id, date) — the composite natural PK keeps + the table TimescaleDB-hypertable compatible. + """ + + def __init__(self, connection_string: str, db_type: str): + self.db_type = db_type.lower() + if self.db_type not in ("sqlite", "postgresql", "mysql"): + raise ValueError( + f"Reservoir collection requires a SQL database, got '{db_type}'" + ) + self.connection_string = connection_string + self.engine = None + + def connect(self) -> bool: + try: + from sqlalchemy import create_engine + + self.engine = create_engine(self.connection_string, pool_pre_ping=True) + self._create_tables() + return True + except Exception as e: + logger.error(f"RidReservoirStore failed to connect: {e}") + self.engine = None + return False + + def _create_tables(self): + from sqlalchemy import text + + ddl = [ + """ + CREATE TABLE IF NOT EXISTS rid_dams ( + dam_id VARCHAR(10) PRIMARY KEY, + region VARCHAR(40), + name_th VARCHAR(255), + latitude NUMERIC(10,6), + longitude NUMERIC(10,6), + capacity_max_mcm NUMERIC(10,2), + capacity_normal_mcm NUMERIC(10,2), + updated_at TIMESTAMP + ) + """, + """ + CREATE TABLE IF NOT EXISTS rid_reservoir_daily ( + dam_id VARCHAR(10) NOT NULL, + date DATE NOT NULL, + storage_mcm NUMERIC(10,2), + storage_pct NUMERIC(6,2), + inflow_mcm NUMERIC(10,2), + outflow_mcm NUMERIC(10,2), + level_msl NUMERIC(8,2), + created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + PRIMARY KEY (dam_id, date) + ) + """, + ] + if self.db_type != "mysql": + ddl.append( + "CREATE INDEX IF NOT EXISTS idx_rid_reservoir_date " + "ON rid_reservoir_daily(date)" + ) + with self.engine.begin() as conn: + for statement in ddl: + conn.execute(text(statement)) + + def _upsert(self, table: str, key_cols: List[str], value_cols: List[str]) -> str: + cols = key_cols + value_cols + col_list = ", ".join(cols) + params = ", ".join(f":{c}" for c in cols) + if self.db_type == "sqlite": + return f"INSERT OR REPLACE INTO {table} ({col_list}) VALUES ({params})" + if self.db_type == "postgresql": + updates = ", ".join(f"{c} = EXCLUDED.{c}" for c in value_cols) + conflict = ", ".join(key_cols) + return ( + f"INSERT INTO {table} ({col_list}) VALUES ({params}) " + f"ON CONFLICT ({conflict}) DO UPDATE SET {updates}" + ) + updates = ", ".join(f"{c} = VALUES({c})" for c in value_cols) + return ( + f"INSERT INTO {table} ({col_list}) VALUES ({params}) " + f"ON DUPLICATE KEY UPDATE {updates}" + ) + + def save(self, records: List[Dict]) -> int: + """Upsert one day's dam rows (metadata + measurements); idempotent.""" + if not records: + return 0 + if not self.engine and not self.connect(): + return 0 + from sqlalchemy import text + + dam_cols = [ + "region", + "name_th", + "latitude", + "longitude", + "capacity_max_mcm", + "capacity_normal_mcm", + ] + measure_cols = [ + "storage_mcm", + "storage_pct", + "inflow_mcm", + "outflow_mcm", + "level_msl", + ] + dam_sql = self._upsert("rid_dams", ["dam_id"], dam_cols + ["updated_at"]) + measure_sql = self._upsert( + "rid_reservoir_daily", ["dam_id", "date"], measure_cols + ) + now = datetime.datetime.now() + dams = {} + measurements = [] + for record in records: + dam_row = {c: record.get(c) for c in dam_cols} + dam_row.update({"dam_id": record["dam_id"], "updated_at": now}) + dams[record["dam_id"]] = dam_row + measure_row = {c: record.get(c) for c in measure_cols} + measure_row.update( + {"dam_id": record["dam_id"], "date": record["date"]} + ) + measurements.append(measure_row) + try: + with self.engine.begin() as conn: + conn.execute(text(dam_sql), list(dams.values())) + conn.execute(text(measure_sql), measurements) + return len(measurements) + except Exception as e: + logger.error(f"RidReservoirStore save failed: {e}") + return 0 + + def present_dates( + self, start: datetime.date, end: datetime.date + ) -> "set[datetime.date]": + """Dates in [start, end] that already have rows, for backfill skipping.""" + if not self.engine and not self.connect(): + return set() + from sqlalchemy import text + + with self.engine.begin() as conn: + values = conn.execute( + text( + "SELECT DISTINCT date FROM rid_reservoir_daily " + "WHERE date >= :start AND date <= :end" + ), + {"start": start, "end": end}, + ).fetchall() + dates = set() + for (value,) in values: + if isinstance(value, str): # sqlite returns ISO strings + value = datetime.date.fromisoformat(value[:10]) + if isinstance(value, datetime.datetime): + value = value.date() + dates.add(value) + return dates + + +class RidReservoirCollector: + """Fetch + persist the daily dam snapshot (today and yesterday).""" + + def __init__(self, db_config: Dict, client: Optional[RidReservoirClient] = None): + self.client = client or RidReservoirClient() + self.store = RidReservoirStore( + connection_string=db_config["connection_string"], + db_type=db_config["type"], + ) + + def run_cycle(self) -> int: + """Collect today's snapshot plus yesterday's (late daily revisions).""" + saved = 0 + today = datetime.date.today() + for date in (None, today - datetime.timedelta(days=1)): + try: + saved += self.store.save(self.client.fetch_day(date)) + except Exception as e: + logger.error(f"RID reservoir collection failed for {date}: {e}") + logger.info(f"RID reservoir collection: {saved} dam-day rows saved") + return saved + + +def backfill( + store: RidReservoirStore, + start: datetime.date, + end: Optional[datetime.date] = None, + client: Optional[RidReservoirClient] = None, + throttle_seconds: float = 0.4, +) -> int: + """Fetch every MISSING day in [start, end]; one polite request per day. + + Only dates absent from rid_reservoir_daily are requested, so a rerun + repairs holes left by transient failures instead of resuming past them + (the hourly collector writes today's rows immediately, which makes any + newest-row cursor useless as a resume point). A save that persists + nothing counts as a failure too — a broken DB must not burn thousands + of requests against the RID API. + """ + client = client or RidReservoirClient() + end = end or datetime.date.today() + if not store.engine and not store.connect(): + logger.error("backfill aborted: database connection failed") + return 0 + span = [ + start + datetime.timedelta(days=i) for i in range((end - start).days + 1) + ] + present = store.present_dates(start, end) + targets = [d for d in span if d not in present] + logger.info( + f"backfill: {len(targets)} of {len(span)} days missing in [{start}, {end}]" + ) + total = 0 + failures = 0 + for i, date in enumerate(targets): + try: + records = client.fetch_day(date) + saved = store.save(records) + if records and not saved: + raise RuntimeError("database save persisted 0 rows") + total += saved + failures = 0 + except Exception as e: + failures += 1 + logger.warning(f"backfill {date} failed ({failures} in a row): {e}") + if failures >= 5: + logger.error("5 consecutive failures — aborting backfill") + break + if i % 100 == 0: + logger.info(f"backfill progress: {date} ({total} rows)") + time.sleep(throttle_seconds) + return total + + +def create_collector_from_config() -> Optional[RidReservoirCollector]: + """Build a collector from app Config; None when disabled or non-SQL DB.""" + from .config import Config + + if not Config.ENABLE_RESERVOIR_COLLECTION: + return None + db_config = Config.get_database_config() + if db_config["type"] not in ("sqlite", "postgresql", "mysql"): + logger.warning( + f"Reservoir collection skipped: DB_TYPE '{db_config['type']}' is not SQL" + ) + return None + return RidReservoirCollector(db_config) diff --git a/src/web_api.py b/src/web_api.py index 2060fe3..731caac 100644 --- a/src/web_api.py +++ b/src/web_api.py @@ -186,6 +186,17 @@ async def lifespan(app: FastAPI): app_state["hii_collector"] = None logger.error(f"HII collector initialization failed: {e}") + # Initialize RID reservoir collector (Mae Ngat + all large dams, daily) + try: + from .rid_reservoir import create_collector_from_config as create_rsv + + app_state["reservoir_collector"] = create_rsv() + if app_state["reservoir_collector"]: + logger.info("RID reservoir collection enabled (large-dam daily status)") + except Exception as e: + app_state["reservoir_collector"] = None + logger.error(f"Reservoir collector initialization failed: {e}") + # Initialize health checks health_manager = HealthCheckManager() health_manager.add_check(DatabaseHealthCheck(app_state["scraper"].db_adapter)) @@ -372,6 +383,17 @@ async def background_scraping_task(): except Exception as e: logger.error(f"HII collection failed: {e}") + # RID large-dam daily status (Mae Ngat storage/inflow/outflow) + reservoir_collector = app_state.get("reservoir_collector") + if reservoir_collector: + try: + rsv_saved = await asyncio.get_event_loop().run_in_executor( + None, reservoir_collector.run_cycle + ) + set_gauge("reservoir_rows_saved", rsv_saved) + except Exception as e: + logger.error(f"Reservoir collection failed: {e}") + # Persist the Open-Meteo catchment rain (observed tail + # 48h forecast) so the DB carries the weather context too await _persist_rain() @@ -798,6 +820,25 @@ def _hii_engine(): return store.engine +def _aux_stats_engine(): + """Any available engine for counting auxiliary tables in /api/stats. + + The HII and reservoir collectors are gated by independent config flags; + either store's engine can run the guarded COUNT queries, so falling back + keeps the stats honest when one collector is disabled. + """ + engine = _hii_engine() + if engine is not None: + return engine + collector = app_state.get("reservoir_collector") + if not collector: + return None + store = collector.store + if not store.engine and not store.connect(): + return None + return store.engine + + def _hii_rows(sql: str, params: Dict[str, Any]) -> List[Dict[str, Any]]: """Run a read query against the hii_* tables, JSON-normalizing numerics.""" engine = _hii_engine() @@ -1222,21 +1263,26 @@ async def get_database_stats(): except Exception as e: logger.warning(f"HII stats unavailable: {e}") - openmeteo_n = 0 - try: # table appears with the 2026-08 rain work; older DBs lack it - engine = _hii_engine() - if engine is not None: - from sqlalchemy import text as _text + # Tables that appear with 2026-08 feature work; older DBs lack them, + # so each count is independently best-effort. + extra_counts = {"openmeteo_rain": 0, "rid_reservoir_daily": 0} + engine = _aux_stats_engine() + if engine is not None: + from sqlalchemy import text as _text - with engine.connect() as conn: - openmeteo_n = int( - conn.execute( - _text("SELECT COUNT(*) FROM openmeteo_rain") - ).scalar() - or 0 - ) - except Exception: - pass + for table in extra_counts: + try: + with engine.connect() as conn: + extra_counts[table] = int( + conn.execute( + _text(f"SELECT COUNT(*) FROM {table}") + ).scalar() + or 0 + ) + except Exception: + pass + openmeteo_n = extra_counts["openmeteo_rain"] + reservoir_n = extra_counts["rid_reservoir_daily"] def as_dt(value): if isinstance(value, str): @@ -1261,11 +1307,13 @@ async def get_database_stats(): "total_measurements": stats["total_measurements"] + rain_n + wl_n - + openmeteo_n, + + openmeteo_n + + reservoir_n, "rid_measurements": stats["total_measurements"], "hii_rainfall_measurements": rain_n, "hii_waterlevel_measurements": wl_n, "openmeteo_rain_measurements": openmeteo_n, + "reservoir_measurements": reservoir_n, "station_count": stats["station_count"] + hii_stations, "rid_station_count": stats["station_count"], "hii_station_count": hii_stations, diff --git a/tests/test_rid_reservoir.py b/tests/test_rid_reservoir.py new file mode 100644 index 0000000..8c03a94 --- /dev/null +++ b/tests/test_rid_reservoir.py @@ -0,0 +1,218 @@ +"""Tests for the RID large-dam daily collector (parsing + persistence).""" + +import datetime + +import pytest + +from src.rid_reservoir import ( + MAE_NGAT_DAM_ID, + RidReservoirCollector, + RidReservoirStore, + backfill, + parse_dam_records, +) + + +def _dams_payload(date="2026-08-13"): + return { + "date_th": date, + "regions": [ + { + "region_name": "เหนือ", + "dams": [ + { + "DAM_ID": "200103", + "DAM_Name": "เขื่อนแม่งัดสมบูรณ์ชล", + "DAM_Lat": 19.16138, + "DAM_Lon": 99.04011, + "DMD_Date": date, + "DAM_QMax": "323.00", + "DAM_QStore": "265.00", + "DMD_QUse": "222.01", + "PERCENT_DMD_QUse": "84.00", + "DMD_Inflow": "3.44", + "DMD_Outflow": "3.57", + "DMD_Q": " - ", # placeholder -> None + }, + # No DAM_ID -> skipped + {"DAM_Name": "broken", "DMD_Date": date}, + # Bad date -> skipped + {"DAM_ID": "200199", "DMD_Date": "n/a"}, + ], + }, + { + "region_name": "กลาง", + "dams": [ + { + "DAM_ID": "200301", + "DAM_Name": "เขื่อนป่าสักชลสิทธิ์", + "DMD_Date": date, + "DMD_QUse": "500.10", + "DMD_Q": "255.57", + } + ], + }, + ], + } + + +class FakeClient: + def __init__(self, payload=None, fail_dates=()): + self.payload = payload or _dams_payload() + self.fail_dates = set(fail_dates) + self.calls = [] + + def fetch_day(self, date=None): + self.calls.append(date) + if date in self.fail_dates: + raise ConnectionError("boom") + if date is not None: + return parse_dam_records(_dams_payload(date.isoformat())) + return parse_dam_records(self.payload) + + +class TestParsing: + def test_parse_dam_records(self): + records = parse_dam_records(_dams_payload()) + assert len(records) == 2 # broken rows skipped + ngat = next(r for r in records if r["dam_id"] == MAE_NGAT_DAM_ID) + assert ngat["region"] == "เหนือ" + assert ngat["date"] == datetime.date(2026, 8, 13) + assert ngat["storage_mcm"] == 222.01 + assert ngat["storage_pct"] == 84.0 + assert ngat["inflow_mcm"] == 3.44 + assert ngat["outflow_mcm"] == 3.57 + assert ngat["level_msl"] is None # ' - ' placeholder + assert ngat["capacity_normal_mcm"] == 265.0 + + def test_parse_empty_payload(self): + assert parse_dam_records({}) == [] + + +class TestStore: + @pytest.fixture + def store(self, tmp_path): + store = RidReservoirStore(f"sqlite:///{tmp_path}/rsv_test.db", "sqlite") + assert store.connect() + return store + + def test_rejects_non_sql(self): + with pytest.raises(ValueError): + RidReservoirStore("http://localhost:8428", "victoriametrics") + + def test_roundtrip_and_upsert(self, store): + records = parse_dam_records(_dams_payload()) + assert store.save(records) == 2 + # Second save of the same day updates in place, no duplicates + records[0]["storage_mcm"] = 230.00 + assert store.save(records) == 2 + from sqlalchemy import text + + with store.engine.begin() as conn: + rows = conn.execute( + text( + "SELECT dam_id, storage_mcm FROM rid_reservoir_daily " + "ORDER BY dam_id" + ) + ).fetchall() + dams = conn.execute(text("SELECT COUNT(*) FROM rid_dams")).scalar() + assert len(rows) == 2 + assert dams == 2 + assert float(rows[0][1]) == 230.00 + + def test_save_empty(self, store): + assert store.save([]) == 0 + + def test_present_dates(self, store): + lo, hi = datetime.date(2026, 8, 1), datetime.date(2026, 8, 31) + assert store.present_dates(lo, hi) == set() + store.save(parse_dam_records(_dams_payload())) + assert store.present_dates(lo, hi) == {datetime.date(2026, 8, 13)} + # Outside the window -> excluded + assert store.present_dates(lo, datetime.date(2026, 8, 12)) == set() + + +class TestCollectorAndBackfill: + def test_run_cycle_today_and_yesterday(self, tmp_path): + collector = RidReservoirCollector( + {"type": "sqlite", "connection_string": f"sqlite:///{tmp_path}/c.db"}, + client=FakeClient(), + ) + saved = collector.run_cycle() + assert saved == 4 # 2 dams x (today + yesterday) + assert collector.client.calls[0] is None # today via empty date + + def test_run_cycle_survives_fetch_failure(self, tmp_path): + yesterday = datetime.date.today() - datetime.timedelta(days=1) + collector = RidReservoirCollector( + {"type": "sqlite", "connection_string": f"sqlite:///{tmp_path}/f.db"}, + client=FakeClient(fail_dates={yesterday}), + ) + assert collector.run_cycle() == 2 # today only + + def test_backfill_range(self, tmp_path): + store = RidReservoirStore(f"sqlite:///{tmp_path}/b.db", "sqlite") + assert store.connect() + client = FakeClient() + saved = backfill( + store, + datetime.date(2024, 9, 24), + datetime.date(2024, 9, 26), + client=client, + throttle_seconds=0, + ) + assert saved == 6 # 3 days x 2 dams + assert client.calls == [ + datetime.date(2024, 9, 24), + datetime.date(2024, 9, 25), + datetime.date(2024, 9, 26), + ] + + def test_backfill_skips_present_days_so_reruns_repair_holes(self, tmp_path): + store = RidReservoirStore(f"sqlite:///{tmp_path}/skip.db", "sqlite") + assert store.connect() + # Day 25 already stored (e.g. by the live collector) + store.save(parse_dam_records(_dams_payload("2024-09-25"))) + client = FakeClient() + saved = backfill( + store, + datetime.date(2024, 9, 24), + datetime.date(2024, 9, 26), + client=client, + throttle_seconds=0, + ) + assert client.calls == [ + datetime.date(2024, 9, 24), + datetime.date(2024, 9, 26), + ] + assert saved == 4 + + def test_backfill_aborts_when_db_saves_nothing(self, tmp_path, monkeypatch): + store = RidReservoirStore(f"sqlite:///{tmp_path}/deaddb.db", "sqlite") + assert store.connect() + monkeypatch.setattr(store, "save", lambda records: 0) # broken DB + client = FakeClient() + backfill( + store, + datetime.date(2024, 1, 1), + datetime.date(2024, 3, 1), + client=client, + throttle_seconds=0, + ) + assert len(client.calls) == 5 # aborted, not one request per day + + def test_backfill_aborts_after_consecutive_failures(self, tmp_path): + store = RidReservoirStore(f"sqlite:///{tmp_path}/a.db", "sqlite") + assert store.connect() + start = datetime.date(2024, 1, 1) + fail_dates = {start + datetime.timedelta(days=i) for i in range(2, 30)} + client = FakeClient(fail_dates=fail_dates) + saved = backfill( + store, + start, + datetime.date(2024, 3, 1), + client=client, + throttle_seconds=0, + ) + assert saved == 4 # first 2 days succeeded, then 5 failures -> abort + assert len(client.calls) == 7