rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-08-18 12:08:35 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(rust): keep explicit AWS credentials atomic

Browse Source
This commit is contained in:
Gatefixer
2026-08-05 23:06:33 +00:00
parent 7357d63e87
commit bae68d45ca
2 changed files with 135 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/lancedb/src/database/listing.rs
+17 -54
View File
@@ -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<String>) -> 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;
rust/lancedb/src/io/object_store.rs
+118 -3
View File
@@ -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<String, String>,
) -> Option<object_store::aws::AwsCredentialProvider> {
let aws_options = storage_options
.iter()
.filter_map(|(key, value)| {
AmazonS3ConfigKey::from_str(&key.to_ascii_lowercase())
.ok()
.map(|key| (key, value))
})
.collect::<HashMap<_, _>>();
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<String, String>,
provider: Option<Arc<dyn StorageOptionsProvider>>,
) {
#[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<String, String>,
) -> ObjectStoreParams {
let mut params = ObjectStoreParams::default();
set_storage_options(&mut params, storage_options, None);
params
}
#[derive(Debug)]
struct MirroringObjectStore {
primary: Arc<dyn ObjectStore>,
@@ -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::<HashMap<_, _>>();
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 {