From 7db2628c876ec4418c1f5ab79169ec90045c5665 Mon Sep 17 00:00:00 2001 From: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Thu, 6 Aug 2026 21:15:17 +0000 Subject: [PATCH] fix(rust): preserve AWS store provider contracts --- rust/lancedb/src/database/listing.rs | 75 ++++++++++++++- rust/lancedb/src/io/object_store.rs | 125 +++++++++++++++++++++++-- rust/lancedb/src/table.rs | 131 ++++++++++++++++++++++++++- 3 files changed, 320 insertions(+), 11 deletions(-) diff --git a/rust/lancedb/src/database/listing.rs b/rust/lancedb/src/database/listing.rs index b57e0e27e..be360b1da 100644 --- a/rust/lancedb/src/database/listing.rs +++ b/rust/lancedb/src/database/listing.rs @@ -24,7 +24,7 @@ use crate::database::ReadConsistency; use crate::database::namespace::LanceNamespaceDatabase; use crate::error::{CreateDirSnafu, Error, Result}; use crate::io::object_store::{ - MirroringObjectStoreWrapper, install_atomic_aws_provider, + MirroringObjectStoreWrapper, install_atomic_aws_provider, is_aws_credential_option, object_store_params_from_storage_options, set_storage_options, }; use crate::table::NativeTable; @@ -739,8 +739,14 @@ impl ListingDatabase { /// Inherit storage options from the connection into the target map fn inherit_storage_options(&self, target: &mut HashMap) { + // Credential precedence applies to the whole family, not individual members. Once an + // operation supplies any member, the lower-precedence connection family must not fill in + // its missing token (or key/secret); the provider boundary validates completeness later. + let operation_has_aws_credentials = target.keys().any(|key| is_aws_credential_option(key)); for (key, value) in self.storage_options.iter() { - if !target.contains_key(key) { + if !target.contains_key(key) + && !(operation_has_aws_credentials && is_aws_credential_option(key)) + { target.insert(key.clone(), value.clone()); } } @@ -1371,6 +1377,71 @@ mod tests { ); } + #[cfg(feature = "aws")] + #[tokio::test] + async fn operation_credential_family_does_not_inherit_connection_token() { + let tempdir = tempdir().unwrap(); + let request = ConnectRequest { + uri: tempdir.path().to_string_lossy().into_owned(), + #[cfg(feature = "remote")] + client_config: Default::default(), + options: HashMap::from([ + ( + "aws_access_key_id".to_string(), + "connection-key".to_string(), + ), + ( + "aws_secret_access_key".to_string(), + "connection-secret".to_string(), + ), + ( + "aws_session_token".to_string(), + "connection-token".to_string(), + ), + ]), + namespace_client_properties: Default::default(), + manifest_enabled: false, + read_consistency_interval: None, + session: None, + }; + let db = ListingDatabase::connect_with_options(&request) + .await + .unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); + let create_request = CreateTableRequest { + name: "operation_credentials".to_string(), + namespace_path: vec![], + data: Box::new(RecordBatch::new_empty(schema)) as Box, + mode: CreateTableMode::Create, + write_options: WriteOptions { + lance_write_params: Some(lance::dataset::WriteParams { + store_params: Some(operation_aws_store_params()), + ..Default::default() + }), + }, + location: None, + namespace_client: None, + }; + + let write_params = db.prepare_write_params(&create_request, None, None, None); + let options = write_params + .store_params + .unwrap() + .storage_options() + .unwrap() + .clone(); + + assert_eq!(options.get("aws_access_key_id").unwrap(), "operation-key"); + assert_eq!( + options.get("aws_secret_access_key").unwrap(), + "operation-secret" + ); + assert!( + !options.contains_key("aws_session_token"), + "an explicit operation family must not inherit a lower-precedence token" + ); + } + #[tokio::test] async fn test_listing_database_root_ops_do_not_create_manifest() { let tempdir = tempdir().unwrap(); diff --git a/rust/lancedb/src/io/object_store.rs b/rust/lancedb/src/io/object_store.rs index 88dfb1c63..48f8759e8 100644 --- a/rust/lancedb/src/io/object_store.rs +++ b/rust/lancedb/src/io/object_store.rs @@ -20,7 +20,7 @@ use lance::io::{ObjectStoreParams, WrappingObjectStore}; #[cfg(feature = "aws")] use lance_io::object_store::{ ObjectStore as LanceObjectStore, ObjectStoreProvider, ObjectStoreRegistry, StorageOptions, - providers::aws::build_aws_credential, + providers::aws::{AwsStoreProvider, build_aws_credential}, throttle::{AimdThrottleConfig, AimdThrottledStore}, }; use lance_io::object_store::{StorageOptionsAccessor, StorageOptionsProvider}; @@ -98,6 +98,17 @@ fn is_aws_credential_key(key: &AmazonS3ConfigKey) -> bool { ) } +#[cfg(feature = "aws")] +pub(crate) fn is_aws_credential_option(key: &str) -> bool { + AmazonS3ConfigKey::from_str(&key.to_ascii_lowercase()) + .is_ok_and(|key| is_aws_credential_key(&key)) +} + +#[cfg(not(feature = "aws"))] +pub(crate) fn is_aws_credential_option(_key: &str) -> bool { + false +} + #[cfg(feature = "aws")] fn canonical_noncredential_options( storage_options: &HashMap, @@ -410,9 +421,27 @@ struct AtomicAwsStoreProvider { } #[cfg(feature = "aws")] -#[async_trait] -impl ObjectStoreProvider for AtomicAwsStoreProvider { - async fn new_store( +impl AtomicAwsStoreProvider { + const CACHE_GENERATION: &'static str = "lancedb-atomic-aws-v1"; + + /// Lance does not expose provider downcasting. Its built-in provider is a unit struct with a + /// stable derived Debug representation, which is the only available way to distinguish the + /// one provider whose OpenDAL implementation needs adaptation from an unknown custom one. + fn is_builtin_aws_provider(&self) -> bool { + format!("{:?}", self.inner.as_ref()) == format!("{:?}", AwsStoreProvider) + } + + fn generated_prefix( + &self, + url: &url::Url, + storage_options: Option<&HashMap>, + ) -> lance_core::Result { + self.inner + .calculate_object_store_prefix(url, storage_options) + .map(|prefix| format!("{prefix}${}", Self::CACHE_GENERATION)) + } + + async fn new_store_inner( &self, base_path: url::Url, params: &ObjectStoreParams, @@ -423,6 +452,13 @@ impl ObjectStoreProvider for AtomicAwsStoreProvider { .is_some_and(|value| value == "true"); if use_opendal { + // A registered custom provider owns its store construction contract. In particular, + // its returned store may add encryption, authorization, or wrapping behavior that + // cannot be reconstructed from ObjectStoreParams. Only adapt Lance's known built-in + // provider, whose OpenDAL path ignores both supported credential-provider fields. + if !self.is_builtin_aws_provider() { + return self.inner.new_store(base_path, params).await; + } if storage_options .get("aws_provider_scheme") .is_some_and(|scheme| !scheme.is_empty()) @@ -532,6 +568,23 @@ impl ObjectStoreProvider for AtomicAwsStoreProvider { atomic_params.aws_credentials = Some(credential_provider); self.inner.new_store(base_path, &atomic_params).await } +} + +#[cfg(feature = "aws")] +#[async_trait] +impl ObjectStoreProvider for AtomicAwsStoreProvider { + async fn new_store( + &self, + base_path: url::Url, + params: &ObjectStoreParams, + ) -> lance_core::Result { + let store_prefix = self.generated_prefix(&base_path, params.storage_options())?; + let mut store = self.new_store_inner(base_path, params).await?; + // The registry cache and the returned store must use the same identity. Lance compares + // these values when resolving external blob bases. + store.store_prefix = store_prefix; + Ok(store) + } fn extract_path(&self, url: &url::Url) -> lance_core::Result { self.inner.extract_path(url) @@ -542,9 +595,7 @@ impl ObjectStoreProvider for AtomicAwsStoreProvider { url: &url::Url, storage_options: Option<&HashMap>, ) -> lance_core::Result { - self.inner - .calculate_object_store_prefix(url, storage_options) - .map(|prefix| format!("{prefix}$lancedb-atomic-aws-v1")) + self.generated_prefix(url, storage_options) } } @@ -869,6 +920,24 @@ mod credential_tests { } } + #[derive(Debug)] + struct CustomStoreProvider { + marker: Arc, + } + + #[async_trait] + impl ObjectStoreProvider for CustomStoreProvider { + async fn new_store( + &self, + base_path: url::Url, + params: &ObjectStoreParams, + ) -> lance_core::Result { + let mut store = AwsStoreProvider.new_store(base_path, params).await?; + store.inner = self.marker.clone(); + Ok(store) + } + } + #[derive(Debug, PartialEq, Eq)] struct ObservedCredential { key_id: String, @@ -1266,6 +1335,48 @@ mod credential_tests { ); } + #[tokio::test] + async fn opendal_preserves_custom_provider_store_behavior() { + let marker: Arc = Arc::new(object_store::memory::InMemory::new()); + let provider = AtomicAwsStoreProvider { + inner: Arc::new(CustomStoreProvider { + marker: marker.clone(), + }), + }; + let mut options = local_s3_options(); + options.insert("use_opendal".to_string(), "true".to_string()); + + let store = provider + .new_store( + url::Url::parse("s3://bucket/table").unwrap(), + &object_store_params_from_storage_options(options), + ) + .await + .unwrap(); + + assert!( + Arc::ptr_eq(&store.inner, &marker), + "the wrapper must not discard custom provider store behavior" + ); + } + + #[tokio::test] + async fn wrapper_store_prefix_matches_registry_identity() { + let registry = Arc::new(ObjectStoreRegistry::default()); + let session = lance::session::Session::new(16, 16, registry.clone()); + install_atomic_aws_provider(&session); + let uri = "s3://bucket/table"; + let url = url::Url::parse(uri).unwrap(); + let params = object_store_params_from_storage_options(local_s3_options()); + + let store = registry.get_store(url, ¶ms).await.unwrap(); + let registry_prefix = registry + .calculate_object_store_prefix(uri, params.storage_options()) + .unwrap(); + + assert_eq!(store.store_prefix, registry_prefix); + } + #[tokio::test] async fn identical_explicit_options_reuse_the_session_store() { let registry = Arc::new(ObjectStoreRegistry::default()); diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index 7b15b3c88..5f1b36962 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -2256,7 +2256,13 @@ impl NativeTable { pushdown_operations: HashSet, managed_versioning: Option, ) -> Result { - let params = params.unwrap_or_default(); + let mut params = params.unwrap_or_default(); + let effective_session = params + .session + .clone() + .unwrap_or_else(|| Arc::new(lance::session::Session::default())); + crate::io::object_store::install_atomic_aws_provider(&effective_session); + params.session(effective_session); // patch the params if we have a write store wrapper let params = match write_store_wrapper.clone() { Some(wrapper) => params.patch_with_store_wrapper(wrapper)?, @@ -2530,9 +2536,15 @@ impl NativeTable { pushdown_operations: HashSet, ) -> Result { // Default params uses format v1. - let params = params.unwrap_or(WriteParams { + let mut params = params.unwrap_or(WriteParams { ..Default::default() }); + let effective_session = params + .session + .clone() + .unwrap_or_else(|| Arc::new(lance::session::Session::default())); + crate::io::object_store::install_atomic_aws_provider(&effective_session); + params.session = Some(effective_session); // patch the params if we have a write store wrapper let params = match write_store_wrapper.clone() { Some(wrapper) => params.patch_with_store_wrapper(wrapper)?, @@ -3581,6 +3593,10 @@ mod tests { use lance::Dataset; use lance::io::{ObjectStoreParams, WrappingObjectStore}; use lance_core::datatypes::LANCE_UNENFORCED_PRIMARY_KEY_POSITION; + #[cfg(feature = "aws")] + use lance_io::object_store::{ + ObjectStore as LanceObjectStore, ObjectStoreProvider, ObjectStoreRegistry, + }; use tempfile::tempdir; use super::*; @@ -3590,6 +3606,53 @@ mod tests { use crate::query::{ExecutableQuery, QueryBase}; use crate::test_utils::connection::new_test_connection; + #[cfg(feature = "aws")] + #[derive(Debug)] + struct RecordingS3Provider { + saw_atomic_credentials: Arc, + } + + #[cfg(feature = "aws")] + #[async_trait] + impl ObjectStoreProvider for RecordingS3Provider { + async fn new_store( + &self, + _base_path: url::Url, + params: &ObjectStoreParams, + ) -> lance_core::Result { + self.saw_atomic_credentials + .store(params.aws_credentials.is_some(), Ordering::SeqCst); + Err(lance_core::Error::invalid_input("recorded test request")) + } + } + + #[cfg(feature = "aws")] + fn recording_s3_session() -> (Arc, Arc) { + let saw_atomic_credentials = Arc::new(AtomicBool::new(false)); + let registry = Arc::new(ObjectStoreRegistry::default()); + registry.insert( + "s3", + Arc::new(RecordingS3Provider { + saw_atomic_credentials: saw_atomic_credentials.clone(), + }), + ); + ( + Arc::new(lance::session::Session::new(16, 16, registry)), + saw_atomic_credentials, + ) + } + + #[cfg(feature = "aws")] + fn explicit_s3_store_params() -> ObjectStoreParams { + crate::io::object_store::object_store_params_from_storage_options(HashMap::from([ + ("aws_access_key_id".to_string(), "explicit-key".to_string()), + ( + "aws_secret_access_key".to_string(), + "explicit-secret".to_string(), + ), + ])) + } + #[test] fn test_tokenize_uses_explicit_simple_tokenizer() { let params = @@ -3651,6 +3714,70 @@ mod tests { ); } + #[cfg(feature = "aws")] + #[tokio::test] + async fn direct_native_open_installs_the_atomic_provider() { + let (session, saw_atomic_credentials) = recording_s3_session(); + let params = ReadParams { + session: Some(session), + store_options: Some(explicit_s3_store_params()), + ..Default::default() + }; + + let error = NativeTable::open_with_params( + "s3://bucket/table", + "table", + vec![], + None, + Some(params), + None, + None, + HashSet::new(), + None, + ) + .await + .unwrap_err(); + + assert!(error.to_string().contains("recorded test request")); + assert!( + saw_atomic_credentials.load(Ordering::SeqCst), + "the public direct open path must install the wrapper on its session" + ); + } + + #[cfg(feature = "aws")] + #[tokio::test] + async fn direct_native_create_installs_the_atomic_provider() { + let (session, saw_atomic_credentials) = recording_s3_session(); + let params = WriteParams { + session: Some(session), + store_params: Some(explicit_s3_store_params()), + ..Default::default() + }; + let batch = make_test_batches(); + let reader = RecordBatchIterator::new(vec![Ok(batch.clone())], batch.schema()); + + let error = NativeTable::create( + "s3://bucket/table", + "table", + vec![], + reader, + None, + Some(params), + None, + None, + HashSet::new(), + ) + .await + .unwrap_err(); + + assert!(error.to_string().contains("recorded test request")); + assert!( + saw_atomic_credentials.load(Ordering::SeqCst), + "the public direct create path must install the wrapper on its session" + ); + } + /// Write a table and then break it, leaving the `.lance` directory in place. /// /// `remove_all` reproduces an interrupted drop + re-create (the directory is left