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.
This commit is contained in:
somil jain
2026-10-31 07:17:07 +08:00
committed by GitHub
parent e410a108f0
commit bd6e907c73
4 changed files with 219 additions and 31 deletions
+3 -2
View File
@@ -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
+1 -2
View File
@@ -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))
+32 -19
View File
@@ -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():
+183 -8
View File
@@ -1516,14 +1516,21 @@ impl<S: HttpSend> Database for RemoteDatabase<S> {
})
}
#[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(|_| {