Added ExpireSnapshots Feature

pyiceberg/table/__init__.py

@@ -115,11 +115,7 @@
     update_table_metadata,
 )
 from pyiceberg.table.update.schema import UpdateSchema
-from pyiceberg.table.update.snapshot import (
-    ManageSnapshots,
-    UpdateSnapshot,
-    _FastAppendFiles,
-)
+from pyiceberg.table.update.snapshot import ExpireSnapshots, ManageSnapshots, UpdateSnapshot, _FastAppendFiles
 from pyiceberg.table.update.spec import UpdateSpec
 from pyiceberg.table.update.statistics import UpdateStatistics
 from pyiceberg.transforms import IdentityTransform
@@ -1079,6 +1075,15 @@ def manage_snapshots(self) -> ManageSnapshots:
         """
         return ManageSnapshots(transaction=Transaction(self, autocommit=True))
 
+    def expire_snapshots(self) -> ExpireSnapshots:
+        """
+        Shorthand to run expire snapshots by id or by a timestamp.
+
+        Use table.expire_snapshots().<operation>().commit() to run a specific operation.
+        Use table.expire_snapshots().<operation-one>().<operation-two>().commit() to run multiple operations.
+        """
+        return ExpireSnapshots(transaction=Transaction(self, autocommit=True))
+
     def update_statistics(self) -> UpdateStatistics:
         """
         Shorthand to run statistics management operations like add statistics and remove statistics.

pyiceberg/table/update/snapshot.py

@@ -55,6 +55,7 @@
 from pyiceberg.partitioning import (
     PartitionSpec,
 )
