Implement snapshot expiration with protection for branch/tag heads

Added ExpireSnapshots.expire_snapshots_older_than and expire_snapshots_by_ids methods to support expiring snapshots by timestamp and by multiple IDs.

Ensured that protected snapshots (branch/tag heads) cannot be expired, both at the API and commit stages.

Updated the expiration logic to always skip protected snapshots, even if they are accidentally included.

Added and fixed tests to verify that protected snapshots are never expired and that expiration works as expected for unprotected snapshots.

Improved test setup to accurately reflect post-commit metadata and to assert correct expiration behavior.

Author: ForeverAngry

pyiceberg/table/update/snapshot.py

@@ -864,12 +864,15 @@ def _commit(self) -> UpdatesAndRequirements:
         """
         Commit the staged updates and requirements.
 
-        This will remove the snapshots with the given IDs.
+        This will remove the snapshots with the given IDs, but will always skip protected snapshots (branch/tag heads).
 
         Returns:
             Tuple of updates and requirements to be committed,
             as required by the calling parent apply functions.
         """
+        # Remove any protected snapshot IDs from the set to expire, just in case
+        protected_ids = self._get_protected_snapshot_ids()
+        self._snapshot_ids_to_expire -= protected_ids
         update = RemoveSnapshotsUpdate(snapshot_ids=self._snapshot_ids_to_expire)
         self._updates += (update,)
         return self._updates, self._requirements
@@ -911,3 +914,34 @@ def expire_snapshot_by_id(self, snapshot_id: int) -> ExpireSnapshots:
         self._snapshot_ids_to_expire.add(snapshot_id)
 
         return self
+
+    def expire_snapshots_by_ids(self, snapshot_ids: List[int]) -> "ExpireSnapshots":
+        """
+        Expire multiple snapshots by their IDs.
+
+        This will mark the snapshots for expiration.
+
+        Args:
+            snapshot_ids (List[int]): List of snapshot IDs to expire.
+        Returns:
+            This for method chaining.
+        """
+        for snapshot_id in snapshot_ids:
+            self.expire_snapshot_by_id(snapshot_id)
+        return self
+
+    def expire_snapshots_older_than(self, timestamp_ms: int) -> "ExpireSnapshots":
+        """
+        Expire all unprotected snapshots with a timestamp older than a given value.
+
+        Args:
+            timestamp_ms (int): Only snapshots with timestamp_ms < this value will be expired.
+
+        Returns:
+            This for method chaining.
+        """
+        protected_ids = self._get_protected_snapshot_ids()
+        for snapshot in self._transaction.table_metadata.snapshots:
+            if snapshot.timestamp_ms < timestamp_ms and snapshot.snapshot_id not in protected_ids:
+                self._snapshot_ids_to_expire.add(snapshot.snapshot_id)
+        return self

tests/table/test_expire_snapshots.py

@@ -117,3 +117,108 @@ def test_expire_nonexistent_snapshot_raises(table_v2: Table) -> None:
         table_v2.expire_snapshots().expire_snapshot_by_id(NONEXISTENT_SNAPSHOT).commit()
 
     table_v2.catalog.commit_table.assert_not_called()
+
+
+def test_expire_snapshots_by_timestamp_skips_protected(table_v2: Table) -> None:
+    # Setup: two snapshots; both are old, but one is head/tag protected
+    HEAD_SNAPSHOT = 3051729675574597004
+    TAGGED_SNAPSHOT = 3055729675574597004
+
+    # Add snapshots to metadata for timestamp/protected test
+    from types import SimpleNamespace
+
+    table_v2.metadata = table_v2.metadata.model_copy(
+        update={
+            "refs": {
+                "main": MagicMock(snapshot_id=HEAD_SNAPSHOT, snapshot_ref_type="branch"),
+                "mytag": MagicMock(snapshot_id=TAGGED_SNAPSHOT, snapshot_ref_type="tag"),
+            },
+            "snapshots": [
+                SimpleNamespace(snapshot_id=HEAD_SNAPSHOT, timestamp_ms=1, parent_snapshot_id=None),
+                SimpleNamespace(snapshot_id=TAGGED_SNAPSHOT, timestamp_ms=1, parent_snapshot_id=None),
+            ],
+        }
+    )
+    table_v2.catalog = MagicMock()
+
+    # Attempt to expire all snapshots before a future timestamp (so both are candidates)
+    future_timestamp = 9999999999999  # Far in the future, after any real snapshot
+
+    # Mock the catalog's commit_table to return the current metadata (simulate no change)
+    mock_response = CommitTableResponse(
+        metadata=table_v2.metadata,  # protected snapshots remain
+        metadata_location="mock://metadata/location",
+        uuid=uuid4(),
+    )
+    table_v2.catalog.commit_table.return_value = mock_response
+
+    table_v2.expire_snapshots().expire_snapshots_older_than(future_timestamp).commit()
+    # Update metadata to reflect the commit (as in other tests)
+    table_v2.metadata = mock_response.metadata
+
+    # Both protected snapshots should remain
+    remaining_ids = {s.snapshot_id for s in table_v2.metadata.snapshots}
+    assert HEAD_SNAPSHOT in remaining_ids
+    assert TAGGED_SNAPSHOT in remaining_ids
+
+    # No snapshots should have been expired (commit_table called, but with empty snapshot_ids)
+    args, kwargs = table_v2.catalog.commit_table.call_args
+    updates = args[2] if len(args) > 2 else ()
+    # Find RemoveSnapshotsUpdate in updates
+    remove_update = next((u for u in updates if getattr(u, "action", None) == "remove-snapshots"), None)
+    assert remove_update is not None
+    assert remove_update.snapshot_ids == []
+
+
+def test_expire_snapshots_by_ids(table_v2: Table) -> None:
+    """Test that multiple unprotected snapshots can be expired by IDs."""
+    EXPIRE_SNAPSHOT_1 = 3051729675574597004
+    EXPIRE_SNAPSHOT_2 = 3051729675574597005
+    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
+
+    # Remove any refs that protect the snapshots to be expired
+    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"),
+            }
+        }
+    )
+
+    # Add snapshots to metadata for multi-id test
+    from types import SimpleNamespace
+
+    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"),
+            },
+            "snapshots": [
+                SimpleNamespace(snapshot_id=EXPIRE_SNAPSHOT_1, timestamp_ms=1, parent_snapshot_id=None),
+                SimpleNamespace(snapshot_id=EXPIRE_SNAPSHOT_2, timestamp_ms=1, parent_snapshot_id=None),
+                SimpleNamespace(snapshot_id=KEEP_SNAPSHOT, timestamp_ms=2, parent_snapshot_id=None),
+            ],
+        }
+    )
+
+    # Assert fixture data
+    assert all(ref.snapshot_id not in (EXPIRE_SNAPSHOT_1, EXPIRE_SNAPSHOT_2) for ref in table_v2.metadata.refs.values())
+
+    # Expire the snapshots
+    table_v2.expire_snapshots().expire_snapshots_by_ids([EXPIRE_SNAPSHOT_1, EXPIRE_SNAPSHOT_2]).commit()
+
+    table_v2.catalog.commit_table.assert_called_once()
+    remaining_snapshots = table_v2.metadata.snapshots
+    assert EXPIRE_SNAPSHOT_1 not in remaining_snapshots
+    assert EXPIRE_SNAPSHOT_2 not in remaining_snapshots
+    assert len(table_v2.metadata.snapshots) == 1