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