updates

Author: chinmay-bhat

pyiceberg/table/__init__.py

@@ -1291,21 +1291,29 @@ def snapshot_by_name(self, name: str) -> Optional[Snapshot]:
         return None
 
     def rollback_to_snapshot(self, snapshot_id: int) -> ManageSnapshots:
+        """Rollback the table to the given snapshot id, whose snapshot needs to be an ancestor of the current table state."""
         if self.snapshot_by_id(snapshot_id) is None:
             raise ValidationError(f"Cannot roll back to unknown snapshot id: {snapshot_id}")
         if snapshot_id not in set(ancestor.ancestor_id for ancestor in ancestors_of(self.current_snapshot(), self._transaction._table.metadata)):
             raise ValidationError(f"Cannot roll back to snapshot, not an ancestor of the current state: {snapshot_id}")
         return self._transaction.set_ref_snapshot(snapshot_id=snapshot_id, ref_name="main", type="branch")
 
     def rollback_to_timestamp(self, timestamp: int) -> ManageSnapshots:
-        if (snapshot := self.snapshot_as_of_timestamp(timestamp)) is None:
+        """Rollback the table to the snapshot right before the given timestamp."""
+        if (snapshot := self.snapshot_as_of_timestamp(timestamp, inclusive=False)) is None:
             raise ValidationError(f"Cannot roll back, no valid snapshot older than: {timestamp}")
         return self._transaction.set_ref_snapshot(snapshot_id=snapshot.snapshot_id, ref_name="main", type="branch")
 
-    def set_current_snapshot(self, snapshot_id: int) -> ManageSnapshots:
-        if self.snapshot_by_id(snapshot_id) is None:
-            raise ValidationError(f"Cannot roll back to unknown snapshot id: {snapshot_id}")
-        return self._transaction.set_ref_snapshot(snapshot_id=snapshot_id, ref_name="main", type="branch")
+    def set_current_snapshot(self, snapshot_id: Optional[int] = None, ref_name: Optional[str] = None) -> ManageSnapshots:
+        """Set the table to a specific snapshot identified either by its id or the branch/tag its on, not both."""
+        if (snapshot_id or ref_name) and not (snapshot_id or ref_name):
+            target_snapshot_id = snapshot_id if snapshot_id is not None else self.metadata.refs[ref_name].snapshot_id
+            snapshot = self.snapshot_by_id(target_snapshot_id)
+        else:
+            raise ValidationError("Either snapshot_id or ref must be provided")
+        if snapshot is None:
+            raise ValidationError(f"Cannot set snapshot current with snapshot id: {snapshot_id} or ref_name: {ref_name}")
+        return self._transaction.set_ref_snapshot(snapshot_id=snapshot.snapshot_id, ref_name="main", type="branch")
 
     def history(self) -> List[SnapshotLogEntry]:
         """Get the snapshot history of this table."""

tests/integration/test_writes/test_writes.py

@@ -38,7 +38,7 @@
 from pyiceberg.exceptions import NoSuchTableError
 from pyiceberg.partitioning import PartitionField, PartitionSpec
 from pyiceberg.schema import Schema
-from pyiceberg.table import TableProperties, _dataframe_to_data_files
+from pyiceberg.table import TableProperties, _dataframe_to_data_files, SnapshotRef
 from pyiceberg.transforms import IdentityTransform
 from pyiceberg.types import IntegerType, NestedField
 from tests.conftest import TEST_DATA_WITH_NULL
