refactor: consolidate snapshot expiration into MaintenanceTable

ForeverAngry · ForeverAngry · commit 0c1e4c7bfc55 · 2025-06-23T18:40:07.000-04:00
- Move ExpireSnapshots functionality from standalone class to MaintenanceTable
- Replace fluent API (table.expire_snapshots().method().commit()) with direct execution (table.maintenance.method())
- Remove ExpireSnapshots class and integrate logic into maintenance operations
- Update all tests to use new unified maintenance API
- Maintain all existing validation and protection logic for snapshots

This change consolidates table maintenance operations under a single interface
and simplifies the API by removing the need for explicit commit calls.

Breaking change: table.expire_snapshots() API is replaced with table.maintenance.expire_*() methods
diff --git a/pyiceberg/table/__init__.py b/pyiceberg/table/__init__.py
@@ -116,7 +116,7 @@
     update_table_metadata,
 )
 from pyiceberg.table.update.schema import UpdateSchema
-from pyiceberg.table.update.snapshot import ExpireSnapshots, ManageSnapshots, UpdateSnapshot, _FastAppendFiles
+from pyiceberg.table.update.snapshot import ManageSnapshots, UpdateSnapshot, _FastAppendFiles
 from pyiceberg.table.update.spec import UpdateSpec
 from pyiceberg.table.update.statistics import UpdateStatistics
 from pyiceberg.transforms import IdentityTransform
@@ -1220,10 +1220,6 @@ 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."""
-        return ExpireSnapshots(transaction=Transaction(self, autocommit=True))
-
     def update_statistics(self) -> UpdateStatistics:
         """
         Shorthand to run statistics management operations like add statistics and remove statistics.
diff --git a/pyiceberg/table/maintenance.py b/pyiceberg/table/maintenance.py
@@ -19,7 +19,7 @@
 import logging
 from datetime import datetime, timedelta, timezone
 from functools import reduce
-from typing import TYPE_CHECKING, Optional, Set
+from typing import TYPE_CHECKING, List, Optional, Set
 
 from pyiceberg.utils.concurrent import ExecutorFactory
 
@@ -117,3 +117,88 @@ def _delete(file: str) -> None:
                     logger.warning(f"Files:\n{failed_to_delete_files}")
         else:
             logger.info(f"No orphaned files found at {location}!")
