thermograph/infra/lake-iceberg/test_sync.py

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