Added the builder method to __init__.py, updated the snapshot api with a new Expired Snapshot class. updated tests.

ForeverAngry · ForeverAngry · commit 65365e1d3ebe · 2025-04-01T00:27:36.000-04:00
diff --git a/pyiceberg/table/update/__init__.py b/pyiceberg/table/update/__init__.py
@@ -575,6 +575,20 @@ 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,
diff --git a/pyiceberg/table/update/note.md b/pyiceberg/table/update/note.md
@@ -0,0 +1,2 @@
+in the snapshot.py class, you define the "api" or logic to collect the changes and then stage, them.  Then the
+__init__.py has a decorator that calls the type to actually apply the metadata changes.
diff --git a/pyiceberg/table/update/snapshot.py b/pyiceberg/table/update/snapshot.py
@@ -66,6 +66,7 @@
     AddSnapshotUpdate,
     AssertRefSnapshotId,
     RemoveSnapshotRefUpdate,
+    RemoveSnapshotsUpdate,
     SetSnapshotRefUpdate,
     TableRequirement,
     TableMetadata,
@@ -84,9 +85,8 @@
 from pyiceberg.utils.properties import property_as_bool, property_as_int
 
 if TYPE_CHECKING:
-    from pyiceberg.table import Transaction
+    from pyiceberg.table import Table
 
-from pyiceberg.table import Table
 from pyiceberg.table.metadata import Snapshot
 from pyiceberg.table.update import UpdateTableMetadata
 from typing import Optional, Set
@@ -748,13 +748,13 @@ 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, ...] = ()
 
     def _commit(self) -> UpdatesAndRequirements:
         """Apply the pending changes and commit."""
+        if not hasattr(self._transaction, "_apply"):
+            raise AttributeError("Transaction object is not properly initialized.")
         return self._updates, self._requirements
 
     def _remove_ref_snapshot(self, ref_name: str) -> ManageSnapshots:
@@ -855,80 +855,55 @@ def remove_branch(self, branch_name: str) -> ManageSnapshots:
         """
         return self._remove_ref_snapshot(ref_name=branch_name)
 
+
+class ExpireSnapshots(UpdateTableMetadata["ExpireSnapshots"]):
+    """
+    API for removing old snapshots from the table.
+    """
+
+    _ids_to_remove: List[int] = []
+
+    _updates: Tuple[TableUpdate, ...] = ()
+    _requirements: Tuple[TableRequirement, ...] = ()
+
+    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 _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."""
+        """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
 
-    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)
+    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)
         return self
 
-    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),
-            )
-
-            # 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
-
-            # clean up the the unused files
-
+    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)
         return self
 
-    def cleanup_files(self):
-        """Remove files no longer referenced by any snapshots."""
-        # 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
+    # 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
 
 
diff --git a/tests/table/test_expire_snapshots.py b/tests/table/test_expire_snapshots.py
@@ -1,77 +1,117 @@
-# pylint:disable=redefined-outer-name
-# pylint:disable=redefined-outer-name
-from unittest.mock import Mock
+from datetime import datetime, timezone
+from pathlib import PosixPath
+from random import randint
+import time
+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
+from pyiceberg.table.update.snapshot import ExpireSnapshots
+from pyiceberg.transforms import IdentityTransform
+from pyiceberg.types import BooleanType, FloatType, IntegerType, ListType, LongType, MapType, StructType
+from tests.catalog.test_base import InMemoryCatalog, Table
 from pyiceberg.table import Table
-from pyiceberg.table.metadata import new_table_metadata
-from pyiceberg.table.snapshots import Snapshot, SnapshotLogEntry
-from pyiceberg.table.update.snapshot import ManageSnapshots
-
 from pyiceberg.schema import Schema
-from pyiceberg.partitioning import PartitionSpec
-from pyiceberg.table.sorting import SortOrder
+from pyiceberg.types import NestedField, LongType, StringType
+from pyiceberg.table.snapshots import Snapshot
+from pyiceberg.table.metadata import TableMetadata, TableMetadataV2, new_table_metadata
 
 
 
-@pytest.fixture
-def mock_table():
-    """
-    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"),
-        Snapshot(snapshot_id=3, timestamp_ms=3000, manifest_list="manifest3.avro"),
-    ]
-    snapshot_log = [
-        SnapshotLogEntry(snapshot_id=1, timestamp_ms=1000),
-        SnapshotLogEntry(snapshot_id=2, timestamp_ms=2000),
-        SnapshotLogEntry(snapshot_id=3, timestamp_ms=3000),
-    ]
 
