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(|_| {