kosiew commented on code in PR #23325:
URL: https://github.com/apache/datafusion/pull/23325#discussion_r3549509756
##########
datafusion-cli/src/object_storage.rs:
##########
@@ -589,14 +660,342 @@ 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, path::PathBuf, process};
+ #[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 = unique_temp_dir("datafusion-s3-credential-cache")?;
Review Comment:
Small nit: this test creates a temp directory with `unique_temp_dir`, but it
does not remove it afterward. It might be worth using `tempfile::TempDir` or
another cleanup guard so repeated local or CI runs do not leave
`datafusion-s3-credential-cache-*` directories behind.
--
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]
