diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 2b2691139..78826df92 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 list_functions(self) -> List[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]: ... diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index ecaae42f8..2554e9908 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -712,6 +712,24 @@ class DBConnection(EnforceOverrides): "Function catalog operations are not supported for this connection type" ) + def list_functions(self) -> List[FunctionVersion]: + """List every published immutable Function version. + + Results are ordered by Function name then version. Local connections + raise ``NotImplementedError``. + + Examples + -------- + List the identities available to use in Function-backed columns: + + ```python + [(function.name, function.version) for function in db.list_functions()] + ``` + """ + raise NotImplementedError( + "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. @@ -1423,6 +1441,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 list_functions(self) -> List[FunctionVersion]: + return LOOP.run(self._conn.list_functions()) + @override def drop_function(self, name: str, *, version: str) -> bool: return LOOP.run(self._conn.drop_function(name, version=version)) @@ -2257,6 +2279,17 @@ 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 list_functions(self) -> List[FunctionVersion]: + """List every published immutable Function version. + + Results are ordered by Function name then version. Local connections + raise ``NotImplementedError``. + """ + return [ + FunctionVersion.from_json(value) + for value in await self._inner.list_functions() + ] + 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) diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index 27e21d200..0e95e035b 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 list_functions(self) -> List[FunctionVersion]: + return LOOP.run(self._conn.list_functions()) + @override def drop_function(self, name: str, *, version: str) -> bool: return LOOP.run(self._conn.drop_function(name, version=version)) diff --git a/python/python/tests/test_first_class_function_slice2.py b/python/python/tests/test_first_class_function_slice2.py index 57b08e18d..b1308a2bc 100644 --- a/python/python/tests/test_first_class_function_slice2.py +++ b/python/python/tests/test_first_class_function_slice2.py @@ -1017,6 +1017,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.list_functions() with pytest.raises(NotImplementedError, match=message): db.drop_function("normalize_score", version="fv_exact") @@ -1064,6 +1066,22 @@ def _mock_remote_function_catalog(): "version": "fv_exact", } response = state["version"] + elif self.path == "/v1/functions/list": + assert body["include_definition"] is True + if "page_token" not in body: + response = { + "functions": [ + { + "name": "normalize_score", + "version": "fv_exact", + "definition": state["version"], + } + ], + "page_token": "next", + } + else: + assert body["page_token"] == "next" + response = {"functions": []} elif self.path == "/v1/functions/drop": assert body == { "name": "normalize_score", @@ -1130,6 +1148,49 @@ def test_blocking_remote_registration_returns_function_version(): ] +def test_remote_list_functions_paginates_and_returns_typed_versions(): + 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}}, + ) + created = db.create_function(normalize_score) + state["requests"].clear() + functions = db.list_functions() + + assert functions == [created] + assert state["requests"] == [ + ("/v1/functions/list", {"include_definition": True}), + ( + "/v1/functions/list", + {"include_definition": True, "page_token": "next"}, + ), + ] + + +@pytest.mark.asyncio +async def test_async_remote_list_functions_returns_typed_versions(): + 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}}, + ) + registration = await db.create_function_async(normalize_score) + created = await registration.wait() + state["requests"].clear() + functions = await db.list_functions() + + assert functions == [created] + assert [path for path, _ in state["requests"]] == [ + "/v1/functions/list", + "/v1/functions/list", + ] + + def test_remote_drop_function_sends_exact_version(): with _mock_remote_function_catalog() as (host, state): db = lancedb.connect( diff --git a/python/src/connection.rs b/python/src/connection.rs index fc835f805..5477ab3d2 100644 --- a/python/src/connection.rs +++ b/python/src/connection.rs @@ -629,6 +629,19 @@ impl Connection { }) } + pub fn list_functions(self_: PyRef<'_, Self>) -> PyResult> { + let inner = self_.get_inner()?.clone(); + future_into_py(self_.py(), async move { + inner + .list_functions() + .await + .infer_error()? + .into_iter() + .map(|function| function.to_canonical_json().infer_error()) + .collect::>>() + }) + } + pub fn drop_function( self_: PyRef<'_, Self>, name: String, diff --git a/rust/lancedb/src/connection.rs b/rust/lancedb/src/connection.rs index 943ad51b7..35c5c0737 100644 --- a/rust/lancedb/src/connection.rs +++ b/rust/lancedb/src/connection.rs @@ -523,6 +523,28 @@ impl Connection { .await } + /// List every published immutable Function version in the remote catalog. + /// + /// Results are ordered by Function name then version. The client walks all + /// server pages before returning. Local databases return + /// [`Error::NotSupported`]. + /// + /// # Example + /// + /// ```no_run + /// # async fn list_functions( + /// # connection: &lancedb::Connection, + /// # ) -> Result<(), Box> { + /// for function in connection.list_functions().await? { + /// println!("{} {}", function.name(), function.version()); + /// } + /// # Ok(()) + /// # } + /// ``` + pub async fn list_functions(&self) -> Result> { + self.internal.list_functions().await + } + /// Drop one exact immutable Function version from the remote catalog. /// /// Returns `true` when the server appended a Dropped transition and diff --git a/rust/lancedb/src/database.rs b/rust/lancedb/src/database.rs index 775b0b579..61424bb05 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() } + /// List every published immutable Function version in the remote catalog. + async fn list_functions(&self) -> 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() diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index 39b258a63..ecdadf464 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::sync::Arc; use async_trait::async_trait; @@ -533,6 +533,19 @@ struct RemoteListJobsResponse { page_token: Option, } +#[derive(serde::Deserialize)] +struct RemoteListedFunctionVersion { + definition: FunctionVersion, +} + +#[derive(serde::Deserialize)] +struct RemoteListFunctionsResponse { + #[serde(default)] + functions: Vec, + #[serde(default)] + page_token: Option, +} + #[derive(serde::Deserialize)] struct RemoteDropFunctionResponse { dropped: bool, @@ -588,6 +601,43 @@ impl Database for RemoteDatabase { response.json().await.err_to_http(request_id) } + async fn list_functions(&self) -> Result> { + let mut functions = Vec::new(); + let mut page_token: Option = None; + let mut seen_page_tokens = HashSet::new(); + loop { + let mut body = serde_json::json!({ "include_definition": true }); + if let Some(token) = &page_token { + body["page_token"] = serde_json::Value::String(token.clone()); + } + let req = self.client.post("/v1/functions/list").json(&body); + let (request_id, response) = self.client.send(req).await?; + let response = self.client.check_response(&request_id, response).await?; + let status = response.status(); + let response: RemoteListFunctionsResponse = + response.json().await.err_to_http(request_id.clone())?; + functions.extend( + response + .functions + .into_iter() + .map(|listed| listed.definition), + ); + let Some(next_page_token) = response.page_token.filter(|token| !token.is_empty()) + else { + break; + }; + if !seen_page_tokens.insert(next_page_token.clone()) { + return Err(Error::Http { + source: "Function listing response repeated a page_token".into(), + request_id, + status_code: Some(status), + }); + } + page_token = Some(next_page_token); + } + Ok(functions) + } + async fn drop_function(&self, name: &str, version: &str) -> Result { let req = self .client @@ -2708,6 +2758,119 @@ mod tests { assert_eq!(version.version(), "fv_01K3EXACT"); } + #[tokio::test] + async fn test_list_functions_requests_definitions_and_paginates() { + const VERSION: &str = include_str!( + "../../tests/fixtures/first_class_functions/v1/remote_function_version.canonical.json" + ); + let version: serde_json::Value = serde_json::from_str(VERSION).unwrap(); + let page = Arc::new(AtomicUsize::new(0)); + let conn = Connection::new_with_handler(move |request| { + assert_eq!(request.method(), &reqwest::Method::POST); + assert_eq!(request.url().path(), "/v1/functions/list"); + let body: serde_json::Value = + serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap(); + assert_eq!(body["include_definition"], true); + match page.fetch_add(1, Ordering::SeqCst) { + 0 => { + assert!(body.get("page_token").is_none()); + http::Response::builder() + .status(200) + .body(r#"{"functions": [], "page_token": "next"}"#.to_string()) + .unwrap() + } + _ => { + assert_eq!(body["page_token"], "next"); + http::Response::builder() + .status(200) + .body( + serde_json::json!({ + "functions": [{ + "name": "embed", + "version": "fv_01K3EXACT", + "definition": version.clone(), + }], + }) + .to_string(), + ) + .unwrap() + } + } + }); + let functions = conn.list_functions().await.unwrap(); + assert_eq!(functions.len(), 1); + assert_eq!(functions[0].name(), "embed"); + assert_eq!(functions[0].version(), "fv_01K3EXACT"); + } + + #[tokio::test] + async fn test_list_functions_stops_on_an_empty_page_token() { + let requests = Arc::new(AtomicUsize::new(0)); + let seen = requests.clone(); + let conn = Connection::new_with_handler(move |request| { + seen.fetch_add(1, Ordering::SeqCst); + let body: serde_json::Value = + serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap(); + assert!(body.get("page_token").is_none()); + http::Response::builder() + .status(200) + .body(r#"{"functions": [], "page_token": ""}"#) + .unwrap() + }); + + let functions = conn.list_functions().await.unwrap(); + assert!(functions.is_empty()); + assert_eq!(requests.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn test_list_functions_rejects_a_page_token_cycle() { + let page = Arc::new(AtomicUsize::new(0)); + let requests = page.clone(); + let conn = Connection::new_with_handler(move |request| { + let body: serde_json::Value = + serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap(); + let next_page_token = match page.fetch_add(1, Ordering::SeqCst) { + 0 => { + assert!(body.get("page_token").is_none()); + "one" + } + 1 => { + assert_eq!(body["page_token"], "one"); + "two" + } + 2 => { + assert_eq!(body["page_token"], "two"); + "one" + } + page => panic!("unexpected page: {page}"), + }; + http::Response::builder() + .status(200) + .body( + serde_json::json!({ + "functions": [], + "page_token": next_page_token, + }) + .to_string(), + ) + .unwrap() + }); + + let error = conn.list_functions().await.unwrap_err(); + assert!( + matches!( + &error, + Error::Http { + status_code: Some(http::StatusCode::OK), + .. + } + ), + "got {error:?}" + ); + assert_eq!(requests.load(Ordering::SeqCst), 3); + } + #[tokio::test] async fn test_drop_function_sends_exact_version_and_decodes_replay() { let conn = Connection::new_with_handler(|request| {