thermograph/backend/data/store.py

"""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
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
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 (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 _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
    for stmt in _PG_SCHEMA:
        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 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
        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."""
    try:
        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


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()
    if cache:
        try:
            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


# --- 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."""
    try:
        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:
    try:
        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."""
    out = {"db_path": os.path.abspath(DB_PATH), "derived": 0, "revgeo": 0, "bytes": 0}
    try:
        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