thermograph/backend/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 api/payloads.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 contextlib
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 / api/payloads.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 ""

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

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


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)."""
    try:
        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 history_max_date(cell_id: str) -> "str | None":
    """The ISO date (YYYY-MM-DD) of the cell's newest archived day, or None when
    there is no cached history (or Postgres is off/unreachable).

    Backs the content-page cache token (api/payloads.content_token), so it must be
    CHEAP — an indexed ``MAX(date)`` over the ``(cell_id, date)`` primary key (see
    the 0002 migration), never a full-history load. Fail-soft: any error reads as
    None, so the caller falls back to a 'none' bucket rather than raising."""
    try:
        with _conn() as conn:
            if conn is None:
                return None
            row = conn.execute(
                "SELECT MAX(date) FROM climate_history WHERE cell_id = %s",
                (cell_id,),
            ).fetchone()
            if not row or row[0] is None:
                return None
            return row[0].isoformat()
    except Exception:  # noqa: BLE001
        return None


def cols_version(cell_id: str) -> int:
    """The stored schema version for a cell's history; 0 if absent."""
    try:
        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


# --- 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)."""
    try:
        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


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


# --- 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."""
    try:
        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


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."""
    try:
        with _conn() as conn:
            if conn is None:
                return
            _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."""
    try:
        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