This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-23325-95398f07f26f098188472385f808ddcd49160eda in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit d30b7f58c681b9eb37b79d5359ac8ebca59e2ba2 Author: Alex Metelli <[email protected]> AuthorDate: Sat Jul 11 16:38:26 2026 +0800 fix: cache CLI AWS credentials until expiry (#23325) ## Which issue does this PR close? - Closes #23321. ## Rationale for this change The CLI's S3 credential provider wrapped AWS SDK credentials but called `SharedCredentialsProvider::provide_credentials()` for every object-store credential request. That bypasses the AWS SDK request pipeline identity cache and can repeatedly invoke expensive providers such as `credential_process` or SSO-backed providers even when the returned credentials have not expired. ## What changes are included in this PR? - Cache the converted `object_store::aws::AwsCredential` inside `S3CredentialProvider`. - Seed that cache from the initial credential fetch already performed while building the S3 object store. - Reuse cached credentials until they are expired or within a small refresh buffer. - Add regression coverage for seeded credentials, non-expiring credentials, expired and near-expiry credentials, and an end-to-end `credential_process` flow that fails on `main` without this fix. ## Are these changes tested? Yes: - `cargo test -p datafusion-cli s3_object_store_reuses_fetched_credentials_until_expiry -- --nocapture` - `cargo test -p datafusion-cli object_storage` - `cargo test -p datafusion-cli` - `cargo clippy --all-targets --all-features -- -D warnings` I also verified the new public-flow regression test against detached `main`; it fails there because the credential process runs multiple times instead of once. ## Are there any user-facing changes? No API or configuration changes. The CLI should make fewer redundant AWS credential-provider calls for S3 object stores while still refreshing expiring credentials. --- Cargo.lock | 1 + datafusion-cli/Cargo.toml | 1 + datafusion-cli/src/object_storage.rs | 412 ++++++++++++++++++++++++++++++++++- 3 files changed, 403 insertions(+), 11 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 5b43435ec0..bead33e99e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1857,6 +1857,7 @@ dependencies = [ "rstest", "rustyline", "serde_json", + "tempfile", "testcontainers-modules", "tokio", "url", diff --git a/datafusion-cli/Cargo.toml b/datafusion-cli/Cargo.toml index 62eedafe79..be449189cc 100644 --- a/datafusion-cli/Cargo.toml +++ b/datafusion-cli/Cargo.toml @@ -77,6 +77,7 @@ ctor = { workspace = true } insta = { workspace = true } insta-cmd = "0.7.0" rstest = { workspace = true } +tempfile = { workspace = true } testcontainers-modules = { workspace = true, features = ["minio"] } # Makes sure `test_display_pg_json` behaves in a consistent way regardless of # feature unification with dependencies diff --git a/datafusion-cli/src/object_storage.rs b/datafusion-cli/src/object_storage.rs index 4293788e0c..73fd3b47a1 100644 --- a/datafusion-cli/src/object_storage.rs +++ b/datafusion-cli/src/object_storage.rs @@ -22,8 +22,9 @@ pub use stdin::{StdinCarriesCommands, is_stdin_location}; use async_trait::async_trait; use aws_config::BehaviorVersion; -use aws_credential_types::provider::{ - ProvideCredentials, SharedCredentialsProvider, error::CredentialsError, +use aws_credential_types::{ + Credentials as AwsSdkCredentials, + provider::{ProvideCredentials, SharedCredentialsProvider, error::CredentialsError}, }; use datafusion::{ common::{ @@ -48,7 +49,9 @@ use std::{ error::Error, fmt::{Debug, Display}, sync::Arc, + time::{Duration, SystemTime}, }; +use tokio::sync::RwLock; use url::Url; #[cfg(not(test))] @@ -117,7 +120,10 @@ async fn get_s3_object_store_builder_inner( builder = builder.with_region(region); } if let Some(credentials) = credentials { - let credentials = Arc::new(S3CredentialProvider { credentials }); + let credentials = Arc::new(S3CredentialProvider::new( + credentials.provider, + &credentials.initial_credentials, + )); builder = builder.with_credentials(credentials); } else { debug!("No credentials found, defaulting to skip signature "); @@ -170,7 +176,12 @@ async fn get_s3_object_store_builder_inner( /// Credentials from the AWS SDK struct CredentialsFromConfig { region: Option<String>, - credentials: Option<SharedCredentialsProvider>, + credentials: Option<CredentialsFromConfigProvider>, +} + +struct CredentialsFromConfigProvider { + provider: SharedCredentialsProvider, + initial_credentials: AwsSdkCredentials, } impl CredentialsFromConfig { @@ -193,7 +204,10 @@ impl CredentialsFromConfig { // until they are needed. To ensure that the credentials are valid, // we can call `provide_credentials` here. let credentials = match credentials.provide_credentials().await { - Ok(_) => Some(credentials), + Ok(initial_credentials) => Some(CredentialsFromConfigProvider { + provider: credentials, + initial_credentials, + }), Err(CredentialsError::CredentialsNotLoaded(_)) => { debug!("Could not use AWS SDK to get credentials"); None @@ -225,9 +239,51 @@ impl CredentialsFromConfig { } } +const S3_CREDENTIAL_CACHE_EXPIRY_BUFFER: Duration = Duration::from_secs(10); + #[derive(Debug)] struct S3CredentialProvider { credentials: SharedCredentialsProvider, + cache: RwLock<Option<CachedAwsCredential>>, +} + +impl S3CredentialProvider { + fn new( + credentials: SharedCredentialsProvider, + initial_credentials: &AwsSdkCredentials, + ) -> Self { + Self { + credentials, + cache: RwLock::new(Some(CachedAwsCredential::new(initial_credentials))), + } + } +} + +#[derive(Debug)] +struct CachedAwsCredential { + credential: Arc<AwsCredential>, + expires_at: Option<SystemTime>, +} + +impl CachedAwsCredential { + fn new(credentials: &AwsSdkCredentials) -> Self { + Self { + expires_at: credentials.expiry(), + credential: Arc::new(AwsCredential { + key_id: credentials.access_key_id().to_string(), + secret_key: credentials.secret_access_key().to_string(), + token: credentials.session_token().map(ToString::to_string), + }), + } + } + + fn is_valid(&self, now: SystemTime) -> bool { + self.expires_at.is_none_or(|expires_at| { + expires_at + .duration_since(now) + .is_ok_and(|ttl| ttl > S3_CREDENTIAL_CACHE_EXPIRY_BUFFER) + }) + } } #[async_trait] @@ -235,7 +291,22 @@ impl CredentialProvider for S3CredentialProvider { type Credential = AwsCredential; async fn get_credential(&self) -> object_store::Result<Arc<Self::Credential>> { - let creds = + let now = SystemTime::now(); + if let Some(cached) = self.cache.read().await.as_ref() + && cached.is_valid(now) + { + return Ok(Arc::clone(&cached.credential)); + } + + let mut cache = self.cache.write().await; + let now = SystemTime::now(); + if let Some(cached) = cache.as_ref() + && cached.is_valid(now) + { + return Ok(Arc::clone(&cached.credential)); + } + + let credentials = self.credentials .provide_credentials() .await @@ -243,11 +314,11 @@ impl CredentialProvider for S3CredentialProvider { store: "S3", source: Box::new(e), })?; - Ok(Arc::new(AwsCredential { - key_id: creds.access_key_id().to_string(), - secret_key: creds.secret_access_key().to_string(), - token: creds.session_token().map(ToString::to_string), - })) + let cached = CachedAwsCredential::new(&credentials); + let credential = Arc::clone(&cached.credential); + *cache = Some(cached); + + Ok(credential) } } @@ -589,6 +660,7 @@ mod tests { use crate::cli_context::CliSessionContext; use super::*; + use aws_credential_types::provider::future; use datafusion::{ datasource::listing::ListingTableUrl, @@ -596,7 +668,325 @@ mod tests { prelude::SessionContext, }; + #[cfg(unix)] + use object_store::{ObjectStoreExt, path::Path as ObjectStorePath}; use object_store::{aws::AmazonS3ConfigKey, gcp::GoogleConfigKey}; + use std::{ + collections::VecDeque, + sync::{ + Mutex, + atomic::{AtomicUsize, Ordering}, + }, + }; + #[cfg(unix)] + use std::{ffi::OsString, fs}; + #[cfg(unix)] + use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::TcpListener, + }; + + #[derive(Debug)] + struct CountingCredentialsProvider { + credentials: Mutex<VecDeque<AwsSdkCredentials>>, + calls: AtomicUsize, + } + + impl CountingCredentialsProvider { + fn new(credentials: Vec<AwsSdkCredentials>) -> Self { + assert!(!credentials.is_empty()); + Self { + credentials: Mutex::new(credentials.into()), + calls: AtomicUsize::new(0), + } + } + + fn calls(&self) -> usize { + self.calls.load(Ordering::SeqCst) + } + } + + impl ProvideCredentials for CountingCredentialsProvider { + fn provide_credentials<'a>(&'a self) -> future::ProvideCredentials<'a> + where + Self: 'a, + { + self.calls.fetch_add(1, Ordering::SeqCst); + + let credentials = { + let mut credentials = self.credentials.lock().unwrap(); + if credentials.len() > 1 { + credentials.pop_front().unwrap() + } else { + credentials.front().unwrap().clone() + } + }; + + future::ProvideCredentials::ready(Ok(credentials)) + } + } + + fn shared_counting_provider( + credentials: Vec<AwsSdkCredentials>, + ) -> (SharedCredentialsProvider, Arc<CountingCredentialsProvider>) { + let provider = Arc::new(CountingCredentialsProvider::new(credentials)); + let shared_provider = SharedCredentialsProvider::from( + Arc::clone(&provider) as Arc<dyn ProvideCredentials> + ); + (shared_provider, provider) + } + + fn sdk_credentials( + key_id: &'static str, + expires_after: Option<SystemTime>, + ) -> AwsSdkCredentials { + AwsSdkCredentials::new( + key_id, + format!("{key_id}_secret"), + Some(format!("{key_id}_token")), + expires_after, + "test", + ) + } + + #[tokio::test] + async fn s3_credential_provider_uses_seeded_credentials_without_refetching() + -> Result<()> { + let initial = sdk_credentials( + "initial", + Some( + SystemTime::now() + + S3_CREDENTIAL_CACHE_EXPIRY_BUFFER + + Duration::from_secs(60), + ), + ); + let (shared_provider, counting_provider) = + shared_counting_provider(vec![sdk_credentials("refreshed", None)]); + let provider = S3CredentialProvider::new(shared_provider, &initial); + + let first = provider.get_credential().await?; + let second = provider.get_credential().await?; + + assert_eq!(first.key_id, "initial"); + assert_eq!(second.key_id, "initial"); + assert!(Arc::ptr_eq(&first, &second)); + assert_eq!(counting_provider.calls(), 0); + + Ok(()) + } + + #[tokio::test] + async fn s3_credential_provider_reuses_non_expiring_credentials() -> Result<()> { + let initial = sdk_credentials("static", None); + let (shared_provider, counting_provider) = + shared_counting_provider(vec![sdk_credentials("refreshed", None)]); + let provider = S3CredentialProvider::new(shared_provider, &initial); + + let first = provider.get_credential().await?; + let second = provider.get_credential().await?; + + assert_eq!(first.key_id, "static"); + assert_eq!(second.key_id, "static"); + assert!(Arc::ptr_eq(&first, &second)); + assert_eq!(counting_provider.calls(), 0); + + Ok(()) + } + + #[tokio::test] + async fn s3_credential_provider_refreshes_near_expiry_credentials() -> Result<()> { + let initial = sdk_credentials( + "initial", + Some(SystemTime::now() + S3_CREDENTIAL_CACHE_EXPIRY_BUFFER), + ); + let refreshed = sdk_credentials( + "refreshed", + Some( + SystemTime::now() + + S3_CREDENTIAL_CACHE_EXPIRY_BUFFER + + Duration::from_secs(60), + ), + ); + let (shared_provider, counting_provider) = + shared_counting_provider(vec![refreshed]); + let provider = S3CredentialProvider::new(shared_provider, &initial); + + let first = provider.get_credential().await?; + let second = provider.get_credential().await?; + + assert_eq!(first.key_id, "refreshed"); + assert_eq!(second.key_id, "refreshed"); + assert!(Arc::ptr_eq(&first, &second)); + assert_eq!(counting_provider.calls(), 1); + + Ok(()) + } + + #[tokio::test] + async fn s3_credential_provider_refreshes_expired_credentials() -> Result<()> { + let initial = + sdk_credentials("initial", Some(SystemTime::now() - Duration::from_secs(1))); + let refreshed = sdk_credentials( + "refreshed", + Some( + SystemTime::now() + + S3_CREDENTIAL_CACHE_EXPIRY_BUFFER + + Duration::from_secs(60), + ), + ); + let (shared_provider, counting_provider) = + shared_counting_provider(vec![refreshed]); + let provider = S3CredentialProvider::new(shared_provider, &initial); + + let first = provider.get_credential().await?; + let second = provider.get_credential().await?; + + assert_eq!(first.key_id, "refreshed"); + assert_eq!(second.key_id, "refreshed"); + assert!(Arc::ptr_eq(&first, &second)); + assert_eq!(counting_provider.calls(), 1); + + Ok(()) + } + + #[cfg(unix)] + #[tokio::test] + async fn s3_object_store_reuses_fetched_credentials_until_expiry() -> Result<()> { + use std::os::unix::fs::PermissionsExt; + + let test_dir_guard = tempfile::Builder::new() + .prefix("datafusion-s3-credential-cache") + .tempdir()?; + let test_dir = test_dir_guard.path(); + let count_path = test_dir.join("credential_process_count"); + let process_path = test_dir.join("credential_process.sh"); + let config_path = test_dir.join("config"); + let credentials_path = test_dir.join("credentials"); + + fs::write( + &process_path, + r#"#!/bin/sh +count_file="$1" +if [ -f "$count_file" ]; then + count=$(cat "$count_file") +else + count=0 +fi +count=$((count + 1)) +printf "%s" "$count" > "$count_file" +cat <<'JSON' +{"Version":1,"AccessKeyId":"test_access_key","SecretAccessKey":"test_secret_key","SessionToken":"test_session_token","Expiration":"2099-01-01T00:00:00Z"} +JSON +"#, + )?; + let mut permissions = fs::metadata(&process_path)?.permissions(); + permissions.set_mode(0o700); + fs::set_permissions(&process_path, permissions)?; + + fs::write( + &config_path, + format!( + "[profile datafusion-cache-test]\nregion = us-east-1\ncredential_process = {} {}\n", + process_path.display(), + count_path.display() + ), + )?; + fs::write(&credentials_path, "")?; + + let _env = EnvGuard::set([ + ("AWS_CONFIG_FILE", Some(config_path.into_os_string())), + ( + "AWS_SHARED_CREDENTIALS_FILE", + Some(credentials_path.into_os_string()), + ), + ("AWS_PROFILE", Some(OsString::from("datafusion-cache-test"))), + ("AWS_EC2_METADATA_DISABLED", Some(OsString::from("true"))), + ("AWS_ACCESS_KEY_ID", None), + ("AWS_SECRET_ACCESS_KEY", None), + ("AWS_SESSION_TOKEN", None), + ("AWS_REGION", None), + ]); + + let listener = TcpListener::bind("127.0.0.1:0").await?; + let endpoint = format!("http://{}";, listener.local_addr()?); + let server = async move { + for _ in 0..2 { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = [0; 1024]; + let _ = socket.read(&mut request).await; + socket + .write_all( + b"HTTP/1.1 404 Not Found\r\nContent-Length: 0\r\nConnection: close\r\n\r\n", + ) + .await + .unwrap(); + } + }; + + let url = Url::parse("s3://bucket/path/file.parquet").unwrap(); + let aws_options = AwsOptions { + region: Some("us-east-1".to_string()), + endpoint: Some(endpoint), + allow_http: Some(true), + ..Default::default() + }; + let store = get_s3_object_store_builder(&url, &aws_options, false) + .await? + .build()?; + + let path = ObjectStorePath::from("path/file.parquet"); + let requests = async { + let _ = store.head(&path).await; + let _ = store.head(&path).await; + }; + tokio::join!(server, requests); + + assert_eq!(fs::read_to_string(count_path)?, "1"); + + Ok(()) + } + + #[cfg(unix)] + struct EnvGuard { + saved: Vec<(&'static str, Option<OsString>)>, + } + + #[cfg(unix)] + impl EnvGuard { + fn set<const N: usize>(updates: [(&'static str, Option<OsString>); N]) -> Self { + let saved = updates + .iter() + .map(|(key, _)| (*key, std::env::var_os(key))) + .collect(); + + for (key, value) in updates { + unsafe { + if let Some(value) = value { + std::env::set_var(key, value); + } else { + std::env::remove_var(key); + } + } + } + + Self { saved } + } + } + + #[cfg(unix)] + impl Drop for EnvGuard { + fn drop(&mut self) { + for (key, value) in self.saved.drain(..) { + unsafe { + if let Some(value) = value { + std::env::set_var(key, value); + } else { + std::env::remove_var(key); + } + } + } + } + } #[tokio::test] async fn s3_object_store_builder_default() -> Result<()> { --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
