feat(rust): ping-collector, phase 1 of the Rust rewrite
CI / Format & lint (push) Successful in 19s
CI / Test suite (push) Successful in 47s
Docs / Validate documentation (push) Successful in 19s
Rust / fmt, clippy, test, build (push) Successful in 7m35s
Security / Dependency vulnerabilities (push) Successful in 1m37s
Security / Static analysis (push) Successful in 17s
Security / License report (push) Successful in 1m30s

rust/collector replaces the collection half of the Python web app: RID hourly
gauges (with the late-report semantics: store everything, retry every 10 min
while RID is behind, re-fetch yesterday's holes, startup gap fill), HII rain +
water level, RID dams (today + yesterday) and Open-Meteo rain. Same tables and
upsert semantics, one UNNEST statement per table, Bangkok wall-clock time
regardless of the server time zone.

Python side: COLLECTION_ENABLED=false makes the API leader stop collecting and
instead re-forecast + notify when the newest RID reading in the database
advances (polled every FORECAST_POLL_SECONDS) and once per clock hour;
/scraping/status then serves the collector's status file; /scrape/trigger
answers 409.

Verified: 20 Rust tests (parsers, gap/schedule logic, database tests against
PostgreSQL); live parity run of both collectors into two empty databases:
all 9 tables row-for-row identical (scripts/collector_parity.py); the Python
API served health, measurements, stats, status and model forecasts from a
database filled only by ping-collector. 177 Python tests pass.

reqwest uses native-tls: app.rid.go.th offers only TLS 1.2 CBC suites, which
rustls refuses.

Deploy/cutover/rollback: scripts/install_collector.sh, scripts/ping-collector.service;
CI: .gitea/workflows/rust.yml (fmt, clippy, tests with a PostgreSQL service,
release artifact).
This commit is contained in:
grabowski committed 2026-10-17 22:20:02 +07:00
1 parent e3694d7591
commit 4e619aa2ae
26 files changed
+5836 -15

No files matched your search

+4
View File
@@ -58,6 +58,10 @@ RID_POLL_OFFSET_MINUTES=5
RID_RETRY_MINUTES=10 RID_RETRY_MINUTES=10
# Days of RID holes the leader re-fetches at startup (0 = off) # Days of RID holes the leader re-fetches at startup (0 = off)
GAP_FILL_DAYS_ON_START=14 GAP_FILL_DAYS_ON_START=14
# false when the Rust ping-collector service collects (rust/collector);
# share the status file between both services for /scraping/status
COLLECTION_ENABLED=true
# SCRAPING_STATUS_FILE=/opt/thailand-water-monitor/collector-status.json
REQUEST_TIMEOUT=30 REQUEST_TIMEOUT=30
MAX_RETRIES=3 MAX_RETRIES=3
RETRY_DELAY_SECONDS=60 RETRY_DELAY_SECONDS=60
+71
View File
@@ -0,0 +1,71 @@
name: Rust
# rust/ workspace (ping-collector): format, clippy, tests (including the
# database tests against a throwaway PostgreSQL service), release build.
# The Linux binary is uploaded as an artifact for scripts/install_collector.sh.
on:
push:
branches: [master, develop]
paths:
- "rust/**"
- "src/data/stations.json"
- ".gitea/workflows/rust.yml"
pull_request:
branches: [master]
paths:
- "rust/**"
- "src/data/stations.json"
- ".gitea/workflows/rust.yml"
workflow_dispatch:
jobs:
rust:
name: fmt, clippy, test, build
runs-on: ubuntu-latest
services:
postgres:
image: postgres:16
env:
POSTGRES_PASSWORD: ping
POSTGRES_DB: ping
options: >-
--health-cmd pg_isready
--health-interval 5s
--health-timeout 5s
--health-retries 20
env:
CARGO_TERM_COLOR: always
TEST_DATABASE_URL: postgres://postgres:ping@postgres:5432/ping
defaults:
run:
working-directory: rust
steps:
- uses: actions/checkout@v4
- name: System packages (OpenSSL headers for native-tls)
run: |
SUDO=$(command -v sudo || true)
$SUDO apt-get update -qq
$SUDO apt-get install -y -qq pkg-config libssl-dev
- uses: dtolnay/rust-toolchain@stable
with:
components: rustfmt, clippy
- name: fmt
run: cargo fmt --all --check
- name: clippy
run: cargo clippy --all-targets --locked -- -D warnings
- name: test
run: cargo test --locked
- name: release build
run: cargo build --release --locked
- uses: actions/upload-artifact@v3
with:
name: ping-collector-linux-x86_64
path: rust/target/release/ping-collector
+7
View File
@@ -28,6 +28,11 @@ Production: one systemd unit on a small VPS, `/opt/thailand-water-monitor`, user
- Do not add Co-Authored-By trailers. - Do not add Co-Authored-By trailers.
- The dashboard is a single file, `src/static/dashboard.html`, EN + TH via the `t()` - The dashboard is a single file, `src/static/dashboard.html`, EN + TH via the `t()`
table: every user-visible string needs both languages. table: every user-visible string needs both languages.
- Rust (`rust/`, see rust/README.md): `cargo fmt --all`, `cargo clippy --all-targets
-- -D warnings`, `cargo test` (database tests need `TEST_DATABASE_URL`). The Rust
collector writes the same tables as the Python collectors: a change to a table,
upsert or parser on one side must be made on the other until the Python side is
retired (check with a live run of both into two databases).
## Where things are ## Where things are
@@ -36,5 +41,7 @@ Production: one systemd unit on a small VPS, `/opt/thailand-water-monitor`, user
`rain.py` (Open-Meteo), `dam.py`, `hii_rain.py`. `rain.py` (Open-Meteo), `dam.py`, `hii_rain.py`.
- `scripts/retrain.sh` + `water-monitor-retrain.timer`: monthly retrain with staged - `scripts/retrain.sh` + `water-monitor-retrain.timer`: monthly retrain with staged
promote. `scripts/dev_proxy.py`: serve the working-copy dashboard against the live API. promote. `scripts/dev_proxy.py`: serve the working-copy dashboard against the live API.
- `rust/collector` (`ping-collector`): Rust data collector; with it deployed the
Python app runs `COLLECTION_ENABLED=false`. Rewrite plan: `docs/RUST_REWRITE.md`.
- `docs/FLOOD_FORECASTING.md` is the authoritative model write-up; `docs/DATA_SOURCES.md` - `docs/FLOOD_FORECASTING.md` is the authoritative model write-up; `docs/DATA_SOURCES.md`
the source catalog. the source catalog.
+23 -11
View File
@@ -1,6 +1,7 @@
# A Rust rewrite: assessment # A Rust rewrite: assessment
Status: evaluation, nothing decided. Spike: [spikes/rust-forecast](../spikes/rust-forecast/README.md). Status: **in progress.** Phase 1 (collector) is built and verified, deploy pending; see
[rust/README.md](../rust/README.md). Spike: [spikes/rust-forecast](../spikes/rust-forecast/README.md).
## Short version ## Short version
@@ -76,23 +77,34 @@ This slows model iteration, which is the most active part of the repo.
| One static binary; no venv on the server for the API | Server still needs Python + sklearn for the monthly retrain | | One static binary; no venv on the server for the API | Server still needs Python + sklearn for the monthly retrain |
| Same stack as proplan | Feature work = Python + Rust change every time | | Same stack as proplan | Feature work = Python + Rust change every time |
## Plan, if pursued (each phase ships on its own) ## Plan (each phase ships on its own)
| Phase | Scope | Status |
| --- | --- | --- |
| 0 | Spike: P.1 features + HGB inference | done, bit-identical |
| 1 | `ping-collector`: RID, HII, dams, Open-Meteo | built; live parity with Python on all 9 tables; deploy pending |
| 2 | Forecast crate + ntfy | next |
| 3 | Read API | |
| 4 | Cleanup | |
0. **Spike: done.** P.1 features + HGB inference, bit-identical. 0. **Spike: done.** P.1 features + HGB inference, bit-identical.
1. **Collector + notifier daemon** (~1.5k LOC): RID (the fixed semantics: store 1. **Collector daemon** (`rust/collector`, ~1.5k LOC incl. tests): RID with the
late readings, retry while behind, re-fetch yesterday's holes), HII, dams, fixed semantics (store late readings, retry while behind, re-fetch
Open-Meteo, ntfy, `notification_state`. The Python API runs with collection yesterday's holes), HII, dams, Open-Meteo, into the same tables. The Python
off and becomes a stateless reader. Shadow one week first: compare stored API runs with `COLLECTION_ENABLED=false` and re-forecasts when the newest
hours against RID upstream. RID reading in the database advances. ntfy moved to phase 2: the outlook
2. **Forecast crate** (~2k LOC): all 15 stations, dam/HII/rain inputs, heuristic alert needs the forecasts, which are still Python. Cutover and rollback
fallback, probabilities and stages, `forecast_history`; parity CI for every are one command each (`scripts/install_collector.sh`).
station on a flood window and a dry window. 2. **Forecast crate + ntfy** (~2k LOC): all 15 stations, dam/HII/rain inputs,
heuristic fallback, probabilities and stages, `forecast_history`,
`notification_state`; parity CI for every station on a flood window and a
dry window. The Python API then only serves.
3. **Read API** (~2k LOC): routes recorded from the Python API as golden files, 3. **Read API** (~2k LOC): routes recorded from the Python API as golden files,
then a Caddy switch. Python keeps only training/evaluation. then a Caddy switch. Python keeps only training/evaluation.
4. **Cleanup:** drop the multi-backend adapters, demo databases and the 4. **Cleanup:** drop the multi-backend adapters, demo databases and the
Python web stack. Python web stack.
Open questions before phase 1: is Matrix alerting still used? Must `/docs` Open questions before phase 3: is Matrix alerting still used? Must `/docs`
(OpenAPI UI) stay? Is SQLite dev mode needed? Where does the monthly retrain (OpenAPI UI) stay? Is SQLite dev mode needed? Where does the monthly retrain
run if the server stops having a Python environment for the API? run if the server stops having a Python environment for the API?
+1
View File
@@ -0,0 +1 @@
target/
+2639
View File
File diff suppressed because it is too large. Load diff
+29
View File
@@ -0,0 +1,29 @@
[workspace]
resolver = "3"
members = ["collector"]
[workspace.package]
edition = "2024"
license = "MIT"
publish = false
repository = "https://git.b4l.co.th/B4L/Northern-Thailand-Ping-River-Monitor"
[workspace.dependencies]
anyhow = "1"
chrono = { version = "0.4", default-features = false, features = ["std", "clock", "serde"] }
clap = { version = "4", features = ["derive"] }
dotenvy = "0.15"
# native-tls (OpenSSL on Linux, SChannel on Windows), not rustls: app.rid.go.th (dams)
# offers only TLS 1.2 CBC suites (ECDHE-RSA-AES256-SHA384), which rustls refuses.
reqwest = { version = "0.12", default-features = false, features = ["native-tls", "json"] }
serde = { version = "1", features = ["derive"] }
serde_json = { version = "1", features = ["float_roundtrip"] }
sqlx = { version = "0.8", default-features = false, features = ["runtime-tokio", "tls-rustls", "postgres", "chrono", "macros"] }
tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "signal"] }
tracing = "0.1"
tracing-subscriber = { version = "0.3", features = ["env-filter"] }
[profile.release]
lto = "thin"
codegen-units = 1
strip = true
+85
View File
@@ -0,0 +1,85 @@
# Rust services
Rewrite of the Ping River Monitor in Rust, one service at a time. The plan,
costs and the parity rules are in [docs/RUST_REWRITE.md](../docs/RUST_REWRITE.md).
| Crate | Replaces | Status |
| --- | --- | --- |
| `collector` (`ping-collector`) | RID, HII/ThaiWater, RID dams and Open-Meteo collection in the Python web app | phase 1: ready to deploy |
## ping-collector
Collects into the same PostgreSQL tables as the Python app, so the two can be
swapped (and swapped back) with no data migration. The Python web app keeps
serving the API, forecasting and sending ntfy alerts with
`COLLECTION_ENABLED=false`, and re-forecasts whenever the newest RID reading in
the database advances.
```text
ping-collector run # the service: startup gap fill, RID at :05 (+10 min retries
# while RID is behind), HII/dams/rain once per clock hour
ping-collector once [--dry-run] # one cycle of every source; --dry-run writes nothing
# and compares the RID fetch with the database
ping-collector fill-gaps <days|all> # re-fetch RID reports with holes
ping-collector check-db # connect, create missing tables, row counts
```
Configuration is the Python app's `.env` (read from the working directory, or
`--env-file`); real environment variables win. Used keys:
`POSTGRES_CONNECTION_STRING` (or `POSTGRES_HOST/PORT/DB/USER/PASSWORD`),
`RID_POLL_OFFSET_MINUTES` (5), `RID_RETRY_MINUTES` (10),
`GAP_FILL_DAYS_ON_START` (14), `ENABLE_HII_COLLECTION`,
`ENABLE_RESERVOIR_COLLECTION`, `ENABLE_RAIN_COLLECTION`, `HII_BASIN_CODE` (6),
`STATION_CONFIG_PATH` (falls back to the bundled `src/data/stations.json`),
`SCRAPING_STATUS_FILE` (status JSON served by the API's `/scraping/status`),
`LOG_LEVEL` / `RUST_LOG`. All timestamps are Asia/Bangkok wall-clock,
whatever the server's time zone.
Not ported: the one-off archive backfills (`scripts/backfill_*.py`), which
stay in Python.
### Verified
- 20 tests: parsers, gap/schedule logic, and database tests against a real
PostgreSQL (late reports, midnight recovery, upsert revisions, NUMERIC
rounding, dam metadata preservation).
- Live parity: the Python collectors and `ping-collector` run back-to-back
against the real upstream APIs into two empty databases produced identical
rows in all 9 tables (16 stations, 1,104 RID readings, 103 + 76 HII rows,
35 dams x 2 days, 168 rain hours). Rerun with
`scripts/collector_parity.py` after changing a parser, table or upsert on
either side.
- The Python API with `COLLECTION_ENABLED=false` served `/health`,
`/measurements/latest`, `/api/stats`, `/scraping/status` and model
forecasts from a database filled only by `ping-collector`.
### Develop
```bash
cd rust
cargo fmt --all && cargo clippy --all-targets -- -D warnings
docker run -d --name ping-pg -e POSTGRES_PASSWORD=ping -e POSTGRES_DB=ping -p 55432:5432 postgres:16
TEST_DATABASE_URL=postgres://postgres:ping@localhost:55432/ping cargo test
```
Without `TEST_DATABASE_URL` the database tests skip. CI
(`.gitea/workflows/rust.yml`) runs them against a PostgreSQL service and
uploads the Linux binary as the `ping-collector-linux-x86_64` artifact.
HTTP uses native-tls (OpenSSL on Linux), not rustls: app.rid.go.th (dams)
only offers TLS 1.2 CBC cipher suites, which rustls refuses. Building on
Linux needs `pkg-config` and `libssl-dev`.
### Deploy
On the server, after `git pull` (the Python side of the switch ships in the
same commit):
```bash
cargo build --release --locked --manifest-path rust/Cargo.toml # as a normal user,
# or BIN=<CI artifact>
sudo bash scripts/install_collector.sh # install + dry run, nothing switched
sudo bash scripts/install_collector.sh --cutover # Rust collects from now on
sudo bash scripts/install_collector.sh --rollback # back to Python collection
journalctl -u ping-collector -f
```
+21
View File
@@ -0,0 +1,21 @@
[package]
name = "ping-collector"
version = "0.1.0"
description = "Ping River Monitor data collector: RID gauges, HII/ThaiWater, RID dams, Open-Meteo rain -> PostgreSQL"
edition.workspace = true
license.workspace = true
publish.workspace = true
repository.workspace = true
[dependencies]
anyhow.workspace = true
chrono.workspace = true
clap.workspace = true
dotenvy.workspace = true
reqwest.workspace = true
serde.workspace = true
serde_json.workspace = true
sqlx.workspace = true
tokio.workspace = true
tracing.workspace = true
tracing-subscriber.workspace = true
+175
View File
@@ -0,0 +1,175 @@
//! Configuration from the same environment / `.env` file as the Python app,
//! so both services can share `/opt/thailand-water-monitor/.env`.
use anyhow::{Context, Result, bail};
use std::env;
use std::path::PathBuf;
pub const RID_API_URL: &str =
"https://hyd-app-db.rid.go.th/webservice/getGroupHourlyWaterLevelReportAllHL.ashx";
pub const HII_API_BASE: &str = "https://api-v3.thaiwater.net/api/v1/thaiwater30/public";
pub const RID_DAMS_URL: &str = "https://app.rid.go.th/reservoir/api/dams";
pub const OPENMETEO_URL: &str = "https://api.open-meteo.com/v1/forecast";
/// The RID endpoint serves the same JSON to any client; keep the UA the
/// Python collector has used for years.
pub const USER_AGENT: &str = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 \
(KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36";
#[derive(Clone, Debug)]
pub struct Config {
pub database_url: String,
pub rid_api_url: String,
pub hii_api_base: String,
pub dams_url: String,
pub openmeteo_url: String,
pub hii_enabled: bool,
pub dams_enabled: bool,
pub rain_enabled: bool,
pub hii_basin_code: i64,
/// RID is polled this many minutes past each hour.
pub poll_offset_minutes: u32,
/// While the current hour is not published yet, retry this often.
pub retry_minutes: u32,
/// Days of RID holes re-fetched at startup (0 = off).
pub gap_fill_days_on_start: u32,
pub station_config_path: PathBuf,
/// Status JSON read by the Python API's /scraping/status.
pub status_file: Option<PathBuf>,
}
fn var(name: &str) -> Option<String> {
env::var(name)
.ok()
.map(|v| v.trim().to_string())
.filter(|v| !v.is_empty())
}
fn flag(name: &str, default: bool) -> bool {
var(name).map_or(default, |v| {
matches!(v.to_lowercase().as_str(), "1" | "true" | "yes")
})
}
fn number<T: std::str::FromStr>(name: &str, default: T) -> Result<T>
where
T::Err: std::fmt::Display,
{
match var(name) {
None => Ok(default),
Some(v) => v.parse().map_err(|e| anyhow::anyhow!("{name}={v:?}: {e}")),
}
}
/// Percent-encode everything outside RFC 3986 "unreserved" (as Python's
/// `urllib.parse.quote(password, safe="")`).
fn encode_userinfo(s: &str) -> String {
s.bytes()
.map(|b| match b {
b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'.' | b'_' | b'~' => {
(b as char).to_string()
}
_ => format!("%{b:02X}"),
})
.collect()
}
/// SQLAlchemy URLs carry a driver suffix (`postgresql+psycopg2://`) that
/// sqlx does not understand.
pub fn normalize_database_url(url: &str) -> String {
match url.split_once("://") {
Some((scheme, rest))
if scheme.starts_with("postgresql") || scheme.starts_with("postgres") =>
{
format!("postgres://{rest}")
}
_ => url.to_string(),
}
}
impl Config {
pub fn from_env() -> Result<Self> {
let database_url = match var("POSTGRES_CONNECTION_STRING") {
Some(url) => normalize_database_url(&url),
None => {
let password = var("POSTGRES_PASSWORD")
.context("POSTGRES_CONNECTION_STRING or POSTGRES_PASSWORD must be set")?;
format!(
"postgres://{}:{}@{}:{}/{}",
encode_userinfo(&var("POSTGRES_USER").unwrap_or_else(|| "postgres".into())),
encode_userinfo(&password),
var("POSTGRES_HOST").unwrap_or_else(|| "localhost".into()),
var("POSTGRES_PORT").unwrap_or_else(|| "5432".into()),
var("POSTGRES_DB").unwrap_or_else(|| "water_monitoring".into()),
)
}
};
if let Some(db_type) = var("DB_TYPE")
&& db_type != "postgresql"
{
bail!("ping-collector supports PostgreSQL only (DB_TYPE={db_type})");
}
let cfg = Config {
database_url,
rid_api_url: var("RID_API_URL").unwrap_or_else(|| RID_API_URL.into()),
hii_api_base: var("HII_API_BASE").unwrap_or_else(|| HII_API_BASE.into()),
dams_url: var("RID_DAMS_URL").unwrap_or_else(|| RID_DAMS_URL.into()),
openmeteo_url: var("OPENMETEO_URL").unwrap_or_else(|| OPENMETEO_URL.into()),
hii_enabled: flag("ENABLE_HII_COLLECTION", true),
dams_enabled: flag("ENABLE_RESERVOIR_COLLECTION", true),
rain_enabled: flag("ENABLE_RAIN_COLLECTION", true),
hii_basin_code: number("HII_BASIN_CODE", 6)?,
poll_offset_minutes: number("RID_POLL_OFFSET_MINUTES", 5)?,
retry_minutes: number("RID_RETRY_MINUTES", 10)?,
gap_fill_days_on_start: number("GAP_FILL_DAYS_ON_START", 14)?,
station_config_path: var("STATION_CONFIG_PATH")
.unwrap_or_else(|| "stations.json".into())
.into(),
status_file: var("SCRAPING_STATUS_FILE").map(PathBuf::from),
};
if cfg.poll_offset_minutes >= 60 {
bail!("RID_POLL_OFFSET_MINUTES must be between 0 and 59");
}
if cfg.retry_minutes == 0 {
bail!("RID_RETRY_MINUTES must be positive");
}
Ok(cfg)
}
/// The URL with the password masked, for logs.
pub fn redacted_database_url(&self) -> String {
match (self.database_url.find("://"), self.database_url.rfind('@')) {
(Some(s), Some(at)) if at > s => {
let creds = &self.database_url[s + 3..at];
let user = creds.split(':').next().unwrap_or("");
format!(
"{}{user}:***{}",
&self.database_url[..s + 3],
&self.database_url[at..]
)
}
_ => self.database_url.clone(),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn sqlalchemy_urls_are_normalized() {
assert_eq!(
normalize_database_url("postgresql+psycopg2://u:p@h:5432/db"),
"postgres://u:p@h:5432/db"
);
assert_eq!(
normalize_database_url("postgresql://u@h/db"),
"postgres://u@h/db"
);
}
#[test]
fn passwords_are_percent_encoded() {
assert_eq!(encode_userinfo("p@ss:w/rd"), "p%40ss%3Aw%2Frd");
}
}
+143
View File
@@ -0,0 +1,143 @@
//! RID large-dam daily status (app.rid.go.th/reservoir): storage, inflow,
//! outflow for every large dam. Port of the hourly path of
//! src/rid_reservoir.py (today + yesterday's late revisions). The one-off
//! archive backfills stay in Python (scripts/backfill_rid_reservoir.py).
use crate::util::{lenient_f64, scalar_string};
use anyhow::{Context, Result};
use chrono::NaiveDate;
use serde_json::Value;
#[derive(Clone, Debug, PartialEq)]
pub struct DamRecord {
pub dam_id: String,
pub region: Option<String>,
pub name_th: Option<String>,
pub latitude: Option<f64>,
pub longitude: Option<f64>,
pub capacity_max_mcm: Option<f64>,
pub capacity_normal_mcm: Option<f64>,
pub date: NaiveDate,
pub storage_mcm: Option<f64>,
pub storage_pct: Option<f64>,
pub inflow_mcm: Option<f64>,
pub outflow_mcm: Option<f64>,
pub level_msl: Option<f64>,
}
/// Values at or beyond the column's NUMERIC capacity become NULL instead of
/// failing the whole daily batch (the source publishes e.g. 87798 %).
fn bounded(v: Option<f64>, limit: f64) -> Option<f64> {
v.filter(|x| x.abs() < limit)
}
fn string(v: Option<&Value>) -> Option<String> {
v.and_then(Value::as_str)
.filter(|s| !s.is_empty())
.map(str::to_string)
}
pub fn parse_dams(payload: &Value) -> Vec<DamRecord> {
let mut out = Vec::new();
for region in payload
.get("regions")
.and_then(Value::as_array)
.into_iter()
.flatten()
{
for dam in region
.get("dams")
.and_then(Value::as_array)
.into_iter()
.flatten()
{
let Some(date) = dam
.get("DMD_Date")
.and_then(Value::as_str)
.and_then(|s| NaiveDate::parse_from_str(s, "%Y-%m-%d").ok())
else {
continue;
};
let Some(dam_id) = scalar_string(dam.get("DAM_ID")) else {
continue;
};
out.push(DamRecord {
dam_id,
region: string(region.get("region_name")),
name_th: string(dam.get("DAM_Name")),
latitude: lenient_f64(dam.get("DAM_Lat")),
longitude: lenient_f64(dam.get("DAM_Lon")),
capacity_max_mcm: lenient_f64(dam.get("DAM_QMax")),
capacity_normal_mcm: lenient_f64(dam.get("DAM_QStore")),
date,
storage_mcm: bounded(lenient_f64(dam.get("DMD_QUse")), 1e8),
storage_pct: bounded(lenient_f64(dam.get("PERCENT_DMD_QUse")), 1e6),
inflow_mcm: bounded(lenient_f64(dam.get("DMD_Inflow")), 1e8),
outflow_mcm: bounded(lenient_f64(dam.get("DMD_Outflow")), 1e8),
level_msl: bounded(lenient_f64(dam.get("DMD_Q")), 1e6),
});
}
}
out
}
pub struct DamsClient {
http: reqwest::Client,
url: String,
}
impl DamsClient {
pub fn new(http: reqwest::Client, url: String) -> Self {
Self { http, url }
}
/// Every large dam on `date` (None = today). POST only; the body may be
/// empty but must be sent.
pub async fn fetch_day(&self, date: Option<NaiveDate>) -> Result<Vec<DamRecord>> {
let date = date
.map(|d| d.format("%Y-%m-%d").to_string())
.unwrap_or_default();
let payload: Value = self
.http
.post(&self.url)
.header("X-Requested-With", "XMLHttpRequest")
.form(&[("date", date.as_str())])
.timeout(std::time::Duration::from_secs(120))
.send()
.await
.context("RID dams request failed")?
.error_for_status()?
.json()
.await
.context("RID dams response is not JSON")?;
Ok(parse_dams(&payload))
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn parses_regions_and_cleans_numbers() {
let payload = json!({"regions": [{"region_name": "ภาคเหนือ", "dams": [
{"DAM_ID": "200103", "DAM_Name": "แม่งัด", "DMD_Date": "2026-10-17", "DAM_Lat": "19.17",
"DAM_QMax": "265.00", "DMD_QUse": "1,234.5", "PERCENT_DMD_QUse": "87798000",
"DMD_Inflow": " - ", "DMD_Outflow": 1.2, "DMD_Q": " - "},
{"DAM_ID": "", "DMD_Date": "2026-10-17"},
{"DAM_ID": 100602, "DMD_Date": "not a date"}
]}]});
let r = parse_dams(&payload);
assert_eq!(r.len(), 1);
let d = &r[0];
assert_eq!(d.dam_id, "200103");
assert_eq!(d.region.as_deref(), Some("ภาคเหนือ"));
assert_eq!(d.storage_mcm, Some(1234.5));
assert_eq!(d.storage_pct, None); // junk beyond NUMERIC(8,2)
assert_eq!(d.inflow_mcm, None);
assert_eq!(d.outflow_mcm, Some(1.2));
assert_eq!(d.level_msl, None);
assert_eq!(d.latitude, Some(19.17));
}
}
+476
View File
@@ -0,0 +1,476 @@
//! PostgreSQL access: the same tables and upsert semantics as the Python
//! collectors (src/database_adapters.py, hii_collector.py, rid_reservoir.py,
//! ml/rain.py), written as one set-based statement per table (UNNEST) instead
//! of a statement per row.
use crate::dams::DamRecord;
use crate::hii::{RainRecord, WaterLevelRecord};
use crate::rain::RainRow;
use crate::rid::Reading;
use crate::stations::Stations;
use anyhow::{Context, Result};
use chrono::{NaiveDate, NaiveDateTime, NaiveTime};
use sqlx::PgPool;
use sqlx::postgres::PgPoolOptions;
use std::collections::{BTreeMap, HashMap, HashSet};
pub async fn connect(url: &str) -> Result<PgPool> {
PgPoolOptions::new()
.max_connections(4)
.acquire_timeout(std::time::Duration::from_secs(30))
.connect(url)
.await
.context("connecting to PostgreSQL")
}
/// Idempotent DDL, identical to what the Python app creates (so either side
/// can start first against an empty database).
const SCHEMA: &[&str] = &[
"CREATE TABLE IF NOT EXISTS stations (
id SERIAL PRIMARY KEY,
station_code VARCHAR(10) UNIQUE NOT NULL,
thai_name VARCHAR(255) NOT NULL,
english_name VARCHAR(255) NOT NULL,
latitude DECIMAL(10,8),
longitude DECIMAL(11,8),
geohash VARCHAR(20),
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)",
"CREATE TABLE IF NOT EXISTS water_measurements (
id BIGSERIAL PRIMARY KEY,
timestamp TIMESTAMP NOT NULL,
station_id INTEGER NOT NULL,
water_level NUMERIC(10,3),
discharge NUMERIC(10,2),
discharge_percent NUMERIC(5,2),
status VARCHAR(20) DEFAULT 'active',
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
FOREIGN KEY (station_id) REFERENCES stations(id),
UNIQUE(timestamp, station_id))",
"CREATE INDEX IF NOT EXISTS idx_timestamp ON water_measurements(timestamp)",
"CREATE INDEX IF NOT EXISTS idx_station_timestamp ON water_measurements(station_id, timestamp DESC)",
"CREATE TABLE IF NOT EXISTS hii_rain_stations (
id INTEGER PRIMARY KEY,
oldcode VARCHAR(60),
name_th VARCHAR(255),
name_en VARCHAR(255),
latitude NUMERIC(10,6),
longitude NUMERIC(10,6),
sub_basin_id VARCHAR(10),
agency VARCHAR(40),
updated_at TIMESTAMP)",
"CREATE TABLE IF NOT EXISTS hii_rainfall (
station_id INTEGER NOT NULL,
timestamp TIMESTAMP NOT NULL,
rain_1h NUMERIC(7,2),
rain_24h NUMERIC(8,2),
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
PRIMARY KEY (station_id, timestamp))",
"CREATE TABLE IF NOT EXISTS hii_wl_stations (
id INTEGER PRIMARY KEY,
oldcode VARCHAR(60),
rid_code VARCHAR(10),
name_th VARCHAR(255),
name_en VARCHAR(255),
latitude NUMERIC(10,6),
longitude NUMERIC(10,6),
agency VARCHAR(40),
river_name VARCHAR(255),
offset_msl NUMERIC(8,3),
ground_level_msl NUMERIC(8,3),
min_bank_msl NUMERIC(8,3),
critical_level_msl NUMERIC(8,3),
critical_level_m NUMERIC(8,3),
qmax NUMERIC(10,2),
is_key_station BOOLEAN,
updated_at TIMESTAMP)",
"CREATE TABLE IF NOT EXISTS hii_waterlevel (
station_id INTEGER NOT NULL,
timestamp TIMESTAMP NOT NULL,
wl_msl NUMERIC(8,3),
wl_m NUMERIC(8,3),
discharge NUMERIC(10,2),
flow_rate NUMERIC(10,2),
storage_percent NUMERIC(6,2),
situation_level INTEGER,
diff_wl_bank NUMERIC(8,3),
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
PRIMARY KEY (station_id, timestamp))",
"CREATE INDEX IF NOT EXISTS idx_hii_rainfall_ts ON hii_rainfall(timestamp)",
"CREATE INDEX IF NOT EXISTS idx_hii_waterlevel_ts ON hii_waterlevel(timestamp)",
"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(8,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))",
"CREATE INDEX IF NOT EXISTS idx_rid_reservoir_date ON rid_reservoir_daily(date)",
"CREATE TABLE IF NOT EXISTS openmeteo_rain (
timestamp TIMESTAMP PRIMARY KEY,
chiang_dao NUMERIC(6,2), mae_taeng NUMERIC(6,2), mae_ngat NUMERIC(6,2),
mae_rim NUMERIC(6,2), chiang_mai NUMERIC(6,2),
catchment_mean NUMERIC(6,2),
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP)",
];
pub async fn ensure_schema(pool: &PgPool) -> Result<()> {
for ddl in SCHEMA {
sqlx::query(ddl)
.execute(pool)
.await
.with_context(|| format!("DDL failed: {ddl}"))?;
}
Ok(())
}
/// Keep the last record per key (a later report revision wins), preserving
/// first-seen order; Postgres refuses to upsert the same key twice in one
/// statement.
fn dedupe_last<T, K: std::hash::Hash + Eq + Clone>(items: &[T], key: impl Fn(&T) -> K) -> Vec<&T> {
let mut index: HashMap<K, usize> = HashMap::new();
let mut out: Vec<&T> = Vec::new();
for item in items {
let k = key(item);
match index.get(&k) {
Some(&i) => out[i] = item,
None => {
index.insert(k, out.len());
out.push(item);
}
}
}
out
}
pub async fn save_rid(pool: &PgPool, readings: &[Reading], stations: &Stations) -> Result<usize> {
let rows = dedupe_last(readings, |r| (r.timestamp, r.station_num));
let station_nums: Vec<u32> = rows
.iter()
.map(|r| r.station_num)
.collect::<std::collections::BTreeSet<_>>()
.into_iter()
.collect();
let infos: Vec<_> = station_nums.iter().map(|&n| stations.get(n)).collect();
let mut tx = pool.begin().await?;
sqlx::query(
"INSERT INTO stations (id, station_code, thai_name, english_name, latitude, longitude, geohash, updated_at)
SELECT *, NOW() FROM UNNEST($1::int4[], $2::text[], $3::text[], $4::text[], $5::float8[], $6::float8[], $7::text[])
ON CONFLICT (id) DO UPDATE SET
thai_name = EXCLUDED.thai_name, english_name = EXCLUDED.english_name,
latitude = EXCLUDED.latitude, longitude = EXCLUDED.longitude,
geohash = EXCLUDED.geohash, updated_at = NOW()",
)
.bind(station_nums.iter().map(|&n| n as i32).collect::<Vec<_>>())
.bind(infos.iter().map(|s| s.code.clone()).collect::<Vec<_>>())
.bind(infos.iter().map(|s| s.thai_name.clone()).collect::<Vec<_>>())
.bind(infos.iter().map(|s| s.english_name.clone()).collect::<Vec<_>>())
.bind(infos.iter().map(|s| s.latitude).collect::<Vec<_>>())
.bind(infos.iter().map(|s| s.longitude).collect::<Vec<_>>())
.bind(infos.iter().map(|s| s.geohash.clone()).collect::<Vec<_>>())
.execute(&mut *tx)
.await
.context("upserting stations")?;
sqlx::query(
"INSERT INTO water_measurements (timestamp, station_id, water_level, discharge, discharge_percent, status)
SELECT *, 'active' FROM UNNEST($1::timestamp[], $2::int4[], $3::float8[], $4::float8[], $5::float8[])
ON CONFLICT (timestamp, station_id) DO UPDATE SET
water_level = EXCLUDED.water_level, discharge = EXCLUDED.discharge,
discharge_percent = EXCLUDED.discharge_percent, status = EXCLUDED.status",
)
.bind(rows.iter().map(|r| r.timestamp).collect::<Vec<_>>())
.bind(rows.iter().map(|r| r.station_num as i32).collect::<Vec<_>>())
.bind(rows.iter().map(|r| r.water_level).collect::<Vec<_>>())
.bind(rows.iter().map(|r| r.discharge).collect::<Vec<_>>())
.bind(rows.iter().map(|r| r.discharge_percent).collect::<Vec<_>>())
.execute(&mut *tx)
.await
.context("upserting water_measurements")?;
tx.commit().await?;
Ok(rows.len())
}
/// Day -> hours (0-23) that have at least one RID reading, for [start, end].
pub async fn recorded_hours_by_day(
pool: &PgPool,
start: NaiveDate,
end: NaiveDate,
) -> Result<HashMap<NaiveDate, HashSet<u32>>> {
let rows: Vec<(NaiveDate, i32)> = sqlx::query_as(
"SELECT CAST(timestamp AS DATE) AS day, CAST(EXTRACT(HOUR FROM timestamp) AS INTEGER) AS hour
FROM water_measurements WHERE timestamp >= $1 AND timestamp < $2
GROUP BY 1, 2",
)
.bind(start.and_time(NaiveTime::MIN))
.bind((end + chrono::Duration::days(1)).and_time(NaiveTime::MIN))
.fetch_all(pool)
.await?;
let mut out: HashMap<NaiveDate, HashSet<u32>> = HashMap::new();
for (day, hour) in rows {
out.entry(day).or_default().insert(hour as u32);
}
Ok(out)
}
pub async fn earliest_rid_timestamp(pool: &PgPool) -> Result<Option<NaiveDateTime>> {
Ok(
sqlx::query_scalar("SELECT MIN(timestamp) FROM water_measurements")
.fetch_one(pool)
.await?,
)
}
pub async fn latest_rid_timestamp(pool: &PgPool) -> Result<Option<NaiveDateTime>> {
Ok(
sqlx::query_scalar("SELECT MAX(timestamp) FROM water_measurements")
.fetch_one(pool)
.await?,
)
}
/// How many of `readings` are already stored (dry-run comparison).
pub async fn count_stored_rid(pool: &PgPool, readings: &[Reading]) -> Result<i64> {
Ok(sqlx::query_scalar(
"SELECT COUNT(*) FROM water_measurements w
JOIN UNNEST($1::timestamp[], $2::int4[]) AS k(ts, sid) ON w.timestamp = k.ts AND w.station_id = k.sid",
)
.bind(readings.iter().map(|r| r.timestamp).collect::<Vec<_>>())
.bind(readings.iter().map(|r| r.station_num as i32).collect::<Vec<_>>())
.fetch_one(pool)
.await?)
}
pub async fn save_hii_rain(
pool: &PgPool,
records: &[RainRecord],
now: NaiveDateTime,
) -> Result<usize> {
if records.is_empty() {
return Ok(0);
}
let st = dedupe_last(records, |r| r.station_id);
let ms = dedupe_last(records, |r| (r.station_id, r.timestamp));
let mut tx = pool.begin().await?;
sqlx::query(
"INSERT INTO hii_rain_stations (id, oldcode, name_th, name_en, latitude, longitude, sub_basin_id, agency, updated_at)
SELECT * FROM UNNEST($1::int4[], $2::text[], $3::text[], $4::text[], $5::float8[], $6::float8[], $7::text[], $8::text[], $9::timestamp[])
ON CONFLICT (id) DO UPDATE SET
oldcode = EXCLUDED.oldcode, name_th = EXCLUDED.name_th, name_en = EXCLUDED.name_en,
latitude = EXCLUDED.latitude, longitude = EXCLUDED.longitude,
sub_basin_id = EXCLUDED.sub_basin_id, agency = EXCLUDED.agency, updated_at = EXCLUDED.updated_at",
)
.bind(st.iter().map(|r| r.station_id as i32).collect::<Vec<_>>())
.bind(st.iter().map(|r| r.oldcode.clone()).collect::<Vec<_>>())
.bind(st.iter().map(|r| r.name_th.clone()).collect::<Vec<_>>())
.bind(st.iter().map(|r| r.name_en.clone()).collect::<Vec<_>>())
.bind(st.iter().map(|r| r.latitude).collect::<Vec<_>>())
.bind(st.iter().map(|r| r.longitude).collect::<Vec<_>>())
.bind(st.iter().map(|r| r.sub_basin_id.clone()).collect::<Vec<_>>())
.bind(st.iter().map(|r| r.agency.clone()).collect::<Vec<_>>())
.bind(vec![now; st.len()])
.execute(&mut *tx)
.await
.context("upserting hii_rain_stations")?;
sqlx::query(
"INSERT INTO hii_rainfall (station_id, timestamp, rain_1h, rain_24h)
SELECT * FROM UNNEST($1::int4[], $2::timestamp[], $3::float8[], $4::float8[])
ON CONFLICT (station_id, timestamp) DO UPDATE SET
rain_1h = EXCLUDED.rain_1h, rain_24h = EXCLUDED.rain_24h",
)
.bind(ms.iter().map(|r| r.station_id as i32).collect::<Vec<_>>())
.bind(ms.iter().map(|r| r.timestamp).collect::<Vec<_>>())
.bind(ms.iter().map(|r| r.rain_1h).collect::<Vec<_>>())
.bind(ms.iter().map(|r| r.rain_24h).collect::<Vec<_>>())
.execute(&mut *tx)
.await
.context("upserting hii_rainfall")?;
tx.commit().await?;
Ok(ms.len())
}
pub async fn save_hii_waterlevel(
pool: &PgPool,
records: &[WaterLevelRecord],
now: NaiveDateTime,
) -> Result<usize> {
if records.is_empty() {
return Ok(0);
}
let st = dedupe_last(records, |r| r.station_id);
let ms = dedupe_last(records, |r| (r.station_id, r.timestamp));
let mut tx = pool.begin().await?;
sqlx::query(
"INSERT INTO hii_wl_stations (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, updated_at)
SELECT * FROM UNNEST($1::int4[], $2::text[], $3::text[], $4::text[], $5::text[], $6::float8[],
$7::float8[], $8::text[], $9::text[], $10::float8[], $11::float8[], $12::float8[], $13::float8[],
$14::float8[], $15::float8[], $16::bool[], $17::timestamp[])
ON CONFLICT (id) DO UPDATE SET
oldcode = EXCLUDED.oldcode, rid_code = EXCLUDED.rid_code, name_th = EXCLUDED.name_th,
name_en = EXCLUDED.name_en, latitude = EXCLUDED.latitude, longitude = EXCLUDED.longitude,
agency = EXCLUDED.agency, river_name = EXCLUDED.river_name, offset_msl = EXCLUDED.offset_msl,
ground_level_msl = EXCLUDED.ground_level_msl, min_bank_msl = EXCLUDED.min_bank_msl,
critical_level_msl = EXCLUDED.critical_level_msl, critical_level_m = EXCLUDED.critical_level_m,
qmax = EXCLUDED.qmax, is_key_station = EXCLUDED.is_key_station, updated_at = EXCLUDED.updated_at",
)
.bind(st.iter().map(|r| r.station_id as i32).collect::<Vec<_>>())
.bind(st.iter().map(|r| r.oldcode.clone()).collect::<Vec<_>>())
.bind(st.iter().map(|r| r.rid_code.clone()).collect::<Vec<_>>())
.bind(st.iter().map(|r| r.name_th.clone()).collect::<Vec<_>>())
.bind(st.iter().map(|r| r.name_en.clone()).collect::<Vec<_>>())
.bind(st.iter().map(|r| r.latitude).collect::<Vec<_>>())
.bind(st.iter().map(|r| r.longitude).collect::<Vec<_>>())
.bind(st.iter().map(|r| r.agency.clone()).collect::<Vec<_>>())
.bind(st.iter().map(|r| r.river_name.clone()).collect::<Vec<_>>())
.bind(st.iter().map(|r| r.offset_msl).collect::<Vec<_>>())
.bind(st.iter().map(|r| r.ground_level_msl).collect::<Vec<_>>())
.bind(st.iter().map(|r| r.min_bank_msl).collect::<Vec<_>>())
.bind(st.iter().map(|r| r.critical_level_msl).collect::<Vec<_>>())
.bind(st.iter().map(|r| r.critical_level_m).collect::<Vec<_>>())
.bind(st.iter().map(|r| r.qmax).collect::<Vec<_>>())
.bind(st.iter().map(|r| r.is_key_station).collect::<Vec<_>>())
.bind(vec![now; st.len()])
.execute(&mut *tx)
.await
.context("upserting hii_wl_stations")?;
sqlx::query(
"INSERT INTO hii_waterlevel (station_id, timestamp, wl_msl, wl_m, discharge, flow_rate,
storage_percent, situation_level, diff_wl_bank)
SELECT * FROM UNNEST($1::int4[], $2::timestamp[], $3::float8[], $4::float8[], $5::float8[],
$6::float8[], $7::float8[], $8::int4[], $9::float8[])
ON CONFLICT (station_id, timestamp) DO UPDATE SET
wl_msl = EXCLUDED.wl_msl, wl_m = EXCLUDED.wl_m, discharge = EXCLUDED.discharge,
flow_rate = EXCLUDED.flow_rate, storage_percent = EXCLUDED.storage_percent,
situation_level = EXCLUDED.situation_level, diff_wl_bank = EXCLUDED.diff_wl_bank",
)
.bind(ms.iter().map(|r| r.station_id as i32).collect::<Vec<_>>())
.bind(ms.iter().map(|r| r.timestamp).collect::<Vec<_>>())
.bind(ms.iter().map(|r| r.wl_msl).collect::<Vec<_>>())
.bind(ms.iter().map(|r| r.wl_m).collect::<Vec<_>>())
.bind(ms.iter().map(|r| r.discharge).collect::<Vec<_>>())
.bind(ms.iter().map(|r| r.flow_rate).collect::<Vec<_>>())
.bind(ms.iter().map(|r| r.storage_percent).collect::<Vec<_>>())
.bind(ms.iter().map(|r| r.situation_level).collect::<Vec<_>>())
.bind(ms.iter().map(|r| r.diff_wl_bank).collect::<Vec<_>>())
.execute(&mut *tx)
.await
.context("upserting hii_waterlevel")?;
tx.commit().await?;
Ok(ms.len())
}
/// rid_dams metadata never blanks a stored value (COALESCE), and the daily
/// `level_msl` is preserved too: only the range backfill carries a level, the
/// daily snapshot would NULL it on every cycle.
pub async fn save_dams(pool: &PgPool, records: &[DamRecord], now: NaiveDateTime) -> Result<usize> {
if records.is_empty() {
return Ok(0);
}
let dams = dedupe_last(records, |r| r.dam_id.clone());
let ms = dedupe_last(records, |r| (r.dam_id.clone(), r.date));
let mut tx = pool.begin().await?;
sqlx::query(
"INSERT INTO rid_dams (dam_id, region, name_th, latitude, longitude, capacity_max_mcm,
capacity_normal_mcm, updated_at)
SELECT * FROM UNNEST($1::text[], $2::text[], $3::text[], $4::float8[], $5::float8[],
$6::float8[], $7::float8[], $8::timestamp[])
ON CONFLICT (dam_id) DO UPDATE SET
region = COALESCE(EXCLUDED.region, rid_dams.region),
name_th = COALESCE(EXCLUDED.name_th, rid_dams.name_th),
latitude = COALESCE(EXCLUDED.latitude, rid_dams.latitude),
longitude = COALESCE(EXCLUDED.longitude, rid_dams.longitude),
capacity_max_mcm = COALESCE(EXCLUDED.capacity_max_mcm, rid_dams.capacity_max_mcm),
capacity_normal_mcm = COALESCE(EXCLUDED.capacity_normal_mcm, rid_dams.capacity_normal_mcm),
updated_at = EXCLUDED.updated_at",
)
.bind(dams.iter().map(|r| r.dam_id.clone()).collect::<Vec<_>>())
.bind(dams.iter().map(|r| r.region.clone()).collect::<Vec<_>>())
.bind(dams.iter().map(|r| r.name_th.clone()).collect::<Vec<_>>())
.bind(dams.iter().map(|r| r.latitude).collect::<Vec<_>>())
.bind(dams.iter().map(|r| r.longitude).collect::<Vec<_>>())
.bind(dams.iter().map(|r| r.capacity_max_mcm).collect::<Vec<_>>())
.bind(dams.iter().map(|r| r.capacity_normal_mcm).collect::<Vec<_>>())
.bind(vec![now; dams.len()])
.execute(&mut *tx)
.await
.context("upserting rid_dams")?;
sqlx::query(
"INSERT INTO rid_reservoir_daily (dam_id, date, storage_mcm, storage_pct, inflow_mcm, outflow_mcm, level_msl)
SELECT * FROM UNNEST($1::text[], $2::date[], $3::float8[], $4::float8[], $5::float8[], $6::float8[], $7::float8[])
ON CONFLICT (dam_id, date) DO UPDATE SET
storage_mcm = EXCLUDED.storage_mcm, storage_pct = EXCLUDED.storage_pct,
inflow_mcm = EXCLUDED.inflow_mcm, outflow_mcm = EXCLUDED.outflow_mcm,
level_msl = COALESCE(EXCLUDED.level_msl, rid_reservoir_daily.level_msl)",
)
.bind(ms.iter().map(|r| r.dam_id.clone()).collect::<Vec<_>>())
.bind(ms.iter().map(|r| r.date).collect::<Vec<_>>())
.bind(ms.iter().map(|r| r.storage_mcm).collect::<Vec<_>>())
.bind(ms.iter().map(|r| r.storage_pct).collect::<Vec<_>>())
.bind(ms.iter().map(|r| r.inflow_mcm).collect::<Vec<_>>())
.bind(ms.iter().map(|r| r.outflow_mcm).collect::<Vec<_>>())
.bind(ms.iter().map(|r| r.level_msl).collect::<Vec<_>>())
.execute(&mut *tx)
.await
.context("upserting rid_reservoir_daily")?;
tx.commit().await?;
Ok(ms.len())
}
pub async fn save_rain(pool: &PgPool, rows: &[RainRow]) -> Result<usize> {
if rows.is_empty() {
return Ok(0);
}
let rows = dedupe_last(rows, |r| r.timestamp);
let col = |i: usize| rows.iter().map(|r| r.points[i]).collect::<Vec<_>>();
sqlx::query(
"INSERT INTO openmeteo_rain (timestamp, chiang_dao, mae_taeng, mae_ngat, mae_rim, chiang_mai, catchment_mean)
SELECT * FROM UNNEST($1::timestamp[], $2::float8[], $3::float8[], $4::float8[], $5::float8[], $6::float8[], $7::float8[])
ON CONFLICT (timestamp) DO UPDATE SET
chiang_dao = EXCLUDED.chiang_dao, mae_taeng = EXCLUDED.mae_taeng, mae_ngat = EXCLUDED.mae_ngat,
mae_rim = EXCLUDED.mae_rim, chiang_mai = EXCLUDED.chiang_mai, catchment_mean = EXCLUDED.catchment_mean",
)
.bind(rows.iter().map(|r| r.timestamp).collect::<Vec<_>>())
.bind(col(0))
.bind(col(1))
.bind(col(2))
.bind(col(3))
.bind(col(4))
.bind(rows.iter().map(|r| r.catchment_mean).collect::<Vec<_>>())
.execute(pool)
.await
.context("upserting openmeteo_rain")?;
Ok(rows.len())
}
/// Row counts per table, for `check-db`.
pub async fn table_counts(pool: &PgPool) -> Result<BTreeMap<&'static str, i64>> {
let mut out = BTreeMap::new();
for table in [
"water_measurements",
"hii_rainfall",
"hii_waterlevel",
"rid_reservoir_daily",
"openmeteo_rain",
] {
let n: i64 = sqlx::query_scalar(&format!("SELECT COUNT(*) FROM {table}"))
.fetch_one(pool)
.await?;
out.insert(table, n);
}
Ok(out)
}
+332
View File
@@ -0,0 +1,332 @@
//! Database tests against a real PostgreSQL. Each test runs in its own
//! throwaway schema. Set TEST_DATABASE_URL to run them, e.g.
//! `postgres://postgres:ping@localhost:55432/ping`; without it they are
//! skipped (reported as passed with a note on stderr).
use crate::dams::DamRecord;
use crate::db;
use crate::hii::{RainRecord, WaterLevelRecord};
use crate::rain::RainRow;
use crate::rid::{self, Reading, RidSource};
use crate::stations::Stations;
use anyhow::Result;
use chrono::{Duration, NaiveDate, NaiveDateTime, NaiveTime};
use sqlx::PgPool;
use sqlx::postgres::{PgConnectOptions, PgPoolOptions};
use std::collections::HashMap;
use std::str::FromStr;
use std::sync::Mutex;
struct TestDb {
pool: PgPool,
admin: PgPool,
schema: String,
}
impl TestDb {
async fn new() -> Option<Self> {
let Ok(url) = std::env::var("TEST_DATABASE_URL") else {
eprintln!("TEST_DATABASE_URL not set: skipping database test");
return None;
};
let schema = format!(
"t_{}_{}",
std::process::id(),
chrono::Utc::now().timestamp_nanos_opt().unwrap()
);
let admin = PgPool::connect(&url).await.expect("connect");
sqlx::query(&format!("CREATE SCHEMA {schema}"))
.execute(&admin)
.await
.unwrap();
let opts = PgConnectOptions::from_str(&url)
.unwrap()
.options([("search_path", schema.as_str())]);
let pool = PgPoolOptions::new()
.max_connections(2)
.connect_with(opts)
.await
.unwrap();
db::ensure_schema(&pool).await.expect("schema");
Some(TestDb {
pool,
admin,
schema,
})
}
async fn drop(self) {
self.pool.close().await;
sqlx::query(&format!("DROP SCHEMA {} CASCADE", self.schema))
.execute(&self.admin)
.await
.unwrap();
}
}
fn d(y: i32, m: u32, day: u32) -> NaiveDate {
NaiveDate::from_ymd_opt(y, m, day).unwrap()
}
fn t(y: i32, m: u32, day: u32, h: u32, min: u32) -> NaiveDateTime {
d(y, m, day).and_hms_opt(h, min, 0).unwrap()
}
/// Canned RID: date -> last published hour (1..=24); P.1 (column 8) and P.20
/// (column 1). Records which reports were requested.
struct FakeRid {
published: HashMap<NaiveDate, i64>,
calls: Mutex<Vec<NaiveDate>>,
}
impl FakeRid {
fn new(published: &[(NaiveDate, i64)]) -> Self {
FakeRid {
published: published.iter().copied().collect(),
calls: Mutex::new(Vec::new()),
}
}
fn calls(&self) -> Vec<NaiveDate> {
self.calls.lock().unwrap().clone()
}
}
impl RidSource for FakeRid {
async fn fetch_report(&self, date: NaiveDate) -> Result<Vec<Reading>> {
self.calls.lock().unwrap().push(date);
let last = self.published.get(&date).copied().unwrap_or(0);
Ok((1..=last)
.flat_map(|h| {
let ts = date.and_time(NaiveTime::MIN) + Duration::hours(h);
[8u32, 1].map(|n| Reading {
timestamp: ts,
station_num: n,
water_level: 1.0 + h as f64 / 100.0,
discharge: Some(50.0),
discharge_percent: Some(12.5),
})
})
.collect())
}
}
async fn stored_hours(pool: &PgPool, day: NaiveDate) -> Vec<u32> {
let rec = db::recorded_hours_by_day(pool, day, day).await.unwrap();
let mut h: Vec<u32> = rec
.get(&day)
.cloned()
.unwrap_or_default()
.into_iter()
.collect();
h.sort();
h
}
#[tokio::test]
async fn late_report_is_stored_and_reported_behind() {
let Some(db) = TestDb::new().await else {
return;
};
let today = d(2026, 10, 17);
let yesterday = today - Duration::days(1);
let stations = Stations::bundled();
// yesterday complete already
let full = FakeRid::new(&[(yesterday, 24)]);
rid::fill_gaps(&full, &db.pool, &stations, Some(1), t(2026, 10, 17, 0, 30))
.await
.unwrap();
// 21:30 and RID has published up to 20:00
let src = FakeRid::new(&[(today, 20)]);
let o = rid::run_cycle(&src, &db.pool, &stations, t(2026, 10, 17, 21, 30)).await;
assert!(!o.current);
assert_eq!(o.latest, Some(t(2026, 10, 17, 20, 0)));
assert_eq!(o.saved, 40);
assert_eq!(src.calls(), vec![today]);
assert_eq!(
stored_hours(&db.pool, today).await,
(0..=20).collect::<Vec<_>>()
);
// the 21:00 reading lands; the next poll is current
let src = FakeRid::new(&[(today, 21)]);
let o = rid::run_cycle(&src, &db.pool, &stations, t(2026, 10, 17, 21, 40)).await;
assert!(o.current);
db.drop().await;
}
#[tokio::test]
async fn yesterdays_late_hours_are_recovered_after_midnight() {
let Some(db) = TestDb::new().await else {
return;
};
let today = d(2026, 10, 17);
let yesterday = today - Duration::days(1);
let stations = Stations::bundled();
// stored up to 21:00 yesterday (the evening RID published late)
let partial = FakeRid::new(&[(yesterday, 21)]);
rid::run_cycle(&partial, &db.pool, &stations, t(2026, 10, 16, 21, 5)).await;
assert_eq!(
stored_hours(&db.pool, yesterday).await,
(1..=21).collect::<Vec<_>>()
);
// 02:10: RID now has all of yesterday and today up to 02:00
let src = FakeRid::new(&[(yesterday, 24), (today, 2)]);
let o = rid::run_cycle(&src, &db.pool, &stations, t(2026, 10, 17, 2, 10)).await;
assert!(o.current);
assert_eq!(src.calls(), vec![yesterday, today]);
assert_eq!(
stored_hours(&db.pool, yesterday).await,
(1..=23).collect::<Vec<_>>()
);
assert_eq!(stored_hours(&db.pool, today).await, vec![0, 1, 2]);
// complete now: the next cycle fetches only today's report
let src = FakeRid::new(&[(yesterday, 24), (today, 3)]);
rid::run_cycle(&src, &db.pool, &stations, t(2026, 10, 17, 3, 5)).await;
assert_eq!(src.calls(), vec![today]);
db.drop().await;
}
#[tokio::test]
async fn rid_upsert_revises_values_and_rounds_like_numeric() {
let Some(db) = TestDb::new().await else {
return;
};
let stations = Stations::bundled();
let r = |level: f64| Reading {
timestamp: t(2026, 10, 17, 9, 0),
station_num: 8,
water_level: level,
discharge: None,
discharge_percent: None,
};
db::save_rid(&db.pool, &[r(1.0700000524520874)], &stations)
.await
.unwrap();
db::save_rid(&db.pool, &[r(1.0900000333786)], &stations)
.await
.unwrap();
let rows: Vec<(String, f64, Option<f64>, String)> = sqlx::query_as(
"SELECT s.station_code, w.water_level::float8, w.discharge::float8, w.status
FROM water_measurements w JOIN stations s ON s.id = w.station_id",
)
.fetch_all(&db.pool)
.await
.unwrap();
assert_eq!(rows, vec![("P.1".into(), 1.09, None, "active".into())]);
db.drop().await;
}
#[tokio::test]
async fn hii_dams_and_rain_round_trip() {
let Some(db) = TestDb::new().await else {
return;
};
let now = t(2026, 10, 17, 21, 5);
let rain = RainRecord {
station_id: 11,
oldcode: Some("CHM001".into()),
name_th: None,
name_en: Some("Chiang Mai 1".into()),
latitude: Some(18.79),
longitude: Some(99.0),
sub_basin_id: Some("601".into()),
agency: Some("HII".into()),
timestamp: t(2026, 10, 17, 21, 0),
rain_1h: Some(1.234),
rain_24h: Some(12.0),
};
assert_eq!(
db::save_hii_rain(&db.pool, &[rain.clone(), rain.clone()], now)
.await
.unwrap(),
1
);
let (r1h,): (f64,) = sqlx::query_as("SELECT rain_1h::float8 FROM hii_rainfall")
.fetch_one(&db.pool)
.await
.unwrap();
assert_eq!(r1h, 1.23);
let wl = WaterLevelRecord {
station_id: 5,
oldcode: Some("ridhydro_P.1".into()),
rid_code: Some("P.1".into()),
name_th: Some("สะพานนวรัฐ".into()),
name_en: None,
latitude: None,
longitude: None,
agency: None,
river_name: None,
offset_msl: Some(300.0),
ground_level_msl: None,
min_bank_msl: None,
critical_level_msl: None,
critical_level_m: None,
qmax: None,
is_key_station: true,
timestamp: t(2026, 10, 17, 21, 0),
wl_msl: Some(301.2),
wl_m: None,
discharge: None,
flow_rate: None,
storage_percent: None,
situation_level: Some(2),
diff_wl_bank: None,
};
db::save_hii_waterlevel(&db.pool, &[wl], now).await.unwrap();
let (code, key, sit): (String, bool, i32) = sqlx::query_as(
"SELECT s.rid_code, s.is_key_station, w.situation_level
FROM hii_waterlevel w JOIN hii_wl_stations s ON s.id = w.station_id",
)
.fetch_one(&db.pool)
.await
.unwrap();
assert_eq!((code.as_str(), key, sit), ("P.1", true, 2));
let dam = |name: Option<&str>, level: Option<f64>| DamRecord {
dam_id: "200103".into(),
region: Some("north".into()),
name_th: name.map(str::to_string),
latitude: Some(19.17),
longitude: None,
capacity_max_mcm: Some(265.0),
capacity_normal_mcm: None,
date: d(2026, 10, 17),
storage_mcm: Some(200.0),
storage_pct: Some(75.5),
inflow_mcm: None,
outflow_mcm: Some(1.2),
level_msl: level,
};
db::save_dams(&db.pool, &[dam(Some("แม่งัด"), Some(395.1))], now)
.await
.unwrap();
// a later snapshot without name / level must not blank them
db::save_dams(&db.pool, &[dam(None, None)], now)
.await
.unwrap();
let (name, level): (Option<String>, Option<f64>) = sqlx::query_as(
"SELECT d.name_th, r.level_msl::float8 FROM rid_reservoir_daily r JOIN rid_dams d USING (dam_id)",
)
.fetch_one(&db.pool)
.await
.unwrap();
assert_eq!(name.as_deref(), Some("แม่งัด"));
assert_eq!(level, Some(395.1));
let rows = vec![RainRow {
timestamp: t(2026, 10, 18, 3, 0),
points: [Some(1.0), None, Some(2.0), None, Some(0.333)],
catchment_mean: Some(1.111),
}];
db::save_rain(&db.pool, &rows).await.unwrap();
let (mean, missing): (f64, Option<f64>) =
sqlx::query_as("SELECT catchment_mean::float8, mae_taeng::float8 FROM openmeteo_rain")
.fetch_one(&db.pool)
.await
.unwrap();
assert_eq!((mean, missing), (1.11, None));
db.drop().await;
}
+262
View File
@@ -0,0 +1,262 @@
//! HII / ThaiWater open api-v3 (no key): Ping-basin rain gauges and an
//! independent water-level network (m MSL). Port of src/hii_collector.py.
use crate::util::{lenient_f64, lenient_i64, scalar_string, truthy};
use anyhow::{Context, Result};
use chrono::NaiveDateTime;
use serde_json::Value;
#[derive(Clone, Debug, PartialEq)]
pub struct RainRecord {
pub station_id: i64,
pub oldcode: Option<String>,
pub name_th: Option<String>,
pub name_en: Option<String>,
pub latitude: Option<f64>,
pub longitude: Option<f64>,
pub sub_basin_id: Option<String>,
pub agency: Option<String>,
pub timestamp: NaiveDateTime,
pub rain_1h: Option<f64>,
pub rain_24h: Option<f64>,
}
#[derive(Clone, Debug, PartialEq)]
pub struct WaterLevelRecord {
pub station_id: i64,
pub oldcode: Option<String>,
pub rid_code: Option<String>,
pub name_th: Option<String>,
pub name_en: Option<String>,
pub latitude: Option<f64>,
pub longitude: Option<f64>,
pub agency: Option<String>,
pub river_name: Option<String>,
pub offset_msl: Option<f64>,
pub ground_level_msl: Option<f64>,
pub min_bank_msl: Option<f64>,
pub critical_level_msl: Option<f64>,
pub critical_level_m: Option<f64>,
pub qmax: Option<f64>,
pub is_key_station: bool,
pub timestamp: NaiveDateTime,
pub wl_msl: Option<f64>,
pub wl_m: Option<f64>,
pub discharge: Option<f64>,
pub flow_rate: Option<f64>,
pub storage_percent: Option<f64>,
pub situation_level: Option<i32>,
pub diff_wl_bank: Option<f64>,
}
/// Thai local time, '2026-08-11 13:00' or with seconds.
fn parse_datetime(v: Option<&Value>) -> Option<NaiveDateTime> {
let s = v?.as_str()?;
NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M")
.or_else(|_| NaiveDateTime::parse_from_str(s, "%Y-%m-%d %H:%M:%S"))
.ok()
}
/// `tele_station_name` is {th, en} or a bare (Thai) string.
fn name(station: &Value, lang: &str) -> Option<String> {
match station.get("tele_station_name")? {
Value::Object(o) => o.get(lang).and_then(Value::as_str).map(str::to_string),
Value::String(s) if lang == "th" => Some(s.clone()),
_ => None,
}
}
fn string(v: Option<&Value>) -> Option<String> {
v.and_then(Value::as_str).map(str::to_string)
}
fn agency(row: &Value) -> Option<String> {
string(row.pointer("/agency/agency_shortname/en"))
}
/// 'P.1', 'ridhydro_P.67' -> 'P.1' / 'P.67'; 'ridtele_TUP.14' -> None.
/// (Python: `(?:^|_)(P\.\d+[A-Z]?)$`.)
pub fn rid_code_from_oldcode(oldcode: &str) -> Option<String> {
let seg = oldcode.rsplit('_').next()?;
let digits = seg.strip_prefix("P.")?;
let (num, suffix) = match digits.char_indices().find(|(_, c)| !c.is_ascii_digit()) {
Some((i, _)) => digits.split_at(i),
None => (digits, ""),
};
let suffix_ok = suffix.is_empty()
|| (suffix.chars().count() == 1 && suffix.chars().all(|c| c.is_ascii_uppercase()));
(!num.is_empty() && suffix_ok).then(|| seg.to_string())
}
fn in_basin(row: &Value, basin_code: i64) -> bool {
row.pointer("/basin/basin_code").and_then(Value::as_i64) == Some(basin_code)
}
pub fn parse_rain(payload: &Value, basin_code: i64) -> Vec<RainRecord> {
let rows = payload
.get("data")
.and_then(Value::as_array)
.cloned()
.unwrap_or_default();
rows.iter()
.filter(|row| in_basin(row, basin_code))
.filter_map(|row| {
let station = row.get("station").unwrap_or(&Value::Null);
Some(RainRecord {
station_id: lenient_i64(station.get("id"))?,
timestamp: parse_datetime(row.get("rainfall_datetime"))?,
oldcode: string(station.get("tele_station_oldcode")),
name_th: name(station, "th"),
name_en: name(station, "en"),
latitude: lenient_f64(station.get("tele_station_lat")),
longitude: lenient_f64(station.get("tele_station_long")),
// Python: str(x or "") or None
sub_basin_id: truthy(station.get("sub_basin_id"))
.then(|| scalar_string(station.get("sub_basin_id")))
.flatten(),
agency: agency(row),
rain_1h: lenient_f64(row.get("rain_1h")),
rain_24h: lenient_f64(row.get("rain_24h")),
})
})
.collect()
}
pub fn parse_waterlevel(payload: &Value, basin_code: i64) -> Vec<WaterLevelRecord> {
let rows = payload
.pointer("/waterlevel_data/data")
.and_then(Value::as_array)
.cloned()
.unwrap_or_default();
rows.iter()
.filter(|row| in_basin(row, basin_code))
.filter_map(|row| {
let station = row.get("station").unwrap_or(&Value::Null);
let oldcode = string(station.get("tele_station_oldcode"));
Some(WaterLevelRecord {
station_id: lenient_i64(station.get("id"))?,
timestamp: parse_datetime(row.get("waterlevel_datetime"))?,
rid_code: oldcode.as_deref().and_then(rid_code_from_oldcode),
oldcode,
name_th: name(station, "th"),
name_en: name(station, "en"),
latitude: lenient_f64(station.get("tele_station_lat")),
longitude: lenient_f64(station.get("tele_station_long")),
agency: agency(row),
river_name: string(row.get("river_name")),
offset_msl: lenient_f64(station.get("offset")),
ground_level_msl: lenient_f64(station.get("ground_level")),
min_bank_msl: lenient_f64(station.get("min_bank")),
critical_level_msl: lenient_f64(station.get("critical_level_msl")),
critical_level_m: lenient_f64(station.get("critical_level_m")),
qmax: lenient_f64(station.get("qmax")),
is_key_station: truthy(station.get("is_key_station")),
wl_msl: lenient_f64(row.get("waterlevel_msl")),
wl_m: lenient_f64(row.get("waterlevel_m")),
discharge: lenient_f64(row.get("discharge")),
flow_rate: lenient_f64(row.get("flow_rate")),
storage_percent: lenient_f64(row.get("storage_percent")),
situation_level: lenient_i64(row.get("situation_level")).map(|v| v as i32),
diff_wl_bank: lenient_f64(row.get("diff_wl_bank")),
})
})
.collect()
}
pub struct HiiClient {
http: reqwest::Client,
base: String,
}
impl HiiClient {
pub fn new(http: reqwest::Client, base: String) -> Self {
Self {
http,
base: base.trim_end_matches('/').to_string(),
}
}
async fn get(&self, endpoint: &str) -> Result<Value> {
self.http
.get(format!("{}/{endpoint}", self.base))
.timeout(std::time::Duration::from_secs(90))
.send()
.await
.with_context(|| format!("HII {endpoint} request failed"))?
.error_for_status()?
.json()
.await
.with_context(|| format!("HII {endpoint} response is not JSON"))
}
pub async fn fetch_rain(&self, basin_code: i64) -> Result<Vec<RainRecord>> {
Ok(parse_rain(&self.get("rain_24h").await?, basin_code))
}
pub async fn fetch_waterlevel(&self, basin_code: i64) -> Result<Vec<WaterLevelRecord>> {
Ok(parse_waterlevel(
&self.get("waterlevel_load").await?,
basin_code,
))
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn rid_codes() {
assert_eq!(rid_code_from_oldcode("P.1").as_deref(), Some("P.1"));
assert_eq!(
rid_code_from_oldcode("ridhydro_P.67").as_deref(),
Some("P.67")
);
assert_eq!(rid_code_from_oldcode("x_P.4A").as_deref(), Some("P.4A"));
assert_eq!(rid_code_from_oldcode("ridtele_TUP.14"), None);
assert_eq!(rid_code_from_oldcode("P.4AB"), None);
assert_eq!(rid_code_from_oldcode("CHM001"), None);
}
#[test]
fn rain_rows_filtered_to_basin() {
let payload = json!({"data": [
{"basin": {"basin_code": 6}, "station": {"id": 11, "tele_station_oldcode": "CHM001",
"tele_station_name": {"th": "ชม1", "en": "Chiang Mai 1"}, "tele_station_lat": "18.79",
"tele_station_long": 99.0, "sub_basin_id": 601},
"agency": {"agency_shortname": {"en": "HII"}},
"rainfall_datetime": "2026-08-11 13:00", "rain_1h": "1.5", "rain_24h": 12},
{"basin": {"basin_code": 7}, "station": {"id": 12}, "rainfall_datetime": "2026-08-11 13:00"},
{"basin": {"basin_code": 6}, "station": {"id": 13}, "rainfall_datetime": "bad"}
]});
let r = parse_rain(&payload, 6);
assert_eq!(r.len(), 1);
assert_eq!(r[0].station_id, 11);
assert_eq!(r[0].name_en.as_deref(), Some("Chiang Mai 1"));
assert_eq!(r[0].sub_basin_id.as_deref(), Some("601"));
assert_eq!(r[0].agency.as_deref(), Some("HII"));
assert_eq!(r[0].rain_1h, Some(1.5));
assert_eq!(r[0].latitude, Some(18.79));
}
#[test]
fn waterlevel_rows() {
let payload = json!({"waterlevel_data": {"data": [
{"basin": {"basin_code": 6},
"station": {"id": 5, "tele_station_oldcode": "ridhydro_P.1", "tele_station_name": "สะพานนวรัฐ",
"offset": "300.0", "is_key_station": 1},
"waterlevel_datetime": "2026-08-11 13:00:00", "waterlevel_msl": "301.2",
"situation_level": "2", "river_name": "แม่น้ำปิง"}
]}});
let r = parse_waterlevel(&payload, 6);
assert_eq!(r.len(), 1);
assert_eq!(r[0].rid_code.as_deref(), Some("P.1"));
assert_eq!(r[0].name_th.as_deref(), Some("สะพานนวรัฐ"));
assert_eq!(r[0].name_en, None);
assert!(r[0].is_key_station);
assert_eq!(r[0].situation_level, Some(2));
assert_eq!(r[0].wl_msl, Some(301.2));
assert_eq!(r[0].offset_msl, Some(300.0));
}
}
+117
View File
@@ -0,0 +1,117 @@
//! ping-collector: the Ping River Monitor's data collection service.
//!
//! Replaces the collection half of the Python web app (RID hourly gauges,
//! HII/ThaiWater rain + water level, RID dams, Open-Meteo rain). Writes the
//! same PostgreSQL tables; the Python API keeps forecasting and notifying
//! with `COLLECTION_ENABLED=false`. See docs/RUST_REWRITE.md.
mod config;
mod dams;
mod db;
#[cfg(test)]
mod db_tests;
mod hii;
mod rain;
mod rid;
mod service;
mod stations;
mod util;
use anyhow::{Result, bail};
use clap::{Parser, Subcommand};
use std::path::PathBuf;
#[derive(Parser)]
#[command(version, about)]
struct Cli {
/// .env file to load (default: ./.env when present). Real environment
/// variables take precedence.
#[arg(long, global = true)]
env_file: Option<PathBuf>,
#[command(subcommand)]
command: Command,
}
#[derive(Subcommand)]
enum Command {
/// Run the collection service (what systemd starts)
Run,
/// One cycle of every source, then exit
Once {
/// Fetch and compare with the database, write nothing
#[arg(long)]
dry_run: bool,
},
/// Re-fetch RID reports that have holes: a number of days back, or "all"
FillGaps { days: String },
/// Connect, create missing tables, print row counts and the newest RID reading
CheckDb,
}
#[tokio::main]
async fn main() -> Result<()> {
let cli = Cli::parse();
match &cli.env_file {
Some(path) => {
dotenvy::from_path(path)?;
}
None => {
dotenvy::dotenv().ok();
}
}
let level = std::env::var("RUST_LOG")
.or_else(|_| std::env::var("LOG_LEVEL").map(|l| l.to_lowercase()))
.unwrap_or_else(|_| "info".into());
tracing_subscriber::fmt()
.with_env_filter(tracing_subscriber::EnvFilter::new(format!(
"warn,ping_collector={level}"
)))
.with_target(false)
.init();
let cfg = config::Config::from_env()?;
tracing::info!(
"ping-collector {} -> {}",
env!("CARGO_PKG_VERSION"),
cfg.redacted_database_url()
);
let pool = db::connect(&cfg.database_url).await?;
let stations = stations::Stations::load(&cfg.station_config_path)?;
let collector = service::Collector::new(cfg, pool, stations)?;
match cli.command {
Command::Run => collector.run().await,
Command::Once { dry_run } => collector.once(dry_run).await,
Command::FillGaps { days } => {
let days = match days.as_str() {
"all" => None,
n => match n.parse::<u32>() {
Ok(d) => Some(d),
Err(_) => bail!("fill-gaps takes a number of days or \"all\", got {n:?}"),
},
};
db::ensure_schema(&collector.pool).await?;
let filled = rid::fill_gaps(
&collector.rid,
&collector.pool,
&collector.stations,
days,
util::bangkok_now(),
)
.await?;
println!("{filled} readings re-fetched");
Ok(())
}
Command::CheckDb => {
db::ensure_schema(&collector.pool).await?;
for (table, n) in db::table_counts(&collector.pool).await? {
println!("{table:<22} {n:>10} rows");
}
println!(
"newest RID reading: {:?}",
db::latest_rid_timestamp(&collector.pool).await?
);
Ok(())
}
}
}
+133
View File
@@ -0,0 +1,133 @@
//! Open-Meteo hourly precipitation at five upper-Ping catchment points
//! (recent days + next 48 h) into `openmeteo_rain`. Port of the live path of
//! src/ml/rain.py (`fetch_forecast` + `save_to_db`). Times are Asia/Bangkok.
use anyhow::{Context, Result};
use chrono::NaiveDateTime;
use serde_json::Value;
use std::collections::BTreeMap;
/// (column, lat, lon), headwaters to city. Order = column order in the table.
pub const CATCHMENT_POINTS: [(&str, f64, f64); 5] = [
("chiang_dao", 19.37, 98.97),
("mae_taeng", 19.12, 98.94),
("mae_ngat", 19.17, 99.05),
("mae_rim", 18.92, 98.92),
("chiang_mai", 18.79, 99.00),
];
#[derive(Clone, Debug, PartialEq)]
pub struct RainRow {
pub timestamp: NaiveDateTime,
pub points: [Option<f64>; 5],
/// Mean of the points that have a value (pandas `mean(axis=1)`).
pub catchment_mean: Option<f64>,
}
/// A multi-coordinate request returns a list, one result per point.
pub fn parse(payload: &Value) -> Vec<RainRow> {
let results: Vec<&Value> = match payload {
Value::Array(a) => a.iter().collect(),
other => vec![other],
};
let mut by_time: BTreeMap<NaiveDateTime, [Option<f64>; 5]> = BTreeMap::new();
for (i, result) in results.iter().take(CATCHMENT_POINTS.len()).enumerate() {
let hourly = result.get("hourly");
let times = hourly.and_then(|h| h.get("time")).and_then(Value::as_array);
let values = hourly
.and_then(|h| h.get("precipitation"))
.and_then(Value::as_array);
let (Some(times), Some(values)) = (times, values) else {
continue;
};
for (t, v) in times.iter().zip(values) {
let Some(ts) = t
.as_str()
.and_then(|s| NaiveDateTime::parse_from_str(s, "%Y-%m-%dT%H:%M").ok())
else {
continue;
};
by_time.entry(ts).or_insert([None; 5])[i] = v.as_f64().filter(|x| x.is_finite());
}
}
by_time
.into_iter()
.map(|(timestamp, points)| {
let present: Vec<f64> = points.iter().flatten().copied().collect();
let catchment_mean =
(!present.is_empty()).then(|| present.iter().sum::<f64>() / present.len() as f64);
RainRow {
timestamp,
points,
catchment_mean,
}
})
.collect()
}
pub struct OpenMeteoClient {
http: reqwest::Client,
url: String,
}
impl OpenMeteoClient {
pub fn new(http: reqwest::Client, url: String) -> Self {
Self { http, url }
}
pub async fn fetch_forecast(&self, past_days: u32, forecast_days: u32) -> Result<Vec<RainRow>> {
let join = |f: fn(&(&str, f64, f64)) -> f64| {
CATCHMENT_POINTS
.iter()
.map(|p| f(p).to_string())
.collect::<Vec<_>>()
.join(",")
};
let payload: Value = self
.http
.get(&self.url)
.query(&[
("latitude", join(|p| p.1)),
("longitude", join(|p| p.2)),
("hourly", "precipitation".into()),
("timezone", "Asia/Bangkok".into()),
("past_days", past_days.to_string()),
("forecast_days", forecast_days.to_string()),
])
.timeout(std::time::Duration::from_secs(60))
.send()
.await
.context("Open-Meteo request failed")?
.error_for_status()?
.json()
.await
.context("Open-Meteo response is not JSON")?;
Ok(parse(&payload))
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn aligns_points_and_averages_present_values() {
let point = |vals: Value| json!({"hourly": {"time": ["2026-10-17T00:00", "2026-10-17T01:00"], "precipitation": vals}});
let payload = json!([
point(json!([1.0, null])),
point(json!([3.0, null])),
point(json!([null, null])),
point(json!([2.0, null])),
point(json!([0.0, 4.0])),
]);
let rows = parse(&payload);
assert_eq!(rows.len(), 2);
assert_eq!(
rows[0].points,
[Some(1.0), Some(3.0), None, Some(2.0), Some(0.0)]
);
assert_eq!(rows[0].catchment_mean, Some(1.5));
assert_eq!(rows[1].catchment_mean, Some(4.0));
}
}
+435
View File
@@ -0,0 +1,435 @@
//! RID hourly water-level/discharge report (16 Ping gauges).
//!
//! `POST getGroupHourlyWaterLevelReportAllHL.ashx` with a Buddhist-era date
//! returns 24 rows (`hourlytime` "1.00".."24.00"; hour 24 is 00:00 of the next
//! day) with `wlvaluesN` / `qvaluesN` / `QPercentN` for report column N=1..16.
//! Hours not published yet are null.
//!
//! The report often lags 1-2 h in the evening. Each cycle therefore stores
//! every reading it gets (not only "the current hour"), and re-fetches
//! yesterday's report while the store still has holes in it: hours 22-24 of a
//! day are reachable only through that day's report.
use crate::db;
use crate::stations::Stations;
use crate::util::{floor_hour, lenient_f64};
use anyhow::{Context, Result};
use chrono::{Datelike, Duration, NaiveDate, NaiveDateTime, NaiveTime};
use serde_json::Value;
use sqlx::PgPool;
use std::collections::{BTreeSet, HashMap, HashSet};
use std::future::Future;
pub const N_STATIONS: u32 = 16;
// Plausibility ranges (src/validators.py)
const WATER_LEVEL_RANGE: (f64, f64) = (-10.0, 50.0);
const DISCHARGE_RANGE: (f64, f64) = (0.0, 10_000.0);
// % of channel capacity; Oct 2024 peaked at 226 %. Out of range -> NULL,
// the reading itself is kept.
const DISCHARGE_PERCENT_RANGE: (f64, f64) = (0.0, 500.0);
#[derive(Clone, Debug, PartialEq)]
pub struct Reading {
pub timestamp: NaiveDateTime,
/// RID report column 1..16 (= stations.id in the database)
pub station_num: u32,
pub water_level: f64,
pub discharge: Option<f64>,
pub discharge_percent: Option<f64>,
}
fn within(x: f64, (lo, hi): (f64, f64)) -> bool {
lo <= x && x <= hi
}
/// Parse one report. Rows with an unparsable hour, a missing/implausible
/// water level or an implausible discharge are dropped, as in Python.
pub fn parse_report(report_date: NaiveDate, payload: &Value) -> Vec<Reading> {
let mut out = Vec::new();
let Some(rows) = payload.get("rows").and_then(Value::as_array) else {
return out;
};
for row in rows {
let hour = match row.get("hourlytime") {
Some(Value::String(s)) => s.trim().parse::<f64>().ok(),
Some(Value::Number(n)) => n.as_f64(),
_ => None,
};
let Some(hour) = hour
.map(|h| h.trunc() as i64)
.filter(|h| (1..=24).contains(h))
else {
continue;
};
let timestamp = report_date.and_time(NaiveTime::MIN) + Duration::hours(hour);
for n in 1..=N_STATIONS {
let Some(water_level) = lenient_f64(row.get(format!("wlvalues{n}"))) else {
continue;
};
if !within(water_level, WATER_LEVEL_RANGE) {
tracing::warn!(
"RID {timestamp} column {n}: water level {water_level} out of range"
);
continue;
}
let discharge = lenient_f64(row.get(format!("qvalues{n}")));
if let Some(q) = discharge
&& !within(q, DISCHARGE_RANGE)
{
tracing::warn!("RID {timestamp} column {n}: discharge {q} out of range");
continue;
}
let discharge_percent = discharge
.and(lenient_f64(row.get(format!("QPercent{n}"))))
.filter(|p| within(*p, DISCHARGE_PERCENT_RANGE));
out.push(Reading {
timestamp,
station_num: n,
water_level,
discharge,
discharge_percent,
});
}
}
out
}
/// Anything that can produce a parsed RID report (HTTP in production,
/// canned rows in tests).
pub trait RidSource {
fn fetch_report(&self, date: NaiveDate) -> impl Future<Output = Result<Vec<Reading>>> + Send;
}
pub struct RidClient {
http: reqwest::Client,
url: String,
}
impl RidClient {
pub fn new(http: reqwest::Client, url: String) -> Self {
Self { http, url }
}
}
impl RidSource for RidClient {
async fn fetch_report(&self, date: NaiveDate) -> Result<Vec<Reading>> {
let thai_date = format!(
"{:02}/{:02}/{}",
date.day(),
date.month(),
date.year() + 543
);
let nd = chrono::Utc::now().timestamp_millis().to_string();
let form = [
("DW[UtokID]", "1"),
("DW[BasinID]", "6"),
("DW[TimeCurrent]", thai_date.as_str()),
("_search", "false"),
("nd", nd.as_str()),
("rows", "100"),
("page", "1"),
("sidx", "indexhourly"),
("sord", "asc"),
];
let payload: Value = self
.http
.post(&self.url)
.header("X-Requested-With", "XMLHttpRequest")
.header("Accept", "application/json, text/javascript, */*; q=0.01")
.form(&form)
.timeout(std::time::Duration::from_secs(30))
.send()
.await
.context("RID request failed")?
.error_for_status()?
.json()
.await
.context("RID response is not JSON")?;
Ok(parse_report(date, &payload))
}
}
/// Report dates whose re-fetch would fill holes in [start, end].
///
/// Hours 1-23 of day D come from D's report, hour 0 from the previous day's
/// report ("hour 24"). For today only hours already passed are expected.
pub fn gap_fetch_dates(
start: NaiveDate,
end: NaiveDate,
now: NaiveDateTime,
recorded: &HashMap<NaiveDate, HashSet<u32>>,
) -> BTreeSet<NaiveDate> {
let mut dates = BTreeSet::new();
let empty = HashSet::new();
let mut day = start;
while day <= end {
let expected: Vec<u32> = if day == now.date() {
(0..chrono::Timelike::hour(&now)).collect()
} else {
(0..24).collect()
};
let have = recorded.get(&day).unwrap_or(&empty);
let missing: Vec<u32> = expected.into_iter().filter(|h| !have.contains(h)).collect();
if missing.iter().any(|&h| h >= 1) {
dates.insert(day);
}
if missing.contains(&0) && day > start {
dates.insert(day - Duration::days(1));
}
day += Duration::days(1);
}
dates
}
/// Reports to fetch in one cycle: the one holding the current hour (today's,
/// or before 01:00 yesterday's) plus yesterday's while it has holes.
pub async fn cycle_dates(pool: &PgPool, now: NaiveDateTime) -> Vec<NaiveDate> {
let today = now.date();
let yesterday = today - Duration::days(1);
let mut dates = BTreeSet::from([if chrono::Timelike::hour(&now) >= 1 {
today
} else {
yesterday
}]);
match db::recorded_hours_by_day(pool, yesterday, today).await {
Ok(recorded) => dates.extend(gap_fetch_dates(yesterday, today, now, &recorded)),
Err(e) => tracing::warn!("gap detection failed, fetching the current report only: {e:#}"),
}
dates.into_iter().collect()
}
#[derive(Clone, Debug, Default)]
pub struct CycleOutcome {
pub dates: Vec<NaiveDate>,
pub fetched: usize,
pub saved: usize,
pub latest: Option<NaiveDateTime>,
/// The current hour's reading is in: poll again next hour.
pub current: bool,
pub error: Option<String>,
}
pub async fn fetch_dates(
src: &impl RidSource,
dates: &[NaiveDate],
) -> (Vec<Reading>, Option<String>) {
let mut readings = Vec::new();
let mut error = None;
for &d in dates {
match src.fetch_report(d).await {
Ok(r) => readings.extend(r),
Err(e) => {
tracing::error!("RID report {d}: {e:#}");
error = Some(format!("{e:#}"));
}
}
}
(readings, error)
}
/// One RID poll: fetch, store everything, report whether the feed is current.
pub async fn run_cycle(
src: &impl RidSource,
pool: &PgPool,
stations: &Stations,
now: NaiveDateTime,
) -> CycleOutcome {
let mut out = CycleOutcome {
dates: cycle_dates(pool, now).await,
..Default::default()
};
let (readings, fetch_error) = fetch_dates(src, &out.dates).await;
out.fetched = readings.len();
out.error = fetch_error;
if readings.is_empty() {
out.error.get_or_insert_with(|| "no data fetched".into());
tracing::warn!("RID: no data fetched for {:?}", out.dates);
return out;
}
if let Err(e) = db::save_rid(pool, &readings, stations).await {
tracing::error!("RID save failed: {e:#}");
out.error = Some(format!("save failed: {e:#}"));
return out;
}
out.saved = readings.len();
let latest = readings
.iter()
.map(|r| r.timestamp)
.max()
.expect("non-empty");
let current_hour = floor_hour(now);
out.latest = Some(latest);
out.current = latest >= current_hour;
if out.current {
tracing::info!("RID: {} readings saved, latest {latest}", out.saved);
} else {
tracing::warn!(
"RID: {} readings saved up to {latest}; {current_hour} not published yet ({} h behind), will retry",
out.saved,
(current_hour - latest).num_hours()
);
}
out
}
/// Re-fetch every report with holes in the last `days` days (None = whole
/// history). Returns readings re-saved.
pub async fn fill_gaps(
src: &impl RidSource,
pool: &PgPool,
stations: &Stations,
days: Option<u32>,
now: NaiveDateTime,
) -> Result<usize> {
let end = now.date();
let start = match days {
Some(d) => end - Duration::days(d as i64),
None => match db::earliest_rid_timestamp(pool).await? {
Some(t) => t.date(),
None => {
tracing::warn!("database is empty; nothing to gap-fill");
return Ok(0);
}
},
};
let recorded = db::recorded_hours_by_day(pool, start, end).await?;
let dates = gap_fetch_dates(start, end, now, &recorded);
tracing::info!(
"gap fill {start}..{end}: {} report(s) to re-fetch",
dates.len()
);
let mut filled = 0;
for (i, d) in dates.iter().enumerate() {
if i > 0 {
tokio::time::sleep(std::time::Duration::from_secs(1)).await; // be polite
}
match src.fetch_report(*d).await {
Ok(readings) if !readings.is_empty() => {
db::save_rid(pool, &readings, stations).await?;
filled += readings.len();
}
Ok(_) => tracing::info!("gap fill {d}: RID has no data"),
Err(e) => tracing::warn!("gap fill {d}: {e:#}"),
}
}
tracing::info!("gap fill done: {filled} readings re-saved");
Ok(filled)
}
/// Next RID poll: `offset` minutes past the next hour when current, else in
/// `retry` minutes (never later than that hourly poll).
pub fn next_poll_time(
now: NaiveDateTime,
current: bool,
offset_min: u32,
retry_min: u32,
) -> NaiveDateTime {
let on_schedule = floor_hour(now) + Duration::hours(1) + Duration::minutes(offset_min as i64);
if current {
on_schedule
} else {
on_schedule.min(now + Duration::minutes(retry_min as i64))
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
fn d(y: i32, m: u32, day: u32) -> NaiveDate {
NaiveDate::from_ymd_opt(y, m, day).unwrap()
}
fn t(y: i32, m: u32, day: u32, h: u32, min: u32) -> NaiveDateTime {
d(y, m, day).and_hms_opt(h, min, 0).unwrap()
}
#[test]
fn parses_levels_discharge_and_percent() {
let p = json!({"rows": [{"hourlytime": "9.00", "wlvalues1": "3.50", "qvalues1": "120.5", "QPercent1": 45.2}]});
let r = parse_report(d(2026, 7, 22), &p);
assert_eq!(
r,
vec![Reading {
timestamp: t(2026, 7, 22, 9, 0),
station_num: 1,
water_level: 3.5,
discharge: Some(120.5),
discharge_percent: Some(45.2),
}]
);
}
#[test]
fn hour_24_is_next_day_midnight_and_junk_hours_are_skipped() {
let p = json!({"rows": [
{"hourlytime": "24.00", "wlvalues1": 3.0},
{"hourlytime": "25.00", "wlvalues1": 3.0},
{"hourlytime": "0.00", "wlvalues1": 3.0},
{"hourlytime": "", "wlvalues1": 3.0}
]});
let r = parse_report(d(2026, 7, 22), &p);
assert_eq!(r.len(), 1);
assert_eq!(r[0].timestamp, t(2026, 7, 23, 0, 0));
}
#[test]
fn missing_or_junk_values() {
let p = json!({"rows": [{"hourlytime": "10.00",
"wlvalues1": null,
"wlvalues2": 2.5, "qvalues2": "***", "QPercent2": 10,
"wlvalues3": 80.0,
"wlvalues4": 1.0, "qvalues4": "20", "QPercent4": 900,
"wlvalues5": 1.0, "qvalues5": "-5"}]});
let r = parse_report(d(2026, 7, 22), &p);
let by: HashMap<u32, &Reading> = r.iter().map(|x| (x.station_num, x)).collect();
assert!(!by.contains_key(&1)); // no level
assert_eq!(by[&2].discharge, None); // "***"
assert_eq!(by[&2].discharge_percent, None); // percent needs a discharge
assert!(!by.contains_key(&3)); // level out of range
assert_eq!(by[&4].discharge_percent, None); // 900 % nulled, reading kept
assert!(!by.contains_key(&5)); // negative discharge drops the reading
}
#[test]
fn missing_rows_key_is_empty() {
assert!(parse_report(d(2026, 7, 22), &json!({"unexpected": true})).is_empty());
}
fn hours(spec: &[(NaiveDate, std::ops::Range<u32>)]) -> HashMap<NaiveDate, HashSet<u32>> {
spec.iter()
.map(|(day, r)| (*day, r.clone().collect()))
.collect()
}
#[test]
fn gap_dates_late_evening_and_midnight() {
let today = d(2026, 10, 17);
let yesterday = d(2026, 10, 16);
// 02:10, yesterday has 0..22, today only hour 1: yesterday's report holds
// 22:00, 23:00 and today's 00:00 ("hour 24")
let rec = hours(&[(yesterday, 0..22), (today, 1..2)]);
let got = gap_fetch_dates(yesterday, today, t(2026, 10, 17, 2, 10), &rec);
assert_eq!(got.into_iter().collect::<Vec<_>>(), vec![yesterday]);
// today's 01:00 missing as well -> today's report too
let rec = hours(&[(yesterday, 0..22)]);
let got = gap_fetch_dates(yesterday, today, t(2026, 10, 17, 2, 10), &rec);
assert_eq!(got.into_iter().collect::<Vec<_>>(), vec![yesterday, today]);
// complete yesterday, today complete up to now -> nothing
let rec = hours(&[(yesterday, 0..24), (today, 0..15)]);
assert!(gap_fetch_dates(yesterday, today, t(2026, 10, 17, 15, 5), &rec).is_empty());
}
#[test]
fn poll_schedule() {
let now = t(2026, 10, 17, 21, 6);
assert_eq!(next_poll_time(now, true, 5, 10), t(2026, 10, 17, 22, 5));
assert_eq!(next_poll_time(now, false, 5, 10), t(2026, 10, 17, 21, 16));
assert_eq!(
next_poll_time(t(2026, 10, 17, 21, 58), false, 5, 10),
t(2026, 10, 17, 22, 5)
);
}
}
+273
View File
@@ -0,0 +1,273 @@
//! The long-running collection loop and the one-shot commands.
use crate::config::{Config, USER_AGENT};
use crate::dams::DamsClient;
use crate::db;
use crate::hii::HiiClient;
use crate::rain::OpenMeteoClient;
use crate::rid::{self, CycleOutcome, RidClient};
use crate::stations::Stations;
use crate::util::{bangkok_now, floor_hour};
use anyhow::Result;
use chrono::{Duration, NaiveDateTime};
use serde_json::json;
use sqlx::PgPool;
pub struct Collector {
pub cfg: Config,
pub pool: PgPool,
pub stations: Stations,
pub rid: RidClient,
hii: HiiClient,
dams: DamsClient,
rain: OpenMeteoClient,
}
#[derive(Default)]
struct Stats {
total_runs: u64,
successful_runs: u64,
failed_runs: u64,
last_run: Option<NaiveDateTime>,
next_run: Option<NaiveDateTime>,
feed_current: Option<bool>,
latest_reading: Option<NaiveDateTime>,
readings_saved: usize,
}
fn iso(t: Option<NaiveDateTime>) -> serde_json::Value {
t.map_or(serde_json::Value::Null, |t| {
json!(t.format("%Y-%m-%dT%H:%M:%S").to_string())
})
}
impl Collector {
pub fn new(cfg: Config, pool: PgPool, stations: Stations) -> Result<Self> {
let http = reqwest::Client::builder()
.user_agent(USER_AGENT)
.connect_timeout(std::time::Duration::from_secs(20))
.build()?;
Ok(Self {
rid: RidClient::new(http.clone(), cfg.rid_api_url.clone()),
hii: HiiClient::new(http.clone(), cfg.hii_api_base.clone()),
dams: DamsClient::new(http.clone(), cfg.dams_url.clone()),
rain: OpenMeteoClient::new(http, cfg.openmeteo_url.clone()),
cfg,
pool,
stations,
})
}
/// HII rain + water level, RID dams (today and yesterday), Open-Meteo.
/// Each source is failure-isolated.
pub async fn collect_hourly(&self, now: NaiveDateTime) {
if self.cfg.hii_enabled {
match self.hii.fetch_rain(self.cfg.hii_basin_code).await {
Ok(r) => match db::save_hii_rain(&self.pool, &r, now).await {
Ok(n) => tracing::info!("HII rainfall: {n} rows saved"),
Err(e) => tracing::error!("HII rainfall save failed: {e:#}"),
},
Err(e) => tracing::error!("HII rainfall fetch failed: {e:#}"),
}
match self.hii.fetch_waterlevel(self.cfg.hii_basin_code).await {
Ok(r) => match db::save_hii_waterlevel(&self.pool, &r, now).await {
Ok(n) => tracing::info!("HII water level: {n} rows saved"),
Err(e) => tracing::error!("HII water level save failed: {e:#}"),
},
Err(e) => tracing::error!("HII water level fetch failed: {e:#}"),
}
}
if self.cfg.dams_enabled {
for date in [None, Some(now.date() - Duration::days(1))] {
match self.dams.fetch_day(date).await {
Ok(r) => match db::save_dams(&self.pool, &r, now).await {
Ok(n) => tracing::info!("RID dams {date:?}: {n} dam-day rows saved"),
Err(e) => tracing::error!("RID dams save failed: {e:#}"),
},
Err(e) => tracing::error!("RID dams fetch {date:?} failed: {e:#}"),
}
}
}
if self.cfg.rain_enabled {
match self.rain.fetch_forecast(5, 2).await {
Ok(r) => match db::save_rain(&self.pool, &r).await {
Ok(n) => tracing::info!("Open-Meteo rain: {n} hourly rows upserted"),
Err(e) => tracing::error!("Open-Meteo save failed: {e:#}"),
},
Err(e) => tracing::error!("Open-Meteo fetch failed: {e:#}"),
}
}
}
fn write_status(&self, stats: &Stats, running: bool) {
let Some(path) = &self.cfg.status_file else {
return;
};
let body = json!({
"is_running": running,
"total_runs": stats.total_runs,
"successful_runs": stats.successful_runs,
"failed_runs": stats.failed_runs,
"last_run": iso(stats.last_run),
"next_run": iso(stats.next_run),
"feed_current": stats.feed_current,
"latest_reading": iso(stats.latest_reading),
"readings_saved": stats.readings_saved,
"collector": concat!("ping-collector ", env!("CARGO_PKG_VERSION")),
});
let tmp = path.with_extension("json.tmp");
let result =
std::fs::write(&tmp, body.to_string()).and_then(|_| std::fs::rename(&tmp, path));
if let Err(e) = result {
tracing::warn!("could not write status file {}: {e}", path.display());
}
}
/// The service: startup gap fill, then RID polls on the hour (+offset)
/// with retries while RID is behind, and the hourly sources once per
/// clock hour. Returns when a shutdown signal arrives.
pub async fn run(&self) -> Result<()> {
db::ensure_schema(&self.pool).await?;
let shutdown = shutdown_signal();
tokio::pin!(shutdown);
if self.cfg.gap_fill_days_on_start > 0 {
let days = Some(self.cfg.gap_fill_days_on_start);
tokio::select! {
r = rid::fill_gaps(&self.rid, &self.pool, &self.stations, days, bangkok_now()) => {
if let Err(e) = r { tracing::error!("startup gap fill failed: {e:#}") }
}
_ = &mut shutdown => return Ok(()),
}
}
let mut stats = Stats::default();
let mut last_hourly_slot: Option<NaiveDateTime> = None;
loop {
let start = bangkok_now();
self.write_status(&stats, true);
let outcome: CycleOutcome =
rid::run_cycle(&self.rid, &self.pool, &self.stations, start).await;
stats.total_runs += 1;
stats.last_run = Some(start);
if outcome.current {
stats.successful_runs += 1;
} else {
stats.failed_runs += 1;
}
stats.feed_current = Some(outcome.current);
stats.readings_saved = outcome.saved;
if outcome.latest.is_some() {
stats.latest_reading = outcome.latest.max(stats.latest_reading);
}
let slot = floor_hour(bangkok_now());
if last_hourly_slot != Some(slot) {
last_hourly_slot = Some(slot);
self.collect_hourly(bangkok_now()).await;
}
let now = bangkok_now();
let next = rid::next_poll_time(
now,
outcome.current,
self.cfg.poll_offset_minutes,
self.cfg.retry_minutes,
);
stats.next_run = Some(next);
self.write_status(&stats, false);
tracing::info!(
"next RID poll at {}{}",
next.format("%H:%M"),
if outcome.current {
""
} else {
" (current hour not published yet)"
}
);
let wait = (next - now)
.to_std()
.unwrap_or_default()
.max(std::time::Duration::from_secs(30));
tokio::select! {
_ = tokio::time::sleep(wait) => {}
_ = &mut shutdown => {
tracing::info!("shutdown signal received");
return Ok(());
}
}
}
}
/// One cycle of every source; with `dry_run` nothing is written and the
/// RID fetch is compared against what the database already holds.
pub async fn once(&self, dry_run: bool) -> Result<()> {
let now = bangkok_now();
if !dry_run {
db::ensure_schema(&self.pool).await?;
let o = rid::run_cycle(&self.rid, &self.pool, &self.stations, now).await;
println!(
"RID: reports {:?}, {} fetched, {} saved, latest {:?}, current hour in: {}{}",
o.dates,
o.fetched,
o.saved,
o.latest,
o.current,
o.error.map(|e| format!(", error: {e}")).unwrap_or_default()
);
self.collect_hourly(now).await;
return Ok(());
}
let dates = rid::cycle_dates(&self.pool, now).await;
let (readings, error) = rid::fetch_dates(&self.rid, &dates).await;
let stored = db::count_stored_rid(&self.pool, &readings).await?;
let latest = readings.iter().map(|r| r.timestamp).max();
let db_latest = db::latest_rid_timestamp(&self.pool).await?;
println!(
"RID (dry run): reports {dates:?}: {} readings fetched, {stored} already stored, {} would be added; latest fetched {latest:?}, latest stored {db_latest:?}{}",
readings.len(),
readings.len() as i64 - stored,
error.map(|e| format!(", error: {e}")).unwrap_or_default()
);
if self.cfg.hii_enabled {
match self.hii.fetch_rain(self.cfg.hii_basin_code).await {
Ok(r) => println!("HII rainfall (dry run): {} Ping-basin stations", r.len()),
Err(e) => println!("HII rainfall (dry run): {e:#}"),
}
match self.hii.fetch_waterlevel(self.cfg.hii_basin_code).await {
Ok(r) => println!("HII water level (dry run): {} Ping-basin stations", r.len()),
Err(e) => println!("HII water level (dry run): {e:#}"),
}
}
if self.cfg.dams_enabled {
match self.dams.fetch_day(None).await {
Ok(r) => println!("RID dams (dry run): {} dams today", r.len()),
Err(e) => println!("RID dams (dry run): {e:#}"),
}
}
if self.cfg.rain_enabled {
match self.rain.fetch_forecast(5, 2).await {
Ok(r) => println!("Open-Meteo (dry run): {} hourly rows", r.len()),
Err(e) => println!("Open-Meteo (dry run): {e:#}"),
}
}
Ok(())
}
}
async fn shutdown_signal() {
#[cfg(unix)]
{
use tokio::signal::unix::{SignalKind, signal};
let mut term = signal(SignalKind::terminate()).expect("SIGTERM handler");
tokio::select! {
_ = tokio::signal::ctrl_c() => {}
_ = term.recv() => {}
}
}
#[cfg(not(unix))]
{
let _ = tokio::signal::ctrl_c().await;
}
}
+84
View File
@@ -0,0 +1,84 @@
//! RID station mapping: report column number (1..16) -> station metadata.
//! Same precedence as the Python scraper: the runtime file
//! (STATION_CONFIG_PATH, edited through the station CRUD API) if it exists,
//! else the defaults bundled in src/data/stations.json.
use anyhow::{Context, Result};
use serde::Deserialize;
use std::collections::BTreeMap;
use std::path::Path;
const BUNDLED: &str = include_str!("../../../src/data/stations.json");
#[derive(Clone, Debug, Deserialize)]
pub struct StationInfo {
pub code: String,
pub thai_name: String,
pub english_name: String,
pub latitude: Option<f64>,
pub longitude: Option<f64>,
pub geohash: Option<String>,
}
#[derive(Clone, Debug)]
pub struct Stations(BTreeMap<u32, StationInfo>);
impl Stations {
pub fn load(runtime_path: &Path) -> Result<Self> {
if runtime_path.exists() {
let text = std::fs::read_to_string(runtime_path)
.with_context(|| format!("reading {}", runtime_path.display()))?;
let stations = Self::parse(&text)
.with_context(|| format!("parsing {}", runtime_path.display()))?;
tracing::info!(
"loaded {} stations from {}",
stations.0.len(),
runtime_path.display()
);
return Ok(stations);
}
let stations = Self::bundled();
tracing::info!("loaded {} bundled default stations", stations.0.len());
Ok(stations)
}
pub fn bundled() -> Self {
Self::parse(BUNDLED).expect("bundled stations.json is valid")
}
fn parse(text: &str) -> Result<Self> {
let raw: BTreeMap<String, StationInfo> = serde_json::from_str(text)?;
let map = raw
.into_iter()
.map(|(k, v)| Ok((k.trim().parse::<u32>()?, v)))
.collect::<Result<BTreeMap<_, _>>>()?;
Ok(Stations(map))
}
/// Metadata for report column `n`, with the Python fallback for unknown
/// columns (`P.{19+n}` / `Station n`).
pub fn get(&self, n: u32) -> StationInfo {
self.0.get(&n).cloned().unwrap_or_else(|| StationInfo {
code: format!("P.{}", 19 + n),
thai_name: format!("Station {n}"),
english_name: format!("Station {n}"),
latitude: None,
longitude: None,
geohash: None,
})
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn bundled_mapping_has_the_16_rid_gauges() {
let s = Stations::bundled();
assert_eq!(s.0.len(), 16);
assert_eq!(s.get(1).code, "P.20");
assert_eq!(s.get(8).code, "P.1");
assert_eq!(s.get(99).code, "P.118"); // Python fallback
}
}
+86
View File
@@ -0,0 +1,86 @@
//! Small shared helpers: Bangkok wall-clock time and lenient JSON numerics.
use chrono::{Duration, NaiveDateTime, Timelike, Utc};
use serde_json::Value;
/// Asia/Bangkok wall-clock time (UTC+7, no DST), without an offset. Every
/// timestamp in the database is stored this way, independent of the
/// server's own time zone.
pub fn bangkok_now() -> NaiveDateTime {
(Utc::now() + Duration::hours(7)).naive_utc()
}
pub fn floor_hour(t: NaiveDateTime) -> NaiveDateTime {
t.with_minute(0)
.and_then(|t| t.with_second(0))
.and_then(|t| t.with_nanosecond(0))
.expect("valid time")
}
/// Upstream APIs send numbers as JSON numbers, numeric strings ("335.00",
/// "1,234.5"), placeholders ("", "-", "***") or null. Anything that is not a
/// finite number becomes None.
pub fn lenient_f64(v: Option<&Value>) -> Option<f64> {
let x = match v? {
Value::Number(n) => n.as_f64()?,
Value::String(s) => s.trim().replace(',', "").parse::<f64>().ok()?,
_ => return None,
};
x.is_finite().then_some(x)
}
pub fn lenient_i64(v: Option<&Value>) -> Option<i64> {
match v? {
Value::Number(n) => n.as_i64().or_else(|| n.as_f64().map(|f| f as i64)),
Value::String(s) => s.trim().parse::<f64>().ok().map(|f| f as i64),
_ => None,
}
}
/// A JSON scalar as a string (ids arrive as numbers or strings).
pub fn scalar_string(v: Option<&Value>) -> Option<String> {
match v? {
Value::String(s) if !s.is_empty() => Some(s.clone()),
Value::Number(n) => Some(n.to_string()),
_ => None,
}
}
/// Python truthiness of a JSON value (`bool(x)`).
pub fn truthy(v: Option<&Value>) -> bool {
match v {
None | Some(Value::Null) => false,
Some(Value::Bool(b)) => *b,
Some(Value::Number(n)) => n.as_f64() != Some(0.0),
Some(Value::String(s)) => !s.is_empty(),
Some(Value::Array(a)) => !a.is_empty(),
Some(Value::Object(o)) => !o.is_empty(),
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn lenient_numbers() {
assert_eq!(lenient_f64(Some(&json!("335.00"))), Some(335.0));
assert_eq!(lenient_f64(Some(&json!(" 1,234.5 "))), Some(1234.5));
assert_eq!(lenient_f64(Some(&json!(2.5))), Some(2.5));
for junk in [json!(""), json!("-"), json!("***"), json!(null), json!([1])] {
assert_eq!(lenient_f64(Some(&junk)), None, "{junk}");
}
assert_eq!(lenient_i64(Some(&json!("2"))), Some(2));
assert_eq!(lenient_i64(Some(&json!(3))), Some(3));
}
#[test]
fn python_truthiness() {
assert!(truthy(Some(&json!(true))));
assert!(truthy(Some(&json!(1))));
assert!(!truthy(Some(&json!(0))));
assert!(!truthy(Some(&json!(""))));
assert!(!truthy(None));
}
}
+150
View File
@@ -0,0 +1,150 @@
"""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()
+77
View File
@@ -0,0 +1,77 @@
#!/usr/bin/env bash
# Install the Rust collector (rust/collector, `ping-collector`) and switch
# data collection over to it, or back.
#
# sudo bash scripts/install_collector.sh # install binary + unit, dry-run check
# sudo bash scripts/install_collector.sh --cutover # install, then Rust collects:
# # COLLECTION_ENABLED=false, start
# # ping-collector, restart water-monitor
# sudo bash scripts/install_collector.sh --rollback # Python collects again
#
# The binary comes from BIN=/path/to/ping-collector (e.g. the "ping-collector-
# linux-x86_64" CI artifact) or from rust/target/release/ping-collector, built
# beforehand as a normal user:
# cargo build --release --locked --manifest-path rust/Cargo.toml
# (needs rustup, pkg-config and libssl-dev).
set -euo pipefail
APP_DIR="${APP_DIR:-/opt/thailand-water-monitor}"
ENV_FILE="${APP_DIR}/.env"
STATUS_FILE="${APP_DIR}/collector-status.json"
REPO_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
MODE="${1:-install}"
log() { printf '\033[1;32m==>\033[0m %s\n' "$*"; }
die() { printf '\033[1;31m[error]\033[0m %s\n' "$*" >&2; exit 1; }
[ "$(id -u)" -eq 0 ] || die "run as root (sudo)"
[ -f "$ENV_FILE" ] || die "$ENV_FILE not found"
set_env() {
local key="$1" value="$2"
if grep -q "^${key}=" "$ENV_FILE"; then
sed -i "s|^${key}=.*|${key}=${value}|" "$ENV_FILE"
else
printf '%s=%s\n' "$key" "$value" >>"$ENV_FILE"
fi
}
if [ "$MODE" = "--rollback" ]; then
log "Stopping ping-collector; the Python app collects again"
systemctl disable --now ping-collector.service 2>/dev/null || true
set_env COLLECTION_ENABLED true
systemctl restart water-monitor.service
log "Done. Check: journalctl -u water-monitor -n 50"
exit 0
fi
case "$MODE" in install | --cutover) ;; *) die "unknown option: $MODE" ;; esac
BIN="${BIN:-${REPO_DIR}/rust/target/release/ping-collector}"
[ -x "$BIN" ] || die "no binary at $BIN: build it (see header) or pass BIN=/path/to/ping-collector"
log "Installing $BIN -> ${APP_DIR}/bin/ping-collector ($("$BIN" --version))"
install -D -m 0755 -o root -g root "$BIN" "${APP_DIR}/bin/ping-collector"
install -m 0644 "${REPO_DIR}/scripts/ping-collector.service" /etc/systemd/system/ping-collector.service
systemctl daemon-reload
# Both services read and write /scraping/status through this file (inside
# the app dir: the units use PrivateTmp).
set_env SCRAPING_STATUS_FILE "$STATUS_FILE"
touch "$STATUS_FILE"
chown "$(stat -c %U:%G "$APP_DIR")" "$STATUS_FILE"
log "Dry run against the configured database (writes nothing)"
(cd "$APP_DIR" && "${APP_DIR}/bin/ping-collector" --env-file "$ENV_FILE" once --dry-run)
if [ "$MODE" = "--cutover" ]; then
log "Cutover: ping-collector collects, the Python app forecasts and notifies"
set_env COLLECTION_ENABLED false
systemctl enable --now ping-collector.service
systemctl restart water-monitor.service
sleep 20
systemctl --no-pager --lines=0 status ping-collector.service || true
journalctl -u ping-collector -n 20 --no-pager || true
log "Rollback if needed: sudo bash scripts/install_collector.sh --rollback"
else
log "Installed, not started. Switch over with: sudo bash scripts/install_collector.sh --cutover"
fi
+35
View File
@@ -0,0 +1,35 @@
[Unit]
Description=Ping River Monitor data collector (Rust: RID, HII, dams, Open-Meteo)
Documentation=https://git.b4l.co.th/B4L/Northern-Thailand-Ping-River-Monitor/src/branch/master/rust
After=network-online.target
Wants=network-online.target
[Service]
Type=simple
User=water-monitor
Group=water-monitor
WorkingDirectory=/opt/thailand-water-monitor
ExecStart=/opt/thailand-water-monitor/bin/ping-collector run
Restart=always
RestartSec=30
TimeoutStopSec=20
# Same .env as water-monitor.service (POSTGRES_*, RID_*, GAP_FILL_*,
# SCRAPING_STATUS_FILE). The Python app must run with COLLECTION_ENABLED=false.
EnvironmentFile=/opt/thailand-water-monitor/.env
# Security settings
NoNewPrivileges=true
PrivateTmp=true
ProtectSystem=strict
ProtectHome=true
ReadWritePaths=/opt/thailand-water-monitor
CapabilityBoundingSet=
MemoryMax=256M
StandardOutput=journal
StandardError=journal
SyslogIdentifier=ping-collector
[Install]
WantedBy=multi-user.target
+14 -1
View File
@@ -142,8 +142,20 @@ class Config:
# On startup the collection leader re-fetches RID reports with holes from # On startup the collection leader re-fetches RID reports with holes from
# the last N days (0 = off). Older history: `run.py --fill-gaps all`. # the last N days (0 = off). Older history: `run.py --fill-gaps all`.
GAP_FILL_DAYS_ON_START = int(os.getenv("GAP_FILL_DAYS_ON_START", "14")) GAP_FILL_DAYS_ON_START = int(os.getenv("GAP_FILL_DAYS_ON_START", "14"))
# Leader -> other workers hand-off for /scraping/status (default: temp dir) # Leader -> other workers hand-off for /scraping/status (default: temp dir).
# With ping-collector, point both services at the same file inside the app
# dir (the systemd units use PrivateTmp).
SCRAPING_STATUS_FILE = os.getenv("SCRAPING_STATUS_FILE", "") SCRAPING_STATUS_FILE = os.getenv("SCRAPING_STATUS_FILE", "")
# false: data collection runs in the Rust ping-collector service
# (rust/collector); this app then only serves, forecasts and notifies,
# re-forecasting when the newest RID reading advances (polled every
# FORECAST_POLL_SECONDS) and at least once per clock hour.
COLLECTION_ENABLED = os.getenv("COLLECTION_ENABLED", "true").lower() in (
"1",
"true",
"yes",
)
FORECAST_POLL_SECONDS = int(os.getenv("FORECAST_POLL_SECONDS", "60"))
# Logging settings # Logging settings
LOG_LEVEL = os.getenv("LOG_LEVEL", "INFO") LOG_LEVEL = os.getenv("LOG_LEVEL", "INFO")
@@ -293,6 +305,7 @@ class Config:
"RID_POLL_OFFSET_MINUTES": cls.RID_POLL_OFFSET_MINUTES, "RID_POLL_OFFSET_MINUTES": cls.RID_POLL_OFFSET_MINUTES,
"RID_RETRY_MINUTES": cls.RID_RETRY_MINUTES, "RID_RETRY_MINUTES": cls.RID_RETRY_MINUTES,
"GAP_FILL_DAYS_ON_START": cls.GAP_FILL_DAYS_ON_START, "GAP_FILL_DAYS_ON_START": cls.GAP_FILL_DAYS_ON_START,
"COLLECTION_ENABLED": cls.COLLECTION_ENABLED,
"LOG_LEVEL": cls.LOG_LEVEL, "LOG_LEVEL": cls.LOG_LEVEL,
"LOG_FILE": cls.LOG_FILE, "LOG_FILE": cls.LOG_FILE,
"DATA_RETENTION_DAYS": cls.DATA_RETENTION_DAYS, "DATA_RETENTION_DAYS": cls.DATA_RETENTION_DAYS,
+68 -3
View File
@@ -218,8 +218,15 @@ async def lifespan(app: FastAPI):
app_state["notify"] = None app_state["notify"] = None
if app_state["leader_lock"]: if app_state["leader_lock"]:
app_state["notify"] = _init_notifications() app_state["notify"] = _init_notifications()
app_state["scraping_task"] = asyncio.create_task(background_scraping_task()) if Config.COLLECTION_ENABLED:
logger.info("This worker is the background-collection leader") app_state["scraping_task"] = asyncio.create_task(background_scraping_task())
logger.info("This worker is the background-collection leader")
else:
app_state["scraping_task"] = asyncio.create_task(forecast_notify_task())
logger.info(
"COLLECTION_ENABLED=false: data collection runs in ping-collector "
"(rust/collector); this worker only refreshes forecasts and notifies"
)
else: else:
logger.info( logger.info(
"Another worker holds collection leadership; " "Another worker holds collection leadership; "
@@ -587,6 +594,58 @@ async def background_scraping_task():
await asyncio.sleep(60) # Wait a minute before retrying await asyncio.sleep(60) # Wait a minute before retrying
def _latest_rid_reading():
"""Newest RID reading in the database (collected by another process)."""
from sqlalchemy import text
engine = app_state["scraper"].db_adapter.engine
with engine.connect() as conn:
return conn.execute(
text("SELECT MAX(timestamp) FROM water_measurements")
).scalar()
async def forecast_notify_task():
"""Leader loop when COLLECTION_ENABLED=false (ping-collector collects).
Polls the newest RID reading every FORECAST_POLL_SECONDS; when it
advances, and once per clock hour regardless (so a stopped feed is still
noticed and reported), refreshes the forecast cache/history and pushes
ntfy transitions.
"""
last_hourly_slot = None
newest_reading = None
while True:
try:
try:
latest = await asyncio.to_thread(_latest_rid_reading)
except Exception as e:
logger.warning(f"Could not read the newest RID reading: {e}")
latest = None
new_readings = latest is not None and (
newest_reading is None or latest > newest_reading
)
if new_readings:
newest_reading = latest
slot = datetime.now().replace(minute=0, second=0, microsecond=0)
hourly_due = slot != last_hourly_slot
if hourly_due:
last_hourly_slot = slot
if new_readings or hourly_due:
await _precompute_forecasts()
await _notify_transitions()
await asyncio.sleep(Config.FORECAST_POLL_SECONDS)
except asyncio.CancelledError:
logger.info("Forecast/notify task cancelled")
break
except Exception as e:
logger.error(f"Error in forecast/notify task: {e}")
await asyncio.sleep(60)
# Umami server-side API tracking (fire-and-forget; never blocks a response) # Umami server-side API tracking (fire-and-forget; never blocks a response)
_UMAMI_TRACK_PREFIXES = ( _UMAMI_TRACK_PREFIXES = (
@@ -1658,6 +1717,11 @@ async def trigger_scraping(background_tasks: BackgroundTasks):
"""Trigger manual data scraping""" """Trigger manual data scraping"""
increment_counter("api_requests", labels={"endpoint": "scrape_trigger"}) increment_counter("api_requests", labels={"endpoint": "scrape_trigger"})
if not Config.COLLECTION_ENABLED:
raise HTTPException(
status_code=409,
detail="Collection runs in ping-collector (COLLECTION_ENABLED=false)",
)
if app_state["is_scraping"]: if app_state["is_scraping"]:
raise HTTPException(status_code=409, detail="Scraping already in progress") raise HTTPException(status_code=409, detail="Scraping already in progress")
@@ -1696,7 +1760,8 @@ async def get_scraping_status():
increment_counter("api_requests", labels={"endpoint": "scraping_status"}) increment_counter("api_requests", labels={"endpoint": "scraping_status"})
# Only the leader worker collects; the others report what it published. # Only the leader worker collects; the others report what it published.
if app_state.get("leader_lock"): # With COLLECTION_ENABLED=false ping-collector writes the same file.
if Config.COLLECTION_ENABLED and app_state.get("leader_lock"):
stats = dict(app_state["scraping_stats"], is_running=app_state["is_scraping"]) stats = dict(app_state["scraping_stats"], is_running=app_state["is_scraping"])
else: else:
stats = _read_scraping_status() or {"is_running": False} stats = _read_scraping_status() or {"is_running": False}
+96
View File
@@ -0,0 +1,96 @@
"""COLLECTION_ENABLED=false: the Rust ping-collector collects, this app only
re-forecasts and notifies when the database's newest RID reading advances
(and once per clock hour), and /scraping/status reports the collector's file."""
import asyncio
import datetime
import json
from src import web_api
from src.config import Config
def test_forecasts_follow_new_readings_and_the_clock_hour(monkeypatch):
clock = {"now": datetime.datetime(2026, 10, 17, 21, 30)}
class FixedClock(datetime.datetime):
@classmethod
def now(cls, tz=None):
return clock["now"]
# newest reading per poll: unchanged, unchanged, advances, next hour
latest = iter(
[
datetime.datetime(2026, 10, 17, 20),
datetime.datetime(2026, 10, 17, 20),
datetime.datetime(2026, 10, 17, 21),
datetime.datetime(2026, 10, 17, 21),
]
)
calls = {"precompute": 0, "notify": 0}
sleeps = []
async def precompute():
calls["precompute"] += 1
async def notify():
calls["notify"] += 1
async def fake_sleep(seconds):
sleeps.append(seconds)
if len(sleeps) == 3: # the last poll happens in the next clock hour
clock["now"] = datetime.datetime(2026, 10, 17, 22, 1)
if len(sleeps) == 4:
raise asyncio.CancelledError
monkeypatch.setattr(web_api, "datetime", FixedClock)
monkeypatch.setattr(web_api.asyncio, "sleep", fake_sleep)
monkeypatch.setattr(web_api, "_latest_rid_reading", lambda: next(latest))
monkeypatch.setattr(web_api, "_precompute_forecasts", precompute)
monkeypatch.setattr(web_api, "_notify_transitions", notify)
monkeypatch.setattr(Config, "FORECAST_POLL_SECONDS", 60)
asyncio.run(web_api.forecast_notify_task())
# first poll (new + hourly), third (new reading), fourth (new clock hour)
assert calls == {"precompute": 3, "notify": 3}
assert sleeps == [60, 60, 60, 60]
def test_status_comes_from_the_collector_file(tmp_path, monkeypatch):
status = tmp_path / "collector-status.json"
status.write_text(
json.dumps(
{
"is_running": False,
"total_runs": 4,
"successful_runs": 3,
"failed_runs": 1,
"last_run": "2026-10-17T22:05:01",
"next_run": "2026-10-17T23:05:00",
"feed_current": True,
"latest_reading": "2026-10-17T22:00:00",
"readings_saved": 352,
"collector": "ping-collector 0.1.0",
}
)
)
monkeypatch.setattr(Config, "COLLECTION_ENABLED", False)
monkeypatch.setattr(Config, "SCRAPING_STATUS_FILE", str(status))
monkeypatch.setitem(web_api.app_state, "leader_lock", object())
resp = asyncio.run(web_api.get_scraping_status())
assert resp.total_runs == 4
assert resp.feed_current is True
assert resp.latest_reading == datetime.datetime(2026, 10, 17, 22)
def test_manual_trigger_is_refused_when_collection_runs_elsewhere(monkeypatch):
import pytest
from fastapi import BackgroundTasks, HTTPException
monkeypatch.setattr(Config, "COLLECTION_ENABLED", False)
with pytest.raises(HTTPException) as err:
asyncio.run(web_api.trigger_scraping(BackgroundTasks()))
assert err.value.status_code == 409