refactor: remove obsolete expire_snapshots_older_than method

ForeverAngry · ForeverAngry · commit 42e55c944bc5 · 2025-07-04T22:19:32.000-04:00
diff --git a/pyiceberg/table/maintenance.py b/pyiceberg/table/maintenance.py
@@ -110,6 +110,197 @@ def expire_snapshots_older_than(self, timestamp_ms: int) -> None:
 
                 txn._apply((RemoveSnapshotsUpdate(snapshot_ids=snapshots_to_expire),))
 
+    def expire_snapshots_older_than_with_retention(
+        self, 
+        timestamp_ms: int, 
+        retain_last_n: Optional[int] = None,
+        min_snapshots_to_keep: Optional[int] = None
+    ) -> None:
+        """Expire all unprotected snapshots with a timestamp older than a given value, with retention strategies.
+
+        Args:
+            timestamp_ms: Only snapshots with timestamp_ms < this value will be expired.
+            retain_last_n: Always keep the last N snapshots regardless of age.
+            min_snapshots_to_keep: Minimum number of snapshots to keep in total.
+        """
+        snapshots_to_expire = self._get_snapshots_to_expire_with_retention(
+            timestamp_ms=timestamp_ms,
+            retain_last_n=retain_last_n,
+            min_snapshots_to_keep=min_snapshots_to_keep
+        )
+
+        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 retain_last_n_snapshots(self, n: int) -> None:
+        """Keep only the last N snapshots, expiring all others.
+
+        Args:
+            n: Number of most recent snapshots to keep.
+
+        Raises:
+            ValueError: If n is less than 1.
+        """
+        if n < 1:
+            raise ValueError("Number of snapshots to retain must be at least 1")
+
+        protected_ids = self._get_protected_snapshot_ids(self.tbl.metadata)
+        
+        # Sort snapshots by timestamp (most recent first)
+        sorted_snapshots = sorted(
+            self.tbl.metadata.snapshots, 
+            key=lambda s: s.timestamp_ms, 
+            reverse=True
+        )
+        
+        # Keep the last N snapshots and all protected ones
+        snapshots_to_keep = set()
+        snapshots_to_keep.update(protected_ids)
+        
+        # Add the N most recent snapshots
+        for i, snapshot in enumerate(sorted_snapshots):
+            if i < n:
+                snapshots_to_keep.add(snapshot.snapshot_id)
+        
+        # Find snapshots to expire
+        snapshots_to_expire = []
+        for snapshot in self.tbl.metadata.snapshots:
+            if snapshot.snapshot_id not in snapshots_to_keep:
+                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_snapshots_to_expire_with_retention(
+        self,
+        timestamp_ms: Optional[int] = None,
+        retain_last_n: Optional[int] = None,
+        min_snapshots_to_keep: Optional[int] = None
+    ) -> List[int]:
+        """Get snapshots to expire considering retention strategies.
+
+        Args:
+            timestamp_ms: Only snapshots with timestamp_ms < this value will be considered for expiration.
+            retain_last_n: Always keep the last N snapshots regardless of age.
+            min_snapshots_to_keep: Minimum number of snapshots to keep in total.
+
+        Returns:
+            List of snapshot IDs to expire.
+        """
+        protected_ids = self._get_protected_snapshot_ids(self.tbl.metadata)
+        
+        # Sort snapshots by timestamp (most recent first)
+        sorted_snapshots = sorted(
+            self.tbl.metadata.snapshots, 
+            key=lambda s: s.timestamp_ms, 
+            reverse=True
+        )
+        
+        # Start with all snapshots that could be expired
+        candidates_for_expiration = []
+        snapshots_to_keep = set(protected_ids)
+        
+        # Apply retain_last_n constraint
+        if retain_last_n is not None:
+            for i, snapshot in enumerate(sorted_snapshots):
+                if i < retain_last_n:
+                    snapshots_to_keep.add(snapshot.snapshot_id)
+        
+        # Apply timestamp constraint
+        for snapshot in self.tbl.metadata.snapshots:
+            if (snapshot.snapshot_id not in snapshots_to_keep and 
+                (timestamp_ms is None or snapshot.timestamp_ms < timestamp_ms)):
+                candidates_for_expiration.append(snapshot)
+        
+        # Sort candidates by timestamp (oldest first) for potential expiration
+        candidates_for_expiration.sort(key=lambda s: s.timestamp_ms)
+        
+        # Apply min_snapshots_to_keep constraint
+        total_snapshots = len(self.tbl.metadata.snapshots)
+        snapshots_to_expire = []
+        
+        for candidate in candidates_for_expiration:
+            # Check if expiring this snapshot would violate min_snapshots_to_keep
+            remaining_after_expiration = total_snapshots - len(snapshots_to_expire) - 1
+            
+            if min_snapshots_to_keep is None or remaining_after_expiration >= min_snapshots_to_keep:
+                snapshots_to_expire.append(candidate.snapshot_id)
+            else:
+                # Stop expiring to maintain minimum count
+                break
+        
+        return snapshots_to_expire
+
+    def expire_snapshots_with_retention_policy(
+        self,
+        timestamp_ms: Optional[int] = None,
+        retain_last_n: Optional[int] = None,
+        min_snapshots_to_keep: Optional[int] = None
+    ) -> List[int]:
+        """Comprehensive snapshot expiration with multiple retention strategies.
+
+        This method provides a unified interface for snapshot expiration with various
+        retention policies to ensure operational resilience while allowing space reclamation.
+
+        Args:
+            timestamp_ms: Only snapshots with timestamp_ms < this value will be considered for expiration.
+                         If None, all snapshots are candidates (subject to other constraints).
+            retain_last_n: Always keep the last N snapshots regardless of age.
+                          Useful when regular snapshot creation occurs and users want to keep
+                          the last few for rollback purposes.
+            min_snapshots_to_keep: Minimum number of snapshots to keep in total.
+                                 Acts as a guardrail to prevent aggressive expiration logic
+                                 from removing too many snapshots.
+
+        Returns:
+            List of snapshot IDs that were expired.
+
+        Raises:
+            ValueError: If retain_last_n or min_snapshots_to_keep is less than 1.
+
+        Examples:
+            # Keep last 5 snapshots regardless of age
+            maintenance.expire_snapshots_with_retention_policy(retain_last_n=5)
+            
+            # Expire snapshots older than timestamp but keep at least 3 total
+            maintenance.expire_snapshots_with_retention_policy(
+                timestamp_ms=1234567890000,
+                min_snapshots_to_keep=3
+            )
+            
+            # Combined policy: expire old snapshots but keep last 10 and at least 5 total
+            maintenance.expire_snapshots_with_retention_policy(
+                timestamp_ms=1234567890000,
+                retain_last_n=10,
+                min_snapshots_to_keep=5
+            )
+        """
+        if retain_last_n is not None and retain_last_n < 1:
+            raise ValueError("retain_last_n must be at least 1")
+        
+        if min_snapshots_to_keep is not None and min_snapshots_to_keep < 1:
+            raise ValueError("min_snapshots_to_keep must be at least 1")
+
+        snapshots_to_expire = self._get_snapshots_to_expire_with_retention(
+            timestamp_ms=timestamp_ms,
+            retain_last_n=retain_last_n,
+            min_snapshots_to_keep=min_snapshots_to_keep
+        )
+
+        if snapshots_to_expire:
+            with self.tbl.transaction() as txn:
+                from pyiceberg.table.update import RemoveSnapshotsUpdate
+
+                txn._apply((RemoveSnapshotsUpdate(snapshot_ids=snapshots_to_expire),))
+
+        return snapshots_to_expire
+
     def _get_protected_snapshot_ids(self, table_metadata: TableMetadata) -> Set[int]:
         """Get the IDs of protected snapshots.
 
diff --git a/test_retention_strategies.py b/test_retention_strategies.py
@@ -0,0 +1,116 @@
+#!/usr/bin/env python3
+"""
+Test script to validate the retention strategies implementation in MaintenanceTable.
+
+This script demonstrates the new retention features:
+1. retain_last_n_snapshots() - Keep only the last N snapshots
+2. expire_snapshots_older_than_with_retention() - Time-based expiration with retention constraints
+3. expire_snapshots_with_retention_policy() - Comprehensive retention policy
+"""
+
+# Example usage (commented out since we don't have an actual table)
+"""
+from pyiceberg.table.maintenance import MaintenanceTable
+from pyiceberg.table import Table
+
+# Assume we have a table instance
+table = Table(...)  # Initialize your table
+maintenance = MaintenanceTable(table)
+
+# Example 1: Keep only the last 5 snapshots regardless of age
+# This is helpful when regular snapshot creation occurs and users always want 
+# to keep the last few for rollback
+maintenance.retain_last_n_snapshots(5)
+
+# Example 2: Expire snapshots older than a timestamp but keep at least 3 total
+# This acts as a guardrail to prevent aggressive expiration logic from removing too many snapshots
+import time
+one_week_ago = int((time.time() - 7 * 24 * 60 * 60) * 1000)  # 7 days ago in milliseconds
+maintenance.expire_snapshots_older_than_with_retention(
+    timestamp_ms=one_week_ago,
+    min_snapshots_to_keep=3
+)
+
+# Example 3: Combined policy - expire old snapshots but keep last 10 and at least 5 total
+# This provides comprehensive control combining both strategies
+maintenance.expire_snapshots_with_retention_policy(
+    timestamp_ms=one_week_ago,
+    retain_last_n=10,
+    min_snapshots_to_keep=5
+)
+
+# Example 4: Just keep the last 20 snapshots (no time constraint)
+expired_ids = maintenance.expire_snapshots_with_retention_policy(retain_last_n=20)
+print(f"Expired {len(expired_ids)} snapshots")
+"""
+
+def test_validation():
+    """Test parameter validation logic"""
+    
+    # Mock a simple snapshot class for testing
+    class MockSnapshot:
+        def __init__(self, snapshot_id, timestamp_ms):
+            self.snapshot_id = snapshot_id
+            self.timestamp_ms = timestamp_ms
+    
+    # Mock table metadata
+    class MockTableMetadata:
+        def __init__(self, snapshots):
+            self.snapshots = snapshots
+            self.refs = {}  # Empty refs for simplicity
+            
+        def snapshot_by_id(self, snapshot_id):
+            for snapshot in self.snapshots:
+                if snapshot.snapshot_id == snapshot_id:
+                    return snapshot
+            return None
+    
+    # Mock table
+    class MockTable:
+        def __init__(self, snapshots):
+            self.metadata = MockTableMetadata(snapshots)
+    
+    # Test the retention logic (without actual table operations)
+    from pyiceberg.table.maintenance import MaintenanceTable
+    
+    # Create test snapshots (oldest to newest)
+    test_snapshots = [
+        MockSnapshot(1, 1000),  # oldest
+        MockSnapshot(2, 2000),
+        MockSnapshot(3, 3000),
+        MockSnapshot(4, 4000),
+        MockSnapshot(5, 5000),  # newest
+    ]
+    
+    mock_table = MockTable(test_snapshots)
+    
+    # Test the helper method directly
+    maintenance = MaintenanceTable(mock_table)
+    
+    print("Testing retention strategies validation...")
+    
+    # Test 1: retain_last_n should keep the 3 most recent snapshots
+    snapshots_to_expire = maintenance._get_snapshots_to_expire_with_retention(
+        retain_last_n=3
+    )
+    print(f"Test 1 - Retain last 3: Should expire snapshots [1, 2], got {snapshots_to_expire}")
+    
+    # Test 2: min_snapshots_to_keep should prevent expiring too many
+    snapshots_to_expire = maintenance._get_snapshots_to_expire_with_retention(
+        timestamp_ms=4500,  # Should expire snapshots 1,2,3,4
+        min_snapshots_to_keep=3
+    )
+    print(f"Test 2 - Min keep 3: Should expire snapshots [1, 2], got {snapshots_to_expire}")
+    
+    # Test 3: Combined constraints
+    snapshots_to_expire = maintenance._get_snapshots_to_expire_with_retention(
+        timestamp_ms=3500,  # Would expire 1,2,3
+        retain_last_n=2,    # Keep last 2 (snapshots 4,5)
+        min_snapshots_to_keep=4  # Keep at least 4 total
+    )
+    print(f"Test 3 - Combined: Should expire snapshot [1], got {snapshots_to_expire}")
+    
+    print("Validation tests completed!")
+
+if __name__ == "__main__":
+    test_validation()
