From a4d34d7b5d0ebf8e71a509a0aa2ce18fe9e4778e Mon Sep 17 00:00:00 2001 From: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Thu, 6 Aug 2026 19:56:39 +0000 Subject: [PATCH] fix(rust): preserve namespace credential providers --- rust/lancedb/src/database/namespace.rs | 25 ++- rust/lancedb/src/io/object_store.rs | 221 +++++++++++++++++++++++-- 2 files changed, 221 insertions(+), 25 deletions(-) diff --git a/rust/lancedb/src/database/namespace.rs b/rust/lancedb/src/database/namespace.rs index 740e11645..4d063a584 100644 --- a/rust/lancedb/src/database/namespace.rs +++ b/rust/lancedb/src/database/namespace.rs @@ -34,6 +34,7 @@ use crate::database::read_freshness::{ FreshnessBaselines, ReadFreshnessContextProvider, TableFreshness, }; use crate::error::{Error, Result}; +use crate::io::object_store::install_atomic_aws_provider; use crate::table::{NativeTable, map_namespace_lance_error}; use lance::dataset::WriteMode; @@ -101,6 +102,8 @@ impl LanceNamespaceDatabase { session: Option>, namespace_client_pushdown_operations: HashSet, ) -> Self { + let session = session.unwrap_or_else(|| Arc::new(lance::session::Session::default())); + install_atomic_aws_provider(&session); // Client is pre-built, so we can't install the freshness provider here; // baselines are still tracked for a uniform bump path. let delimiter = resolve_delimiter(&namespace_client_properties); @@ -108,7 +111,7 @@ impl LanceNamespaceDatabase { namespace: namespace_client, storage_options, read_consistency_interval, - session, + session: Some(session), uri: format!("namespace://{}", namespace_client_impl), pushdown_operations: namespace_client_pushdown_operations, ns_impl: namespace_client_impl, @@ -153,13 +156,13 @@ impl LanceNamespaceDatabase { pushdown_operations: HashSet, new_table_config: NewTableConfig, ) -> Result { + let session = session.unwrap_or_else(|| Arc::new(lance::session::Session::default())); + install_atomic_aws_provider(&session); let mut builder = ConnectBuilder::new(ns_impl); for (key, value) in ns_properties.clone() { builder = builder.property(key, value); } - if let Some(ref sess) = session { - builder = builder.session(sess.clone()); - } + builder = builder.session(session.clone()); // Install the read-freshness provider before building the client. let freshness_baselines: FreshnessBaselines = Arc::new(Mutex::new(HashMap::new())); @@ -177,7 +180,7 @@ impl LanceNamespaceDatabase { namespace, storage_options, read_consistency_interval, - session, + session: Some(session), uri: format!("namespace://{}", ns_impl), pushdown_operations, ns_impl: ns_impl.to_string(), @@ -654,9 +657,17 @@ mod tests { properties.insert("root".to_string(), root_path); // This should succeed with directory-based namespace - let result = connect_namespace("dir", properties).execute().await; + let connection = connect_namespace("dir", properties) + .execute() + .await + .unwrap(); + let database = connection + .database() + .as_any() + .downcast_ref::() + .unwrap(); - assert!(result.is_ok()); + assert!(database.session.is_some()); } #[tokio::test] diff --git a/rust/lancedb/src/io/object_store.rs b/rust/lancedb/src/io/object_store.rs index 2c7ec442f..8db6dada4 100644 --- a/rust/lancedb/src/io/object_store.rs +++ b/rust/lancedb/src/io/object_store.rs @@ -155,6 +155,17 @@ impl ObjectStoreProvider for AtomicAwsStoreProvider { base_path: url::Url, params: &ObjectStoreParams, ) -> lance_core::Result { + // Caller-supplied credential providers and refreshable storage options are already + // atomic credential authorities. Preserve Lance's precedence and refresh behavior. + if params.aws_credentials.is_some() + || params + .storage_options_accessor + .as_ref() + .is_some_and(|accessor| accessor.has_provider()) + { + return self.inner.new_store(base_path, params).await; + } + let storage_options = params.storage_options().cloned().unwrap_or_default(); let Some(credential) = explicit_aws_credential(&storage_options)? else { return self.inner.new_store(base_path, params).await; @@ -169,17 +180,10 @@ impl ObjectStoreProvider for AtomicAwsStoreProvider { // the correctly initialized Lance ObjectStore shell used below for OpenDAL. let mut native_options = storage_options.clone(); native_options.insert("use_opendal".to_string(), "false".to_string()); - let provider = params - .storage_options_accessor - .as_ref() - .and_then(|accessor| accessor.provider().cloned()); - let accessor = if let Some(provider) = provider { - StorageOptionsAccessor::with_initial_and_provider(native_options, provider) - } else { - StorageOptionsAccessor::with_static_options(native_options) - }; let mut native_params = params.clone(); - native_params.storage_options_accessor = Some(Arc::new(accessor)); + native_params.storage_options_accessor = Some(Arc::new( + StorageOptionsAccessor::with_static_options(native_options), + )); native_params.aws_credentials = Some(Arc::new(StaticCredentialProvider::new(credential))); let mut store = self .inner @@ -267,16 +271,21 @@ pub(crate) fn install_atomic_aws_provider(_session: &lance::session::Session) {} /// Apply storage options to object store parameters. /// -/// Explicit AWS credentials are also installed as a single credential provider. This keeps -/// an ambient `AWS_SESSION_TOKEN` from being combined with an explicitly supplied access key -/// and secret key when Lance adds missing options from the environment. +/// Static credentials for Lance's native S3 backend are installed directly. OpenDAL options are +/// left for the session's credential-safe provider because that backend ignores this field. +/// Caller-supplied and refreshable credential providers retain their existing precedence. pub(crate) fn set_storage_options( params: &mut ObjectStoreParams, storage_options: HashMap, provider: Option>, ) { #[cfg(feature = "aws")] - if provider.is_none() && params.aws_credentials.is_none() { + if provider.is_none() + && params.aws_credentials.is_none() + && !storage_options + .get("use_opendal") + .is_some_and(|value| value == "true") + { params.aws_credentials = explicit_aws_credentials(&storage_options); } @@ -475,7 +484,10 @@ impl WrappingObjectStore for MirroringObjectStoreWrapper { mod credential_tests { use super::*; use lance_io::object_store::providers::aws::build_aws_credential; - use std::sync::atomic::{AtomicBool, Ordering}; + use std::sync::{ + Mutex, + atomic::{AtomicBool, AtomicUsize, Ordering}, + }; use std::time::Duration; #[derive(Debug)] @@ -496,6 +508,78 @@ mod credential_tests { } } + #[derive(Debug)] + struct RotatingOptionsProvider { + fetches: Arc, + } + + #[async_trait] + impl StorageOptionsProvider for RotatingOptionsProvider { + async fn fetch_storage_options( + &self, + ) -> lance_core::Result>> { + self.fetches.fetch_add(1, Ordering::SeqCst); + Ok(Some(HashMap::from([ + ("aws_access_key_id".to_string(), "refreshed-key".to_string()), + ( + "aws_secret_access_key".to_string(), + "refreshed-secret".to_string(), + ), + ]))) + } + + fn provider_id(&self) -> String { + "rotating-test-provider".to_string() + } + } + + #[derive(Debug, PartialEq, Eq)] + struct ObservedCredential { + key_id: String, + token: Option, + } + + #[derive(Debug)] + struct ResolvingProvider { + resolved_credential: Arc>>, + } + + #[async_trait] + impl ObjectStoreProvider for ResolvingProvider { + async fn new_store( + &self, + _base_path: url::Url, + params: &ObjectStoreParams, + ) -> lance_core::Result { + let storage_options = params + .storage_options() + .cloned() + .unwrap_or_default() + .into_iter() + .filter_map(|(key, value)| { + AmazonS3ConfigKey::from_str(&key) + .ok() + .map(|key| (key, value)) + }) + .collect::>(); + let (provider, _) = build_aws_credential( + Duration::from_secs(60), + params.aws_credentials.clone(), + Some(&storage_options), + Some("us-east-1".to_string()), + params.storage_options_accessor.clone(), + None, + ) + .await?; + let credential = provider.get_credential().await?; + *self.resolved_credential.lock().unwrap() = Some(ObservedCredential { + key_id: credential.key_id.clone(), + token: credential.token.clone(), + }); + Err(lance_core::Error::invalid_input("recorded test request")) + } + } + #[tokio::test] async fn explicit_aws_credentials_do_not_inherit_an_ambient_session_token() { let storage_options = HashMap::from([ @@ -553,11 +637,13 @@ mod credential_tests { ), ("AWS_REGION".to_string(), "us-east-1".to_string()), ]; + let params = object_store_params_from_storage_options(storage_options.clone()); let options = atomic_opendal_options(&storage_options, environment) .unwrap() .unwrap(); + assert!(params.aws_credentials.is_none()); assert_eq!(options.get("aws_access_key_id").unwrap(), "explicit-key"); assert_eq!( options.get("aws_secret_access_key").unwrap(), @@ -622,7 +708,7 @@ mod credential_tests { } #[tokio::test] - async fn namespace_and_opendal_params_reach_the_atomic_provider_boundary() { + async fn public_namespace_connection_installs_the_atomic_provider() { let saw_atomic_credentials = Arc::new(AtomicBool::new(false)); let registry = Arc::new(ObjectStoreRegistry::default()); registry.insert( @@ -631,8 +717,24 @@ mod credential_tests { saw_atomic_credentials: saw_atomic_credentials.clone(), }), ); - let session = lance::session::Session::new(16, 16, registry.clone()); - install_atomic_aws_provider(&session); + let session = Arc::new(lance::session::Session::new(16, 16, registry.clone())); + let root = tempfile::tempdir().unwrap(); + + crate::connect_namespace( + "dir", + HashMap::from([( + "root".to_string(), + root.path().to_string_lossy().into_owned(), + )]), + ) + .storage_options([ + ("aws_access_key_id", "explicit-key"), + ("aws_secret_access_key", "explicit-secret"), + ]) + .session(session) + .execute() + .await + .unwrap(); // DirectoryNamespaceBuilder constructs fresh params with only a static accessor. The // OpenDAL selector is included to verify that both paths cross the installed boundary. @@ -660,6 +762,89 @@ mod credential_tests { assert!(error.to_string().contains("recorded test request")); assert!(saw_atomic_credentials.load(Ordering::SeqCst)); } + + #[tokio::test] + async fn dynamic_storage_options_provider_remains_the_credential_authority() { + let fetches = Arc::new(AtomicUsize::new(0)); + let resolved_credential = Arc::new(Mutex::new(None)); + let provider = AtomicAwsStoreProvider { + inner: Arc::new(ResolvingProvider { + resolved_credential: resolved_credential.clone(), + }), + }; + let params = ObjectStoreParams { + storage_options_accessor: Some(Arc::new( + StorageOptionsAccessor::with_initial_and_provider( + HashMap::from([ + ("aws_access_key_id".to_string(), "expired-key".to_string()), + ( + "aws_secret_access_key".to_string(), + "expired-secret".to_string(), + ), + ("expires_at_millis".to_string(), "0".to_string()), + ]), + Arc::new(RotatingOptionsProvider { + fetches: fetches.clone(), + }), + ), + )), + ..Default::default() + }; + + provider + .new_store(url::Url::parse("s3://bucket/table").unwrap(), ¶ms) + .await + .unwrap_err(); + + assert_eq!(fetches.load(Ordering::SeqCst), 1); + assert_eq!( + *resolved_credential.lock().unwrap(), + Some(ObservedCredential { + key_id: "refreshed-key".to_string(), + token: None, + }) + ); + } + + #[tokio::test] + async fn caller_supplied_aws_provider_remains_the_credential_authority() { + let resolved_credential = Arc::new(Mutex::new(None)); + let provider = AtomicAwsStoreProvider { + inner: Arc::new(ResolvingProvider { + resolved_credential: resolved_credential.clone(), + }), + }; + let params = ObjectStoreParams { + aws_credentials: Some(Arc::new(StaticCredentialProvider::new(AwsCredential { + key_id: "provider-key".to_string(), + secret_key: "provider-secret".to_string(), + token: None, + }))), + storage_options_accessor: Some(Arc::new(StorageOptionsAccessor::with_static_options( + HashMap::from([ + ("aws_access_key_id".to_string(), "option-key".to_string()), + ( + "aws_secret_access_key".to_string(), + "option-secret".to_string(), + ), + ]), + ))), + ..Default::default() + }; + + provider + .new_store(url::Url::parse("s3://bucket/table").unwrap(), ¶ms) + .await + .unwrap_err(); + + assert_eq!( + *resolved_credential.lock().unwrap(), + Some(ObservedCredential { + key_id: "provider-key".to_string(), + token: None, + }) + ); + } } // windows pathing can't be simply concatenated