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): preserve namespace credential providers

Browse Source
This commit is contained in:
Gatefixer
2026-08-06 19:56:39 +00:00
parent 261e9272f6
commit a4d34d7b5d
2 changed files with 221 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/lancedb/src/database/namespace.rs
+18 -7
View File
@@ -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<Arc<lance::session::Session>>,
namespace_client_pushdown_operations: HashSet<NamespaceClientPushdownOperation>,
) -> 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<NamespaceClientPushdownOperation>,
new_table_config: NewTableConfig,
) -> Result<Self> {
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::<LanceNamespaceDatabase>()
.unwrap();
assert!(result.is_ok());
assert!(database.session.is_some());
}
#[tokio::test]
rust/lancedb/src/io/object_store.rs
+203 -18
View File
@@ -155,6 +155,17 @@ impl ObjectStoreProvider for AtomicAwsStoreProvider {
base_path: url::Url,
params: &ObjectStoreParams,
) -> lance_core::Result<LanceObjectStore> {
// 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<String, String>,
provider: Option<Arc<dyn StorageOptionsProvider>>,
) {
#[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<AtomicUsize>,
}
#[async_trait]
impl StorageOptionsProvider for RotatingOptionsProvider {
async fn fetch_storage_options(
&self,
) -> lance_core::Result<Option<HashMap<String, String>>> {
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<String>,
}
#[derive(Debug)]
struct ResolvingProvider {
resolved_credential: Arc<Mutex<Option<ObservedCredential>>>,
}
#[async_trait]
impl ObjectStoreProvider for ResolvingProvider {
async fn new_store(
&self,
_base_path: url::Url,
params: &ObjectStoreParams,
) -> lance_core::Result<LanceObjectStore> {
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::<HashMap<_, _>>();
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(), &params)
.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(), &params)
.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