alamb commented on code in PR #23325:
URL: https://github.com/apache/datafusion/pull/23325#discussion_r3590794653
##########
datafusion-cli/src/object_storage.rs:
##########
@@ -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,
Review Comment:
Any reason we can't use the same field name (provider)?
##########
datafusion-cli/src/object_storage.rs:
##########
@@ -589,14 +660,333 @@ mod tests {
use crate::cli_context::CliSessionContext;
use super::*;
+ use aws_credential_types::provider::future;
use datafusion::{
datasource::listing::ListingTableUrl,
logical_expr::{DdlStatement, LogicalPlan},
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?;
Review Comment:
What does this do? Mimic a AWS endpoint with a shell script? I am not sure
that is any better than a rust mock, and it won't work on windows , will it?
##########
datafusion-cli/src/object_storage.rs:
##########
@@ -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 {
Review Comment:
can we add some comments expalining this?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
