Compare commits
2
Commits
811af1625b
...
ba781465a9
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
ba781465a9 | ||
|
|
6eafb353b1 |
@@ -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 |
|
||||
|
||||
@@ -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"):
|
||||
|
||||
@@ -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())
|
||||
@@ -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)
|
||||
|
||||
+128
-1
@@ -97,6 +97,128 @@ def _fetch_from_db(
|
||||
return _normalize_long(df)
|
||||
|
||||
|
||||
# Stations whose HII mirror is the SAME telemetry (corr ≈ 1.000, median diff
|
||||
# == station offset exactly — validated 2026-08-11) plus P.81, where the HII
|
||||
# twin reads the same river with a bias (corr 0.906, MAE 19 cm) that the
|
||||
# dynamic overlap offset corrects. P.76/P.77/P.85/P.87 HII twins are DIFFERENT
|
||||
# physical sensors (corr 0.25-0.62) and must never be merged into RID series.
|
||||
HII_FILL_STATIONS = (
|
||||
"P.1",
|
||||
"P.103",
|
||||
"P.20",
|
||||
"P.4A",
|
||||
"P.67",
|
||||
"P.75",
|
||||
"P.82",
|
||||
"P.84",
|
||||
"P.92",
|
||||
"P.81",
|
||||
)
|
||||
_HII_EXACT_MIRRORS = frozenset(HII_FILL_STATIONS) - {"P.81"}
|
||||
_HII_MIN_OVERLAP_HOURS = 168
|
||||
|
||||
|
||||
def _fetch_hii_levels(
|
||||
db_url: str,
|
||||
stations: List[str],
|
||||
start: Optional[datetime.datetime],
|
||||
end: Optional[datetime.datetime],
|
||||
) -> pd.DataFrame:
|
||||
engine = create_engine(db_url, pool_pre_ping=True)
|
||||
query = (
|
||||
"SELECT m.timestamp, s.rid_code AS station_code, m.wl_msl, m.discharge "
|
||||
"FROM hii_waterlevel m JOIN hii_wl_stations s ON s.id = m.station_id "
|
||||
"WHERE s.rid_code IS NOT NULL"
|
||||
)
|
||||
params: Dict = {}
|
||||
if start is not None:
|
||||
query += " AND m.timestamp >= :start_time"
|
||||
params["start_time"] = start
|
||||
if end is not None:
|
||||
query += " AND m.timestamp <= :end_time"
|
||||
params["end_time"] = end
|
||||
placeholders = ", ".join(f":station_{i}" for i in range(len(stations)))
|
||||
query += f" AND s.rid_code IN ({placeholders})"
|
||||
for i, code in enumerate(stations):
|
||||
params[f"station_{i}"] = code
|
||||
|
||||
with engine.connect() as connection:
|
||||
df = pd.read_sql(text(query), connection, params=params)
|
||||
df = df.dropna(subset=["wl_msl"])
|
||||
if df.empty:
|
||||
return df
|
||||
df["timestamp"] = pd.to_datetime(df["timestamp"]).dt.floor("h")
|
||||
df["wl_msl"] = pd.to_numeric(df["wl_msl"], errors="coerce")
|
||||
df["discharge"] = pd.to_numeric(df["discharge"], errors="coerce")
|
||||
df = df.sort_values("timestamp").drop_duplicates(
|
||||
subset=["station_code", "timestamp"], keep="last"
|
||||
)
|
||||
return df
|
||||
|
||||
|
||||
def fill_from_hii(
|
||||
df: pd.DataFrame,
|
||||
db_url: str,
|
||||
start: Optional[datetime.datetime] = None,
|
||||
end: Optional[datetime.datetime] = None,
|
||||
stations: Optional[List[str]] = None,
|
||||
min_overlap_hours: int = _HII_MIN_OVERLAP_HOURS,
|
||||
) -> pd.DataFrame:
|
||||
"""Fill missing (station, hour) rows from the HII mirror telemetry.
|
||||
|
||||
In-memory only — water_measurements is never written. Each station's
|
||||
MSL→gauge offset is derived from the overlap between the two series
|
||||
(median of wl_msl − water_level over ≥ `min_overlap_hours` shared hours),
|
||||
which reproduces the published offset for exact mirrors and bias-corrects
|
||||
P.81. Discharge is copied only for exact mirrors; P.81 fills get NaN
|
||||
discharge (its discharge bias was never validated). Failures degrade to
|
||||
returning `df` unchanged, so DBs without hii_* tables keep working.
|
||||
"""
|
||||
codes = [c for c in (stations or HII_FILL_STATIONS) if c in set(df["station_code"])]
|
||||
if not codes:
|
||||
return df
|
||||
try:
|
||||
hii = _fetch_hii_levels(db_url, codes, start, end)
|
||||
except Exception as error:
|
||||
logger.warning(f"HII gap-fill skipped (fetch failed): {error}")
|
||||
return df
|
||||
if hii.empty:
|
||||
return df
|
||||
|
||||
fills = []
|
||||
for code, mirror in hii.groupby("station_code"):
|
||||
base = df[df["station_code"] == code]
|
||||
overlap = base.merge(
|
||||
mirror[["timestamp", "wl_msl"]], on="timestamp", how="inner"
|
||||
).dropna(subset=["water_level", "wl_msl"])
|
||||
if len(overlap) < min_overlap_hours:
|
||||
continue
|
||||
offset = (overlap["wl_msl"] - overlap["water_level"]).median()
|
||||
# Hours the RID series lacks entirely OR carries only a NaN level;
|
||||
# _normalize_long keeps the later (fill) row on collision.
|
||||
present = base.loc[base["water_level"].notna(), "timestamp"]
|
||||
missing = mirror[~mirror["timestamp"].isin(present)]
|
||||
if missing.empty:
|
||||
continue
|
||||
fill = pd.DataFrame(
|
||||
{
|
||||
"timestamp": missing["timestamp"],
|
||||
"station_code": code,
|
||||
"water_level": missing["wl_msl"] - offset,
|
||||
"discharge": missing["discharge"]
|
||||
if code in _HII_EXACT_MIRRORS
|
||||
else float("nan"),
|
||||
}
|
||||
)
|
||||
fills.append(fill)
|
||||
logger.info(
|
||||
f"HII gap-fill {code}: +{len(fill)} hours (offset {offset:.3f} m)"
|
||||
)
|
||||
if not fills:
|
||||
return df
|
||||
return _normalize_long(pd.concat([df] + fills, ignore_index=True))
|
||||
|
||||
|
||||
def _fetch_station_from_api(
|
||||
api_url: str, station_code: str, hours: int, limit: int = 100000
|
||||
) -> pd.DataFrame:
|
||||
@@ -177,18 +299,23 @@ def load_measurements(
|
||||
use_cache: bool = True,
|
||||
cache_dir: Path = CACHE_DIR,
|
||||
api_url: str = DEFAULT_API_URL,
|
||||
hii_fill: bool = True,
|
||||
) -> pd.DataFrame:
|
||||
"""Load the long-format [timestamp, station_code, water_level, discharge] history.
|
||||
|
||||
Tries PostgreSQL first, then the HTTP API, then the on-disk cache as a last
|
||||
resort. A successful DB/API fetch refreshes the cache; the cache itself is
|
||||
never treated as a source of fresh data.
|
||||
never treated as a source of fresh data. With `hii_fill` (DB path only),
|
||||
gaps are patched in memory from the HII mirror telemetry — training and
|
||||
serving both flow through here, so the two sides see identical series.
|
||||
"""
|
||||
resolved_db_url = resolve_db_url(db_url)
|
||||
|
||||
if resolved_db_url:
|
||||
try:
|
||||
df = _fetch_from_db(resolved_db_url, stations, start, end)
|
||||
if hii_fill:
|
||||
df = fill_from_hii(df, resolved_db_url, start=start, end=end)
|
||||
if use_cache:
|
||||
_write_cache(
|
||||
df, cache_dir, source="postgres", discharge_maybe_synthetic=False
|
||||
|
||||
@@ -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)
|
||||
+63
-15
@@ -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,
|
||||
|
||||
@@ -0,0 +1,124 @@
|
||||
"""Tests for the in-memory HII gap-fill in the ML data loader."""
|
||||
|
||||
import datetime
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pytest
|
||||
|
||||
from src.hii_collector import HiiStore
|
||||
from src.ml.data import fill_from_hii
|
||||
|
||||
START = datetime.datetime(2024, 9, 1, 0, 0)
|
||||
|
||||
|
||||
def _hours(n, offset=0):
|
||||
return [START + datetime.timedelta(hours=offset + i) for i in range(n)]
|
||||
|
||||
|
||||
def _base_frame(code="P.20", n=200, level=2.0):
|
||||
return pd.DataFrame(
|
||||
{
|
||||
"timestamp": _hours(n),
|
||||
"station_code": code,
|
||||
"water_level": level,
|
||||
"discharge": 100.0,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def hii_db(tmp_path):
|
||||
"""SQLite DB with hii_* tables; returns (db_url, store)."""
|
||||
db_url = f"sqlite:///{tmp_path}/hii_fill.db"
|
||||
store = HiiStore(db_url, "sqlite")
|
||||
assert store.connect()
|
||||
return db_url, store
|
||||
|
||||
def _seed_mirror(store, station_id, rid_code, hours, wl_msl, discharge=250.0):
|
||||
from sqlalchemy import text
|
||||
|
||||
with store.engine.begin() as conn:
|
||||
conn.execute(
|
||||
text("INSERT INTO hii_wl_stations (id, rid_code) VALUES (:i, :c)"),
|
||||
{"i": station_id, "c": rid_code},
|
||||
)
|
||||
conn.execute(
|
||||
text(
|
||||
"INSERT INTO hii_waterlevel (station_id, timestamp, wl_msl, discharge) "
|
||||
"VALUES (:i, :t, :w, :d)"
|
||||
),
|
||||
[
|
||||
{"i": station_id, "t": t, "w": w, "d": discharge}
|
||||
for t, w in zip(hours, wl_msl)
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
class TestFillFromHii:
|
||||
def test_exact_mirror_fills_missing_hours(self, hii_db):
|
||||
db_url, store = hii_db
|
||||
base = _base_frame("P.20", n=200, level=2.0)
|
||||
# Mirror covers the base window plus 48 extra hours, at MSL offset +300
|
||||
_seed_mirror(store, 1, "P.20", _hours(248), [302.0] * 248)
|
||||
|
||||
filled = fill_from_hii(base, db_url, min_overlap_hours=168)
|
||||
p20 = filled[filled["station_code"] == "P.20"]
|
||||
assert len(p20) == 248
|
||||
new_rows = p20[p20["timestamp"] >= START + datetime.timedelta(hours=200)]
|
||||
assert len(new_rows) == 48
|
||||
# MSL converted back to gauge datum via the overlap-derived offset
|
||||
assert new_rows["water_level"].round(6).eq(2.0).all()
|
||||
# Exact mirrors copy discharge
|
||||
assert new_rows["discharge"].eq(250.0).all()
|
||||
|
||||
def test_p81_bias_corrected_without_discharge(self, hii_db):
|
||||
db_url, store = hii_db
|
||||
base = _base_frame("P.81", n=200, level=1.5)
|
||||
# Biased mirror: offset 310.2 (not a published offset — derived only)
|
||||
_seed_mirror(store, 2, "P.81", _hours(230), [311.7] * 230)
|
||||
|
||||
filled = fill_from_hii(base, db_url, min_overlap_hours=168)
|
||||
p81 = filled[filled["station_code"] == "P.81"]
|
||||
new_rows = p81[p81["timestamp"] >= START + datetime.timedelta(hours=200)]
|
||||
assert len(new_rows) == 30
|
||||
assert new_rows["water_level"].round(6).eq(1.5).all()
|
||||
assert new_rows["discharge"].isna().all()
|
||||
|
||||
def test_insufficient_overlap_skips_station(self, hii_db):
|
||||
db_url, store = hii_db
|
||||
base = _base_frame("P.20", n=50) # only 50 shared hours
|
||||
_seed_mirror(store, 1, "P.20", _hours(100), [302.0] * 100)
|
||||
|
||||
filled = fill_from_hii(base, db_url, min_overlap_hours=168)
|
||||
assert len(filled) == len(base)
|
||||
|
||||
def test_nan_level_rows_are_replaced(self, hii_db):
|
||||
db_url, store = hii_db
|
||||
base = _base_frame("P.20", n=200, level=2.0)
|
||||
base.loc[10, "water_level"] = np.nan
|
||||
_seed_mirror(store, 1, "P.20", _hours(200), [302.5] * 200)
|
||||
|
||||
filled = fill_from_hii(base, db_url, min_overlap_hours=168)
|
||||
p20 = filled[filled["station_code"] == "P.20"]
|
||||
assert len(p20) == 200 # no duplicate hour
|
||||
replaced = p20[p20["timestamp"] == START + datetime.timedelta(hours=10)]
|
||||
assert replaced["water_level"].round(6).eq(2.0).all()
|
||||
|
||||
def test_never_merged_station_untouched(self, hii_db):
|
||||
db_url, store = hii_db
|
||||
base = _base_frame("P.76", n=200, level=1.0)
|
||||
_seed_mirror(store, 3, "P.76", _hours(300), [301.0] * 300)
|
||||
|
||||
filled = fill_from_hii(base, db_url, min_overlap_hours=1)
|
||||
assert len(filled) == len(base) # P.76 not in HII_FILL_STATIONS
|
||||
|
||||
def test_missing_tables_degrade_gracefully(self, tmp_path):
|
||||
base = _base_frame("P.20")
|
||||
filled = fill_from_hii(base, f"sqlite:///{tmp_path}/empty.db")
|
||||
assert filled.equals(base)
|
||||
|
||||
def test_no_station_overlap_returns_input(self, hii_db):
|
||||
db_url, _ = hii_db
|
||||
base = _base_frame("P.5") # not fillable
|
||||
assert fill_from_hii(base, db_url).equals(base)
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user