move apis from Transaction to Table

Author: chinmay-bhat

pyiceberg/table/__init__.py

@@ -340,23 +340,6 @@ def set_properties(self, properties: Properties = EMPTY_DICT, **kwargs: Any) ->
         updates = properties or kwargs
         return self._apply((SetPropertiesUpdate(updates=updates),))
 
-    def rollback_to_snapshot(self, snapshot_id: int) -> Transaction:
-        if self.table_metadata.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()]:
-            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, ref_name="main", type="branch")
-
-    def rollback_to_timestamp(self, timestamp: int) -> Transaction:
-        if (snapshot := self.latest_snapshot_before_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, ref_name="main", type="branch")
-
-    def set_current_snapshot(self, snapshot_id: int) -> None:
-        if self.table_metadata.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, ref_name="main", type="branch")
-
     def update_schema(self, allow_incompatible_changes: bool = False, case_sensitive: bool = True) -> UpdateSchema:
         """Create a new UpdateSchema to alter the columns of this table.
 
@@ -1307,6 +1290,23 @@ 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:
+        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()]:
+            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")
+
+    def rollback_to_timestamp(self, timestamp: int) -> Transaction:
+        if (snapshot := self.latest_snapshot_before_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")
+
+    def set_current_snapshot(self, snapshot_id: int) -> Transaction:
+        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")
+
     def history(self) -> List[SnapshotLogEntry]:
         """Get the snapshot history of this table."""
         return self.metadata.snapshot_log