501 lines
18 KiB
Python
501 lines
18 KiB
Python
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
import sqlite3
|
||
|
|
import threading
|
||
|
|
import time
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Any, Dict, List, Optional
|
||
|
|
|
||
|
|
from loguru import logger
|
||
|
|
|
||
|
|
from src.database.db_manager import DBManager
|
||
|
|
|
||
|
|
STATE_STORAGE_FILE = "file"
|
||
|
|
STATE_STORAGE_DUAL = "dual"
|
||
|
|
STATE_STORAGE_SQLITE = "sqlite"
|
||
|
|
VALID_STATE_STORAGE_MODES = {
|
||
|
|
STATE_STORAGE_FILE,
|
||
|
|
STATE_STORAGE_DUAL,
|
||
|
|
STATE_STORAGE_SQLITE,
|
||
|
|
}
|
||
|
|
|
||
|
|
_LOGGED_MODES: set[str] = set()
|
||
|
|
|
||
|
|
|
||
|
|
def get_state_storage_mode() -> str:
|
||
|
|
raw = str(os.getenv("POLYWEATHER_STATE_STORAGE_MODE") or STATE_STORAGE_DUAL).strip().lower()
|
||
|
|
if raw not in VALID_STATE_STORAGE_MODES:
|
||
|
|
logger.warning(
|
||
|
|
f"invalid POLYWEATHER_STATE_STORAGE_MODE={raw!r}, fallback to {STATE_STORAGE_DUAL}"
|
||
|
|
)
|
||
|
|
raw = STATE_STORAGE_DUAL
|
||
|
|
if raw not in _LOGGED_MODES:
|
||
|
|
logger.info(f"runtime state storage mode={raw}")
|
||
|
|
_LOGGED_MODES.add(raw)
|
||
|
|
return raw
|
||
|
|
|
||
|
|
|
||
|
|
class RuntimeStateDB:
|
||
|
|
_instance: Optional["RuntimeStateDB"] = None
|
||
|
|
_instance_lock = threading.Lock()
|
||
|
|
|
||
|
|
def __init__(self, db_path: Optional[str] = None):
|
||
|
|
self.db_path = DBManager(db_path).db_path
|
||
|
|
self._init_tables()
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def instance(cls) -> "RuntimeStateDB":
|
||
|
|
with cls._instance_lock:
|
||
|
|
if cls._instance is None:
|
||
|
|
cls._instance = cls()
|
||
|
|
return cls._instance
|
||
|
|
|
||
|
|
def connect(self) -> sqlite3.Connection:
|
||
|
|
conn = sqlite3.connect(self.db_path)
|
||
|
|
conn.row_factory = sqlite3.Row
|
||
|
|
return conn
|
||
|
|
|
||
|
|
def _init_tables(self) -> None:
|
||
|
|
with self.connect() as conn:
|
||
|
|
conn.execute(
|
||
|
|
"""
|
||
|
|
CREATE TABLE IF NOT EXISTS daily_records_store (
|
||
|
|
city TEXT NOT NULL,
|
||
|
|
target_date TEXT NOT NULL,
|
||
|
|
actual_high REAL,
|
||
|
|
deb_prediction REAL,
|
||
|
|
mu REAL,
|
||
|
|
updated_at REAL NOT NULL,
|
||
|
|
payload_json TEXT NOT NULL,
|
||
|
|
PRIMARY KEY (city, target_date)
|
||
|
|
)
|
||
|
|
"""
|
||
|
|
)
|
||
|
|
conn.execute(
|
||
|
|
"""
|
||
|
|
CREATE TABLE IF NOT EXISTS telegram_alert_last_by_city (
|
||
|
|
city TEXT PRIMARY KEY,
|
||
|
|
signature TEXT,
|
||
|
|
trigger_key TEXT,
|
||
|
|
severity TEXT,
|
||
|
|
ts INTEGER,
|
||
|
|
active INTEGER DEFAULT 0,
|
||
|
|
cleared_ts INTEGER,
|
||
|
|
evidence_json TEXT
|
||
|
|
)
|
||
|
|
"""
|
||
|
|
)
|
||
|
|
conn.execute(
|
||
|
|
"""
|
||
|
|
CREATE TABLE IF NOT EXISTS telegram_alert_signature_state (
|
||
|
|
signature TEXT PRIMARY KEY,
|
||
|
|
ts INTEGER NOT NULL
|
||
|
|
)
|
||
|
|
"""
|
||
|
|
)
|
||
|
|
conn.execute(
|
||
|
|
"""
|
||
|
|
CREATE TABLE IF NOT EXISTS probability_training_snapshots_store (
|
||
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
|
|
city TEXT NOT NULL,
|
||
|
|
target_date TEXT NOT NULL,
|
||
|
|
timestamp TEXT NOT NULL,
|
||
|
|
raw_mu REAL,
|
||
|
|
raw_sigma REAL,
|
||
|
|
max_so_far REAL,
|
||
|
|
peak_status TEXT,
|
||
|
|
probability_mode TEXT,
|
||
|
|
legacy_top_bucket INTEGER,
|
||
|
|
shadow_top_bucket INTEGER,
|
||
|
|
payload_json TEXT NOT NULL
|
||
|
|
)
|
||
|
|
"""
|
||
|
|
)
|
||
|
|
conn.execute(
|
||
|
|
"CREATE INDEX IF NOT EXISTS idx_probability_snapshot_city_date ON probability_training_snapshots_store(city, target_date, id DESC)"
|
||
|
|
)
|
||
|
|
conn.execute(
|
||
|
|
"""
|
||
|
|
CREATE TABLE IF NOT EXISTS open_meteo_cache_store (
|
||
|
|
source_kind TEXT NOT NULL,
|
||
|
|
cache_key TEXT NOT NULL,
|
||
|
|
updated_at REAL NOT NULL,
|
||
|
|
expires_at REAL,
|
||
|
|
payload_json TEXT NOT NULL,
|
||
|
|
PRIMARY KEY (source_kind, cache_key)
|
||
|
|
)
|
||
|
|
"""
|
||
|
|
)
|
||
|
|
conn.execute(
|
||
|
|
"CREATE INDEX IF NOT EXISTS idx_open_meteo_cache_expires ON open_meteo_cache_store(source_kind, expires_at)"
|
||
|
|
)
|
||
|
|
conn.commit()
|
||
|
|
|
||
|
|
|
||
|
|
class DailyRecordRepository:
|
||
|
|
def __init__(self, db: Optional[RuntimeStateDB] = None):
|
||
|
|
self.db = db or RuntimeStateDB.instance()
|
||
|
|
|
||
|
|
def load_all(self) -> Dict[str, Dict[str, Dict[str, Any]]]:
|
||
|
|
out: Dict[str, Dict[str, Dict[str, Any]]] = {}
|
||
|
|
with self.db.connect() as conn:
|
||
|
|
rows = conn.execute(
|
||
|
|
"SELECT city, target_date, payload_json FROM daily_records_store ORDER BY city, target_date"
|
||
|
|
).fetchall()
|
||
|
|
for row in rows:
|
||
|
|
try:
|
||
|
|
payload = json.loads(row["payload_json"])
|
||
|
|
except Exception:
|
||
|
|
continue
|
||
|
|
city = str(row["city"])
|
||
|
|
date_str = str(row["target_date"])
|
||
|
|
out.setdefault(city, {})[date_str] = payload
|
||
|
|
return out
|
||
|
|
|
||
|
|
def upsert_record(self, city: str, target_date: str, record: Dict[str, Any]) -> None:
|
||
|
|
payload_json = json.dumps(record, ensure_ascii=False)
|
||
|
|
updated_at = time.time()
|
||
|
|
with self.db.connect() as conn:
|
||
|
|
conn.execute(
|
||
|
|
"""
|
||
|
|
INSERT INTO daily_records_store (
|
||
|
|
city, target_date, actual_high, deb_prediction, mu, updated_at, payload_json
|
||
|
|
) VALUES (?, ?, ?, ?, ?, ?, ?)
|
||
|
|
ON CONFLICT(city, target_date) DO UPDATE SET
|
||
|
|
actual_high = excluded.actual_high,
|
||
|
|
deb_prediction = excluded.deb_prediction,
|
||
|
|
mu = excluded.mu,
|
||
|
|
updated_at = excluded.updated_at,
|
||
|
|
payload_json = excluded.payload_json
|
||
|
|
""",
|
||
|
|
(
|
||
|
|
city,
|
||
|
|
target_date,
|
||
|
|
record.get("actual_high"),
|
||
|
|
record.get("deb_prediction"),
|
||
|
|
record.get("mu"),
|
||
|
|
updated_at,
|
||
|
|
payload_json,
|
||
|
|
),
|
||
|
|
)
|
||
|
|
conn.commit()
|
||
|
|
|
||
|
|
def replace_all(self, data: Dict[str, Dict[str, Dict[str, Any]]]) -> int:
|
||
|
|
count = 0
|
||
|
|
with self.db.connect() as conn:
|
||
|
|
conn.execute("DELETE FROM daily_records_store")
|
||
|
|
for city, city_rows in (data or {}).items():
|
||
|
|
if not isinstance(city_rows, dict):
|
||
|
|
continue
|
||
|
|
for target_date, record in city_rows.items():
|
||
|
|
payload_json = json.dumps(record, ensure_ascii=False)
|
||
|
|
conn.execute(
|
||
|
|
"""
|
||
|
|
INSERT INTO daily_records_store (
|
||
|
|
city, target_date, actual_high, deb_prediction, mu, updated_at, payload_json
|
||
|
|
) VALUES (?, ?, ?, ?, ?, ?, ?)
|
||
|
|
""",
|
||
|
|
(
|
||
|
|
city,
|
||
|
|
target_date,
|
||
|
|
record.get("actual_high"),
|
||
|
|
record.get("deb_prediction"),
|
||
|
|
record.get("mu"),
|
||
|
|
time.time(),
|
||
|
|
payload_json,
|
||
|
|
),
|
||
|
|
)
|
||
|
|
count += 1
|
||
|
|
conn.commit()
|
||
|
|
return count
|
||
|
|
|
||
|
|
def delete_older_than(self, cutoff_date: str) -> int:
|
||
|
|
with self.db.connect() as conn:
|
||
|
|
cur = conn.execute(
|
||
|
|
"DELETE FROM daily_records_store WHERE target_date < ?",
|
||
|
|
(cutoff_date,),
|
||
|
|
)
|
||
|
|
conn.commit()
|
||
|
|
return int(cur.rowcount or 0)
|
||
|
|
|
||
|
|
|
||
|
|
class TelegramAlertStateRepository:
|
||
|
|
def __init__(self, db: Optional[RuntimeStateDB] = None):
|
||
|
|
self.db = db or RuntimeStateDB.instance()
|
||
|
|
|
||
|
|
def load_state(self) -> Dict[str, Any]:
|
||
|
|
state = {"last_by_city": {}, "by_signature": {}}
|
||
|
|
with self.db.connect() as conn:
|
||
|
|
city_rows = conn.execute(
|
||
|
|
"SELECT city, signature, trigger_key, severity, ts, active, cleared_ts, evidence_json FROM telegram_alert_last_by_city"
|
||
|
|
).fetchall()
|
||
|
|
sig_rows = conn.execute(
|
||
|
|
"SELECT signature, ts FROM telegram_alert_signature_state"
|
||
|
|
).fetchall()
|
||
|
|
for row in city_rows:
|
||
|
|
entry = {
|
||
|
|
"signature": row["signature"],
|
||
|
|
"trigger_key": row["trigger_key"],
|
||
|
|
"severity": row["severity"],
|
||
|
|
"ts": row["ts"],
|
||
|
|
"active": bool(row["active"]),
|
||
|
|
}
|
||
|
|
if row["cleared_ts"] is not None:
|
||
|
|
entry["cleared_ts"] = row["cleared_ts"]
|
||
|
|
if row["evidence_json"]:
|
||
|
|
try:
|
||
|
|
entry["evidence"] = json.loads(row["evidence_json"])
|
||
|
|
except Exception:
|
||
|
|
pass
|
||
|
|
state["last_by_city"][str(row["city"])] = entry
|
||
|
|
for row in sig_rows:
|
||
|
|
state["by_signature"][str(row["signature"])] = int(row["ts"] or 0)
|
||
|
|
return state
|
||
|
|
|
||
|
|
def save_state(self, state: Dict[str, Any]) -> None:
|
||
|
|
last_by_city = state.get("last_by_city") or {}
|
||
|
|
by_signature = state.get("by_signature") or {}
|
||
|
|
with self.db.connect() as conn:
|
||
|
|
conn.execute("DELETE FROM telegram_alert_last_by_city")
|
||
|
|
conn.execute("DELETE FROM telegram_alert_signature_state")
|
||
|
|
for city, row in last_by_city.items():
|
||
|
|
if not isinstance(row, dict):
|
||
|
|
continue
|
||
|
|
conn.execute(
|
||
|
|
"""
|
||
|
|
INSERT INTO telegram_alert_last_by_city (
|
||
|
|
city, signature, trigger_key, severity, ts, active, cleared_ts, evidence_json
|
||
|
|
) VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
||
|
|
""",
|
||
|
|
(
|
||
|
|
city,
|
||
|
|
row.get("signature"),
|
||
|
|
row.get("trigger_key"),
|
||
|
|
row.get("severity"),
|
||
|
|
int(row.get("ts") or 0),
|
||
|
|
1 if row.get("active") else 0,
|
||
|
|
row.get("cleared_ts"),
|
||
|
|
json.dumps(row.get("evidence"), ensure_ascii=False)
|
||
|
|
if row.get("evidence") is not None
|
||
|
|
else None,
|
||
|
|
),
|
||
|
|
)
|
||
|
|
for signature, ts in by_signature.items():
|
||
|
|
conn.execute(
|
||
|
|
"INSERT INTO telegram_alert_signature_state (signature, ts) VALUES (?, ?)",
|
||
|
|
(signature, int(ts or 0)),
|
||
|
|
)
|
||
|
|
conn.commit()
|
||
|
|
|
||
|
|
def replace_from_state(self, state: Dict[str, Any]) -> int:
|
||
|
|
self.save_state(state)
|
||
|
|
return len((state.get("last_by_city") or {})) + len((state.get("by_signature") or {}))
|
||
|
|
|
||
|
|
|
||
|
|
class ProbabilitySnapshotRepository:
|
||
|
|
def __init__(self, db: Optional[RuntimeStateDB] = None):
|
||
|
|
self.db = db or RuntimeStateDB.instance()
|
||
|
|
|
||
|
|
def append_snapshot(self, payload: Dict[str, Any]) -> None:
|
||
|
|
legacy_top = _top_bucket(payload.get("prob_snapshot"))
|
||
|
|
shadow_top = _top_bucket(payload.get("shadow_prob_snapshot"))
|
||
|
|
with self.db.connect() as conn:
|
||
|
|
conn.execute(
|
||
|
|
"""
|
||
|
|
INSERT INTO probability_training_snapshots_store (
|
||
|
|
city, target_date, timestamp, raw_mu, raw_sigma, max_so_far,
|
||
|
|
peak_status, probability_mode, legacy_top_bucket, shadow_top_bucket, payload_json
|
||
|
|
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||
|
|
""",
|
||
|
|
(
|
||
|
|
payload.get("city"),
|
||
|
|
payload.get("date"),
|
||
|
|
payload.get("timestamp"),
|
||
|
|
payload.get("raw_mu"),
|
||
|
|
payload.get("raw_sigma"),
|
||
|
|
payload.get("max_so_far"),
|
||
|
|
payload.get("peak_status"),
|
||
|
|
payload.get("probability_mode"),
|
||
|
|
legacy_top,
|
||
|
|
shadow_top,
|
||
|
|
json.dumps(payload, ensure_ascii=False),
|
||
|
|
),
|
||
|
|
)
|
||
|
|
conn.commit()
|
||
|
|
|
||
|
|
def load_recent_rows(self, city: str, target_date: str, limit: int) -> List[Dict[str, Any]]:
|
||
|
|
with self.db.connect() as conn:
|
||
|
|
rows = conn.execute(
|
||
|
|
"""
|
||
|
|
SELECT payload_json
|
||
|
|
FROM probability_training_snapshots_store
|
||
|
|
WHERE city = ? AND target_date = ?
|
||
|
|
ORDER BY id DESC
|
||
|
|
LIMIT ?
|
||
|
|
""",
|
||
|
|
(city, target_date, int(limit)),
|
||
|
|
).fetchall()
|
||
|
|
out = []
|
||
|
|
for row in rows:
|
||
|
|
try:
|
||
|
|
out.append(json.loads(row["payload_json"]))
|
||
|
|
except Exception:
|
||
|
|
continue
|
||
|
|
return out
|
||
|
|
|
||
|
|
def load_all_rows(self) -> List[Dict[str, Any]]:
|
||
|
|
with self.db.connect() as conn:
|
||
|
|
rows = conn.execute(
|
||
|
|
"SELECT payload_json FROM probability_training_snapshots_store ORDER BY id"
|
||
|
|
).fetchall()
|
||
|
|
out = []
|
||
|
|
for row in rows:
|
||
|
|
try:
|
||
|
|
out.append(json.loads(row["payload_json"]))
|
||
|
|
except Exception:
|
||
|
|
continue
|
||
|
|
return out
|
||
|
|
|
||
|
|
def replace_all(self, rows: List[Dict[str, Any]]) -> int:
|
||
|
|
count = 0
|
||
|
|
with self.db.connect() as conn:
|
||
|
|
conn.execute("DELETE FROM probability_training_snapshots_store")
|
||
|
|
for payload in rows or []:
|
||
|
|
if not isinstance(payload, dict):
|
||
|
|
continue
|
||
|
|
legacy_top = _top_bucket(payload.get("prob_snapshot"))
|
||
|
|
shadow_top = _top_bucket(payload.get("shadow_prob_snapshot"))
|
||
|
|
conn.execute(
|
||
|
|
"""
|
||
|
|
INSERT INTO probability_training_snapshots_store (
|
||
|
|
city, target_date, timestamp, raw_mu, raw_sigma, max_so_far,
|
||
|
|
peak_status, probability_mode, legacy_top_bucket, shadow_top_bucket, payload_json
|
||
|
|
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||
|
|
""",
|
||
|
|
(
|
||
|
|
payload.get("city"),
|
||
|
|
payload.get("date"),
|
||
|
|
payload.get("timestamp"),
|
||
|
|
payload.get("raw_mu"),
|
||
|
|
payload.get("raw_sigma"),
|
||
|
|
payload.get("max_so_far"),
|
||
|
|
payload.get("peak_status"),
|
||
|
|
payload.get("probability_mode"),
|
||
|
|
legacy_top,
|
||
|
|
shadow_top,
|
||
|
|
json.dumps(payload, ensure_ascii=False),
|
||
|
|
),
|
||
|
|
)
|
||
|
|
count += 1
|
||
|
|
conn.commit()
|
||
|
|
return count
|
||
|
|
|
||
|
|
|
||
|
|
class OpenMeteoCacheRepository:
|
||
|
|
def __init__(self, db: Optional[RuntimeStateDB] = None):
|
||
|
|
self.db = db or RuntimeStateDB.instance()
|
||
|
|
|
||
|
|
def replace_payload(self, payload: Dict[str, Any], max_age: int) -> int:
|
||
|
|
count = 0
|
||
|
|
now = time.time()
|
||
|
|
with self.db.connect() as conn:
|
||
|
|
conn.execute("DELETE FROM open_meteo_cache_store")
|
||
|
|
for source_kind in ("forecast", "ensemble", "multi_model"):
|
||
|
|
bucket = payload.get(source_kind) or {}
|
||
|
|
if not isinstance(bucket, dict):
|
||
|
|
continue
|
||
|
|
for cache_key, entry in bucket.items():
|
||
|
|
if not isinstance(entry, dict):
|
||
|
|
continue
|
||
|
|
updated_at = float(entry.get("t") or now)
|
||
|
|
expires_at = updated_at + max_age
|
||
|
|
conn.execute(
|
||
|
|
"""
|
||
|
|
INSERT INTO open_meteo_cache_store (
|
||
|
|
source_kind, cache_key, updated_at, expires_at, payload_json
|
||
|
|
) VALUES (?, ?, ?, ?, ?)
|
||
|
|
""",
|
||
|
|
(
|
||
|
|
source_kind,
|
||
|
|
cache_key,
|
||
|
|
updated_at,
|
||
|
|
expires_at,
|
||
|
|
json.dumps(entry, ensure_ascii=False),
|
||
|
|
),
|
||
|
|
)
|
||
|
|
count += 1
|
||
|
|
conn.commit()
|
||
|
|
return count
|
||
|
|
|
||
|
|
def load_payload(self, max_age: int) -> Dict[str, Any]:
|
||
|
|
now = time.time()
|
||
|
|
payload: Dict[str, Any] = {
|
||
|
|
"forecast": {},
|
||
|
|
"ensemble": {},
|
||
|
|
"multi_model": {},
|
||
|
|
"saved_at": now,
|
||
|
|
}
|
||
|
|
with self.db.connect() as conn:
|
||
|
|
rows = conn.execute(
|
||
|
|
"SELECT source_kind, cache_key, updated_at, payload_json FROM open_meteo_cache_store"
|
||
|
|
).fetchall()
|
||
|
|
for row in rows:
|
||
|
|
updated_at = float(row["updated_at"] or 0)
|
||
|
|
if now - updated_at >= max(600, max_age):
|
||
|
|
continue
|
||
|
|
try:
|
||
|
|
entry = json.loads(row["payload_json"])
|
||
|
|
except Exception:
|
||
|
|
continue
|
||
|
|
payload.setdefault(str(row["source_kind"]), {})[str(row["cache_key"])] = entry
|
||
|
|
return payload
|
||
|
|
|
||
|
|
def latest_updated_at(self) -> float:
|
||
|
|
with self.db.connect() as conn:
|
||
|
|
row = conn.execute(
|
||
|
|
"SELECT MAX(updated_at) AS max_updated_at FROM open_meteo_cache_store"
|
||
|
|
).fetchone()
|
||
|
|
if not row:
|
||
|
|
return 0.0
|
||
|
|
try:
|
||
|
|
return float(row["max_updated_at"] or 0.0)
|
||
|
|
except Exception:
|
||
|
|
return 0.0
|
||
|
|
|
||
|
|
|
||
|
|
def _top_bucket(snapshot: Optional[List[Dict[str, Any]]]) -> Optional[int]:
|
||
|
|
best_value = None
|
||
|
|
best_prob = -1.0
|
||
|
|
for row in snapshot or []:
|
||
|
|
if not isinstance(row, dict):
|
||
|
|
continue
|
||
|
|
value = row.get("v")
|
||
|
|
if value is None:
|
||
|
|
value = row.get("value")
|
||
|
|
try:
|
||
|
|
ivalue = int(value)
|
||
|
|
except Exception:
|
||
|
|
continue
|
||
|
|
prob = row.get("p")
|
||
|
|
if prob is None:
|
||
|
|
prob = row.get("probability")
|
||
|
|
try:
|
||
|
|
fprob = float(prob)
|
||
|
|
except Exception:
|
||
|
|
continue
|
||
|
|
if fprob > best_prob:
|
||
|
|
best_prob = fprob
|
||
|
|
best_value = ivalue
|
||
|
|
return best_value
|
||
|
|
|
||
|
|
|
||
|
|
def get_runtime_data_dir() -> str:
|
||
|
|
raw = str(os.getenv("POLYWEATHER_RUNTIME_DATA_DIR") or "").strip()
|
||
|
|
if raw:
|
||
|
|
return raw
|
||
|
|
project_root = Path(__file__).resolve().parents[2]
|
||
|
|
return str(project_root / "data")
|