From 215c46cea75eaa0546c875bdc854f31823ec3ae2 Mon Sep 17 00:00:00 2001 From: emi Date: Thu, 23 Jul 2026 04:41:26 +0000 Subject: [PATCH] Bound Postgres connections, add rate-limit/timeout guards, move revgeo off the threadpool (#8) --- accounts/db.py | 32 +++- data/climate.py | 292 ++++++++++++++++++++++-------- data/climate_store.py | 215 +++++++++++++--------- data/places.py | 9 +- data/store.py | 194 +++++++++++++------- requirements.txt | 6 +- tests/accounts/test_db_engines.py | 25 +++ tests/data/test_climate.py | 197 ++++++++++++++++++++ tests/data/test_climate_store.py | 11 +- tests/data/test_store.py | 88 +++++++++ tests/test_warm_cities.py | 38 ++++ tests/web/test_gzip.py | 12 ++ warm_cities.py | 21 +++ web/app.py | 11 +- 14 files changed, 912 insertions(+), 239 deletions(-) create mode 100644 tests/test_warm_cities.py create mode 100644 tests/web/test_gzip.py diff --git a/accounts/db.py b/accounts/db.py index 680c15f..8b97ff8 100644 --- a/accounts/db.py +++ b/accounts/db.py @@ -56,18 +56,40 @@ def _apply_sqlite_pragmas(dbapi_conn, _rec): cur.close() +# Postgres-only pool sizing + sync-path timeouts, defined at module scope (not +# only inside `if IS_POSTGRES` below) so the values themselves are unit-testable +# without a real Postgres connection. +# +# pool_size=1/max_overflow=1 (the original setting on both) meant a 3rd +# concurrent request per worker had to wait out pool_timeout (~30s default) for a +# slot — easy to hit with 4 uvicorn workers x real traffic. 5+5 gives each async +# engine headroom for a burst of concurrent account requests without coming close +# to Postgres's own connection ceiling (two async engines here plus the sync +# engine below, times N workers, all share that budget). +_ASYNC_POOL_KWARGS = dict(pool_size=5, max_overflow=5, pool_pre_ping=True, + pool_recycle=1800, future=True) +# Smaller than the async engines: only the notifier leader (one process +# cluster-wide, see core/singleton.py) uses the sync engine, driving its periodic +# sweep from a plain thread with no event loop. connect_timeout + a +# statement_timeout (via psycopg's `options`) bound how long a wedged Postgres +# can hang the notifier — it runs for the lifetime of that process's leader +# flock, so an indefinite hang here would silently stop every subscription +# notification until the process is restarted. +_SYNC_POOL_KWARGS = dict(pool_size=2, max_overflow=3, pool_pre_ping=True, + pool_recycle=1800, future=True, + connect_args={"connect_timeout": 5, + "options": "-c statement_timeout=30000"}) + if IS_POSTGRES: DB_PATH = None - _COMMON = dict(pool_size=1, max_overflow=1, pool_pre_ping=True, pool_recycle=1800, future=True) # Read-write engine (default) + a read-only engine that pins every session to a # read-only transaction on the server, so a pure-read endpoint can't accidentally # write and a long read stays off the write connection. - async_engine = create_async_engine(DATABASE_URL, **_COMMON) + async_engine = create_async_engine(DATABASE_URL, **_ASYNC_POOL_KWARGS) async_engine_ro = create_async_engine( - DATABASE_URL, **_COMMON, + DATABASE_URL, **_ASYNC_POOL_KWARGS, connect_args={"server_settings": {"default_transaction_read_only": "on"}}) - sync_engine = create_engine(_sync_url(DATABASE_URL), pool_size=1, max_overflow=1, - pool_pre_ping=True, pool_recycle=1800, future=True) + sync_engine = create_engine(_sync_url(DATABASE_URL), **_SYNC_POOL_KWARGS) else: DB_PATH = os.environ.get("THERMOGRAPH_ACCOUNTS_DB") or os.path.join(paths.DATA_DIR, "accounts.sqlite") os.makedirs(os.path.dirname(DB_PATH), exist_ok=True) diff --git a/data/climate.py b/data/climate.py index f88fc43..8f22ce1 100644 --- a/data/climate.py +++ b/data/climate.py @@ -5,8 +5,11 @@ a zstd-compressed parquet file keyed by cell id. Percentiles are computed on the fly from this raw record (see grading.py), which keeps the cache small and lets us handle ties (e.g. many zero-precip days) correctly. """ +import concurrent.futures import datetime import os +import queue +import tempfile import threading import time @@ -20,8 +23,30 @@ import paths from data import climate_store from data import store +# Reused across every fetch (archive/forecast/geocode/revgeo) instead of a bare +# httpx.get per call opening a fresh TCP+TLS connection every time -- mirrors +# web/app.py's reused _frontend_client. Every call still passes its own explicit +# per-call timeout (see _request), so this only changes connection reuse, never +# retry/backoff semantics. Never explicitly closed: this is a long-lived module +# used by both the server process and one-shot scripts (warm_cities.py, +# migrate_cache_to_pg.py), and the OS reclaims the sockets at process exit either +# way. +_client = httpx.Client() + CACHE_DIR = os.path.join(paths.DATA_DIR, "cache") MAX_ATTEMPTS = 3 # per upstream call, before giving up + +# _fetch_history / _fetch_history_nasa run under _load_history's per-cell lock +# (see _cell_lock), so every attempt's timeout adds directly to how long a +# same-cell waiter queues behind this thread. A flat 180s x MAX_ATTEMPTS could +# pin the thread for up to 9 minutes on one slow-but-eventually-ok upstream. The +# full archive response is genuinely large (~16k daily rows from START_DATE, see +# below), so it can legitimately need more than a topup fetch's 60s -- but that +# headroom belongs on the LAST attempt only: the first attempts fail fast so a +# truly wedged upstream frees the lock for a waiter sooner, and only the final, +# worth-waiting-out attempt gets the longer allowance. Worst case under the lock +# drops from 540s (9 min) to 60+60+150=270s (4.5 min). +ARCHIVE_FETCH_TIMEOUTS = (60, 60, 150) START_DATE = "1980-01-01" # ERA5 reaches back to 1940; 1980 = 45 yrs, fast + robust ARCHIVE_LATENCY_DAYS = 6 # ERA5 archive lags real time by a few days # A real archive spans decades (~16k daily rows from START_DATE). A self-hosted @@ -111,6 +136,17 @@ _archive_cooldown_until = 0.0 _archive_limit_daily = False # is the current cooldown a daily-quota exhaustion? +# The forecast (recent+forward) endpoint's own cooldown, tracked separately from +# the archive one above: Open-Meteo's archive and forecast APIs have independent +# quotas, so a 429 on one must not gate (or be masked by) a cooldown meant for the +# other. Without this, a forecast brownout let every subscribed cell AND every +# live /cell request independently retry-storm the endpoint -- each paying +# MAX_ATTEMPTS x up to 60s plus the MET Norway fallback's own retry cycle, with no +# circuit breaker (see _load_recent_forecast). +FORECAST_COOLDOWN = 120 # seconds +_forecast_cooldown_until = 0.0 +_forecast_limit_daily = False # is the current forecast cooldown a daily-quota exhaustion? + class WeatherUnavailable(RuntimeError): """Weather data can't be fetched right now (rate limit, upstream outage with @@ -157,7 +193,8 @@ def _seconds_to_utc_reset() -> float: def _note_rate_limit(e) -> None: - """Record a rate limit: back off ~2 min, or until tomorrow for a daily quota.""" + """Record an archive rate limit: back off ~2 min, or until tomorrow for a + daily quota.""" global _archive_cooldown_until, _archive_limit_daily reason = _rate_limit_reason(e).lower() _archive_limit_daily = "daily" in reason or "tomorrow" in reason @@ -165,6 +202,16 @@ def _note_rate_limit(e) -> None: _seconds_to_utc_reset() if _archive_limit_daily else ARCHIVE_COOLDOWN) +def _note_forecast_rate_limit(e) -> None: + """Record a forecast-endpoint rate limit. Mirrors _note_rate_limit exactly, + but against the separate _forecast_cooldown_until (see its module comment).""" + global _forecast_cooldown_until, _forecast_limit_daily + reason = _rate_limit_reason(e).lower() + _forecast_limit_daily = "daily" in reason or "tomorrow" in reason + _forecast_cooldown_until = time.time() + ( + _seconds_to_utc_reset() if _forecast_limit_daily else FORECAST_COOLDOWN) + + def _cell_lock(cell_id: str) -> threading.Lock: with _LOCKS_GUARD: lk = _CELL_LOCKS.get(cell_id) @@ -176,11 +223,18 @@ def _cell_lock(cell_id: str) -> threading.Lock: def _request(url, params, timeout, *, phase, headers=None, attempts=MAX_ATTEMPTS): """GET with bounded retries; every retry and the final failure are logged to the errors folder (tagged ``retry`` / ``error``). A 429 (rate limit) fails fast - without retrying, so we don't hammer the limit and make it worse.""" + without retrying, so we don't hammer the limit and make it worse. + + ``timeout`` is either one number applied to every attempt (the common case), + or a sequence of per-attempt values (see ARCHIVE_FETCH_TIMEOUTS) so a caller + that holds a lock across every attempt can fail fast on the early ones and + reserve a longer wait for the last.""" last = None for attempt in range(1, attempts + 1): + attempt_timeout = (timeout[min(attempt - 1, len(timeout) - 1)] + if isinstance(timeout, (tuple, list)) else timeout) try: - r = httpx.get(url, params=params, timeout=timeout, headers=headers) + r = _client.get(url, params=params, timeout=attempt_timeout, headers=headers) r.raise_for_status() metrics.record_outbound(phase, "ok") return r @@ -283,9 +337,28 @@ def _with_doy(df: pl.DataFrame) -> pl.DataFrame: def _write_cache(df: pl.DataFrame, path: str) -> None: """Persist a daily record to the parquet cache — the raw record only, never - the derived doy column (it's recomputed at read time).""" - os.makedirs(CACHE_DIR, exist_ok=True) - df.drop("doy", strict=False).write_parquet(path, compression="zstd") + the derived doy column (it's recomputed at read time). + + Written to a tempfile in the same directory, then renamed into place + (os.replace is atomic on the same filesystem) — mirrors the pattern in + api/homepage.py's refresh() and data/places.py's _fetch. A reader (or a + second writer — warm_cities.py guards against overlapping runs, but a live + request racing a topup is normal) must never see a partially-written parquet + file; writing straight to ``path`` (the old behavior) could hand a truncated + file to a concurrent read.""" + directory = os.path.dirname(path) + os.makedirs(directory, exist_ok=True) + fd, tmp = tempfile.mkstemp(dir=directory, suffix=".tmp") + try: + os.close(fd) + df.drop("doy", strict=False).write_parquet(tmp, compression="zstd") + os.replace(tmp, path) + except BaseException: + try: + os.unlink(tmp) + except OSError: + pass + raise # --- cache backend dispatch -------------------------------------------------- @@ -462,7 +535,7 @@ def _to_frame(daily: dict) -> pl.DataFrame: def _fetch_history(cell: dict) -> pl.DataFrame: end = (datetime.date.today() - datetime.timedelta(days=ARCHIVE_LATENCY_DAYS)).isoformat() params = _om_daily_params(cell, start_date=START_DATE, end_date=end, models=ARCHIVE_MODEL) - r = _request(ARCHIVE_URL, params, 180, phase="history_fetch") + r = _request(ARCHIVE_URL, params, ARCHIVE_FETCH_TIMEOUTS, phase="history_fetch") return _to_frame(r.json()["daily"]) @@ -523,7 +596,7 @@ def _fetch_history_nasa(cell: dict) -> pl.DataFrame: "end": end, "format": "JSON", } - r = _request(NASA_POWER_URL, params, 180, phase="history_nasa") + r = _request(NASA_POWER_URL, params, ARCHIVE_FETCH_TIMEOUTS, phase="history_nasa") return _nasa_to_frame(r.json()["properties"]["parameter"]) @@ -798,14 +871,27 @@ def _load_recent_forecast(cell: dict) -> pl.DataFrame: "past_days": RECENT_PAST_DAYS, "forecast_days": FORECAST_DAYS, } - try: - r = _request(FORECAST_URL, params, 60, phase="recent_forecast_fetch") - df = _to_frame(r.json()["daily"]) - except Exception as e: # noqa: BLE001 - # Open-Meteo forecast unavailable (rate limit or outage). Fall back to MET - # Norway (yr.no) — global + keyless — mirroring the archive's NASA POWER - # backup. MET Norway is forecast-only (no recent past days), so this keeps - # the forecast / day-ahead views working in a degraded form. + # Skip Open-Meteo's forecast endpoint entirely while it's in its own + # rate-limit cooldown (see _note_forecast_rate_limit) -- mirrors + # _load_history's archive-cooldown check, so a forecast brownout doesn't + # retry-storm the endpoint from every subscribed cell and every live request + # independently; it falls straight through to the MET Norway backup instead. + df = None + primary_error = None + if time.time() >= _forecast_cooldown_until: + try: + r = _request(FORECAST_URL, params, 60, phase="recent_forecast_fetch") + df = _to_frame(r.json()["daily"]) + except Exception as e: # noqa: BLE001 + if is_rate_limit(e): + _note_forecast_rate_limit(e) + primary_error = e + + if df is None: + # Open-Meteo forecast unavailable (rate limit, cooldown, or outage). Fall + # back to MET Norway (yr.no) — global + keyless — mirroring the archive's + # NASA POWER backup. MET Norway is forecast-only (no recent past days), + # so this keeps the forecast / day-ahead views working in a degraded form. try: df = _fetch_forecast_metno(cell) except Exception: # noqa: BLE001 - backup unavailable too @@ -818,11 +904,16 @@ def _load_recent_forecast(cell: dict) -> pl.DataFrame: return _with_doy(hit[0]) # No cache either: surface the original error, classifying an Open-Meteo # rate limit as the typed, daily-aware WeatherUnavailable (the archive - # path classifies its own inside _load_history). - if is_rate_limit(e): - daily = "daily" in _rate_limit_reason(e).lower() - raise WeatherUnavailable(limit_message(daily), daily=daily) from e - raise + # path classifies its own inside _load_history). primary_error is None + # when the primary fetch was skipped outright (cooldown active) -- + # report the cooldown itself in that case. + if primary_error is not None and is_rate_limit(primary_error): + daily = "daily" in _rate_limit_reason(primary_error).lower() + raise WeatherUnavailable(limit_message(daily), daily=daily) from primary_error + if primary_error is not None: + raise primary_error + raise WeatherUnavailable(limit_message(_forecast_limit_daily), + daily=_forecast_limit_daily) _write_recent_backed(cell_id, df) return df @@ -830,13 +921,21 @@ def _load_recent_forecast(cell: dict) -> pl.DataFrame: _REVGEO_CACHE: dict[str, str | None] = {} # Nominatim's usage policy allows ~1 request/second and rejects bursts. The compare -# page loads several locations at once, so serialize the reverse lookups here and -# space them out — otherwise the burst is rate-limited, a null label gets cached, and -# those locations are left showing bare coordinates. -_REVGEO_LOCK = threading.Lock() +# page loads several locations at once, so the actual fetch + pacing run on ONE +# dedicated worker thread (see _revgeo_worker), never on a caller's own thread — +# in the server, a caller's thread is one of Starlette's shared sync threadpool +# threads, and the old design serialized AND slept (up to _REVGEO_MIN_INTERVAL) +# right there, pinning a threadpool thread for the whole wait. A single worker +# draining a queue gets the same ~1/sec pacing for free (nothing else ever calls +# Nominatim) without blocking anything but itself. _REVGEO_MIN_INTERVAL = 1.1 # seconds between successive Nominatim reverse calls _revgeo_last = 0.0 +_REVGEO_QUEUE: "queue.Queue[tuple[float, float, str, concurrent.futures.Future]]" = queue.Queue() +_REVGEO_WORKER_LOCK = threading.Lock() +_revgeo_worker_started = False +_REVGEO_WAIT_TIMEOUT = 10.0 # seconds a caller waits for ITS OWN request before giving up + def reverse_geocode_cached(lat: float, lon: float) -> tuple[bool, str | None]: """(found, label) from the in-memory + SQLite revgeo caches only — never calls @@ -850,64 +949,113 @@ def reverse_geocode_cached(lat: float, lon: float) -> tuple[bool, str | None]: return (found, label) +def _fetch_revgeo_label(lat: float, lon: float) -> str | None: + """The actual Nominatim reverse-geocode HTTP call + label assembly. Called + ONLY from _revgeo_worker (never on a caller's thread) — returns None (never + raises) so a bad response degrades to bare coordinates instead of killing the + worker thread.""" + try: + r = _request( + "https://nominatim.openstreetmap.org/reverse", + # zoom 14 resolves to the suburb/neighbourhood level so we can lead + # with it when OSM has one (zoom 10 only ever returns the city). + # accept-language=en keeps labels in one script worldwide (matches + # the forward geocoder's language=en). + {"lat": lat, "lon": lon, "format": "jsonv2", "zoom": 14, + "addressdetails": 1, "accept-language": "en"}, + 15, + phase="reverse_geocode", + headers={"User-Agent": "Thermograph/0.1 (local weather grading app)"}, + ) + a = r.json().get("address", {}) or {} + # Lead with the neighbourhood when available, then the city, then region. + neighborhood = (a.get("neighbourhood") or a.get("suburb") + or a.get("quarter") or a.get("city_district") + or a.get("borough")) + city = (a.get("city") or a.get("town") or a.get("village") + or a.get("hamlet") or a.get("county")) + region = a.get("state") or a.get("province") or a.get("region") + country = a.get("country") + # dict.fromkeys drops duplicates (e.g. neighbourhood == city) in order. + # Country trails the neighbourhood/city/region so the label reads as a full + # hierarchy ("West Seattle, Seattle, Washington, United States"); the + # frontend leads with the first part and mutes the rest. + parts = dict.fromkeys(p for p in (neighborhood, city, region, country) if p) + return ", ".join(parts) or None + except Exception: # noqa: BLE001 - reverse geocoding is a nicety, never fatal + return None + + +def _revgeo_worker() -> None: + """Drains _REVGEO_QUEUE one request at a time, forever. The sole caller of + _fetch_revgeo_label, so the ~1/sec pacing below is enforced just by doing the + work serially — no lock needed, since nothing else ever touches Nominatim. + Re-checks the cache before fetching (a request queued behind an identical one + is answered from what the earlier request just cached, no duplicate call), + and always finishes the fetch + persists the result even if the original + caller already gave up waiting (see reverse_geocode's timeout).""" + global _revgeo_last + while True: + lat, lon, key, fut = _REVGEO_QUEUE.get() + try: + found, label = reverse_geocode_cached(lat, lon) + if not found: + wait = _REVGEO_MIN_INTERVAL - (time.monotonic() - _revgeo_last) + if wait > 0: + time.sleep(wait) + label = _fetch_revgeo_label(lat, lon) + _revgeo_last = time.monotonic() + _REVGEO_CACHE[key] = label + store.put_revgeo(key, label) # survive restarts (a None label retries after its TTL) + if not fut.done(): + fut.set_result(label) + except Exception: # noqa: BLE001 - never let a bad request kill the worker + if not fut.done(): + fut.set_result(None) + finally: + _REVGEO_QUEUE.task_done() + + +def _start_revgeo_worker() -> None: + """Start the dedicated reverse-geocode worker thread, once per process + (idempotent, thread-safe). Lazy (on first use) rather than at import, since + climate.py has no app-lifespan hook of its own — mirrors data/places.py's + start_loading().""" + global _revgeo_worker_started + with _REVGEO_WORKER_LOCK: + if _revgeo_worker_started: + return + _revgeo_worker_started = True + threading.Thread(target=_revgeo_worker, name="revgeo-worker", daemon=True).start() + + def reverse_geocode(lat: float, lon: float) -> str | None: """Best-effort neighbourhood/city label for a point (OpenStreetMap Nominatim). Cached per ~cell — in memory for the process and in SQLite across restarts — so panning around (or redeploying) doesn't re-hit the service, and failures return None so the caller can fall back to bare coordinates. + + The actual fetch (and Nominatim's pacing) happens on the dedicated + _revgeo_worker thread, not here — this just enqueues the request and waits up + to _REVGEO_WAIT_TIMEOUT for its result, so a burst of concurrent uncached + lookups (the compare page fires several at once) queues behind the ~1/sec + limit without pinning a caller's (in the server, a shared threadpool) thread + for the wait. A timed-out wait returns None like any other failed lookup — + the caller falls back to bare coordinates — but the worker keeps going and + still caches the answer for the next call. """ key = store.revgeo_key(lat, lon) found, label = reverse_geocode_cached(lat, lon) if found: return label - # Serialize + rate-limit the upstream call. Concurrent callers (the compare page - # loads several locations at once) would otherwise burst past Nominatim's ~1/sec - # limit and get rate-limited, caching a null label. Under the lock we re-check the - # cache (a peer may have just resolved this cell) and space successive calls out. - global _revgeo_last - with _REVGEO_LOCK: - found, label = reverse_geocode_cached(lat, lon) - if found: - return label - wait = _REVGEO_MIN_INTERVAL - (time.monotonic() - _revgeo_last) - if wait > 0: - time.sleep(wait) - label = None - try: - r = _request( - "https://nominatim.openstreetmap.org/reverse", - # zoom 14 resolves to the suburb/neighbourhood level so we can lead - # with it when OSM has one (zoom 10 only ever returns the city). - # accept-language=en keeps labels in one script worldwide (matches - # the forward geocoder's language=en). - {"lat": lat, "lon": lon, "format": "jsonv2", "zoom": 14, - "addressdetails": 1, "accept-language": "en"}, - 15, - phase="reverse_geocode", - headers={"User-Agent": "Thermograph/0.1 (local weather grading app)"}, - ) - a = r.json().get("address", {}) or {} - # Lead with the neighbourhood when available, then the city, then region. - neighborhood = (a.get("neighbourhood") or a.get("suburb") - or a.get("quarter") or a.get("city_district") - or a.get("borough")) - city = (a.get("city") or a.get("town") or a.get("village") - or a.get("hamlet") or a.get("county")) - region = a.get("state") or a.get("province") or a.get("region") - country = a.get("country") - # dict.fromkeys drops duplicates (e.g. neighbourhood == city) in order. - # Country trails the neighbourhood/city/region so the label reads as a full - # hierarchy ("West Seattle, Seattle, Washington, United States"); the - # frontend leads with the first part and mutes the rest. - parts = dict.fromkeys(p for p in (neighborhood, city, region, country) if p) - label = ", ".join(parts) or None - except Exception: # noqa: BLE001 - reverse geocoding is a nicety, never fatal - label = None - _revgeo_last = time.monotonic() - _REVGEO_CACHE[key] = label - store.put_revgeo(key, label) # survive restarts (a None label retries after its TTL) - return label + _start_revgeo_worker() + fut: "concurrent.futures.Future[str | None]" = concurrent.futures.Future() + _REVGEO_QUEUE.put((lat, lon, key, fut)) + try: + return fut.result(timeout=_REVGEO_WAIT_TIMEOUT) + except concurrent.futures.TimeoutError: + return None def geocode(name: str, count: int = 5) -> list[dict]: diff --git a/data/climate_store.py b/data/climate_store.py index 69d77fb..19c883c 100644 --- a/data/climate_store.py +++ b/data/climate_store.py @@ -24,6 +24,7 @@ recompute-from-parquet), a miss here degrades to *fetch-from-upstream* — prod no parquet fallback, and ``climate.py``'s loaders already treat "no cache" as "go fetch". """ +import contextlib import os import threading @@ -68,29 +69,69 @@ def _libpq_dsn(url: str) -> str: _PG_DSN = _libpq_dsn(DATABASE_URL) if IS_POSTGRES else "" -# Per-thread connection: uvicorn serves on a thread pool and psycopg connections -# aren't thread-safe, so each thread lazily opens its own (mirrors store.py). -_local = threading.local() +# Bounded pool, replacing the old thread-local connection (one psycopg.connect per +# THREAD, opened lazily, never closed). Starlette's sync threadpool can spin up +# dozens of worker threads per uvicorn worker; with N workers x this module x +# data/store.py all doing the same thing, that had no ceiling and could exhaust +# Postgres max_connections, taking down every module sharing the database at +# once. min_size=1 keeps an idle deploy cheap; max_size=8 is a modest per-process +# ceiling that still lets several concurrent cell fetches proceed without +# serializing on one connection. connect_timeout + statement_timeout bound how +# long a caller (or, worse, the notifier's leader-election flock) can be wedged +# by an unresponsive Postgres. +_POOL_MIN_SIZE = 1 +_POOL_MAX_SIZE = 8 + +_pool_lock = threading.Lock() +_pool_obj = None # type: "psycopg_pool.ConnectionPool | None" -def _conn(): - """Thread-local autocommit psycopg connection, or None when Postgres is off or - unreachable. A dropped connection is retired so the next call reconnects.""" +def _pool(): + """Lazily open (once) the bounded connection pool, or None when Postgres is + off or the pool can't be opened — callers treat that exactly like the old + _conn() returning None.""" + global _pool_obj if not IS_POSTGRES: return None - conn = getattr(_local, "conn", None) - if conn is not None: - if getattr(conn, "closed", False): - _local.conn = None - else: - return conn - try: - import psycopg # local import: only the Postgres path needs it - conn = psycopg.connect(_PG_DSN, autocommit=True) - _local.conn = conn - return conn - except Exception: # noqa: BLE001 - the store is optional; callers fetch on miss - return None + if _pool_obj is not None: + return _pool_obj + with _pool_lock: + if _pool_obj is None: + try: + import psycopg_pool # local import: only the Postgres path needs it + _pool_obj = psycopg_pool.ConnectionPool( + _PG_DSN, + min_size=_POOL_MIN_SIZE, + max_size=_POOL_MAX_SIZE, + kwargs={ + "autocommit": True, + "connect_timeout": 5, + "options": "-c statement_timeout=30000", + }, + open=True, + ) + except Exception: # noqa: BLE001 - the store is optional; callers fetch on miss + return None + return _pool_obj + + +@contextlib.contextmanager +def _conn(): + """Borrow a connection from the bounded pool for the duration of one + operation, or yield None when Postgres is off or unreachable. Every caller + frames exactly one operation in a ``with _conn() as conn:`` block, so the + connection is returned to the pool (or discarded, if broken) the moment that + operation finishes — replacing the old pattern where a thread opened one raw + connection and held it open, unclosed, for its entire lifetime.""" + if not IS_POSTGRES: + yield None + return + pool = _pool() + if pool is None: + yield None + return + with pool.connection() as conn: + yield conn # --- freshness (climate_sync) ----------------------------------------------- @@ -124,12 +165,12 @@ def _bump_sync(conn, cell_id, *, history=None, recent=None, cols_version=None) - def history_synced_at(cell_id: str) -> float: """Epoch seconds of the cell's last history write/topup; 0.0 if absent.""" - conn = _conn() - if conn is None: - return 0.0 try: - row = _get_sync(conn, cell_id) - return float(row[0]) if row and row[0] is not None else 0.0 + with _conn() as conn: + if conn is None: + return 0.0 + row = _get_sync(conn, cell_id) + return float(row[0]) if row and row[0] is not None else 0.0 except Exception: # noqa: BLE001 return 0.0 @@ -137,24 +178,24 @@ def history_synced_at(cell_id: str) -> float: def recent_synced_at(cell_id: str) -> float: """Epoch seconds of the cell's last recent-bundle write; 0.0 if absent. Backs climate.recent_stamp, so it must advance only on a real rewrite (see write_recent).""" - conn = _conn() - if conn is None: - return 0.0 try: - row = _get_sync(conn, cell_id) - return float(row[1]) if row and row[1] is not None else 0.0 + with _conn() as conn: + if conn is None: + return 0.0 + row = _get_sync(conn, cell_id) + return float(row[1]) if row and row[1] is not None else 0.0 except Exception: # noqa: BLE001 return 0.0 def cols_version(cell_id: str) -> int: """The stored schema version for a cell's history; 0 if absent.""" - conn = _conn() - if conn is None: - return 0 try: - row = _get_sync(conn, cell_id) - return int(row[2]) if row and row[2] is not None else 0 + with _conn() as conn: + if conn is None: + return 0 + row = _get_sync(conn, cell_id) + return int(row[2]) if row and row[2] is not None else 0 except Exception: # noqa: BLE001 return 0 @@ -175,19 +216,19 @@ def read_history(cell_id: str) -> "tuple[pl.DataFrame, float] | None": """(history frame, history_synced_at) for a cell, or None when there is no schema-complete cached record. A stored cols_version below COLS_VERSION reads as absent, so climate.py refetches to add the new column(s).""" - conn = _conn() - if conn is None: - return None try: - sync = _get_sync(conn, cell_id) - if sync is None or sync[0] is None: - return None - if (sync[2] or 0) < COLS_VERSION: - return None - df = _read_frame(conn, "climate_history", cell_id) - if df.is_empty(): - return None - return df, float(sync[0]) + with _conn() as conn: + if conn is None: + return None + sync = _get_sync(conn, cell_id) + if sync is None or sync[0] is None: + return None + if (sync[2] or 0) < COLS_VERSION: + return None + df = _read_frame(conn, "climate_history", cell_id) + if df.is_empty(): + return None + return df, float(sync[0]) except Exception: # noqa: BLE001 return None @@ -197,29 +238,29 @@ def read_history_raw(cell_id: str) -> "pl.DataFrame | None": fallback when every upstream fetch fails (the Postgres analogue of reading the parquet file directly, bypassing the NEW_COLS completeness check). None when there are no rows.""" - conn = _conn() - if conn is None: - return None try: - df = _read_frame(conn, "climate_history", cell_id) - return None if df.is_empty() else df + with _conn() as conn: + if conn is None: + return None + df = _read_frame(conn, "climate_history", cell_id) + return None if df.is_empty() else df except Exception: # noqa: BLE001 return None def read_recent(cell_id: str) -> "tuple[pl.DataFrame, float] | None": """(recent+forecast frame, recent_synced_at) for a cell, or None when absent.""" - conn = _conn() - if conn is None: - return None try: - sync = _get_sync(conn, cell_id) - if sync is None or sync[1] is None: - return None - df = _read_frame(conn, "climate_recent", cell_id) - if df.is_empty(): - return None - return df, float(sync[1]) + with _conn() as conn: + if conn is None: + return None + sync = _get_sync(conn, cell_id) + if sync is None or sync[1] is None: + return None + df = _read_frame(conn, "climate_recent", cell_id) + if df.is_empty(): + return None + return df, float(sync[1]) except Exception: # noqa: BLE001 return None @@ -238,21 +279,21 @@ def write_history(cell_id: str, df: pl.DataFrame, synced_at: float) -> None: """Upsert a cell's history rows and stamp history_synced_at. Idempotent: a re-fetched day supersedes its stored duplicate (ON CONFLICT DO UPDATE), which mirrors the parquet unique(subset="date", keep="last") merge. Best-effort.""" - conn = _conn() - if conn is None: - return try: - with conn.transaction(): - conn.execute( - "CREATE TEMP TABLE _hist_stage (LIKE climate_history) ON COMMIT DROP") - with conn.cursor() as cur: - _copy_rows(cur, "_hist_stage", cell_id, df) - set_cols = ", ".join(f"{c} = EXCLUDED.{c}" for c in COLS if c != "date") - conn.execute( - f"INSERT INTO climate_history (cell_id, {', '.join(COLS)}) " - f"SELECT cell_id, {', '.join(COLS)} FROM _hist_stage " - f"ON CONFLICT (cell_id, date) DO UPDATE SET {set_cols}") - _bump_sync(conn, cell_id, history=synced_at, cols_version=COLS_VERSION) + with _conn() as conn: + if conn is None: + return + with conn.transaction(): + conn.execute( + "CREATE TEMP TABLE _hist_stage (LIKE climate_history) ON COMMIT DROP") + with conn.cursor() as cur: + _copy_rows(cur, "_hist_stage", cell_id, df) + set_cols = ", ".join(f"{c} = EXCLUDED.{c}" for c in COLS if c != "date") + conn.execute( + f"INSERT INTO climate_history (cell_id, {', '.join(COLS)}) " + f"SELECT cell_id, {', '.join(COLS)} FROM _hist_stage " + f"ON CONFLICT (cell_id, date) DO UPDATE SET {set_cols}") + _bump_sync(conn, cell_id, history=synced_at, cols_version=COLS_VERSION) except Exception: # noqa: BLE001 - failing to cache must not fail the request pass @@ -261,11 +302,11 @@ def touch_history(cell_id: str, synced_at: float) -> None: """Bump only history_synced_at — the analogue of ``os.utime`` on the parquet file when the archive tail is already current, so the hourly topup timer resets without a rewrite. Best-effort.""" - conn = _conn() - if conn is None: - return try: - _bump_sync(conn, cell_id, history=synced_at) + with _conn() as conn: + if conn is None: + return + _bump_sync(conn, cell_id, history=synced_at) except Exception: # noqa: BLE001 pass @@ -276,14 +317,14 @@ def write_recent(cell_id: str, df: pl.DataFrame, synced_at: float) -> None: refresh, so a full replace is the simplest correct write. Call ONLY on a real rewrite — recent_synced_at feeds recent_stamp, which must not advance when a stale bundle is served without a refetch. Best-effort.""" - conn = _conn() - if conn is None: - return try: - with conn.transaction(): - conn.execute("DELETE FROM climate_recent WHERE cell_id = %s", (cell_id,)) - with conn.cursor() as cur: - _copy_rows(cur, "climate_recent", cell_id, df) - _bump_sync(conn, cell_id, recent=synced_at) + with _conn() as conn: + if conn is None: + return + with conn.transaction(): + conn.execute("DELETE FROM climate_recent WHERE cell_id = %s", (cell_id,)) + with conn.cursor() as cur: + _copy_rows(cur, "climate_recent", cell_id, df) + _bump_sync(conn, cell_id, recent=synced_at) except Exception: # noqa: BLE001 pass diff --git a/data/places.py b/data/places.py index c2faadb..c80a9f1 100644 --- a/data/places.py +++ b/data/places.py @@ -34,6 +34,13 @@ import paths GEO_DIR = os.path.join(paths.DATA_DIR, "geonames") GEONAMES_URL = "https://download.geonames.org/export/dump/" +# Reused connection instead of a bare httpx.get per call -- mirrors +# web/app.py's reused _frontend_client / data/climate.py's _client. Only a +# handful of calls ever happen (the dump files are cached forever, see _fetch), +# but a shared client costs nothing and keeps the pattern consistent everywhere +# this module talks HTTP. +_client = httpx.Client() + # Which GeoNames cities dump to index. cities1000 (~170k places, population # ≥ 1000) is small enough to hold in memory and big enough to cover the tiny # vacation towns people actually search for; set THERMOGRAPH_CITIES=cities5000 @@ -74,7 +81,7 @@ def _fetch(name: str) -> str: path = os.path.join(GEO_DIR, name) if os.path.exists(path) and os.path.getsize(path) > 0: return path - r = httpx.get(GEONAMES_URL + name, timeout=120, follow_redirects=True) + r = _client.get(GEONAMES_URL + name, timeout=120, follow_redirects=True) r.raise_for_status() tmp = path + ".part" with open(tmp, "wb") as f: diff --git a/data/store.py b/data/store.py index 4ebc288..d025994 100644 --- a/data/store.py +++ b/data/store.py @@ -21,9 +21,11 @@ their own errors), and deleting data/thermograph.sqlite is a safe full reset. On Postgres (``THERMOGRAPH_DATABASE_URL=postgresql://…``) the same two tables live as UNLOGGED tables — fast and non-durable, which suits a rebuildable cache — and -the per-thread connection uses psycopg instead of sqlite3. Every fail-soft path is -identical: a connection or query that fails just degrades to recompute-from-parquet. +psycopg (via a bounded pool, see _get_pg_pool) replaces sqlite3. Every fail-soft +path is identical: a connection or query that fails just degrades to +recompute-from-parquet. """ +import contextlib import json import os import sqlite3 @@ -118,61 +120,122 @@ else: _SQL_PUT_REVGEO = "INSERT OR REPLACE INTO revgeo VALUES (?,?,?)" # SQLite connections aren't shareable across threads; uvicorn serves requests on -# a thread pool, so each thread lazily opens its own connection. WAL lets those -# readers run concurrently with the (serialized) writers. The Postgres path keeps -# the same per-thread model (psycopg connections likewise aren't thread-safe). +# a thread pool, so each thread lazily opens (and keeps, for its lifetime) its own +# connection. WAL lets those readers run concurrently with the (serialized) +# writers. There's no server-side connection limit to protect against here, so +# this thread-local model stays as-is for SQLite. _local = threading.local() +# Postgres instead uses a bounded pool (see _get_pg_pool), not a thread-local +# connection: the old pattern opened one raw psycopg connection per THREAD, +# lazily, and never closed it. Starlette's sync threadpool can spin up dozens of +# worker threads per uvicorn worker; with N workers x this module x +# data/climate_store.py doing the same thing, that had no ceiling and could +# exhaust Postgres max_connections. min_size=1 keeps an idle deploy cheap; +# max_size=8 is a modest per-process ceiling that still lets several concurrent +# requests proceed without serializing on one connection. connect_timeout + +# statement_timeout bound how long a caller (or the notifier's leader-election +# flock) can be wedged by an unresponsive Postgres. +_pg_pool_lock = threading.Lock() +_pg_pool = None # type: "psycopg_pool.ConnectionPool | None" -def _open_pg(): - """Open a thread-local psycopg connection and ensure the schema. autocommit so - reads never leave an idle-in-transaction and each write lands immediately (the - shared put/get code still calls ``.commit()``, a harmless no-op under autocommit).""" + +def _configure_pg(conn) -> None: + """Run once per NEW physical connection the pool opens (psycopg_pool's + ``configure=``), not on every checkout — mirrors the old _open_pg's + connect-then-ensure-schema sequence without re-running the (idempotent) + DDL on every borrow. + + Multiple connections can legitimately open around the same moment (the + pool's min_size warm-up racing an immediate caller's on-demand open; two + uvicorn workers' pools both cold-starting at once) and each runs this same + CREATE TABLE IF NOT EXISTS. Despite IF NOT EXISTS, two sessions creating the + same brand-new table at the same instant can still hit a duplicate-key error + on Postgres's catalog — a documented race, not a real failure: it just means + the table exists either way, so it's swallowed here rather than failing this + connection's open (which put_payload/get_payload etc. would otherwise + silently treat as "store unavailable" and drop the operation).""" import psycopg # local import: only the Postgres path needs it - conn = psycopg.connect(_PG_DSN, autocommit=True) for stmt in _PG_SCHEMA: - conn.execute(stmt) - return conn + try: + conn.execute(stmt) + except psycopg.errors.UniqueViolation: + pass +def _get_pg_pool(): + """Lazily open (once) the bounded Postgres pool, or None if it can't be + opened — callers treat that exactly like the old _conn() returning None.""" + global _pg_pool + if _pg_pool is not None: + return _pg_pool + with _pg_pool_lock: + if _pg_pool is None: + try: + import psycopg_pool # local import: only the Postgres path needs it + _pg_pool = psycopg_pool.ConnectionPool( + _PG_DSN, + min_size=1, + max_size=8, + kwargs={ + "autocommit": True, + "connect_timeout": 5, + "options": "-c statement_timeout=30000", + }, + configure=_configure_pg, + open=True, + ) + except Exception: # noqa: BLE001 - the store is optional; readers fall back + return None + return _pg_pool + + +@contextlib.contextmanager def _conn(): + """Yield a working connection, or None when the store is unavailable. + + Postgres: borrowed from the bounded pool above for the duration of one + operation and returned automatically on exit. SQLite: the thread's + persistent connection (opened once, kept for the thread's lifetime, same as + before). Either way every caller frames exactly one operation in a + ``with _conn() as conn:`` block.""" + if IS_POSTGRES: + pool = _get_pg_pool() + if pool is None: + yield None + return + with pool.connection() as conn: + yield conn + return conn = getattr(_local, "conn", None) - if conn is not None: - # A dropped Postgres connection would otherwise disable the cache for the - # life of the thread; retire it so the next call reconnects. - if IS_POSTGRES and getattr(conn, "closed", False): - _local.conn = None - else: - return conn - try: - if IS_POSTGRES: - conn = _open_pg() - else: + if conn is None: + try: os.makedirs(os.path.dirname(DB_PATH), exist_ok=True) conn = sqlite3.connect(DB_PATH, timeout=5.0) conn.execute("PRAGMA journal_mode=WAL") conn.execute("PRAGMA synchronous=NORMAL") conn.executescript(_SCHEMA) - _local.conn = conn - return conn - except Exception: # noqa: BLE001 - the store is optional; readers fall back - return None + _local.conn = conn + except Exception: # noqa: BLE001 - the store is optional; readers fall back + yield None + return + yield conn # --- derived payloads ------------------------------------------------------- def get_payload(kind: str, cell_id: str, key: str, token: str) -> bytes | None: """Raw JSON bytes of a cached payload, or None when absent or token-invalid.""" - conn = _conn() - if conn is None: - return None try: - row = conn.execute(_SQL_GET_PAYLOAD, (kind, cell_id, key)).fetchone() - if row is None or row[0] != token: - return None - # psycopg returns bytea as bytes/memoryview; bytes() normalizes both. - data = bytes(row[1]) if IS_POSTGRES else row[1] - return zlib.decompress(data) + with _conn() as conn: + if conn is None: + return None + row = conn.execute(_SQL_GET_PAYLOAD, (kind, cell_id, key)).fetchone() + if row is None or row[0] != token: + return None + # psycopg returns bytea as bytes/memoryview; bytes() normalizes both. + data = bytes(row[1]) if IS_POSTGRES else row[1] + return zlib.decompress(data) except Exception: # noqa: BLE001 return None @@ -196,14 +259,15 @@ def put_payload(kind: str, cell_id: str, key: str, token: str, payload: dict, grading should fail loudly here, not be cached as JS-unparseable `NaN`.""" body = json.dumps(payload, default=str, allow_nan=False, separators=(",", ":")).encode() - conn = _conn() - if cache and conn is not None: + if cache: try: - conn.execute( - _SQL_PUT_PAYLOAD, - (kind, cell_id, key, token, zlib.compress(body, 6), time.time()), - ) - conn.commit() + with _conn() as conn: + if conn is not None: + conn.execute( + _SQL_PUT_PAYLOAD, + (kind, cell_id, key, token, zlib.compress(body, 6), time.time()), + ) + conn.commit() except Exception: # noqa: BLE001 - failing to cache must not fail the request pass return body @@ -223,43 +287,43 @@ def revgeo_key(lat: float, lon: float) -> str: def get_revgeo(key: str) -> tuple[bool, str | None]: """(found, label). A stored NULL label counts as found (a cached miss) until it ages past REVGEO_MISS_TTL, when the lookup becomes retryable.""" - conn = _conn() - if conn is None: - return (False, None) try: - row = conn.execute(_SQL_GET_REVGEO, (key,)).fetchone() - if row is None: - return (False, None) - if row[0] is None and time.time() - row[1] > REVGEO_MISS_TTL: - return (False, None) - return (True, row[0]) + with _conn() as conn: + if conn is None: + return (False, None) + row = conn.execute(_SQL_GET_REVGEO, (key,)).fetchone() + if row is None: + return (False, None) + if row[0] is None and time.time() - row[1] > REVGEO_MISS_TTL: + return (False, None) + return (True, row[0]) except Exception: # noqa: BLE001 return (False, None) def put_revgeo(key: str, label: str | None) -> None: - conn = _conn() - if conn is None: - return try: - conn.execute(_SQL_PUT_REVGEO, (key, label, time.time())) - conn.commit() + with _conn() as conn: + if conn is None: + return + conn.execute(_SQL_PUT_REVGEO, (key, label, time.time())) + conn.commit() except Exception: # noqa: BLE001 pass def stats() -> dict: """Row counts + on-disk size, for the migrate script's summary output.""" - conn = _conn() out = {"db_path": os.path.abspath(DB_PATH), "derived": 0, "revgeo": 0, "bytes": 0} - if conn is None: - return out try: - out["derived"] = conn.execute("SELECT COUNT(*) FROM derived").fetchone()[0] - out["revgeo"] = conn.execute("SELECT COUNT(*) FROM revgeo").fetchone()[0] - # Recent writes may still live in the WAL sidecar — count both files. - out["bytes"] = sum(os.path.getsize(p) for p in (DB_PATH, DB_PATH + "-wal") - if os.path.exists(p)) + with _conn() as conn: + if conn is None: + return out + out["derived"] = conn.execute("SELECT COUNT(*) FROM derived").fetchone()[0] + out["revgeo"] = conn.execute("SELECT COUNT(*) FROM revgeo").fetchone()[0] + # Recent writes may still live in the WAL sidecar — count both files. + out["bytes"] = sum(os.path.getsize(p) for p in (DB_PATH, DB_PATH + "-wal") + if os.path.exists(p)) except Exception: # noqa: BLE001 pass return out diff --git a/requirements.txt b/requirements.txt index b8bb685..c12d96d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -6,11 +6,13 @@ numpy==2.2.1 # Accounts + notification subscriptions (see accounts/db.py, users.py, notify.py). # Postgres in prod/containers: asyncpg (async web/store) + psycopg (sync notifier); # aiosqlite is kept for the SQLite fallback the test suite runs on. alembic manages -# the accounts schema. +# the accounts schema. The `pool` extra pulls in psycopg_pool, used by +# data/climate_store.py + data/store.py to bound the raw connections those modules +# open on Postgres (see their module docstrings) instead of one-per-thread forever. fastapi-users[sqlalchemy]==15.0.5 aiosqlite==0.22.1 asyncpg==0.30.0 -psycopg[binary]==3.2.3 +psycopg[binary,pool]==3.2.3 alembic==1.14.0 # Web Push (VAPID) delivery of notifications (see notifications/push.py). Pulls in # py-vapid, cryptography, and http-ece. diff --git a/tests/accounts/test_db_engines.py b/tests/accounts/test_db_engines.py index 8bc5224..29454d5 100644 --- a/tests/accounts/test_db_engines.py +++ b/tests/accounts/test_db_engines.py @@ -28,6 +28,31 @@ def test_sync_url_derivation(): assert db._sync_url("sqlite+aiosqlite:////tmp/x.sqlite") == "sqlite:////tmp/x.sqlite" +def test_postgres_pool_sizing_and_sync_timeouts(): + """The Postgres engine-construction kwargs, defined at module scope so they're + testable without a real Postgres connection (the engines built from them are + only actually constructed when IS_POSTGRES, exercised by the docker-compose + stack -- see the module docstring). pool_size=1/max_overflow=1 (the old + setting) let a 3rd concurrent request per worker wait out pool_timeout for a + slot; both async engines need real headroom, and the sync engine (the + notifier's, driven off-event-loop) must bound connect + statement time so a + wedged Postgres can't hang it forever under its leader flock.""" + assert db._ASYNC_POOL_KWARGS["pool_size"] >= 5 + assert db._ASYNC_POOL_KWARGS["max_overflow"] >= 5 + + assert db._SYNC_POOL_KWARGS["pool_size"] < db._ASYNC_POOL_KWARGS["pool_size"] + connect_args = db._SYNC_POOL_KWARGS["connect_args"] + assert connect_args["connect_timeout"] > 0 + assert "statement_timeout" in connect_args["options"] + + # Build a real (unconnected -- SQLAlchemy engines are lazy) engine from these + # exact kwargs to confirm they're accepted by the psycopg dialect. + import sqlalchemy + engine = sqlalchemy.create_engine("postgresql+psycopg://u:p@h:5432/d", **db._SYNC_POOL_KWARGS) + assert engine.pool.size() == db._SYNC_POOL_KWARGS["pool_size"] + engine.dispose() + + def test_read_session_yields_a_working_session(): # On SQLite the read session is fully usable (no read-only enforcement); this # just confirms the dependency wiring resolves to a live AsyncSession. diff --git a/tests/data/test_climate.py b/tests/data/test_climate.py index e8e5782..fe4dd4d 100644 --- a/tests/data/test_climate.py +++ b/tests/data/test_climate.py @@ -324,3 +324,200 @@ def test_full_archive_is_accepted(monkeypatch, tmp_path): df, meta = climate._load_history(cell) assert meta["source"] == "open-meteo" assert df.height == full.height + + +# --- forecast cooldown (mirrors the archive path's, tracked separately) ----- + +def test_forecast_cooldown_skips_open_meteo_and_falls_to_metno(monkeypatch, tmp_path): + """While the forecast cooldown is active, the forecast fetch must not call + Open-Meteo at all -- straight to the MET Norway backup, mirroring + _load_history's archive-cooldown skip.""" + import time as time_mod + monkeypatch.setattr(climate, "CACHE_DIR", str(tmp_path)) + monkeypatch.setattr(climate, "_forecast_cooldown_until", time_mod.time() + 60) + cell = {"id": "fc_cooldown_cell", "center_lat": 47.6, "center_lon": -122.3} + + met = {"timeseries": [_metno_step("2026-07-16T00:00:00Z", 18.0, 55.0, 3.0, p1=0.0)]} + + class Resp: + def json(self): return {"properties": met} + + def fake_request(url, params, timeout, *, phase, headers=None, attempts=climate.MAX_ATTEMPTS): + assert url != climate.FORECAST_URL, "must not call Open-Meteo during cooldown" + assert url == climate.METNO_URL + return Resp() + + monkeypatch.setattr(climate, "_request", fake_request) + df = climate._load_recent_forecast(cell) + assert len(df) == 1 + + +def test_forecast_429_sets_its_own_cooldown(monkeypatch, tmp_path): + """A 429 from the forecast endpoint sets _forecast_cooldown_until (NOT + _archive_cooldown_until -- the two upstream endpoints have independent + quotas) and surfaces as WeatherUnavailable when the backup is also down.""" + import time as time_mod + monkeypatch.setattr(climate, "CACHE_DIR", str(tmp_path)) + monkeypatch.setattr(climate, "_forecast_cooldown_until", 0.0) + monkeypatch.setattr(climate, "_archive_cooldown_until", 0.0) + cell = {"id": "fc_429_cell", "center_lat": 47.6, "center_lon": -122.3} + + class FakeResponse: + status_code = 429 + def json(self): return {"reason": "rate limited"} + + class FakeRateLimitError(Exception): + def __init__(self): + self.response = FakeResponse() + + def fake_request(url, params, timeout, *, phase, headers=None, attempts=climate.MAX_ATTEMPTS): + if url == climate.FORECAST_URL: + raise FakeRateLimitError() + raise RuntimeError("metno also down") + + monkeypatch.setattr(climate, "_request", fake_request) + with pytest.raises(climate.WeatherUnavailable): + climate._load_recent_forecast(cell) + assert climate._forecast_cooldown_until > time_mod.time() + assert climate._archive_cooldown_until == 0.0 # the archive cooldown is untouched + + +def test_forecast_cooldown_does_not_block_archive_fetch(monkeypatch, tmp_path): + """The two cooldowns are independent state: an active forecast cooldown must + not gate _load_history's archive fetch.""" + import time as time_mod + monkeypatch.setattr(climate, "CACHE_DIR", str(tmp_path)) + monkeypatch.setattr(climate, "_forecast_cooldown_until", time_mod.time() + 60) + monkeypatch.setattr(climate, "_archive_cooldown_until", 0.0) + cell = {"id": "fc_indep_cell", "center_lat": 47.6, "center_lon": -122.3} + + full = _full_history_frame() + monkeypatch.setattr(climate, "_fetch_history", lambda c: full) + def _nasa_should_not_run(c): raise AssertionError("NASA should not be called") + monkeypatch.setattr(climate, "_fetch_history_nasa", _nasa_should_not_run) + + df, meta = climate._load_history(cell) + assert meta["source"] == "open-meteo" + + +# --- per-attempt timeouts (the cell-lock finding) --------------------------- + +def test_request_applies_a_per_attempt_timeout_sequence(monkeypatch): + """A timeout sequence (e.g. ARCHIVE_FETCH_TIMEOUTS) is applied per attempt, + not the same value on every retry -- the early attempts fail fast and only + the last gets the longer allowance, bounding how long a caller holding a + lock across every attempt (see _load_history) can be pinned.""" + import time as time_mod + seen = [] + + class FakeResp: + def raise_for_status(self): + pass + + class FakeClient: + def get(self, url, params=None, timeout=None, headers=None): + seen.append(timeout) + if len(seen) < 3: + raise RuntimeError("simulated transient failure") + return FakeResp() + + monkeypatch.setattr(climate, "_client", FakeClient()) + monkeypatch.setattr(time_mod, "sleep", lambda s: None) # skip retry backoff + climate._request("http://example.invalid", {}, (60, 60, 150), phase="test") + assert seen == [60, 60, 150] + + +def test_request_scalar_timeout_is_unchanged_for_every_attempt(monkeypatch): + """Every other _request caller passes a single number -- confirm that path is + untouched (same value on every attempt, not just the first).""" + import time as time_mod + seen = [] + + class FakeResp: + def raise_for_status(self): + pass + + class FakeClient: + def get(self, url, params=None, timeout=None, headers=None): + seen.append(timeout) + if len(seen) < 3: + raise RuntimeError("simulated transient failure") + return FakeResp() + + monkeypatch.setattr(climate, "_client", FakeClient()) + monkeypatch.setattr(time_mod, "sleep", lambda s: None) + climate._request("http://example.invalid", {}, 30, phase="test") + assert seen == [30, 30, 30] + + +# --- atomic parquet writes (warm_cities.py overlap-run finding) ------------- + +def test_write_cache_leaves_no_tempfile_behind_on_success(tmp_path): + df = climate._to_frame(_om_daily()) + path = str(tmp_path / "cell.parquet") + climate._write_cache(df, path) + import os + assert sorted(os.listdir(tmp_path)) == ["cell.parquet"] + + +def test_write_cache_cleans_up_tempfile_and_leaves_no_partial_target_on_failure( + monkeypatch, tmp_path): + """A write failure (disk full, killed process, ...) must not leave a + half-written file at the real path, and must not leak the tempfile either -- + a concurrent reader (or an overlapping warm_cities.py run) must only ever see + the old complete file or the new complete file, never a partial one.""" + df = climate._to_frame(_om_daily()) + path = str(tmp_path / "cell.parquet") + + def boom(self, *a, **kw): + raise RuntimeError("disk full") + + monkeypatch.setattr(pl.DataFrame, "write_parquet", boom) + with pytest.raises(RuntimeError): + climate._write_cache(df, path) + + import os + assert not os.path.exists(path) + assert os.listdir(tmp_path) == [] + + +# --- reverse-geocode off the shared threadpool ------------------------------ + +def test_reverse_geocode_cache_hit_never_touches_the_worker_queue(monkeypatch): + key = climate.store.revgeo_key(10.0, 20.0) + monkeypatch.setitem(climate._REVGEO_CACHE, key, "Cached Place") + + def boom(*a, **kw): + raise AssertionError("a cache hit must not enqueue a worker request") + + monkeypatch.setattr(climate._REVGEO_QUEUE, "put", boom) + assert climate.reverse_geocode(10.0, 20.0) == "Cached Place" + + +def test_reverse_geocode_miss_resolves_via_the_worker_thread(monkeypatch): + """An uncached lookup is answered by the dedicated worker thread (not the + calling thread), and the result is cached for the next call.""" + monkeypatch.setattr(climate, "_fetch_revgeo_label", lambda lat, lon: "Worker Place") + label = climate.reverse_geocode(11.111, 22.222) + assert label == "Worker Place" + found, cached = climate.reverse_geocode_cached(11.111, 22.222) + assert found and cached == "Worker Place" + + +def test_reverse_geocode_timeout_returns_none_without_blocking_the_caller(monkeypatch): + """A caller waits only up to _REVGEO_WAIT_TIMEOUT for ITS OWN request, even if + the worker is still busy on it -- it must not block for as long as the fetch + itself takes (that was exactly the old threadpool-pinning problem).""" + import time as time_mod + monkeypatch.setattr(climate, "_REVGEO_WAIT_TIMEOUT", 0.05) + + def slow_fetch(lat, lon): + time_mod.sleep(0.3) + return "Too Slow" + + monkeypatch.setattr(climate, "_fetch_revgeo_label", slow_fetch) + t0 = time_mod.monotonic() + label = climate.reverse_geocode(33.333, 44.444) + elapsed = time_mod.monotonic() - t0 + assert label is None + assert elapsed < 0.2 # returned near the wait timeout, not after the 0.3s fetch diff --git a/tests/data/test_climate_store.py b/tests/data/test_climate_store.py index b31bf35..a6fbcd3 100644 --- a/tests/data/test_climate_store.py +++ b/tests/data/test_climate_store.py @@ -81,11 +81,10 @@ def pg_store(): pytest.skip("set THERMOGRAPH_TEST_DATABASE_URL to run the Postgres bridge test") import psycopg - saved = (climate_store.IS_POSTGRES, climate_store._PG_DSN) + saved = (climate_store.IS_POSTGRES, climate_store._PG_DSN, climate_store._pool_obj) climate_store.IS_POSTGRES = True climate_store._PG_DSN = climate_store._libpq_dsn(url) - if hasattr(climate_store._local, "conn"): - del climate_store._local.conn + climate_store._pool_obj = None # force a fresh pool bound to the test DSN ddl_cols = ", ".join(f"{c} DOUBLE PRECISION" for c in climate_store.COLS if c != "date") with psycopg.connect(climate_store._PG_DSN, autocommit=True) as conn: @@ -100,9 +99,9 @@ def pg_store(): yield climate_store - climate_store.IS_POSTGRES, climate_store._PG_DSN = saved - if hasattr(climate_store._local, "conn"): - del climate_store._local.conn + if climate_store._pool_obj is not None: + climate_store._pool_obj.close() + climate_store.IS_POSTGRES, climate_store._PG_DSN, climate_store._pool_obj = saved def test_pg_history_roundtrip_and_dtypes(pg_store): diff --git a/tests/data/test_store.py b/tests/data/test_store.py index 6cc5f47..fb0b6ac 100644 --- a/tests/data/test_store.py +++ b/tests/data/test_store.py @@ -1,7 +1,12 @@ +"""data/store.py: hermetic SQLite tests (always run, via the tmp_store fixture in +conftest.py) plus a gated Postgres round-trip through the bounded connection pool +(skipped unless THERMOGRAPH_TEST_DATABASE_URL points at a reachable Postgres).""" import json import sqlite3 import time +import pytest + def test_payload_round_trip(tmp_store): payload = {"cell": "1_2", "days": [1, 2, 3]} @@ -63,3 +68,86 @@ def test_store_is_a_pure_accelerator_when_db_unavailable(tmp_store, monkeypatch) assert json.loads(body) == {"v": 1} # still serves the encoded payload assert tmp_store.get_revgeo("x") == (False, None) tmp_store.put_revgeo("x", "label") # no-op, no exception + + +# --- gated Postgres integration (the bounded pool, _get_pg_pool) ----------- + +@pytest.fixture +def pg_store(): + """Point store.py at THERMOGRAPH_TEST_DATABASE_URL for one test -- exercises + the real bounded ConnectionPool (borrow-per-operation, configure=_configure_pg + running the schema DDL) rather than the SQLite fallback. + + store.py picks its dialect-specific SQL text (``?`` vs ``%s`` placeholders, + INSERT OR REPLACE vs ON CONFLICT) once, at import time, from IS_POSTGRES -- + correct in production (THERMOGRAPH_DATABASE_URL is set before the module is + ever imported) but not something flipping IS_POSTGRES back on afterwards + retroactively fixes, so this fixture swaps the SQL constants too, not just + IS_POSTGRES/_PG_DSN. Restores all of it afterwards.""" + import os + + from data import store + + url = os.environ.get("THERMOGRAPH_TEST_DATABASE_URL", "").strip() + if not url: + pytest.skip("set THERMOGRAPH_TEST_DATABASE_URL to run the Postgres pool test") + + saved = (store.IS_POSTGRES, store._PG_DSN, store._pg_pool, + store._SQL_GET_PAYLOAD, store._SQL_PUT_PAYLOAD, + store._SQL_GET_REVGEO, store._SQL_PUT_REVGEO) + store.IS_POSTGRES = True + store._PG_DSN = store._libpq_dsn(url) + store._pg_pool = None # force a fresh pool bound to the test DSN + store._SQL_GET_PAYLOAD = "SELECT token, payload FROM derived WHERE kind=%s AND cell_id=%s AND key=%s" + store._SQL_PUT_PAYLOAD = ( + "INSERT INTO derived (kind, cell_id, key, token, payload, updated_at) " + "VALUES (%s,%s,%s,%s,%s,%s) " + "ON CONFLICT (kind, cell_id, key) DO UPDATE SET " + "token=EXCLUDED.token, payload=EXCLUDED.payload, updated_at=EXCLUDED.updated_at" + ) + store._SQL_GET_REVGEO = "SELECT label, updated_at FROM revgeo WHERE key=%s" + store._SQL_PUT_REVGEO = ( + "INSERT INTO revgeo (key, label, updated_at) VALUES (%s,%s,%s) " + "ON CONFLICT (key) DO UPDATE SET label=EXCLUDED.label, updated_at=EXCLUDED.updated_at" + ) + + import psycopg + with psycopg.connect(store._PG_DSN, autocommit=True) as conn: + conn.execute("DROP TABLE IF EXISTS derived") + conn.execute("DROP TABLE IF EXISTS revgeo") + + yield store + + if store._pg_pool is not None: + store._pg_pool.close() + (store.IS_POSTGRES, store._PG_DSN, store._pg_pool, + store._SQL_GET_PAYLOAD, store._SQL_PUT_PAYLOAD, + store._SQL_GET_REVGEO, store._SQL_PUT_REVGEO) = saved + + +def test_pg_payload_round_trip_through_the_pool(pg_store): + """Two consecutive calls borrow (and return) separate pooled connections -- + unlike the old thread-local model there is no single connection object tying + the write and the read together, so this also exercises that the schema + (applied via the pool's configure= callback) is visible across borrows.""" + body = pg_store.put_payload("grade", "1_2", "k", "tok", {"v": 1}) + assert json.loads(body) == {"v": 1} + assert pg_store.get_payload("grade", "1_2", "k", "tok") == body + assert pg_store.get_payload("grade", "1_2", "k", "tok-mismatch") is None + + +def test_pg_revgeo_round_trip_through_the_pool(pg_store): + key = pg_store.revgeo_key(47.6, -122.3) + assert pg_store.get_revgeo(key) == (False, None) + pg_store.put_revgeo(key, "Somewhere, WA") + assert pg_store.get_revgeo(key) == (True, "Somewhere, WA") + + +def test_pg_pool_survives_many_sequential_borrows(pg_store): + """A crude stand-in for the finding this pool exists to fix: many operations + in a row must keep working off a small bounded pool (max_size=8, well below + this loop) instead of piling up unclosed connections.""" + for i in range(25): + pg_store.put_payload("grade", "cell", f"k{i}", "tok", {"i": i}) + for i in range(25): + assert pg_store.get_json("grade", "cell", f"k{i}", "tok") == {"i": i} diff --git a/tests/test_warm_cities.py b/tests/test_warm_cities.py new file mode 100644 index 0000000..71ad619 --- /dev/null +++ b/tests/test_warm_cities.py @@ -0,0 +1,38 @@ +"""warm_cities.py's cross-process single-instance guard (LOCK_PATH + the +core/singleton.py flock it reuses) -- an overlapping deploy launching a second +copy while the first is still warming must stand down, not double-spend +archive-fetch quota.""" +import fcntl +import os +import subprocess +import sys + +import warm_cities + + +def test_lock_path_lives_under_data_dir(): + assert os.path.dirname(warm_cities.LOCK_PATH) == warm_cities.paths.DATA_DIR + assert os.path.basename(warm_cities.LOCK_PATH) == "warm_cities.lock" + + +def test_second_instance_stands_down_and_exits_zero(tmp_path): + """With the lock already held (simulating an overlapping deploy invocation), a + second run of the script must log and exit 0 immediately -- never reaching + main()'s network/parquet-write path.""" + lock = str(tmp_path / "warm_cities.lock") + fd = os.open(lock, os.O_RDWR | os.O_CREAT, 0o644) + fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) + try: + env = dict(os.environ, THERMOGRAPH_DATA_DIR=str(tmp_path)) + # If the guard didn't work, this would try a real network fetch (or hang) + # instead of returning almost instantly with exit code 0. + result = subprocess.run( + [sys.executable, warm_cities.__file__, "--limit", "0"], + cwd=os.path.dirname(os.path.abspath(warm_cities.__file__)), + env=env, capture_output=True, text=True, timeout=20, + ) + assert result.returncode == 0 + assert "already running" in result.stdout + finally: + fcntl.flock(fd, fcntl.LOCK_UN) + os.close(fd) diff --git a/tests/web/test_gzip.py b/tests/web/test_gzip.py new file mode 100644 index 0000000..13bfae7 --- /dev/null +++ b/tests/web/test_gzip.py @@ -0,0 +1,12 @@ +"""GZipMiddleware's compresslevel: it runs synchronously on the event loop, so a +lower level (6, the standard speed/ratio sweet spot) trades a little ratio for +noticeably less per-request CPU competing with every other async route.""" +from starlette.middleware.gzip import GZipMiddleware + +from web import app as appmod + + +def test_gzip_middleware_uses_a_lowered_compresslevel(): + mw = next(m for m in appmod.app.user_middleware if m.cls is GZipMiddleware) + assert mw.kwargs.get("compresslevel") == 6 + assert mw.kwargs.get("compresslevel") != 9 # not Starlette's CPU-heavy default diff --git a/warm_cities.py b/warm_cities.py index 76f39f3..199a726 100644 --- a/warm_cities.py +++ b/warm_cities.py @@ -8,14 +8,29 @@ Idempotent: a cell whose archive is already cached is skipped. Fetches are paced (default 2s) to stay well under the archive API's rate limit. A cell that still has no cached archive when its page is first requested self-heals via get_history, so this is an optimization, not a hard dependency. + +Deploy launches this script detached on every deploy with no overlap check, so a +slow previous run (still mid-warm) can still be going when the next deploy's +invocation starts. Two overlapping runs would double-spend archive-fetch quota on +the same cells and race climate._write_cache's parquet writes, so the CLI entry +point below claims a single-instance flock (see LOCK_PATH) before calling main(). """ +import os import sys import time from api import homepage +from core import singleton from data import cities from data import climate from data import grid +import paths + +# core/singleton.py's flock guard, reused here for cross-*process* (not +# cross-worker) exclusion: the first invocation holds this for its process +# lifetime; a second one (an overlapping deploy) fails the non-blocking flock +# and stands down immediately. +LOCK_PATH = os.path.join(paths.DATA_DIR, "warm_cities.lock") def main(limit: int | None = None, pace: float = 2.0) -> None: @@ -53,6 +68,12 @@ def main(limit: int | None = None, pace: float = 2.0) -> None: if __name__ == "__main__": + if not singleton.claim(LOCK_PATH): + # Not an error: an overlapping run just means a previous deploy's warm is + # still in flight. Exit 0 so the deploy script doesn't treat this as a + # failure. + print("warm_cities: another instance is already running -- exiting") + sys.exit(0) args = sys.argv[1:] lim = int(args[args.index("--limit") + 1]) if "--limit" in args else None pc = float(args[args.index("--pace") + 1]) if "--pace" in args else 2.0 diff --git a/web/app.py b/web/app.py index 0de9cf8..59217ae 100644 --- a/web/app.py +++ b/web/app.py @@ -181,7 +181,16 @@ app = FastAPI(title="Thermograph", version="0.2.0", lifespan=_lifespan) # Compress every sizeable response (the 2-year calendar JSON shrinks ~6-8×). # Applies to API JSON and static assets alike; tiny responses are left alone. -app.add_middleware(GZipMiddleware, minimum_size=1024) +# compresslevel: GZipMiddleware compresses synchronously ON THE EVENT LOOP, so +# this is CPU work competing with every other async route in the worker while it +# runs — zlib's default level 9 spends a lot of that time chasing the last couple +# percent of ratio for a payload this size. 6 is the standard speed/ratio sweet +# spot (roughly zlib's own recommended default): most of level 9's compression +# but a fraction of the CPU, so a big calendar/day/cell-bundle response doesn't +# stall other requests on this worker for as long. Not worth a custom threaded +# compression middleware for payloads this size — the level drop captures most +# of the win for near-zero complexity. +app.add_middleware(GZipMiddleware, minimum_size=1024, compresslevel=6) # CORS (repo-split Stage 5): unset (default) means no CORSMiddleware at all -- # today's exact behavior, browsers already allow same-origin fetch with no