diff --git a/mkdocs/docs/api.md b/mkdocs/docs/api.md index 0a2a72b6f0..8af8d5c17b 100644 --- a/mkdocs/docs/api.md +++ b/mkdocs/docs/api.md @@ -1483,6 +1483,39 @@ Remove an existing branch: table.manage_snapshots().remove_branch("dev").commit() ``` +### Write-Audit-Publish + +Stage a write on a branch, validate it, then publish it to the table with +`cherry_pick_snapshot`. The staged data is invisible to readers of the table until it is +published. + +```python +# Write: stage the changes on an audit branch +table.manage_snapshots().create_branch( + snapshot_id=table.metadata.current_snapshot_id, + branch_name="audit-2024-01-15", +).commit() + +table = catalog.load_table("db.table") +table.append(new_rows, branch="audit-2024-01-15", snapshot_properties={"wap.id": "etl-2024-01-15"}) + +# Audit: validate the staged data without affecting readers of the table +table = catalog.load_table("db.table") +staged = table.metadata.refs["audit-2024-01-15"].snapshot_id +assert len(table.scan(snapshot_id=staged).to_arrow()) > 0 + +# Publish: replay the staged changes onto the current table state +table.manage_snapshots().cherry_pick_snapshot(staged).commit() +``` + +The published snapshot records `source-snapshot-id`, and `published-wap-id` when the staged +snapshot carried a `wap.id`. A given `wap.id` can only be published once. + +Append snapshots are always replayed, so the wap trail is recorded even when the table has not +changed since the branch was cut. A snapshot with any other operation is fast-forwarded to when +its parent is already the current snapshot, and raises otherwise. Picking a snapshot that is +already an ancestor of the current state does nothing. + ## Table Maintenance PyIceberg provides table maintenance operations through the `table.maintenance` API. This provides a clean interface for performing maintenance tasks like snapshot expiration. diff --git a/pyiceberg/table/update/snapshot.py b/pyiceberg/table/update/snapshot.py index 7931edacdd..e78af4b959 100644 --- a/pyiceberg/table/update/snapshot.py +++ b/pyiceberg/table/update/snapshot.py @@ -80,6 +80,12 @@ if TYPE_CHECKING: from pyiceberg.table import Transaction +# Snapshot summary properties for Write-Audit-Publish, mirroring Java's SnapshotSummary: +# https://github.com/apache/iceberg/blob/0d60b2bb8780d781f4ceb69032418d222b649ae5/core/src/main/java/org/apache/iceberg/SnapshotSummary.java#L60-L62 +_STAGED_WAP_ID_PROP = "wap.id" +_PUBLISHED_WAP_ID_PROP = "published-wap-id" +_SOURCE_SNAPSHOT_ID_PROP = "source-snapshot-id" + def _new_manifest_file_name(num: int, commit_uuid: uuid.UUID) -> str: return f"{commit_uuid}-m{num}.avro" @@ -974,6 +980,87 @@ def set_current_snapshot(self, snapshot_id: int | None = None, ref_name: str | N self._transaction._stage(update, requirement) return self + def cherry_pick_snapshot(self, snapshot_id: int) -> ManageSnapshots: + """Apply the changes in a snapshot to the current table state. + + Creates a new snapshot on top of the current one carrying the data files the picked + snapshot added, which is how a Write-Audit-Publish staged write is published. The picked + snapshot itself is left in place. + + Append snapshots are always replayed, matching Java's ``CherryPickOperation``. A snapshot + with another operation is fast-forwarded to when its parent is already the current + snapshot, and rejected otherwise. Picking a snapshot that is already an ancestor of the + current state does nothing. + + A snapshot staged with a ``wap.id`` can only be published once. The new snapshot records + ``source-snapshot-id``, and ``published-wap-id`` when the picked snapshot carried one. + + Args: + snapshot_id: The ID of the snapshot to cherry-pick. + + Returns: + This for method chaining. + + Raises: + ValueError: If the snapshot does not exist, its ``wap.id`` was already published, or + its operation cannot be cherry-picked. + """ + self._commit_if_ref_updates_exist() + + metadata = self._transaction.table_metadata + picked = metadata.snapshot_by_id(snapshot_id) + if picked is None: + raise ValueError(f"Cannot cherry-pick unknown snapshot id: {snapshot_id}") + + if self._is_current_ancestor(snapshot_id): + return self + + wap_id = self._validate_wap_publish(picked) + operation = picked.summary.operation if picked.summary else None + + if operation == Operation.APPEND: + snapshot_properties = {_SOURCE_SNAPSHOT_ID_PROP: str(snapshot_id)} + if wap_id is not None: + snapshot_properties[_PUBLISHED_WAP_ID_PROP] = wap_id + + with self._transaction.update_snapshot(snapshot_properties=snapshot_properties).fast_append() as append: + for data_file in self._added_data_files(picked): + append.append_data_file(data_file) + return self + + if picked.parent_snapshot_id == metadata.current_snapshot_id: + return self.set_current_snapshot(snapshot_id=snapshot_id) + + raise ValueError( + f"Cannot cherry-pick snapshot {snapshot_id}: not append, dynamic overwrite, or fast-forward " + f"(operation: {operation}). Only append snapshots can be replayed." + ) + + def _added_data_files(self, snapshot: Snapshot) -> list[DataFile]: + """Return the data files the given snapshot added, excluding entries carried over from its ancestors.""" + io = self._transaction._table.io + return [ + entry.data_file + for manifest in snapshot.manifests(io) + if manifest.added_snapshot_id == snapshot.snapshot_id + for entry in manifest.fetch_manifest_entry(io, discard_deleted=True) + if entry.status == ManifestEntryStatus.ADDED and entry.snapshot_id == snapshot.snapshot_id + ] + + def _validate_wap_publish(self, picked: Snapshot) -> str | None: + """Return the picked snapshot's staged wap id, rejecting one that an ancestor already published.""" + wap_id = picked.summary.additional_properties.get(_STAGED_WAP_ID_PROP) if picked.summary else None + if not wap_id: + return None + + metadata = self._transaction.table_metadata + for ancestor in ancestors_of(metadata.current_snapshot(), metadata): + props = ancestor.summary.additional_properties if ancestor.summary else {} + if wap_id in (props.get(_STAGED_WAP_ID_PROP), props.get(_PUBLISHED_WAP_ID_PROP)): + raise ValueError(f"Duplicate request to cherry pick wap id that was published already: {wap_id}") + + return wap_id + def rollback_to_snapshot(self, snapshot_id: int) -> ManageSnapshots: """Rollback the table to the given snapshot id. diff --git a/tests/integration/test_snapshot_operations.py b/tests/integration/test_snapshot_operations.py index 07fb77edbb..ec8fd3f745 100644 --- a/tests/integration/test_snapshot_operations.py +++ b/tests/integration/test_snapshot_operations.py @@ -332,3 +332,48 @@ def test_rollback_to_timestamp_chained_with_tag(table_with_snapshots: Table) -> assert table_with_snapshots.metadata.refs[tag_name] == SnapshotRef( snapshot_id=current_snapshot.snapshot_id, snapshot_ref_type="tag" ) + + +@pytest.mark.integration +@pytest.mark.parametrize("catalog", [lf("session_catalog_hive"), lf("session_catalog")]) +def test_cherry_pick_snapshot_publishes_staged_write(catalog: Catalog) -> None: + """Publish a branch-staged write onto a main that has moved on since the branch was cut.""" + catalog.create_namespace_if_not_exists("default") + identifier = f"default.test_cherry_pick_{uuid.uuid4().hex[:8]}" + arrow_schema = pa.schema([pa.field("id", pa.int64(), nullable=False)]) + tbl = catalog.create_table(identifier=identifier, schema=arrow_schema) + + tbl.append(pa.Table.from_pylist([{"id": 1}], schema=arrow_schema)) + tbl = catalog.load_table(identifier) + current_snapshot_id = tbl.metadata.current_snapshot_id + assert current_snapshot_id is not None + tbl.manage_snapshots().create_branch(snapshot_id=current_snapshot_id, branch_name="audit").commit() + + tbl = catalog.load_table(identifier) + tbl.append( + pa.Table.from_pylist([{"id": 2}, {"id": 3}], schema=arrow_schema), + branch="audit", + snapshot_properties={"wap.id": "etl-001"}, + ) + tbl = catalog.load_table(identifier) + staged = tbl.metadata.refs["audit"].snapshot_id + + # main advances independently, so this is a replay rather than a fast-forward + tbl.append(pa.Table.from_pylist([{"id": 9}], schema=arrow_schema)) + tbl = catalog.load_table(identifier) + assert sorted(tbl.scan().to_arrow().column("id").to_pylist()) == [1, 9] + + tbl.manage_snapshots().cherry_pick_snapshot(staged).commit() + + tbl = catalog.load_table(identifier) + assert sorted(tbl.scan().to_arrow().column("id").to_pylist()) == [1, 2, 3, 9] + published = tbl.current_snapshot() + assert published is not None and published.summary is not None + summary = published.summary.additional_properties + assert summary["source-snapshot-id"] == str(staged) + assert summary["published-wap-id"] == "etl-001" + + with pytest.raises(ValueError, match="Duplicate request to cherry pick wap id"): + tbl.manage_snapshots().cherry_pick_snapshot(staged).commit() + + catalog.drop_table(identifier) diff --git a/tests/table/test_manage_snapshots.py b/tests/table/test_manage_snapshots.py index 93301a01c7..307cfb4d83 100644 --- a/tests/table/test_manage_snapshots.py +++ b/tests/table/test_manage_snapshots.py @@ -17,8 +17,10 @@ from unittest.mock import MagicMock from uuid import uuid4 +import pyarrow as pa import pytest +from pyiceberg.catalog import Catalog from pyiceberg.table import CommitTableResponse, Table from pyiceberg.table.update import SetSnapshotRefUpdate, TableUpdate @@ -177,3 +179,172 @@ def test_set_current_snapshot_chained_with_create_tag(table_v2: Table) -> None: # The main branch should point to the same snapshot as the tag main_update = next(u for u in set_ref_updates if u.ref_name == "main") assert main_update.snapshot_id == snapshot_one + + +def _staged_wap_table( + catalog_with_warehouse: Catalog, + namespace: str, + advance_main: bool = True, + wap_id: str | None = None, +) -> tuple[Table, int]: + """Build the WAP shape: a staged branch write, with main optionally moved on since. + + Returns the reloaded table and the staged snapshot id. + """ + catalog_with_warehouse.create_namespace(namespace) + schema = pa.schema([pa.field("id", pa.int64())]) + tbl = catalog_with_warehouse.create_table(f"{namespace}.tbl", schema=schema) + + tbl.append(pa.table({"id": [1]}, schema=schema)) + tbl = catalog_with_warehouse.load_table(f"{namespace}.tbl") + + tbl.manage_snapshots().create_branch(snapshot_id=_current_snapshot_id(tbl), branch_name="audit").commit() + tbl = catalog_with_warehouse.load_table(f"{namespace}.tbl") + + tbl.append( + pa.table({"id": [2, 3]}, schema=schema), + branch="audit", + snapshot_properties={"wap.id": wap_id} if wap_id else {}, + ) + tbl = catalog_with_warehouse.load_table(f"{namespace}.tbl") + staged = tbl.metadata.refs["audit"].snapshot_id + + if advance_main: + tbl.append(pa.table({"id": [9]}, schema=schema)) + tbl = catalog_with_warehouse.load_table(f"{namespace}.tbl") + + return tbl, staged + + +def _ids(tbl: Table) -> list[int]: + return sorted(tbl.scan().to_arrow().column("id").to_pylist()) + + +def _current_summary_props(tbl: Table) -> dict[str, str]: + snapshot = tbl.current_snapshot() + assert snapshot is not None and snapshot.summary is not None + return snapshot.summary.additional_properties + + +def _current_snapshot_id(tbl: Table) -> int: + snapshot_id = tbl.metadata.current_snapshot_id + assert snapshot_id is not None + return snapshot_id + + +def test_cherry_pick_snapshot_replays_staged_append(catalog_with_warehouse: Catalog) -> None: + """The staged rows land on main even though main moved on after the branch was cut.""" + tbl, staged = _staged_wap_table(catalog_with_warehouse, "cp_replay") + assert _ids(tbl) == [1, 9] + + tbl.manage_snapshots().cherry_pick_snapshot(staged).commit() + + tbl = catalog_with_warehouse.load_table("cp_replay.tbl") + assert _ids(tbl) == [1, 2, 3, 9] + assert _current_summary_props(tbl)["source-snapshot-id"] == str(staged) + + +def test_cherry_pick_snapshot_records_published_wap_id(catalog_with_warehouse: Catalog) -> None: + tbl, staged = _staged_wap_table(catalog_with_warehouse, "cp_wap", wap_id="etl-001") + + tbl.manage_snapshots().cherry_pick_snapshot(staged).commit() + + tbl = catalog_with_warehouse.load_table("cp_wap.tbl") + assert _current_summary_props(tbl)["published-wap-id"] == "etl-001" + + +def test_cherry_pick_snapshot_rejects_duplicate_wap_publish(catalog_with_warehouse: Catalog) -> None: + """A wap id may only be published once.""" + tbl, staged = _staged_wap_table(catalog_with_warehouse, "cp_dup", wap_id="etl-001") + tbl.manage_snapshots().cherry_pick_snapshot(staged).commit() + + tbl = catalog_with_warehouse.load_table("cp_dup.tbl") + with pytest.raises(ValueError, match="Duplicate request to cherry pick wap id"): + tbl.manage_snapshots().cherry_pick_snapshot(staged).commit() + + +def test_cherry_pick_snapshot_replays_append_even_when_parent_is_current( + catalog_with_warehouse: Catalog, +) -> None: + """Appends are always replayed, so the wap trail is recorded even on an unmoved table.""" + tbl, staged = _staged_wap_table(catalog_with_warehouse, "cp_ff", advance_main=False, wap_id="etl-001") + assert _ids(tbl) == [1] + + tbl.manage_snapshots().cherry_pick_snapshot(staged).commit() + + tbl = catalog_with_warehouse.load_table("cp_ff.tbl") + assert _ids(tbl) == [1, 2, 3] + assert tbl.metadata.current_snapshot_id != staged + assert _current_summary_props(tbl)["published-wap-id"] == "etl-001" + + +def test_cherry_pick_snapshot_fast_forwards_non_append(catalog_with_warehouse: Catalog) -> None: + """A non-append whose parent is current cannot be replayed, so the ref advances to it.""" + catalog_with_warehouse.create_namespace("cp_ffna") + schema = pa.schema([pa.field("id", pa.int64())]) + tbl = catalog_with_warehouse.create_table("cp_ffna.tbl", schema=schema) + tbl.append(pa.table({"id": [1, 2]}, schema=schema)) + tbl = catalog_with_warehouse.load_table("cp_ffna.tbl") + + tbl.manage_snapshots().create_branch(snapshot_id=_current_snapshot_id(tbl), branch_name="audit").commit() + tbl = catalog_with_warehouse.load_table("cp_ffna.tbl") + tbl.delete("id = 1", branch="audit") + tbl = catalog_with_warehouse.load_table("cp_ffna.tbl") + staged = tbl.metadata.refs["audit"].snapshot_id + + tbl.manage_snapshots().cherry_pick_snapshot(staged).commit() + + tbl = catalog_with_warehouse.load_table("cp_ffna.tbl") + assert tbl.metadata.current_snapshot_id == staged + assert _ids(tbl) == [2] + + +def test_cherry_pick_snapshot_is_noop_for_ancestor(catalog_with_warehouse: Catalog) -> None: + """Picking a snapshot already in main's history changes nothing.""" + catalog_with_warehouse.create_namespace("cp_anc") + schema = pa.schema([pa.field("id", pa.int64())]) + tbl = catalog_with_warehouse.create_table("cp_anc.tbl", schema=schema) + tbl.append(pa.table({"id": [1]}, schema=schema)) + tbl = catalog_with_warehouse.load_table("cp_anc.tbl") + first = _current_snapshot_id(tbl) + tbl.append(pa.table({"id": [2]}, schema=schema)) + tbl = catalog_with_warehouse.load_table("cp_anc.tbl") + before = _current_snapshot_id(tbl) + + tbl.manage_snapshots().cherry_pick_snapshot(first).commit() + + tbl = catalog_with_warehouse.load_table("cp_anc.tbl") + assert tbl.metadata.current_snapshot_id == before + assert _ids(tbl) == [1, 2] + + +def test_cherry_pick_snapshot_rejects_unknown_snapshot(catalog_with_warehouse: Catalog) -> None: + catalog_with_warehouse.create_namespace("cp_unknown") + schema = pa.schema([pa.field("id", pa.int64())]) + tbl = catalog_with_warehouse.create_table("cp_unknown.tbl", schema=schema) + tbl.append(pa.table({"id": [1]}, schema=schema)) + tbl = catalog_with_warehouse.load_table("cp_unknown.tbl") + + with pytest.raises(ValueError, match="Cannot cherry-pick unknown snapshot id"): + tbl.manage_snapshots().cherry_pick_snapshot(1234567890).commit() + + +def test_cherry_pick_snapshot_rejects_non_append(catalog_with_warehouse: Catalog) -> None: + """Only append snapshots can be replayed; anything else must say so rather than silently skip.""" + catalog_with_warehouse.create_namespace("cp_nonappend") + schema = pa.schema([pa.field("id", pa.int64())]) + tbl = catalog_with_warehouse.create_table("cp_nonappend.tbl", schema=schema) + tbl.append(pa.table({"id": [1, 2]}, schema=schema)) + tbl = catalog_with_warehouse.load_table("cp_nonappend.tbl") + + tbl.manage_snapshots().create_branch(snapshot_id=_current_snapshot_id(tbl), branch_name="audit").commit() + tbl = catalog_with_warehouse.load_table("cp_nonappend.tbl") + tbl.delete("id = 1", branch="audit") + tbl = catalog_with_warehouse.load_table("cp_nonappend.tbl") + staged = tbl.metadata.refs["audit"].snapshot_id + + tbl.append(pa.table({"id": [9]}, schema=schema)) + tbl = catalog_with_warehouse.load_table("cp_nonappend.tbl") + + with pytest.raises(ValueError, match="not append, dynamic overwrite, or fast-forward"): + tbl.manage_snapshots().cherry_pick_snapshot(staged).commit()