+from pyiceberg.table.refs import SnapshotRefType
 from pyiceberg.table.snapshots import (
     Operation,
     Snapshot,
@@ -66,6 +67,7 @@
     AddSnapshotUpdate,
     AssertRefSnapshotId,
     RemoveSnapshotRefUpdate,
+    RemoveSnapshotsUpdate,
     SetSnapshotRefUpdate,
     TableRequirement,
     TableUpdate,
@@ -739,6 +741,7 @@ class ManageSnapshots(UpdateTableMetadata["ManageSnapshots"]):
        ms.create_tag(snapshot_id1, "Tag_A").create_tag(snapshot_id2, "Tag_B")
     """
 
+    _snapshot_ids_to_expire: Set[int] = set()
     _updates: Tuple[TableUpdate, ...] = ()
     _requirements: Tuple[TableRequirement, ...] = ()
 
@@ -843,3 +846,69 @@ def remove_branch(self, branch_name: str) -> ManageSnapshots:
             This for method chaining
         """
         return self._remove_ref_snapshot(ref_name=branch_name)
+
+
+class ExpireSnapshots(UpdateTableMetadata["ExpireSnapshots"]):
+    """
+    Expire snapshots by ID.
+
+    Use table.expire_snapshots().<operation>().commit() to run a specific operation.
+    Use table.expire_snapshots().<operation-one>().<operation-two>().commit() to run multiple operations.
+    Pending changes are applied on commit.
+    """
+
+    _snapshot_ids_to_expire: Set[int] = set()
+    _updates: Tuple[TableUpdate, ...] = ()
+    _requirements: Tuple[TableRequirement, ...] = ()
+
+    def _commit(self) -> UpdatesAndRequirements:
+        """
+        Commit the staged updates and requirements.
+
+        This will remove the snapshots with the given IDs.
+
+        Returns:
+            Tuple of updates and requirements to be committed,
+            as required by the calling parent apply functions.
+        """
+        update = RemoveSnapshotsUpdate(snapshot_ids=self._snapshot_ids_to_expire)
+        self._updates += (update,)
+        return self._updates, self._requirements
+
+    def _get_protected_snapshot_ids(self) -> Set[int]:
+        """
+        Get the IDs of protected snapshots.
+
+        These are the HEAD snapshots of all branches and all tagged snapshots.  These ids are to be excluded from expiration.
+
+        Returns:
+            Set of protected snapshot IDs to exclude from expiration.
+        """
+        protected_ids: Set[int] = set()
+
+        for ref in self._transaction.table_metadata.refs.values():
+            if ref.snapshot_ref_type in [SnapshotRefType.TAG, SnapshotRefType.BRANCH]:
+                protected_ids.add(ref.snapshot_id)
+
+        return protected_ids
+
+    def expire_snapshot_by_id(self, snapshot_id: int) -> ExpireSnapshots:
+        """
+        Expire a snapshot by its ID.
+
+        This will mark the snapshot for expiration.
+
+        Args:
+            snapshot_id (int): The ID of the snapshot to expire.
+        Returns:
+            This for method chaining.
+        """
+        if self._transaction.table_metadata.snapshot_by_id(snapshot_id) is None:
+            raise ValueError(f"Snapshot with ID {snapshot_id} does not exist.")
+
+        if snapshot_id in self._get_protected_snapshot_ids():
+            raise ValueError(f"Snapshot with ID {snapshot_id} is protected and cannot be expired.")
+
+        self._snapshot_ids_to_expire.add(snapshot_id)
+
+        return self

tests/table/test_expire_snapshots.py

@@ -0,0 +1,119 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from unittest.mock import MagicMock
+from uuid import uuid4
+
+import pytest
+
+from pyiceberg.table import CommitTableResponse, Table
+
+
+def test_cannot_expire_protected_head_snapshot(table_v2: Table) -> None:
+    """Test that a HEAD (branch) snapshot cannot be expired."""
+    HEAD_SNAPSHOT = 3051729675574597004
+    KEEP_SNAPSHOT = 3055729675574597004
+
+    # Mock the catalog's commit_table method
+    table_v2.catalog = MagicMock()
+    # Simulate refs protecting HEAD_SNAPSHOT as a branch
+    table_v2.metadata = table_v2.metadata.model_copy(
+        update={
+            "refs": {
+                "main": MagicMock(snapshot_id=HEAD_SNAPSHOT, snapshot_ref_type="branch"),
+                "tag1": MagicMock(snapshot_id=KEEP_SNAPSHOT, snapshot_ref_type="tag"),
+            }
+        }
+    )
+    # Assert fixture data
+    assert any(ref.snapshot_id == HEAD_SNAPSHOT for ref in table_v2.metadata.refs.values())
+
+    # Attempt to expire the HEAD snapshot and expect a ValueError
+    with pytest.raises(ValueError, match=f"Snapshot with ID {HEAD_SNAPSHOT} is protected and cannot be expired."):
+        table_v2.expire_snapshots().expire_snapshot_by_id(HEAD_SNAPSHOT).commit()
+
+    table_v2.catalog.commit_table.assert_not_called()
+
+
+def test_cannot_expire_tagged_snapshot(table_v2: Table) -> None:
+    """Test that a tagged snapshot cannot be expired."""
+    TAGGED_SNAPSHOT = 3051729675574597004
+    KEEP_SNAPSHOT = 3055729675574597004
+
+    table_v2.catalog = MagicMock()
+    # Simulate refs protecting TAGGED_SNAPSHOT as a tag
+    table_v2.metadata = table_v2.metadata.model_copy(
+        update={
+            "refs": {
+                "tag1": MagicMock(snapshot_id=TAGGED_SNAPSHOT, snapshot_ref_type="tag"),
+                "main": MagicMock(snapshot_id=KEEP_SNAPSHOT, snapshot_ref_type="branch"),
+            }
+        }
+    )
+    assert any(ref.snapshot_id == TAGGED_SNAPSHOT for ref in table_v2.metadata.refs.values())
+
+    with pytest.raises(ValueError, match=f"Snapshot with ID {TAGGED_SNAPSHOT} is protected and cannot be expired."):
+        table_v2.expire_snapshots().expire_snapshot_by_id(TAGGED_SNAPSHOT).commit()
+
+    table_v2.catalog.commit_table.assert_not_called()
+
+
+def test_expire_unprotected_snapshot(table_v2: Table) -> None:
+    """Test that an unprotected snapshot can be expired."""
+    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
+
+    # Remove any refs that protect the snapshot 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"),
+            }
+        }
+    )
+
+    # Assert fixture data
+    assert all(ref.snapshot_id != EXPIRE_SNAPSHOT for ref in table_v2.metadata.refs.values())
+
+    # Expire the snapshot
+    table_v2.expire_snapshots().expire_snapshot_by_id(EXPIRE_SNAPSHOT).commit()
+
+    table_v2.catalog.commit_table.assert_called_once()
+    remaining_snapshots = table_v2.metadata.snapshots
+    assert EXPIRE_SNAPSHOT not in remaining_snapshots
+    assert len(table_v2.metadata.snapshots) == 1
+
+
+def test_expire_nonexistent_snapshot_raises(table_v2: Table) -> None:
+    """Test that trying to expire a non-existent snapshot raises an error."""
+    NONEXISTENT_SNAPSHOT = 9999999999999999999
+
+    table_v2.catalog = MagicMock()
+    table_v2.metadata = table_v2.metadata.model_copy(update={"refs": {}})
+
+    with pytest.raises(ValueError, match=f"Snapshot with ID {NONEXISTENT_SNAPSHOT} does not exist."):
+        table_v2.expire_snapshots().expire_snapshot_by_id(NONEXISTENT_SNAPSHOT).commit()
+
+    table_v2.catalog.commit_table.assert_not_called()