thermograph/data/climate_store.py

"""Raw climate record store — TimescaleDB hypertables (Postgres) behind polars.

The per-cell daily record (the multi-decade archive history and the hourly
recent+forecast bundle) is the *source of truth* for grading. On Postgres it
lives in two tables managed by Alembic (``backend/alembic/versions``):

* ``climate_history`` — a TimescaleDB **hypertable** (durable/LOGGED) of the full
  daily archive, keyed ``(cell_id, date)``. Expensive to refetch (45 years,
  rate-limited upstream), so it is durable, not a throwaway cache.
* ``climate_recent`` — a plain table holding the recent-observations + forward
  forecast bundle (includes *future* dates), fully rewritten each refresh.
* ``climate_sync`` — per-cell freshness (epoch seconds), replacing the parquet
  file mtimes that used to drive the topup cadence, the forecast TTL, and the
  ``recent_stamp`` token embedded in derived-payload validity (see web/views.py).

This module is the Postgres backend only. ``climate.py`` keeps the parquet cache
as the backend when ``THERMOGRAPH_DATABASE_URL`` is *not* a Postgres URL (dev,
tests, offline tooling) — the dialect switch mirrors ``data/store.py`` and
``accounts/db.py`` exactly, which is what keeps the test suite Postgres-free.

Every helper is **fail-soft**: a connection or query that fails is swallowed and
returns "absent" (``None`` / ``0.0``). Unlike ``store.py`` (which degrades to
recompute-from-parquet), a miss here degrades to *fetch-from-upstream* — prod has
no parquet fallback, and ``climate.py``'s loaders already treat "no cache" as
"go fetch".
"""
import os
import threading

import polars as pl

# Bump when a new persisted metric column is added to the daily record. A cell
# whose stored rows predate the bump reads back as schema-incomplete (see
# read_history), so climate.py refetches to fill the new column — the Postgres
# analogue of the parquet NEW_COLS presence check.
COLS_VERSION = 1

# The persisted daily columns, in a stable order. ``humid`` is the RAW relative
# humidity the archive returns; absolute humidity + wet-bulb are derived at the
# read boundary in climate._derive_metrics, never stored (so the store matches
# what the parquet write path persisted). ``feels`` IS stored. ``doy`` is not —
# climate._with_doy recomputes it at read time.
COLS = ("date", "tmax", "tmin", "precip", "wind", "gust", "humid", "fmax", "fmin", "feels")

# Explicit polars schema so a frame read back from Postgres has the exact dtypes
# the parquet path produced (pl.Date + Float64), even when the result is empty —
# so grading.py / scoring.py / web/views.py are untouched by the backend switch.
_SCHEMA: dict[str, pl.DataType] = {
    "date": pl.Date,
    **{c: pl.Float64 for c in COLS if c != "date"},
}

DATABASE_URL = os.environ.get("THERMOGRAPH_DATABASE_URL", "").strip()
IS_POSTGRES = DATABASE_URL.startswith("postgresql")


def is_postgres() -> bool:
    """True when the climate record is served from Postgres (prod), False when it
    falls back to the parquet cache (dev / tests / offline)."""
    return IS_POSTGRES


def _libpq_dsn(url: str) -> str:
    """A plain libpq URL psycopg 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 ""

# 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()


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."""
    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


# --- freshness (climate_sync) -----------------------------------------------

def _get_sync(conn, cell_id):
    """(history_synced_at, recent_synced_at, cols_version) for a cell, or None."""
    row = conn.execute(
        "SELECT history_synced_at, recent_synced_at, cols_version "
        "FROM climate_sync WHERE cell_id = %s",
        (cell_id,),
    ).fetchone()
    return row


def _bump_sync(conn, cell_id, *, history=None, recent=None, cols_version=None) -> None:
    """Upsert one cell's freshness. A None field leaves the stored value unchanged
    on update (COALESCE), so a recent-bundle write never disturbs history_synced_at
    and vice-versa. On insert, cols_version defaults to 1."""
    conn.execute(
        """
        INSERT INTO climate_sync (cell_id, history_synced_at, recent_synced_at, cols_version)
        VALUES (%(c)s, %(h)s, %(r)s, COALESCE(%(v)s, 1))
        ON CONFLICT (cell_id) DO UPDATE SET
            history_synced_at = COALESCE(%(h)s, climate_sync.history_synced_at),
            recent_synced_at  = COALESCE(%(r)s, climate_sync.recent_synced_at),
            cols_version      = COALESCE(%(v)s, climate_sync.cols_version)
        """,
        {"c": cell_id, "h": history, "r": recent, "v": cols_version},
    )


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
    except Exception:  # noqa: BLE001
        return 0.0


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
    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
    except Exception:  # noqa: BLE001
        return 0


# --- reads ------------------------------------------------------------------

def _read_frame(conn, table: str, cell_id: str) -> pl.DataFrame:
    """A cell's rows from ``table`` as a polars frame with the canonical schema.
    ``table`` is a module constant, never user input."""
    rows = conn.execute(
        f"SELECT {', '.join(COLS)} FROM {table} WHERE cell_id = %s ORDER BY date",
        (cell_id,),
    ).fetchall()
    return pl.DataFrame(rows, schema=_SCHEMA, orient="row")


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])
    except Exception:  # noqa: BLE001
        return None


def read_history_raw(cell_id: str) -> "pl.DataFrame | None":
    """A cell's history rows regardless of schema version — for the stale-serve
    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
    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])
    except Exception:  # noqa: BLE001
        return None


# --- writes -----------------------------------------------------------------

def _copy_rows(cur, table: str, cell_id: str, df: pl.DataFrame) -> None:
    """COPY a cell's frame into ``table`` (text format). NaN was already folded to
    null upstream in climate._finalize_frame, so nulls copy cleanly."""
    with cur.copy(f"COPY {table} (cell_id, {', '.join(COLS)}) FROM STDIN") as copy:
        for row in df.select(COLS).iter_rows():
            copy.write_row((cell_id, *row))


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)
    except Exception:  # noqa: BLE001 - failing to cache must not fail the request
        pass


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)
    except Exception:  # noqa: BLE001
        pass


def write_recent(cell_id: str, df: pl.DataFrame, synced_at: float) -> None:
    """Replace a cell's recent+forecast bundle (delete + COPY in one transaction)
    and stamp recent_synced_at. The window slides and forecast values change each
    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)
    except Exception:  # noqa: BLE001
        pass