From bae68d45cab2d2c24a2df4d29a3a3ba2bc1dde3f Mon Sep 17 00:00:00 2001 From: Gatefixer <313497061+lancedb-gatefixer[bot]@users.noreply.github.com> Date: Wed, 5 Aug 2026 23:06:33 +0000 Subject: [PATCH] fix(rust): keep explicit AWS credentials atomic --- rust/lancedb/src/database/listing.rs | 71 ++++------------ rust/lancedb/src/io/object_store.rs | 121 ++++++++++++++++++++++++++- 2 files changed, 135 insertions(+), 57 deletions(-) diff --git a/rust/lancedb/src/database/listing.rs b/rust/lancedb/src/database/listing.rs index 454498d54..455c63101 100644 --- a/rust/lancedb/src/database/listing.rs +++ b/rust/lancedb/src/database/listing.rs @@ -13,7 +13,7 @@ use lance::dataset::{ReadParams, WriteMode, builder::DatasetBuilder}; use lance::io::{ObjectStore, ObjectStoreParams, WrappingObjectStore}; use lance_datafusion::utils::StreamingWriteSource; use lance_encoding::version::LanceFileVersion; -use lance_io::object_store::{StorageOptionsAccessor, StorageOptionsProvider}; +use lance_io::object_store::StorageOptionsProvider; use lance_table::io::commit::commit_handler_from_url; use object_store::local::LocalFileSystem; use snafu::ResultExt; @@ -23,7 +23,9 @@ use crate::connection::ConnectRequest; use crate::database::ReadConsistency; use crate::database::namespace::LanceNamespaceDatabase; use crate::error::{CreateDirSnafu, Error, Result}; -use crate::io::object_store::MirroringObjectStoreWrapper; +use crate::io::object_store::{ + MirroringObjectStoreWrapper, object_store_params_from_storage_options, set_storage_options, +}; use crate::table::NativeTable; use crate::utils::validate_table_name; @@ -399,16 +401,7 @@ impl ListingDatabase { url.set_query(None); let plain_uri = url.to_string(); - let os_params = ObjectStoreParams { - storage_options_accessor: if storage_options.is_empty() { - None - } else { - Some(Arc::new(StorageOptionsAccessor::with_static_options( - storage_options.clone(), - ))) - }, - ..Default::default() - }; + let os_params = object_store_params_from_storage_options(storage_options.clone()); let (object_store, _) = ObjectStore::from_uri_and_params( session.store_registry(), &plain_uri, @@ -551,16 +544,8 @@ impl ListingDatabase { .session .clone() .unwrap_or_else(|| Arc::new(lance::session::Session::default())); - let os_params = ObjectStoreParams { - storage_options_accessor: if options.storage_options.is_empty() { - None - } else { - Some(Arc::new(StorageOptionsAccessor::with_static_options( - options.storage_options.clone(), - ))) - }, - ..Default::default() - }; + let os_params = + object_store_params_from_storage_options(options.storage_options.clone()); let (object_store, base_path) = ObjectStore::from_uri_and_params( session.store_registry(), &storage_base_uri, @@ -719,16 +704,8 @@ impl ListingDatabase { } async fn drop_tables(&self, names: Vec) -> Result<()> { - let object_store_params = ObjectStoreParams { - storage_options_accessor: if self.storage_options.is_empty() { - None - } else { - Some(Arc::new(StorageOptionsAccessor::with_static_options( - self.storage_options.clone(), - ))) - }, - ..Default::default() - }; + let object_store_params = + object_store_params_from_storage_options(self.storage_options.clone()); let mut uri = self.uri.clone(); if let Some(query_string) = &self.query_string { uri.push_str(&format!("?{}", query_string)); @@ -834,12 +811,11 @@ impl ListingDatabase { if !self.storage_options.is_empty() { self.inherit_storage_options(&mut storage_options); } - let accessor = if let Some(ref provider) = self.storage_options_provider { - StorageOptionsAccessor::with_initial_and_provider(storage_options, provider.clone()) - } else { - StorageOptionsAccessor::with_static_options(storage_options) - }; - store_params.storage_options_accessor = Some(Arc::new(accessor)); + set_storage_options( + store_params, + storage_options, + self.storage_options_provider.clone(), + ); } write_params.data_storage_version = storage_version_override @@ -1102,16 +1078,7 @@ impl Database for ListingDatabase { validate_table_name(&request.target_table_name)?; - let storage_params = ObjectStoreParams { - storage_options_accessor: if self.storage_options.is_empty() { - None - } else { - Some(Arc::new(StorageOptionsAccessor::with_static_options( - self.storage_options.clone(), - ))) - }, - ..Default::default() - }; + let storage_params = object_store_params_from_storage_options(self.storage_options.clone()); let read_params = ReadParams { store_options: Some(storage_params.clone()), session: Some(self.session.clone()), @@ -1188,12 +1155,7 @@ impl Database for ListingDatabase { .as_ref() .and_then(|a| a.provider().cloned()); let provider = self.storage_options_provider.clone().or(request_provider); - let accessor = if let Some(provider) = provider { - StorageOptionsAccessor::with_initial_and_provider(storage_options, provider) - } else { - StorageOptionsAccessor::with_static_options(storage_options) - }; - store_params.storage_options_accessor = Some(Arc::new(accessor)); + set_storage_options(store_params, storage_options, provider); } // Some ReadParams are exposed in the OpenTableBuilder, but we also @@ -1297,6 +1259,7 @@ mod tests { use crate::table::WriteOptions; use arrow_array::{Int32Array, RecordBatch, StringArray}; use arrow_schema::{DataType, Field, Schema}; + use lance_io::object_store::StorageOptionsAccessor; use std::path::PathBuf; use tempfile::tempdir; diff --git a/rust/lancedb/src/io/object_store.rs b/rust/lancedb/src/io/object_store.rs index d27357b82..f0389e5e0 100644 --- a/rust/lancedb/src/io/object_store.rs +++ b/rust/lancedb/src/io/object_store.rs @@ -1,23 +1,92 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The LanceDB Authors -//! A mirroring object store that mirror writes to a secondary object store +//! Object store helpers and a store that mirrors writes to a secondary store -use std::{fmt::Formatter, sync::Arc}; +use std::{collections::HashMap, fmt::Formatter, sync::Arc}; use futures::{StreamExt, TryFutureExt, stream::BoxStream}; -use lance::io::WrappingObjectStore; +use lance::io::{ObjectStoreParams, WrappingObjectStore}; +use lance_io::object_store::{StorageOptionsAccessor, StorageOptionsProvider}; use object_store::{ CopyOptions, Error, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta, ObjectStore, ObjectStoreExt, PutMultipartOptions, PutOptions, PutPayload, PutResult, Result, UploadPart, path::Path, }; +#[cfg(feature = "aws")] +use object_store::{ + StaticCredentialProvider, + aws::{AmazonS3ConfigKey, AwsCredential}, +}; +#[cfg(feature = "aws")] +use std::str::FromStr; use async_trait::async_trait; #[cfg(test)] pub mod io_tracking; +#[cfg(feature = "aws")] +fn explicit_aws_credentials( + storage_options: &HashMap, +) -> Option { + let aws_options = storage_options + .iter() + .filter_map(|(key, value)| { + AmazonS3ConfigKey::from_str(&key.to_ascii_lowercase()) + .ok() + .map(|key| (key, value)) + }) + .collect::>(); + + let key_id = aws_options.get(&AmazonS3ConfigKey::AccessKeyId)?; + let secret_key = aws_options.get(&AmazonS3ConfigKey::SecretAccessKey)?; + let token = aws_options + .get(&AmazonS3ConfigKey::Token) + .map(|token| (*token).clone()); + + Some(Arc::new(StaticCredentialProvider::new(AwsCredential { + key_id: (*key_id).clone(), + secret_key: (*secret_key).clone(), + token, + }))) +} + +/// 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. +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() { + params.aws_credentials = explicit_aws_credentials(&storage_options); + } + + params.storage_options_accessor = match (storage_options.is_empty(), provider) { + (true, None) => None, + (true, Some(provider)) => Some(Arc::new(StorageOptionsAccessor::with_provider(provider))), + (false, None) => Some(Arc::new(StorageOptionsAccessor::with_static_options( + storage_options, + ))), + (false, Some(provider)) => Some(Arc::new( + StorageOptionsAccessor::with_initial_and_provider(storage_options, provider), + )), + }; +} + +pub(crate) fn object_store_params_from_storage_options( + storage_options: HashMap, +) -> ObjectStoreParams { + let mut params = ObjectStoreParams::default(); + set_storage_options(&mut params, storage_options, None); + params +} + #[derive(Debug)] struct MirroringObjectStore { primary: Arc, @@ -184,6 +253,52 @@ impl WrappingObjectStore for MirroringObjectStoreWrapper { } } +#[cfg(all(test, feature = "aws"))] +mod credential_tests { + use super::*; + use lance_io::object_store::providers::aws::build_aws_credential; + use std::time::Duration; + + #[tokio::test] + async fn explicit_aws_credentials_do_not_inherit_an_ambient_session_token() { + let storage_options = HashMap::from([ + ("aws_access_key_id".to_string(), "explicit-key".to_string()), + ( + "aws_secret_access_key".to_string(), + "explicit-secret".to_string(), + ), + ]); + let params = object_store_params_from_storage_options(storage_options.clone()); + + // Simulate Lance's environment merge, which adds AWS_SESSION_TOKEN when running in + // Lambda. The explicit provider must remain an atomic two-part credential and take + // precedence over the mixed storage options. + let mut merged_options = storage_options + .into_iter() + .map(|(key, value)| (AmazonS3ConfigKey::from_str(&key).unwrap(), value)) + .collect::>(); + merged_options.insert( + AmazonS3ConfigKey::Token, + "lambda-execution-role-token".to_string(), + ); + + let (provider, _) = build_aws_credential( + Duration::from_secs(60), + params.aws_credentials, + Some(&merged_options), + Some("us-east-1".to_string()), + None, + ) + .await + .unwrap(); + let credential = provider.get_credential().await.unwrap(); + + assert_eq!(credential.key_id, "explicit-key"); + assert_eq!(credential.secret_key, "explicit-secret"); + assert_eq!(credential.token, None); + } +} + // windows pathing can't be simply concatenated #[cfg(all(test, not(windows)))] mod test {