"""Shared external-incremental engine: pluggable strategies, delta, drift gate (ADR 0018). The external loader (used by both gcs_bulk_load and s3_bulk_load) loads buckets we do not extract: no manifest, schema inferred + drift-gated, snapshots loaded incrementally per a pluggable IncrementalStrategy. This pins the two strategies (filename / last_modified), the shared delta selection, and — with a recording fake cursor — the SQL the last_modified path generates (so the S3 port is exercised without live AWS). Runnable without a test harness:: .venv/bin/python tests/test_external_incremental_load.py """ from __future__ import annotations import pathlib import sys import tempfile from types import SimpleNamespace sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[1])) from ingestion.common.external import infer from ingestion.common.external import load as extload from ingestion.common.external.load import SchemaDriftError, external_load_run, relativize, sample_object, select_increments from ingestion.common.external.strategy import FilenameStrategy, LastModifiedStrategy, build_strategy from ingestion.common.load_config import LoadConfigBase, TableTargets from ingestion.common.raw_product import raw_product _OBJECTS = [ {"name": "p/billing2_invoices/run-1778601958447-part-r-00000-snappy.parquet", "size": 1, "last_modified_ms": 100}, {"name": "p/billing2_invoices/run-1778601958447-part-r-00001-snappy.parquet", "size": 1, "last_modified_ms": 110}, {"name": "p/billing2_invoices/run-1778688388773-part-r-00000-snappy.parquet", "size": 1, "last_modified_ms": 200}, {"name": "p/billing2_invoices/_SUCCESS", "size": 0, "last_modified_ms": 210}, ] def test_filename_strategy_groups_by_ts_ignores_nonparquet() -> None: incs = FilenameStrategy(r"run-(?P\d+)-").increments(_OBJECTS) by_key = {i.key: len(i.files) for i in incs} assert by_key == {1778601958447: 2, 1778688388773: 1} # _SUCCESS ignored assert FilenameStrategy(r"run-(?P\d+)-").watermark_column == "_extraction_ts" def test_last_modified_strategy_is_per_file_keyed_by_mtime() -> None: incs = LastModifiedStrategy().increments(_OBJECTS) assert sorted(i.key for i in incs) == [100, 110, 200] # one per parquet, _SUCCESS ignored assert all(len(i.files) == 1 for i in incs) assert LastModifiedStrategy().watermark_column == "_source_last_modified_ms" def test_build_strategy_and_provenance_columns() -> None: assert build_strategy({"strategy": "filename"}).name == "filename" assert build_strategy({"strategy": "last_modified"}).name == "last_modified" assert build_strategy({}).name == "filename" # default assert dict(FilenameStrategy(r"run-(?P\d+)-").provenance_columns()) == { "_extraction_ts": "BIGINT", "_extraction_at": "TIMESTAMP" } assert dict(LastModifiedStrategy().provenance_columns()) == { "_source_file": "STRING", "_source_last_modified_ms": "BIGINT", "_source_last_modified": "TIMESTAMP" } def _inc(key): from ingestion.common.external.strategy import Increment return Increment(key=key, label=str(key), files=[f"f{key}"]) def test_select_increments_delta_semantics() -> None: incs = [_inc(3), _inc(1), _inc(2)] assert [i.key for i in select_increments(incs, None, True, None)] == [1, 2, 3] # first load: all assert [i.key for i in select_increments(incs, 2, True, None)] == [3] # only newer assert [i.key for i in select_increments(incs, 3, True, None)] == [] # nothing newer assert [i.key for i in select_increments(incs, 3, True, 1)] == [1] # reload override assert [i.key for i in select_increments(incs, 2, False, None)] == [1, 2, 3] # delta off def test_relativize_strips_stage_prefix() -> None: rel = relativize("s3://bucket/events/", ["events/a.parquet", "events/sub/b.parquet"]) assert rel == {"events/a.parquet": "a.parquet", "events/sub/b.parquet": "sub/b.parquet"} def test_schema_changed() -> None: committed = [{"name": "id", "type": "DECIMAL(20, 0)", "nullable": True}] assert infer.schema_changed(list(committed), committed) is False assert infer.schema_changed([{"name": "id", "type": "VARCHAR", "nullable": True}], committed) is True # --- Integration: drive external_load_run with fakes (no live warehouse / cloud) ---------- class RecordingCursor: def __init__(self, infer_rows, max_wm, count): self.sql: list[str] = [] self._infer_rows = infer_rows self._max_wm = max_wm self._count = count self._last = "" def execute(self, sql): self.sql.append(sql) self._last = sql def fetchall(self): return self._infer_rows if "INFER_SCHEMA" in self._last else [] def fetchone(self): if "max(" in self._last: return (self._max_wm,) if "coalesce(sum" in self._last: return (0,) return (self._count,) class FakeStorage: def list_objects(self, uri): return [ {"name": "events/a.parquet", "size": 1, "last_modified_ms": 100}, {"name": "events/b.parquet", "size": 1, "last_modified_ms": 200}, ] def create_stage_sql(self, stage_name, uri): return f"CREATE OR REPLACE STAGE {stage_name} URL = '{uri}'" def credential_env_names(self): return ["EVENTS_S3_SECRET_ACCESS_KEY"] def _fake_config(root: pathlib.Path): return SimpleNamespace( name="events", root=root, external_uri="s3://bucket/events/", external_table="events", external_primary_key=["event_id"], external_domain="telemetry", external_owner="data-platform", external_system="external-s3", external_bucket="bucket", stage_name="events_ext_stage", delta=True, classification_rules={"baseline": ["internal"], "patterns": {}, "name_rules": []}, resolve_contracts=lambda: {"events": TableTargets(table="events", raw_table="events_raw.events", classifications=["internal"], product_ids=[])}, ) def test_external_source_contract_applies_configured_overrides() -> None: config = _fake_config(pathlib.Path("/tmp/events")) config.external_source_contract_overrides = { "purpose": "Source-faithful one-off provided Parquet export of events.", "source": {"description": "One-off provided object-store snapshot export."}, "assets": [ { "produced_by": {"kind": "other", "ref": "one-off audit helper"}, } ], "freshness": {"max_age": "P3650D", "degrade_after": "P3650D"}, "tests": [{"id": "events_source_one_off", "type": "sql", "severity": "informational", "sql": "select 1"}], "governance": {"retention": {"period": "P3650D"}}, "extensions": {"lifecycle": "one_off_ad_hoc", "managed_extract": False, "row_estimate": 5}, } contract = infer.source_contract( config, [ {"name": "event_id", "type": "VARCHAR", "nullable": True}, {"name": "payload", "type": "VARCHAR", "nullable": True}, ], ) assert contract["purpose"] == "Source-faithful one-off provided Parquet export of events." assert contract["source"]["description"] == "One-off provided object-store snapshot export." assert contract["assets"][0]["ref"] == "s3://bucket/events/" assert contract["assets"][0]["produced_by"]["ref"] == "one-off audit helper" assert contract["freshness"] == {"max_age": "P3650D", "degrade_after": "P3650D"} assert contract["tests"][0]["id"] == "events_source_one_off" assert contract["governance"]["retention"] == {"period": "P3650D"} assert contract["extensions"]["primary_key"] == ["event_id"] assert contract["extensions"]["lifecycle"] == "one_off_ad_hoc" assert contract["extensions"]["managed_extract"] is False assert contract["extensions"]["row_estimate"] == 5 def test_contract_overrides_reject_invariant_paths() -> None: config = _fake_config(pathlib.Path("/tmp/events")) config.external_source_contract_overrides = { "schema": {"columns": []}, } raised = False try: infer.source_contract(config, [{"name": "event_id", "type": "VARCHAR", "nullable": True}]) except ValueError as exc: raised = "schema" in str(exc) assert raised, "source contract overrides must not replace inferred schema" def test_raw_product_applies_table_contract_overrides() -> None: source = infer.source_contract( _fake_config(pathlib.Path("/tmp/events")), [{"name": "event_id", "type": "VARCHAR", "nullable": True}], ) source["extensions"]["row_estimate"] = 5 config = SimpleNamespace( name="events", cloud_lake_databases={"raw": "events_raw"}, raw_contract_override_for=lambda table: { "freshness": {"max_age": "P3650D", "degrade_after": "P3650D"}, "tests": [{"id": f"{table}_raw_expected_rows", "type": "sql", "severity": "blocking", "sql": "select 1"}], "governance": {"retention": {"period": "P3650D"}}, }, ) contract = raw_product(config, source, "ingestion.load.s3_bulk_load.engine.load") assert contract["freshness"] == {"max_age": "P3650D", "degrade_after": "P3650D"} assert contract["tests"][0]["id"] == "events_raw_expected_rows" assert contract["governance"]["retention"] == {"period": "P3650D"} assert contract["extensions"]["primary_key"] == ["event_id"] assert contract["extensions"]["row_estimate"] == 5 def test_raw_product_overrides_reject_layer_changes() -> None: source = infer.source_contract( _fake_config(pathlib.Path("/tmp/events")), [{"name": "event_id", "type": "VARCHAR", "nullable": True}], ) config = SimpleNamespace( name="events", cloud_lake_databases={"raw": "events_raw"}, raw_contract_override_for=lambda table: {"layer": "staging"}, ) raised = False try: raw_product(config, source, "ingestion.load.s3_bulk_load.engine.load") except ValueError as exc: raised = "layer" in str(exc) assert raised, "raw contract overrides must not change load-layer ownership" def test_load_config_merges_external_contract_overrides() -> None: config = LoadConfigBase( path=pathlib.Path("/tmp/source.yaml"), data={ "load": {"name": "events"}, "external": { "uri": "s3://bucket/events/", "table": "events", "domain": "telemetry", "source_contract": {"freshness": {"max_age": "P3650D"}}, "raw_contract": {"tests": [{"id": "events_raw_expected_rows"}]}, }, "raw_contracts": {"events": {"governance": {"retention": {"period": "P3650D"}}}}, }, ) assert config.external_source_contract_overrides == {"freshness": {"max_age": "P3650D"}} assert config.raw_contract_override_for("events") == { "governance": {"retention": {"period": "P3650D"}}, "tests": [{"id": "events_raw_expected_rows"}], } def _write_committed(root, columns): import yaml path = root / "contracts" / "upstream" / "events.events.external.source.contract.yaml" path.parent.mkdir(parents=True) path.write_text(yaml.safe_dump({"schema": {"columns": columns}})) def _patch(cursor): orig = extload.client extload.client = lambda config: SimpleNamespace(cursor=lambda: cursor) return orig def test_last_modified_happy_path_generates_expected_sql() -> None: cols = [{"name": "event_id", "type": "VARCHAR", "nullable": True}, {"name": "payload", "type": "VARCHAR", "nullable": True}] with tempfile.TemporaryDirectory() as tmp: root = pathlib.Path(tmp) _write_committed(root, cols) cursor = RecordingCursor(infer_rows=[(c["name"], c["type"], c["nullable"]) for c in cols], max_wm=None, count=5) orig = _patch(cursor) try: results = external_load_run(_fake_config(root), FakeStorage(), LastModifiedStrategy(), reload_key=None) finally: extload.client = orig assert [r["cursor"] for r in results] == [100, 200], "first load takes both files, ascending by mtime" sql = "\n".join(cursor.sql) create = next(s for s in cursor.sql if "CREATE TABLE" in s) for col in ("_source_file", "_source_last_modified_ms", "_source_last_modified", "_loaded_at", "_source_uri", "_source_table"): assert col in create, f"raw table must carry {col}" copies = [s for s in cursor.sql if "COPY INTO" in s] assert len(copies) == 2 and all("FILES = (" in c for c in copies), "per-file COPY via FILES" # COPY FILES uses the stage-relative path... assert "'a.parquet'" in copies[0] and "'b.parquet'" in copies[1] # ...while _source_file stamps the full object key, and clear+stamp use the SAME value # (regression guard: a mismatch silently zeroed loaded_rows + broke idempotency). assert "`_source_last_modified_ms` = 200" in sql and "`_source_file` = 'events/b.parquet'" in sql deletes = [s for s in cursor.sql if s.strip().startswith("DELETE")] assert any("'events/b.parquet'" in d for d in deletes), "clear predicate must match the stamped _source_file" def test_drift_gate_writes_proposed_and_aborts() -> None: committed = [{"name": "event_id", "type": "VARCHAR", "nullable": True}] with tempfile.TemporaryDirectory() as tmp: root = pathlib.Path(tmp) _write_committed(root, committed) # Live infers an extra column -> drift. live_rows = [("event_id", "VARCHAR", True), ("new_col", "VARCHAR", True)] cursor = RecordingCursor(infer_rows=live_rows, max_wm=None, count=0) orig = _patch(cursor) raised = False try: external_load_run(_fake_config(root), FakeStorage(), LastModifiedStrategy(), reload_key=None) except SchemaDriftError: raised = True finally: extload.client = orig assert raised, "schema drift must raise (contract update process)" assert (root / "contracts" / "proposed" / "events.events.external.source.contract.yaml").exists() assert not any("COPY INTO" in s for s in cursor.sql), "must not load on drift" if __name__ == "__main__": cases = [value for name, value in sorted(globals().items()) if name.startswith("test_") and callable(value)] for case in cases: case() print(f"ok {case.__name__}") print(f"\n{len(cases)} passed")