wip

Author: chinmay-bhat

pyiceberg/table/__init__.py

@@ -1290,22 +1290,22 @@ def snapshot_by_name(self, name: str) -> Optional[Snapshot]:
             return self.snapshot_by_id(ref.snapshot_id)
         return None
 
-    def rollback_to_snapshot(self, snapshot_id: int) -> Transaction:
+    def rollback_to_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}")
-        if snapshot_id not in [ancestor.ancestor_id for ancestor in self.current_ancestors()]:
+        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.set_ref_snapshot(snapshot_id=snapshot_id, parent_snapshot_id=self.current_snapshot().snapshot_id, ref_name="main", type="branch")
+        return self._transaction.set_ref_snapshot(snapshot_id=snapshot_id, ref_name="main", type="branch")
 
-    def rollback_to_timestamp(self, timestamp: int) -> Transaction:
-        if (snapshot := self.latest_snapshot_before_timestamp(timestamp)) is None:
+    def rollback_to_timestamp(self, timestamp: int) -> ManageSnapshots:
+        if (snapshot := self.snapshot_as_of_timestamp(timestamp)) is None:
             raise ValidationError(f"Cannot roll back, no valid snapshot older than: {timestamp}")
-        return self.set_ref_snapshot(snapshot_id=snapshot.snapshot_id, parent_snapshot_id=self.current_snapshot().snapshot_id, ref_name="main", type="branch")
+        return self._transaction.set_ref_snapshot(snapshot_id=snapshot.snapshot_id, ref_name="main", type="branch")
 
-    def set_current_snapshot(self, snapshot_id: int) -> Transaction:
+    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.set_ref_snapshot(snapshot_id=snapshot_id, parent_snapshot_id=self.current_snapshot().snapshot_id, ref_name="main", type="branch")
+        return self._transaction.set_ref_snapshot(snapshot_id=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

@@ -869,3 +869,43 @@ def table_write_subset_of_schema(session_catalog: Catalog, arrow_table_with_null
     tbl.append(arrow_table_without_some_columns)
     # overwrite and then append should produce twice the data
     assert len(tbl.scan().to_arrow()) == len(arrow_table_without_some_columns) * 2
+
+
+@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,
+    )
+
+
+@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,
+    )

tests/table/test_init.py

@@ -652,42 +652,6 @@ def test_update_metadata_add_snapshot(table_v2: Table) -> None:
     assert new_metadata.last_updated_ms == new_snapshot.timestamp_ms
 
 
-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_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_update_metadata_set_snapshot_ref(table_v2: Table) -> None:
     update = SetSnapshotRefUpdate(
         ref_name="main",