Added methods needed to expire snapshots by id, and optionally cleanup data

ForeverAngry · ForeverAngry · commit 5f0b62bbd647 · 2025-03-30T21:20:00.000-04:00
diff --git a/.gitignore b/.gitignore
@@ -50,3 +50,5 @@ htmlcov
 pyiceberg/avro/decoder_fast.c
 pyiceberg/avro/*.html
 pyiceberg/avro/*.so
+.vscode/settings.json
+pyiceberg/table/update/expire_snapshot.md
diff --git a/pyiceberg/table/update/snapshot.py b/pyiceberg/table/update/snapshot.py
@@ -247,7 +247,7 @@ def _summary(self, snapshot_properties: Dict[str, str] = EMPTY_DICT) -> Summary:
             truncate_full_table=self._operation == Operation.OVERWRITE,
         )
 
-    def _commit(self) -> UpdatesAndRequirements:
+    def commit(self, base_metadata: TableMetadata) -> UpdatesAndRequirements:
         new_manifests = self._manifests()
         next_sequence_number = self._transaction.table_metadata.next_sequence_number()
 
@@ -748,6 +748,8 @@ 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, ...] = ()
 
@@ -853,102 +855,80 @@ def remove_branch(self, branch_name: str) -> ManageSnapshots:
         """
         return self._remove_ref_snapshot(ref_name=branch_name)
 
-
-class ExpireSnapshots(UpdateTableMetadata["ExpireSnapshots"]):
-    def __init__(self, table: Table):
-        super().__init__(table)
-        self._expire_older_than: Optional[int] = None
-        self._snapshot_ids_to_expire: Set[int] = set()
-        self._retain_last: Optional[int] = None
-        self._delete_func: Optional[Callable[[str], None]] = None
-
-    def expire_older_than(self, timestamp_ms: int) -> "ExpireSnapshots":
-        """Expire snapshots older than the given timestamp."""
-        self._expire_older_than = timestamp_ms
-        return self
-
-    def expire_snapshot_id(self, snapshot_id: int) -> "ExpireSnapshots":
+    def _get_snapshot_ref_name(self, snapshot_id: int) -> Optional[str]:
+        """Get the reference name of a snapshot."""
+        for ref_name, snapshot in self._transaction.table_metadata.refs.items():
+            if snapshot.snapshot_id == snapshot_id:
+                return ref_name
+        return None
+
+    def _check_forward_ref(self, snapshot_id: int) -> bool:
+        """Check if the snapshot ID is a forward reference."""
+        # Ensure that remaining snapshots correctly reference their parent
+        for ref in self._transaction.table_metadata.refs.values():
+            if ref.snapshot_id == snapshot_id:
+                parent_snapshot_id = ref.parent_snapshot_id
+                if parent_snapshot_id is not None and parent_snapshot_id not in self._transaction.table_metadata.snapshots:
+                    return False
+        return True
+    
+    def _find_dependant_snapshot(self, snapshot_id: int) -> Optional[int]:
+        """Find any dependant snapshot."""
+        for ref in self._transaction.table_metadata.refs.values():
+            if ref.snapshot_id == snapshot_id:
+                return ref.parent_snapshot_id
+        return None
+
+    def exipre_snapshot_by_id(self, snapshot_id: int) -> ManageSnapshots:
         """Explicitly expire a snapshot by its ID."""
         self._snapshot_ids_to_expire.add(snapshot_id)
         return self
 
-    def retain_last(self, num_snapshots: int) -> "ExpireSnapshots":
-        """Retain the last N snapshots."""
-        if num_snapshots < 1:
-            raise ValueError("Number of snapshots to retain must be at least 1.")
-        self._retain_last = num_snapshots
-        return self
-
-    def delete_with(self, delete_func: Callable[[str], None]) -> "ExpireSnapshots":
-        """Set a custom delete function for cleaning up files."""
-        self._delete_func = delete_func
-        return self
-
-    def _commit(self, base_metadata: TableMetadata) -> UpdatesAndRequirements:
-        snapshots_to_expire = set()
-
-        # Identify snapshots by timestamp
-        if self._expire_older_than is not None:
-            snapshots_to_expire.update(
-                s.snapshot_id for s in base_metadata.snapshots
-                if s.timestamp_ms < self._expire_older_than
+    def expire_snapshots(self) -> ManageSnapshots:
+        """Expire the snapshots that are marked for expiration."""
+        # iterate over each snapshot requested to be expired
+        for snapshot_id in self._snapshot_ids_to_expire:
+            # remove the reference to the snapshot in the table metadata
+            # and stage the chagnes
+            update, requirement = self._remove_ref_snapshot(
+                ref_name=self._get_snapshot_ref_name(snapshot_id=snapshot_id),
             )
 
-        # Explicitly added snapshot IDs
-        snapshots_to_expire.update(self._snapshot_ids_to_expire)
-
-        # Retain the last N snapshots
-        if self._retain_last is not None:
-            sorted_snapshots = sorted(base_metadata.snapshots, key=lambda s: s.timestamp_ms, reverse=True)
-            retained_snapshots = {s.snapshot_id for s in sorted_snapshots[:self._retain_last]}
-            snapshots_to_expire.difference_update(retained_snapshots)
-
-        if not snapshots_to_expire:
-            print("No snapshots identified for expiration.")
-            return base_metadata  # No change, return original metadata
-
-        print(f"Expiring snapshots: {snapshots_to_expire}")
-
-        # Filter snapshots
-        remaining_snapshots = [
-            snapshot for snapshot in base_metadata.snapshots
-            if snapshot.snapshot_id not in snapshots_to_expire
-        ]
-
-        # Update snapshot log
-        remaining_snapshot_log = [
-            log for log in base_metadata.snapshot_log
-            if log.snapshot_id not in snapshots_to_expire
-        ]
-
-        # Determine the new current snapshot ID
-        new_current_snapshot_id = (
-            max(remaining_snapshots, key=lambda s: s.timestamp_ms).snapshot_id
-            if remaining_snapshots else None
-        )
-
-        # Return new metadata object reflecting the expired snapshots
-        updated_metadata = base_metadata.model_copy(
-            update={
-                "snapshots": remaining_snapshots,
-                "snapshot_log": remaining_snapshot_log,
-                "current_snapshot_id": new_current_snapshot_id
-            }
-        )
+            # return the updates and requirements to be committed
+            self._updates += update
+            self._requirements += requirement
+            
+            # check if there is a dependant snapshot
+            dependant_snapshot_id = self._find_dependant_snapshot(snapshot_id=snapshot_id)
+            if dependant_snapshot_id is not None:
+                # remove the reference to the dependant snapshot in the table metadata
+                # and stage the changes
+                update, requirement = self._transaction._set_ref_snapshot(
+                    ref_name=self._get_snapshot_ref_name(snapshot_id=dependant_snapshot_id),
+                    snapshot_id=dependant_snapshot_id
+                )
+                self._updates += update
+                self._requirements += requirement
 
-        # Cleanup orphaned files (manifests/data files)
-        self._cleanup_files(snapshots_to_expire, base_metadata)
+            # clean up the the unused files
 
-        return updated_metadata
+        return self
 
-    def _cleanup_files(self, expired_snapshot_ids: Set[int], metadata: TableMetadata):
+    def cleanup_files(self):
         """Remove files no longer referenced by any snapshots."""
-        print(f"Cleaning up resources for expired snapshots: {expired_snapshot_ids}")
-        if self._delete_func:
-            # Use the custom delete function if provided
-            for snapshot_id in expired_snapshot_ids:
-                self._delete_func(f"Snapshot {snapshot_id}")
-        else:
-            # Default cleanup logic (placeholder)
-            for snapshot_id in expired_snapshot_ids:
-                print(f"Default cleanup for snapshot {snapshot_id}")
+        # Remove the manifest files for the expired snapshots
+        for entry in self._snapshot_ids_to_expire:
+
+            # remove the manifest files for the expired snapshots
+            for manifest in self._transaction._table.snapshot_by_id(entry).manifests(self._transaction._table.io):
+                # get a list of all parquette files in the manifest that are orphaned
+                data_files = manifest.fetch_manifest_entry(io=self._transaction._table.io, discard_deleted=True)
+                
+                # remove the manfiest
+                self._transaction._table.io.delete(manifest.manifest_path)
+                
+                # remove the data files
+                [self._transaction._table.io.delete(file.data_file.file_path) for file in data_files if file.data_file.file_path is not None]
+        return self
+
+
diff --git a/tests/table/test_expire_snapshots.py b/tests/table/test_expire_snapshots.py
@@ -1,26 +1,31 @@
-# pylint:disable=redefined-outer-name,eval-used
+# pylint:disable=redefined-outer-name
+# pylint:disable=redefined-outer-name
 from unittest.mock import Mock
 import pytest
 
-from pyiceberg.manifest import DataFile, DataFileContent, ManifestContent, ManifestFile
-from pyiceberg.partitioning import PartitionField, PartitionSpec
-from pyiceberg.schema import Schema
 from pyiceberg.table import Table
-from pyiceberg.table.metadata import TableMetadata
-from pyiceberg.table.snapshots import Operation, Snapshot, SnapshotLogEntry, SnapshotSummaryCollector, Summary, update_snapshot_summaries
+from pyiceberg.table.metadata import new_table_metadata
+from pyiceberg.table.snapshots import Snapshot, SnapshotLogEntry
 from pyiceberg.table.update.snapshot import ExpireSnapshots
-from pyiceberg.transforms import IdentityTransform
-from pyiceberg.typedef import Record
-from pyiceberg.types import (
-    BooleanType,
-    IntegerType,
-    NestedField,
-    StringType,
-)
+
+from pyiceberg.schema import Schema
+from pyiceberg.partitioning import PartitionSpec
+from pyiceberg.table.sorting import SortOrder
+
+
 
 @pytest.fixture
 def mock_table():
-    """Fixture to create a mock table with metadata and snapshots."""
+    """
+    Creates a mock Iceberg table with predefined metadata, snapshots, and snapshot log entries.
+    The mock table includes:
+    - Snapshots with unique IDs, timestamps, and manifest lists.
+    - A snapshot log that tracks the history of snapshots with their IDs and timestamps.
+    - Table metadata including schema, partition spec, sort order, location, properties, and UUID.
+    - A current snapshot ID and last updated timestamp.
+    Returns:
+        Mock: A mock object representing an Iceberg table with the specified metadata and attributes.
+    """
     snapshots = [
         Snapshot(snapshot_id=1, timestamp_ms=1000, manifest_list="manifest1.avro"),
         Snapshot(snapshot_id=2, timestamp_ms=2000, manifest_list="manifest2.avro"),
@@ -31,56 +36,61 @@ def mock_table():
         SnapshotLogEntry(snapshot_id=2, timestamp_ms=2000),
         SnapshotLogEntry(snapshot_id=3, timestamp_ms=3000),
     ]
-    metadata = TableMetadata(
-        snapshots=snapshots,
-        snapshot_log=snapshot_log,
-        current_snapshot_id=3,
+
+    metadata = new_table_metadata(
+        schema=Schema(fields=[]),
+        partition_spec=PartitionSpec(spec_id=0, fields=[]),
+        sort_order=SortOrder(order_id=0, fields=[]),
+        location="s3://example-bucket/path/",
+        properties={},
+        table_uuid="12345678-1234-1234-1234-123456789abc",
+    ).model_copy(
+        update={
+            "snapshots": snapshots,
+            "snapshot_log": snapshot_log,
+            "current_snapshot_id": 3,
+            "last_updated_ms": 3000,
+        }
     )
+
     table = Mock(spec=Table)
     table.metadata = metadata
-    return table
-
-
-def test_expire_older_than(mock_table):
-    """Test expiring snapshots older than a given timestamp."""
-    expire_snapshots = ExpireSnapshots(mock_table)
-    expire_snapshots.expire_older_than(2500)._commit(mock_table.metadata)
-
-    remaining_snapshot_ids = {s.snapshot_id for s in mock_table.metadata.snapshots}
-    assert remaining_snapshot_ids == {3}, "Only the latest snapshot should remain."
+    table.identifier = ("db", "table") 
 
 
-def test_retain_last(mock_table):
-    """Test retaining the last N snapshots."""
-    expire_snapshots = ExpireSnapshots(mock_table)
-    expire_snapshots.retain_last(2)._commit(mock_table.metadata)
-
-    remaining_snapshot_ids = {s.snapshot_id for s in mock_table.metadata.snapshots}
-    assert remaining_snapshot_ids == {2, 3}, "The last two snapshots should remain."
-
-
-def test_expire_specific_snapshot(mock_table):
-    """Test explicitly expiring a specific snapshot."""
-    expire_snapshots = ExpireSnapshots(mock_table)
-    expire_snapshots.expire_snapshot_id(2)._commit(mock_table.metadata)
-
-    remaining_snapshot_ids = {s.snapshot_id for s in mock_table.metadata.snapshots}
-    assert remaining_snapshot_ids == {1, 3}, "Snapshot 2 should be expired."
-
-
-def test_custom_delete_function(mock_table):
-    """Test using a custom delete function for cleanup."""
-    delete_func = Mock()
-    expire_snapshots = ExpireSnapshots(mock_table)
-    expire_snapshots.expire_snapshot_id(1).delete_with(delete_func)._commit(mock_table.metadata)
-
-    delete_func.assert_called_once_with("Snapshot 1"), "Custom delete function should be called for expired snapshot."
-
-
-def test_no_snapshots_to_expire(mock_table):
-    """Test when no snapshots are identified for expiration."""
-    expire_snapshots = ExpireSnapshots(mock_table)
-    expire_snapshots.expire_older_than(500)._commit(mock_table.metadata)
+    return table
 
-    remaining_snapshot_ids = {s.snapshot_id for s in mock_table.metadata.snapshots}
-    assert remaining_snapshot_ids == {1, 2, 3}, "No snapshots should be expired."
+def test_expire_snapshots_removes_correct_snapshots(mock_table: Mock):
+    """
+    Test case for the `ExpireSnapshots` class to ensure that the correct snapshots
+    are removed and the delete function is called the expected number of times.
+    Args:
+        mock_table (Mock): A mock object representing the table.
+    Test Steps:
+    1. Create a mock delete function and a mock transaction.
+    2. Instantiate the `ExpireSnapshots` class with the mock transaction.
+    3. Configure the `ExpireSnapshots` instance to expire snapshots with IDs 1 and 2,
+       and set the delete function to the mock delete function.
+    4. Commit the changes using the `_commit` method with the mock table's metadata.
+    5. Validate that the mock delete function is called for the correct snapshots.
+    6. Verify that the delete function is called exactly twice.
+    7. Ensure that the updated metadata returned by `_commit` is not `None`.
+    """
+    mock_delete_func = Mock()
+    mock_transaction = Mock()
+
+    expire_snapshots = ExpireSnapshots(mock_transaction)
+    expire_snapshots \
+        .expire_snapshot_id(1) \
+        .expire_snapshot_id(2) \
+        .delete_with(mock_delete_func)
+
+    updated_metadata = expire_snapshots._commit(mock_table.metadata)
+
+    # Validate delete calls
+    mock_delete_func.assert_any_call(mock_table.return_value.snapshots[0])
+    mock_delete_func.assert_any_call(mock_table.metadata.snapshots[1])
+    assert mock_delete_func.call_count == 2
+
+    # Verify updated metadata returned
+    assert updated_metadata is not None
