From 5ec21bd05a453b71a838dc0f8e94ac2a807e0acf Mon Sep 17 00:00:00 2001 From: Mansi Dhruv Date: Mon, 2 Nov 2026 14:19:02 +0530 Subject: [PATCH 1/2] Fix: assert ref snapshot id requirements during snapshot expiration (#4200) Signed-off-by: Mansi Dhruv --- pyiceberg/table/update/snapshot.py | 2 + tests/table/test_expire_snapshots.py | 63 +++++++++++++++++++++++++++- 2 files changed, 64 insertions(+), 1 deletion(-) diff --git a/pyiceberg/table/update/snapshot.py b/pyiceberg/table/update/snapshot.py index 8feebb9901..5698b58ccc 100644 --- a/pyiceberg/table/update/snapshot.py +++ b/pyiceberg/table/update/snapshot.py @@ -1357,6 +1357,8 @@ def _commit(self) -> UpdatesAndRequirements: self._snapshot_ids_to_expire -= protected_ids update = RemoveSnapshotsUpdate(snapshot_ids=self._snapshot_ids_to_expire) self._updates += (update,) + for ref_name, ref in self._transaction.table_metadata.refs.items(): + self._requirements += (AssertRefSnapshotId(ref=ref_name, snapshot_id=ref.snapshot_id),) return self._updates, self._requirements def _get_protected_snapshot_ids(self) -> set[int]: diff --git a/tests/table/test_expire_snapshots.py b/tests/table/test_expire_snapshots.py index 106e5b786c..eb08c7959c 100644 --- a/tests/table/test_expire_snapshots.py +++ b/tests/table/test_expire_snapshots.py @@ -22,7 +22,7 @@ import pytest from pyiceberg.table import CommitTableResponse, Table -from pyiceberg.table.update import RemoveSnapshotsUpdate, update_table_metadata +from pyiceberg.table.update import AssertRefSnapshotId, RemoveSnapshotsUpdate, update_table_metadata from pyiceberg.table.update.snapshot import ExpireSnapshots @@ -316,3 +316,64 @@ def test_update_remove_snapshots_with_statistics(table_v2_with_statistics: Table assert not any(stat.snapshot_id == REMOVE_SNAPSHOT for stat in new_metadata.statistics), ( "Statistics for removed snapshot should be gone" ) + + +def test_expire_snapshots_asserts_ref_snapshot_id_requirements(table_v2: Table) -> None: + """Test that expiring snapshots stages AssertRefSnapshotId for protected refs.""" + EXPIRE_SNAPSHOT = 3051729675574597004 + KEEP_SNAPSHOT = 3055729675574597004 + + mock_response = CommitTableResponse( + metadata=table_v2.metadata.model_copy(update={"snapshots": [KEEP_SNAPSHOT]}), + metadata_location="mock://metadata/location", + uuid=uuid4(), + ) + table_v2.catalog = MagicMock() + table_v2.catalog.commit_table.return_value = mock_response + + table_v2.metadata = table_v2.metadata.model_copy( + update={ + "refs": { + "main": MagicMock(snapshot_id=KEEP_SNAPSHOT, snapshot_ref_type="branch"), + "tag1": MagicMock(snapshot_id=KEEP_SNAPSHOT, snapshot_ref_type="tag"), + } + } + ) + + table_v2.maintenance.expire_snapshots().by_id(EXPIRE_SNAPSHOT).commit() + + table_v2.catalog.commit_table.assert_called_once() + args, _ = table_v2.catalog.commit_table.call_args + requirements = args[1] if len(args) > 1 else () + ref_reqs = [r for r in requirements if isinstance(r, AssertRefSnapshotId)] + assert {r.ref: r.snapshot_id for r in ref_reqs} == { + "main": KEEP_SNAPSHOT, + "tag1": KEEP_SNAPSHOT, + } + + +def test_expire_snapshots_concurrent_rollback_fails(tmp_path_factory: pytest.TempPathFactory) -> None: + """Test that expiring snapshots fails when a protected reference changes concurrently (#4200).""" + import pyarrow as pa + + from pyiceberg.catalog.sql import SqlCatalog + from pyiceberg.exceptions import CommitFailedException + + tmp_dir = tmp_path_factory.mktemp("iceberg_expire") + warehouse = f"file://{tmp_dir.as_posix()}" + cat = SqlCatalog("c", uri="sqlite:///:memory:", warehouse=warehouse) + cat.create_namespace("ns") + t = cat.create_table("ns.t", schema=pa.schema([("x", pa.int64())])) + t.append(pa.table({"x": [1]})) + s1 = t.current_snapshot().snapshot_id + t.append(pa.table({"x": [2]})) + + cat.load_table("ns.t").manage_snapshots().rollback_to_snapshot(s1).commit() + + with pytest.raises(CommitFailedException, match="Requirement failed: branch main has changed"): + t.maintenance.expire_snapshots().by_id(s1).commit() + + reloaded = cat.load_table("ns.t") + assert "main" in reloaded.metadata.refs + assert reloaded.metadata.refs["main"].snapshot_id == s1 + assert reloaded.scan().to_arrow()["x"].to_pylist() == [1] From a2e0cfebe0112c6b22ddc968ab7ca8d101d8337b Mon Sep 17 00:00:00 2001 From: Mansi Dhruv Date: Mon, 2 Nov 2026 14:53:46 +0530 Subject: [PATCH 2/2] test: assert current_snapshot is not None for mypy narrowing Signed-off-by: Mansi Dhruv --- tests/table/test_expire_snapshots.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/table/test_expire_snapshots.py b/tests/table/test_expire_snapshots.py index eb08c7959c..9bfcb1deda 100644 --- a/tests/table/test_expire_snapshots.py +++ b/tests/table/test_expire_snapshots.py @@ -365,7 +365,9 @@ def test_expire_snapshots_concurrent_rollback_fails(tmp_path_factory: pytest.Tem cat.create_namespace("ns") t = cat.create_table("ns.t", schema=pa.schema([("x", pa.int64())])) t.append(pa.table({"x": [1]})) - s1 = t.current_snapshot().snapshot_id + current_snapshot = t.current_snapshot() + assert current_snapshot is not None + s1 = current_snapshot.snapshot_id t.append(pa.table({"x": [2]})) cat.load_table("ns.t").manage_snapshots().rollback_to_snapshot(s1).commit()