create rollback_to_snapshot, rollback_to_timestamp, set_current_snapshot apis

Author: chinmay-bhat

pyiceberg/table/__init__.py

@@ -340,6 +340,23 @@ 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.
 
@@ -1881,7 +1898,8 @@ def union_by_name(self, new_schema: Union[Schema, "pa.Schema"]) -> UpdateSchema:
         visit_with_partner(
             Catalog._convert_schema_if_needed(new_schema),
             -1,
-            UnionByNameVisitor(update_schema=self, existing_schema=self._schema, case_sensitive=self._case_sensitive),  # type: ignore
+            UnionByNameVisitor(update_schema=self, existing_schema=self._schema, case_sensitive=self._case_sensitive),
+            # type: ignore
             PartnerIdByNameAccessor(partner_schema=self._schema, case_sensitive=self._case_sensitive),
         )
         return self