thermograph/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
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