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 io
import os
import re
import threading
import time

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

from api.internal_routes import _require_internal_token
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
        try:
            os.makedirs(CACHE_DIR, exist_ok=True)
            tmp = f"{path}.{os.getpid()}.partial"
            df.write_parquet(tmp)
            os.replace(tmp, path)   # atomic: concurrent readers see whole files only
        except OSError:
            # An unwritable cache (bad volume ownership, full disk) must never
            # fail the hot path — serve from memory; the cache is an
            # accelerator, not a dependency.
            buf = io.BytesIO()
            df.write_parquet(buf)
            return Response(buf.getvalue(),
                            media_type="application/vnd.apache.parquet",
                            headers={"X-Lake-Point": f"{i},{j}"})
    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)

# The verb denylist above is not a security boundary on its own: DuckDB's file
# readers and setting introspection are plain SELECTs, so every one of these used
# to sail past it. `SELECT current_setting('s3_secret_access_key')` handed the
# bucket credentials straight back to the caller -- the precise outcome the
# module docstring says this service exists to prevent -- and read_csv_auto()/
# glob() turned /query into an arbitrary read of the container filesystem.
#
# A denylist over a full query language is not winnable, so this is defence in
# depth only. The real boundaries are, in order: the shared-secret gate on the
# route, DuckDB SECRETs (opaque to current_setting) instead of SET s3_*, and
# disabled_filesystems in bucket mode. See _connect().
_DENIED_FUNCS = re.compile(
    r"\b(read_csv|read_csv_auto|read_parquet|parquet_scan|parquet_metadata"
    r"|parquet_schema|read_json|read_json_auto|read_json_objects|read_ndjson"
    r"|read_ndjson_auto|read_text|read_blob|glob|iceberg_scan|iceberg_metadata"
    r"|delta_scan|sniff_csv|current_setting|getenv|duckdb_settings"
    r"|duckdb_secrets|duckdb_extensions|duckdb_functions|sql_auto_complete"
    r"|which_secret|open|shell)\s*\(", re.I)

# `FROM 'some/path'` is DuckDB's replacement scan: a bare string literal in table
# position reads that path as a file. Legitimate queries name era5_daily or
# manifest, never a literal.
_FROM_LITERAL = re.compile(r"\b(from|join)\s*\(?\s*['\"]", re.I)

_GUARD_DETAIL = "only reads of the era5_daily and manifest views are allowed"


def _connect():
    """A fresh in-memory DuckDB with the lake views registered. Returns None
    when no lake is configured.

    Bucket mode reads era5_daily through the ICEBERG table, not a hive glob:
    a glob over era5/daily/ must LIST every object under it (~1k per tile —
    minutes at a few hundred tiles, seen live timing out /query at 390),
    while Iceberg gets the file list from table metadata in a handful of
    reads. Local mode (tests, small trees) keeps the plain hive glob — no
    Iceberg table exists in the fixtures and listing a local dir is free."""
    import duckdb  # noqa: PLC0415 - only the lake role pays the import

    mode, cfg = _source()
    if mode == "none":
        return None
    con = duckdb.connect()
    if mode == "local":
        daily_glob = os.path.join(era5lake.local_dir(), era5lake.PREFIX,
                                  "daily/*/*/*/*.parquet")
        con.execute(f"""
            CREATE VIEW era5_daily AS
            SELECT * FROM read_parquet('{daily_glob}', hive_partitioning=1)""")
        manifest_src = os.path.join(era5lake.local_dir(), era5lake.MANIFEST_KEY)
        con.execute(
            f"CREATE VIEW manifest AS SELECT * FROM read_parquet('{manifest_src}')")
        return con

    con.execute("LOAD httpfs")   # both baked into the image at build time,
    con.execute("LOAD iceberg")  # as the runtime user (see Dockerfile)
    host = cfg["endpoint"].split("://", 1)[-1]
    # A SECRET, not `SET s3_secret_access_key`: settings are readable back out
    # with current_setting(), so the old form let any caller of /query retrieve
    # the bucket credentials. Secret values are opaque -- duckdb_secrets() lists
    # the name and scope but never the key material.
    con.execute(
        "CREATE OR REPLACE SECRET lake (TYPE S3, KEY_ID ?, SECRET ?, "
        "REGION ?, ENDPOINT ?, URL_STYLE 'path')",
        [cfg["access_key"], cfg["secret_key"], cfg["region"], host])
    # Newest metadata JSON wins (one single-page LIST of metadata/ — cheap);
    # pyiceberg names them NNNNN-uuid.metadata.json, so lexicographic max is
    # the numeric max.
    meta = con.execute(
        f"SELECT max(file) FROM glob('s3://{cfg['bucket']}/iceberg/era5_daily/"
        f"metadata/*.metadata.json')").fetchone()[0]
    if meta is None:
        return None   # lake bucket exists but the Iceberg table doesn't yet
    con.execute(f"CREATE VIEW era5_daily AS SELECT * FROM iceberg_scan('{meta}')")
    con.execute(f"""CREATE VIEW manifest AS SELECT * FROM
        read_parquet('s3://{cfg['bucket']}/{era5lake.MANIFEST_KEY}')""")
    # Views are resolved lazily, so this must come after they are created --
    # and only in bucket mode, where every scan goes over httpfs and nothing
    # legitimate touches local disk. It makes read_csv_auto('/etc/passwd') and
    # glob('/etc/*') a PermissionException in the engine rather than relying on
    # the regex above. lock_configuration then stops a caller re-enabling it
    # (belt and braces: `set` is already a forbidden verb).
    con.execute("SET disabled_filesystems='LocalFileSystem'")
    con.execute("SET lock_configuration=true")
    return con


@app.post("/query", dependencies=[Depends(_require_internal_token)])
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")
    if _DENIED_FUNCS.search(sql) or _FROM_LITERAL.search(sql):
        raise HTTPException(status_code=400, detail=_GUARD_DETAIL)
    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}