thermograph/backend/lake_app.py

"""The lake service: low-latency SQL-ish reads over the ERA5 bucket (prod-only).

A small FastAPI app run as the Swarm ``lake`` service (same backend image,
``THERMOGRAPH_ROLE=lake`` — the entrypoint launches this module instead of the
web app; no database, no migrations). Web replicas call it instead of touching
the bucket themselves, so a cell's history costs one LAN round-trip against a
warm disk cache instead of an S3 fetch, and analytical queries get
partition-pruned scans without shipping S3 credentials to every web task.

Endpoints:
- ``GET /healthz``            liveness (also reports manifest row count)
- ``GET /history?lat=&lon=``  the hot path: nearest ERA5 point's full daily
                              history as parquet bytes, disk-cached per point
                              (a plain single-file read — no engine needed)
- ``POST /query {"sql": ..}`` SELECT-only SQL over two views: ``era5_daily``
                              (the hive table, partition-pruned on
                              tile/year/month predicates) and ``manifest``.
                              Runs on DuckDB — a real lake engine (full SQL,
                              parallel scans, hive + row-group pushdown); the
                              httpfs extension is baked into the image so S3
                              needs no runtime extension fetch.

Scale-out safe: replicas share nothing — each keeps its own cache under
THERMOGRAPH_LAKE_CACHE (a per-task local volume), so 1..2 replicas behind the
Swarm VIP need no coordination.
"""
import os
import re
import threading
import time

import polars as pl
from fastapi import FastAPI, HTTPException
from fastapi.responses import FileResponse
from pydantic import BaseModel

from data import era5lake

CACHE_DIR = os.environ.get("THERMOGRAPH_LAKE_CACHE", "/state/lake-cache")
MANIFEST_TTL = 6 * 3600        # re-pull the manifest at most this often
QUERY_ROW_CAP = 10_000         # /query responses are JSON; keep them bounded

app = FastAPI(title="thermograph-lake")

_manifest_lock = threading.Lock()
_manifest: "pl.DataFrame | None" = None
_manifest_at = 0.0


def _source():
    """(mode, cfg) for this process: a local lake dir (tests/dev) or the bucket."""
    if era5lake.local_dir():
        return "local", None
    cfg = era5lake.s3_config()
    return ("bucket", cfg) if cfg else ("none", None)


def _manifest_frame() -> "pl.DataFrame | None":
    global _manifest, _manifest_at
    with _manifest_lock:
        if _manifest is not None and time.time() - _manifest_at < MANIFEST_TTL:
            return _manifest
        mode, cfg = _source()
        try:
            if mode == "local":
                df = pl.read_parquet(
                    os.path.join(era5lake.local_dir(), era5lake.MANIFEST_KEY))
            elif mode == "bucket":
                df = pl.read_parquet(
                    f"s3://{cfg['bucket']}/{era5lake.MANIFEST_KEY}",
                    storage_options=era5lake.storage_options(cfg))
            else:
                return None
        except Exception:  # noqa: BLE001 - empty/unreachable lake: serve health, 404 reads
            return None
        _manifest, _manifest_at = df, time.time()
        return df


@app.get("/healthz")
def healthz():
    mode, _ = _source()
    m = _manifest_frame()
    return {"status": "ok", "source": mode,
            "points": None if m is None else m.height}


@app.get("/history")
def history(lat: float, lon: float):
    """One ERA5 point's full daily history, as parquet bytes (cached on disk —
    the cache file IS the response body, so a warm hit is a sendfile)."""
    i, j = era5lake.to_idx(lat, lon)
    path = os.path.join(CACHE_DIR, f"p{i}_{j}.parquet")
    if not os.path.exists(path):
        mode, cfg = _source()
        if mode == "none":
            raise HTTPException(status_code=503, detail="lake not configured")
        m = _manifest_frame()
        if m is not None and m.filter(
                (pl.col("lat_idx") == i) & (pl.col("lon_idx") == j)).is_empty():
            raise HTTPException(status_code=404, detail="point not in the lake")
        try:
            if mode == "local":
                df = era5lake._read_local(i, j)
            else:
                df = era5lake._read_bucket(i, j, cfg)
        except era5lake.LakeUnavailable as e:
            raise HTTPException(status_code=404, detail=str(e)) from e
        os.makedirs(CACHE_DIR, exist_ok=True)
        tmp = f"{path}.{os.getpid()}.partial"
        df.write_parquet(tmp)
        os.replace(tmp, path)   # atomic: concurrent requests see whole files only
    return FileResponse(path, media_type="application/vnd.apache.parquet",
                        headers={"X-Lake-Point": f"{i},{j}"})


class Query(BaseModel):
    sql: str


_FORBIDDEN = re.compile(
    r"\b(insert|update|delete|drop|create|alter|attach|copy|truncate|install|load"
    r"|export|import|pragma|set|call)\b", re.I)


def _connect():
    """A fresh in-memory DuckDB with the lake views registered. Per-request:
    connections are cheap, and a throwaway one means a query can't leave state
    behind for the next. Returns None when no lake is configured."""
    import duckdb  # noqa: PLC0415 - only the lake role pays the import

    mode, cfg = _source()
    if mode == "local":
        daily_glob = os.path.join(era5lake.local_dir(), era5lake.PREFIX,
                                  "daily/*/*/*/*.parquet")
        manifest_src = os.path.join(era5lake.local_dir(), era5lake.MANIFEST_KEY)
    elif mode == "bucket":
        daily_glob = f"s3://{cfg['bucket']}/{era5lake.PREFIX}/daily/*/*/*/*.parquet"
        manifest_src = f"s3://{cfg['bucket']}/{era5lake.MANIFEST_KEY}"
    else:
        return None
    con = duckdb.connect()
    if mode == "bucket":
        con.execute("LOAD httpfs")   # baked into the image at build time
        host = cfg["endpoint"].split("://", 1)[-1]
        con.execute(f"SET s3_endpoint='{host}'")
        con.execute("SET s3_url_style='path'")
        con.execute(f"SET s3_region='{cfg['region']}'")
        con.execute("SET s3_access_key_id=?", [cfg["access_key"]])
        con.execute("SET s3_secret_access_key=?", [cfg["secret_key"]])
    con.execute(f"""
        CREATE VIEW era5_daily AS
        SELECT * FROM read_parquet('{daily_glob}', hive_partitioning=1)""")
    con.execute(f"CREATE VIEW manifest AS SELECT * FROM read_parquet('{manifest_src}')")
    return con


@app.post("/query")
def query(q: Query):
    sql = q.sql.strip().rstrip(";").strip()
    if not sql.lower().startswith(("select", "with")) or ";" in sql \
            or _FORBIDDEN.search(sql):
        raise HTTPException(status_code=400,
                            detail="only a single SELECT statement is allowed")
    con = _connect()
    if con is None:
        raise HTTPException(status_code=503, detail="lake not configured")
    try:
        out = con.execute(
            f"SELECT * FROM ({sql}) LIMIT {QUERY_ROW_CAP}").fetchall()
        cols = [d[0] for d in con.description]
    except Exception as e:  # noqa: BLE001 - surface the engine's message to the caller
        raise HTTPException(status_code=400, detail=f"query failed: {e}") from e
    finally:
        con.close()
    return {"columns": cols, "rows": out, "row_cap": QUERY_ROW_CAP}