diff --git a/rust/lancedb/src/connection.rs b/rust/lancedb/src/connection.rs index c21c4772a..e00d899e0 100644 --- a/rust/lancedb/src/connection.rs +++ b/rust/lancedb/src/connection.rs @@ -28,7 +28,7 @@ use crate::database::{ }; use crate::embeddings::{EmbeddingRegistry, MemoryRegistry}; use crate::error::{Error, Result}; -use crate::function::RegisterFunctionJobSpec; +use crate::function::{Function, FunctionId, RegisterFunctionJobSpec}; #[cfg(feature = "remote")] use crate::remote::{ client::ClientConfig, @@ -563,6 +563,33 @@ impl Connection { self.internal.register_function(spec).await } + /// Look up the Function currently bound to a database-scoped name. + /// + /// The name is lookup indirection only and is never part of the returned + /// [`Function`]. Empty names return [`Error::InvalidInput`] before backend + /// dispatch. Only remote databases support enterprise catalog lookup; + /// nonempty local lookups return [`Error::NotSupported`]. + pub async fn lookup_function_by_name(&self, name: impl AsRef) -> Result { + let name = name.as_ref(); + // Public nonempty invariant: validate before any Database backend sees + // the call so local and remote Connections agree on InvalidInput. + if name.is_empty() { + return Err(Error::InvalidInput { + message: "function lookup name must be non-empty".into(), + }); + } + self.internal.lookup_function_by_name(name).await + } + + /// Look up an immutable Function by exact opaque [`FunctionId`]. + /// + /// Exact-ID lookup is independent of later catalog name changes. Only + /// remote databases support enterprise catalog lookup; local databases + /// return [`Error::NotSupported`]. + pub async fn lookup_function_by_id(&self, function_id: &FunctionId) -> Result { + self.internal.lookup_function_by_id(function_id).await + } + /// Drop a table in the database. /// /// # Arguments diff --git a/rust/lancedb/src/database.rs b/rust/lancedb/src/database.rs index 8b16009c5..3d4534aa3 100644 --- a/rust/lancedb/src/database.rs +++ b/rust/lancedb/src/database.rs @@ -30,7 +30,7 @@ use lance_namespace::models::{ use crate::data::scannable::Scannable; use crate::error::Result; -use crate::function::RegisterFunctionJobSpec; +use crate::function::{Function, FunctionId, RegisterFunctionJobSpec}; use crate::table::{BaseTable, WriteOptions}; pub mod listing; @@ -325,6 +325,31 @@ pub trait Database: async fn register_function(&self, _spec: RegisterFunctionJobSpec) -> Result { job_op_not_supported("register_function") } + /// Look up the Function currently bound to a database-scoped name. + /// + /// The name is lookup indirection only and is never part of the returned + /// [`Function`]. Empty names return [`crate::Error::InvalidInput`] before + /// the unsupported fallback so local and remote backends agree. Nonempty + /// names on databases without enterprise catalog lookup return + /// [`crate::Error::NotSupported`]. + async fn lookup_function_by_name(&self, name: &str) -> Result { + // Public nonempty invariant on the Database trait seam itself: + // Connection::database() exposes Arc, so empty-name + // rejection cannot rely solely on Connection prevalidation. + if name.is_empty() { + return Err(crate::error::Error::InvalidInput { + message: "function lookup name must be non-empty".into(), + }); + } + job_op_not_supported("lookup_function_by_name") + } + /// Look up an immutable Function by exact opaque [`FunctionId`]. + /// + /// Exact-ID lookup is independent of later catalog name changes. Local + /// databases do not support enterprise catalog lookup. + async fn lookup_function_by_id(&self, _function_id: &FunctionId) -> Result { + job_op_not_supported("lookup_function_by_id") + } /// Open a table in the database async fn open_table(&self, request: OpenTableRequest) -> Result>; /// Rename a table in the database diff --git a/rust/lancedb/src/remote.rs b/rust/lancedb/src/remote.rs index 4b5f8832f..57718370e 100644 --- a/rust/lancedb/src/remote.rs +++ b/rust/lancedb/src/remote.rs @@ -8,6 +8,7 @@ pub(crate) mod client; pub(crate) mod db; +pub(crate) mod function; pub(crate) mod job; pub mod oauth; mod retry; diff --git a/rust/lancedb/src/remote/client.rs b/rust/lancedb/src/remote/client.rs index a993b0876..2abac50d2 100644 --- a/rust/lancedb/src/remote/client.rs +++ b/rust/lancedb/src/remote/client.rs @@ -791,6 +791,41 @@ impl RestfulLanceDbClient { Ok((request_id, response)) } + /// Send one attempt with a caller-owned request id. + /// + /// Used by Function lookup so the helper can inspect non-success bodies for + /// an explicit `error_code` before deciding whether to retry. Does not log + /// request bodies or headers (lookup selectors must stay out of logs) and + /// does not interpret HTTP status. + pub(crate) async fn send_attempt_with_request_id( + &self, + req_builder: RequestBuilder, + request_id: &str, + ) -> Result { + let (client, request) = req_builder.build_split(); + let mut request = request.map_err(|e| Error::Runtime { + message: format!("Failed to build request: {}", e), + })?; + self.set_request_id(&mut request, request_id); + request = self.apply_dynamic_headers(request).await?; + if log::log_enabled!(log::Level::Debug) { + debug!( + "{}", + format_request_log(&request, request_id, RequestPrivacy::Sensitive) + ); + } + let response = self + .sender + .send(&client, request) + .await + .err_to_http(request_id.to_string())?; + debug!( + "Received response for request_id={}: {:?}", + request_id, response + ); + Ok(response) + } + /// Send the request using retries configured in the RetryConfig. /// If retry_5xx is false, 5xx requests will not be retried regardless of the statuses configured /// in the RetryConfig. diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index c87653e26..8871d537d 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -23,7 +23,7 @@ use crate::database::{ JobDescription, JobInfo, OpenTableRequest, ReadConsistency, TableNamesRequest, }; use crate::error::Result; -use crate::function::RegisterFunctionJobSpec; +use crate::function::{Function, FunctionId, RegisterFunctionJobSpec}; use crate::remote::util::stream_as_body; use crate::table::BaseTable; @@ -617,6 +617,16 @@ impl Database for RemoteDatabase { )))) } + async fn lookup_function_by_name(&self, name: &str) -> Result { + let selector = super::function::FunctionLookupSelector::by_name(name)?; + super::function::lookup_function(&self.client, selector).await + } + + async fn lookup_function_by_id(&self, function_id: &FunctionId) -> Result { + let selector = super::function::FunctionLookupSelector::by_function_id(function_id); + super::function::lookup_function(&self.client, selector).await + } + async fn table_names(&self, request: TableNamesRequest) -> Result> { let mut req = if !request.namespace_path.is_empty() { let namespace_id = @@ -3406,4 +3416,766 @@ mod tests { "flow must be submit then describe only, with no Function name lookup" ); } + + // ------------------------------------------------------------------------- + // Function lookup transport (RED until lookup_function_by_{name,id}) + // ------------------------------------------------------------------------- + + const LOOKUP_CATALOG_NAME: &str = "text.normalize.lookup-name"; + const LOOKUP_FUNCTION_ID: &str = "fn.exact.lookup-handle"; + const LOOKUP_SERVER_MESSAGE_MARKER: &str = + "SERVER_LOOKUP_DIAGNOSTIC_MARKER name=text.normalize.lookup-name id=fn.exact.lookup-handle"; + + fn sample_lookup_function() -> Function { + let id = FunctionId::try_new(LOOKUP_FUNCTION_ID).expect("valid FunctionId"); + let signature = FunctionSignature::try_new( + vec![ + FunctionParameter::new("text", DataType::Utf8), + FunctionParameter::new("limit", DataType::Int32), + ], + FunctionOutput::new(DataType::Utf8, true), + ) + .expect("valid FunctionSignature"); + Function::new(id, signature) + } + + fn lookup_function_wire(function: &Function) -> Value { + serde_json::to_value(function).expect("serialize Function wire") + } + + fn lookup_success_body(function: &Function, extra_outer: Option) -> String { + let mut body = json!({ + "function": lookup_function_wire(function), + }); + if let Some(Value::Object(extra)) = extra_outer { + let object = body.as_object_mut().expect("lookup success must be object"); + for (k, v) in extra { + object.insert(k, v); + } + } + body.to_string() + } + + fn lookup_error_chain_text(err: &Error) -> String { + let mut text = format!("{err}\n{err:?}"); + let mut current: Option<&(dyn std::error::Error + 'static)> = Some(err); + while let Some(e) = current { + text.push('\n'); + text.push_str(&e.to_string()); + text.push('\n'); + text.push_str(&format!("{e:?}")); + current = e.source(); + } + text + } + + fn assert_lookup_payload_free(err: &Error) { + let text = lookup_error_chain_text(err); + assert!( + !text.contains(LOOKUP_SERVER_MESSAGE_MARKER), + "server diagnostic marker must be absent from error/debug/source chain: {text}" + ); + assert!( + !text.contains(LOOKUP_CATALOG_NAME), + "catalog name must be absent from error/debug/source chain: {text}" + ); + assert!( + !text.contains(LOOKUP_FUNCTION_ID), + "FunctionId must be absent from error/debug/source chain: {text}" + ); + assert!( + !text.contains("SENSITIVE_LOOKUP_BODY_MARKER"), + "non-success/malformed body marker must be absent from error/debug/source chain: {text}" + ); + } + + fn assert_lookup_name_request(request: &reqwest::Request, expected_name: &str) { + assert_eq!(request.method(), &reqwest::Method::POST); + assert_eq!(request.url().path(), "/v1/functions/lookup"); + let body = request + .body() + .and_then(|b| b.as_bytes()) + .expect("lookup request must carry a JSON body"); + let actual: Value = serde_json::from_slice(body).expect("lookup body must be JSON"); + assert_eq!( + actual, + json!({ "name": expected_name }), + "name lookup body must be exactly {{\"name\":...}}" + ); + assert!( + actual.get("function_id").is_none(), + "name lookup must not send function_id: {actual}" + ); + } + + fn assert_lookup_id_request(request: &reqwest::Request, expected_id: &str) { + assert_eq!(request.method(), &reqwest::Method::POST); + assert_eq!(request.url().path(), "/v1/functions/lookup"); + let body = request + .body() + .and_then(|b| b.as_bytes()) + .expect("lookup request must carry a JSON body"); + let actual: Value = serde_json::from_slice(body).expect("lookup body must be JSON"); + assert_eq!( + actual, + json!({ "function_id": expected_id }), + "id lookup body must be exactly {{\"function_id\":...}}" + ); + assert!( + actual.get("name").is_none(), + "id lookup must not send name: {actual}" + ); + } + + /// Name lookup uses exact path/body and returns the immutable Function value. + #[tokio::test] + async fn lookup_function_by_name_posts_exact_body_and_returns_function_without_name() { + let expected = sample_lookup_function(); + let body = lookup_success_body(&expected, None); + let conn = Connection::new_with_handler(move |request| { + assert_lookup_name_request(&request, LOOKUP_CATALOG_NAME); + http::Response::builder() + .status(200) + .body(body.clone()) + .unwrap() + }); + + let function = conn + .lookup_function_by_name(LOOKUP_CATALOG_NAME) + .await + .expect("name lookup must succeed"); + assert_exact_function(&function, &expected); + + let debug = format!("{function:?}"); + let wire = serde_json::to_value(&function).expect("Function serializes"); + assert!( + !debug.contains(LOOKUP_CATALOG_NAME), + "returned Function debug must not carry the catalog name: {debug}" + ); + assert!( + wire.get("name").is_none(), + "returned Function wire must not include a name field: {wire}" + ); + assert_eq!(function.id().as_str(), LOOKUP_FUNCTION_ID); + } + + /// Exact-ID lookup uses exact path/body and is independent of catalog name. + #[tokio::test] + async fn lookup_function_by_id_posts_exact_body_and_returns_function() { + let expected = sample_lookup_function(); + let body = lookup_success_body(&expected, None); + let id = FunctionId::try_new(LOOKUP_FUNCTION_ID).expect("valid FunctionId"); + let conn = Connection::new_with_handler(move |request| { + assert_lookup_id_request(&request, LOOKUP_FUNCTION_ID); + http::Response::builder() + .status(200) + .body(body.clone()) + .unwrap() + }); + + let function = conn + .lookup_function_by_id(&id) + .await + .expect("id lookup must succeed"); + assert_exact_function(&function, &expected); + let debug = format!("{function:?}"); + assert!( + !debug.contains(LOOKUP_CATALOG_NAME), + "id lookup Function must not invent a catalog name: {debug}" + ); + } + + /// Unknown outer success fields are accepted; Function decoding stays strict. + #[tokio::test] + async fn lookup_function_accepts_additive_outer_success_fields() { + let expected = sample_lookup_function(); + let body = lookup_success_body( + &expected, + Some(json!({ + "server_extra": {"ok": true}, + "request_echo_name": LOOKUP_CATALOG_NAME, + })), + ); + let conn = Connection::new_with_handler(move |request| { + assert_lookup_name_request(&request, LOOKUP_CATALOG_NAME); + http::Response::builder() + .status(200) + .body(body.clone()) + .unwrap() + }); + + let function = conn + .lookup_function_by_name(LOOKUP_CATALOG_NAME) + .await + .expect("additive outer success must decode"); + assert_exact_function(&function, &expected); + } + + /// Empty selectors fail closed before any HTTP request is issued. + #[tokio::test] + async fn lookup_function_empty_selectors_fail_before_transport() { + let attempts = Arc::new(AtomicUsize::new(0)); + let attempts_ref = attempts.clone(); + let conn = Connection::new_with_handler(move |_request| -> http::Response { + attempts_ref.fetch_add(1, Ordering::SeqCst); + panic!("empty selector must not issue an HTTP request"); + }); + + let err = conn + .lookup_function_by_name("") + .await + .expect_err("empty name must fail before transport"); + assert!( + matches!(err, Error::InvalidInput { .. }), + "empty name must be InvalidInput, got {err:?}" + ); + assert_eq!( + attempts.load(Ordering::SeqCst), + 0, + "empty name must not touch transport" + ); + + // Empty FunctionId cannot be constructed; lookup by id is therefore + // unreachable with an empty selector. Keep the contract explicit. + let empty_id = FunctionId::try_new("").expect_err("empty FunctionId rejected"); + assert!( + matches!(empty_id, Error::InvalidInput { .. }), + "empty FunctionId must be InvalidInput, got {empty_id:?}" + ); + assert_eq!( + attempts.load(Ordering::SeqCst), + 0, + "empty FunctionId construction must not touch transport" + ); + } + + /// Known explicit not-found code becomes Error::Function; status/message ignored. + #[tokio::test] + async fn lookup_function_explicit_not_found_code_is_function_error() { + let body = json!({ + "error_code": "name_or_function_not_found", + "message": LOOKUP_SERVER_MESSAGE_MARKER, + "looks_like": "definition_validation_failure", + }) + .to_string(); + let conn = Connection::new_with_handler(move |request| { + assert_lookup_name_request(&request, LOOKUP_CATALOG_NAME); + http::Response::builder() + .status(404) + .body(body.clone()) + .unwrap() + }); + + let err = conn + .lookup_function_by_name(LOOKUP_CATALOG_NAME) + .await + .expect_err("missing name must fail"); + match &err { + Error::Function { code, message } => { + assert_eq!(code.as_str(), "name_or_function_not_found"); + assert_ne!(code.as_str(), "definition_validation_failure"); + assert!( + !message.contains(LOOKUP_SERVER_MESSAGE_MARKER), + "Function error message must be sanitized, got {message}" + ); + assert!( + !message.contains(LOOKUP_CATALOG_NAME), + "Function error message must not echo the catalog name, got {message}" + ); + } + other => panic!("expected Error::Function, got {other:?}"), + } + assert_lookup_payload_free(&err); + } + + /// Unknown nonempty explicit code is preserved; HTTP status does not override it. + #[tokio::test] + async fn lookup_function_preserves_unknown_explicit_code_despite_status_and_message() { + let raw = "enterprise_future_lookup_category_xyz"; + let body = json!({ + "error_code": raw, + "message": format!( + "{LOOKUP_SERVER_MESSAGE_MARKER} revoked_function name_or_function_not_found" + ), + }) + .to_string(); + let conn = Connection::new_with_handler(move |request| { + assert_lookup_id_request(&request, LOOKUP_FUNCTION_ID); + http::Response::builder() + .status(409) + .body(body.clone()) + .unwrap() + }); + let id = FunctionId::try_new(LOOKUP_FUNCTION_ID).expect("valid FunctionId"); + + let err = conn + .lookup_function_by_id(&id) + .await + .expect_err("unknown explicit code must surface"); + match &err { + Error::Function { code, message } => { + assert_eq!(code.as_str(), raw); + assert!( + matches!(code, FunctionErrorCode::Unrecognized(_)), + "unknown code must stay Unrecognized, got {code:?}" + ); + assert!( + !message.contains(LOOKUP_SERVER_MESSAGE_MARKER), + "diagnostic message must be sanitized" + ); + assert_ne!(code.as_str(), "revoked_function"); + assert_ne!(code.as_str(), "name_or_function_not_found"); + } + other => panic!("expected Error::Function, got {other:?}"), + } + assert_lookup_payload_free(&err); + } + + /// Missing/empty/wrong-type error_code on non-success stays payload-free Http. + #[tokio::test] + async fn lookup_function_missing_or_invalid_error_code_is_payload_free_http() { + let cases: Vec<(&str, u16, String)> = vec![ + ( + "missing_code_404", + 404, + json!({ + "message": LOOKUP_SERVER_MESSAGE_MARKER, + "SENSITIVE_LOOKUP_BODY_MARKER": true, + }) + .to_string(), + ), + ( + "empty_code", + 400, + json!({ + "error_code": "", + "message": LOOKUP_SERVER_MESSAGE_MARKER, + "SENSITIVE_LOOKUP_BODY_MARKER": true, + }) + .to_string(), + ), + ( + "wrong_type_code", + 400, + json!({ + "error_code": 123, + "message": LOOKUP_SERVER_MESSAGE_MARKER, + "SENSITIVE_LOOKUP_BODY_MARKER": true, + }) + .to_string(), + ), + ( + "null_code", + 404, + json!({ + "error_code": null, + "message": LOOKUP_SERVER_MESSAGE_MARKER, + "SENSITIVE_LOOKUP_BODY_MARKER": true, + }) + .to_string(), + ), + ( + "non_json", + 404, + format!("not-json {LOOKUP_SERVER_MESSAGE_MARKER} SENSITIVE_LOOKUP_BODY_MARKER"), + ), + ]; + + let mut unexpected = Vec::new(); + for (label, status, response_body) in cases { + let body_for_handler = response_body.clone(); + let conn = Connection::new_with_handler(move |request| { + assert_lookup_name_request(&request, LOOKUP_CATALOG_NAME); + http::Response::builder() + .status(status) + .body(body_for_handler.clone()) + .unwrap() + }); + match conn.lookup_function_by_name(LOOKUP_CATALOG_NAME).await { + Err(err @ Error::Http { .. }) => assert_lookup_payload_free(&err), + Err(Error::Function { .. }) => unexpected.push(format!( + "{label}: must not invent Error::Function without explicit code" + )), + other => unexpected.push(format!("{label}: {other:?}")), + } + } + assert!( + unexpected.is_empty(), + "invalid/missing error_code must stay payload-free Http: {unexpected:?}" + ); + } + + /// Malformed/missing/invalid success payloads stay payload-free Http. + #[tokio::test] + async fn lookup_function_invalid_success_payload_is_payload_free_http() { + let cases: Vec<(&str, String)> = vec![ + ( + "missing_function", + json!({ + "server_extra": true, + "SENSITIVE_LOOKUP_BODY_MARKER": LOOKUP_SERVER_MESSAGE_MARKER, + }) + .to_string(), + ), + ( + "null_function", + json!({ + "function": null, + "SENSITIVE_LOOKUP_BODY_MARKER": LOOKUP_SERVER_MESSAGE_MARKER, + }) + .to_string(), + ), + ( + "wrong_type_function", + json!({ + "function": "not-an-object", + "SENSITIVE_LOOKUP_BODY_MARKER": LOOKUP_SERVER_MESSAGE_MARKER, + }) + .to_string(), + ), + ( + "invalid_function_unknown_field", + json!({ + "function": { + "format_version": 1, + "id": LOOKUP_FUNCTION_ID, + "signature": { + "parameters": [], + "output": { + "data_type_ipc": lookup_function_wire(&sample_lookup_function()) + ["signature"]["output"]["data_type_ipc"].clone(), + "nullable": true + } + }, + "name": LOOKUP_CATALOG_NAME, + "SENSITIVE_LOOKUP_BODY_MARKER": true + } + }) + .to_string(), + ), + ( + "malformed_json", + format!("not-json {LOOKUP_SERVER_MESSAGE_MARKER} SENSITIVE_LOOKUP_BODY_MARKER"), + ), + ]; + + let mut unexpected = Vec::new(); + for (label, response_body) in cases { + let body_for_handler = response_body.clone(); + let conn = Connection::new_with_handler(move |request| { + assert_lookup_name_request(&request, LOOKUP_CATALOG_NAME); + http::Response::builder() + .status(200) + .body(body_for_handler.clone()) + .unwrap() + }); + match conn.lookup_function_by_name(LOOKUP_CATALOG_NAME).await { + Err(err @ Error::Http { .. }) => assert_lookup_payload_free(&err), + other => unexpected.push(format!("{label}: {other:?}")), + } + } + assert!( + unexpected.is_empty(), + "invalid success payloads must fail closed as payload-free Http: {unexpected:?}" + ); + } + + /// Local databases reject both lookup seams without mutating state. + #[tokio::test] + async fn lookup_function_local_database_returns_not_supported_without_mutation() { + let dir = tempfile::tempdir().expect("tempdir"); + let conn = ConnectBuilder::new(dir.path().to_str().unwrap()) + .execute() + .await + .expect("local connect"); + let before = conn + .table_names() + .execute() + .await + .expect("table_names before"); + assert!(before.is_empty()); + + let err_name = conn + .lookup_function_by_name(LOOKUP_CATALOG_NAME) + .await + .expect_err("local name lookup must be unsupported"); + assert!( + matches!(err_name, Error::NotSupported { .. }), + "expected NotSupported for name lookup, got {err_name:?}" + ); + + let id = FunctionId::try_new(LOOKUP_FUNCTION_ID).expect("valid FunctionId"); + let err_id = conn + .lookup_function_by_id(&id) + .await + .expect_err("local id lookup must be unsupported"); + assert!( + matches!(err_id, Error::NotSupported { .. }), + "expected NotSupported for id lookup, got {err_id:?}" + ); + + let after = conn + .table_names() + .execute() + .await + .expect("table_names after"); + assert_eq!(before, after, "unsupported lookup must not mutate tables"); + } + + /// Empty name is a public Connection invariant: InvalidInput on local too. + #[tokio::test] + async fn lookup_function_local_empty_name_is_invalid_input_without_mutation() { + let dir = tempfile::tempdir().expect("tempdir"); + let conn = ConnectBuilder::new(dir.path().to_str().unwrap()) + .execute() + .await + .expect("local connect"); + let before = conn + .table_names() + .execute() + .await + .expect("table_names before"); + assert!(before.is_empty()); + + let err = conn + .lookup_function_by_name("") + .await + .expect_err("empty name must fail before backend dispatch"); + assert!( + matches!(err, Error::InvalidInput { .. }), + "empty name must be InvalidInput on local Connection, got {err:?}" + ); + + let after = conn + .table_names() + .execute() + .await + .expect("table_names after"); + assert_eq!(before, after, "empty-name rejection must not mutate tables"); + } + + /// Database trait seam must reject empty names as InvalidInput (not NotSupported). + /// + /// `Connection::database()` exposes `Arc`; callers that bypass + /// Connection prevalidation must still get backend-independent InvalidInput. + #[tokio::test] + async fn lookup_function_local_database_trait_empty_name_is_invalid_input_without_mutation() { + let dir = tempfile::tempdir().expect("tempdir"); + let conn = ConnectBuilder::new(dir.path().to_str().unwrap()) + .execute() + .await + .expect("local connect"); + let before = conn + .table_names() + .execute() + .await + .expect("table_names before"); + assert!(before.is_empty()); + + let err = conn + .database() + .lookup_function_by_name("") + .await + .expect_err("empty name must fail on Database trait default"); + assert!( + matches!(err, Error::InvalidInput { .. }), + "empty name via Database trait must be InvalidInput, got {err:?}" + ); + + let after = conn + .table_names() + .execute() + .await + .expect("table_names after"); + assert_eq!( + before, after, + "Database-trait empty-name rejection must not mutate tables" + ); + } + + /// One retry keeps the SDK-generated request id and exact body before success. + #[tokio::test] + async fn lookup_function_retry_preserves_request_id_and_body() { + let expected = sample_lookup_function(); + let success_body = lookup_success_body(&expected, None); + let seen_request_id = Arc::new(OnceLock::new()); + let seen_request_id_ref = seen_request_id.clone(); + let attempts = Arc::new(AtomicUsize::new(0)); + let attempts_ref = attempts.clone(); + + let conn = Connection::new_with_handler_and_config( + move |request| { + assert_lookup_name_request(&request, LOOKUP_CATALOG_NAME); + let request_id = request.headers()["x-request-id"] + .to_str() + .unwrap() + .to_string(); + assert!(!request_id.is_empty(), "SDK must generate a request id"); + let seen = seen_request_id_ref.get_or_init(|| request_id.clone()); + assert_eq!( + &request_id, seen, + "request id must be identical across retries" + ); + + let n = attempts_ref.fetch_add(1, Ordering::SeqCst); + if n == 0 { + http::Response::builder() + .status(500) + .body(format!( + "{LOOKUP_SERVER_MESSAGE_MARKER} SENSITIVE_LOOKUP_BODY_MARKER" + )) + .unwrap() + } else { + http::Response::builder() + .status(200) + .body(success_body.clone()) + .unwrap() + } + }, + ClientConfig { + retry_config: RetryConfig { + retries: Some(2), + backoff_factor: Some(0.0), + backoff_jitter: Some(0.0), + ..Default::default() + }, + ..Default::default() + }, + ); + + let function = conn + .lookup_function_by_name(LOOKUP_CATALOG_NAME) + .await + .expect("lookup must succeed after one retry"); + assert_exact_function(&function, &expected); + assert_eq!(attempts.load(Ordering::SeqCst), 2); + assert!(seen_request_id.get().is_some()); + } + + /// Response-body read failures consume the read budget, keep request id/body. + #[tokio::test] + async fn lookup_function_retries_response_body_read_failure_with_read_budget() { + let expected = sample_lookup_function(); + let success_body = lookup_success_body(&expected, None); + let seen_request_id = Arc::new(OnceLock::new()); + let seen_request_id_ref = seen_request_id.clone(); + let attempts = Arc::new(AtomicUsize::new(0)); + let attempts_ref = attempts.clone(); + + let conn = Connection::new_with_handler_and_config( + move |request| { + assert_lookup_name_request(&request, LOOKUP_CATALOG_NAME); + let request_id = request.headers()["x-request-id"] + .to_str() + .unwrap() + .to_string(); + assert!(!request_id.is_empty(), "SDK must generate a request id"); + let seen = seen_request_id_ref.get_or_init(|| request_id.clone()); + assert_eq!( + &request_id, seen, + "request id must be identical across retries" + ); + + let n = attempts_ref.fetch_add(1, Ordering::SeqCst); + if n == 0 { + let stream = futures::stream::once(async { + Err::(std::io::Error::other( + "simulated lookup response body read failure", + )) + }); + http::Response::builder() + .status(200) + .body(reqwest::Body::wrap_stream(stream)) + .unwrap() + } else { + http::Response::builder() + .status(200) + .body(reqwest::Body::from(success_body.clone())) + .unwrap() + } + }, + ClientConfig { + // retries=1 would exhaust immediately if body-read were misclassified + // as a request failure; read_retries=2 allows one read failure then success. + retry_config: RetryConfig { + retries: Some(1), + read_retries: Some(2), + connect_retries: Some(1), + backoff_factor: Some(0.0), + backoff_jitter: Some(0.0), + ..Default::default() + }, + ..Default::default() + }, + ); + + let function = conn + .lookup_function_by_name(LOOKUP_CATALOG_NAME) + .await + .expect("lookup must succeed after one response-body read retry"); + assert_exact_function(&function, &expected); + assert_eq!( + attempts.load(Ordering::SeqCst), + 2, + "body-read failure must consume read budget and retry once" + ); + assert!(seen_request_id.get().is_some()); + } + + /// Nonretryable client/header errors must not be repeated. + #[tokio::test] + async fn lookup_function_nonretryable_client_error_is_not_repeated() { + let calls = Arc::new(AtomicUsize::new(0)); + let calls_ref = calls.clone(); + + #[derive(Debug)] + struct CountingErrorHeaderProvider { + calls: Arc, + } + + #[async_trait::async_trait] + impl HeaderProvider for CountingErrorHeaderProvider { + async fn get_headers(&self) -> crate::Result> { + self.calls.fetch_add(1, Ordering::SeqCst); + Err(Error::Runtime { + message: "Failed to fetch auth token".to_string(), + }) + } + } + + let conn = Connection::new_with_handler_and_config( + move |_request| -> http::Response<&'static str> { + panic!("lookup must not reach transport when header provider fails"); + }, + ClientConfig { + header_provider: Some(Arc::new(CountingErrorHeaderProvider { calls: calls_ref }) + as Arc), + retry_config: RetryConfig { + retries: Some(3), + connect_retries: Some(3), + read_retries: Some(3), + backoff_factor: Some(0.0), + backoff_jitter: Some(0.0), + ..Default::default() + }, + ..Default::default() + }, + ); + + let err = conn + .lookup_function_by_name(LOOKUP_CATALOG_NAME) + .await + .expect_err("header provider failure must surface"); + match err { + Error::Runtime { message } => { + assert_eq!(message, "Failed to fetch auth token"); + } + other => panic!("expected Runtime from header provider, got {other:?}"), + } + assert_eq!( + calls.load(Ordering::SeqCst), + 1, + "nonretryable client/header error must not be repeated" + ); + } } diff --git a/rust/lancedb/src/remote/function.rs b/rust/lancedb/src/remote/function.rs new file mode 100644 index 000000000..a46b5d2d8 --- /dev/null +++ b/rust/lancedb/src/remote/function.rs @@ -0,0 +1,209 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +//! Remote first-class Function catalog lookup wire helper. +//! +//! POST `/v1/functions/lookup` resolves a database-scoped name or exact +//! [`FunctionId`] to an immutable [`Function`] value. Name is lookup +//! indirection only and never becomes part of the returned handle. + +use serde::Deserialize; +use serde_json::Value; + +use crate::error::{Error, FunctionErrorCode, Result}; +use crate::function::{Function, FunctionId}; + +use super::client::{HttpSend, RestfulLanceDbClient}; +use super::retry::RetryCounter; + +const LOOKUP_PATH: &str = "/v1/functions/lookup"; + +/// Fixed client diagnostic for [`Error::Function`]. Never carry server text, +/// selector values, or response payload bytes. +const LOOKUP_FUNCTION_ERROR_MESSAGE: &str = "function lookup failed"; + +/// Fixed client diagnostic for protocol / HTTP failures. Never include response +/// payload bytes or selector values. +const LOOKUP_HTTP_ERROR_MESSAGE: &str = "function lookup request failed"; + +/// Fixed client diagnostic for malformed success payloads. +const LOOKUP_INVALID_SUCCESS_MESSAGE: &str = "function lookup response missing or invalid function"; + +/// One exact lookup selector. Exactly one variant is serialized on the wire. +pub enum FunctionLookupSelector { + Name(String), + FunctionId(String), +} + +impl FunctionLookupSelector { + pub fn by_name(name: impl Into) -> Result { + let name = name.into(); + if name.is_empty() { + return Err(Error::InvalidInput { + message: "function lookup name must be non-empty".into(), + }); + } + Ok(Self::Name(name)) + } + + pub fn by_function_id(function_id: &FunctionId) -> Self { + Self::FunctionId(function_id.as_str().to_string()) + } + + fn to_wire(&self) -> Value { + match self { + Self::Name(name) => serde_json::json!({ "name": name }), + Self::FunctionId(function_id) => { + serde_json::json!({ "function_id": function_id }) + } + } + } +} + +#[derive(Deserialize)] +struct LookupSuccessResponse { + function: Function, +} + +/// Resolve a Function via POST `/v1/functions/lookup`. +/// +/// Transport classification matches [`RestfulLanceDbClient::send_with_retry`]: +/// connect → connect_retries; timeout/body/decode (including response-byte +/// reads) → read_retries; configured retryable statuses without an explicit +/// `error_code` → request retries; all other transport/client errors return +/// immediately. An explicit nonempty `error_code` is terminal and wins over +/// HTTP status. Request/response payload bytes never enter error chains. +pub async fn lookup_function( + client: &RestfulLanceDbClient, + selector: FunctionLookupSelector, +) -> Result { + let req_builder = client.post(LOOKUP_PATH).json(&selector.to_wire()); + + let tmp_req = req_builder.try_clone().ok_or_else(|| Error::Runtime { + message: "Attempted to retry a request that cannot be cloned".to_string(), + })?; + let (_, built) = tmp_req.build_split(); + let mut built = built.map_err(|e| Error::Runtime { + message: format!("Failed to build request: {}", e), + })?; + let request_id = client.extract_request_id(&mut built); + let mut retry_counter = RetryCounter::new(&client.retry_config, request_id.clone()); + + loop { + let attempt = req_builder.try_clone().ok_or_else(|| Error::Runtime { + message: "Attempted to retry a request that cannot be cloned".to_string(), + })?; + + let rsp = match client + .send_attempt_with_request_id(attempt, &retry_counter.request_id) + .await + { + Ok(rsp) => rsp, + Err(err) => { + classify_lookup_send_error(&mut retry_counter, err)?; + tokio::time::sleep(retry_counter.next_sleep_time()).await; + continue; + } + }; + + let status = rsp.status(); + let bytes = match rsp.bytes().await { + Ok(bytes) => bytes, + Err(err) => { + // Response body/decode failures share the read budget with + // send-time timeout/body/decode errors. + retry_counter.increment_read_failures(err)?; + tokio::time::sleep(retry_counter.next_sleep_time()).await; + continue; + } + }; + + if status.is_success() { + return decode_lookup_success(&bytes, retry_counter.request_id); + } + + // Explicit nonempty error_code wins over HTTP status and precludes retry. + if let Some(code) = explicit_error_code(&bytes) { + return Err(Error::Function { + code, + message: LOOKUP_FUNCTION_ERROR_MESSAGE.to_string(), + }); + } + + if client.retry_config.statuses.contains(&status) { + let source = Error::Http { + source: LOOKUP_HTTP_ERROR_MESSAGE.into(), + request_id: retry_counter.request_id.clone(), + status_code: Some(status), + }; + retry_counter.increment_request_failures(source)?; + tokio::time::sleep(retry_counter.next_sleep_time()).await; + continue; + } + + return Err(Error::Http { + source: LOOKUP_HTTP_ERROR_MESSAGE.into(), + request_id: retry_counter.request_id, + status_code: Some(status), + }); + } +} + +/// Classify a send-attempt error using the same buckets as `send_with_retry`. +/// +/// Returns `Ok(())` when the caller should sleep and retry. Returns `Err` for +/// nonretryable failures or when a retry budget is exhausted (no extra attempt). +fn classify_lookup_send_error(retry_counter: &mut RetryCounter<'_>, err: Error) -> Result<()> { + match err { + Error::Http { + source, + request_id, + status_code, + } => match source.downcast::() { + Ok(reqwest_err) if reqwest_err.is_connect() => { + retry_counter.increment_connect_failures(*reqwest_err) + } + Ok(reqwest_err) + if reqwest_err.is_timeout() || reqwest_err.is_body() || reqwest_err.is_decode() => + { + retry_counter.increment_read_failures(*reqwest_err) + } + Ok(reqwest_err) => Err(Error::Http { + source: Box::new(*reqwest_err), + request_id, + status_code, + }), + Err(source) => Err(Error::Http { + source, + request_id, + status_code, + }), + }, + // Header-provider / client failures are not transport retries. + other => Err(other), + } +} + +fn decode_lookup_success(bytes: &[u8], request_id: String) -> Result { + match serde_json::from_slice::(bytes) { + Ok(body) => Ok(body.function), + Err(_) => Err(Error::Http { + source: LOOKUP_INVALID_SUCCESS_MESSAGE.into(), + request_id, + status_code: None, + }), + } +} + +/// Decode a stable category only from an explicit nonempty string `error_code`. +/// Missing, empty, wrong-type, or non-JSON bodies yield [`None`]. +fn explicit_error_code(bytes: &[u8]) -> Option { + let value: Value = serde_json::from_slice(bytes).ok()?; + let code = value.get("error_code")?; + match code { + Value::String(raw) if !raw.is_empty() => { + serde_json::from_value(Value::String(raw.clone())).ok() + } + _ => None, + } +}