From c1d9d6f70b757a38dc4bf5671ea95368b86184cb Mon Sep 17 00:00:00 2001
From: Will Jones <willjones127@gmail.com>
Date: Fri, 27 Sep 2024 09:37:54 -0700
Subject: [PATCH] feat(rust): remote rename table (#1703)

Adds rename to remote table. Pre-requisite for
https://github.com/lancedb/lancedb/pull/1701
---
 rust/lancedb/src/connection.rs | 20 ++++++++++++++++++++
 rust/lancedb/src/remote/db.rs  | 34 +++++++++++++++++++++++++++++++++-
 2 files changed, 53 insertions(+), 1 deletion(-)

diff --git a/rust/lancedb/src/connection.rs b/rust/lancedb/src/connection.rs
index 8d073764..144b3134 100644
--- a/rust/lancedb/src/connection.rs
+++ b/rust/lancedb/src/connection.rs
@@ -431,6 +431,7 @@ pub(crate) trait ConnectionInternal:
         data: Box<dyn RecordBatchReader + Send>,
     ) -> Result<Table>;
     async fn do_open_table(&self, options: OpenTableBuilder) -> Result<Table>;
+    async fn rename_table(&self, old_name: &str, new_name: &str) -> Result<()>;
     async fn drop_table(&self, name: &str) -> Result<()>;
     async fn drop_db(&self) -> Result<()>;
 
@@ -513,6 +514,19 @@ impl Connection {
         OpenTableBuilder::new(self.internal.clone(), name.into())
     }
 
+    /// Rename a table in the database.
+    ///
+    /// This is only supported in LanceDB Cloud.
+    pub async fn rename_table(
+        &self,
+        old_name: impl AsRef<str>,
+        new_name: impl AsRef<str>,
+    ) -> Result<()> {
+        self.internal
+            .rename_table(old_name.as_ref(), new_name.as_ref())
+            .await
+    }
+
     /// Drop a table in the database.
     ///
     /// # Arguments
@@ -1066,6 +1080,12 @@ impl ConnectionInternal for Database {
         Ok(Table::new(native_table))
     }
 
+    async fn rename_table(&self, _old_name: &str, _new_name: &str) -> Result<()> {
+        Err(Error::NotSupported {
+            message: "rename_table is not supported in LanceDB OSS".to_string(),
+        })
+    }
+
     async fn drop_table(&self, name: &str) -> Result<()> {
         let dir_name = format!("{}.{}", name, LANCE_EXTENSION);
         let full_path = self.base_path.child(dir_name.clone());
diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs
index 93fbe86e..db26b181 100644
--- a/rust/lancedb/src/remote/db.rs
+++ b/rust/lancedb/src/remote/db.rs
@@ -149,6 +149,16 @@ impl<S: HttpSend> ConnectionInternal for RemoteDatabase<S> {
         ))))
     }
 
+    async fn rename_table(&self, current_name: &str, new_name: &str) -> Result<()> {
+        let req = self
+            .client
+            .post(&format!("/v1/table/{}/rename/", current_name));
+        let req = req.json(&serde_json::json!({ "new_table_name": new_name }));
+        let resp = self.client.send(req).await?;
+        self.client.check_response(resp).await?;
+        Ok(())
+    }
+
     async fn drop_table(&self, name: &str) -> Result<()> {
         let req = self.client.post(&format!("/v1/table/{}/drop/", name));
         let resp = self.client.send(req).await?;
@@ -174,7 +184,10 @@ mod tests {
     use arrow_array::{Int32Array, RecordBatch, RecordBatchIterator};
     use arrow_schema::{DataType, Field, Schema};
 
-    use crate::{remote::db::ARROW_STREAM_CONTENT_TYPE, Connection};
+    use crate::{
+        remote::{ARROW_STREAM_CONTENT_TYPE, JSON_CONTENT_TYPE},
+        Connection,
+    };
 
     #[tokio::test]
     async fn test_table_names() {
@@ -334,4 +347,23 @@ mod tests {
         conn.drop_table("table1").await.unwrap();
         // NOTE: the API will return 200 even if the table does not exist. So we shouldn't expect 404.
     }
+
+    #[tokio::test]
+    async fn test_rename_table() {
+        let conn = Connection::new_with_handler(|request| {
+            assert_eq!(request.method(), &reqwest::Method::POST);
+            assert_eq!(request.url().path(), "/v1/table/table1/rename/");
+            assert_eq!(
+                request.headers().get("Content-Type").unwrap(),
+                JSON_CONTENT_TYPE
+            );
+
+            let body = request.body().unwrap().as_bytes().unwrap();
+            let body: serde_json::Value = serde_json::from_slice(body).unwrap();
+            assert_eq!(body["new_table_name"], "table2");
+
+            http::Response::builder().status(200).body("").unwrap()
+        });
+        conn.rename_table("table1", "table2").await.unwrap();
+    }
 }