"""Persistent derived-data store — SQLite (WAL) at data/thermograph.sqlite. The raw per-cell parquet records stay the source of truth; this database holds only what's *derived* from them, so expensive work is done once per cell and survives restarts/deploys instead of being recomputed per request, per view: * ``derived`` — finished API response payloads (grade / calendar / day / forecast), zlib-compressed JSON keyed by (kind, cell, request key) and guarded by a validity ``token``. The token encodes everything the payload depends on (payload schema version, history end date, recent-fetch stamp…); a row whose stored token doesn't match the caller's current token is simply a miss, so stale data can never be served — freshness is driven by the token advancing, never by clock-based expiry. * ``revgeo`` — reverse-geocode labels ("City, State" per ~cell), previously an in-memory dict that re-hit Nominatim after every restart. The store is a pure accelerator: every reader falls back to recomputing from parquet when the database is missing, locked, or corrupt (all helpers swallow 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. """ import json import os import sqlite3 import threading import time import zlib import paths DB_PATH = os.path.join(paths.DATA_DIR, "thermograph.sqlite") # Dialect selection mirrors accounts/db.py: a ``postgresql://`` URL in # THERMOGRAPH_DATABASE_URL switches the store to Postgres; anything else keeps the # SQLite file above. (This module uses the raw driver, not SQLAlchemy, so the DSN # is a plain libpq URL — the ``+asyncpg`` / ``+psycopg`` SQLAlchemy suffix stripped.) DATABASE_URL = os.environ.get("THERMOGRAPH_DATABASE_URL", "").strip() IS_POSTGRES = DATABASE_URL.startswith("postgresql") def _libpq_dsn(url: str) -> str: """A plain libpq URL that ``psycopg.connect`` accepts: drop the SQLAlchemy driver suffix (``postgresql+asyncpg://`` / ``+psycopg://`` -> ``postgresql://``).""" return url.replace("+asyncpg", "").replace("+psycopg", "") _PG_DSN = _libpq_dsn(DATABASE_URL) if IS_POSTGRES else "" _SCHEMA = """ CREATE TABLE IF NOT EXISTS derived ( kind TEXT NOT NULL, -- payload family: grade | calendar | day | forecast cell_id TEXT NOT NULL, key TEXT NOT NULL, -- request identity within the kind (dates/spans/params) token TEXT NOT NULL, -- validity: mismatched token == miss (see module docstring) payload BLOB NOT NULL, -- zlib-compressed JSON bytes updated_at REAL NOT NULL, PRIMARY KEY (kind, cell_id, key) ); CREATE TABLE IF NOT EXISTS revgeo ( key TEXT PRIMARY KEY, -- "lat,lon" rounded to ~cell precision label TEXT, -- NULL == lookup ran but found no label updated_at REAL NOT NULL ); """ # Postgres mirror of _SCHEMA. UNLOGGED = fast, crash-non-durable (the cache is # rebuildable, so "delete = safe reset" still holds); BLOB becomes BYTEA and REAL # becomes DOUBLE PRECISION. psycopg's extended protocol runs one statement per # execute, so the DDL is a tuple executed in turn (not one script like SQLite). _PG_SCHEMA = ( """ CREATE UNLOGGED TABLE IF NOT EXISTS derived ( kind TEXT NOT NULL, cell_id TEXT NOT NULL, key TEXT NOT NULL, token TEXT NOT NULL, payload BYTEA NOT NULL, updated_at DOUBLE PRECISION NOT NULL, PRIMARY KEY (kind, cell_id, key) ) """, """ CREATE UNLOGGED TABLE IF NOT EXISTS revgeo ( key TEXT PRIMARY KEY, label TEXT, updated_at DOUBLE PRECISION NOT NULL ) """, ) # Statement text differs by dialect (placeholder style + upsert syntax); the SQLite # strings are byte-for-byte the originals. Postgres uses %s params, an explicit # column list, and ON CONFLICT ... DO UPDATE — which, like SQLite's INSERT OR # REPLACE, overwrites *every* non-key column (token included, so a refreshed token # supersedes the stored one and freshness keeps tracking the token, never the clock). if IS_POSTGRES: _SQL_GET_PAYLOAD = "SELECT token, payload FROM derived WHERE kind=%s AND cell_id=%s AND key=%s" _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" ) _SQL_GET_REVGEO = "SELECT label, updated_at FROM revgeo WHERE key=%s" _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" ) else: _SQL_GET_PAYLOAD = "SELECT token, payload FROM derived WHERE kind=? AND cell_id=? AND key=?" _SQL_PUT_PAYLOAD = "INSERT OR REPLACE INTO derived VALUES (?,?,?,?,?,?)" _SQL_GET_REVGEO = "SELECT label, updated_at FROM revgeo WHERE key=?" _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). _local = threading.local() 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).""" 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 def _conn(): 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: 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 # --- 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) except Exception: # noqa: BLE001 return None def get_json(kind: str, cell_id: str, key: str, token: str) -> dict | None: """Like get_payload but decoded — for callers assembling composite responses.""" body = get_payload(kind, cell_id, key, token) return json.loads(body) if body is not None else None def put_payload(kind: str, cell_id: str, key: str, token: str, payload: dict, cache: bool = True) -> bytes: """Encode ``payload`` to compact JSON, persist it, and return the encoded bytes (so the caller can serve the exact same bytes without re-serializing). Pass ``cache=False`` to serve the payload without persisting it — used when a best-effort field (the reverse-geocoded place name) failed to resolve, so the incomplete result isn't cached for the life of the token. allow_nan=False mirrors Starlette's JSONResponse: a NaN slipping through 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: try: 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 # --- reverse-geocode labels -------------------------------------------------- REVGEO_MISS_TTL = 3600.0 # a "no label" result may be transient (rate limit) — retry hourly def revgeo_key(lat: float, lon: float) -> str: # The "v2:" prefix versions the label format: v1 rows (no country) are bypassed # and re-fetched so names pick up the trailing country component. return f"v2:{round(lat, 3)},{round(lon, 3)}" 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]) 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() 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)) except Exception: # noqa: BLE001 pass return out