From c196d033e932591bb696772ebb3490cde49011b7 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Mon, 31 Aug 2026 22:17:11 +0800 Subject: [PATCH] feat: add drop_function client APIs (#4097) --- python/python/lancedb/_lancedb.pyi | 1 + python/python/lancedb/db.py | 18 ++++++++ python/python/lancedb/remote/db.py | 4 ++ .../tests/test_first_class_function_slice2.py | 45 +++++++++++++++++++ python/src/connection.rs | 11 +++++ rust/lancedb/src/connection.rs | 15 +++++++ rust/lancedb/src/database.rs | 4 ++ rust/lancedb/src/remote/db.rs | 38 ++++++++++++++++ .../tests/first_class_function_slice2.rs | 6 ++- 9 files changed, 141 insertions(+), 1 deletion(-) diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 7d7ca7f2a..05ece3043 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -150,6 +150,7 @@ class Connection(object): def job(self, job_id: str) -> Job: ... async def create_function_async(self, request_json: str) -> Job: ... async def get_function(self, name: str, version: str) -> str: ... + async def drop_function(self, name: str, version: str) -> bool: ... async def list_jobs(self) -> List[JobInfo]: ... async def get_job(self, job_id: str) -> Optional[JobDescription]: ... async def cancel_job(self, job_id: str) -> bool: ... diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index 51b8d9993..ecaae42f8 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -712,6 +712,16 @@ class DBConnection(EnforceOverrides): "Function catalog operations are not supported for this connection type" ) + def drop_function(self, name: str, *, version: str) -> bool: + """Drop one exact immutable Function version from the remote catalog. + + Returns True when the version changed to Dropped and False for an + idempotent replay. Local connections raise NotImplementedError. + """ + raise NotImplementedError( + "Function catalog operations are not supported for this connection type" + ) + def job(self, job_id: str) -> Job: """A [Job][lancedb.job.Job] handle for a server-side job by id. @@ -1413,6 +1423,10 @@ class LanceDBConnection(DBConnection): def get_function(self, name: str, *, version: str) -> FunctionVersion: return LOOP.run(self._conn.get_function(name, version=version)) + @override + def drop_function(self, name: str, *, version: str) -> bool: + return LOOP.run(self._conn.drop_function(name, version=version)) + @override def list_jobs(self) -> List[JobInfo]: """List server-side jobs across the database's tables.""" @@ -2243,6 +2257,10 @@ class AsyncConnection(object): """Open one exact immutable Function version from the remote catalog.""" return FunctionVersion.from_json(await self._inner.get_function(name, version)) + async def drop_function(self, name: str, *, version: str) -> bool: + """Drop one exact immutable Function version from the remote catalog.""" + return await self._inner.drop_function(name, version) + async def list_jobs(self) -> List[JobInfo]: """List server-side jobs across the database's tables.""" return await self._inner.list_jobs() diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index b228cfb5b..27e21d200 100644 --- a/python/python/lancedb/remote/db.py +++ b/python/python/lancedb/remote/db.py @@ -749,6 +749,10 @@ class RemoteDBConnection(DBConnection): def get_function(self, name: str, *, version: str) -> FunctionVersion: return LOOP.run(self._conn.get_function(name, version=version)) + @override + def drop_function(self, name: str, *, version: str) -> bool: + return LOOP.run(self._conn.drop_function(name, version=version)) + @override def list_jobs(self) -> List["JobInfo"]: """List server-side jobs across the database's tables.""" diff --git a/python/python/tests/test_first_class_function_slice2.py b/python/python/tests/test_first_class_function_slice2.py index 415ecbe0f..bab78316c 100644 --- a/python/python/tests/test_first_class_function_slice2.py +++ b/python/python/tests/test_first_class_function_slice2.py @@ -941,6 +941,8 @@ def test_local_function_catalog_operations_are_not_supported(tmp_path): db.create_function_async(normalize_score) with pytest.raises(NotImplementedError, match=message): db.get_function("normalize_score", version="fv_exact") + with pytest.raises(NotImplementedError, match=message): + db.drop_function("normalize_score", version="fv_exact") @contextlib.contextmanager @@ -986,6 +988,12 @@ def _mock_remote_function_catalog(): "version": "fv_exact", } response = state["version"] + elif self.path == "/v1/functions/drop": + assert body == { + "name": "normalize_score", + "version": "fv_exact", + } + response = {"dropped": True} else: status = 404 response = {"error": "not found"} @@ -1044,3 +1052,40 @@ def test_blocking_remote_registration_returns_function_version(): "/v1/functions/create", "/v1/jobs/describe", ] + + +def test_remote_drop_function_sends_exact_version(): + with _mock_remote_function_catalog() as (host, state): + db = lancedb.connect( + "db://dev", + api_key="fake", + host_override=host, + client_config={"retry_config": {"retries": 0}}, + ) + assert db.drop_function("normalize_score", version="fv_exact") is True + + assert state["requests"] == [ + ( + "/v1/functions/drop", + {"name": "normalize_score", "version": "fv_exact"}, + ) + ] + + +@pytest.mark.asyncio +async def test_async_remote_drop_function_sends_exact_version(): + with _mock_remote_function_catalog() as (host, state): + db = await lancedb.connect_async( + "db://dev", + api_key="fake", + host_override=host, + client_config={"retry_config": {"retries": 0}}, + ) + assert await db.drop_function("normalize_score", version="fv_exact") is True + + assert state["requests"] == [ + ( + "/v1/functions/drop", + {"name": "normalize_score", "version": "fv_exact"}, + ) + ] diff --git a/python/src/connection.rs b/python/src/connection.rs index 902489f4f..fc835f805 100644 --- a/python/src/connection.rs +++ b/python/src/connection.rs @@ -629,6 +629,17 @@ impl Connection { }) } + pub fn drop_function( + self_: PyRef<'_, Self>, + name: String, + version: String, + ) -> PyResult> { + let inner = self_.get_inner()?.clone(); + future_into_py(self_.py(), async move { + inner.drop_function(name, version).await.infer_error() + }) + } + pub fn list_jobs(self_: PyRef<'_, Self>) -> PyResult> { let inner = self_.get_inner()?.clone(); future_into_py(self_.py(), async move { diff --git a/rust/lancedb/src/connection.rs b/rust/lancedb/src/connection.rs index 5f66d9dee..943ad51b7 100644 --- a/rust/lancedb/src/connection.rs +++ b/rust/lancedb/src/connection.rs @@ -523,6 +523,21 @@ impl Connection { .await } + /// Drop one exact immutable Function version from the remote catalog. + /// + /// Returns `true` when the server appended a Dropped transition and + /// `false` for an idempotent replay. Local databases return + /// [`Error::NotSupported`]. + pub async fn drop_function( + &self, + name: impl AsRef, + version: impl AsRef, + ) -> Result { + self.internal + .drop_function(name.as_ref(), version.as_ref()) + .await + } + /// Rename a table in the database. /// /// This is only supported in LanceDB Cloud. diff --git a/rust/lancedb/src/database.rs b/rust/lancedb/src/database.rs index 6c4537972..775b0b579 100644 --- a/rust/lancedb/src/database.rs +++ b/rust/lancedb/src/database.rs @@ -307,6 +307,10 @@ pub trait Database: ) -> Result { function_catalog_not_supported() } + /// Drop one exact immutable Function version from the remote catalog. + async fn drop_function(&self, _name: &str, _version: &str) -> Result { + function_catalog_not_supported() + } /// A [`crate::job::Job`] handle for a server-side job by id, suitable for /// waiting on or cancelling the job. The handle is constructed without a /// server round trip; an unknown id surfaces when the handle is used. diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index da9a4b09b..39b258a63 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -533,6 +533,11 @@ struct RemoteListJobsResponse { page_token: Option, } +#[derive(serde::Deserialize)] +struct RemoteDropFunctionResponse { + dropped: bool, +} + /// Bound on `list_jobs` page walking; a warning is logged when the listing /// is truncated at this many pages. const MAX_LIST_JOBS_PAGES: usize = 100; @@ -583,6 +588,20 @@ impl Database for RemoteDatabase { response.json().await.err_to_http(request_id) } + async fn drop_function(&self, name: &str, version: &str) -> Result { + let req = self + .client + .post("/v1/functions/drop") + .json(&serde_json::json!({ + "name": name, + "version": version, + })); + let (request_id, response) = self.client.send(req).await?; + let response = self.client.check_response(&request_id, response).await?; + let response: RemoteDropFunctionResponse = response.json().await.err_to_http(request_id)?; + Ok(response.dropped) + } + fn job(&self, job_id: &str) -> Result { Ok(crate::job::Job::new(Box::new(super::job::RemoteJob::new( self.client.clone(), @@ -2689,6 +2708,25 @@ mod tests { assert_eq!(version.version(), "fv_01K3EXACT"); } + #[tokio::test] + async fn test_drop_function_sends_exact_version_and_decodes_replay() { + let conn = Connection::new_with_handler(|request| { + assert_eq!(request.method(), &reqwest::Method::POST); + assert_eq!(request.url().path(), "/v1/functions/drop"); + let body: serde_json::Value = + serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap(); + assert_eq!( + body, + serde_json::json!({"name": "embed", "version": "fv_01K3EXACT"}) + ); + http::Response::builder() + .status(200) + .body(r#"{"dropped":false}"#) + .unwrap() + }); + assert!(!conn.drop_function("embed", "fv_01K3EXACT").await.unwrap()); + } + #[tokio::test] async fn test_conn_job_waits_to_done() { let polls = Arc::new(AtomicUsize::new(0)); diff --git a/rust/lancedb/tests/first_class_function_slice2.rs b/rust/lancedb/tests/first_class_function_slice2.rs index 93252dde4..6d046d9d1 100644 --- a/rust/lancedb/tests/first_class_function_slice2.rs +++ b/rust/lancedb/tests/first_class_function_slice2.rs @@ -45,7 +45,11 @@ async fn local_function_catalog_operations_return_stable_not_supported() { .get_function("normalize_score", "fv_exact") .await .unwrap_err(); - for error in [create_error, lookup_error] { + let drop_error = connection + .drop_function("normalize_score", "fv_exact") + .await + .unwrap_err(); + for error in [create_error, lookup_error, drop_error] { assert!(matches!( error, Error::NotSupported { message }