Snapshots are not being transacted on, but need to re-assign refs

 ValueError: Cannot expire snapshot IDs {3051729675574597004} as they are currently referenced by table refs.

Author: ForeverAngry

pyiceberg/table/__init__.py

@@ -115,6 +115,7 @@
 )
 from pyiceberg.table.update.schema import UpdateSchema
 from pyiceberg.table.update.snapshot import (
+    ExpireSnapshots,
     ManageSnapshots,
     UpdateSnapshot,
     _FastAppendFiles,
@@ -1068,6 +1069,23 @@ def manage_snapshots(self) -> ManageSnapshots:
            ms.create_tag(snapshot_id1, "Tag_A").create_tag(snapshot_id2, "Tag_B")
         """
         return ManageSnapshots(transaction=Transaction(self, autocommit=True))
+    
+    def expire_snapshots(self) -> ExpireSnapshots:
+        """
+        Shorthand to expire snapshots.
+
+        Use table.expire_snapshots().expire_snapshot_id(...).commit() or
+        table.expire_snapshots().expire_older_than(...).commit()
+
+        You can also use it inside a transaction context:
+            with table.transaction() as tx:
+                tx.expire_snapshots().expire_older_than(...)
+
+        """
+        return ExpireSnapshots(Transaction(self, autocommit=True))
+
+
+
 
     def update_statistics(self) -> UpdateStatistics:
         """

pyiceberg/table/update/__init__.py

@@ -575,20 +575,6 @@ def _(update: RemoveStatisticsUpdate, base_metadata: TableMetadata, context: _Ta
 
     return base_metadata.model_copy(update={"statistics": statistics})
 
-@_apply_table_update.register(RemoveSnapshotsUpdate)
-def _(update: RemoveSnapshotsUpdate, base_metadata: TableMetadata, context: _TableMetadataUpdateContext) -> TableMetadata:
-    if len(update.snapshot_ids) == 0 or len(base_metadata.snapshots) == 0:
-        return base_metadata
-
-    retained_snapshots = []
-    ids_to_remove = set(update.snapshot_ids)
-    for snapshot in base_metadata.snapshots:
-        if snapshot.snapshot_id not in ids_to_remove:
-            retained_snapshots.append(snapshot)
-
-    context.add_update(update)
-    return base_metadata.model_copy(update={"snapshots": retained_snapshots})
-
 
 def update_table_metadata(
     base_metadata: TableMetadata,

pyiceberg/table/update/snapshot.py

@@ -92,6 +92,16 @@
 from typing import Optional, Set
 from datetime import datetime, timezone
 
+from typing import Dict, Optional, Set
+import uuid
+from pyiceberg.table.metadata import TableMetadata
+from pyiceberg.table.snapshots import Snapshot
+from pyiceberg.table.update import (
+    UpdateTableMetadata,
+    RemoveSnapshotsUpdate,
+    UpdatesAndRequirements,
+    AssertRefSnapshotId,
+)
 
 def _new_manifest_file_name(num: int, commit_uuid: uuid.UUID) -> str:
     return f"{commit_uuid}-m{num}.avro"
@@ -860,50 +870,56 @@ class ExpireSnapshots(UpdateTableMetadata["ExpireSnapshots"]):
     """
     API for removing old snapshots from the table.
     """
-
-    _ids_to_remove: List[int] = []
+    _updates: Tuple[TableUpdate, ...] = ()
+    _requirements: Tuple[TableRequirement, ...] = ()
 
     _updates: Tuple[TableUpdate, ...] = ()
     _requirements: Tuple[TableRequirement, ...] = ()
 
+    def __init__(self, transaction: Transaction) -> None:
+        super().__init__(transaction)
+        self._transaction = transaction
+        self._ids_to_remove: Set[int] = set()
+
     def _commit(self) -> UpdatesAndRequirements:
-        return (RemoveSnapshotsUpdate(snapshot_ids=self._ids_to_remove),), ()
-
-    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 _find_dependant_snapshot(self, snapshot_id: int) -> Optional[int]:
-        """Find any dependent snapshot."""
-        for ref in self._transaction.table_metadata.refs.values():
-            if ref.snapshot_id == snapshot_id:
-                return ref.parent_snapshot_id
-        return None
+        """Apply the pending changes and commit."""
+        if not hasattr(self, "_transaction") or not self._transaction:
+            raise AttributeError("Transaction object is not properly initialized.")
+
+        if not self._ids_to_remove:
+            raise ValueError("No snapshot IDs marked for expiration.")
+
+        # Ensure current snapshots in refs are not marked for removal
+        current_snapshot_ids = {ref.snapshot_id for ref in self._transaction.table_metadata.refs.values()}
+        conflicting_ids = self._ids_to_remove.intersection(current_snapshot_ids)
+        if conflicting_ids:
+            raise ValueError(f"Cannot expire snapshot IDs {conflicting_ids} as they are currently referenced by table refs.")
+
+        updates = (RemoveSnapshotsUpdate(snapshot_ids=list(self._ids_to_remove)),)
+
+        # Ensure refs haven't changed (snapshot ID consistency check)
+        requirements = tuple(
+            AssertRefSnapshotId(snapshot_id=ref.snapshot_id, ref=ref_name)
+            for ref_name, ref in self._transaction.table_metadata.refs.items()
+        )
+        self._updates += updates
+        self._requirements += requirements
+        return self
 
     def expire_snapshot_id(self, snapshot_id_to_expire: int) -> ExpireSnapshots:
         """Mark a specific snapshot ID for expiration."""
-        if self._transaction._table.snapshot_by_id(snapshot_id_to_expire):
-            self._ids_to_remove.append(snapshot_id_to_expire)
+        snapshot = self._transaction._table.snapshot_by_id(snapshot_id_to_expire)
+        if snapshot:
+            self._ids_to_remove.add(snapshot_id_to_expire)
+        else:
+            raise ValueError(f"Snapshot ID {snapshot_id_to_expire} does not exist.")
         return self
 
     def expire_older_than(self, timestamp_ms: int) -> ExpireSnapshots:
         """Mark snapshots older than the given timestamp for expiration."""
         for snapshot in self._transaction.table_metadata.snapshots:
             if snapshot.timestamp_ms < timestamp_ms:
-                self._ids_to_remove.append(snapshot.snapshot_id)
+                self._ids_to_remove.add(snapshot.snapshot_id)
         return self
 
-    # Uncomment and implement cleanup_files if file cleanup is required
-    # def cleanup_files(self):
-    #     """Remove files no longer referenced by any snapshots."""
-    #     for entry in self._ids_to_remove:
-    #         for manifest in self._transaction._table.snapshot_by_id(entry).manifests(self._transaction._table.io):
-    #             data_files = manifest.fetch_manifest_entry(io=self._transaction._table.io, discard_deleted=True)
-    #             self._transaction._table.io.delete(manifest.manifest_path)
-    #             [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
-
 

tests/table/test_expire_snapshots.py

@@ -5,7 +5,6 @@
 from typing import Any, Dict, Optional
 import pytest
 from pyiceberg.catalog.memory import InMemoryCatalog
-from pyiceberg.catalog.noop import NoopCatalog
 from pyiceberg.io import load_file_io
 from pyiceberg.table import Table
 from pyiceberg.table.sorting import NullOrder, SortDirection, SortField, SortOrder
@@ -17,9 +16,46 @@
 from pyiceberg.schema import Schema
 from pyiceberg.types import NestedField, LongType, StringType
 from pyiceberg.table.snapshots import Snapshot
-from pyiceberg.table.metadata import TableMetadata, TableMetadataV2, new_table_metadata
+from pyiceberg.table.metadata import TableMetadata, TableMetadataUtil, TableMetadataV2, new_table_metadata
 
 
+@pytest.fixture
+def mock_table():
+    """Fixture to create a mock Table instance with proper metadata for testing."""
+    # Create mock metadata with empty snapshots list
+    metadata_dict = {
+        "format-version": 2,
+        "table-uuid": "9c12d441-03fe-4693-9a96-a0705ddf69c1",
+        "location": "s3://bucket/test/location",
+        "last-sequence-number": 0,
+        "last-updated-ms": int(time.time() * 1000),
+        "last-column-id": 3,
+        "current-schema-id": 1,
+        "schemas": [
+            {
+                "type": "struct",
+                "schema-id": 1,
+                "fields": [{"id": 1, "name": "x", "required": True, "type": "long"}]
+            }
+        ],
+        "default-spec-id": 0,
+        "partition-specs": [{"spec-id": 0, "fields": []}],
+        "last-partition-id": 0,
+        "default-sort-order-id": 0,
+        "sort-orders": [{"order-id": 0, "fields": []}],
+        "snapshots": [],
+        "refs": {},
+    }
+    
+    metadata = TableMetadataUtil.parse_obj(metadata_dict)
+    
+    return Table(
+        identifier=("mock_database", "mock_table"),
+        metadata=metadata,
+        metadata_location="mock_location",
+        io=load_file_io(),
+        catalog=InMemoryCatalog("InMemoryCatalog"),
+    )
 
 
 @pytest.fixture
@@ -43,19 +79,24 @@ def generate_snapshot(
     snapshot_log = []
     initial_snapshot_id = 3051729675574597004
 
-    for i in range(2000):
+    for i in range(5):
         snapshot_id = initial_snapshot_id + i
         parent_snapshot_id = snapshot_id - 1 if i > 0 else None
         timestamp_ms = int(time.time() * 1000) - randint(0, 1000000)
         snapshots.append(generate_snapshot(snapshot_id, parent_snapshot_id, timestamp_ms, i))
         snapshot_log.append({"snapshot-id": snapshot_id, "timestamp-ms": timestamp_ms})
 
-    metadata = {
+    metadata_dict = {
         "format-version": 2,
         "table-uuid": "9c12d441-03fe-4693-9a96-a0705ddf69c1",
         "location": "s3://bucket/test/location",
         "last-sequence-number": 34,
-        "last-updated-ms": 1602638573590,
+        "last-updated-ms": snapshots[-1]["timestamp-ms"],
+        "metadata-log": [
+            {"metadata-file": "s3://bucket/test/location/metadata/v1.json", "timestamp-ms": 1700000000000},
+            {"metadata-file": "s3://bucket/test/location/metadata/v2.json", "timestamp-ms": 1700003600000},
+            {"metadata-file": "s3://bucket/test/location/metadata/v3.json", "timestamp-ms": snapshots[-1]["timestamp-ms"]},
+        ],
         "last-column-id": 3,
         "current-schema-id": 1,
         "schemas": [
@@ -72,46 +113,41 @@ def generate_snapshot(
             },
         ],
         "default-spec-id": 0,
-        "partition-specs": [{"spec-id": 0, "fields": [{"name": "x", "transform": "identity", "source-id": 1, "field-id": 1000}]}],
+        "partition-specs": [{"spec-id": 0, "fields": []}],
         "last-partition-id": 1000,
         "default-sort-order-id": 3,
-        "sort-orders": [
-            {
-                "order-id": 3,
-                "fields": [
-                    {"transform": "identity", "source-id": 2, "direction": "asc", "null-order": "nulls-first"},
-                    {"transform": "bucket[4]", "source-id": 3, "direction": "desc", "null-order": "nulls-last"},
-                ],
-            }
-        ],
+        "sort-orders": [{"order-id": 3, "fields": []}],
         "properties": {"read.split.target.size": "134217728"},
-        "current-snapshot-id": initial_snapshot_id + 1999,
+        "current-snapshot-id": initial_snapshot_id + 4,
         "snapshots": snapshots,
         "snapshot-log": snapshot_log,
-        "metadata-log": [{"metadata-file": "s3://bucket/.../v1.json", "timestamp-ms": 1515100}],
         "refs": {"test": {"snapshot-id": initial_snapshot_id, "type": "tag", "max-ref-age-ms": 10000000}},
     }
 
+    metadata = TableMetadataUtil.parse_obj(metadata_dict)
+
     return Table(
         identifier=("database", "table"),
         metadata=metadata,
-        metadata_location=f"{metadata['location']}/uuid.metadata.json",
+        metadata_location=f"{metadata.location}/uuid.metadata.json",
         io=load_file_io(),
-        catalog=NoopCatalog("NoopCatalog"),
+        catalog=InMemoryCatalog("InMemoryCatalog"),
     )
 
 
 
-def test_expire_snapshots_removes_correct_snapshots(generate_test_table):
+def test_expire_snapshots_removes_correct_snapshots(generate_test_table: Table):
     """
     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.
     """
-
+    
     # Use the fixture-provided table
-    with ExpireSnapshots(generate_test_table.transaction()) as manage_snapshots:
-        manage_snapshots.expire_snapshot_id(3051729675574597004)
+    with generate_test_table.expire_snapshots() as transaction:
+        transaction.expire_snapshot_id(3051729675574597004).commit()
 
     # Check the remaining snapshots
     remaining_snapshot_ids = {snapshot.snapshot_id for snapshot in generate_test_table.metadata.snapshots}
-    assert not remaining_snapshot_ids.issubset({3051729675574597004})
+    
+    # Assert that the expired snapshot ID is not in the remaining snapshots
+    assert 3051729675574597004 not in remaining_snapshot_ids