@@ -873,39 +873,51 @@ def table_write_subset_of_schema(session_catalog: Catalog, arrow_table_with_null
 
 @pytest.mark.integration
 @pytest.mark.parametrize("catalog", [pytest.lazy_fixture("session_catalog_hive"), pytest.lazy_fixture("session_catalog")])
-def test_update_metadata_rollback_to_snapshot(table_v2: Table) -> None:
-    assert table_v2.rollback_to_snapshot(snapshot_id=3051729675574597004) == Snapshot(
-        snapshot_id=3051729675574597004,
-        parent_snapshot_id=None,
-        sequence_number=0,
-        timestamp_ms=1515100955770,
-        manifest_list='s3://a/b/1.avro',
-        summary=Summary(Operation.APPEND),
-        schema_id=None,
-    )
+def test_rollback_to_snapshot(catalog: Catalog) -> None:
+    identifier = "default.test_table_snapshot_operations"
+    tbl = catalog.load_table(identifier)
+    assert len(tbl.history()) > 3
+    rollback_snapshot_id = tbl.history()[-3].snapshot_id
+    current_snapshot_id = tbl.current_snapshot().snapshot_id
+    ms = tbl.manage_snapshots().rollback_to_snapshot(snapshot_id=rollback_snapshot_id)
+    ms.commit()
+    tbl.refresh()
+    assert tbl.current_snapshot().snapshot_id is not current_snapshot_id
+    assert tbl.metadata.refs["main"] == SnapshotRef(snapshot_id=rollback_snapshot_id)
 
 
 @pytest.mark.integration
 @pytest.mark.parametrize("catalog", [pytest.lazy_fixture("session_catalog_hive"), pytest.lazy_fixture("session_catalog")])
-def test_update_metadata_rollback_to_timestamp(table_v2: Table) -> None:
-    assert table_v2.rollback_to_timestamp(timestamp=1555100955771) == Snapshot(
-        snapshot_id=3055729675574597004,
-        parent_snapshot_id=3051729675574597004,
-        sequence_number=1,
-        timestamp_ms=1555100955770,
-        manifest_list='s3://a/b/2.avro',
-        summary=Summary(Operation.APPEND),
-        schema_id=1,
-    )
-
-
-def test_update_metadata_set_current_snapshot(table_v2: Table) -> None:
-    assert table_v2.set_current_snapshot(snapshot_id=3051729675574597004) == Snapshot(
-        snapshot_id=3051729675574597004,
-        parent_snapshot_id=None,
-        sequence_number=0,
-        timestamp_ms=1515100955770,
-        manifest_list='s3://a/b/1.avro',
-        summary=Summary(Operation.APPEND),
-        schema_id=None,
-    )
+def test_rollback_to_timestamp(catalog: Catalog) -> None:
+    identifier = "default.test_table_snapshot_operations"
+    tbl = catalog.load_table(identifier)
+    assert len(tbl.history()) > 3
+    timestamp = tbl.history()[-2].timestamp_ms
+    current_snapshot_id = tbl.current_snapshot().snapshot_id
+    # not inclusive of rollback_timestamp
+    ms = tbl.manage_snapshots().rollback_to_timestamp(timestamp=timestamp)
+    ms.commit()
+    tbl.refresh()
+    assert tbl.current_snapshot().snapshot_id is not current_snapshot_id
+    assert tbl.metadata.refs["main"] == SnapshotRef(snapshot_id=tbl.history()[-4].snapshot_id)
+
+
+def test_set_current_snapshot(catalog: Catalog) -> None:
+    identifier = "default.test_table_snapshot_operations"
+    tbl = catalog.load_table(identifier)
+    assert len(tbl.history()) > 3
+    # test with snapshot_id
+    target_snapshot_id = tbl.history()[-4].snapshot_id
+    current_snapshot_id = tbl.current_snapshot().snapshot_id
+    tbl.manage_snapshots().set_current_snapshot(snapshot_id=target_snapshot_id).commit()
+    tbl.refresh()
+    assert tbl.current_snapshot().snapshot_id is not current_snapshot_id
+    assert tbl.metadata.refs["main"] == SnapshotRef(snapshot_id=target_snapshot_id)
+    # test with ref_name
+    new_current_snapshot_id = tbl.current_snapshot().snapshot_id
+    expected_snapshot_id = tbl.history()[-3].snapshot_id
+    tbl.manage_snapshots().create_tag(snapshot_id=expected_snapshot_id, tag_name="test-tag").commit()
+    tbl.manage_snapshots().set_current_snapshot(ref_name="test-tag").commit()
+    tbl.refresh()
+    assert tbl.current_snapshot().snapshot_id is not new_current_snapshot_id
+    assert tbl.metadata.refs["main"] == SnapshotRef(snapshot_id=expected_snapshot_id)