+
+    def expire_snapshot_by_id(self, snapshot_id: int) -> None:
+        """Expire a single snapshot by its ID.
+
+        Args:
+            snapshot_id: The ID of the snapshot to expire.
+
+        Raises:
+            ValueError: If the snapshot does not exist or is protected.
+        """
+        with self.tbl.transaction() as txn:
+            # Check if snapshot exists
+            if txn.table_metadata.snapshot_by_id(snapshot_id) is None:
+                raise ValueError(f"Snapshot with ID {snapshot_id} does not exist.")
+
+            # Check if snapshot is protected
+            protected_ids = self._get_protected_snapshot_ids(txn.table_metadata)
+            if snapshot_id in protected_ids:
+                raise ValueError(f"Snapshot with ID {snapshot_id} is protected and cannot be expired.")
+
+            # Remove the snapshot
+            from pyiceberg.table.update import RemoveSnapshotsUpdate
+            txn._apply((RemoveSnapshotsUpdate(snapshot_ids=[snapshot_id]),))
+
+    def expire_snapshots_by_ids(self, snapshot_ids: List[int]) -> None:
+        """Expire multiple snapshots by their IDs.
+
+        Args:
+            snapshot_ids: List of snapshot IDs to expire.
+
+        Raises:
+            ValueError: If any snapshot does not exist or is protected.
+        """
+        with self.tbl.transaction() as txn:
+            protected_ids = self._get_protected_snapshot_ids(txn.table_metadata)
+            
+            # Validate all snapshots before expiring any
+            for snapshot_id in snapshot_ids:
+                if txn.table_metadata.snapshot_by_id(snapshot_id) is None:
+                    raise ValueError(f"Snapshot with ID {snapshot_id} does not exist.")
+                if snapshot_id in protected_ids:
+                    raise ValueError(f"Snapshot with ID {snapshot_id} is protected and cannot be expired.")
+
+            # Remove all snapshots
+            from pyiceberg.table.update import RemoveSnapshotsUpdate
+            txn._apply((RemoveSnapshotsUpdate(snapshot_ids=snapshot_ids),))
+
+    def expire_snapshots_older_than(self, timestamp_ms: int) -> None:
+        """Expire all unprotected snapshots with a timestamp older than a given value.
+
+        Args:
+            timestamp_ms: Only snapshots with timestamp_ms < this value will be expired.
+        """
+        # First check if there are any snapshots to expire to avoid unnecessary transactions
+        protected_ids = self._get_protected_snapshot_ids(self.tbl.metadata)
+        snapshots_to_expire = []
+        
+        for snapshot in self.tbl.metadata.snapshots:
+            if snapshot.timestamp_ms < timestamp_ms and snapshot.snapshot_id not in protected_ids:
+                snapshots_to_expire.append(snapshot.snapshot_id)
+        
+        if snapshots_to_expire:
+            with self.tbl.transaction() as txn:
+                from pyiceberg.table.update import RemoveSnapshotsUpdate
+                txn._apply((RemoveSnapshotsUpdate(snapshot_ids=snapshots_to_expire),))
+
+    def _get_protected_snapshot_ids(self, table_metadata) -> 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.
+
+        Args:
+            table_metadata: The table metadata to check for protected snapshots.
+
+        Returns:
+            Set of protected snapshot IDs to exclude from expiration.
+        """
+        from pyiceberg.table.refs import SnapshotRefType
+        
+        protected_ids: Set[int] = set()
+        for ref in table_metadata.refs.values():
+            if ref.snapshot_ref_type in [SnapshotRefType.TAG, SnapshotRefType.BRANCH]:
+                protected_ids.add(ref.snapshot_id)
+        return protected_ids
diff --git a/pyiceberg/table/update/snapshot.py b/pyiceberg/table/update/snapshot.py
@@ -858,103 +858,3 @@ 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, 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
-
-    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
-
-    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
diff --git a/tests/table/test_expire_snapshots.py b/tests/table/test_expire_snapshots.py
@@ -43,7 +43,7 @@ def test_cannot_expire_protected_head_snapshot(table_v2: Table) -> None:
 
     # 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.maintenance.expire_snapshot_by_id(HEAD_SNAPSHOT)
 
     table_v2.catalog.commit_table.assert_not_called()
 
@@ -66,7 +66,7 @@ def test_cannot_expire_tagged_snapshot(table_v2: Table) -> None:
     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.maintenance.expire_snapshot_by_id(TAGGED_SNAPSHOT)
 
     table_v2.catalog.commit_table.assert_not_called()
 
@@ -98,9 +98,11 @@ def test_expire_unprotected_snapshot(table_v2: Table) -> None:
     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.maintenance.expire_snapshot_by_id(EXPIRE_SNAPSHOT)
 
     table_v2.catalog.commit_table.assert_called_once()
+    # Update metadata to reflect the commit
+    table_v2.metadata = mock_response.metadata
     remaining_snapshots = table_v2.metadata.snapshots
     assert EXPIRE_SNAPSHOT not in remaining_snapshots
     assert len(table_v2.metadata.snapshots) == 1
@@ -114,7 +116,7 @@ def test_expire_nonexistent_snapshot_raises(table_v2: Table) -> None:
     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.maintenance.expire_snapshot_by_id(NONEXISTENT_SNAPSHOT)
 
     table_v2.catalog.commit_table.assert_not_called()
 
@@ -152,7 +154,7 @@ def test_expire_snapshots_by_timestamp_skips_protected(table_v2: Table) -> None:
     )
     table_v2.catalog.commit_table.return_value = mock_response
 
-    table_v2.expire_snapshots().expire_snapshots_older_than(future_timestamp).commit()
+    table_v2.maintenance.expire_snapshots_older_than(future_timestamp)
     # Update metadata to reflect the commit (as in other tests)
     table_v2.metadata = mock_response.metadata
 
@@ -161,13 +163,8 @@ def test_expire_snapshots_by_timestamp_skips_protected(table_v2: Table) -> None:
     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 == []
+    # Since all snapshots were protected, commit_table should not be called
+    table_v2.catalog.commit_table.assert_not_called()
 
 
 def test_expire_snapshots_by_ids(table_v2: Table) -> None:
@@ -215,9 +212,11 @@ def test_expire_snapshots_by_ids(table_v2: Table) -> None:
     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.maintenance.expire_snapshots_by_ids([EXPIRE_SNAPSHOT_1, EXPIRE_SNAPSHOT_2])
 
     table_v2.catalog.commit_table.assert_called_once()
+    # Update metadata to reflect the commit
+    table_v2.metadata = mock_response.metadata
     remaining_snapshots = table_v2.metadata.snapshots
     assert EXPIRE_SNAPSHOT_1 not in remaining_snapshots
     assert EXPIRE_SNAPSHOT_2 not in remaining_snapshots
