Fix snapshot history after expiring snapshots

pyiceberg/table/update/__init__.py

@@ -544,11 +544,14 @@ def _(update: RemoveSnapshotsUpdate, base_metadata: TableMetadata, context: _Tab
         for snapshot in base_metadata.snapshots
         if snapshot.snapshot_id not in update.snapshot_ids
     ]
-    snapshot_log = [
-        snapshot_log_entry
-        for snapshot_log_entry in base_metadata.snapshot_log
-        if snapshot_log_entry.snapshot_id not in update.snapshot_ids
-    ]
+    expired_ids = set(update.snapshot_ids)
+    snapshot_log: list[SnapshotLogEntry] = []
+    for entry in base_metadata.snapshot_log:
+        if entry.snapshot_id in expired_ids:
+            # Earlier entries would make time travel resolve across an expired gap.
+            snapshot_log.clear()
+        else:
+            snapshot_log.append(entry)
 
     remove_ref_updates = (
         RemoveSnapshotRefUpdate(ref_name=ref_name)

tests/table/test_expire_snapshots.py

@@ -14,18 +14,69 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+import json
 import threading
-from datetime import datetime, timedelta
+from datetime import UTC, datetime, timedelta
+from types import SimpleNamespace
 from unittest.mock import MagicMock, Mock
 from uuid import uuid4
 
 import pytest
 
+from pyiceberg.catalog import Catalog
 from pyiceberg.table import CommitTableResponse, Table
 from pyiceberg.table.update import RemoveSnapshotsUpdate, update_table_metadata
 from pyiceberg.table.update.snapshot import ExpireSnapshots
 
 
+@pytest.mark.parametrize("expiration", ["by_id", "older_than"])
+def test_expiration_does_not_substitute_older_history(catalog: Catalog, monkeypatch: pytest.MonkeyPatch, expiration: str) -> None:
+    import pyarrow as pa
+
+    catalog.create_namespace("history")
+    table = catalog.create_table("history.events", schema=pa.schema([("x", pa.int64())]))
+    # Keep snapshot and metadata timestamps consistent; commits and file IO remain real.
+    clock_ms = table.metadata.last_updated_ms
+    monkeypatch.setattr("pyiceberg.table.snapshots.time", SimpleNamespace(time=lambda: clock_ms / 1000))
+    monkeypatch.setattr(
+        "pyiceberg.table.update.datetime", SimpleNamespace(now=lambda: datetime.fromtimestamp(clock_ms / 1000, UTC))
+    )
+    for value in (1, 2, 3):
+        clock_ms += 1000
+        table.append(pa.table({"x": [value]}))
+    first, middle, latest = table.history()
+    assert first.timestamp_ms < middle.timestamp_ms < latest.timestamp_ms
+    timestamp_ms = latest.timestamp_ms - 1
+    before = table.snapshot_as_of_timestamp(timestamp_ms)
+    assert before is not None and before.snapshot_id == middle.snapshot_id
+
+    clock_ms += 1000
+    table.manage_snapshots().create_tag(first.snapshot_id, "keep_first").commit()
+    clock_ms += 1000
+    if expiration == "by_id":
+        table.maintenance.expire_snapshots().by_id(middle.snapshot_id).commit()
+    else:
+        cutoff = datetime.fromtimestamp(latest.timestamp_ms / 1000, UTC)
+        table.maintenance.expire_snapshots().older_than(cutoff).commit()
+
+    reloaded = catalog.load_table("history.events")
+    for actual in (table, reloaded):
+        assert actual.metadata.last_updated_ms == clock_ms
+        assert actual.history() == [latest]
+        assert actual.snapshot_as_of_timestamp(timestamp_ms) is None
+        assert actual.snapshot_as_of_timestamp(latest.timestamp_ms, inclusive=False) is None
+        at_latest = actual.snapshot_as_of_timestamp(latest.timestamp_ms)
+        assert at_latest is not None and at_latest.snapshot_id == latest.snapshot_id
+        assert actual.snapshot_by_id(middle.snapshot_id) is None
+        assert actual.refs()["keep_first"].snapshot_id == first.snapshot_id
+        assert actual.scan(snapshot_id=first.snapshot_id).to_arrow()["x"].to_pylist() == [1]
+        assert sorted(actual.scan().to_arrow()["x"].to_pylist()) == [1, 2, 3]
+
+    with reloaded.io.new_input(reloaded.metadata_location).open() as stream:
+        persisted_metadata = json.load(stream)
+    assert persisted_metadata["snapshot-log"] == [{"timestamp-ms": latest.timestamp_ms, "snapshot-id": latest.snapshot_id}]
+
+
 def test_cannot_expire_protected_head_snapshot(table_v2: Table) -> None:
     """Test that a HEAD (branch) snapshot cannot be expired."""
     HEAD_SNAPSHOT = 3051729675574597004

tests/table/test_init.py

@@ -1186,6 +1186,46 @@ def test_update_remove_snapshots_doesnt_exist(table_v2: Table) -> None:
         update_table_metadata(table_v2.metadata, (update,))
 
 
+@pytest.mark.parametrize("table", [lf("table_v1"), lf("table_v2"), lf("table_v3")])
+def test_remove_snapshots_retains_contiguous_history(table: Table) -> None:
+    snapshots = [
+        Snapshot(snapshot_id=snapshot_id, timestamp_ms=snapshot_id, manifest_list=f"file:///manifest-{snapshot_id}.avro")
+        for snapshot_id in (10, 20, 30, 40)
+    ]
+    cases: list[tuple[str, list[int], list[int], list[int]]] = [
+        ("middle", [10, 20, 30, 40], [20], [30, 40]),
+        ("prefix", [10, 20, 30, 40], [10], [20, 30, 40]),
+        ("separated-gaps", [10, 20, 30, 40], [10, 30], [40]),
+        ("adjacent-gaps", [10, 20, 30, 40], [20, 30], [40]),
+        ("repeated-expired-snapshot", [10, 20, 30, 20, 40], [20], [40]),
+        ("rollback-to-retained-snapshot", [10, 20, 10, 30, 40], [20], [10, 30, 40]),
+        ("no-expiration", [10, 20, 30, 40], [], [10, 20, 30, 40]),
+        ("snapshot-outside-main-history", [10, 30, 40], [20], [10, 30, 40]),
+        ("empty-history", [], [20], []),
+        ("last-entry", [10, 20, 30, 40], [40], []),
+    ]
+    for case, history_ids, expired_ids, expected_ids in cases:
+        history = [
+            SnapshotLogEntry(snapshot_id=snapshot_id, timestamp_ms=100 + index) for index, snapshot_id in enumerate(history_ids)
+        ]
+        original = table.metadata.model_copy(
+            update={
+                "snapshots": snapshots,
+                "snapshot_log": history,
+                "current_snapshot_id": 40,
+                "refs": {"main": SnapshotRef(snapshot_id=40, snapshot_ref_type="branch")},
+            }
+        )
+        updated = update_table_metadata(original, (RemoveSnapshotsUpdate(snapshot_ids=expired_ids),), enforce_validation=True)
+
+        assert [entry.snapshot_id for entry in updated.snapshot_log] == expected_ids, case
+        expected_history = history[-len(expected_ids) :] if expected_ids else []
+        assert updated.snapshot_log == expected_history, case
+        assert {snapshot.snapshot_id for snapshot in updated.snapshots} == {10, 20, 30, 40} - set(expired_ids), case
+        assert original.snapshot_log == history, case
+        assert [entry.snapshot_id for entry in original.snapshot_log] == history_ids, case
+
+
 def test_update_remove_snapshots_remove_current_snapshot_id(table_v2: Table) -> None:
     update = RemoveSnapshotsUpdate(snapshot_ids=[3055729675574597004])
     new_metadata = update_table_metadata(table_v2.metadata, (update,))