From 9e03bf97d400572e38d265f5c59acd990c4d818a Mon Sep 17 00:00:00 2001 From: Jack Ye Date: Thu, 12 Nov 2026 08:03:11 -0800 Subject: [PATCH] feat(python): support exist_ok when creating secrets (#4978) Add `exist_ok=False` to synchronous and asynchronous Python `create_secret`. Passing `True` uses `CREATE SECRET IF NOT EXISTS`, preserving an existing credential. The default continues to fail on an existing name; rotation stays explicit through `alter_secret`. ```python db.create_secret("api-key", value, exist_ok=True) await async_db.create_secret("api-key", value, exist_ok=True) ``` The Rust API gains an awaitable builder with `.exist_ok(true)` while existing `create_secret(...).await` calls keep working. Includes Python SQL round-trip coverage and Rust coverage for direct and deferred SQL responses. --- python/python/lancedb/_lancedb.pyi | 7 ++- python/python/lancedb/db.py | 48 ++++++++++++--- python/python/lancedb/remote/db.py | 13 ++++- python/python/tests/test_sql.py | 50 ++++++++++++++++ python/src/connection.rs | 4 +- rust/lancedb/src/connection.rs | 29 +++++++--- rust/lancedb/src/connection/create_secret.rs | 61 ++++++++++++++++++++ rust/lancedb/src/database.rs | 5 +- rust/lancedb/src/remote/db.rs | 12 +++- rust/lancedb/src/remote/sql.rs | 1 + rust/lancedb/src/remote/sql/secret.rs | 29 +++++++--- rust/lancedb/src/remote/sql_test.rs | 48 +++++++++++++-- 12 files changed, 269 insertions(+), 38 deletions(-) create mode 100644 rust/lancedb/src/connection/create_secret.rs diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 11798a520..a64014cb7 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -228,7 +228,12 @@ class Connection(object): self, name: str, namespace_path: Optional[List[str]] = None ) -> None: ... async def create_secret( - self, name: str, value: str, namespace_path: Optional[List[str]] = None + self, + name: str, + value: str, + namespace_path: Optional[List[str]] = None, + *, + exist_ok: bool = False, ) -> None: ... async def alter_secret( self, name: str, value: str, namespace_path: Optional[List[str]] = None diff --git a/python/python/lancedb/db.py b/python/python/lancedb/db.py index e2138520a..59fa5b4c6 100644 --- a/python/python/lancedb/db.py +++ b/python/python/lancedb/db.py @@ -1158,14 +1158,24 @@ class DBConnection(EnforceOverrides): ) def create_secret( - self, name: str, value: str, *, namespace_path: Optional[List[str]] = None + self, + name: str, + value: str, + *, + exist_ok: bool = False, + namespace_path: Optional[List[str]] = None, ) -> None: """Create a named Secret in this database. - Fails if the name is taken, so a create never silently becomes a - rotation. Nothing reads the value back: it is bound to a Function by + Fails if the name is taken unless ``exist_ok=True``, which leaves the + existing value unchanged. Use :meth:`alter_secret` to rotate a Secret. + Nothing reads the value back: it is bound to a Function by name and resolved by the service when that Function runs. Local connections raise ``NotImplementedError``. + + Examples + -------- + >>> db.create_secret("api-key", "credential", exist_ok=True) # doctest: +SKIP """ raise NotImplementedError( _append_doc_url( @@ -2492,9 +2502,18 @@ class LanceDBConnection(DBConnection): @override def create_secret( - self, name: str, value: str, *, namespace_path: Optional[List[str]] = None + self, + name: str, + value: str, + *, + exist_ok: bool = False, + namespace_path: Optional[List[str]] = None, ) -> None: - LOOP.run(self._conn.create_secret(name, value, namespace_path=namespace_path)) + LOOP.run( + self._conn.create_secret( + name, value, exist_ok=exist_ok, namespace_path=namespace_path + ) + ) @override def alter_secret( @@ -3949,17 +3968,30 @@ class AsyncConnection(object): ) async def create_secret( - self, name: str, value: str, *, namespace_path: Optional[List[str]] = None + self, + name: str, + value: str, + *, + exist_ok: bool = False, + namespace_path: Optional[List[str]] = None, ) -> None: """Create a named Secret in this database. - Fails if the name is taken, so a create never silently becomes a - rotation. Nothing reads the value back. + Fails if the name is taken unless ``exist_ok=True``, which leaves the + existing value unchanged. Use :meth:`alter_secret` to rotate a Secret. + Nothing reads the value back. + + Examples + -------- + >>> await db.create_secret( # doctest: +SKIP + ... "api-key", "credential", exist_ok=True + ... ) """ await self._inner.create_secret( validate_secret_name(name), value, list(validate_namespace_path(namespace_path)), + exist_ok=exist_ok, ) async def alter_secret( diff --git a/python/python/lancedb/remote/db.py b/python/python/lancedb/remote/db.py index 4219ca40b..a06af5d62 100644 --- a/python/python/lancedb/remote/db.py +++ b/python/python/lancedb/remote/db.py @@ -1129,9 +1129,18 @@ class RemoteDBConnection(DBConnection): @override def create_secret( - self, name: str, value: str, *, namespace_path: Optional[List[str]] = None + self, + name: str, + value: str, + *, + exist_ok: bool = False, + namespace_path: Optional[List[str]] = None, ) -> None: - LOOP.run(self._conn.create_secret(name, value, namespace_path=namespace_path)) + LOOP.run( + self._conn.create_secret( + name, value, exist_ok=exist_ok, namespace_path=namespace_path + ) + ) @override def alter_secret( diff --git a/python/python/tests/test_sql.py b/python/python/tests/test_sql.py index b4c555635..a479fa5d1 100644 --- a/python/python/tests/test_sql.py +++ b/python/python/tests/test_sql.py @@ -293,6 +293,56 @@ def test_parameters_reach_the_server_as_the_values_that_were_sent(): assert echoed.column("k").to_pylist() == [10] +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize( + ("exist_ok", "status"), + [ + (None, "created"), + (False, "created"), + (True, "created"), + (True, "already_exists"), + ], +) +async def test_create_secret_exist_ok_reaches_sql(asynchronous, exist_ok, status): + flight = pytest.importorskip("pyarrow.flight") + + class SecretServer(flight.FlightServerBase): + def __init__(self): + super().__init__("grpc://127.0.0.1:0") + self.writes = [] + + def do_exchange(self, context, descriptor, reader, writer): + self.writes.append((descriptor.command, reader.read_all())) + result = pa.table({"status": [status]}) + writer.begin(result.schema) + writer.write_table(result) + writer.close() + + with SecretServer() as server: + options = { + "api_key": "test-key", + "host_override": "http://localhost:10024", + "sql_host_override": f"grpc://127.0.0.1:{server.port}", + } + kwargs = {"namespace_path": ["team", "production"]} + if exist_ok is not None: + kwargs["exist_ok"] = exist_ok + if asynchronous: + db = await lancedb.connect_async("db://analytics", **options) + await db.create_secret("api-key", "credential's value", **kwargs) + else: + db = lancedb.connect("db://analytics", **options) + db.create_secret("api-key", "credential's value", **kwargs) + ((command, parameters),) = server.writes + + prefix = "CREATE SECRET IF NOT EXISTS" if exist_ok else "CREATE SECRET" + sql = f'{prefix} "api-key" TYPE = GENERIC_STRING SECRET_STRING = $1' + assert sql.encode() in command + assert b"credential's value" not in command + assert parameters.column(0).to_pylist() == ["credential's value"] + + @pytest.mark.parametrize("job_id", [None, "sql-job-1"]) def test_sql_input_error_metadata_crosses_native_boundary(job_id): flight = pytest.importorskip("pyarrow.flight") diff --git a/python/src/connection.rs b/python/src/connection.rs index 69cf54c99..1a99c1372 100644 --- a/python/src/connection.rs +++ b/python/src/connection.rs @@ -1036,18 +1036,20 @@ impl Connection { }) } - #[pyo3(signature = (name, value, namespace_path=None))] + #[pyo3(signature = (name, value, namespace_path=None, *, exist_ok=false))] pub fn create_secret( self_: PyRef<'_, Self>, name: String, value: String, namespace_path: Option>, + exist_ok: bool, ) -> PyResult> { let inner = self_.get_inner()?.clone(); let namespace_path = namespace_path.unwrap_or_default(); future_into_py(self_.py(), async move { inner .create_secret(name, value, &namespace_path) + .exist_ok(exist_ok) .await .infer_error() }) diff --git a/rust/lancedb/src/connection.rs b/rust/lancedb/src/connection.rs index a360818b5..fa35dc459 100644 --- a/rust/lancedb/src/connection.rs +++ b/rust/lancedb/src/connection.rs @@ -51,8 +51,11 @@ pub use lance_file::version::LanceFileVersion; use lance_io::object_store::StorageOptions; use lance_io::object_store::{StorageOptionsAccessor, StorageOptionsProvider}; +mod create_secret; mod create_table; +pub use create_secret::CreateSecretBuilder; + pub(crate) fn merge_storage_options( store_params: &mut ObjectStoreParams, pairs: impl IntoIterator, @@ -853,20 +856,30 @@ impl Connection { /// Create a named Secret in this database. /// - /// Fails if the name is taken, so a create can never silently become a - /// rotation. There is no API that reads a stored credential back; the only + /// Fails if the name is taken unless [`CreateSecretBuilder::exist_ok`] is true, + /// in which case the existing value is left unchanged. There is no API that + /// reads a stored credential back; the only /// consumer is a Function that binds the Secret by name. LanceDB Embedded databases /// return [`Error::NotSupported`]. - pub async fn create_secret( + /// + /// ``` + /// # async fn example(connection: &lancedb::Connection) -> lancedb::Result<()> { + /// connection.create_secret("api-key", "credential", &[]).exist_ok(true).await?; + /// # Ok(()) + /// # } + /// ``` + pub fn create_secret( &self, name: impl AsRef, value: impl AsRef, namespace_path: &[String], - ) -> Result<()> { - validate_secret_reference(name.as_ref(), namespace_path)?; - self.internal - .create_secret(name.as_ref(), value.as_ref(), namespace_path) - .await + ) -> CreateSecretBuilder { + CreateSecretBuilder::new( + self.internal.clone(), + name.as_ref().to_owned(), + value.as_ref().to_owned(), + namespace_path.to_vec(), + ) } /// Replace the credential behind an existing Secret. diff --git a/rust/lancedb/src/connection/create_secret.rs b/rust/lancedb/src/connection/create_secret.rs new file mode 100644 index 000000000..265b97571 --- /dev/null +++ b/rust/lancedb/src/connection/create_secret.rs @@ -0,0 +1,61 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright The LanceDB Authors + +use std::future::{Future, IntoFuture}; +use std::pin::Pin; +use std::sync::Arc; + +use crate::Result; +use crate::database::Database; +use crate::utils::validate_secret_reference; + +/// Configure creation of a named Secret with [`super::Connection::create_secret`]. +#[must_use] +pub struct CreateSecretBuilder { + database: Arc, + name: String, + value: String, + namespace_path: Vec, + exist_ok: bool, +} + +impl CreateSecretBuilder { + pub(super) fn new( + database: Arc, + name: String, + value: String, + namespace_path: Vec, + ) -> Self { + Self { + database, + name, + value, + namespace_path, + exist_ok: false, + } + } + + /// Succeed without changing the stored value when the Secret already exists. + /// Defaults to false, which reports an error for an existing name. + pub fn exist_ok(mut self, exist_ok: bool) -> Self { + self.exist_ok = exist_ok; + self + } + + /// Create the Secret, preserving any existing value when `exist_ok` is true. + pub async fn execute(self) -> Result<()> { + validate_secret_reference(&self.name, &self.namespace_path)?; + self.database + .create_secret(&self.name, &self.value, &self.namespace_path, self.exist_ok) + .await + } +} + +impl IntoFuture for CreateSecretBuilder { + type Output = Result<()>; + type IntoFuture = Pin + Send>>; + + fn into_future(self) -> Self::IntoFuture { + Box::pin(self.execute()) + } +} diff --git a/rust/lancedb/src/database.rs b/rust/lancedb/src/database.rs index bcdfd94b8..7b4859faa 100644 --- a/rust/lancedb/src/database.rs +++ b/rust/lancedb/src/database.rs @@ -635,13 +635,14 @@ pub trait Database: secret_catalog_not_supported() } - /// Create a named Secret in this database. Fails if the name is taken, so - /// a create can never silently become a rotation. + /// Create a named Secret in this database. If the name is taken, fail unless + /// `exist_ok` is true, in which case the stored value is left unchanged. async fn create_secret( &self, _name: &str, _value: &str, _namespace_path: &[String], + _exist_ok: bool, ) -> Result<()> { secret_catalog_not_supported() } diff --git a/rust/lancedb/src/remote/db.rs b/rust/lancedb/src/remote/db.rs index 549baf345..31808f244 100644 --- a/rust/lancedb/src/remote/db.rs +++ b/rust/lancedb/src/remote/db.rs @@ -44,8 +44,8 @@ use super::client::{ ClientConfig, HeaderProvider, HttpSend, ID_DELIMITER, RequestResultExt, RestfulLanceDbClient, Sender, }; -use super::sql::SqlClient; use super::sql::materialized_view::quote_identifier; +use super::sql::{SecretWriteOperation, SqlClient}; use super::storage_integration as integration_sql; use super::table::RemoteTable; use super::util::parse_server_version; @@ -1044,11 +1044,17 @@ impl Database for RemoteDatabase { name: &str, value: &str, namespace_path: &[String], + exist_ok: bool, ) -> Result<()> { crate::utils::validate_secret_name(name)?; validate_namespace(namespace_path)?; self.sql_client()? - .write_secret(name, value, namespace_path, true) + .write_secret( + name, + value, + namespace_path, + SecretWriteOperation::Create { exist_ok }, + ) .await .map_err(direct_catalog_resource_error) } @@ -1056,7 +1062,7 @@ impl Database for RemoteDatabase { crate::utils::validate_secret_name(name)?; validate_namespace(namespace_path)?; self.sql_client()? - .write_secret(name, value, namespace_path, false) + .write_secret(name, value, namespace_path, SecretWriteOperation::Alter) .await .map_err(direct_catalog_resource_error) } diff --git a/rust/lancedb/src/remote/sql.rs b/rust/lancedb/src/remote/sql.rs index 0b87ff5a5..949b078e3 100644 --- a/rust/lancedb/src/remote/sql.rs +++ b/rust/lancedb/src/remote/sql.rs @@ -40,6 +40,7 @@ mod function; pub(super) mod materialized_view; mod objects; mod secret; +pub(super) use secret::SecretWriteOperation; mod stream; mod view; diff --git a/rust/lancedb/src/remote/sql/secret.rs b/rust/lancedb/src/remote/sql/secret.rs index bdb70ef6f..6ddc78645 100644 --- a/rust/lancedb/src/remote/sql/secret.rs +++ b/rust/lancedb/src/remote/sql/secret.rs @@ -13,6 +13,12 @@ use arrow_array::{Array, Int64Array, RecordBatch, StringArray}; use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD}; use futures::TryStreamExt; use std::sync::Arc; + +pub(in crate::remote) enum SecretWriteOperation { + Create { exist_ok: bool }, + Alter, +} + fn secret_info(batches: Vec) -> Result { if batches.iter().map(RecordBatch::num_rows).sum::() != 1 { return Err(Error::Runtime { @@ -47,10 +53,21 @@ impl SqlClient { name: &str, value: &str, namespace: &[String], - create: bool, + operation: SecretWriteOperation, ) -> Result<()> { validate_secret_name(name)?; validate_namespace(namespace)?; + let (prefix, assignment) = match operation { + SecretWriteOperation::Create { exist_ok } => ( + if exist_ok { + "CREATE SECRET IF NOT EXISTS" + } else { + "CREATE SECRET" + }, + "TYPE = GENERIC_STRING ", + ), + SecretWriteOperation::Alter => ("ALTER SECRET", "SET "), + }; let parameters = RecordBatch::try_from_iter(vec![( "", Arc::new(StringArray::from(vec![value])) as arrow_array::ArrayRef, @@ -61,14 +78,10 @@ impl SqlClient { self.object_job( ExecuteQueryRequest { query: format!( - "{} SECRET {} {}SECRET_STRING = $1", - if create { "CREATE" } else { "ALTER" }, + "{} {} {}SECRET_STRING = $1", + prefix, quote_identifier(name), - if create { - "TYPE = GENERIC_STRING " - } else { - "SET " - } + assignment ), default_namespace_path: namespace.to_vec(), parameters: Some(parameters), diff --git a/rust/lancedb/src/remote/sql_test.rs b/rust/lancedb/src/remote/sql_test.rs index cfaf2a625..471966731 100644 --- a/rust/lancedb/src/remote/sql_test.rs +++ b/rust/lancedb/src/remote/sql_test.rs @@ -3214,29 +3214,55 @@ async fn secret_writes_bind_credentials_and_never_put_them_in_sql() { }; let exchanges = service.exchanges.clone(); let statements = service.statements.clone(); + let headers = service.headers.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()]) + let connection = Connection::new( + Arc::new(db), + Arc::new(crate::embeddings::MemoryRegistry::new()), + ); + connection + .create_secret("api-key", "credential's value", &["team".into()]) .await .unwrap(); - db.alter_secret("api-key", "rotated-value", &["team".into()]) + connection + .create_secret("api-key", "ignored-value", &["team".into()]) + .exist_ok(true) + .await + .unwrap(); + connection + .alter_secret("api-key", "rotated-value", &["team".into()]) + .await + .unwrap(); + connection + .drop_secret("api-key", &["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.len(), 3); assert_eq!( seen[0].0, "CREATE SECRET \"api-key\" TYPE = GENERIC_STRING SECRET_STRING = $1" ); - assert_eq!(seen[1].0, "ALTER SECRET \"api-key\" SET SECRET_STRING = $1"); + assert_eq!( + seen[1].0, + "CREATE SECRET IF NOT EXISTS \"api-key\" TYPE = GENERIC_STRING SECRET_STRING = $1" + ); + assert_eq!(seen[2].0, "ALTER SECRET \"api-key\" SET SECRET_STRING = $1"); + assert!( + headers + .lock() + .unwrap() + .iter() + .all(|header| header.namespace_path == "team") + ); assert_eq!( seen[0] .1 @@ -3259,6 +3285,18 @@ async fn secret_writes_bind_credentials_and_never_put_them_in_sql() { .downcast_ref::() .unwrap() .value(0), + "ignored-value" + ); + assert_eq!( + seen[2] + .1 + .as_ref() + .unwrap() + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), "rotated-value" ); shutdown.send(()).unwrap();