infra/forgejo: give the LAN runner's PATH ~/.local/bin #84
28 changed files with 1900 additions and 41 deletions
|
|
@ -70,4 +70,10 @@ jobs:
|
|||
# thermograph-lan is host-native (the runner job runs directly on the
|
||||
# LAN box, not over SSH like beta/prod), so plain env vars on the
|
||||
# command reach deploy-dev.sh directly -- no ssh-action needed here.
|
||||
# The lake creds ride the Actions S3 secrets: dev renders no vault
|
||||
# (deploy-dev.sh's no-op secrets path), and compose interpolates
|
||||
# ${THERMOGRAPH_LAKE_S3_*} straight from the deploy environment.
|
||||
env:
|
||||
THERMOGRAPH_LAKE_S3_ACCESS_KEY: ${{ secrets.S3_ACCESS_KEY }}
|
||||
THERMOGRAPH_LAKE_S3_SECRET_KEY: ${{ secrets.S3_SECRET_KEY }}
|
||||
run: SERVICE=backend BACKEND_IMAGE_TAG=${{ steps.tag.outputs.tag }} bash "$HOME/thermograph-dev/infra/deploy/deploy-dev.sh"
|
||||
|
|
|
|||
|
|
@ -33,6 +33,9 @@ RUN apt-get update \
|
|||
# Install deps first so this layer caches across code-only changes.
|
||||
COPY requirements.txt /tmp/requirements.txt
|
||||
RUN pip install --no-cache-dir -r /tmp/requirements.txt
|
||||
# Bake DuckDB's httpfs extension so the lake role's S3 reads need no runtime
|
||||
# extension download (tasks may roll while egress is degraded).
|
||||
RUN python -c "import duckdb; duckdb.connect().execute('INSTALL httpfs')"
|
||||
|
||||
# The daemon binary lands before the app tree: it changes far less often than
|
||||
# the Python code, so this layer usually cache-hits and only the COPY below
|
||||
|
|
@ -49,8 +52,8 @@ RUN chmod +x /app/deploy/entrypoint.sh
|
|||
# Docker seeds them from this image content -- including this ownership --
|
||||
# so they stay writable.
|
||||
RUN useradd --system --create-home --uid 10001 thermograph \
|
||||
&& mkdir -p /app/data /app/logs \
|
||||
&& chown -R thermograph:thermograph /app
|
||||
&& mkdir -p /app/data /app/logs /state \
|
||||
&& chown -R thermograph:thermograph /app /state
|
||||
|
||||
USER thermograph
|
||||
WORKDIR /app
|
||||
|
|
|
|||
|
|
@ -21,6 +21,7 @@ from core import audit
|
|||
from core import metrics
|
||||
import paths
|
||||
from data import climate_store
|
||||
from data import era5lake
|
||||
from data import meteostat
|
||||
from data import store
|
||||
|
||||
|
|
@ -789,21 +790,32 @@ def _load_history(cell: dict) -> tuple[pl.DataFrame, dict]:
|
|||
def _serve_stale():
|
||||
return _with_doy(stale), {"cached": True, "stale": True, "cache_age_days": None}
|
||||
|
||||
# Fetch. NASA POWER is the primary source (keyless, independent of Open-Meteo);
|
||||
# the gusts it lacks are filled from Meteostat inside _fetch_history_nasa.
|
||||
# Open-Meteo is the fallback when NASA is unavailable — still guarded by its
|
||||
# rate-limit cooldown. Both sources must return a plausibly-full span before
|
||||
# being accepted and cached as a complete record (a short/partial response is
|
||||
# rejected rather than held indefinitely).
|
||||
# Fetch. The ERA5 lake is first when configured (true ERA5 with measured
|
||||
# gusts, served from our own bucket — no third-party API in the path);
|
||||
# unconfigured or missing-point cells fall through at zero cost. NASA
|
||||
# POWER is next (keyless, independent of Open-Meteo; its missing gusts
|
||||
# are filled from Meteostat inside _fetch_history_nasa), then Open-Meteo
|
||||
# — still guarded by its rate-limit cooldown. Every source must return a
|
||||
# plausibly-full span before being accepted and cached as a complete
|
||||
# record (a short/partial response is rejected rather than held
|
||||
# indefinitely).
|
||||
df = None
|
||||
source = None
|
||||
try:
|
||||
fetched = _fetch_history_nasa(cell)
|
||||
fetched = era5lake.fetch_history_lake(cell, start=START_DATE)
|
||||
if fetched.height >= MIN_ARCHIVE_DAYS:
|
||||
df = fetched
|
||||
source = "nasa-power"
|
||||
except Exception: # noqa: BLE001 - primary unavailable; fall through to Open-Meteo
|
||||
source = "era5-lake"
|
||||
except Exception: # noqa: BLE001 - lake unconfigured/miss; fall through to NASA
|
||||
pass
|
||||
if df is None:
|
||||
try:
|
||||
fetched = _fetch_history_nasa(cell)
|
||||
if fetched.height >= MIN_ARCHIVE_DAYS:
|
||||
df = fetched
|
||||
source = "nasa-power"
|
||||
except Exception: # noqa: BLE001 - primary unavailable; fall through to Open-Meteo
|
||||
pass
|
||||
if df is None and time.time() >= _archive_cooldown_until:
|
||||
try:
|
||||
fetched = _fetch_history(cell)
|
||||
|
|
|
|||
173
backend/data/era5lake.py
Normal file
173
backend/data/era5lake.py
Normal file
|
|
@ -0,0 +1,173 @@
|
|||
"""The ERA5 lake: layout, grid math, and read clients for the bucket-hosted archive.
|
||||
|
||||
The lake is a set of parquet files in an S3-compatible bucket (Contabo Object
|
||||
Storage, ``era5-thermograph``), extracted once from the public Earthmover ERA5
|
||||
Icechunk archive by ``gen_era5_lake.py`` and read three ways:
|
||||
|
||||
- the **lake service** (``lake_app.py``, prod-only Swarm service) serves point
|
||||
histories and SQL-style queries with a local disk cache in front;
|
||||
- the **web/worker backend** asks the lake service first
|
||||
(``THERMOGRAPH_LAKE_URL``), or reads the bucket directly when only S3
|
||||
credentials are configured (beta/LAN, no lake service);
|
||||
- ad-hoc analysis hits the hive table straight from the bucket.
|
||||
|
||||
Two layouts, one manifest (all under ``era5/``):
|
||||
|
||||
- ``daily/tile=<ti>_<tj>/year=<Y>/month=<M>/part.parquet`` — the analytical
|
||||
table, hive-partitioned by location/year/month. The location grain is a
|
||||
12x12-point tile (~3°x3°), aligned to the Earthmover store's chunk grid; a
|
||||
per-point x per-month grain would be ~170M objects, which no object store
|
||||
enjoys.
|
||||
- ``points/lat=<i>/lon=<j>.parquet`` — the serving projection: one file per
|
||||
ERA5 grid point holding its full daily history, so the hot path (one cell's
|
||||
history) is a single GET.
|
||||
- ``manifest.parquet`` — Iceberg-manifest-style index of every point file:
|
||||
grid indices, coordinates, row count, date span. Readers prune with it and
|
||||
the extractor resumes from it.
|
||||
|
||||
Columns everywhere: the finalized history-store schema (date, tmax, tmin,
|
||||
precip, wind, gust, humid, fmax, fmin, feels, doy — °F/in/mph/%, as
|
||||
``climate._finalize_approximated`` emits it), so a lake frame drops straight
|
||||
into ``climate._write_history_backed`` exactly like a NASA fetch does.
|
||||
"""
|
||||
import io
|
||||
import os
|
||||
|
||||
import polars as pl
|
||||
|
||||
# --- ERA5 native grid (0.25°, 721 x 1440, lat 90..-90, lon 0..359.75) ---------
|
||||
LAT_N = 721
|
||||
LON_N = 1440
|
||||
STEP = 0.25
|
||||
TILE = 12 # matches the Earthmover chunk grid (8736h x 12 x 12)
|
||||
|
||||
PREFIX = "era5"
|
||||
MANIFEST_KEY = f"{PREFIX}/manifest.parquet"
|
||||
|
||||
MIN_LAKE_DAYS = 10_000 # a plausibly-full 1980+ record, same spirit as
|
||||
# climate.MIN_ARCHIVE_DAYS (rejects partial files)
|
||||
|
||||
|
||||
def to_idx(lat: float, lon: float) -> tuple[int, int]:
|
||||
"""Nearest ERA5 grid indices for a lat/lon. lat_idx runs north->south."""
|
||||
i = round((90.0 - lat) / STEP)
|
||||
j = round((lon % 360.0) / STEP) % LON_N
|
||||
return min(max(i, 0), LAT_N - 1), j
|
||||
|
||||
|
||||
def to_coords(i: int, j: int) -> tuple[float, float]:
|
||||
"""Center lat/lon ([-180, 180) longitude) of grid indices."""
|
||||
lon = j * STEP
|
||||
return 90.0 - i * STEP, lon - 360.0 if lon >= 180.0 else lon
|
||||
|
||||
|
||||
def tile_of(i: int, j: int) -> tuple[int, int]:
|
||||
return i // TILE, j // TILE
|
||||
|
||||
|
||||
def point_key(i: int, j: int) -> str:
|
||||
return f"{PREFIX}/points/lat={i}/lon={j}.parquet"
|
||||
|
||||
|
||||
def daily_part_key(ti: int, tj: int, year: int, month: int) -> str:
|
||||
return f"{PREFIX}/daily/tile={ti}_{tj}/year={year}/month={month:02d}/part.parquet"
|
||||
|
||||
|
||||
# --- Config (all optional: unset means "no lake configured") ------------------
|
||||
|
||||
def lake_url() -> str:
|
||||
"""Lake service base URL (prod: the Swarm VIP), '' when not deployed."""
|
||||
return os.environ.get("THERMOGRAPH_LAKE_URL", "").rstrip("/")
|
||||
|
||||
|
||||
def s3_config() -> "dict | None":
|
||||
"""Bucket access for direct reads/writes, None unless fully configured.
|
||||
Contabo is path-style; region is a signing formality there ('default')."""
|
||||
key = os.environ.get("THERMOGRAPH_LAKE_S3_ACCESS_KEY", "")
|
||||
secret = os.environ.get("THERMOGRAPH_LAKE_S3_SECRET_KEY", "")
|
||||
if not (key and secret):
|
||||
return None
|
||||
return {
|
||||
"endpoint": os.environ.get("THERMOGRAPH_LAKE_S3_ENDPOINT",
|
||||
"https://eu2.contabostorage.com"),
|
||||
"bucket": os.environ.get("THERMOGRAPH_LAKE_S3_BUCKET", "era5-thermograph"),
|
||||
"region": os.environ.get("THERMOGRAPH_LAKE_S3_REGION", "default"),
|
||||
"access_key": key,
|
||||
"secret_key": secret,
|
||||
}
|
||||
|
||||
|
||||
def local_dir() -> str:
|
||||
"""Filesystem mirror of the lake layout — tests and offline dev reads."""
|
||||
return os.environ.get("THERMOGRAPH_LAKE_LOCAL_DIR", "")
|
||||
|
||||
|
||||
def storage_options(cfg: dict) -> dict:
|
||||
"""polars/object_store options for scan_parquet against the bucket."""
|
||||
return {
|
||||
"aws_access_key_id": cfg["access_key"],
|
||||
"aws_secret_access_key": cfg["secret_key"],
|
||||
"aws_region": cfg["region"],
|
||||
"endpoint_url": cfg["endpoint"],
|
||||
# Contabo serves buckets on the path, not as a subdomain.
|
||||
"aws_virtual_hosted_style_request": "false",
|
||||
}
|
||||
|
||||
|
||||
# --- Point-history reads (the app's hot path) ----------------------------------
|
||||
|
||||
class LakeUnavailable(Exception):
|
||||
"""No lake configured, point not in the lake, or the read failed."""
|
||||
|
||||
|
||||
def _read_local(i: int, j: int) -> pl.DataFrame:
|
||||
path = os.path.join(local_dir(), point_key(i, j))
|
||||
if not os.path.exists(path):
|
||||
raise LakeUnavailable(f"no local lake file for point ({i}, {j})")
|
||||
return pl.read_parquet(path)
|
||||
|
||||
|
||||
def _read_service(i: int, j: int) -> pl.DataFrame:
|
||||
import httpx # noqa: PLC0415 - keep module import light for the extractor venv
|
||||
|
||||
lat, lon = to_coords(i, j)
|
||||
r = httpx.get(f"{lake_url()}/history", params={"lat": lat, "lon": lon},
|
||||
timeout=httpx.Timeout(20.0, connect=3.0))
|
||||
if r.status_code == 404:
|
||||
raise LakeUnavailable(f"point ({i}, {j}) not in the lake")
|
||||
r.raise_for_status()
|
||||
return pl.read_parquet(io.BytesIO(r.content))
|
||||
|
||||
|
||||
def _read_bucket(i: int, j: int, cfg: dict) -> pl.DataFrame:
|
||||
try:
|
||||
return pl.read_parquet(f"s3://{cfg['bucket']}/{point_key(i, j)}",
|
||||
storage_options=storage_options(cfg))
|
||||
except Exception as e: # noqa: BLE001 - object_store errors are not one type
|
||||
raise LakeUnavailable(f"bucket read failed for point ({i}, {j}): {e}") from e
|
||||
|
||||
|
||||
def fetch_history_lake(cell: dict, start: "str | None" = None) -> pl.DataFrame:
|
||||
"""Finalized daily history for the ERA5 point nearest a grid cell, from the
|
||||
first configured source: local dir -> lake service -> bucket. The lake
|
||||
stores each point's whole 1940+ record; pass ``start`` (the caller's
|
||||
grading window, e.g. climate.START_DATE) to slice — grading semantics stay
|
||||
the caller's, the deep record stays in the lake. Raises LakeUnavailable
|
||||
when unconfigured, missing, or implausibly short."""
|
||||
import datetime # noqa: PLC0415
|
||||
|
||||
i, j = to_idx(cell["center_lat"], cell["center_lon"])
|
||||
if local_dir():
|
||||
df = _read_local(i, j)
|
||||
elif lake_url():
|
||||
df = _read_service(i, j)
|
||||
else:
|
||||
cfg = s3_config()
|
||||
if cfg is None:
|
||||
raise LakeUnavailable("no lake configured")
|
||||
df = _read_bucket(i, j, cfg)
|
||||
if start is not None:
|
||||
df = df.filter(pl.col("date") >= datetime.date.fromisoformat(start))
|
||||
if df.height < MIN_LAKE_DAYS:
|
||||
raise LakeUnavailable(f"lake point ({i}, {j}) is short: {df.height} rows")
|
||||
return df.drop([c for c in ("lat_idx", "lon_idx") if c in df.columns])
|
||||
|
|
@ -62,6 +62,13 @@ run_migrations() {
|
|||
done
|
||||
}
|
||||
|
||||
# The lake role serves the ERA5 bucket only — no database, so no migrations: a
|
||||
# lake replica must come up (and stay up) even when Postgres is down.
|
||||
if [ "${THERMOGRAPH_ROLE:-}" = "lake" ]; then
|
||||
echo "==> Starting lake service on 0.0.0.0:${PORT:-8141} with ${WORKERS:-1} worker(s)"
|
||||
exec uvicorn lake_app:app --host 0.0.0.0 --port "${PORT:-8141}" --workers "${WORKERS:-1}"
|
||||
fi
|
||||
|
||||
# One-shot migrate mode (`entrypoint.sh migrate`, or THERMOGRAPH_MIGRATE_ONLY=1): run
|
||||
# the migration and exit — a one-shot task, so it runs once against the DB rather
|
||||
# than racing per web replica.
|
||||
|
|
|
|||
237
backend/gen_era5_lake.py
Normal file
237
backend/gen_era5_lake.py
Normal file
|
|
@ -0,0 +1,237 @@
|
|||
"""Extract the ERA5 lake: Earthmover Icechunk archive -> the Thermograph bucket.
|
||||
|
||||
Builds the lake that ``data/era5lake.py`` reads (layout documented there):
|
||||
per-point serving files, the tile/year/month hive table, and the manifest.
|
||||
Pulls are tile-aligned to the source's chunk grid (8736h x 12 x 12) — one
|
||||
46-year pull (~19s measured) decodes all 144 points of a tile, so per-point
|
||||
pulls would waste 144x the wire. Roughly: ~600 tiles cover the curated
|
||||
cities (hours), ~3000 cover land (a day or two); both resume from the
|
||||
manifest, so ctrl-C and re-run is always safe.
|
||||
|
||||
pip install -r requirements-seed.txt 'numcodecs[pcodec]' boto3
|
||||
python gen_era5_lake.py --cities --dry-run # one tile, no writes
|
||||
python gen_era5_lake.py --cities # tiles covering cities.json
|
||||
python gen_era5_lake.py --tiles 20:30,40:60 # explicit tile-index window
|
||||
python gen_era5_lake.py --land # every land-containing tile
|
||||
|
||||
Destination: the bucket from THERMOGRAPH_LAKE_S3_* (see era5lake.s3_config),
|
||||
or a local mirror when THERMOGRAPH_LAKE_LOCAL_DIR is set (tests/rehearsal).
|
||||
Seed-only deps (icechunk/xarray/zarr) stay out of the app image, same policy
|
||||
as seed_era5.py.
|
||||
"""
|
||||
import argparse
|
||||
import concurrent.futures
|
||||
import datetime
|
||||
import io
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
|
||||
import polars as pl
|
||||
|
||||
from data import cities
|
||||
from data import climate
|
||||
from data import era5lake
|
||||
from seed_era5 import ERA5_BUCKET, ERA5_GROUP, ERA5_PREFIX, ERA5_REGION, ERA5_VARS, \
|
||||
hourly_to_daily
|
||||
|
||||
# The lake keeps each location's WHOLE available record (the store reaches
|
||||
# back to 1940); the app slices to its own grading window (climate.START_DATE)
|
||||
# at read time, so the extra decades are analytical surface, not a grading
|
||||
# change.
|
||||
START_DATE = "1940-01-01"
|
||||
|
||||
|
||||
def _open_store():
|
||||
import icechunk # noqa: PLC0415 - seed-only dep
|
||||
import xarray as xr # noqa: PLC0415
|
||||
|
||||
storage = icechunk.s3_storage(
|
||||
bucket=ERA5_BUCKET, prefix=ERA5_PREFIX, region=ERA5_REGION, anonymous=True)
|
||||
repo = icechunk.Repository.open(storage)
|
||||
return xr.open_zarr(repo.readonly_session("main").store,
|
||||
group=ERA5_GROUP, consolidated=False)
|
||||
|
||||
|
||||
class _Sink:
|
||||
"""Writes lake objects to the bucket (boto3) or a local mirror dir."""
|
||||
|
||||
def __init__(self):
|
||||
self.local = era5lake.local_dir()
|
||||
self.cfg = era5lake.s3_config()
|
||||
if not self.local and self.cfg is None:
|
||||
sys.exit("no destination: set THERMOGRAPH_LAKE_S3_* or "
|
||||
"THERMOGRAPH_LAKE_LOCAL_DIR")
|
||||
self._s3 = None
|
||||
if not self.local:
|
||||
import boto3 # noqa: PLC0415 - extractor-only dep
|
||||
from botocore.config import Config # noqa: PLC0415
|
||||
# Contabo throttles bursts with 429s (seen live at 16 concurrent
|
||||
# PUTs); adaptive mode rate-limits client-side and retries with
|
||||
# backoff instead of dying after botocore's default 4 attempts.
|
||||
self._s3 = boto3.client(
|
||||
"s3", endpoint_url=self.cfg["endpoint"],
|
||||
region_name=self.cfg["region"],
|
||||
aws_access_key_id=self.cfg["access_key"],
|
||||
aws_secret_access_key=self.cfg["secret_key"],
|
||||
config=Config(retries={"max_attempts": 10, "mode": "adaptive"}))
|
||||
|
||||
def put(self, key: str, df: pl.DataFrame) -> None:
|
||||
if self.local:
|
||||
path = os.path.join(self.local, key)
|
||||
os.makedirs(os.path.dirname(path), exist_ok=True)
|
||||
df.write_parquet(path)
|
||||
return
|
||||
buf = io.BytesIO()
|
||||
df.write_parquet(buf)
|
||||
self._s3.put_object(Bucket=self.cfg["bucket"], Key=key,
|
||||
Body=buf.getvalue())
|
||||
|
||||
def read_manifest(self) -> "pl.DataFrame | None":
|
||||
try:
|
||||
if self.local:
|
||||
return pl.read_parquet(
|
||||
os.path.join(self.local, era5lake.MANIFEST_KEY))
|
||||
obj = self._s3.get_object(Bucket=self.cfg["bucket"],
|
||||
Key=era5lake.MANIFEST_KEY)
|
||||
return pl.read_parquet(io.BytesIO(obj["Body"].read()))
|
||||
except Exception: # noqa: BLE001 - a fresh lake has no manifest yet
|
||||
return None
|
||||
|
||||
|
||||
def _tile_frame(ds, ti: int, tj: int, end: str) -> "pl.DataFrame | None":
|
||||
"""Pull one 12x12 tile's hourly series and return the finalized daily frame
|
||||
for all its points (lat_idx, lon_idx, date, ...). None for all-sea tiles."""
|
||||
t = era5lake.TILE
|
||||
i0, j0 = ti * t, tj * t
|
||||
lat_sl = slice(i0, min(i0 + t, era5lake.LAT_N))
|
||||
lon_sl = slice(j0, j0 + t)
|
||||
if "lsm" in ds: # land-sea mask: skip all-sea tiles outright
|
||||
lsm = ds["lsm"].isel(latitude=lat_sl, longitude=lon_sl)
|
||||
if "valid_time" in lsm.dims: # static field, stored with a time axis
|
||||
lsm = lsm.isel(valid_time=-1)
|
||||
if float(lsm.max()) < 0.05:
|
||||
return None
|
||||
sub = (ds[list(ERA5_VARS)]
|
||||
.isel(latitude=lat_sl, longitude=lon_sl)
|
||||
.sel(valid_time=slice(START_DATE, end))
|
||||
.load())
|
||||
|
||||
frames = []
|
||||
for a, lat in enumerate(sub.latitude.values):
|
||||
for b, lon in enumerate(sub.longitude.values):
|
||||
cols = {"time": sub.valid_time.values}
|
||||
for src, dst in ERA5_VARS.items():
|
||||
cols[dst] = sub[src].isel(latitude=a, longitude=b).values
|
||||
daily = climate._finalize_approximated(hourly_to_daily(pl.DataFrame(cols)))
|
||||
i, j = era5lake.to_idx(float(lat), float(lon))
|
||||
frames.append(daily.with_columns(
|
||||
pl.lit(i, dtype=pl.Int32).alias("lat_idx"),
|
||||
pl.lit(j, dtype=pl.Int32).alias("lon_idx")))
|
||||
return pl.concat(frames)
|
||||
|
||||
|
||||
def _write_tile(sink: _Sink, ti: int, tj: int, tile_df: pl.DataFrame) -> pl.DataFrame:
|
||||
"""Write one tile's point files + hive partitions; return its manifest rows.
|
||||
Uploads run on a thread pool: a tile is ~1.2k small objects and sequential
|
||||
PUTs are pure round-trip latency (measured 2.5 obj/s -> 8 min/tile; the
|
||||
pool brings that to well under a minute). boto3 clients are thread-safe."""
|
||||
puts = []
|
||||
rows = []
|
||||
for (i, j), pdf in tile_df.group_by(["lat_idx", "lon_idx"]):
|
||||
pdf = pdf.sort("date")
|
||||
puts.append((era5lake.point_key(i, j), pdf.drop(["lat_idx", "lon_idx"])))
|
||||
lat, lon = era5lake.to_coords(i, j)
|
||||
rows.append({"lat_idx": i, "lon_idx": j, "lat": lat, "lon": lon,
|
||||
"tile": f"{ti}_{tj}", "rows": pdf.height,
|
||||
"date_min": str(pdf["date"].min()),
|
||||
"date_max": str(pdf["date"].max())})
|
||||
parts = tile_df.with_columns(
|
||||
pl.col("date").dt.year().alias("year"),
|
||||
pl.col("date").dt.month().alias("month"))
|
||||
for (year, month), mdf in parts.group_by(["year", "month"]):
|
||||
puts.append((era5lake.daily_part_key(ti, tj, year, month),
|
||||
mdf.drop(["year", "month"]).sort(["lat_idx", "lon_idx", "date"])))
|
||||
with concurrent.futures.ThreadPoolExecutor(max_workers=8) as pool:
|
||||
for f in [pool.submit(sink.put, k, df) for k, df in puts]:
|
||||
f.result() # surface the first failure instead of hiding it
|
||||
return pl.DataFrame(rows)
|
||||
|
||||
|
||||
def _city_tiles() -> list:
|
||||
from data import grid # noqa: PLC0415
|
||||
|
||||
tiles = set()
|
||||
for c in cities.all_cities():
|
||||
cell = grid.snap(c["lat"], c["lon"])
|
||||
tiles.add(era5lake.tile_of(
|
||||
*era5lake.to_idx(cell["center_lat"], cell["center_lon"])))
|
||||
return sorted(tiles)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--cities", action="store_true")
|
||||
ap.add_argument("--land", action="store_true")
|
||||
ap.add_argument("--tiles", help="ti0:ti1,tj0:tj1 tile-index window")
|
||||
ap.add_argument("--limit", type=int)
|
||||
ap.add_argument("--overwrite", action="store_true")
|
||||
ap.add_argument("--dry-run", action="store_true")
|
||||
args = ap.parse_args()
|
||||
|
||||
if args.cities:
|
||||
todo = _city_tiles()
|
||||
elif args.tiles:
|
||||
(a, b), (c, d) = (p.split(":") for p in args.tiles.split(","))
|
||||
todo = [(ti, tj) for ti in range(int(a), int(b))
|
||||
for tj in range(int(c), int(d))]
|
||||
elif args.land:
|
||||
todo = [(ti, tj) for ti in range(era5lake.LAT_N // era5lake.TILE + 1)
|
||||
for tj in range(era5lake.LON_N // era5lake.TILE)]
|
||||
else:
|
||||
ap.error("pick one of --cities / --tiles / --land")
|
||||
if args.limit:
|
||||
todo = todo[:args.limit]
|
||||
|
||||
sink = _Sink()
|
||||
ds = _open_store()
|
||||
end = (datetime.date.today()
|
||||
- datetime.timedelta(days=climate.ARCHIVE_LATENCY_DAYS)).isoformat()
|
||||
|
||||
manifest = sink.read_manifest()
|
||||
done = set() if manifest is None or args.overwrite else \
|
||||
set(manifest["tile"].unique().to_list())
|
||||
|
||||
written = skipped = empty = failed = 0
|
||||
for n, (ti, tj) in enumerate(todo, 1):
|
||||
if f"{ti}_{tj}" in done:
|
||||
skipped += 1
|
||||
continue
|
||||
t0 = time.time()
|
||||
try:
|
||||
tile_df = _tile_frame(ds, ti, tj, end)
|
||||
if tile_df is None:
|
||||
empty += 1
|
||||
continue
|
||||
if args.dry_run:
|
||||
print(f"[dry-run] tile {ti}_{tj}: {tile_df.height} rows "
|
||||
f"({time.time() - t0:.1f}s); nothing written")
|
||||
print(tile_df.head())
|
||||
return
|
||||
rows = _write_tile(sink, ti, tj, tile_df)
|
||||
except Exception as e: # noqa: BLE001 - keep going; the tile isn't in the
|
||||
failed += 1 # manifest, so the next run retries it
|
||||
print(f"[{n}/{len(todo)}] FAILED tile {ti}_{tj}: {e}")
|
||||
continue
|
||||
manifest = rows if manifest is None else pl.concat(
|
||||
[manifest.filter(pl.col("tile") != f"{ti}_{tj}"), rows])
|
||||
sink.put(era5lake.MANIFEST_KEY, manifest) # after every tile: resumable
|
||||
written += 1
|
||||
print(f"[{n}/{len(todo)}] tile {ti}_{tj}: {rows.height} points "
|
||||
f"in {time.time() - t0:.1f}s")
|
||||
print(f"done: tiles written={written} skipped(done)={skipped} sea={empty} "
|
||||
f"failed={failed}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
187
backend/lake_app.py
Normal file
187
backend/lake_app.py
Normal file
|
|
@ -0,0 +1,187 @@
|
|||
"""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 FastAPI, HTTPException
|
||||
from fastapi.responses import FileResponse, Response
|
||||
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
|
||||
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)
|
||||
|
||||
|
||||
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}
|
||||
|
|
@ -7,3 +7,10 @@
|
|||
icechunk
|
||||
xarray
|
||||
zarr
|
||||
# The live store compresses with PCodec; without this extra every read fails
|
||||
# with "codec not available: 'pcodec'".
|
||||
numcodecs[pcodec]
|
||||
# gen_era5_lake.py's bucket writer.
|
||||
boto3
|
||||
numcodecs[pcodec]
|
||||
boto3
|
||||
|
|
|
|||
|
|
@ -24,3 +24,6 @@ PyNaCl==1.5.0
|
|||
# api/internal_routes.py — so the direct `websockets` pin (only discord_bot.py
|
||||
# imported it; uvicorn[standard] still pulls it transitively for its own ws
|
||||
# protocol support) and `apscheduler` are gone from this list.
|
||||
# The lake role's SQL engine (lake_app.py /query): embedded, parallel
|
||||
# partition-pruned scans over the ERA5 hive table. Only lake_app imports it.
|
||||
duckdb==1.4.2
|
||||
|
|
|
|||
|
|
@ -16,10 +16,10 @@ Idempotent by default: a cell that already has a cached history is skipped (pass
|
|||
--overwrite to reseed). Not run in CI or the app image — the icechunk/xarray/zarr
|
||||
stack lives only in requirements-seed.txt.
|
||||
|
||||
IMPORTANT — the S3/Icechunk *access* layer (_open_store / fetch_point) targets a
|
||||
young, fast-moving API and the store's exact variable/coordinate names; it is NOT
|
||||
exercised by the test suite. VERIFY it with `--dry-run` on the seed box before a
|
||||
full run. The hourly→daily transform (hourly_to_daily) IS unit-tested and is the
|
||||
The access layer was VERIFIED against the live store 2026-07-23 (prefix
|
||||
`icechunkV2`, group `single/temporal`, coord `valid_time`, ECMWF short variable
|
||||
names, and the pcodec codec — install `numcodecs[pcodec]`). Re-verify with
|
||||
`--dry-run` if the store revs. The hourly→daily transform (hourly_to_daily) IS unit-tested and is the
|
||||
part that must be numerically correct.
|
||||
|
||||
Source: https://registry.opendata.aws/earthmover-era5/ (group="temporal", anonymous)
|
||||
|
|
@ -38,19 +38,19 @@ START_DATE = climate.START_DATE # "1980-01-01"
|
|||
|
||||
# Anonymous AWS Open Data Icechunk store; overridable in case the registry moves.
|
||||
ERA5_BUCKET = os.environ.get("THERMOGRAPH_ERA5_BUCKET", "earthmover-icechunk-era5")
|
||||
ERA5_PREFIX = os.environ.get("THERMOGRAPH_ERA5_PREFIX", "") # verify on the seed box
|
||||
ERA5_PREFIX = os.environ.get("THERMOGRAPH_ERA5_PREFIX", "icechunkV2")
|
||||
ERA5_REGION = os.environ.get("THERMOGRAPH_ERA5_REGION", "us-east-1")
|
||||
ERA5_GROUP = os.environ.get("THERMOGRAPH_ERA5_GROUP", "temporal") # time-contiguous layout
|
||||
ERA5_GROUP = os.environ.get("THERMOGRAPH_ERA5_GROUP", "single/temporal") # time-contiguous layout
|
||||
|
||||
# ERA5 store variable name -> our short name. ERA5 gusts (i10fg) are real here, so the
|
||||
# seed keeps genuine measured gusts rather than the Meteostat/estimate fallback.
|
||||
ERA5_VARS = {
|
||||
"2m_temperature": "t2m", # K
|
||||
"2m_dewpoint_temperature": "d2m", # K (→ RH with t2m)
|
||||
"total_precipitation": "tp", # m, hourly accumulation
|
||||
"10m_u_component_of_wind": "u10", # m/s
|
||||
"10m_v_component_of_wind": "v10", # m/s
|
||||
"instantaneous_10m_wind_gust": "i10fg", # m/s
|
||||
"t2m": "t2m", # 2m temperature, K
|
||||
"d2m": "d2m", # 2m dewpoint, K (→ RH with t2m)
|
||||
"tp": "tp", # total precipitation, m, hourly accumulation
|
||||
"u10": "u10", # 10m wind u, m/s
|
||||
"v10": "v10", # 10m wind v, m/s
|
||||
"fg10": "i10fg", # 10m wind gust, m/s (store short name; transform expects i10fg)
|
||||
}
|
||||
|
||||
_M_TO_IN = 39.37007874
|
||||
|
|
@ -117,9 +117,9 @@ def fetch_point(ds, lat: float, lon: float, start: str, end: str) -> pl.DataFram
|
|||
may differ per store revision — verify with --dry-run."""
|
||||
sub = (ds[list(ERA5_VARS)]
|
||||
.sel(latitude=lat, longitude=lon % 360.0, method="nearest")
|
||||
.sel(time=slice(start, end))
|
||||
.sel(valid_time=slice(start, end))
|
||||
.load())
|
||||
data = {"time": sub["time"].values}
|
||||
data = {"time": sub["valid_time"].values}
|
||||
for src, dst in ERA5_VARS.items():
|
||||
data[dst] = sub[src].values
|
||||
return pl.DataFrame(data)
|
||||
|
|
|
|||
|
|
@ -311,6 +311,44 @@ def _full_history_frame(n=4000):
|
|||
}))
|
||||
|
||||
|
||||
def test_lake_is_first_history_source(monkeypatch, tmp_path):
|
||||
"""A configured ERA5 lake short-circuits the whole third-party chain: neither
|
||||
NASA POWER nor Open-Meteo is consulted."""
|
||||
monkeypatch.setattr(climate, "CACHE_DIR", str(tmp_path))
|
||||
monkeypatch.setattr(climate, "_archive_cooldown_until", 0.0)
|
||||
cell = {"id": "lake_primary", "center_lat": 47.6, "center_lon": -122.3}
|
||||
|
||||
full = _full_history_frame()
|
||||
monkeypatch.setattr(climate.era5lake, "fetch_history_lake",
|
||||
lambda c, start=None: full)
|
||||
def _nasa_should_not_run(c): raise AssertionError("NASA should not be called")
|
||||
def _om_should_not_run(c): raise AssertionError("Open-Meteo should not be called")
|
||||
monkeypatch.setattr(climate, "_fetch_history_nasa", _nasa_should_not_run)
|
||||
monkeypatch.setattr(climate, "_fetch_history", _om_should_not_run)
|
||||
|
||||
df, meta = climate._load_history(cell)
|
||||
assert meta["source"] == "era5-lake"
|
||||
assert df.height == full.height
|
||||
|
||||
|
||||
def test_lake_miss_falls_through_to_nasa(monkeypatch, tmp_path):
|
||||
"""An unconfigured lake (or a point outside it) costs nothing: NASA serves
|
||||
exactly as before the lake existed."""
|
||||
monkeypatch.setattr(climate, "CACHE_DIR", str(tmp_path))
|
||||
monkeypatch.setattr(climate, "_archive_cooldown_until", 0.0)
|
||||
cell = {"id": "lake_miss", "center_lat": 47.6, "center_lon": -122.3}
|
||||
|
||||
def _no_lake(c, start=None):
|
||||
raise climate.era5lake.LakeUnavailable("no lake configured")
|
||||
monkeypatch.setattr(climate.era5lake, "fetch_history_lake", _no_lake)
|
||||
full = _full_history_frame()
|
||||
monkeypatch.setattr(climate, "_fetch_history_nasa", lambda c: full)
|
||||
|
||||
df, meta = climate._load_history(cell)
|
||||
assert meta["source"] == "nasa-power"
|
||||
assert df.height == full.height
|
||||
|
||||
|
||||
def test_nasa_is_primary_history_source(monkeypatch, tmp_path):
|
||||
"""NASA POWER is the primary history source: a full NASA frame is accepted and
|
||||
Open-Meteo is never consulted."""
|
||||
|
|
|
|||
90
backend/tests/data/test_era5lake.py
Normal file
90
backend/tests/data/test_era5lake.py
Normal file
|
|
@ -0,0 +1,90 @@
|
|||
"""The ERA5 lake: grid math, layout keys, and the read/slice client (no network
|
||||
— the local-dir source stands in for the bucket)."""
|
||||
import datetime
|
||||
|
||||
import polars as pl
|
||||
import pytest
|
||||
|
||||
from data import era5lake
|
||||
|
||||
|
||||
def _daily_frame(start="1940-01-01", days=32000):
|
||||
d0 = datetime.date.fromisoformat(start)
|
||||
dates = [d0 + datetime.timedelta(days=k) for k in range(days)]
|
||||
n = len(dates)
|
||||
return 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,
|
||||
})
|
||||
|
||||
|
||||
def test_grid_roundtrip_and_bounds():
|
||||
# Nearest-index snap, poles clamped, longitude wrapped to [0, 360).
|
||||
assert era5lake.to_idx(90.0, 0.0) == (0, 0)
|
||||
assert era5lake.to_idx(-90.0, 0.0) == (720, 0)
|
||||
assert era5lake.to_idx(51.5074, -0.1278) == (154, 1439)
|
||||
i, j = era5lake.to_idx(35.6762, 139.6503)
|
||||
lat, lon = era5lake.to_coords(i, j)
|
||||
assert abs(lat - 35.6762) <= era5lake.STEP / 2 + 1e-9
|
||||
assert abs(lon - 139.6503) <= era5lake.STEP / 2 + 1e-9
|
||||
|
||||
|
||||
def test_tile_and_keys_align_with_hive_layout():
|
||||
assert era5lake.tile_of(154, 1439) == (12, 119)
|
||||
assert era5lake.point_key(154, 1439) == "era5/points/lat=154/lon=1439.parquet"
|
||||
assert era5lake.daily_part_key(12, 119, 1994, 3) == \
|
||||
"era5/daily/tile=12_119/year=1994/month=03/part.parquet"
|
||||
|
||||
|
||||
def test_s3_config_requires_both_keys(monkeypatch):
|
||||
monkeypatch.delenv("THERMOGRAPH_LAKE_S3_ACCESS_KEY", raising=False)
|
||||
monkeypatch.delenv("THERMOGRAPH_LAKE_S3_SECRET_KEY", raising=False)
|
||||
assert era5lake.s3_config() is None
|
||||
monkeypatch.setenv("THERMOGRAPH_LAKE_S3_ACCESS_KEY", "k")
|
||||
assert era5lake.s3_config() is None
|
||||
monkeypatch.setenv("THERMOGRAPH_LAKE_S3_SECRET_KEY", "s")
|
||||
cfg = era5lake.s3_config()
|
||||
assert cfg["bucket"] == "era5-thermograph"
|
||||
opts = era5lake.storage_options(cfg)
|
||||
assert opts["aws_virtual_hosted_style_request"] == "false"
|
||||
|
||||
|
||||
def _write_point(tmp_path, cell, df):
|
||||
i, j = era5lake.to_idx(cell["center_lat"], cell["center_lon"])
|
||||
path = tmp_path / era5lake.point_key(i, j)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
df.write_parquet(path)
|
||||
|
||||
|
||||
CELL = {"center_lat": 51.5, "center_lon": -0.125}
|
||||
|
||||
|
||||
def test_fetch_slices_to_start_and_keeps_schema(monkeypatch, tmp_path):
|
||||
monkeypatch.setenv("THERMOGRAPH_LAKE_LOCAL_DIR", str(tmp_path))
|
||||
_write_point(tmp_path, CELL, _daily_frame())
|
||||
df = era5lake.fetch_history_lake(CELL, start="1980-01-01")
|
||||
assert df["date"].min() == datetime.date(1980, 1, 1)
|
||||
assert set(df.columns) == {"date", "tmax", "tmin", "precip", "wind",
|
||||
"gust", "humid", "fmax", "fmin"}
|
||||
# Without a start, the whole 1940+ record comes back.
|
||||
assert era5lake.fetch_history_lake(CELL).height == 32000
|
||||
|
||||
|
||||
def test_fetch_rejects_short_and_missing(monkeypatch, tmp_path):
|
||||
monkeypatch.setenv("THERMOGRAPH_LAKE_LOCAL_DIR", str(tmp_path))
|
||||
with pytest.raises(era5lake.LakeUnavailable):
|
||||
era5lake.fetch_history_lake(CELL) # nothing written
|
||||
_write_point(tmp_path, CELL, _daily_frame(days=500))
|
||||
with pytest.raises(era5lake.LakeUnavailable): # implausibly short record
|
||||
era5lake.fetch_history_lake(CELL)
|
||||
|
||||
|
||||
def test_fetch_unconfigured_raises(monkeypatch):
|
||||
monkeypatch.delenv("THERMOGRAPH_LAKE_LOCAL_DIR", raising=False)
|
||||
monkeypatch.delenv("THERMOGRAPH_LAKE_URL", raising=False)
|
||||
monkeypatch.delenv("THERMOGRAPH_LAKE_S3_ACCESS_KEY", raising=False)
|
||||
monkeypatch.delenv("THERMOGRAPH_LAKE_S3_SECRET_KEY", raising=False)
|
||||
with pytest.raises(era5lake.LakeUnavailable):
|
||||
era5lake.fetch_history_lake(CELL)
|
||||
114
backend/tests/web/test_lake_app.py
Normal file
114
backend/tests/web/test_lake_app.py
Normal file
|
|
@ -0,0 +1,114 @@
|
|||
"""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
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def lake(tmp_path, monkeypatch):
|
||||
"""A tiny two-point lake (points, hive partitions, manifest) on disk."""
|
||||
from data import era5lake
|
||||
|
||||
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", 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", 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", json={"sql": sql}).status_code == 400
|
||||
4
infra/.gitignore
vendored
4
infra/.gitignore
vendored
|
|
@ -15,6 +15,10 @@ age.key
|
|||
|
||||
.DS_Store
|
||||
|
||||
# lake-iceberg's pytest suite runs in place.
|
||||
__pycache__/
|
||||
*.pyc
|
||||
|
||||
# Host-side deploy state (deploy.sh): the live per-service image tags and the
|
||||
# cross-repo deploy lock. Untracked on purpose -- they must survive the
|
||||
# `git reset --hard` at the top of every deploy (deploy.sh's comments already
|
||||
|
|
|
|||
|
|
@ -187,18 +187,19 @@ export FRONTEND_IMAGE_TAG="${FRONTEND_IMAGE_TAG:-local}"
|
|||
|
||||
# Which compose services this run pulls/rolls.
|
||||
#
|
||||
# `daemon` rides with `backend` and is never a target on its own. It runs the
|
||||
# SAME image at the SAME tag (it is a second binary inside the backend image),
|
||||
# and it talks to the web process over the /internal/* contract -- so rolling one
|
||||
# without the other is exactly the version skew that contract has no negotiation
|
||||
# for. Pairing them here is also what makes the service exist at all: a
|
||||
# single-service deploy runs `up -d --no-deps <targets>`, so a `daemon` listed
|
||||
# `daemon` and `lake` ride with `backend` and are never targets on their own.
|
||||
# Both run the SAME image at the SAME tag (a second binary and a second role
|
||||
# inside the backend image): the daemon talks to the web process over the
|
||||
# /internal/* contract, the lake serves it history — rolling one without the
|
||||
# other is exactly the version skew those contracts have no negotiation for.
|
||||
# Pairing them here is also what makes the services exist at all: a
|
||||
# single-service deploy runs `up -d --no-deps <targets>`, so a service listed
|
||||
# only in compose but absent from TARGETS would never be created on a
|
||||
# backend-only deploy, and the gateway bot would silently never start.
|
||||
# backend-only deploy.
|
||||
case "$SERVICE" in
|
||||
backend) TARGETS=(backend daemon) ;;
|
||||
backend) TARGETS=(backend lake daemon) ;;
|
||||
frontend) TARGETS=(frontend) ;;
|
||||
all) TARGETS=(backend frontend daemon) ;;
|
||||
all) TARGETS=(backend lake frontend daemon) ;;
|
||||
esac
|
||||
|
||||
# Login only when a token is supplied. The SSH/CI deploy paths (deploy.yml,
|
||||
|
|
|
|||
|
|
@ -15,7 +15,7 @@ no per-host duplication.
|
|||
| File | Holds | Encrypted? |
|
||||
|------|-------|------------|
|
||||
| `common.yaml` | Secrets shared by every host — Discord tokens/secret, VAPID keypair, `THERMOGRAPH_AUTH_SECRET`, metrics token, indexnow key, SMTP creds | ✅ |
|
||||
| `prod.yaml` | Prod-only overrides — `POSTGRES_PASSWORD`, `REGISTRY_TOKEN` | ✅ |
|
||||
| `prod.yaml` | Prod-only overrides — `POSTGRES_PASSWORD`, `REGISTRY_TOKEN`, ERA5 lake bucket creds (`THERMOGRAPH_LAKE_S3_ACCESS_KEY` / `THERMOGRAPH_LAKE_S3_SECRET_KEY`, consumed by the `lake` Swarm service and `gen_era5_lake.py`) | ✅ |
|
||||
| `beta.yaml` | Beta-only overrides | ✅ |
|
||||
| `example.yaml` | Format reference / CI fixture (fake values) | ✅ |
|
||||
| `../../.sops.yaml` | Which age recipient files are encrypted to (plaintext config) | — |
|
||||
|
|
|
|||
|
|
@ -16,6 +16,8 @@ THERMOGRAPH_VAPID_PRIVATE_KEY: ENC[AES256_GCM,data:sTtkAbMVUtwvRssCZJa2PWUjlfnYS
|
|||
THERMOGRAPH_VAPID_PUBLIC_KEY: ENC[AES256_GCM,data:SnFw+pwNgZlxMEhMUH/1daMZlRB2QFNJHtL+4C3S1eFcJJdT2uNGkY7Ym9FPfx1V7cwt/3urHg1oVgdDzYaIal2wye4l5HhKBP5PjtQkDW+Pck24kxul,iv:Ku4MhsrAfyA2Y3CqhGIsYKrjhKQ4cqAYjCqjFeEvBWE=,tag:hw42iS+79D7gNf+BPTYM0g==,type:str]
|
||||
THERMOGRAPH_VAPID_CONTACT: ENC[AES256_GCM,data:ox+hHiFMTPZqYwICDSwPwcT13ZDSKiGcAVHfVVMNYQ==,iv:DssV/olPXERpdmXNMmzjVYvhn/0zmuSp2XlTs42pIi0=,tag:4YHXI1BFbGZ6JsdXqJuf3w==,type:str]
|
||||
REGISTRY_TOKEN: ENC[AES256_GCM,data:pC1Zzs56la8qjzWcbkuKyeVHwS5WfIe+mrLm5Sdi/LKYziAvRU5cOA==,iv:57PlJEBdgwNJsfFdNOGqAnKk8l6zpmMSWYFEKS6pn+4=,tag:pPftNHUVU6+CY2nWMaXO6w==,type:str]
|
||||
THERMOGRAPH_LAKE_S3_ACCESS_KEY: ENC[AES256_GCM,data:ZEQ278P7azMKlYMjUraP6hIqnjymqL90zEbvTNC9iXU=,iv:5nXXZwx7RfqwCMqx5LWtxQy2nQcq+6Yr79ulBi6BOnE=,tag:jXcCCCcz4NsDoIZwa/ji4A==,type:str]
|
||||
THERMOGRAPH_LAKE_S3_SECRET_KEY: ENC[AES256_GCM,data:p5qCrHWoH1i3smDZZBIO8zMSWWxDARYsTY12AA0E8ao=,iv:PVOEQ0kBFeYaE/FT2UswtSTY7S84U/YOPSbcsXR51f4=,tag:JFwaOwAYInPs199li66U/g==,type:str]
|
||||
sops:
|
||||
age:
|
||||
- enc: |
|
||||
|
|
@ -27,7 +29,7 @@ sops:
|
|||
0+PQdtDxom72BqrAoj5lhepxW8YWPqQIRXsVS/XYQEQfPJuuuiRMKw==
|
||||
-----END AGE ENCRYPTED FILE-----
|
||||
recipient: age1xx4dzs0dxlwvkv9sjuqzsphl7lfrxannkfken374yu2qvvcte9sqzktqt2
|
||||
lastmodified: "2026-07-22T03:17:28Z"
|
||||
mac: ENC[AES256_GCM,data:npF+ngF4Khj3NJ2PNFcQjQYIrZVpdPM0drk43pAI6HfFaynX8JOqRFZRsH0sCfMoPkXB/ffCJxla1Tt3Mu3zgKTNSWcn+mjzl+dYvlruY1VylHLS+GnTO9zKeQcwbrYzTfZNWzQhm6/mw8KCL3791U/xNHwJc6g5VxRUWl6TBdo=,iv:JLzY8rm33NLG8CqyAlU6vBhlODUiRipKu9daoRvgyzg=,tag:yIsxVZAnw8FD7Bt/x55wMQ==,type:str]
|
||||
lastmodified: "2026-07-23T22:30:10Z"
|
||||
mac: ENC[AES256_GCM,data:JWfkQk4Z1ohPha9LID/hvZOycmDy0sUyjiqDdCidfCQEsPx+OwhPatcM31x1ZKFgh3TyYTxrauoSTlIeuwfy7+R9jATDdBBZe9zMqgs6Szu9WUdFiaZzIFpG7SEVw+8IqIyAuqZVSW3k5yHmchcmkNyDekV+1RIF8wxUrOeSCVQ=,iv:6PXUiCnYBiFyNPnCRAs6f7a+RuvWI4jSIMXwuMn7WjA=,tag:y+Ix2m65hQRkSt8FylfKSg==,type:str]
|
||||
unencrypted_suffix: _unencrypted
|
||||
version: 3.13.2
|
||||
|
|
|
|||
|
|
@ -24,6 +24,8 @@ THERMOGRAPH_DISCORD_BOT_TOKEN: ENC[AES256_GCM,data:zfVJLH8ddct2Jvstl/olZLQ9SM15c
|
|||
REGISTRY_TOKEN: ENC[AES256_GCM,data:3VETZR026YXvWA7DLly4/mz+hqEK7vkEMwKlDKU6CQZeZiFpbQXWig==,iv:ixZJi2VWAZak6dFPvVyErVRanp9bzqXyV3MVa4QsBuU=,tag:Pc2Mdw58S0tSdkgJ5yuVNw==,type:str]
|
||||
THERMOGRAPH_DISCORD_BOT: ENC[AES256_GCM,data:Kg==,iv:CPYN8aCu79PHaz3Zj6OgqyJhXdHLt0tLodgML2NuwQM=,tag:1B3IhlfjCwSlPq7FnKGeRg==,type:str]
|
||||
THERMOGRAPH_DISCORD_WEATHER_CHANNEL: ENC[AES256_GCM,data:WIxYcoOnBDdDcNj8uJ5eUqXjVQ==,iv:l6kU5PGVcW2Nv7EqBUGCRxBzaJOJ3AW1nEEH33aAbmI=,tag:/KgdCKmrlQioIWdz21PcTg==,type:str]
|
||||
THERMOGRAPH_LAKE_S3_ACCESS_KEY: ENC[AES256_GCM,data:goKTqruUdKUpSSrhu1Ivbh+yS7/DhGdCBeW/xb7AIkI=,iv:TmuHDgUi7dJZYv6uuUgqpnl5kG0tFNLgn11yizhHAV8=,tag:kkH+qBWA5rJwySIILk1cPg==,type:str]
|
||||
THERMOGRAPH_LAKE_S3_SECRET_KEY: ENC[AES256_GCM,data:7q07fpPgjgJN1NivbhAHn5JRyhI/+f0nV7Oll9bxW3w=,iv:hKZNH4PLmyyhA/crJpy1RA6peL9HzosyiNi+Y4yMKwo=,tag:sb5ailD6K44d4m8TmzZD9g==,type:str]
|
||||
sops:
|
||||
age:
|
||||
- enc: |
|
||||
|
|
@ -35,7 +37,7 @@ sops:
|
|||
PDXCB8EqKwdC3BnKlpD9ptYcuPHVUSgve1wIZYJ7iHIEhW9Kc6RXOw==
|
||||
-----END AGE ENCRYPTED FILE-----
|
||||
recipient: age1xx4dzs0dxlwvkv9sjuqzsphl7lfrxannkfken374yu2qvvcte9sqzktqt2
|
||||
lastmodified: "2026-07-23T04:33:06Z"
|
||||
mac: ENC[AES256_GCM,data:BjNMEudsS9ZViR55OjSXu5qrOJ9yu2kCGx61qu4vE1ZZJIZtnnyZ3CfpsPygvLLhiFRj76xcqFnuSqTdr17BCEHz3UzsHFv0vBi17lZ5ksN5Oz+9wr+3Pg3TvxC0VF2VHJVP84VH61CmB5+G4WOms8WqLtmIGSLYcTVtcIFoUhk=,iv:jTEB8Mxg8qIyNfhpSgJtUvULVx4+XFrPLyQKhNcbkTg=,tag:SIyvNltdGCMfDGvi4kD3MA==,type:str]
|
||||
lastmodified: "2026-07-23T21:26:01Z"
|
||||
mac: ENC[AES256_GCM,data:RHLp+0waxaKQXfAzipchk/vZnE3NtMOHAJwqddTlJYpxl9Sl47IqZcdV05j8kveySNoedZbpmu6+unso9HPaaJ9lmunZVmOPc9wqn+5JJNw8kkPIosrXL8UzKW9gO4OIbTNkEC4xSQP4qyQeUO/RaDP2aH0d+GCa7sKTePrd6Gg=,iv:9027dfABH2VuEJKWzY1+eTns+IdCUXW2KQTAdlj4iP4=,tag:JLLkj0g4gLMseLdUy896/g==,type:str]
|
||||
unencrypted_suffix: _unencrypted
|
||||
version: 3.13.2
|
||||
|
|
|
|||
|
|
@ -25,7 +25,9 @@
|
|||
set -euo pipefail
|
||||
|
||||
STACK_NAME="${STACK_NAME:-thermograph}"
|
||||
SERVICE="${STACK_NAME}_web"
|
||||
# Which stack service to scale (the stack runs one autoscaler instance per
|
||||
# scaled service — web and lake today; same loop, different bounds).
|
||||
SERVICE="${STACK_NAME}_${TARGET_SERVICE:-web}"
|
||||
MIN="${MIN_REPLICAS:-1}"
|
||||
MAX="${MAX_REPLICAS:-3}"
|
||||
UP_AT="${SCALE_UP_CPU:-220}"
|
||||
|
|
|
|||
|
|
@ -163,9 +163,16 @@ if [ "$SERVICE" = "all" ] || ! docker service inspect "${STACK_NAME}_web" >/dev/
|
|||
else
|
||||
case "$SERVICE" in
|
||||
backend)
|
||||
echo "==> Rolling web + worker to $BACKEND_IMAGE"
|
||||
echo "==> Rolling web + worker + lake to $BACKEND_IMAGE"
|
||||
docker service update --with-registry-auth --detach=false --image "$BACKEND_IMAGE" "${STACK_NAME}_web"
|
||||
docker service update --with-registry-auth --detach=false --image "$BACKEND_IMAGE" "${STACK_NAME}_worker"
|
||||
# lake ships in the same image; a stack file predating it has no service
|
||||
# yet — the next SERVICE=all stack deploy creates it, so don't fail here.
|
||||
if docker service inspect "${STACK_NAME}_lake" >/dev/null 2>&1; then
|
||||
docker service update --with-registry-auth --detach=false --image "$BACKEND_IMAGE" "${STACK_NAME}_lake"
|
||||
else
|
||||
echo " (no ${STACK_NAME}_lake service yet; created on the next full stack deploy)"
|
||||
fi
|
||||
;;
|
||||
frontend)
|
||||
echo "==> Rolling frontend to $FRONTEND_IMAGE"
|
||||
|
|
|
|||
|
|
@ -82,6 +82,10 @@ services:
|
|||
# not the compose bridge's 172.19.0.1 (an overlay has no host gateway).
|
||||
# provision-mail.sh's DOCKER_MAIL_GATEWAY/SUBNET cover this listener.
|
||||
THERMOGRAPH_SMTP_HOST: ${STACK_SMTP_HOST:-172.18.0.1}
|
||||
# History reads ask the lake service (Swarm VIP over 1..2 replicas)
|
||||
# before any third-party source; unset on beta/LAN, where compose has no
|
||||
# lake service and climate.py falls straight through to NASA.
|
||||
THERMOGRAPH_LAKE_URL: http://lake:8141
|
||||
volumes:
|
||||
- appdata:/state
|
||||
- applogs:/app/logs
|
||||
|
|
@ -121,6 +125,7 @@ services:
|
|||
THERMOGRAPH_SINGLETON_PG: "1"
|
||||
RUN_MIGRATIONS: "0"
|
||||
THERMOGRAPH_SMTP_HOST: ${STACK_SMTP_HOST:-172.18.0.1}
|
||||
THERMOGRAPH_LAKE_URL: http://lake:8141
|
||||
volumes:
|
||||
- appdata:/state
|
||||
- applogs:/app/logs
|
||||
|
|
@ -138,6 +143,44 @@ services:
|
|||
restart_policy:
|
||||
condition: on-failure
|
||||
|
||||
# The ERA5 lake indexer: SQL-style, partition-pruned reads over the
|
||||
# era5-thermograph bucket with a node-local parquet cache, so a cell's
|
||||
# history is one LAN hop for web instead of an S3 fetch (or a NASA call).
|
||||
# Same backend image; THERMOGRAPH_ROLE=lake makes the entrypoint start
|
||||
# lake_app (no database, no migrations). Its S3 credentials
|
||||
# (THERMOGRAPH_LAKE_S3_ACCESS_KEY/_SECRET_KEY) arrive via the same rendered
|
||||
# /host/thermograph.env as every other secret. Without them it stays
|
||||
# healthy and answers 503/404, and web falls through to NASA — the lake is
|
||||
# an accelerator, never a point of failure.
|
||||
lake:
|
||||
image: ${REGISTRY_HOST:-git.thermograph.org}/${BACKEND_IMAGE_PATH:-emi/thermograph/backend}:${BACKEND_IMAGE_TAG:?required}
|
||||
entrypoint: ["/host/env-entrypoint.sh"]
|
||||
environment:
|
||||
THERMOGRAPH_ROLE: lake
|
||||
PORT: 8141
|
||||
THERMOGRAPH_SERVICE_ROLE: backend
|
||||
WORKERS: "1"
|
||||
THERMOGRAPH_LAKE_CACHE: /state/lake-cache
|
||||
volumes:
|
||||
# Replicas share the cache; lake_app's tmp+rename writes make that safe.
|
||||
- lakecache:/state
|
||||
- /opt/thermograph/infra/deploy/stack/env-entrypoint.sh:/host/env-entrypoint.sh:ro
|
||||
- /etc/thermograph/stack.env:/host/thermograph.env:ro
|
||||
networks:
|
||||
- internal
|
||||
deploy:
|
||||
replicas: ${LAKE_MIN_REPLICAS:-1}
|
||||
placement:
|
||||
constraints: ["node.role == manager"]
|
||||
resources:
|
||||
limits:
|
||||
cpus: "${LAKE_CPUS:-2}"
|
||||
restart_policy:
|
||||
condition: on-failure
|
||||
update_config:
|
||||
order: start-first
|
||||
failure_action: rollback
|
||||
|
||||
daemon:
|
||||
# SAME image and tag as web/worker on purpose: the daemon binary
|
||||
# (/usr/local/bin/thermograph-daemon, built into the backend image) and the
|
||||
|
|
@ -254,6 +297,39 @@ services:
|
|||
restart_policy:
|
||||
condition: any
|
||||
|
||||
# Second autoscaler instance, watching `lake` (1..2 replicas). The lake is
|
||||
# read-only and cache-heavy, so it saturates on decode CPU well below web's
|
||||
# thresholds — scale up sooner, still down slowly.
|
||||
autoscaler-lake:
|
||||
image: docker:27-cli
|
||||
entrypoint: ["/bin/sh", "/host/autoscale.sh"]
|
||||
environment:
|
||||
STACK_NAME: ${STACK_NAME:-thermograph}
|
||||
TARGET_SERVICE: lake
|
||||
MIN_REPLICAS: ${LAKE_MIN_REPLICAS:-1}
|
||||
MAX_REPLICAS: ${LAKE_MAX_REPLICAS:-2}
|
||||
SCALE_UP_CPU: ${LAKE_SCALE_UP_CPU:-120}
|
||||
SCALE_DOWN_CPU: ${LAKE_SCALE_DOWN_CPU:-30}
|
||||
POLL_SECONDS: ${POLL_SECONDS:-15}
|
||||
UP_SAMPLES: ${UP_SAMPLES:-3}
|
||||
DOWN_SAMPLES: ${DOWN_SAMPLES:-20}
|
||||
COOLDOWN_SECONDS: ${COOLDOWN_SECONDS:-180}
|
||||
volumes:
|
||||
- /var/run/docker.sock:/var/run/docker.sock
|
||||
- /opt/thermograph/infra/deploy/stack/autoscale.sh:/host/autoscale.sh:ro
|
||||
networks:
|
||||
- internal
|
||||
deploy:
|
||||
replicas: 1
|
||||
placement:
|
||||
constraints: ["node.role == manager"]
|
||||
resources:
|
||||
limits:
|
||||
cpus: "0.2"
|
||||
memory: 64m
|
||||
restart_policy:
|
||||
condition: any
|
||||
|
||||
networks:
|
||||
internal:
|
||||
driver: overlay
|
||||
|
|
@ -275,3 +351,6 @@ volumes:
|
|||
applogs:
|
||||
external: true
|
||||
name: ${APPLOGS_VOLUME:-thermograph_applogs}
|
||||
# Lake parquet cache: node-local and disposable (a cold cache just refills
|
||||
# from the bucket), so NOT external — the stack owns its lifecycle.
|
||||
lakecache: {}
|
||||
|
|
|
|||
|
|
@ -45,6 +45,9 @@ services:
|
|||
db:
|
||||
cpus: !reset null
|
||||
deploy: !reset null
|
||||
lake:
|
||||
cpus: !reset null
|
||||
deploy: !reset null
|
||||
# Uncapped memory on dev too (prod ceilings it at 8g). The Postgres/DuckDB
|
||||
# memory *budget* is still the ~8 GB derived by deploy/db/init/20-tuning.sh from
|
||||
# the default DB_MEMORY (raise DB_MEMORY to give dev more); shm_size stays (it's
|
||||
|
|
|
|||
|
|
@ -124,6 +124,11 @@ services:
|
|||
# (NOT /app/data) so it never shadows the Python `data/` package -- see the
|
||||
# volumes: note below and thermograph-backend paths.py.
|
||||
THERMOGRAPH_SINGLETON_LOCK: /state/notifier.lock
|
||||
# History reads ask the lake service before any third-party source; if
|
||||
# the lake has no bucket creds (LAN without them exported) it answers
|
||||
# 503 and climate.py falls straight through to NASA — an accelerator,
|
||||
# never a dependency.
|
||||
THERMOGRAPH_LAKE_URL: http://lake:8141
|
||||
# Shared secret for the /internal/* routes the daemon service calls.
|
||||
# Interpolated (like POSTGRES_PASSWORD) so local runs get it from the
|
||||
# repo-root .env; prod/beta have it in /etc/thermograph.env, which
|
||||
|
|
@ -164,6 +169,36 @@ services:
|
|||
- "127.0.0.1:8137:8137"
|
||||
restart: unless-stopped
|
||||
|
||||
# The ERA5 lake service: partition-pruned reads over the era5-thermograph
|
||||
# bucket with a local parquet cache, same backend image with
|
||||
# THERMOGRAPH_ROLE=lake (the entrypoint starts lake_app — no database, no
|
||||
# migrations, so no depends_on: db). Bucket creds come from
|
||||
# /etc/thermograph.env where the vault renders them (beta/prod) or from the
|
||||
# deploy environment (LAN — deploy-dev.yml injects the Actions S3 secrets);
|
||||
# without them the service stays healthy and every read falls through to
|
||||
# NASA. Never published on a host port: only the backend talks to it.
|
||||
lake:
|
||||
image: ${REGISTRY_HOST:-git.thermograph.org}/${BACKEND_IMAGE_PATH:-emi/thermograph/backend}:${BACKEND_IMAGE_TAG:-local}
|
||||
environment:
|
||||
THERMOGRAPH_ROLE: lake
|
||||
PORT: 8141
|
||||
THERMOGRAPH_SERVICE_ROLE: backend
|
||||
WORKERS: "1"
|
||||
THERMOGRAPH_LAKE_CACHE: /state/lake-cache
|
||||
THERMOGRAPH_LAKE_S3_ACCESS_KEY: ${THERMOGRAPH_LAKE_S3_ACCESS_KEY:-}
|
||||
THERMOGRAPH_LAKE_S3_SECRET_KEY: ${THERMOGRAPH_LAKE_S3_SECRET_KEY:-}
|
||||
env_file:
|
||||
- path: /etc/thermograph.env
|
||||
required: false
|
||||
volumes:
|
||||
- lakecache:/state
|
||||
cpus: ${LAKE_CPUS:-2}
|
||||
deploy:
|
||||
resources:
|
||||
limits:
|
||||
cpus: "${LAKE_CPUS:-2}"
|
||||
restart: unless-stopped
|
||||
|
||||
daemon:
|
||||
# The SAME image and tag as backend, on purpose: the daemon (Go, built into
|
||||
# the backend image as /usr/local/bin/thermograph-daemon) and the backend
|
||||
|
|
@ -250,6 +285,7 @@ volumes:
|
|||
pgdata: {}
|
||||
appdata: {}
|
||||
applogs: {}
|
||||
lakecache: {}
|
||||
|
||||
networks:
|
||||
# Pin the default network's subnet + gateway so the host Postfix can rely on a
|
||||
|
|
|
|||
20
infra/lake-iceberg/Dockerfile
Normal file
20
infra/lake-iceberg/Dockerfile
Normal file
|
|
@ -0,0 +1,20 @@
|
|||
# One-shot Iceberg sync for the ERA5 lake: registers new hive tiles from the
|
||||
# era5-thermograph bucket into the iceberg/era5_daily table (see
|
||||
# sync_iceberg.py). Everything ships as manylinux wheels — no toolchain.
|
||||
FROM python:3.12-slim
|
||||
|
||||
# The pyiceberg extras pull sqlalchemy (sqlite catalog), pyarrow (FileIO +
|
||||
# parquet footers) and s3fs (bucket listing/manifest reads).
|
||||
COPY requirements.txt /tmp/requirements.txt
|
||||
RUN pip install --no-cache-dir -r /tmp/requirements.txt
|
||||
|
||||
WORKDIR /app
|
||||
COPY sync_iceberg.py .
|
||||
|
||||
# The sqlite catalog lives on a volume so runs share it; losing it is fine
|
||||
# (the next run re-registers the table from its latest metadata JSON).
|
||||
ENV THERMOGRAPH_ICEBERG_CATALOG_DB=/state/iceberg-catalog.db \
|
||||
PYTHONUNBUFFERED=1
|
||||
VOLUME /state
|
||||
|
||||
ENTRYPOINT ["python", "sync_iceberg.py"]
|
||||
123
infra/lake-iceberg/README.md
Normal file
123
infra/lake-iceberg/README.md
Normal file
|
|
@ -0,0 +1,123 @@
|
|||
# lake-iceberg — Iceberg conversion for the ERA5 lake
|
||||
|
||||
Maintains an Apache Iceberg v2 table `era5_daily` at
|
||||
`s3://era5-thermograph/iceberg/era5_daily` over the lake's existing hive
|
||||
parquet (`era5/daily/tile=<ti>_<tj>/year=<Y>/month=<M>/part.parquet`).
|
||||
|
||||
Key properties:
|
||||
|
||||
- **No data rewrite.** `sync_iceberg.py` uses pyiceberg `add_files`: the
|
||||
Iceberg metadata references the existing part files in place. Storage is
|
||||
not doubled and nothing under `era5/` is ever written, moved, or deleted —
|
||||
the tool only writes under the `iceberg/` prefix.
|
||||
- **Incremental + idempotent.** Each run diffs the tiles in
|
||||
`era5/manifest.parquet` (the source of truth for *complete* tiles — the
|
||||
extractor uploads all of a tile's objects before its manifest rows appear)
|
||||
against the `thermograph.synced-tiles` table property, and registers only
|
||||
the new ones. Files and the property update land in the same commit, so an
|
||||
interrupted run never half-syncs a tile; re-running is always safe, even
|
||||
while an extraction job is writing new tiles.
|
||||
- **Self-contained catalog.** A pyiceberg SqlCatalog on local sqlite
|
||||
(`THERMOGRAPH_ICEBERG_CATALOG_DB`, default `/state/iceberg-catalog.db`).
|
||||
The catalog is only a pointer: if the db is lost, the next run finds the
|
||||
newest `*.metadata.json` under the table location and re-registers it —
|
||||
the synced-tiles set lives in the table metadata itself, so nothing is
|
||||
re-added.
|
||||
- **Partitioning `(tile, year, month)`** via order-preserving transforms of
|
||||
real columns: `truncate[12](lat_idx)` / `truncate[12](lon_idx)` (partition
|
||||
values are `ti*12` / `tj*12` — the part files carry no tile column, the 12
|
||||
matches the lake's 12x12-point tile grid) plus `year(date)` and
|
||||
`month(date)`. Predicates on `lat_idx`/`lon_idx`/`date` prune partitions
|
||||
directly — no synthetic tile/year/month columns needed in queries.
|
||||
|
||||
## Build
|
||||
|
||||
```sh
|
||||
docker build -t thermograph-lake-iceberg infra/lake-iceberg/
|
||||
```
|
||||
|
||||
## Run (one-shot)
|
||||
|
||||
Config is the backend's lake variables (same names and defaults as
|
||||
`backend/data/era5lake.py`): `THERMOGRAPH_LAKE_S3_ACCESS_KEY`,
|
||||
`THERMOGRAPH_LAKE_S3_SECRET_KEY`, and optionally
|
||||
`THERMOGRAPH_LAKE_S3_ENDPOINT` / `_BUCKET` / `_REGION`.
|
||||
|
||||
```sh
|
||||
docker run --rm --env-file /path/to/lake.env \
|
||||
-v lake-iceberg-state:/state \
|
||||
thermograph-lake-iceberg
|
||||
```
|
||||
|
||||
Flags: `--dry-run` (list pending tiles, write nothing), `--limit N` (cap
|
||||
tiles per run), `--batch N` (tiles per Iceberg commit, default 5),
|
||||
`--catalog-db PATH`. `PYICEBERG_MAX_WORKERS` controls the parquet-footer
|
||||
read parallelism inside `add_files` (32 is a good value for the ~1000 small
|
||||
footers per tile).
|
||||
|
||||
Each run prints the row-count reconciliation (records added vs the
|
||||
manifest's per-tile row counts) and ends with the table's current metadata
|
||||
JSON path. It also refreshes `metadata/version-hint.text` (Hadoop-catalog
|
||||
convention) so readers can scan the bare table root without knowing that
|
||||
path. Transient S3 errors (Contabo intermittently returns SLOW_DOWN /
|
||||
connection timeouts under load) are retried per batch with backoff
|
||||
(`--retries`, default 3); a run that still dies resumes at the next batch
|
||||
boundary on re-run.
|
||||
|
||||
## Scheduling
|
||||
|
||||
Meant to run as a one-shot on prod after extraction batches finish — e.g. a
|
||||
cron or an ops-cron workflow step invoking the `docker run` above, with the
|
||||
secrets rendered from the SOPS vault like every other consumer of the
|
||||
`THERMOGRAPH_LAKE_S3_*` pair. Deliberately **not** wired into the Swarm
|
||||
stack or CI: the sync is an offline maintenance job, not a service, and each
|
||||
run costs one footer read per new part file (~1000 per tile), so it belongs
|
||||
after extraction batches, not on a tight loop.
|
||||
|
||||
## Querying from DuckDB
|
||||
|
||||
No catalog server needed — the version-hint pointer lets DuckDB scan the
|
||||
table root directly:
|
||||
|
||||
```sql
|
||||
INSTALL iceberg; LOAD iceberg;
|
||||
INSTALL httpfs; LOAD httpfs;
|
||||
CREATE SECRET lake (
|
||||
TYPE s3,
|
||||
KEY_ID '<THERMOGRAPH_LAKE_S3_ACCESS_KEY>',
|
||||
SECRET '<THERMOGRAPH_LAKE_S3_SECRET_KEY>',
|
||||
ENDPOINT 'eu2.contabostorage.com',
|
||||
REGION 'default',
|
||||
URL_STYLE 'path'
|
||||
);
|
||||
SELECT count(*)
|
||||
FROM iceberg_scan('s3://era5-thermograph/iceberg/era5_daily');
|
||||
```
|
||||
|
||||
To pin an exact table version, scan the metadata JSON path the sync run
|
||||
printed (`iceberg_scan('s3://…/metadata/<NNNNN-uuid>.metadata.json')`).
|
||||
|
||||
Partition pruning works on the raw columns, e.g. one tile-month:
|
||||
|
||||
```sql
|
||||
SELECT lat_idx, lon_idx, avg(tmax)
|
||||
FROM iceberg_scan('s3://era5-thermograph/iceberg/era5_daily')
|
||||
WHERE lat_idx BETWEEN 120 AND 131 AND lon_idx BETWEEN 120 AND 131
|
||||
AND date >= DATE '2020-07-01' AND date < DATE '2020-08-01'
|
||||
GROUP BY 1, 2;
|
||||
```
|
||||
|
||||
The part files carry no Iceberg field ids; the table sets
|
||||
`schema.name-mapping.default` so engines resolve columns by name (validated
|
||||
with pyiceberg 0.11.1 and DuckDB 1.5.5, values byte-identical to the hive
|
||||
side).
|
||||
|
||||
## Tests
|
||||
|
||||
```sh
|
||||
pip install -r infra/lake-iceberg/requirements.txt pytest
|
||||
pytest infra/lake-iceberg/test_sync.py
|
||||
```
|
||||
|
||||
Local-filesystem fixtures only (tiny parquet files, `file://` warehouse,
|
||||
sqlite catalog in tmp) — no network, no credentials.
|
||||
5
infra/lake-iceberg/requirements.txt
Normal file
5
infra/lake-iceberg/requirements.txt
Normal file
|
|
@ -0,0 +1,5 @@
|
|||
# Pinned to the versions the sync was validated with (add_files partition
|
||||
# inference from footer stats, sqlite catalog, path-style S3).
|
||||
pyiceberg[sql-sqlite,pyarrow,s3fs]==0.11.1
|
||||
pyarrow==25.0.0
|
||||
s3fs==2026.6.0
|
||||
451
infra/lake-iceberg/sync_iceberg.py
Normal file
451
infra/lake-iceberg/sync_iceberg.py
Normal file
|
|
@ -0,0 +1,451 @@
|
|||
"""Maintain an Apache Iceberg table over the ERA5 lake's hive parquet.
|
||||
|
||||
The lake extractor (backend/gen_era5_lake.py) writes plain hive-partitioned
|
||||
parquet to the era5-thermograph bucket:
|
||||
|
||||
era5/daily/tile=<ti>_<tj>/year=<Y>/month=<M>/part.parquet
|
||||
|
||||
This tool registers those files, in place, into an Iceberg v2 table
|
||||
``era5_daily`` located at ``s3://<bucket>/iceberg/era5_daily`` via pyiceberg
|
||||
``add_files`` — no data is rewritten or copied, the Iceberg metadata simply
|
||||
points at the existing objects. Everything this tool writes lives under the
|
||||
``iceberg/`` prefix; the ``era5/`` prefixes are read-only to it by design.
|
||||
|
||||
Partitioning: the part files carry no tile/year/month columns (those live in
|
||||
the object path), so the table partitions by order-preserving transforms of
|
||||
real columns instead — ``truncate[12](lat_idx)`` / ``truncate[12](lon_idx)``
|
||||
(a 12x12-point tile, so the partition value is ``ti*12``/``tj*12``) and
|
||||
``year(date)`` / ``month(date)``. Each part file holds exactly one tile and
|
||||
one calendar month, so every partition value is single-valued per file and
|
||||
``add_files`` can derive it from the parquet footer statistics alone.
|
||||
|
||||
Incremental and safe under concurrent extraction: ``era5/manifest.parquet``
|
||||
is the source of truth for complete tiles (the extractor uploads all of a
|
||||
tile's objects before adding it to the manifest), so each run diffs the
|
||||
manifest's tiles against the ``thermograph.synced-tiles`` table property and
|
||||
adds only the difference. The property is updated in the same commit as the
|
||||
files, so a run can die at any point and the next one resumes cleanly.
|
||||
|
||||
Catalog: a local pyiceberg SqlCatalog (sqlite, path from
|
||||
THERMOGRAPH_ICEBERG_CATALOG_DB, default /state/iceberg-catalog.db). The
|
||||
catalog is a convenience pointer only — if it is lost, the next run finds the
|
||||
latest ``*.metadata.json`` under the table location and re-registers it. Each
|
||||
run ends by printing the current metadata JSON path, which DuckDB can read
|
||||
directly: ``iceberg_scan('s3://.../metadata/NNNNN-....metadata.json')``.
|
||||
|
||||
S3 config comes from the same THERMOGRAPH_LAKE_S3_* variables (and defaults)
|
||||
as backend/data/era5lake.py; THERMOGRAPH_LAKE_LOCAL_DIR switches to a local
|
||||
mirror of the lake layout (tests, rehearsal).
|
||||
"""
|
||||
import argparse
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
import time
|
||||
|
||||
import pyarrow.parquet as pq
|
||||
from pyiceberg.catalog.sql import SqlCatalog
|
||||
from pyiceberg.exceptions import NoSuchTableError
|
||||
from pyiceberg.partitioning import PartitionField, PartitionSpec
|
||||
from pyiceberg.schema import Schema
|
||||
from pyiceberg.table.name_mapping import create_mapping_from_schema
|
||||
from pyiceberg.transforms import MonthTransform, TruncateTransform, YearTransform
|
||||
from pyiceberg.types import (
|
||||
DateType,
|
||||
DoubleType,
|
||||
FloatType,
|
||||
IntegerType,
|
||||
NestedField,
|
||||
)
|
||||
|
||||
PREFIX = "era5"
|
||||
MANIFEST_KEY = f"{PREFIX}/manifest.parquet"
|
||||
ICEBERG_PREFIX = "iceberg"
|
||||
|
||||
NAMESPACE = "lake"
|
||||
TABLE_NAME = "era5_daily"
|
||||
IDENTIFIER = f"{NAMESPACE}.{TABLE_NAME}"
|
||||
SYNCED_PROP = "thermograph.synced-tiles"
|
||||
|
||||
TILE = 12 # grid points per tile side, as in era5lake.TILE
|
||||
DEFAULT_BATCH = 5 # tiles per Iceberg commit
|
||||
DEFAULT_RETRIES = 3 # per-batch retries on transient S3 errors
|
||||
BACKOFF_SECONDS = 15 # first retry delay, doubled per attempt
|
||||
DEFAULT_CATALOG_DB = "/state/iceberg-catalog.db"
|
||||
|
||||
|
||||
class LakeInvariantError(RuntimeError):
|
||||
"""The lake violated an extractor invariant; never retried."""
|
||||
|
||||
# The part-file schema as the extractor writes it (verified against the
|
||||
# bucket): float32 for the heat-index-derived and raw met columns, float64
|
||||
# where polars promoted (fmin/sun/feels), int16 doy, int32 grid indices.
|
||||
# Iceberg has no 16-bit type, so doy maps to int; all fields optional.
|
||||
SCHEMA = Schema(
|
||||
NestedField(1, "date", DateType(), required=False),
|
||||
NestedField(2, "tmax", FloatType(), required=False),
|
||||
NestedField(3, "tmin", FloatType(), required=False),
|
||||
NestedField(4, "precip", FloatType(), required=False),
|
||||
NestedField(5, "wind", FloatType(), required=False),
|
||||
NestedField(6, "gust", FloatType(), required=False),
|
||||
NestedField(7, "humid", FloatType(), required=False),
|
||||
NestedField(8, "fmax", FloatType(), required=False),
|
||||
NestedField(9, "fmin", DoubleType(), required=False),
|
||||
NestedField(10, "sun", DoubleType(), required=False),
|
||||
NestedField(11, "feels", DoubleType(), required=False),
|
||||
NestedField(12, "doy", IntegerType(), required=False),
|
||||
NestedField(13, "lat_idx", IntegerType(), required=False),
|
||||
NestedField(14, "lon_idx", IntegerType(), required=False),
|
||||
)
|
||||
|
||||
# (tile, year, month): truncate[12] of the grid indices is the tile (values
|
||||
# are ti*12 / tj*12), year/month come from the date column. All four
|
||||
# transforms preserve order, which is what lets add_files infer the partition
|
||||
# value from footer min/max stats (single value per file by construction).
|
||||
PARTITION_SPEC = PartitionSpec(
|
||||
PartitionField(source_id=13, field_id=1000,
|
||||
transform=TruncateTransform(TILE), name="tile_lat"),
|
||||
PartitionField(source_id=14, field_id=1001,
|
||||
transform=TruncateTransform(TILE), name="tile_lon"),
|
||||
PartitionField(source_id=1, field_id=1002,
|
||||
transform=YearTransform(), name="year"),
|
||||
PartitionField(source_id=1, field_id=1003,
|
||||
transform=MonthTransform(), name="month"),
|
||||
)
|
||||
|
||||
|
||||
def s3_config() -> "dict | None":
|
||||
"""Bucket access, or None when unconfigured. Same variable names and
|
||||
defaults as backend/data/era5lake.py's s3_config (Contabo: path-style
|
||||
addressing, region is a signing formality)."""
|
||||
key = os.environ.get("THERMOGRAPH_LAKE_S3_ACCESS_KEY", "")
|
||||
secret = os.environ.get("THERMOGRAPH_LAKE_S3_SECRET_KEY", "")
|
||||
if not (key and secret):
|
||||
return None
|
||||
return {
|
||||
"endpoint": os.environ.get("THERMOGRAPH_LAKE_S3_ENDPOINT",
|
||||
"https://eu2.contabostorage.com"),
|
||||
"bucket": os.environ.get("THERMOGRAPH_LAKE_S3_BUCKET", "era5-thermograph"),
|
||||
"region": os.environ.get("THERMOGRAPH_LAKE_S3_REGION", "default"),
|
||||
"access_key": key,
|
||||
"secret_key": secret,
|
||||
}
|
||||
|
||||
|
||||
class HiveLake:
|
||||
"""Read-only view of the extracted lake (bucket or local mirror) plus the
|
||||
iceberg/ prefix the table lives under. ``root`` is an fsspec path prefix
|
||||
("<bucket>" or "/abs/dir"), ``uri`` the equivalent URI prefix pyiceberg
|
||||
sees ("s3://<bucket>" or "file:///abs/dir")."""
|
||||
|
||||
def __init__(self, fs, root: str, uri: str, cfg: "dict | None"):
|
||||
self.fs = fs
|
||||
self.root = root.rstrip("/")
|
||||
self.uri = uri.rstrip("/")
|
||||
self.cfg = cfg
|
||||
|
||||
@property
|
||||
def warehouse(self) -> str:
|
||||
return f"{self.uri}/{ICEBERG_PREFIX}"
|
||||
|
||||
def to_uri(self, path: str) -> str:
|
||||
rel = path.removeprefix(self.root).lstrip("/")
|
||||
return f"{self.uri}/{rel}"
|
||||
|
||||
def tile_rows(self) -> "dict[str, int]":
|
||||
"""tile -> row count, from the manifest (the source of truth for
|
||||
complete tiles: the extractor uploads every object of a tile before
|
||||
its manifest rows appear, so listing incomplete tiles is impossible
|
||||
here even while extraction is running)."""
|
||||
with self.fs.open(f"{self.root}/{MANIFEST_KEY}", "rb") as f:
|
||||
m = pq.read_table(f, columns=["tile", "rows"])
|
||||
out: dict[str, int] = {}
|
||||
for tile, rows in zip(m.column("tile").to_pylist(),
|
||||
m.column("rows").to_pylist()):
|
||||
out[tile] = out.get(tile, 0) + rows
|
||||
return out
|
||||
|
||||
def tile_files(self, tile: str) -> "list[str]":
|
||||
"""URIs of one complete tile's hive part files. A flat find() on the
|
||||
tile prefix, not a glob: fsspec expands year=*/month=* by walking
|
||||
every directory level (~1000 LIST calls per tile on S3), while find()
|
||||
is one paginated listing."""
|
||||
prefix = f"{self.root}/{PREFIX}/daily/tile={tile}"
|
||||
part = re.compile(r"/year=\d+/month=\d+/part\.parquet$")
|
||||
try:
|
||||
paths = self.fs.find(prefix)
|
||||
except FileNotFoundError:
|
||||
return []
|
||||
return sorted(self.to_uri(p) for p in paths if part.search(p))
|
||||
|
||||
def latest_metadata_uri(self) -> "str | None":
|
||||
"""Newest *.metadata.json under the table location, for re-registering
|
||||
the table into a fresh catalog db (pyiceberg names them with a
|
||||
zero-padded, monotonically increasing version prefix)."""
|
||||
prefix = f"{self.root}/{ICEBERG_PREFIX}/{TABLE_NAME}/metadata"
|
||||
try:
|
||||
paths = [p for p in self.fs.find(prefix)
|
||||
if p.endswith(".metadata.json")]
|
||||
except FileNotFoundError:
|
||||
return None
|
||||
if not paths:
|
||||
return None
|
||||
|
||||
def version(p: str) -> int:
|
||||
m = re.match(r"(\d+)", os.path.basename(p))
|
||||
return int(m.group(1)) if m else -1
|
||||
|
||||
return self.to_uri(max(paths, key=version))
|
||||
|
||||
|
||||
def lake_from_env() -> HiveLake:
|
||||
local = os.environ.get("THERMOGRAPH_LAKE_LOCAL_DIR", "")
|
||||
if local:
|
||||
from fsspec.implementations.local import LocalFileSystem
|
||||
root = os.path.abspath(local)
|
||||
return HiveLake(LocalFileSystem(), root, f"file://{root}", None)
|
||||
cfg = s3_config()
|
||||
if cfg is None:
|
||||
sys.exit("no lake configured: set THERMOGRAPH_LAKE_S3_ACCESS_KEY/"
|
||||
"_SECRET_KEY (or THERMOGRAPH_LAKE_LOCAL_DIR for a local mirror)")
|
||||
import s3fs
|
||||
fs = s3fs.S3FileSystem(
|
||||
key=cfg["access_key"], secret=cfg["secret_key"],
|
||||
endpoint_url=cfg["endpoint"],
|
||||
client_kwargs={"region_name": cfg["region"]},
|
||||
config_kwargs={"s3": {"addressing_style": "path"}})
|
||||
return HiveLake(fs, cfg["bucket"], f"s3://{cfg['bucket']}", cfg)
|
||||
|
||||
|
||||
def open_catalog(catalog_db: str, lake: HiveLake) -> SqlCatalog:
|
||||
parent = os.path.dirname(os.path.abspath(catalog_db))
|
||||
os.makedirs(parent, exist_ok=True)
|
||||
props = {"uri": f"sqlite:///{os.path.abspath(catalog_db)}",
|
||||
"warehouse": lake.warehouse}
|
||||
if lake.cfg:
|
||||
props |= {
|
||||
"py-io-impl": "pyiceberg.io.pyarrow.PyArrowFileIO",
|
||||
"s3.endpoint": lake.cfg["endpoint"],
|
||||
"s3.region": lake.cfg["region"],
|
||||
"s3.access-key-id": lake.cfg["access_key"],
|
||||
"s3.secret-access-key": lake.cfg["secret_key"],
|
||||
# Contabo serves buckets on the path, not as a subdomain.
|
||||
"s3.force-virtual-addressing": "false",
|
||||
}
|
||||
return SqlCatalog("thermograph", **props)
|
||||
|
||||
|
||||
def ensure_table(catalog: SqlCatalog, lake: HiveLake, create: bool = True):
|
||||
"""Load the table; adopt existing S3 metadata into a fresh catalog db;
|
||||
otherwise create it. ``create=False`` (dry runs) returns None instead of
|
||||
creating anything."""
|
||||
try:
|
||||
return catalog.load_table(IDENTIFIER)
|
||||
except NoSuchTableError:
|
||||
pass
|
||||
existing = lake.latest_metadata_uri()
|
||||
if existing:
|
||||
catalog.create_namespace_if_not_exists(NAMESPACE)
|
||||
print(f"registering existing table metadata: {existing}")
|
||||
return catalog.register_table(IDENTIFIER, existing)
|
||||
if not create:
|
||||
return None
|
||||
catalog.create_namespace_if_not_exists(NAMESPACE)
|
||||
return catalog.create_table(
|
||||
IDENTIFIER,
|
||||
schema=SCHEMA,
|
||||
location=f"{lake.warehouse}/{TABLE_NAME}",
|
||||
partition_spec=PARTITION_SPEC,
|
||||
properties={
|
||||
"format-version": "2",
|
||||
# The part files carry no Iceberg field ids; every reader needs
|
||||
# this name mapping to resolve columns (DuckDB included).
|
||||
"schema.name-mapping.default":
|
||||
create_mapping_from_schema(SCHEMA).model_dump_json(),
|
||||
})
|
||||
|
||||
|
||||
def synced_tiles(table) -> "set[str]":
|
||||
if table is None:
|
||||
return set()
|
||||
raw = table.properties.get(SYNCED_PROP, "")
|
||||
return {t for t in raw.split(",") if t}
|
||||
|
||||
|
||||
def tile_key(tile: str) -> "tuple[int, int]":
|
||||
ti, tj = tile.split("_")
|
||||
return int(ti), int(tj)
|
||||
|
||||
|
||||
def pending_tiles(manifest_tiles, synced: "set[str]") -> "list[str]":
|
||||
return sorted(set(manifest_tiles) - synced, key=tile_key)
|
||||
|
||||
|
||||
def sync_batch(table, lake: HiveLake, chunk: "list[str]",
|
||||
retries: int = DEFAULT_RETRIES) -> "list[str]":
|
||||
"""Commit one batch of tiles (files + synced-tiles property, atomically),
|
||||
retrying transient S3 failures with backoff. Each attempt re-reads the
|
||||
synced set from the (refreshed) table, so an attempt that committed but
|
||||
errored afterwards is not re-added. Returns the tiles actually added."""
|
||||
for attempt in range(retries + 1):
|
||||
try:
|
||||
table.refresh()
|
||||
left = [t for t in chunk if t not in synced_tiles(table)]
|
||||
if not left:
|
||||
return []
|
||||
files = []
|
||||
for t in left:
|
||||
tf = lake.tile_files(t)
|
||||
if not tf:
|
||||
raise LakeInvariantError(
|
||||
f"tile {t} is in the manifest but has no daily part "
|
||||
f"files — refusing to mark it synced")
|
||||
files.extend(tf)
|
||||
new_prop = ",".join(
|
||||
sorted(synced_tiles(table) | set(left), key=tile_key))
|
||||
with table.transaction() as tx:
|
||||
tx.add_files(
|
||||
file_paths=files,
|
||||
snapshot_properties={"thermograph.tiles-added":
|
||||
",".join(left)})
|
||||
tx.set_properties({SYNCED_PROP: new_prop})
|
||||
return left
|
||||
except LakeInvariantError:
|
||||
raise
|
||||
except Exception as e: # noqa: BLE001 - S3 throttling has many faces
|
||||
if attempt >= retries:
|
||||
raise
|
||||
wait = BACKOFF_SECONDS * 2 ** attempt
|
||||
print(f"batch {chunk[0]}..{chunk[-1]} attempt {attempt + 1} "
|
||||
f"failed ({e}); retrying in {wait}s", file=sys.stderr)
|
||||
time.sleep(wait)
|
||||
return [] # unreachable
|
||||
|
||||
|
||||
def write_version_hint(lake: HiveLake, table) -> "str | None":
|
||||
"""Hadoop-catalog-style pointer next to the metadata files so engines can
|
||||
scan the bare table root — DuckDB: iceberg_scan('s3://…/era5_daily') —
|
||||
without knowing the exact metadata path. Content is the current metadata
|
||||
file's basename minus '.metadata.json' (DuckDB's default
|
||||
version_name_format tries '%s.metadata.json', which resolves it). A
|
||||
convenience pointer only: refreshed every run, so a failed write just
|
||||
lags one version."""
|
||||
hint = os.path.basename(table.metadata_location) \
|
||||
.removesuffix(".metadata.json")
|
||||
path = f"{lake.root}/{ICEBERG_PREFIX}/{TABLE_NAME}/metadata/version-hint.text"
|
||||
for attempt in range(3):
|
||||
try:
|
||||
with lake.fs.open(path, "wb") as f:
|
||||
f.write(hint.encode())
|
||||
return hint
|
||||
except Exception as e: # noqa: BLE001
|
||||
if attempt == 2:
|
||||
print(f"WARNING: version-hint.text write failed ({e}); "
|
||||
f"table-root scans will lag one version", file=sys.stderr)
|
||||
return None
|
||||
time.sleep(BACKOFF_SECONDS)
|
||||
return None
|
||||
|
||||
|
||||
def run(lake: HiveLake, catalog_db: str, *, limit: "int | None" = None,
|
||||
batch: int = DEFAULT_BATCH, retries: int = DEFAULT_RETRIES,
|
||||
dry_run: bool = False) -> dict:
|
||||
tile_rows = lake.tile_rows()
|
||||
catalog = open_catalog(catalog_db, lake)
|
||||
table = ensure_table(catalog, lake, create=not dry_run)
|
||||
done = synced_tiles(table)
|
||||
todo = pending_tiles(tile_rows, done)
|
||||
if limit:
|
||||
todo = todo[:limit]
|
||||
print(f"manifest tiles={len(tile_rows)} synced={len(done)} "
|
||||
f"pending this run={len(todo)}")
|
||||
|
||||
added_files = added_records = 0
|
||||
mismatches = []
|
||||
if dry_run:
|
||||
for t in todo:
|
||||
files = lake.tile_files(t)
|
||||
print(f"[dry-run] tile {t}: {len(files)} part files, "
|
||||
f"{tile_rows[t]} manifest rows")
|
||||
else:
|
||||
for k in range(0, len(todo), batch):
|
||||
chunk = todo[k:k + batch]
|
||||
# One commit per batch: the files and the synced-tiles marker land
|
||||
# atomically, so an interrupted run never half-syncs a tile.
|
||||
added = sync_batch(table, lake, chunk, retries)
|
||||
done |= set(chunk)
|
||||
if not added:
|
||||
print(f"[{min(k + batch, len(todo))}/{len(todo)}] tiles "
|
||||
f"{chunk[0]}..{chunk[-1]}: already synced")
|
||||
continue
|
||||
summary = table.current_snapshot().summary
|
||||
got = int(summary.get("added-records", "0"))
|
||||
expected = sum(tile_rows[t] for t in added)
|
||||
n_files = int(summary.get("added-data-files", "0"))
|
||||
added_files += n_files
|
||||
added_records += got
|
||||
status = "ok" if got == expected else "MISMATCH"
|
||||
if got != expected:
|
||||
mismatches.append((added, expected, got))
|
||||
print(f"[{min(k + batch, len(todo))}/{len(todo)}] tiles "
|
||||
f"{added[0]}..{added[-1]}: +{n_files} files, "
|
||||
f"+{got} records (manifest says {expected}: {status})")
|
||||
|
||||
result = {"pending": todo, "added_files": added_files,
|
||||
"added_records": added_records, "mismatches": mismatches,
|
||||
"metadata_location": None, "total_records": 0,
|
||||
"total_files": 0}
|
||||
if table is not None:
|
||||
snap = table.current_snapshot()
|
||||
if snap is not None:
|
||||
result["total_records"] = int(snap.summary.get("total-records", "0"))
|
||||
result["total_files"] = int(snap.summary.get("total-data-files", "0"))
|
||||
result["metadata_location"] = table.metadata_location
|
||||
if not dry_run:
|
||||
result["version_hint"] = write_version_hint(lake, table)
|
||||
expected_total = sum(tile_rows[t] for t in done if t in tile_rows)
|
||||
print(f"table: {result['total_files']} data files, "
|
||||
f"{result['total_records']} records "
|
||||
f"(manifest rows for synced tiles: {expected_total})")
|
||||
print(f"metadata: {table.metadata_location}")
|
||||
if lake.cfg:
|
||||
print("read it from DuckDB (no catalog needed):\n"
|
||||
" LOAD iceberg; LOAD httpfs; -- plus an S3 secret for the "
|
||||
"endpoint\n"
|
||||
f" SELECT count(*) FROM iceberg_scan("
|
||||
f"'{lake.warehouse}/{TABLE_NAME}');\n"
|
||||
f" -- or pin the exact version: iceberg_scan("
|
||||
f"'{table.metadata_location}')")
|
||||
if mismatches:
|
||||
print(f"WARNING: {len(mismatches)} batch(es) added a record count "
|
||||
f"that differs from the manifest — investigate before trusting "
|
||||
f"the table", file=sys.stderr)
|
||||
return result
|
||||
|
||||
|
||||
def main() -> int:
|
||||
ap = argparse.ArgumentParser(
|
||||
description="Register new ERA5 lake tiles into the Iceberg table")
|
||||
ap.add_argument("--catalog-db",
|
||||
default=os.environ.get("THERMOGRAPH_ICEBERG_CATALOG_DB",
|
||||
DEFAULT_CATALOG_DB),
|
||||
help="sqlite catalog path (default: %(default)s)")
|
||||
ap.add_argument("--batch", type=int, default=DEFAULT_BATCH,
|
||||
help="tiles per Iceberg commit (default: %(default)s)")
|
||||
ap.add_argument("--limit", type=int,
|
||||
help="sync at most this many new tiles this run")
|
||||
ap.add_argument("--retries", type=int, default=DEFAULT_RETRIES,
|
||||
help="per-batch retries on transient S3 errors "
|
||||
"(default: %(default)s)")
|
||||
ap.add_argument("--dry-run", action="store_true",
|
||||
help="list pending tiles; write nothing")
|
||||
args = ap.parse_args()
|
||||
|
||||
lake = lake_from_env()
|
||||
result = run(lake, args.catalog_db, limit=args.limit, batch=args.batch,
|
||||
retries=args.retries, dry_run=args.dry_run)
|
||||
return 1 if result["mismatches"] else 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
247
infra/lake-iceberg/test_sync.py
Normal file
247
infra/lake-iceberg/test_sync.py
Normal file
|
|
@ -0,0 +1,247 @@
|
|||
"""Unit tests for the Iceberg lake sync: incremental tile diffing and
|
||||
add_files planning against a local filesystem fixture (a miniature lake in a
|
||||
tmp dir, a file:// warehouse, a sqlite catalog). No network."""
|
||||
import datetime
|
||||
import os
|
||||
import sys
|
||||
|
||||
import pyarrow as pa
|
||||
import pyarrow.parquet as pq
|
||||
import pytest
|
||||
from fsspec.implementations.local import LocalFileSystem
|
||||
|
||||
sys.path.insert(0, os.path.dirname(__file__))
|
||||
import sync_iceberg as si # noqa: E402
|
||||
|
||||
# Mirrors the real part-file schema (float32/float64 split, int16 doy,
|
||||
# all-null float64 sun) so the compatibility check is exercised for real.
|
||||
PART_SCHEMA = pa.schema([
|
||||
("date", pa.date32()),
|
||||
("tmax", pa.float32()),
|
||||
("tmin", pa.float32()),
|
||||
("precip", pa.float32()),
|
||||
("wind", pa.float32()),
|
||||
("gust", pa.float32()),
|
||||
("humid", pa.float32()),
|
||||
("fmax", pa.float32()),
|
||||
("fmin", pa.float64()),
|
||||
("sun", pa.float64()),
|
||||
("feels", pa.float64()),
|
||||
("doy", pa.int16()),
|
||||
("lat_idx", pa.int32()),
|
||||
("lon_idx", pa.int32()),
|
||||
])
|
||||
|
||||
POINTS = 2 # points per fixture tile
|
||||
DAYS = 3 # rows per point per month
|
||||
|
||||
|
||||
def write_tile(root, ti, tj, months):
|
||||
"""Write one tile's hive part files; return its total row count."""
|
||||
total = 0
|
||||
for year, month in months:
|
||||
rows = {name: [] for name in PART_SCHEMA.names}
|
||||
for p in range(POINTS):
|
||||
i, j = ti * si.TILE + p, tj * si.TILE + p
|
||||
for d in range(DAYS):
|
||||
day = datetime.date(year, month, d + 1)
|
||||
rows["date"].append(day)
|
||||
for col in ("tmax", "tmin", "precip", "wind", "gust",
|
||||
"humid", "fmax", "fmin", "feels"):
|
||||
rows[col].append(float(10 * p + d))
|
||||
rows["sun"].append(None)
|
||||
rows["doy"].append(day.timetuple().tm_yday)
|
||||
rows["lat_idx"].append(i)
|
||||
rows["lon_idx"].append(j)
|
||||
part = pa.table(rows, schema=PART_SCHEMA)
|
||||
path = os.path.join(root, si.PREFIX, "daily", f"tile={ti}_{tj}",
|
||||
f"year={year}", f"month={month:02d}",
|
||||
"part.parquet")
|
||||
os.makedirs(os.path.dirname(path), exist_ok=True)
|
||||
pq.write_table(part, path)
|
||||
total += part.num_rows
|
||||
return total
|
||||
|
||||
|
||||
def write_manifest(root, tile_rows):
|
||||
"""Manifest with one row per (tile, point) like the extractor writes;
|
||||
the sync only reads tile+rows."""
|
||||
tiles, rows = [], []
|
||||
for tile, total in tile_rows.items():
|
||||
per_point = total // POINTS
|
||||
for _ in range(POINTS):
|
||||
tiles.append(tile)
|
||||
rows.append(per_point)
|
||||
m = pa.table({
|
||||
"lat_idx": pa.array([0] * len(tiles), pa.int64()),
|
||||
"lon_idx": pa.array([0] * len(tiles), pa.int64()),
|
||||
"lat": pa.array([0.0] * len(tiles)),
|
||||
"lon": pa.array([0.0] * len(tiles)),
|
||||
"tile": pa.array(tiles, pa.large_string()),
|
||||
"rows": pa.array(rows, pa.int64()),
|
||||
"date_min": pa.array(["1940-01-01"] * len(tiles), pa.large_string()),
|
||||
"date_max": pa.array(["2026-01-01"] * len(tiles), pa.large_string()),
|
||||
})
|
||||
path = os.path.join(root, si.MANIFEST_KEY)
|
||||
os.makedirs(os.path.dirname(path), exist_ok=True)
|
||||
pq.write_table(m, path)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def lake(tmp_path):
|
||||
root = str(tmp_path / "lake")
|
||||
counts = {
|
||||
"10_10": write_tile(root, 10, 10, [(2000, 1), (2000, 2)]),
|
||||
"10_11": write_tile(root, 10, 11, [(2000, 1)]),
|
||||
}
|
||||
write_manifest(root, counts)
|
||||
hl = si.HiveLake(LocalFileSystem(), root, f"file://{root}", None)
|
||||
hl.fixture_counts = counts
|
||||
hl.fixture_root = root
|
||||
return hl
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def catalog_db(tmp_path):
|
||||
return str(tmp_path / "catalog.db")
|
||||
|
||||
|
||||
def test_pending_tiles_diff():
|
||||
manifest = {"10_10": 1, "10_11": 1, "2_3": 1}
|
||||
assert si.pending_tiles(manifest, set()) == ["2_3", "10_10", "10_11"]
|
||||
assert si.pending_tiles(manifest, {"10_10"}) == ["2_3", "10_11"]
|
||||
assert si.pending_tiles(manifest, set(manifest)) == []
|
||||
# numeric, not lexicographic, tile ordering
|
||||
assert si.pending_tiles({"10_2": 1, "9_11": 1}, set()) == ["9_11", "10_2"]
|
||||
|
||||
|
||||
def test_initial_sync_registers_all_files_in_place(lake, catalog_db):
|
||||
result = si.run(lake, catalog_db, batch=1)
|
||||
expected = sum(lake.fixture_counts.values())
|
||||
assert result["added_records"] == expected
|
||||
assert result["total_records"] == expected
|
||||
assert result["added_files"] == 3 # 2 months + 1 month
|
||||
assert not result["mismatches"]
|
||||
|
||||
table = si.open_catalog(catalog_db, lake).load_table(si.IDENTIFIER)
|
||||
assert si.synced_tiles(table) == set(lake.fixture_counts)
|
||||
# add_files, not a rewrite: every data file is an original hive part file
|
||||
files = table.inspect.files().column("file_path").to_pylist()
|
||||
assert len(files) == 3
|
||||
assert all(f"/{si.PREFIX}/daily/tile=" in f for f in files)
|
||||
# nothing was written outside iceberg/ (the hive tree is untouched, the
|
||||
# table tree holds only metadata)
|
||||
fs = LocalFileSystem()
|
||||
written = fs.find(os.path.join(lake.fixture_root, si.ICEBERG_PREFIX))
|
||||
assert written and all("/metadata/" in p for p in written)
|
||||
|
||||
|
||||
def test_rerun_is_noop(lake, catalog_db):
|
||||
first = si.run(lake, catalog_db)
|
||||
again = si.run(lake, catalog_db)
|
||||
assert again["pending"] == []
|
||||
assert again["added_records"] == 0
|
||||
assert again["total_records"] == first["total_records"]
|
||||
assert again["metadata_location"] == first["metadata_location"]
|
||||
|
||||
|
||||
def test_incremental_new_tile_only(lake, catalog_db):
|
||||
si.run(lake, catalog_db)
|
||||
counts = dict(lake.fixture_counts)
|
||||
counts["11_10"] = write_tile(lake.fixture_root, 11, 10, [(2001, 6)])
|
||||
write_manifest(lake.fixture_root, counts)
|
||||
|
||||
result = si.run(lake, catalog_db)
|
||||
assert result["pending"] == ["11_10"]
|
||||
assert result["added_records"] == counts["11_10"]
|
||||
assert result["total_records"] == sum(counts.values())
|
||||
|
||||
|
||||
def test_version_hint_tracks_current_metadata(lake, catalog_db):
|
||||
first = si.run(lake, catalog_db)
|
||||
hint_path = os.path.join(lake.fixture_root, si.ICEBERG_PREFIX,
|
||||
si.TABLE_NAME, "metadata", "version-hint.text")
|
||||
with open(hint_path) as f:
|
||||
assert f.read() == os.path.basename(
|
||||
first["metadata_location"]).removesuffix(".metadata.json")
|
||||
|
||||
counts = dict(lake.fixture_counts)
|
||||
counts["11_10"] = write_tile(lake.fixture_root, 11, 10, [(2001, 6)])
|
||||
write_manifest(lake.fixture_root, counts)
|
||||
second = si.run(lake, catalog_db)
|
||||
assert second["metadata_location"] != first["metadata_location"]
|
||||
with open(hint_path) as f:
|
||||
assert f.read() == os.path.basename(
|
||||
second["metadata_location"]).removesuffix(".metadata.json")
|
||||
|
||||
|
||||
def test_partition_spec_and_pruned_scan(lake, catalog_db):
|
||||
si.run(lake, catalog_db)
|
||||
table = si.open_catalog(catalog_db, lake).load_table(si.IDENTIFIER)
|
||||
assert [f.name for f in table.spec().fields] == \
|
||||
["tile_lat", "tile_lon", "year", "month"]
|
||||
|
||||
parts = table.inspect.partitions()
|
||||
keys = {(r["partition"]["tile_lat"], r["partition"]["tile_lon"])
|
||||
for r in parts.to_pylist()}
|
||||
assert keys == {(120, 120), (120, 132)} # ti*12, tj*12
|
||||
|
||||
# a tile-shaped predicate resolves to that tile's rows only
|
||||
scan = table.scan(row_filter="lat_idx >= 120 and lat_idx < 132 "
|
||||
"and lon_idx >= 132 and lon_idx < 144")
|
||||
assert scan.to_arrow().num_rows == lake.fixture_counts["10_11"]
|
||||
# and file planning pruned to that tile's single part file
|
||||
assert len(list(scan.plan_files())) == 1
|
||||
|
||||
|
||||
def test_lost_catalog_db_reregisters_without_readding(lake, catalog_db):
|
||||
first = si.run(lake, catalog_db)
|
||||
os.remove(catalog_db) # simulate losing local state
|
||||
again = si.run(lake, catalog_db)
|
||||
# synced-tiles lives in table metadata, so nothing is re-added
|
||||
assert again["pending"] == []
|
||||
assert again["added_records"] == 0
|
||||
assert again["total_records"] == first["total_records"]
|
||||
|
||||
|
||||
def test_dry_run_writes_nothing(lake, catalog_db):
|
||||
result = si.run(lake, catalog_db, dry_run=True)
|
||||
assert sorted(result["pending"]) == sorted(lake.fixture_counts)
|
||||
assert result["metadata_location"] is None
|
||||
assert not os.path.exists(
|
||||
os.path.join(lake.fixture_root, si.ICEBERG_PREFIX))
|
||||
|
||||
|
||||
def test_manifest_tile_without_files_refuses(lake, catalog_db):
|
||||
counts = dict(lake.fixture_counts)
|
||||
counts["50_50"] = 6 # in the manifest, no files
|
||||
write_manifest(lake.fixture_root, counts)
|
||||
# an invariant break aborts immediately: no retries, nothing marked synced
|
||||
with pytest.raises(si.LakeInvariantError, match="50_50"):
|
||||
si.run(lake, catalog_db)
|
||||
|
||||
|
||||
def test_transient_failure_retries_and_succeeds(lake, catalog_db, monkeypatch):
|
||||
monkeypatch.setattr(si, "BACKOFF_SECONDS", 0)
|
||||
fails = {"n": 2}
|
||||
original = si.HiveLake.tile_files
|
||||
|
||||
def flaky(self, tile):
|
||||
if fails["n"]:
|
||||
fails["n"] -= 1
|
||||
raise OSError("AWS Error SLOW_DOWN during CompleteMultipartUpload")
|
||||
return original(self, tile)
|
||||
|
||||
monkeypatch.setattr(si.HiveLake, "tile_files", flaky)
|
||||
result = si.run(lake, catalog_db, retries=2)
|
||||
assert result["added_records"] == sum(lake.fixture_counts.values())
|
||||
assert not result["mismatches"]
|
||||
|
||||
|
||||
def test_transient_failure_exhausts_retries(lake, catalog_db, monkeypatch):
|
||||
monkeypatch.setattr(si, "BACKOFF_SECONDS", 0)
|
||||
monkeypatch.setattr(
|
||||
si.HiveLake, "tile_files",
|
||||
lambda self, tile: (_ for _ in ()).throw(OSError("SLOW_DOWN")))
|
||||
with pytest.raises(OSError):
|
||||
si.run(lake, catalog_db, retries=1)
|
||||
Loading…
Reference in a new issue