-    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,
+@pytest.fixture
+def generate_test_table() -> Table:
+    def generate_snapshot(
+        snapshot_id: int,
+        parent_snapshot_id: Optional[int] = None,
+        timestamp_ms: Optional[int] = None,
+        sequence_number: int = 0,
+    ) -> Dict[str, Any]:
+        return {
+            "snapshot-id": snapshot_id,
+            "parent-snapshot-id": parent_snapshot_id,
+            "timestamp-ms": timestamp_ms or int(time.time() * 1000),
+            "sequence-number": sequence_number,
+            "summary": {"operation": "append"},
+            "manifest-list": f"s3://a/b/{snapshot_id}.avro",
         }
-    )
 
-    table = Mock(spec=Table)
-    table.metadata = metadata
-    table.identifier = ("db", "table") 
+    snapshots = []
+    snapshot_log = []
+    initial_snapshot_id = 3051729675574597004
+
+    for i in range(2000):
+        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 = {
+        "format-version": 2,
+        "table-uuid": "9c12d441-03fe-4693-9a96-a0705ddf69c1",
+        "location": "s3://bucket/test/location",
+        "last-sequence-number": 34,
+        "last-updated-ms": 1602638573590,
+        "last-column-id": 3,
+        "current-schema-id": 1,
+        "schemas": [
+            {"type": "struct", "schema-id": 0, "fields": [{"id": 1, "name": "x", "required": True, "type": "long"}]},
+            {
+                "type": "struct",
+                "schema-id": 1,
+                "identifier-field-ids": [1, 2],
+                "fields": [
+                    {"id": 1, "name": "x", "required": True, "type": "long"},
+                    {"id": 2, "name": "y", "required": True, "type": "long", "doc": "comment"},
+                    {"id": 3, "name": "z", "required": True, "type": "long"},
+                ],
+            },
+        ],
+        "default-spec-id": 0,
+        "partition-specs": [{"spec-id": 0, "fields": [{"name": "x", "transform": "identity", "source-id": 1, "field-id": 1000}]}],
+        "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"},
+                ],
+            }
+        ],
+        "properties": {"read.split.target.size": "134217728"},
+        "current-snapshot-id": initial_snapshot_id + 1999,
+        "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}},
+    }
+
+    return Table(
+        identifier=("database", "table"),
+        metadata=metadata,
+        metadata_location=f"{metadata['location']}/uuid.metadata.json",
+        io=load_file_io(),
+        catalog=NoopCatalog("NoopCatalog"),
+    )
 
-    return table
 
-def test_expire_snapshots_removes_correct_snapshots(mock_table: Mock):
+
+def test_expire_snapshots_removes_correct_snapshots(generate_test_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.
-
     """
 
-    with ManageSnapshots(mock_table) as transaction:
-        # Mock the transaction to return the mock table
-        transaction.exipre_snapshot_by_id(1).exipre_snapshot_by_id(2).expire_snapshots().cleanup_files()
-
+    # Use the fixture-provided table
+    with ExpireSnapshots(generate_test_table.transaction()) as manage_snapshots:
+        manage_snapshots.expire_snapshot_id(3051729675574597004)
 
-    for snapshot in mock_table.metadata.snapshots:
-        # Verify that the snapshot is removed from the metadata
-        assert snapshot.snapshot_id not in [1, 2]
+    # 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})

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,2 @@`
	`1`	`+in the snapshot.py class, you define the "api" or logic to collect the changes and then stage, them. Then the`
	`2`	`+__init__.py has a decorator that calls the type to actually apply the metadata changes.`