mirror of
https://github.com/lancedb/lancedb.git
synced 2026-11-14 06:36:54 +00:00
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:
@@ -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
|
||||
|
||||
@@ -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))
|
||||
|
||||
|
||||
@@ -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():
|
||||
|
||||
@@ -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(|_| {
|
||||
|
||||
Reference in New Issue
Block a user