feat: add restore remote api (#2282)

python/python/lancedb/_lancedb.pyi

@@ -48,7 +48,7 @@ class Table:
     async def version(self) -> int: ...
     async def checkout(self, version: int): ...
     async def checkout_latest(self): ...
-    async def restore(self): ...
+    async def restore(self, version: Optional[int] = None): ...
     async def list_indices(self) -> list[IndexConfig]: ...
     async def delete(self, filter: str): ...
     async def add_columns(self, columns: list[tuple[str, str]]) -> None: ...

python/python/lancedb/remote/table.py

@@ -87,6 +87,9 @@ class RemoteTable(Table):
     def checkout_latest(self):
         return LOOP.run(self._table.checkout_latest())
 
+    def restore(self, version: Optional[int] = None):
+        return LOOP.run(self._table.restore(version))
+
     def list_indices(self) -> Iterable[IndexConfig]:
         """List all the indices on the table"""
         return LOOP.run(self._table.list_indices())

python/python/lancedb/table.py

@@ -1342,6 +1342,21 @@ class Table(ABC):
         It can also be used to undo a `[Self::checkout]` operation
         """
 
+    @abstractmethod
+    def restore(self, version: Optional[int] = None):
+        """Restore a version of the table. This is an in-place operation.
+
+        This creates a new version where the data is equivalent to the
+        specified previous version. Data is not copied (as of python-v0.2.1).
+
+        Parameters
+        ----------
+        version : int, default None
+            The version to restore. If unspecified then restores the currently
+            checked out version. If the currently checked out version is the
+            latest version then this is a no-op.
+        """
+
     @abstractmethod
     def list_versions(self) -> List[Dict[str, Any]]:
         """List all versions of the table"""
@@ -3613,7 +3628,7 @@ class AsyncTable:
         """
         await self._inner.checkout_latest()
 
-    async def restore(self):
+    async def restore(self, version: Optional[int] = None):
         """
         Restore the table to the currently checked out version
 
@@ -3626,7 +3641,7 @@ class AsyncTable:
         Once the operation concludes the table will no longer be in a checked
         out state and the read_consistency_interval, if any, will apply.
         """
-        await self._inner.restore()
+        await self._inner.restore(version)
 
     async def optimize(
         self,

python/src/table.rs

@@ -303,12 +303,16 @@ impl Table {
         })
     }
 
-    pub fn restore(self_: PyRef<'_, Self>) -> PyResult<Bound<'_, PyAny>> {
+    #[pyo3(signature = (version=None))]
+    pub fn restore(self_: PyRef<'_, Self>, version: Option<u64>) -> PyResult<Bound<'_, PyAny>> {
         let inner = self_.inner_ref()?.clone();
-        future_into_py(
-            self_.py(),
-            async move { inner.restore().await.infer_error() },
-        )
+
+        future_into_py(self_.py(), async move {
+            if let Some(version) = version {
+                inner.checkout(version).await.infer_error()?;
+            }
+            inner.restore().await.infer_error()
+        })
     }
 
     pub fn query(&self) -> Query {

rust/lancedb/src/remote/table.rs

@@ -425,10 +425,16 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
         Ok(())
     }
     async fn restore(&self) -> Result<()> {
-        self.check_mutable().await?;
-        Err(Error::NotSupported {
-            message: "restore is not supported on LanceDB cloud.".into(),
-        })
+        let mut request = self
+            .client
+            .post(&format!("/v1/table/{}/restore/", self.name));
+        let version = self.current_version().await;
+        let body = serde_json::json!({ "version": version });
+        request = request.json(&body);
+
+        let (request_id, response) = self.client.send(request, true).await?;
+        self.check_table_response(&request_id, response).await?;
+        Ok(())
     }
 
     async fn list_versions(&self) -> Result<Vec<Version>> {