thermograph/backend/tests/web/test_lake_app.py

"""The lake service over a local-dir lake fixture: health, the point hot path
(+ its disk cache), and the SELECT-only SQL surface."""
import datetime
import io

import polars as pl
import pytest
from fastapi.testclient import TestClient

# /query is gated on the same shared secret as /internal/* (it is an operator
# surface, not a public one). The fixture provisions it so the existing query
# tests exercise the authorised path; the auth tests below drive the closed ones.
TOKEN = "test-internal-token"
HDR = {"X-Thermograph-Internal-Token": TOKEN}


@pytest.fixture
def lake(tmp_path, monkeypatch):
    """A tiny two-point lake (points, hive partitions, manifest) on disk."""
    from data import era5lake

    monkeypatch.setenv("THERMOGRAPH_INTERNAL_TOKEN", TOKEN)

    points = [(154, 1439), (154, 0)]           # around London
    d0 = datetime.date(1940, 1, 1)
    dates = [d0 + datetime.timedelta(days=k) for k in range(12000)]
    n = len(dates)
    rows = []
    for i, j in points:
        df = pl.DataFrame({
            "date": dates,
            "tmax": [70.0] * n, "tmin": [50.0] * n, "precip": [0.1] * n,
            "wind": [8.0] * n, "gust": [14.0] * n, "humid": [55.0] * n,
            "fmax": [71.0] * n, "fmin": [48.0] * n,
        })
        p = tmp_path / era5lake.point_key(i, j)
        p.parent.mkdir(parents=True, exist_ok=True)
        df.write_parquet(p)
        ti, tj = era5lake.tile_of(i, j)
        part = df.with_columns(pl.lit(i, dtype=pl.Int32).alias("lat_idx"),
                               pl.lit(j, dtype=pl.Int32).alias("lon_idx"))
        hp = tmp_path / era5lake.daily_part_key(ti, tj, 1940, 1)
        hp.parent.mkdir(parents=True, exist_ok=True)
        part.filter((pl.col("date").dt.year() == 1940)
                    & (pl.col("date").dt.month() == 1)).write_parquet(hp)
        lat, lon = era5lake.to_coords(i, j)
        rows.append({"lat_idx": i, "lon_idx": j, "lat": lat, "lon": lon,
                     "tile": f"{ti}_{tj}", "rows": n,
                     "date_min": "1940-01-01", "date_max": str(dates[-1])})
    (tmp_path / "era5").mkdir(exist_ok=True)
    pl.DataFrame(rows).write_parquet(tmp_path / era5lake.MANIFEST_KEY)

    monkeypatch.setenv("THERMOGRAPH_LAKE_LOCAL_DIR", str(tmp_path))
    monkeypatch.setenv("THERMOGRAPH_LAKE_CACHE", str(tmp_path / "cache"))
    import lake_app
    monkeypatch.setattr(lake_app, "CACHE_DIR", str(tmp_path / "cache"))
    monkeypatch.setattr(lake_app, "_manifest", None)
    return TestClient(lake_app.app)


def test_healthz_reports_points(lake):
    body = lake.get("/healthz").json()
    assert body["status"] == "ok"
    assert body["source"] == "local"
    assert body["points"] == 2


def test_history_serves_parquet_and_caches(lake, tmp_path):
    r = lake.get("/history", params={"lat": 51.5074, "lon": -0.1278})
    assert r.status_code == 200
    assert r.headers["x-lake-point"] == "154,1439"
    df = pl.read_parquet(io.BytesIO(r.content))
    assert df.height == 12000
    assert (tmp_path / "cache" / "p154_1439.parquet").exists()
    # Second hit comes off the cache file and is byte-identical.
    assert lake.get("/history",
                    params={"lat": 51.5074, "lon": -0.1278}).content == r.content


def test_history_serves_from_memory_when_cache_unwritable(lake, tmp_path, monkeypatch):
    """An unwritable cache dir (bad volume ownership — seen live) degrades to
    serving from memory, never a 500."""
    import lake_app
    ro = tmp_path / "ro-cache" / "nested"   # parent made read-only below
    (tmp_path / "ro-cache").mkdir()
    (tmp_path / "ro-cache").chmod(0o555)
    monkeypatch.setattr(lake_app, "CACHE_DIR", str(ro))
    r = lake.get("/history", params={"lat": 51.5074, "lon": -0.1278})
    (tmp_path / "ro-cache").chmod(0o755)
    assert r.status_code == 200
    assert r.headers["x-lake-point"] == "154,1439"


def test_history_404_for_point_outside_lake(lake):
    r = lake.get("/history", params={"lat": -45.0, "lon": 170.0})
    assert r.status_code == 404


def test_query_selects_with_partition_predicates(lake):
    r = lake.post("/query", headers=HDR, json={"sql":
        "SELECT year, month, count(*) AS n FROM era5_daily "
        "WHERE year = 1940 AND month = 1 GROUP BY year, month"})
    assert r.status_code == 200
    body = r.json()
    assert body["columns"] == ["year", "month", "n"]
    assert body["rows"][0][2] == 2 * 31          # two points x January 1940


def test_query_manifest_view(lake):
    r = lake.post("/query", headers=HDR, json={"sql": "SELECT count(*) FROM manifest"})
    assert r.status_code == 200
    assert r.json()["rows"][0][0] == 2


