From 45e00d4de75500f42687e176e442981677294106 Mon Sep 17 00:00:00 2001 From: Jack Ye Date: Fri, 30 Oct 2026 09:15:53 -0700 Subject: [PATCH] feat: use parameterized SQL for remote secrets Stacked on #4674. Remote Secret creation, rotation, description, listing, and drop now use the SQL backend. Writes bind the credential to `CREATE SECRET name AS $1` or `ALTER SECRET name AS $1`, keeping values out of SQL text. Description returns metadata; listing uses namespace-scoped continuation tokens. The existing synchronous APIs submit a SQL query and wait internally through the same query-job adapter as function operations. The single query identity is retained internally; no registry job or extra query ID is introduced. Pollable shared results contain status or metadata and never credential values. --------- Co-authored-by: Xuanwo --- docs/web/basics/jobs.mdx | 7 + .../tests/test_first_class_function_slice2.py | 73 +--- python/python/tests/test_listing.py | 2 +- rust/lancedb/src/remote/db.rs | 326 ++---------------- rust/lancedb/src/remote/sql.rs | 1 + rust/lancedb/src/remote/sql/secret.rs | 156 +++++++++ rust/lancedb/src/remote/sql_test.rs | 272 ++++++++++++++- rust/lancedb/src/secrets.rs | 2 +- 8 files changed, 466 insertions(+), 373 deletions(-) create mode 100644 rust/lancedb/src/remote/sql/secret.rs diff --git a/docs/web/basics/jobs.mdx b/docs/web/basics/jobs.mdx index 9f4440de0..baa1c2ae8 100644 --- a/docs/web/basics/jobs.mdx +++ b/docs/web/basics/jobs.mdx @@ -125,6 +125,13 @@ the immutable FunctionVersion. `drop_function_async()` returns the catalog-remov boolean together with a query job; waiting on that handle confirms removal while physical cleanup continues separately. +Remote Secret create, alter, list, describe, and drop use the SQL query backend. +The synchronous APIs wait on query completion internally. Create and alter bind +the credential as a parameter to `CREATE SECRET name AS $1` or +`ALTER SECRET name AS $1`, keeping it out of SQL text. Results contain only +completion status or Secret names and timestamps; they never return credential +values. SQL `DESCRIBE SECRET` additionally exposes rotation history. + The existing `job.job_type` identifies query-backed handles as `"query"` from submission, and `job.id` contains the connection-scoped SQL query UUID string. Use `describe_query(job.id)` on the submitting connection to inspect the query. diff --git a/python/python/tests/test_first_class_function_slice2.py b/python/python/tests/test_first_class_function_slice2.py index 6240c362a..7e60df039 100644 --- a/python/python/tests/test_first_class_function_slice2.py +++ b/python/python/tests/test_first_class_function_slice2.py @@ -246,50 +246,44 @@ def test_bindings_may_not_collide_with_plain_configuration(): _SECRET_DEBUG_LOG_SOURCE = """ -import http.server -import json -import threading +import pyarrow as pa +import pyarrow.flight as flight import lancedb -class Handler(http.server.BaseHTTPRequestHandler): - def log_message(self, *args): - pass +class SecretServer(flight.FlightServerBase): + def __init__(self): + super().__init__("grpc://127.0.0.1:0") - def do_POST(self): - self.rfile.read(int(self.headers.get("Content-Length", "0"))) - payload = json.dumps({}).encode() - self.send_response(200) - self.send_header("Content-Type", "application/json") - self.send_header("Content-Length", str(len(payload))) - self.end_headers() - self.wfile.write(payload) + def do_exchange(self, context, descriptor, reader, writer): + assert b"SECRET_VALUE_SENTINEL" not in descriptor.command + parameters = reader.read_all() + assert parameters.column(0).to_pylist() == ["SECRET_VALUE_SENTINEL"] + result = pa.table({"status": ["created"]}) + writer.begin(result.schema) + writer.write_table(result) + writer.close() -server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), Handler) -threading.Thread( - target=server.serve_forever, kwargs={"poll_interval": 0.01}, daemon=True -).start() -try: +with SecretServer() as server: db = lancedb.connect( "db://dev", api_key="API_KEY_SENTINEL", - host_override="http://127.0.0.1:%d" % server.server_address[1], - client_config={"retry_config": {"retries": 0}}, + host_override="http://localhost:10024", + sql_host_override=f"grpc://127.0.0.1:{server.port}", ) db.create_secret("openai-prod", "SECRET_VALUE_SENTINEL") -finally: - server.shutdown() """ def test_a_credential_never_reaches_a_debug_log(tmp_path): - """The logger sees the serialized body, so no value-side redaction reaches it. + """Credential parameters and authentication stay out of Flight debug logs. Runs in a subprocess because the Rust logger reads ``LANCEDB_LOG`` once, at import. """ + pytest.importorskip("pyarrow.flight") script = tmp_path / "write_secret.py" script.write_text(_SECRET_DEBUG_LOG_SOURCE) @@ -303,7 +297,7 @@ def test_a_credential_never_reaches_a_debug_log(tmp_path): output = result.stdout + result.stderr # Without this the test passes when debug logging is simply off. - assert "Sending request_id=" in output, output + assert "Created client for host:" in output, output assert "SECRET_VALUE_SENTINEL" not in output assert "API_KEY_SENTINEL" not in output @@ -1816,35 +1810,6 @@ def _mock_remote_function_catalog(): thread.join() -def test_remote_secret_verbs_round_trip(): - 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.create_secret("openai-prod", "sk-live-0001") is None - assert db.alter_secret("openai-prod", "sk-live-0002") is None - assert list(db.list_secrets()) == ["openai-prod", "hf-prod"] - assert db.drop_secret("openai-prod") is None - - routes = [path for path, _ in state["requests"]] - assert routes == [ - "/v1/secret/openai-prod/create", - "/v1/secret/openai-prod/alter", - "/v1/namespace/$/secret/list", - "/v1/namespace/$/secret/list", - "/v1/secret/openai-prod/drop", - ] - # The Secret is the path identifier, so the body is the value alone. - assert state["requests"][0][1] == {"value": "sk-live-0001"} - # Listing is a GET: the first page asks for nothing, the second resumes on - # the token the server handed back, and neither carries a body. - assert state["requests"][2][1] == {} - assert state["requests"][3][1] == {"page_token": "next"} - - def test_building_a_binding_contacts_no_server(): """A binding is a local value: it says nothing about whether the Secret exists. diff --git a/python/python/tests/test_listing.py b/python/python/tests/test_listing.py index 7b916dadc..380d6fae8 100644 --- a/python/python/tests/test_listing.py +++ b/python/python/tests/test_listing.py @@ -84,7 +84,7 @@ async def advance(items, asynchronous): return await items.__anext__() if asynchronous else next(items) -KINDS = ["secrets", "views", "jobs"] +KINDS = ["views", "jobs"] @pytest.mark.asyncio diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index b7a641dd3..22db05b06 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -34,8 +34,8 @@ use crate::remote::util::stream_as_body; use crate::secrets::SecretInfo; use crate::table::BaseTable; use crate::utils::{ - validate_database_name, validate_function_name, validate_namespace, validate_secret_name, - validate_table_name, validate_view_name, + validate_database_name, validate_function_name, validate_namespace, validate_table_name, + validate_view_name, }; use crate::view::ViewDescription; @@ -477,24 +477,6 @@ impl RemoteDatabase { }) } - /// Post a request whose body carries a credential. - /// - /// Shared by the create and alter verbs, which declare their own request - /// types: the two mean different things to the service and are free to - /// diverge, so what they share is the posting and not the shape. - /// - /// The value is a request field and never a path segment or query - /// parameter, which keeps it out of access logs and proxy traces. - async fn post_secret_write(&self, route: &str, body: &T) -> Result<()> { - let req = self.client.post(route).json(body); - // This call is what says the body is a credential. Nothing downstream - // can tell from the bytes, and a route list in the transport would have - // to be kept in step with endpoints declared here. - let (request_id, response) = self.client.send_suppressing_body(req).await?; - self.client.check_response(&request_id, response).await?; - Ok(()) - } - async fn submit_drop_table( &self, name: &str, @@ -670,14 +652,6 @@ fn build_namespace_identifier(namespace: &[String]) -> Result { } } -/// Build a full path to a secret. Note that we do not need to percent-encode -/// this when using it in a URL because the validation functions ensure that -/// only URL-safe characters are present. -fn build_secret_identifier(name: &str, namespace: &[String]) -> Result { - validate_secret_name(name)?; - build_compound_identifier(name, namespace) -} - /// Build a full path to a view. Note that we do not need to percent-encode /// this when using it in a URL because the validation functions ensure that /// only URL-safe characters are present. @@ -730,40 +704,6 @@ struct RemoteListJobsResponse { page_token: Option, } -/// Create a Secret under a name the database does not yet hold. -/// -/// Declared separately from the alter request although the two are identical -/// today: they are different operations to the service -- one refuses an -/// existing name, the other requires it -- and either may grow a field the -/// other has no meaning for. -/// -/// The name and its namespace are the path identifier, so neither appears here. -#[derive(serde::Serialize)] -struct RemoteCreateSecretRequest<'a> { - value: &'a str, -} - -/// Replace the credential behind a Secret the database already holds. -#[derive(serde::Serialize)] -struct RemoteAlterSecretRequest<'a> { - value: &'a str, -} - -#[derive(serde::Deserialize)] -struct RemoteListSecretsResponse { - #[serde(default)] - secrets: Vec, - #[serde(default)] - page_token: Option, -} - -/// An object rather than a bare name so a later listing can carry a Secret's -/// type or last-updated time without breaking this one. -#[derive(serde::Deserialize)] -struct RemoteListedSecret { - name: String, -} - /// Define a view from a query. The name and its namespace are the path /// identifier, so neither appears here. #[derive(serde::Serialize)] @@ -1070,67 +1010,40 @@ impl Database for RemoteDatabase { value: &str, namespace_path: &[String], ) -> Result<()> { - let secret_id = build_secret_identifier(name, namespace_path)?; - self.post_secret_write( - &format!("/v1/secret/{secret_id}/create"), - &RemoteCreateSecretRequest { value }, - ) - .await + crate::utils::validate_secret_name(name)?; + validate_namespace(namespace_path)?; + self.sql_client()? + .write_secret(name, value, namespace_path, true) + .await } - async fn alter_secret(&self, name: &str, value: &str, namespace_path: &[String]) -> Result<()> { - let secret_id = build_secret_identifier(name, namespace_path)?; - self.post_secret_write( - &format!("/v1/secret/{secret_id}/alter"), - &RemoteAlterSecretRequest { value }, - ) - .await + crate::utils::validate_secret_name(name)?; + validate_namespace(namespace_path)?; + self.sql_client()? + .write_secret(name, value, namespace_path, false) + .await } - async fn list_secrets( &self, namespace_path: &[String], options: ListingOptions, ) -> Result> { - let namespace_id = build_namespace_identifier(namespace_path)?; - let path = format!("/v1/namespace/{namespace_id}/secret/list"); - let mut req = self.client.get(&path); - if let Some(token) = options.page_token { - req = req.query(&[("page_token", token)]); - } - if let Some(limit) = options.page_limit { - req = req.query(&[("limit", limit)]); - } - let (request_id, response) = self.client.send(req).await?; - let response = self.client.check_response(&request_id, response).await?; - let response: RemoteListSecretsResponse = response.json().await.err_to_http(request_id)?; - let items = response - .secrets - .into_iter() - .map(|secret| secret.name) - .collect(); - Ok(ListingPage { - items, - page_token: response.page_token, - }) + validate_namespace(namespace_path)?; + self.sql_client()? + .list_secrets(namespace_path, options) + .await } - async fn drop_secret(&self, name: &str, namespace_path: &[String]) -> Result<()> { - let secret_id = build_secret_identifier(name, namespace_path)?; - let req = self.client.post(&format!("/v1/secret/{secret_id}/drop")); - let (request_id, response) = self.client.send(req).await?; - self.client.check_response(&request_id, response).await?; - Ok(()) + crate::utils::validate_secret_name(name)?; + validate_namespace(namespace_path)?; + self.sql_client()?.drop_secret(name, namespace_path).await } - async fn describe_secret(&self, name: &str, namespace_path: &[String]) -> Result { - let secret_id = build_secret_identifier(name, namespace_path)?; - let req = self - .client - .post(&format!("/v1/secret/{secret_id}/describe")); - let (request_id, response) = self.client.send(req).await?; - let response = self.client.check_response(&request_id, response).await?; - response.json().await.err_to_http(request_id) + crate::utils::validate_secret_name(name)?; + validate_namespace(namespace_path)?; + self.sql_client()? + .describe_secret(name, namespace_path) + .await } async fn create_view( @@ -3704,7 +3617,7 @@ mod tests { assert!(items.try_next().await.unwrap().is_none()); assert!(items.try_next().await.unwrap().is_none()); } - for resource in ["secret", "view", "jobs"] { + for resource in ["view", "jobs"] { let count = Arc::new(AtomicUsize::new(0)); let conn = Connection::new_with_handler({ let count = count.clone(); @@ -4130,29 +4043,6 @@ mod tests { } } - /// An identifier the service accepts is untouched by the encoding, so the - /// route reads like the table and Function routes beside it. - #[tokio::test] - async fn test_an_admissible_identifier_is_not_encoded() { - use std::sync::{Arc, Mutex}; - let seen = Arc::new(Mutex::new(String::new())); - let captured = seen.clone(); - let conn = Connection::new_with_handler(move |request| { - *captured.lock().unwrap() = request.url().path().to_string(); - http::Response::builder().status(200).body("{}").unwrap() - }); - conn.drop_secret( - "openai-prod.v1", - &["prod".to_string(), "vision_2".to_string()], - ) - .await - .unwrap(); - assert_eq!( - *seen.lock().unwrap(), - "/v1/secret/prod$vision_2$openai-prod.v1/drop" - ); - } - /// A table name is joined into the route the same way a Secret's is, so the /// same two failures are reachable: a dot-only name leaves the table route /// entirely, and a name holding the delimiter is indistinguishable from a @@ -4211,172 +4101,6 @@ mod tests { )); } - /// Only `.` and `..` are relative segments. `...` and longer runs are - /// ordinary and address perfectly well, so refusing them would make an - /// object that works today stop working on upgrade. This pins that. - #[tokio::test] - async fn test_a_longer_run_of_periods_is_an_ordinary_name() { - use std::sync::{Arc, Mutex}; - for name in ["...", "....", "a.", ".a", "a..b"] { - let seen = Arc::new(Mutex::new(String::new())); - let captured = seen.clone(); - let conn = Connection::new_with_handler(move |request| { - *captured.lock().unwrap() = request.url().path().to_string(); - http::Response::builder().status(200).body("{}").unwrap() - }); - conn.drop_secret(name, &[]) - .await - .unwrap_or_else(|error| panic!("{name:?} must remain addressable: {error}")); - assert_eq!( - *seen.lock().unwrap(), - format!("/v1/secret/{name}/drop"), - "{name:?} did not reach its own route" - ); - } - } - - #[tokio::test] - async fn test_create_and_alter_secret_send_the_value_in_the_request_body() { - for (route, call) in [ - ("/v1/secret/openai-prod/create", true), - ("/v1/secret/openai-prod/alter", false), - ] { - let conn = Connection::new_with_handler(move |request| { - assert_eq!(request.method(), &reqwest::Method::POST); - assert_eq!(request.url().path(), route); - // Never a path segment or query parameter, which is what keeps - // it out of access logs and proxy traces. - assert!(request.url().query().is_none(), "{:?}", request.url()); - let body: serde_json::Value = - serde_json::from_slice(request.body().unwrap().as_bytes().unwrap()).unwrap(); - // The name is the path identifier, so the body is the value alone. - assert_eq!(body, serde_json::json!({ "value": "sk-live-0001" })); - http::Response::builder().status(200).body("{}").unwrap() - }); - if call { - conn.create_secret("openai-prod", "sk-live-0001", &[]) - .await - .unwrap(); - } else { - conn.alter_secret("openai-prod", "sk-live-0001", &[]) - .await - .unwrap(); - } - } - } - - #[tokio::test] - async fn test_list_secrets_walks_pages_and_returns_names_only() { - let conn = Connection::new_with_handler(|request| { - assert_eq!(request.method(), &reqwest::Method::GET); - assert_eq!(request.url().path(), "/v1/namespace/$/secret/list"); - let page = request - .url() - .query_pairs() - .find(|(key, _)| key == "page_token") - .map(|(_, value)| value.into_owned()); - let body = match page.as_deref() { - None => r#"{"secrets":[{"name":"openai-prod"}],"page_token":"p2"}"#, - Some("p2") => r#"{"secrets":[{"name":"hf-prod"}]}"#, - Some(other) => panic!("unexpected page token: {other}"), - }; - http::Response::builder().status(200).body(body).unwrap() - }); - assert_eq!( - conn.list_secrets(&[], Default::default()) - .try_collect::>() - .await - .unwrap(), - vec!["openai-prod".to_string(), "hf-prod".to_string()] - ); - } - - /// A server that keeps handing back the same token would otherwise spin - /// forever. - #[tokio::test] - async fn test_list_secrets_rejects_a_repeated_page_token() { - let conn = Connection::new_with_handler(|_| { - http::Response::builder() - .status(200) - .body(r#"{"secrets":[{"name":"openai-prod"}],"page_token":"same"}"#) - .unwrap() - }); - let error = conn - .list_secrets(&[], Default::default()) - .try_collect::>() - .await - .unwrap_err(); - assert!( - error.to_string().contains("repeated a page_token"), - "{error}" - ); - } - - #[tokio::test] - async fn test_drop_and_describe_address_the_secret_in_the_path() { - let conn = Connection::new_with_handler(|request| { - assert_eq!(request.url().path(), "/v1/secret/openai-prod/drop"); - // Nothing is left to say once the path names the Secret. - assert!(request.body().is_none(), "{:?}", request.body()); - http::Response::builder().status(200).body("{}").unwrap() - }); - conn.drop_secret("openai-prod", &[]).await.unwrap(); - - let conn = Connection::new_with_handler(|request| { - assert_eq!(request.url().path(), "/v1/secret/openai-prod/describe"); - assert!(request.body().is_none(), "{:?}", request.body()); - http::Response::builder() - .status(200) - .body(r#"{"name":"openai-prod","created_at_millis":1,"updated_at_millis":2}"#) - .unwrap() - }); - let info = conn.describe_secret("openai-prod", &[]).await.unwrap(); - assert_eq!(info.name, "openai-prod"); - } - - /// The namespace is part of the identifier the path addresses, joined with - /// the client's configured delimiter the way every other object's is. A - /// root Secret is therefore addressed by its bare name, and a namespaced - /// one by the joined path -- there is no body field either way. - #[tokio::test] - async fn test_a_namespace_path_is_addressed_in_the_path() { - let conn = Connection::new_with_handler(|request| { - assert_eq!( - request.url().path(), - "/v1/secret/prod$vision$openai-prod/drop" - ); - http::Response::builder().status(200).body("{}").unwrap() - }); - conn.drop_secret("openai-prod", &["prod".to_string(), "vision".to_string()]) - .await - .unwrap(); - - let conn = Connection::new_with_handler(|request| { - assert_eq!(request.url().path(), "/v1/secret/openai-prod/drop"); - http::Response::builder().status(200).body("{}").unwrap() - }); - conn.drop_secret("openai-prod", &[]).await.unwrap(); - - // Listing is namespace-scoped, so the namespace is the whole identifier. - let conn = Connection::new_with_handler(|request| { - assert_eq!( - request.url().path(), - "/v1/namespace/prod$vision/secret/list" - ); - http::Response::builder() - .status(200) - .body(r#"{"secrets":[]}"#) - .unwrap() - }); - conn.list_secrets( - &["prod".to_string(), "vision".to_string()], - Default::default(), - ) - .try_collect::>() - .await - .unwrap(); - } - /// A view description carries the schema in the namespace spec's JSON /// encoding, so the body a test asserts against is the one that encoding /// produces rather than a hand-written guess at it. diff --git a/rust/lancedb/src/remote/sql.rs b/rust/lancedb/src/remote/sql.rs index 69cba2ef5..ebe1cab49 100644 --- a/rust/lancedb/src/remote/sql.rs +++ b/rust/lancedb/src/remote/sql.rs @@ -38,6 +38,7 @@ use crate::sql::{Query, QueryDescription, QueryHandle, QueryStatus}; mod function; pub(super) mod materialized_view; mod objects; +mod secret; mod stream; use stream::{StreamingQuery, StreamingQueryHandle, Submission}; diff --git a/rust/lancedb/src/remote/sql/secret.rs b/rust/lancedb/src/remote/sql/secret.rs new file mode 100644 index 000000000..709cd6388 --- /dev/null +++ b/rust/lancedb/src/remote/sql/secret.rs @@ -0,0 +1,156 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +use super::{ + SqlClient, + materialized_view::quote_identifier, + objects::{definition, strings}, +}; +use crate::{ + database::ExecuteQueryRequest, + error::{Error, Result}, + listing::{ListingOptions, ListingPage}, + secrets::SecretInfo, + utils::{validate_namespace, validate_secret_name}, +}; +use arrow_array::{RecordBatch, StringArray}; +use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD}; +use futures::TryStreamExt; +use std::sync::Arc; +impl SqlClient { + pub(in crate::remote) async fn write_secret( + &self, + name: &str, + value: &str, + namespace: &[String], + create: bool, + ) -> Result<()> { + validate_secret_name(name)?; + validate_namespace(namespace)?; + let parameters = RecordBatch::try_from_iter(vec![( + "", + Arc::new(StringArray::from(vec![value])) as arrow_array::ArrayRef, + )]) + .map_err(|error| Error::Runtime { + message: error.to_string(), + })?; + self.object_job( + ExecuteQueryRequest { + query: format!( + "{} SECRET {} AS $1", + if create { "CREATE" } else { "ALTER" }, + quote_identifier(name) + ), + default_namespace_path: namespace.to_vec(), + parameters: Some(parameters), + }, + |_| Ok(()), + ) + .await? + .wait() + .await + } + pub(in crate::remote) async fn describe_secret( + &self, + name: &str, + namespace: &[String], + ) -> Result { + validate_secret_name(name)?; + validate_namespace(namespace)?; + self.object_job( + ExecuteQueryRequest { + query: format!("DESCRIBE SECRET {}", quote_identifier(name)), + default_namespace_path: namespace.to_vec(), + parameters: None, + }, + definition, + ) + .await? + .wait() + .await + } + pub(in crate::remote) async fn drop_secret( + &self, + name: &str, + namespace: &[String], + ) -> Result<()> { + validate_secret_name(name)?; + validate_namespace(namespace)?; + self.object_job( + ExecuteQueryRequest { + query: format!("DROP SECRET IF EXISTS {}", quote_identifier(name)), + default_namespace_path: namespace.to_vec(), + parameters: None, + }, + |_| Ok(()), + ) + .await? + .wait() + .await + } + pub(in crate::remote) async fn list_secrets( + &self, + namespace: &[String], + options: ListingOptions, + ) -> Result> { + validate_namespace(namespace)?; + let limit = options.page_limit.unwrap_or(100); + if limit == 0 || limit > i32::MAX as u32 { + return Err(Error::InvalidInput { + message: "page limit must be between 1 and i32::MAX".into(), + }); + } + let invalid_token = || Error::InvalidInput { + message: "invalid secret page token for this database or namespace".into(), + }; + let after = if let Some(token) = options.page_token.filter(|token| !token.is_empty()) { + let decoded = URL_SAFE_NO_PAD.decode(token).map_err(|_| invalid_token())?; + let (kind, database, prefix, path, last): ( + String, + String, + Option, + Vec, + String, + ) = serde_json::from_slice(&decoded).map_err(|_| invalid_token())?; + if kind != "secret" + || database != self.inner.database + || prefix != self.inner.database_prefix + || path != namespace + { + return Err(invalid_token()); + } + format!(" WHERE name > '{}'", last.replace('\'', "''")) + } else { + String::new() + }; + let query = self + .execute(ExecuteQueryRequest { + query: format!("SHOW SECRETS{after} LIMIT {}", u64::from(limit) + 1), + default_namespace_path: namespace.to_vec(), + parameters: None, + }) + .await?; + let mut reader = query.reader().await?; + let mut items = Vec::new(); + while let Some(batch) = reader.try_next().await? { + items.extend(strings(&[batch], "name")?); + } + let page_token = if items.len() > limit as usize { + items.truncate(limit as usize); + let cursor = serde_json::to_vec(&( + "secret", + &self.inner.database, + &self.inner.database_prefix, + namespace, + items.last().unwrap(), + )) + .map_err(|error| Error::Runtime { + message: error.to_string(), + })?; + Some(URL_SAFE_NO_PAD.encode(cursor)) + } else { + None + }; + Ok(ListingPage { items, page_token }) + } +} diff --git a/rust/lancedb/src/remote/sql_test.rs b/rust/lancedb/src/remote/sql_test.rs index 40b6c1b81..61f807d37 100644 --- a/rust/lancedb/src/remote/sql_test.rs +++ b/rust/lancedb/src/remote/sql_test.rs @@ -90,6 +90,7 @@ struct TestSqlService { headers: Arc>>, gate_refresh_release: Arc, exchanges: ReceivedExchanges, + result_overrides: Arc>>, /// How many exchange responses have been dropped, finished or not. exchange_drops: Arc, result: RecordBatch, @@ -127,6 +128,22 @@ impl TestSqlService { .unwrap_or_default(), }); } + + fn result_for_query(&self, query: &str) -> RecordBatch { + if query == "SELECT large message" { + self.large_result.clone() + } else if query == "SELECT dictionary" { + self.dictionary_result.clone() + } else { + self.result_overrides + .lock() + .unwrap() + .get(query) + .cloned() + .unwrap_or_else(|| self.result.clone()) + } + } + fn adaptive_response( &self, query: &str, @@ -223,6 +240,7 @@ impl Default for TestSqlService { headers: Arc::new(std::sync::Mutex::new(Vec::new())), gate_refresh_release: Arc::new(Notify::new()), exchanges: Arc::new(std::sync::Mutex::new(Vec::new())), + result_overrides: Arc::new(std::sync::Mutex::new(HashMap::new())), exchange_drops: Arc::new(AtomicUsize::new(0)), result, large_result, @@ -338,14 +356,8 @@ impl FlightService for TestSqlService { ); } if query != "SELECT empty" { - let schema = if query == "SELECT large message" { - self.large_result.schema_ref() - } else if query == "SELECT dictionary" { - self.dictionary_result.schema_ref() - } else { - self.result.schema_ref() - }; - info = info.try_with_schema(schema).unwrap(); + let schema = self.result_for_query(&query).schema(); + info = info.try_with_schema(&schema).unwrap(); } Ok(Response::new(PollInfo { info: (query != "SELECT no info" || stage > 0).then_some(info), @@ -393,13 +405,7 @@ impl FlightService for TestSqlService { let empty = ticket == b"SELECT empty"; let slow = ticket == b"SELECT slow get"; let large = ticket == b"SELECT large message"; - let result = if large { - self.large_result.clone() - } else if ticket == b"SELECT dictionary" { - self.dictionary_result.clone() - } else { - self.result.clone() - }; + let result = self.result_for_query(std::str::from_utf8(ticket).unwrap_or_default()); let schema = result.schema(); let gate_refresh_release = self.gate_refresh_release.clone(); let input = futures::stream::once(async move { @@ -598,6 +604,14 @@ fn parameterized(query: &str, parameters: RecordBatch) -> ExecuteQueryRequest { async fn start_sql_service( service: TestSqlService, client_config: ClientConfig, +) -> (SqlClient, tokio::sync::oneshot::Sender<()>) { + start_sql_service_with_prefix(service, client_config, None).await +} + +async fn start_sql_service_with_prefix( + service: TestSqlService, + client_config: ClientConfig, + database_prefix: Option<&str>, ) -> (SqlClient, tokio::sync::oneshot::Sender<()>) { let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); let address = listener.local_addr().unwrap(); @@ -615,7 +629,7 @@ async fn start_sql_service( ); let client = SqlClient::new( "analytics".to_string(), - None, + database_prefix.map(str::to_owned), "test-key".to_string(), None, Some(format!("grpc://{address}")), @@ -2372,6 +2386,232 @@ async fn function_drop_keeps_the_query_id_and_reports_catalog_removal() { } } +#[tokio::test] +async fn secret_writes_bind_credentials_and_never_put_them_in_sql() { + for defer in [false, true] { + let service = TestSqlService { + adaptive: true, + defer, + ..Default::default() + }; + let exchanges = service.exchanges.clone(); + let statements = service.statements.clone(); + let (client, shutdown) = start_sql_service(service, ClientConfig::default()).await; + let db = RemoteDatabase::new_mock(|request| -> http::Response { + panic!("unexpected REST request: {}", request.url().path()); + }) + .with_sql_client(client); + db.create_secret("api-key", "credential's value", &["team".into()]) + .await + .unwrap(); + db.alter_secret("api-key", "rotated-value", &["team".into()]) + .await + .unwrap(); + db.drop_secret("api-key", &["team".into()]).await.unwrap(); + assert_eq!( + statements.lock().unwrap().as_slice(), + ["DROP SECRET IF EXISTS \"api-key\""] + ); + let seen = exchanges.lock().unwrap(); + assert_eq!(seen.len(), 2); + assert_eq!(seen[0].0, "CREATE SECRET \"api-key\" AS $1"); + assert_eq!(seen[1].0, "ALTER SECRET \"api-key\" AS $1"); + assert_eq!( + seen[0] + .1 + .as_ref() + .unwrap() + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + "credential's value" + ); + assert_eq!( + seen[1] + .1 + .as_ref() + .unwrap() + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + "rotated-value" + ); + shutdown.send(()).unwrap(); + } +} +#[tokio::test] +async fn secret_description_uses_the_sql_metadata_projection() { + let definition = r#"{"name":"api-key","created_at_millis":1,"updated_at_millis":2,"history":[{"generation":2,"at_millis":2}]}"#; + let result = RecordBatch::try_from_iter(vec![( + "definition", + Arc::new(StringArray::from(vec![definition])) as ArrayRef, + )]) + .unwrap(); + let service = TestSqlService { + adaptive: true, + defer: true, + result, + ..Default::default() + }; + let (client, shutdown) = start_sql_service(service, ClientConfig::default()).await; + let db = RemoteDatabase::new_mock(|_| -> http::Response { + panic!("unexpected REST request"); + }) + .with_sql_client(client); + let info = db.describe_secret("api-key", &[]).await.unwrap(); + assert_eq!(info.name, "api-key"); + assert_eq!(info.updated_at_millis, 2); + shutdown.send(()).unwrap(); +} + +#[tokio::test] +async fn secret_listing_uses_bounded_sql_and_scoped_cursors() { + for defer in [false, true] { + let first_result = RecordBatch::try_from_iter(vec![( + "name", + Arc::new(StringArray::from(vec!["api-key", "hf'prod", "openai-prod"])) as ArrayRef, + )]) + .unwrap(); + let second_result = RecordBatch::try_from_iter(vec![( + "name", + Arc::new(StringArray::from(vec!["openai-prod"])) as ArrayRef, + )]) + .unwrap(); + let result_overrides = Arc::new(std::sync::Mutex::new(HashMap::from([ + ("SHOW SECRETS LIMIT 3".to_string(), first_result), + ( + "SHOW SECRETS WHERE name > 'hf''prod' LIMIT 3".to_string(), + second_result, + ), + ]))); + let service = TestSqlService { + adaptive: true, + defer, + result_overrides, + ..Default::default() + }; + let statements = service.statements.clone(); + let headers = service.headers.clone(); + let (client, shutdown) = start_sql_service_with_prefix( + service, + ClientConfig::default(), + Some("tenant/production"), + ) + .await; + let db = RemoteDatabase::new_mock(|_| -> http::Response { + panic!("unexpected REST request"); + }) + .with_sql_client(client); + let namespace = vec!["team".to_owned(), "reports".to_owned()]; + + let first = db + .list_secrets( + &namespace, + crate::listing::ListingOptions::default().page_limit(2), + ) + .await + .unwrap(); + assert_eq!(first.items, ["api-key", "hf'prod"]); + let token = first.page_token.expect("a third row must produce a token"); + + let second = db + .list_secrets( + &namespace, + crate::listing::ListingOptions::default() + .page_limit(2) + .page_token(&token), + ) + .await + .unwrap(); + assert_eq!(second.items, ["openai-prod"]); + assert!(second.page_token.is_none()); + + assert_eq!( + *statements.lock().unwrap(), + [ + "SHOW SECRETS LIMIT 3", + "SHOW SECRETS WHERE name > 'hf''prod' LIMIT 3" + ] + ); + { + let headers = headers.lock().unwrap(); + assert!(!headers.is_empty()); + for header in headers.iter() { + assert_eq!(header.database, "analytics"); + assert_eq!(header.database_prefix, "tenant/production"); + assert_eq!(header.namespace_path, "team$reports"); + } + } + + for (path, token) in [ + (vec!["other".to_owned()], token), + (namespace.clone(), "invalid".to_owned()), + ] { + assert!(matches!( + db.list_secrets( + &path, + crate::listing::ListingOptions::default().page_token(token) + ) + .await, + Err(Error::InvalidInput { .. }) + )); + } + assert_eq!(statements.lock().unwrap().len(), 2); + shutdown.send(()).unwrap(); + } +} + +#[tokio::test] +async fn secret_listing_handles_empty_pages_and_rejects_invalid_limits() { + let empty = RecordBatch::try_from_iter(vec![( + "name", + Arc::new(StringArray::from(Vec::::new())) as ArrayRef, + )]) + .unwrap(); + let service = TestSqlService { + adaptive: true, + defer: true, + result: empty, + ..Default::default() + }; + let statements = service.statements.clone(); + let (client, shutdown) = start_sql_service(service, ClientConfig::default()).await; + let db = RemoteDatabase::new_mock(|_| -> http::Response { + panic!("unexpected REST request"); + }) + .with_sql_client(client); + let namespace = vec!["team".to_owned(), "reports".to_owned()]; + + let page = db + .list_secrets( + &namespace, + crate::listing::ListingOptions::default() + .page_limit(2) + .page_token(""), + ) + .await + .unwrap(); + assert!(page.items.is_empty()); + assert!(page.page_token.is_none()); + + for limit in [0, i32::MAX as u32 + 1] { + assert!(matches!( + db.list_secrets( + &namespace, + crate::listing::ListingOptions::default().page_limit(limit) + ) + .await, + Err(Error::InvalidInput { .. }) + )); + } + assert_eq!(*statements.lock().unwrap(), ["SHOW SECRETS LIMIT 3"]); + shutdown.send(()).unwrap(); +} + #[tokio::test] async fn function_listing_uses_bounded_sql_and_scoped_version_cursors() { const VERSION: &str = include_str!( diff --git a/rust/lancedb/src/secrets.rs b/rust/lancedb/src/secrets.rs index c980e51ee..c15c7ead9 100644 --- a/rust/lancedb/src/secrets.rs +++ b/rust/lancedb/src/secrets.rs @@ -19,7 +19,7 @@ use serde_json::Value; /// /// Returned by [`crate::connection::Connection::describe_secret`]. There is no /// field for the credential and no method that could produce one. -#[derive(Debug, Clone, PartialEq, Eq, serde::Deserialize)] +#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)] pub struct SecretInfo { /// The Secret's database-scoped name. pub name: String,