From 65c0968c0f0212b956aff66b6b74978957f77f9a Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Wed, 12 Aug 2026 05:03:17 +0800 Subject: [PATCH] feat: submit function registration jobs --- rust/lancedb/src/connection.rs | 13 + rust/lancedb/src/database.rs | 8 + rust/lancedb/src/remote/client.rs | 387 ++++++++++++++++++++++++-- rust/lancedb/src/remote/db.rs | 448 +++++++++++++++++++++++++++++- 4 files changed, 832 insertions(+), 24 deletions(-) diff --git a/rust/lancedb/src/connection.rs b/rust/lancedb/src/connection.rs index dd53a2d2e..c21c4772a 100644 --- a/rust/lancedb/src/connection.rs +++ b/rust/lancedb/src/connection.rs @@ -28,6 +28,7 @@ use crate::database::{ }; use crate::embeddings::{EmbeddingRegistry, MemoryRegistry}; use crate::error::{Error, Result}; +use crate::function::RegisterFunctionJobSpec; #[cfg(feature = "remote")] use crate::remote::{ client::ClientConfig, @@ -550,6 +551,18 @@ impl Connection { self.internal.job_history(job_id).await } + /// Submit a first-class Function registration job. + /// + /// Returns a [`crate::job::Job`] handle for the accepted server-side job. + /// Only remote databases support registration; local databases return + /// [`Error::NotSupported`]. + pub async fn register_function( + &self, + spec: RegisterFunctionJobSpec, + ) -> Result { + self.internal.register_function(spec).await + } + /// Drop a table in the database. /// /// # Arguments diff --git a/rust/lancedb/src/database.rs b/rust/lancedb/src/database.rs index a69c623a2..8b16009c5 100644 --- a/rust/lancedb/src/database.rs +++ b/rust/lancedb/src/database.rs @@ -30,6 +30,7 @@ use lance_namespace::models::{ use crate::data::scannable::Scannable; use crate::error::Result; +use crate::function::RegisterFunctionJobSpec; use crate::table::{BaseTable, WriteOptions}; pub mod listing; @@ -317,6 +318,13 @@ pub trait Database: async fn job_history(&self, _job_id: Option<&str>) -> Result> { job_op_not_supported("job_history") } + /// Submit a first-class Function registration job. + /// + /// Returns a [`crate::job::Job`] handle for the accepted server-side job. + /// Local databases do not support registration. + async fn register_function(&self, _spec: RegisterFunctionJobSpec) -> Result { + job_op_not_supported("register_function") + } /// 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/client.rs b/rust/lancedb/src/remote/client.rs index 57dd89890..a993b0876 100644 --- a/rust/lancedb/src/remote/client.rs +++ b/rust/lancedb/src/remote/client.rs @@ -15,6 +15,51 @@ use crate::remote::retry::{ResolvedRetryConfig, RetryCounter}; const REQUEST_ID_HEADER: HeaderName = HeaderName::from_static("x-request-id"); +/// Privacy mode for request logging and non-success response handling. +/// +/// [`RequestPrivacy::Standard`] preserves the existing harmless JSON body +/// visibility. [`RequestPrivacy::Sensitive`] never includes request bodies or +/// headers in logs, and never folds response bodies into error chains. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum RequestPrivacy { + Standard, + Sensitive, +} + +/// Format a request for debug logging according to [`RequestPrivacy`]. +fn format_request_log(request: &Request, request_id: &str, privacy: RequestPrivacy) -> String { + match privacy { + RequestPrivacy::Standard => { + let content_type = request + .headers() + .get("content-type") + .and_then(|v| v.to_str().ok()); + if content_type == Some("application/json") { + let body = request + .body() + .and_then(|b| b.as_bytes()) + .map(|b| String::from_utf8_lossy(b).into_owned()) + .unwrap_or_default(); + format!( + "Sending request_id={}: {:?} with body {}", + request_id, request, body + ) + } else { + format!("Sending request_id={}: {:?}", request_id, request) + } + } + RequestPrivacy::Sensitive => { + // Safe context only: request id, method, and URL. Never body or headers. + format!( + "Sending request_id={}: {} {}", + request_id, + request.method(), + request.url() + ) + } + } +} + /// Configuration for TLS/mTLS settings. #[derive(Clone, Debug)] pub struct TlsConfig { @@ -753,9 +798,37 @@ impl RestfulLanceDbClient { pub async fn send_with_retry( &self, req_builder: RequestBuilder, - mut make_body: Option Result + Send + 'static>>, + make_body: Option Result + Send + 'static>>, retry_5xx: bool, ) -> Result<(String, Response)> { + self.send_with_retry_inner(req_builder, make_body, retry_5xx, RequestPrivacy::Standard) + .await + } + + /// Like [`Self::send_with_retry`], but never logs request bodies/headers and + /// never folds non-success response bodies into retry or HTTP error chains. + /// + /// Privacy affects only logging and error-body exposure; retry budgets are + /// identical to [`Self::send_with_retry`] for the same [`RetryConfig`]. + pub(crate) async fn send_sensitive_with_retry( + &self, + req_builder: RequestBuilder, + make_body: Option Result + Send + 'static>>, + retry_5xx: bool, + ) -> Result<(String, Response)> { + self.send_with_retry_inner(req_builder, make_body, retry_5xx, RequestPrivacy::Sensitive) + .await + } + + async fn send_with_retry_inner( + &self, + req_builder: RequestBuilder, + mut make_body: Option Result + Send + 'static>>, + retry_5xx: bool, + privacy: RequestPrivacy, + ) -> Result<(String, Response)> { + // Privacy must never alter retry budgets: both Standard and Sensitive + // share the same ResolvedRetryConfig / RetryCounter semantics. let retry_config = &self.retry_config; let non_5xx_statuses = retry_config .statuses @@ -772,6 +845,7 @@ impl RestfulLanceDbClient { let mut r = r.map_err(|e| Error::Runtime { message: format!("Failed to build request: {}", e), })?; + // One SDK-generated request id is reused across every retry attempt. let request_id = self.extract_request_id(&mut r); let mut retry_counter = RetryCounter::new(retry_config, request_id.clone()); @@ -790,12 +864,14 @@ impl RestfulLanceDbClient { let mut request = request.map_err(|e| Error::Runtime { message: format!("Failed to build request: {}", e), })?; - self.set_request_id(&mut request, &request_id.clone()); + self.set_request_id(&mut request, &request_id); // Apply dynamic headers before each retry attempt request = self.apply_dynamic_headers(request).await?; - self.log_request(&request, &request_id); + if log::log_enabled!(log::Level::Debug) { + debug!("{}", format_request_log(&request, &request_id, privacy)); + } let response = self.sender.send(&c, request).await.map(|r| (r.status(), r)); @@ -811,10 +887,16 @@ impl RestfulLanceDbClient { if (retry_5xx && retry_config.statuses.contains(&status)) || non_5xx_statuses.contains(&status) => { - let source = self - .check_response(&retry_counter.request_id, response) - .await - .unwrap_err(); + let source = match privacy { + RequestPrivacy::Standard => self + .check_response(&retry_counter.request_id, response) + .await + .unwrap_err(), + RequestPrivacy::Sensitive => self + .check_sensitive_response(&retry_counter.request_id, response) + .await + .unwrap_err(), + }; retry_counter.increment_request_failures(source)?; } Err(err) if err.is_connect() => { @@ -839,22 +921,12 @@ impl RestfulLanceDbClient { } } - pub(crate) fn log_request(&self, request: &Request, request_id: &String) { + pub(crate) fn log_request(&self, request: &Request, request_id: &str) { if log::log_enabled!(log::Level::Debug) { - let content_type = request - .headers() - .get("content-type") - .map(|v| v.to_str().unwrap()); - if content_type == Some("application/json") { - let body = request.body().as_ref().unwrap().as_bytes().unwrap(); - let body = String::from_utf8_lossy(body); - debug!( - "Sending request_id={}: {:?} with body {}", - request_id, request, body - ); - } else { - debug!("Sending request_id={}: {:?}", request_id, request); - } + debug!( + "{}", + format_request_log(request, request_id, RequestPrivacy::Standard) + ); } } @@ -898,6 +970,27 @@ impl RestfulLanceDbClient { }) } } + + /// Like [`Self::check_response`], but discards the response body on failure + /// so marker-bearing payloads never enter [`Error::Http`] chains. + pub(crate) async fn check_sensitive_response( + &self, + request_id: &str, + response: Response, + ) -> Result { + let status = response.status(); + if status.is_success() { + Ok(response) + } else { + // Discard the body entirely; never fold it into Error::Http. + let _ = response.bytes().await; + Err(Error::Http { + source: status.to_string().into(), + request_id: request_id.into(), + status_code: Some(status), + }) + } + } } pub trait RequestResultExt { @@ -1066,6 +1159,7 @@ pub mod test_utils { mod tests { use super::*; use serial_test::serial; + use std::sync::atomic::{AtomicUsize, Ordering}; use std::time::Duration; // Serializes the env-var-mutating tests below: cargo test runs tests in @@ -1664,4 +1758,253 @@ mod tests { } assert!(matches!(err, Error::InvalidInput { .. }), "got: {err:?}"); } + + // ------------------------------------------------------------------------- + // Sensitive-request privacy mode (generic transport; RED until helpers exist) + // ------------------------------------------------------------------------- + + const PRIVACY_SOURCE_MARKER: &str = "SENSITIVE_PRIVACY_SOURCE_BODY_MARKER_client"; + const PRIVACY_SECRET_MARKER: &str = "secret://team/client-privacy-token"; + + fn privacy_json_request(url: &str, body: &str, request_id: &str) -> Request { + reqwest::Client::new() + .post(url) + .header("content-type", "application/json") + .header("x-request-id", request_id) + .body(body.to_string()) + .build() + .expect("build privacy fixture request") + } + + fn assert_markers_absent(text: &str) { + assert!( + !text.contains(PRIVACY_SOURCE_MARKER), + "source marker must be absent: {text}" + ); + assert!( + !text.contains(PRIVACY_SECRET_MARKER), + "secret marker must be absent: {text}" + ); + } + + fn 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 + } + + /// Standard JSON request logging keeps the current harmless body visibility. + #[test] + fn format_request_log_standard_retains_harmless_json_body() { + let request_id = "req-privacy-standard"; + let body = r#"{"ok":true,"note":"harmless-visible-body"}"#; + let request = privacy_json_request("http://localhost/v1/table/", body, request_id); + + let log = format_request_log(&request, request_id, RequestPrivacy::Standard); + + assert!( + log.contains(request_id), + "standard log must retain request id: {log}" + ); + assert!( + log.contains("POST"), + "standard log must retain method: {log}" + ); + assert!( + log.contains("/v1/table/"), + "standard log must retain URL path: {log}" + ); + assert!( + log.contains("harmless-visible-body"), + "standard JSON logging must retain body visibility: {log}" + ); + assert!( + log.contains(body) || log.contains(r#""note":"harmless-visible-body""#), + "standard JSON logging must include the harmless JSON body: {log}" + ); + } + + /// Sensitive JSON formatting redacts the entire body and keeps only safe context. + #[test] + fn format_request_log_sensitive_redacts_json_body_keeps_safe_context() { + let request_id = "req-privacy-sensitive"; + let body = + format!(r#"{{"source":"{PRIVACY_SOURCE_MARKER}","secret":"{PRIVACY_SECRET_MARKER}"}}"#); + let request = + privacy_json_request("http://localhost/v1/functions/register", &body, request_id); + + let log = format_request_log(&request, request_id, RequestPrivacy::Sensitive); + + assert!( + log.contains(request_id), + "sensitive log must retain request id: {log}" + ); + assert!( + log.contains("POST"), + "sensitive log must retain method: {log}" + ); + assert!( + log.contains("/v1/functions/register"), + "sensitive log must retain URL path: {log}" + ); + assert_markers_absent(&log); + assert!( + !log.contains(&body), + "sensitive JSON formatting must redact the entire body: {log}" + ); + } + + /// Sensitive non-success responses omit the response body from Error::Http text. + #[tokio::test] + async fn check_sensitive_response_omits_non_success_response_body() { + let client = test_utils::client_with_handler(|_| { + http::Response::builder().status(200).body("").unwrap() + }); + let response: Response = http::Response::builder() + .status(400) + .body(format!( + "client error echoed {PRIVACY_SOURCE_MARKER} and {PRIVACY_SECRET_MARKER}" + )) + .unwrap() + .into(); + + let err = client + .check_sensitive_response("req-privacy-check", response) + .await + .expect_err("non-success sensitive response must fail closed"); + + assert!( + matches!(err, Error::Http { .. }), + "expected Error::Http, got {err:?}" + ); + assert_markers_absent(&error_chain_text(&err)); + } + + /// Sensitive send+retry must not leak request/response markers into retry errors. + #[tokio::test] + async fn send_sensitive_with_retry_omits_markers_from_exhausted_retry_errors() { + let call_count = Arc::new(AtomicUsize::new(0)); + let counted = call_count.clone(); + let client = test_utils::client_with_handler_and_config( + move |request| { + counted.fetch_add(1, Ordering::SeqCst); + let body = request.body().and_then(|b| b.as_bytes()).unwrap_or(b""); + let body = std::str::from_utf8(body).unwrap_or(""); + assert!( + body.contains(PRIVACY_SOURCE_MARKER) && body.contains(PRIVACY_SECRET_MARKER), + "trusted wire body must still carry sensitive fields" + ); + http::Response::builder() + .status(500) + .body(format!( + "server echoed {PRIVACY_SOURCE_MARKER} and {PRIVACY_SECRET_MARKER}" + )) + .unwrap() + }, + ClientConfig { + retry_config: RetryConfig { + // RetryCounter treats `retries` as max request failures, so + // retries=2 yields exactly two transport attempts before Error::Retry. + retries: Some(2), + backoff_factor: Some(0.0), + backoff_jitter: Some(0.0), + ..Default::default() + }, + ..Default::default() + }, + ); + + let payload = serde_json::json!({ + "source": PRIVACY_SOURCE_MARKER, + "secret": PRIVACY_SECRET_MARKER, + }); + let req = client.post("/v1/functions/register").json(&payload); + let err = client + .send_sensitive_with_retry(req, None, true) + .await + .expect_err("exhausted sensitive 5xx retries must fail"); + + assert!( + matches!(err, Error::Retry { .. }), + "expected Error::Retry, got {err:?}" + ); + assert_markers_absent(&error_chain_text(&err)); + assert_eq!( + call_count.load(Ordering::SeqCst), + 2, + "RetryCounter max request failures=2 must make exactly two transport attempts" + ); + } + + /// Standard and Sensitive share the same RetryCounter attempt budget. + #[tokio::test] + async fn send_with_retry_standard_and_sensitive_share_attempt_budget() { + async fn exhausted_attempts(sensitive: bool) -> usize { + let call_count = Arc::new(AtomicUsize::new(0)); + let counted = call_count.clone(); + let client = test_utils::client_with_handler_and_config( + move |_| { + counted.fetch_add(1, Ordering::SeqCst); + http::Response::builder() + .status(500) + .body(format!( + "server echoed {PRIVACY_SOURCE_MARKER} and {PRIVACY_SECRET_MARKER}" + )) + .unwrap() + }, + ClientConfig { + retry_config: RetryConfig { + retries: Some(2), + backoff_factor: Some(0.0), + backoff_jitter: Some(0.0), + ..Default::default() + }, + ..Default::default() + }, + ); + + let payload = serde_json::json!({ + "source": PRIVACY_SOURCE_MARKER, + "secret": PRIVACY_SECRET_MARKER, + }); + let req = client.post("/v1/functions/register").json(&payload); + let err = if sensitive { + client + .send_sensitive_with_retry(req, None, true) + .await + .expect_err("exhausted sensitive 5xx retries must fail") + } else { + client + .send_with_retry(req, None, true) + .await + .expect_err("exhausted standard 5xx retries must fail") + }; + assert!( + matches!(err, Error::Retry { .. }), + "expected Error::Retry, got {err:?}" + ); + if sensitive { + assert_markers_absent(&error_chain_text(&err)); + } + call_count.load(Ordering::SeqCst) + } + + let standard_attempts = exhausted_attempts(false).await; + let sensitive_attempts = exhausted_attempts(true).await; + assert_eq!( + standard_attempts, sensitive_attempts, + "Standard and Sensitive must share the same attempt budget for identical RetryConfig" + ); + assert_eq!( + standard_attempts, 2, + "RetryCounter max request failures=2 must make exactly two transport attempts" + ); + } } diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index f3042ec0e..c87653e26 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -23,6 +23,7 @@ use crate::database::{ JobDescription, JobInfo, OpenTableRequest, ReadConsistency, TableNamesRequest, }; use crate::error::Result; +use crate::function::RegisterFunctionJobSpec; use crate::remote::util::stream_as_body; use crate::table::BaseTable; @@ -576,6 +577,46 @@ impl Database for RemoteDatabase { .map_err(Into::into) } + async fn register_function(&self, spec: RegisterFunctionJobSpec) -> Result { + let req = self.client.post("/v1/functions/register").json(&spec); + let (request_id, rsp) = self + .client + .send_sensitive_with_retry(req, None, true) + .await?; + let rsp = self + .client + .check_sensitive_response(&request_id, rsp) + .await?; + + // Payload-free protocol failure: never fold response bytes into Error::Http. + let bytes = rsp.bytes().await.err_to_http(request_id.clone())?; + let value: serde_json::Value = match serde_json::from_slice(&bytes) { + Ok(value) => value, + Err(_) => { + return Err(Error::Http { + source: "register function response is not valid JSON".into(), + request_id, + status_code: None, + }); + } + }; + let job_id = match value.get("job_id") { + Some(serde_json::Value::String(job_id)) if !job_id.is_empty() => job_id.clone(), + _ => { + return Err(Error::Http { + source: "register function response missing or invalid job_id".into(), + request_id, + status_code: None, + }); + } + }; + + Ok(crate::job::Job::new(Box::new(super::job::RemoteJob::new( + self.client.clone(), + job_id, + )))) + } + async fn table_names(&self, request: TableNamesRequest) -> Result> { let mut req = if !request.namespace_path.is_empty() { let namespace_id = @@ -1077,9 +1118,15 @@ mod tests { Connection, Error, database::CreateTableMode, error::FunctionErrorCode, - function::{Function, FunctionId, FunctionOutput, FunctionParameter, FunctionSignature}, + function::{ + Function, FunctionCapability, FunctionDefinition, FunctionId, FunctionOutput, + FunctionParameter, FunctionSignature, PythonFunctionDefinition, + RegisterFunctionJobSpec, + }, job::JobResult, - remote::{ARROW_STREAM_CONTENT_TYPE, ClientConfig, HeaderProvider, JSON_CONTENT_TYPE}, + remote::{ + ARROW_STREAM_CONTENT_TYPE, ClientConfig, HeaderProvider, JSON_CONTENT_TYPE, RetryConfig, + }, }; use serde_json::{Value, json}; @@ -2962,4 +3009,401 @@ mod tests { function_wire ); } + + // ------------------------------------------------------------------------- + // RegisterFunctionJobSpec remote submit transport (RED until register_function) + // ------------------------------------------------------------------------- + + const REGISTER_SOURCE_MARKER: &str = "def normalize(text, limit):\n return text[:limit] # SENSITIVE_REGISTER_SOURCE_MARKER\n"; + const REGISTER_SECRET_MARKER: &str = "secret://team/register-function-privacy-token"; + + fn sample_register_function_job_spec() -> RegisterFunctionJobSpec { + let signature = FunctionSignature::try_new( + vec![ + FunctionParameter::new("text", DataType::Utf8), + FunctionParameter::new("limit", DataType::Int32), + ], + FunctionOutput::new(DataType::Utf8, true), + ) + .expect("valid FunctionSignature"); + let python = PythonFunctionDefinition::try_new( + "normalize_mod", + "normalize", + REGISTER_SOURCE_MARKER, + "3.12", + vec!["Unidecode==1.3.8".to_string()], + ) + .expect("valid PythonFunctionDefinition"); + let capabilities = vec![ + FunctionCapability::try_network("https://api.example.com").expect("network capability"), + FunctionCapability::try_secret(REGISTER_SECRET_MARKER, "API_TOKEN") + .expect("secret capability"), + ]; + let definition = FunctionDefinition::try_new(signature, python, capabilities) + .expect("valid FunctionDefinition"); + RegisterFunctionJobSpec::try_new("text.normalize", definition, None) + .expect("valid RegisterFunctionJobSpec") + } + + fn register_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_register_markers_absent(err: &Error) { + let text = register_error_chain_text(err); + assert!( + !text.contains(REGISTER_SOURCE_MARKER), + "Python source marker must be absent from error/debug/source chain: {text}" + ); + assert!( + !text.contains(REGISTER_SECRET_MARKER), + "secret reference marker must be absent from error/debug/source chain: {text}" + ); + } + + fn assert_register_request( + request: &reqwest::Request, + expected_spec: &RegisterFunctionJobSpec, + ) { + assert_eq!(request.method(), &reqwest::Method::POST); + assert_eq!(request.url().path(), "/v1/functions/register"); + let body = request + .body() + .and_then(|b| b.as_bytes()) + .expect("register request must carry a JSON body"); + let actual: Value = serde_json::from_slice(body).expect("register body must be JSON"); + let expected = + serde_json::to_value(expected_spec).expect("serialize RegisterFunctionJobSpec"); + assert_eq!( + actual, expected, + "POST body must be the exact RegisterFunctionJobSpec wire" + ); + assert!( + actual + .to_string() + .contains("SENSITIVE_REGISTER_SOURCE_MARKER"), + "trusted request body must include full Python source" + ); + assert_eq!( + actual["definition"]["capabilities"][1]["reference"], + Value::String(REGISTER_SECRET_MARKER.into()), + "trusted request body must include secret reference" + ); + } + + /// Successful submit uses exact path/method/body and projects a non-empty Job id. + #[tokio::test] + async fn register_function_submit_posts_exact_spec_and_returns_remote_job() { + let spec = sample_register_function_job_spec(); + let expected_body = serde_json::to_value(&spec).expect("serialize spec"); + let conn = Connection::new_with_handler(move |request| { + assert_eq!(request.method(), &reqwest::Method::POST); + assert_eq!(request.url().path(), "/v1/functions/register"); + let body = request.body().unwrap().as_bytes().unwrap(); + let actual: Value = serde_json::from_slice(body).unwrap(); + assert_eq!(actual, expected_body); + http::Response::builder() + .status(200) + .body(r#"{"job_id":"job-register-transport-1","server_extra":{"ok":true}}"#) + .unwrap() + }); + + let job = conn + .register_function(spec) + .await + .expect("register_function submit must succeed"); + assert_eq!( + job.id(), + Some("job-register-transport-1"), + "successful submit must project the non-empty job_id onto the unified remote Job" + ); + } + + /// One retry keeps the SDK-generated request id and exact body before success. + #[tokio::test] + async fn register_function_submit_retry_preserves_request_id_and_body() { + let spec = sample_register_function_job_spec(); + let expected_body = serde_json::to_value(&spec).expect("serialize spec"); + 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 expected_spec = sample_register_function_job_spec(); + let conn = Connection::new_with_handler_and_config( + move |request| { + assert_register_request(&request, &expected_spec); + assert_eq!( + serde_json::from_slice::(request.body().unwrap().as_bytes().unwrap()) + .unwrap(), + expected_body + ); + + 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("transient register failure") + .unwrap() + } else { + http::Response::builder() + .status(200) + .body(r#"{"job_id":"job-register-retry-1"}"#) + .unwrap() + } + }, + ClientConfig { + retry_config: RetryConfig { + retries: Some(2), + backoff_factor: Some(0.0), + backoff_jitter: Some(0.0), + ..Default::default() + }, + ..Default::default() + }, + ); + + let job = conn + .register_function(spec) + .await + .expect("register_function must succeed after one retry"); + assert_eq!(job.id(), Some("job-register-retry-1")); + assert_eq!(attempts.load(Ordering::SeqCst), 2); + assert!(seen_request_id.get().is_some()); + } + + /// Missing/null/empty/wrong-type/malformed job_id fail closed as Error::Http. + #[tokio::test] + async fn register_function_submit_invalid_job_id_is_http_without_markers() { + let cases: Vec<(&str, String)> = vec![ + ("missing", r#"{"server_extra":true}"#.to_string()), + ("null", r#"{"job_id":null}"#.to_string()), + ("empty", r#"{"job_id":""}"#.to_string()), + ("wrong_type", r#"{"job_id":123}"#.to_string()), + ("malformed", "not-json".to_string()), + ]; + + let mut unexpected = Vec::new(); + for (label, response_body) in cases { + let spec = sample_register_function_job_spec(); + let expected_body = serde_json::to_value(&spec).expect("serialize spec"); + let body_for_handler = response_body.clone(); + let conn = Connection::new_with_handler(move |request| { + assert_eq!(request.url().path(), "/v1/functions/register"); + let actual: Value = + serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap(); + assert_eq!(actual, expected_body); + http::Response::builder() + .status(200) + .body(body_for_handler.clone()) + .unwrap() + }); + + match conn.register_function(spec).await { + Err(err @ Error::Http { .. }) => assert_register_markers_absent(&err), + other => unexpected.push(format!("{label}: {other:?}")), + } + } + assert!( + unexpected.is_empty(), + "invalid job_id shapes must fail closed as Error::Http: {unexpected:?}" + ); + } + + /// Non-retry 4xx and exhausted 5xx bodies that echo markers stay out of error text. + #[tokio::test] + async fn register_function_submit_error_bodies_omit_sensitive_markers() { + let echoed = + format!("register failed with {REGISTER_SOURCE_MARKER} and {REGISTER_SECRET_MARKER}"); + + // Non-retryable 4xx + { + let spec = sample_register_function_job_spec(); + let body = echoed.clone(); + let conn = Connection::new_with_handler(move |request| { + assert_eq!(request.url().path(), "/v1/functions/register"); + http::Response::builder() + .status(400) + .body(body.clone()) + .unwrap() + }); + let err = conn + .register_function(spec) + .await + .expect_err("4xx register submit must fail"); + assert!( + matches!(err, Error::Http { .. }), + "non-retry 4xx must surface as Error::Http, got {err:?}" + ); + assert_register_markers_absent(&err); + } + + // Exhausted retryable 5xx + { + let spec = sample_register_function_job_spec(); + let body = echoed.clone(); + let attempts = Arc::new(AtomicUsize::new(0)); + let attempts_ref = attempts.clone(); + let conn = Connection::new_with_handler_and_config( + move |request| { + attempts_ref.fetch_add(1, Ordering::SeqCst); + assert_eq!(request.url().path(), "/v1/functions/register"); + http::Response::builder() + .status(500) + .body(body.clone()) + .unwrap() + }, + ClientConfig { + retry_config: RetryConfig { + // RetryCounter treats `retries` as max request failures, so + // retries=2 yields exactly two transport attempts before Error::Retry. + retries: Some(2), + backoff_factor: Some(0.0), + backoff_jitter: Some(0.0), + ..Default::default() + }, + ..Default::default() + }, + ); + let err = conn + .register_function(spec) + .await + .expect_err("exhausted 5xx register submit must fail"); + assert!( + matches!(err, Error::Retry { .. }), + "exhausted 5xx must surface as Error::Retry, got {err:?}" + ); + assert_eq!( + attempts.load(Ordering::SeqCst), + 2, + "RetryCounter max request failures=2 must make exactly two transport attempts" + ); + assert_register_markers_absent(&err); + } + } + + /// Local databases reject registration without mutating database state. + #[tokio::test] + async fn register_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 = conn + .register_function(sample_register_function_job_spec()) + .await + .expect_err("local register_function must be unsupported"); + assert!( + matches!(err, Error::NotSupported { .. }), + "expected NotSupported, got {err:?}" + ); + + let after = conn + .table_names() + .execute() + .await + .expect("table_names after"); + assert_eq!(before, after, "unsupported register must not mutate tables"); + } + + /// Submit then existing /v1/jobs/describe returns the exact Function (no name lookup). + #[tokio::test] + async fn register_function_submit_then_describe_returns_exact_function() { + let expected = sample_description_function(); + let function_wire = job_result_function_wire(&expected); + let describe_body = describe_body( + "job-register-wait-1", + "DONE", + JsonField::Present(Value::String("register_function".into())), + JsonField::Present(function_wire), + ); + let spec = sample_register_function_job_spec(); + let expected_spec_body = serde_json::to_value(&spec).expect("serialize spec"); + let paths = Arc::new(std::sync::Mutex::new(Vec::::new())); + let paths_ref = paths.clone(); + + let conn = Connection::new_with_handler(move |request| { + let path = request.url().path().to_string(); + paths_ref.lock().unwrap().push(path.clone()); + match path.as_str() { + "/v1/functions/register" => { + assert_eq!(request.method(), &reqwest::Method::POST); + let actual: Value = + serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()) + .unwrap(); + assert_eq!(actual, expected_spec_body); + http::Response::builder() + .status(200) + .body(r#"{"job_id":"job-register-wait-1"}"#.to_string()) + .unwrap() + } + "/v1/jobs/describe" => { + assert_eq!(request.method(), &reqwest::Method::POST); + let body: Value = + serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()) + .unwrap(); + assert_eq!(body["job_id"], "job-register-wait-1"); + assert!( + body.get("name").is_none(), + "describe must not perform a second name lookup: {body}" + ); + http::Response::builder() + .status(200) + .body(describe_body.clone()) + .unwrap() + } + other => panic!("unexpected path for register+wait flow: {other}"), + } + }); + + let job = conn + .register_function(spec) + .await + .expect("register_function submit must return a Job"); + assert_eq!(job.id(), Some("job-register-wait-1")); + + let waited = job.wait().await.expect("wait via /v1/jobs/describe"); + let function = waited + .function() + .expect("register_function success must be JobResult::Function"); + assert_exact_function(function, &expected); + + let seen = paths.lock().unwrap().clone(); + assert_eq!( + seen, + vec![ + "/v1/functions/register".to_string(), + "/v1/jobs/describe".to_string(), + ], + "flow must be submit then describe only, with no Function name lookup" + ); + } }