@pytest.mark.parametrize("sql", [
    "DROP TABLE era5_daily",
    "SELECT 1; SELECT 2",
    "INSERT INTO era5_daily VALUES (1)",
    "COPY era5_daily TO 'x'",
])
def test_query_rejects_non_select(lake, sql):
    assert lake.post("/query", headers=HDR,
                     json={"sql": sql}).status_code == 400


# --- /query is an operator surface, not a public one -------------------------

def test_query_requires_the_internal_token(lake):
    """No header at all: the caller is not the operator. 401, not a result."""
    r = lake.post("/query", json={"sql": "SELECT count(*) FROM manifest"})
    assert r.status_code == 401


def test_query_rejects_a_wrong_internal_token(lake):
    r = lake.post("/query", headers={"X-Thermograph-Internal-Token": "nope"},
                  json={"sql": "SELECT count(*) FROM manifest"})
    assert r.status_code == 401


def test_query_is_closed_when_no_token_is_provisioned(lake, monkeypatch):
    """Unprovisioned means the surface does not exist (404) rather than falling
    open to no-auth — same posture as /internal/* and /api/v2/metrics."""
    monkeypatch.delenv("THERMOGRAPH_INTERNAL_TOKEN", raising=False)
    monkeypatch.delenv("THERMOGRAPH_AUTH_SECRET", raising=False)
    r = lake.post("/query", headers=HDR,
                  json={"sql": "SELECT count(*) FROM manifest"})
    assert r.status_code == 404


# --- the SELECT guard must reject reads outside the two lake views -----------
#
# The old guard was a denylist of DDL/DML *verbs*, so it passed every one of
# these: DuckDB's reader functions and setting introspection are plain SELECTs.
# In bucket mode `SET s3_secret_access_key` made the bucket credentials readable
# straight back out via current_setting() -- the exact thing this service exists
# to avoid (see the module docstring). These are the attack shapes, not the
# shapes the old regex happened to be written for.

@pytest.mark.parametrize("sql", [
    "SELECT * FROM read_csv_auto('/etc/passwd')",
    "SELECT * FROM read_csv('/etc/passwd')",
    "SELECT * FROM read_parquet('/etc/passwd')",
    "SELECT * FROM read_json_auto('/etc/passwd')",
    "SELECT * FROM read_text('/etc/passwd')",
    "SELECT * FROM read_blob('/etc/passwd')",
    "SELECT * FROM glob('/etc/*')",
    "SELECT * FROM duckdb_settings()",
    "SELECT * FROM duckdb_secrets()",
    "SELECT current_setting('s3_secret_access_key')",
    "SELECT getenv('THERMOGRAPH_LAKE_S3_SECRET')",
    "SELECT * FROM 'file:///etc/passwd'",
    "SELECT * FROM '/etc/passwd'",
    # Nested one level down, so a prefix-only check doesn't pass it.
    "SELECT * FROM (SELECT * FROM read_csv_auto('/etc/passwd'))",
    "WITH x AS (SELECT * FROM glob('/etc/*')) SELECT * FROM x",
])
def test_query_rejects_reads_outside_the_lake_views(lake, sql):
    r = lake.post("/query", headers=HDR, json={"sql": sql})
    assert r.status_code == 400, f"guard let this through: {sql}"
    # Assert the GUARD refused it, not that DuckDB happened to choke on the file
    # format. Several of these (read_parquet on a text file, a bare '/etc/passwd')
    # error in the engine anyway, so a status-only assertion would still pass with
    # the guard deleted -- exactly the "negatives shaped like the regex" trap the
    # original four-case test fell into.
    detail = r.json()["detail"]
    assert not detail.startswith("query failed"), (
        f"rejected by the engine, not the guard: {sql} -> {detail}")


def test_query_still_allows_a_block_comment_inside_a_real_select(lake):
    """Defence in depth must not become a false positive: an inline comment in
    an otherwise ordinary query is fine."""
    r = lake.post("/query", headers=HDR,
                  json={"sql": "SELECT count(*) /* rows */ FROM manifest"})
    assert r.status_code == 200


def test_bucket_mode_uses_a_secret_not_a_readable_setting(monkeypatch, tmp_path):
    """Regression guard for the credential-disclosure shape: the S3 credentials
    must be installed as a DuckDB SECRET (opaque to current_setting) rather than
    via SET s3_secret_access_key, which hands them back to any caller."""
    import lake_app

    statements = []

    class _FakeCon:
        def execute(self, sql, *args):
            statements.append(sql)
            return self

        def fetchone(self):
            return ("s3://b/iceberg/era5_daily/metadata/1.metadata.json",)

        def close(self):
            pass

    monkeypatch.setattr(lake_app, "_source", lambda: ("bucket", {
        "endpoint": "https://example.com", "bucket": "b", "region": "r",
        "access_key": "AK", "secret_key": "SUPER-SECRET-KEY"}))
    monkeypatch.setitem(__import__("sys").modules, "duckdb",
                        type("m", (), {"connect": staticmethod(lambda: _FakeCon())}))

    lake_app._connect()
    joined = "\n".join(statements)
    assert "SECRET" in joined.upper() and "TYPE S3" in joined.upper()
    assert "SET s3_secret_access_key" not in joined
    assert "SET s3_access_key_id" not in joined
    # The key material goes as a bound parameter, never interpolated into SQL.
    assert "SUPER-SECRET-KEY" not in joined
    # The local filesystem is unreachable in bucket mode, and the caller cannot
    # turn it back on.
    assert "disabled_filesystems" in joined
    assert "lock_configuration" in joined