From bd6e907c7313a29fbd1a8e7bd22366549fb0a602 Mon Sep 17 00:00:00 2001 From: somil jain Date: Sat, 31 Oct 2026 04:47:07 +0530 Subject: [PATCH] feat(remote): support dropping all tables in remote databases Fixes #4403 Implement remote drop_all_tables() support by listing tables and dropping them individually, including tables within namespaces. Handle paginated table listings and propagate errors appropriately. Add tests covering root and namespace tables, pagination, empty databases, and listing failures to ensure reliable remote database cleanup. --- python/python/lancedb/db.py | 5 +- python/python/lancedb/remote/db.py | 3 +- python/python/tests/test_remote_db.py | 51 ++++--- rust/lancedb/src/remote/db.rs | 191 ++++++++++++++++++++++++-- 4 files changed, 219 insertions(+), 31 deletions(-) diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index a2988933a..a93444a9e 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -851,8 +851,7 @@ class DBConnection(EnforceOverrides): """ Drop all tables from the database - LanceDB Distributed does not currently support this operation. Use - ``list_tables()`` and ``drop_table()`` to drop tables individually. + Lists the tables in the selected namespace and drops them one at a time. Parameters ---------- @@ -3251,6 +3250,8 @@ class AsyncConnection(object): async def drop_all_tables(self, namespace_path: Optional[List[str]] = None): """Drop all tables from the database. + Lists the tables in the selected namespace and drops them one at a time. + Parameters ---------- namespace_path: List[str], optional diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index 436f4f10e..48ce686b9 100644 --- a/python/python/lancedb/remote/db.py +++ b/python/python/lancedb/remote/db.py @@ -910,8 +910,7 @@ class RemoteDBConnection(DBConnection): def drop_all_tables(self, namespace_path: Optional[List[str]] = None): """Drop all tables from the database. - The remote API does not currently support this operation. Use - ``list_tables()`` and ``drop_table()`` to drop tables individually. + Lists the tables in the selected namespace and drops them one at a time. """ LOOP.run(self._conn.drop_all_tables(namespace_path=namespace_path)) diff --git a/python/python/tests/test_remote_db.py b/python/python/tests/test_remote_db.py index 1b4af3487..8a0e2ea92 100644 --- a/python/python/tests/test_remote_db.py +++ b/python/python/tests/test_remote_db.py @@ -1080,36 +1080,49 @@ async def test_remote_to_lance_rejection( @pytest.mark.asyncio @pytest.mark.parametrize( - "method, kwargs", + "method, kwargs, list_path, table_prefix", [ - ("drop_all_tables", {}), - ("drop_all_tables", {"namespace_path": []}), - ("drop_all_tables", {"namespace_path": ["team"]}), - ("drop_database", {}), + ("drop_all_tables", {}, "/v1/table/", ""), + ("drop_all_tables", {"namespace_path": []}, "/v1/table/", ""), + ( + "drop_all_tables", + {"namespace_path": ["team"]}, + "/v1/namespace/team/table/list", + "team$", + ), + ("drop_database", {}, "/v1/table/", ""), ], ) -async def test_remote_drop_all_tables_not_supported(method, kwargs): +async def test_remote_drop_all_tables(method, kwargs, list_path, table_prefix): requests = [] + drop_paths = { + f"/v1/table/{table_prefix}table1/drop/", + f"/v1/table/{table_prefix}table2/drop/", + } def handler(request): - requests.append(request.path) - request.send_response(500) - request.end_headers() + requests.append((request.command, request.path)) + if request.command == "GET": + assert request.path == list_path + send_json(request, {"tables": ["table1", "table2"]}) + return + + assert request.command == "POST" + assert request.path in drop_paths + send_json(request, {}) with mock_lancedb_connection(handler) as db: - with pytest.raises(NotImplementedError) as sync_error: - getattr(db, method)(**kwargs) + getattr(db, method)(**kwargs) async with mock_lancedb_connection_async(handler) as db: - with pytest.raises(NotImplementedError) as async_error: - await getattr(db, method)(**kwargs) + await getattr(db, method)(**kwargs) - assert "Dropping all tables is not currently supported in the remote API" in str( - sync_error.value - ) - assert "drop_table" in str(sync_error.value) - assert str(sync_error.value) == str(async_error.value) - assert requests == [] + expected_requests = [ + ("GET", list_path), + ("POST", f"/v1/table/{table_prefix}table1/drop/"), + ("POST", f"/v1/table/{table_prefix}table2/drop/"), + ] + assert requests == expected_requests * 2 def test_table_to_pandas_not_supported(): diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index b50c355d5..9e37b304e 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -1516,14 +1516,21 @@ impl Database for RemoteDatabase { }) } + #[allow(deprecated)] async fn drop_all_tables(&self, namespace_path: &[String]) -> Result<()> { - // TODO: Implement namespace-aware drop_all_tables - let _namespace_path = namespace_path; // Suppress unused warning for now - Err(crate::Error::NotSupported { - message: "Dropping all tables is not currently supported in the remote API. \ - List tables and call drop_table for each table instead." - .to_string(), - }) + let tables = self + .table_names(TableNamesRequest { + namespace_path: namespace_path.to_vec(), + start_after: None, + limit: None, + }) + .await?; + + for table in tables { + self.drop_table(&table, namespace_path).await?; + } + + Ok(()) } async fn list_namespaces( @@ -1727,7 +1734,7 @@ mod tests { use futures::TryStreamExt; use std::collections::HashMap; use std::sync::atomic::{AtomicUsize, Ordering}; - use std::sync::{Arc, OnceLock}; + use std::sync::{Arc, Mutex, OnceLock}; use arrow_array::{Int32Array, RecordBatch}; use arrow_schema::{DataType, Field, Schema}; @@ -2645,6 +2652,174 @@ mod tests { // NOTE: the API will return 200 even if the table does not exist. So we shouldn't expect 404. } + #[tokio::test] + async fn test_drop_all_tables() { + let dropped_paths = Arc::new(Mutex::new(Vec::new())); + let dropped_paths_in_handler = Arc::clone(&dropped_paths); + let conn = Connection::new_with_handler(move |request| match request.url().path() { + "/v1/table/" => { + assert_eq!(request.method(), &reqwest::Method::GET); + assert_eq!(request.url().query(), None); + http::Response::builder() + .status(200) + .body(r#"{"tables":["table1","table2"]}"#.to_string()) + .unwrap() + } + "/v1/table/table1/drop/" | "/v1/table/table2/drop/" => { + assert_eq!(request.method(), &reqwest::Method::POST); + assert_eq!(request.url().query(), None); + assert!(request.body().is_none()); + dropped_paths_in_handler + .lock() + .unwrap() + .push(request.url().path().to_string()); + http::Response::builder() + .status(200) + .body(String::new()) + .unwrap() + } + path => panic!("unexpected request path: {path}"), + }); + + conn.drop_all_tables(&[]).await.unwrap(); + + let mut actual = dropped_paths.lock().unwrap().clone(); + actual.sort(); + assert_eq!( + actual, + vec!["/v1/table/table1/drop/", "/v1/table/table2/drop/"] + ); + } + + #[tokio::test] + async fn test_drop_all_tables_with_namespace() { + let dropped_paths = Arc::new(Mutex::new(Vec::new())); + let dropped_paths_in_handler = Arc::clone(&dropped_paths); + let conn = Connection::new_with_handler(move |request| match request.url().path() { + "/v1/namespace/ns1$ns2/table/list" => { + assert_eq!(request.method(), &reqwest::Method::GET); + assert_eq!(request.url().query(), None); + http::Response::builder() + .status(200) + .body(r#"{"tables":["table1","table2"]}"#.to_string()) + .unwrap() + } + "/v1/table/ns1$ns2$table1/drop/" | "/v1/table/ns1$ns2$table2/drop/" => { + assert_eq!(request.method(), &reqwest::Method::POST); + assert_eq!(request.url().query(), None); + assert!(request.body().is_none()); + dropped_paths_in_handler + .lock() + .unwrap() + .push(request.url().path().to_string()); + http::Response::builder() + .status(200) + .body(String::new()) + .unwrap() + } + path => panic!("unexpected request path: {path}"), + }); + + conn.drop_all_tables(&["ns1".to_string(), "ns2".to_string()]) + .await + .unwrap(); + + let mut actual = dropped_paths.lock().unwrap().clone(); + actual.sort(); + assert_eq!( + actual, + vec![ + "/v1/table/ns1$ns2$table1/drop/", + "/v1/table/ns1$ns2$table2/drop/" + ] + ); + } + + #[tokio::test] + async fn test_drop_all_tables_follows_server_pages() { + let listing_page = Arc::new(AtomicUsize::new(0)); + let listing_page_in_handler = Arc::clone(&listing_page); + let dropped_paths = Arc::new(Mutex::new(Vec::new())); + let dropped_paths_in_handler = Arc::clone(&dropped_paths); + let conn = Connection::new_with_handler(move |request| match request.url().path() { + "/v1/table/" => { + assert_eq!(request.method(), &reqwest::Method::GET); + match listing_page_in_handler.fetch_add(1, Ordering::SeqCst) { + 0 => { + assert_eq!(request.url().query(), None); + http::Response::builder() + .status(200) + .body( + r#"{"tables":["table1"],"page_token":"opaque-token"}"#.to_string(), + ) + .unwrap() + } + 1 => { + let query: HashMap<_, _> = request.url().query_pairs().collect(); + assert_eq!( + query.get("page_token").map(|token| token.as_ref()), + Some("opaque-token") + ); + http::Response::builder() + .status(200) + .body(r#"{"tables":["table2"]}"#.to_string()) + .unwrap() + } + _ => panic!("table listing requested too many pages"), + } + } + "/v1/table/table1/drop/" | "/v1/table/table2/drop/" => { + assert_eq!(request.method(), &reqwest::Method::POST); + dropped_paths_in_handler + .lock() + .unwrap() + .push(request.url().path().to_string()); + http::Response::builder() + .status(200) + .body(String::new()) + .unwrap() + } + path => panic!("unexpected request path: {path}"), + }); + + conn.drop_all_tables(&[]).await.unwrap(); + + assert_eq!(listing_page.load(Ordering::SeqCst), 2); + assert_eq!( + *dropped_paths.lock().unwrap(), + vec!["/v1/table/table1/drop/", "/v1/table/table2/drop/"] + ); + } + + #[tokio::test] + async fn test_drop_all_tables_propagates_listing_error() { + let conn = Connection::new_with_handler(|request| { + assert_eq!(request.method(), &reqwest::Method::GET); + assert_eq!(request.url().path(), "/v1/table/"); + http::Response::builder() + .status(500) + .body("internal server error".to_string()) + .unwrap() + }); + + let error = conn.drop_all_tables(&[]).await.unwrap_err(); + assert!(matches!(error, Error::Retry { .. })); + } + + #[tokio::test] + async fn test_drop_all_tables_with_no_tables() { + let conn = Connection::new_with_handler(|request| { + assert_eq!(request.method(), &reqwest::Method::GET); + assert_eq!(request.url().path(), "/v1/table/"); + http::Response::builder() + .status(200) + .body(r#"{"tables":[]}"#.to_string()) + .unwrap() + }); + + conn.drop_all_tables(&[]).await.unwrap(); + } + #[tokio::test] async fn test_drop_table_does_not_read_response_body() { let conn = Connection::new_with_handler(|_| {