thermograph/infra/lake-iceberg/sync_iceberg.py

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