Skip to main content

lance_io/object_store/providers/
aws.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use std::{collections::HashMap, str::FromStr, sync::Arc, time::Duration};
5
6#[cfg(test)]
7use mock_instant::thread_local::{SystemTime, UNIX_EPOCH};
8
9#[cfg(not(test))]
10use std::time::{SystemTime, UNIX_EPOCH};
11
12use object_store::ObjectStore as OSObjectStore;
13use object_store::list::PaginatedListStore;
14use opendal::{Operator, services::S3};
15
16use aws_config::default_provider::credentials::DefaultCredentialsChain;
17use aws_config::ecs::EcsCredentialsProvider;
18use aws_config::provider_config::ProviderConfig;
19use aws_config::web_identity_token::WebIdentityTokenCredentialsProvider;
20use aws_config::{BehaviorVersion, Region, SdkConfig};
21use aws_credential_types::provider::ProvideCredentials;
22use object_store::{
23    ClientOptions, CredentialProvider, Result as ObjectStoreResult, RetryConfig,
24    StaticCredentialProvider,
25    aws::{
26        AmazonS3, AmazonS3Builder, AmazonS3ConfigKey, AwsCredential as ObjectStoreAwsCredential,
27        AwsCredentialProvider,
28    },
29};
30use tokio::sync::RwLock;
31use url::Url;
32
33use crate::object_store::opendal_store::OpendalStore;
34use crate::object_store::{
35    DEFAULT_CLOUD_BLOCK_SIZE, DEFAULT_CLOUD_IO_PARALLELISM, DEFAULT_MAX_IOP_SIZE, ObjectStore,
36    ObjectStoreParams, ObjectStoreProvider, StorageOptions, StorageOptionsAccessor,
37    dynamic_credentials::{NamespaceCredentialsProvider, build_dynamic_credential_provider},
38    throttle::{AimdThrottleConfig, AimdThrottleState, cloud_http_connector, with_throttling},
39};
40use lance_core::error::{Error, Result};
41use lance_core::utils::parse::str_is_truthy;
42
43#[derive(Default, Debug)]
44pub struct AwsStoreProvider;
45
46struct ResolvedS3StorageOptions {
47    options: HashMap<AmazonS3ConfigKey, String>,
48    profile_region: Option<String>,
49}
50
51impl ResolvedS3StorageOptions {
52    fn new(
53        mut options: HashMap<AmazonS3ConfigKey, String>,
54        profile_config: Option<&SdkConfig>,
55    ) -> Self {
56        if effective_s3_endpoint(&options).is_none()
57            && let Some(endpoint) = profile_config.and_then(SdkConfig::endpoint_url)
58        {
59            options.insert(AmazonS3ConfigKey::Endpoint, endpoint.to_string());
60        }
61        let profile_region = profile_config
62            .and_then(SdkConfig::region)
63            .map(|region| region.as_ref().to_string());
64        Self {
65            options,
66            profile_region,
67        }
68    }
69
70    fn effective_endpoint(&self) -> Option<&str> {
71        effective_s3_endpoint(&self.options)
72    }
73
74    fn requires_constant_size_upload_parts(&self) -> bool {
75        self.effective_endpoint()
76            .is_some_and(|endpoint| endpoint.contains("r2.cloudflarestorage.com"))
77    }
78}
79
80impl AwsStoreProvider {
81    async fn build_amazon_s3_store(
82        &self,
83        base_path: &mut Url,
84        params: &ObjectStoreParams,
85        storage_options: &StorageOptions,
86        mut resolved_s3_options: ResolvedS3StorageOptions,
87        is_s3_express: bool,
88        throttle_state: Option<&AimdThrottleState>,
89        // Concrete rather than `dyn`, so the caller keeps the handle a paginated listing
90        // needs: `PaginatedListStore` is a separate trait from `ObjectStore`.
91    ) -> Result<Arc<AmazonS3>> {
92        // Use a low retry count since the AIMD throttle layer handles
93        // throttle recovery with its own retry loop.
94        let retry_config = RetryConfig {
95            backoff: Default::default(),
96            max_retries: storage_options.client_max_retries(),
97            retry_timeout: Duration::from_secs(storage_options.client_retry_timeout()),
98        };
99
100        let region = resolve_s3_region(base_path, &resolved_s3_options).await?;
101
102        // Get accessor from params
103        let accessor = params.get_accessor();
104
105        let provider_scheme = storage_options.aws_provider_scheme()?;
106
107        let (aws_creds, region) = build_aws_credential(
108            params.s3_credentials_refresh_offset,
109            params.aws_credentials.clone(),
110            Some(&resolved_s3_options.options),
111            region,
112            accessor,
113            provider_scheme,
114        )
115        .await?;
116
117        // Set S3Express flag if detected
118        if is_s3_express {
119            resolved_s3_options
120                .options
121                .insert(AmazonS3ConfigKey::S3Express, true.to_string());
122        }
123
124        // Compute the metrics label before rewriting the URL below so it
125        // matches the prefix the registry uses to key this store.
126        let store_prefix =
127            self.calculate_object_store_prefix(base_path, Some(&storage_options.0))?;
128
129        // before creating the OSObjectStore we need to rewrite the url to drop ddb related parts
130        base_path.set_scheme("s3").unwrap();
131        base_path.set_query(None);
132
133        // we can't use parse_url_opts here because we need to manually set the credentials provider
134        let mut builder =
135            AmazonS3Builder::new().with_client_options(storage_options.client_options()?);
136        for (key, value) in resolved_s3_options.options {
137            builder = builder.with_config(key, value);
138        }
139        builder = builder
140            .with_url(base_path.as_ref())
141            .with_credentials(aws_creds)
142            .with_retry(retry_config)
143            .with_region(region);
144
145        builder = builder.with_http_connector(cloud_http_connector(throttle_state, store_prefix));
146
147        Ok(Arc::new(builder.build()?))
148    }
149
150    async fn build_opendal_s3_store(
151        &self,
152        base_path: &Url,
153        storage_options: &StorageOptions,
154    ) -> Result<Arc<dyn OSObjectStore>> {
155        let bucket = base_path
156            .host_str()
157            .ok_or_else(|| Error::invalid_input("S3 URL must contain bucket name"))?
158            .to_string();
159
160        let prefix = base_path.path().trim_start_matches('/').to_string();
161
162        // Start with all storage options as the config map
163        // OpenDAL will handle environment variables through its default credentials chain
164        let mut config_map: HashMap<String, String> = storage_options.0.clone();
165
166        if let Some(provider_scheme) = storage_options.aws_provider_scheme()? {
167            return Result::Err(Error::not_supported(format!(
168                "OpendalStore does not currently support an explicit provider_scheme (currently set to {:?})",
169                provider_scheme
170            )));
171        }
172
173        // Set required OpenDAL configuration
174        config_map.insert("bucket".to_string(), bucket);
175
176        if !prefix.is_empty() {
177            config_map.insert("root".to_string(), "/".to_string());
178        }
179
180        let operator = Operator::from_iter::<S3>(config_map)
181            .map_err(|e| Error::invalid_input(format!("Failed to create S3 operator: {:?}", e)))?;
182
183        Ok(Arc::new(OpendalStore::new(operator)) as Arc<dyn OSObjectStore>)
184    }
185}
186
187#[async_trait::async_trait]
188impl ObjectStoreProvider for AwsStoreProvider {
189    async fn new_store(
190        &self,
191        mut base_path: Url,
192        params: &ObjectStoreParams,
193    ) -> Result<ObjectStore> {
194        let block_size = params.block_size.unwrap_or(DEFAULT_CLOUD_BLOCK_SIZE);
195        let mut storage_options =
196            StorageOptions::new(params.storage_options().cloned().unwrap_or_default());
197        storage_options.with_env_s3();
198        let download_retry_count = storage_options.download_retry_count();
199
200        let use_opendal = storage_options
201            .0
202            .get("use_opendal")
203            .map(|v| str_is_truthy(v))
204            .unwrap_or(false);
205
206        let profile_config = if std::env::var_os("AWS_PROFILE").is_some() {
207            Some(aws_config::load_defaults(BehaviorVersion::latest()).await)
208        } else {
209            None
210        };
211        let resolved_s3_options =
212            ResolvedS3StorageOptions::new(storage_options.as_s3_options(), profile_config.as_ref());
213
214        // Determine S3 Express and constant size upload parts before building the store
215        let is_s3_express = check_s3_express(&base_path, &storage_options);
216
217        let use_constant_size_upload_parts =
218            resolved_s3_options.requires_constant_size_upload_parts();
219
220        let throttle_config = AimdThrottleConfig::from_storage_options(params.storage_options())?;
221        let throttle_state = if throttle_config.is_disabled() {
222            None
223        } else {
224            Some(AimdThrottleState::new(throttle_config)?)
225        };
226
227        let (inner, paginated_lister) = if use_opendal {
228            // Use OpenDAL implementation
229            // Listed in full: no paginated lister covers OpenDAL yet.
230            (
231                self.build_opendal_s3_store(&base_path, &storage_options)
232                    .await?,
233                None,
234            )
235        } else {
236            // Use default Amazon S3 implementation
237            let store = self
238                .build_amazon_s3_store(
239                    &mut base_path,
240                    params,
241                    &storage_options,
242                    resolved_s3_options,
243                    is_s3_express,
244                    throttle_state.as_ref(),
245                )
246                .await?;
247            (
248                store.clone() as Arc<dyn OSObjectStore>,
249                Some(store as Arc<dyn PaginatedListStore>),
250            )
251        };
252        let (inner, paginated_lister) =
253            with_throttling(throttle_state, !use_opendal, inner, paginated_lister);
254
255        Ok(ObjectStore {
256            inner,
257            local_dir_operations: None,
258            scheme: String::from(base_path.scheme()),
259            block_size,
260            max_iop_size: *DEFAULT_MAX_IOP_SIZE,
261            use_constant_size_upload_parts,
262            list_is_lexically_ordered: !is_s3_express,
263            io_parallelism: DEFAULT_CLOUD_IO_PARALLELISM,
264            download_retry_count,
265            io_tracker: Default::default(),
266            store_prefix: self
267                .calculate_object_store_prefix(&base_path, params.storage_options())?,
268            paginated_lister,
269        })
270    }
271}
272
273/// Check if the storage is S3 Express
274fn check_s3_express(url: &Url, storage_options: &StorageOptions) -> bool {
275    storage_options
276        .0
277        .get("s3_express")
278        .map(|v| str_is_truthy(v))
279        .unwrap_or(false)
280        || url.authority().ends_with("--x-s3")
281}
282
283fn effective_s3_endpoint(storage_options: &HashMap<AmazonS3ConfigKey, String>) -> Option<&str> {
284    storage_options
285        .get(&AmazonS3ConfigKey::S3Endpoint)
286        .or_else(|| storage_options.get(&AmazonS3ConfigKey::Endpoint))
287        .map(String::as_str)
288}
289
290/// Figure out the S3 region of the bucket.
291///
292/// This resolves in order of precedence:
293/// 1. The region provided in the storage options
294/// 2. The selected AWS profile's region when a custom endpoint is configured
295/// 3. (If endpoint is not set), the region returned by the S3 API for the bucket
296///
297/// It can return None if no region is provided and the endpoint is set.
298async fn resolve_s3_region(
299    url: &Url,
300    resolved_s3_options: &ResolvedS3StorageOptions,
301) -> Result<Option<String>> {
302    let storage_options = &resolved_s3_options.options;
303    if let Some(region) = storage_options.get(&AmazonS3ConfigKey::Region) {
304        Ok(Some(region.clone()))
305    } else if resolved_s3_options.effective_endpoint().is_none() {
306        // If no endpoint is set, we can assume this is AWS S3 and the region
307        // can be resolved from the bucket.
308        let bucket = url.host_str().ok_or_else(|| {
309            Error::invalid_input(format!("Could not parse bucket from url: {}", url))
310        })?;
311
312        let mut client_options = ClientOptions::default();
313        for (key, value) in storage_options {
314            if let AmazonS3ConfigKey::Client(client_key) = key {
315                client_options = client_options.with_config(*client_key, value.clone());
316            }
317        }
318
319        let bucket_region =
320            object_store::aws::resolve_bucket_region(bucket, &client_options).await?;
321        Ok(Some(bucket_region))
322    } else {
323        Ok(resolved_s3_options.profile_region.clone())
324    }
325}
326
327/// Selects which AWS credential provider to use for a dataset.
328///
329/// When set, overrides automatic credential resolution for everything except an
330/// explicitly-supplied `credentials` provider or `storage_options_accessor`.
331#[derive(Debug, Clone, PartialEq)]
332pub enum AwsProviderScheme {
333    /// Require static access-key credentials (`aws_access_key_id` +
334    /// `aws_secret_access_key`). Returns an error if they are absent.
335    Token,
336    /// Use the ECS/Pod Identity container credential endpoint.
337    /// The endpoint URI is read from the `AWS_CONTAINER_CREDENTIALS_FULL_URI`
338    /// or `AWS_CONTAINER_CREDENTIALS_RELATIVE_URI` environment variables.
339    Ecs,
340    /// Use IRSA (IAM Roles for Service Accounts) web identity token credentials.
341    /// The token file and role ARN are read from the `AWS_WEB_IDENTITY_TOKEN_FILE`
342    /// and `AWS_ROLE_ARN` environment variables.
343    Irsa,
344}
345
346/// Build AWS credentials
347///
348/// This resolves credentials from the following sources in order:
349/// 1. An explicit `credentials` provider
350/// 2. An explicit `storage_options_accessor` with a provider
351/// 3. If `provider_scheme` is set:
352///    - [`AwsProviderScheme::Token`]: static access-key credentials (error if absent)
353///    - [`AwsProviderScheme::Ecs`]: ECS container credential provider
354///    - [`AwsProviderScheme::Irsa`]: web identity token (IRSA) provider
355/// 4. Static access-key credentials from `storage_options`, if present
356/// 5. The default AWS credential provider chain
357///
358/// # Storage Options Accessor
359///
360/// When `storage_options_accessor` is provided and has a dynamic provider,
361/// credentials are fetched and cached by the accessor with automatic refresh
362/// before expiration.
363///
364/// `credentials_refresh_offset` is the amount of time before expiry to refresh credentials.
365pub async fn build_aws_credential(
366    credentials_refresh_offset: Duration,
367    credentials: Option<AwsCredentialProvider>,
368    storage_options: Option<&HashMap<AmazonS3ConfigKey, String>>,
369    region: Option<String>,
370    storage_options_accessor: Option<Arc<StorageOptionsAccessor>>,
371    provider_scheme: Option<AwsProviderScheme>,
372) -> Result<(AwsCredentialProvider, String)> {
373    use aws_config::meta::region::RegionProviderChain;
374    const DEFAULT_REGION: &str = "us-west-2";
375
376    let region = if let Some(region) = region {
377        region
378    } else {
379        RegionProviderChain::default_provider()
380            .or_else(DEFAULT_REGION)
381            .region()
382            .await
383            .map(|r| r.as_ref().to_string())
384            .unwrap_or(DEFAULT_REGION.to_string())
385    };
386
387    // If the user supplied their own credential provider that takes top priority
388    if let Some(creds) = credentials {
389        return Ok((creds, region));
390    }
391
392    // Otherwise, if the user provided a storage_options_accessor, try and use that
393    if let Some(dynamic_creds) = build_dynamic_credential_provider::<ObjectStoreAwsCredential>(
394        storage_options_accessor.clone(),
395    )
396    .await?
397    {
398        return Ok((dynamic_creds, region));
399    }
400
401    // If the user provided a storage_options_accessor, then it must not have matched AWS.
402    // Log a message and ignore it.
403    if storage_options_accessor
404        .as_ref()
405        .is_some_and(|a| a.has_provider())
406    {
407        log::debug!(
408            "Storage options from provider do not contain explicit AWS credentials, \
409             falling back to default AWS credentials chain."
410        );
411    }
412
413    // If the caller specified an explicit provider scheme, use only that provider.
414    if let Some(scheme) = provider_scheme {
415        return match scheme {
416            AwsProviderScheme::Token => {
417                let creds = storage_options
418                    .and_then(extract_static_s3_credentials)
419                    .ok_or_else(|| {
420                        Error::invalid_input(
421                            "aws_provider_scheme=token requires aws_access_key_id \
422                             and aws_secret_access_key to be set",
423                        )
424                    })?;
425                Ok((Arc::new(creds), region))
426            }
427            AwsProviderScheme::Ecs => {
428                let provider = EcsCredentialsProvider::builder().build();
429                Ok((
430                    Arc::new(AwsCredentialAdapter::new(
431                        Arc::new(provider),
432                        credentials_refresh_offset,
433                    )),
434                    region,
435                ))
436            }
437            AwsProviderScheme::Irsa => {
438                let conf = ProviderConfig::default().with_region(Some(Region::new(region.clone())));
439                let provider = WebIdentityTokenCredentialsProvider::builder()
440                    .configure(&conf)
441                    .build();
442                Ok((
443                    Arc::new(AwsCredentialAdapter::new(
444                        Arc::new(provider),
445                        credentials_refresh_offset,
446                    )),
447                    region,
448                ))
449            }
450        };
451    }
452
453    if let Some(opts) = storage_options {
454        // Check for static credentials (access key & secret)
455        if let Some(creds) = extract_static_s3_credentials(opts) {
456            return Ok((Arc::new(creds), region));
457        }
458    }
459
460    let credentials_provider = DefaultCredentialsChain::builder().build().await;
461    Ok((
462        Arc::new(AwsCredentialAdapter::new(
463            Arc::new(credentials_provider),
464            credentials_refresh_offset,
465        )),
466        region,
467    ))
468}
469
470fn extract_static_s3_credentials(
471    options: &HashMap<AmazonS3ConfigKey, String>,
472) -> Option<StaticCredentialProvider<ObjectStoreAwsCredential>> {
473    let key_id = options.get(&AmazonS3ConfigKey::AccessKeyId).cloned();
474    let secret_key = options.get(&AmazonS3ConfigKey::SecretAccessKey).cloned();
475    let token = options.get(&AmazonS3ConfigKey::Token).cloned();
476    match (key_id, secret_key, token) {
477        (Some(key_id), Some(secret_key), token) => {
478            Some(StaticCredentialProvider::new(ObjectStoreAwsCredential {
479                key_id,
480                secret_key,
481                token,
482            }))
483        }
484        _ => None,
485    }
486}
487
488/// Adapt an AWS SDK cred into object_store credentials
489#[derive(Debug)]
490pub struct AwsCredentialAdapter {
491    pub inner: Arc<dyn ProvideCredentials>,
492
493    // RefCell can't be shared across threads, so we use HashMap
494    cache: Arc<RwLock<HashMap<String, Arc<aws_credential_types::Credentials>>>>,
495
496    // The amount of time before expiry to refresh credentials
497    credentials_refresh_offset: Duration,
498}
499
500impl AwsCredentialAdapter {
501    pub fn new(
502        provider: Arc<dyn ProvideCredentials>,
503        credentials_refresh_offset: Duration,
504    ) -> Self {
505        Self {
506            inner: provider,
507            cache: Arc::new(RwLock::new(HashMap::new())),
508            credentials_refresh_offset,
509        }
510    }
511}
512
513const AWS_CREDS_CACHE_KEY: &str = "aws_credentials";
514
515/// Convert std::time::SystemTime from AWS SDK to our mockable SystemTime
516fn to_system_time(time: std::time::SystemTime) -> SystemTime {
517    let duration_since_epoch = time
518        .duration_since(std::time::UNIX_EPOCH)
519        .expect("time should be after UNIX_EPOCH");
520    UNIX_EPOCH + duration_since_epoch
521}
522
523#[async_trait::async_trait]
524impl CredentialProvider for AwsCredentialAdapter {
525    type Credential = ObjectStoreAwsCredential;
526
527    async fn get_credential(&self) -> ObjectStoreResult<Arc<Self::Credential>> {
528        let cached_creds = {
529            let cache_value = self.cache.read().await.get(AWS_CREDS_CACHE_KEY).cloned();
530            let expired = cache_value
531                .clone()
532                .map(|cred| {
533                    cred.expiry()
534                        .map(|exp| {
535                            to_system_time(exp)
536                                .checked_sub(self.credentials_refresh_offset)
537                                .expect("this time should always be valid")
538                                < SystemTime::now()
539                        })
540                        // no expiry is never expire
541                        .unwrap_or(false)
542                })
543                .unwrap_or(true); // no cred is the same as expired;
544            if expired { None } else { cache_value.clone() }
545        };
546
547        if let Some(creds) = cached_creds {
548            Ok(Arc::new(Self::Credential {
549                key_id: creds.access_key_id().to_string(),
550                secret_key: creds.secret_access_key().to_string(),
551                token: creds.session_token().map(|s| s.to_string()),
552            }))
553        } else {
554            let refreshed_creds = Arc::new(
555                self.inner
556                    .provide_credentials()
557                    .await
558                    .map_err(|e| Error::io(format!("Failed to get AWS credentials: {:?}", e)))?,
559            );
560
561            self.cache
562                .write()
563                .await
564                .insert(AWS_CREDS_CACHE_KEY.to_string(), refreshed_creds.clone());
565
566            Ok(Arc::new(Self::Credential {
567                key_id: refreshed_creds.access_key_id().to_string(),
568                secret_key: refreshed_creds.secret_access_key().to_string(),
569                token: refreshed_creds.session_token().map(|s| s.to_string()),
570            }))
571        }
572    }
573}
574
575impl StorageOptions {
576    /// Add values from the environment to storage options.
577    ///
578    /// Only adds keys that are not already present, so explicitly-set options
579    /// (including empty-string sentinels) always take precedence over env vars.
580    pub fn with_env_s3(&mut self) {
581        for (os_key, os_value) in std::env::vars_os() {
582            if let (Some(key), Some(value)) = (os_key.to_str(), os_value.to_str())
583                && let Ok(config_key) = AmazonS3ConfigKey::from_str(&key.to_ascii_lowercase())
584                && !self.0.contains_key(config_key.as_ref())
585            {
586                self.0
587                    .insert(config_key.as_ref().to_string(), value.to_string());
588            }
589        }
590    }
591
592    /// Subset of options relevant for s3 storage
593    pub fn as_s3_options(&self) -> HashMap<AmazonS3ConfigKey, String> {
594        self.0
595            .iter()
596            .filter_map(|(key, value)| {
597                let s3_key = AmazonS3ConfigKey::from_str(&key.to_ascii_lowercase()).ok()?;
598                Some((s3_key, value.clone()))
599            })
600            .collect()
601    }
602
603    /// Parse the `aws_provider_scheme` storage option, if set.
604    pub fn aws_provider_scheme(&self) -> Result<Option<AwsProviderScheme>> {
605        match self.0.get("aws_provider_scheme").map(|s| s.as_str()) {
606            None | Some("") => Ok(None),
607            Some("token") => Ok(Some(AwsProviderScheme::Token)),
608            Some("ecs") => Ok(Some(AwsProviderScheme::Ecs)),
609            Some("irsa") => Ok(Some(AwsProviderScheme::Irsa)),
610            Some(other) => Err(Error::invalid_input(format!(
611                "Invalid aws_provider_scheme '{}'. Valid values are: token, ecs, irsa",
612                other
613            ))),
614        }
615    }
616}
617
618impl ObjectStoreParams {
619    /// Create a new instance of [`ObjectStoreParams`] based on the AWS credentials.
620    pub fn with_aws_credentials(
621        aws_credentials: Option<AwsCredentialProvider>,
622        region: Option<String>,
623    ) -> Self {
624        let storage_options_accessor = region.map(|region| {
625            let opts: HashMap<String, String> =
626                [("region".into(), region)].iter().cloned().collect();
627            Arc::new(StorageOptionsAccessor::with_static_options(opts))
628        });
629        Self {
630            aws_credentials,
631            storage_options_accessor,
632            ..Default::default()
633        }
634    }
635}
636
637pub type DynamicStorageOptionsCredentialProvider =
638    NamespaceCredentialsProvider<ObjectStoreAwsCredential>;
639
640#[cfg(test)]
641mod tests {
642    use crate::object_store::ObjectStoreRegistry;
643    use crate::object_store::StorageOptionsProvider;
644    #[allow(deprecated)]
645    use aws_config::profile::profile_file::{ProfileFileKind, ProfileFiles};
646    use aws_credential_types::provider::error::CredentialsError;
647    use mock_instant::thread_local::MockClock;
648    use object_store::path::Path;
649    use std::sync::atomic::{AtomicBool, Ordering};
650
651    use super::*;
652
653    #[derive(Debug, Default)]
654    struct MockAwsCredentialsProvider {
655        called: AtomicBool,
656    }
657
658    #[async_trait::async_trait]
659    impl CredentialProvider for MockAwsCredentialsProvider {
660        type Credential = ObjectStoreAwsCredential;
661
662        async fn get_credential(&self) -> ObjectStoreResult<Arc<Self::Credential>> {
663            self.called.store(true, Ordering::Relaxed);
664            Ok(Arc::new(Self::Credential {
665                key_id: "".to_string(),
666                secret_key: "".to_string(),
667                token: None,
668            }))
669        }
670    }
671
672    #[allow(deprecated)]
673    async fn load_test_profile(config: &str) -> SdkConfig {
674        let profile_files = ProfileFiles::builder()
675            .with_contents(ProfileFileKind::Config, config)
676            .build();
677        aws_config::defaults(BehaviorVersion::latest())
678            .profile_name("selected")
679            .profile_files(profile_files)
680            .load()
681            .await
682    }
683
684    #[derive(Debug)]
685    struct FailingAwsCredentialsProvider;
686
687    impl ProvideCredentials for FailingAwsCredentialsProvider {
688        fn provide_credentials<'a>(
689            &'a self,
690        ) -> aws_credential_types::provider::future::ProvideCredentials<'a>
691        where
692            Self: 'a,
693        {
694            aws_credential_types::provider::future::ProvideCredentials::new(async {
695                Err(CredentialsError::provider_error(Box::new(
696                    std::io::Error::other("Glue credential endpoint unavailable"),
697                )))
698            })
699        }
700    }
701
702    #[tokio::test]
703    async fn test_aws_credential_failure_is_io_error() {
704        let provider = AwsCredentialAdapter::new(
705            Arc::new(FailingAwsCredentialsProvider),
706            Duration::from_secs(60),
707        );
708
709        let error = provider.get_credential().await.unwrap_err();
710        let object_store::Error::Generic { source, .. } = &error else {
711            panic!("expected a generic object store error, got {error}");
712        };
713        assert!(matches!(
714            source.downcast_ref::<Error>(),
715            Some(Error::IO { .. })
716        ));
717
718        let message = error.to_string();
719        assert!(message.contains("Failed to get AWS credentials"));
720        assert!(message.contains("Glue credential endpoint unavailable"));
721        assert!(!message.contains("Encountered internal error"));
722    }
723
724    #[tokio::test]
725    async fn test_injected_aws_creds_option_is_used() {
726        let mock_provider = Arc::new(MockAwsCredentialsProvider::default());
727        let registry = Arc::new(ObjectStoreRegistry::default());
728
729        let params = ObjectStoreParams {
730            aws_credentials: Some(mock_provider.clone() as AwsCredentialProvider),
731            ..ObjectStoreParams::default()
732        };
733
734        // Not called yet
735        assert!(!mock_provider.called.load(Ordering::Relaxed));
736
737        let (store, _) = ObjectStore::from_uri_and_params(registry, "s3://not-a-bucket", &params)
738            .await
739            .unwrap();
740
741        // fails, but we don't care
742        let _ = store
743            .open(&Path::parse("/").unwrap())
744            .await
745            .unwrap()
746            .get_range(0..1)
747            .await;
748
749        // Not called yet
750        assert!(mock_provider.called.load(Ordering::Relaxed));
751    }
752
753    #[tokio::test]
754    async fn test_resolve_s3_region_from_aws_profile() {
755        let profile_config = load_test_profile(
756            "[profile selected]\n\
757             region = us-west-004\n\
758             endpoint_url = https://s3.us-west-004.backblazeb2.com\n\
759             aws_access_key_id = test-key\n\
760             aws_secret_access_key = test-secret",
761        )
762        .await;
763        let url = Url::parse("s3://test-bucket/path").unwrap();
764
765        let resolved_s3_options =
766            ResolvedS3StorageOptions::new(HashMap::new(), Some(&profile_config));
767        let region = resolve_s3_region(&url, &resolved_s3_options).await.unwrap();
768
769        assert_eq!(region.as_deref(), Some("us-west-004"));
770        assert_eq!(
771            resolved_s3_options
772                .options
773                .get(&AmazonS3ConfigKey::Endpoint),
774            Some(&"https://s3.us-west-004.backblazeb2.com".to_string())
775        );
776
777        let explicit_options = HashMap::from([
778            (AmazonS3ConfigKey::Region, "explicit-region".to_string()),
779            (
780                AmazonS3ConfigKey::Endpoint,
781                "https://explicit.example.com".to_string(),
782            ),
783        ]);
784        let resolved_s3_options =
785            ResolvedS3StorageOptions::new(explicit_options, Some(&profile_config));
786        let region = resolve_s3_region(&url, &resolved_s3_options).await.unwrap();
787
788        assert_eq!(region.as_deref(), Some("explicit-region"));
789        assert_eq!(
790            resolved_s3_options
791                .options
792                .get(&AmazonS3ConfigKey::Endpoint),
793            Some(&"https://explicit.example.com".to_string())
794        );
795    }
796
797    #[tokio::test]
798    async fn test_region_only_aws_profile_preserves_bucket_discovery() {
799        let profile_config = load_test_profile("[profile selected]\nregion = us-east-1").await;
800        let resolved_s3_options =
801            ResolvedS3StorageOptions::new(HashMap::new(), Some(&profile_config));
802
803        let url = Url::parse("s3:///path").unwrap();
804        let error = resolve_s3_region(&url, &resolved_s3_options)
805            .await
806            .unwrap_err();
807
808        assert!(matches!(error, Error::InvalidInput { .. }));
809        assert!(error.to_string().contains("Could not parse bucket"));
810    }
811
812    #[tokio::test]
813    async fn test_r2_aws_profile_requires_constant_size_upload_parts() {
814        let profile_config = load_test_profile(
815            "[profile selected]\n\
816             region = auto\n\
817             endpoint_url = https://account.r2.cloudflarestorage.com",
818        )
819        .await;
820        let resolved_s3_options =
821            ResolvedS3StorageOptions::new(HashMap::new(), Some(&profile_config));
822
823        assert!(resolved_s3_options.requires_constant_size_upload_parts());
824    }
825
826    #[test]
827    fn test_s3_path_parsing() {
828        let provider = AwsStoreProvider;
829
830        let cases = [
831            ("s3://bucket/path/to/file", "path/to/file"),
832            // for non ASCII string tests: the URL encodes them, extract_path must decode back
833            ("s3://bucket/测试path/to/file", "测试path/to/file"),
834            ("s3://bucket/path/&to/file", "path/&to/file"),
835            ("s3://bucket/path/=to/file", "path/=to/file"),
836            (
837                "s3+ddb://bucket/path/to/file?ddbTableName=test",
838                "path/to/file",
839            ),
840        ];
841
842        for (uri, expected_path) in cases {
843            let url = Url::parse(uri).unwrap();
844            let path = provider.extract_path(&url).unwrap();
845            // extract_path decodes url.path(), so the Path stores the raw (decoded)
846            // string. Path::parse keeps its input verbatim, matching that, whereas
847            // Path::from would percent-encode non-ASCII bytes and not match.
848            let expected_path = Path::parse(expected_path).unwrap();
849            assert_eq!(path, expected_path)
850        }
851    }
852
853    // Regression test for https://github.com/lance-format/lance/issues/6643
854    // extract_path must NOT double-encode paths that contain non-ASCII characters.
855    // url.path() returns a percent-encoded string; we must decode it back to raw
856    // UTF-8 before storing it in a Path, so the object store HTTP client can apply
857    // a single, correct percent-encoding when building the request URL.
858    #[test]
859    fn test_s3_non_ascii_path_no_double_encoding() {
860        let provider = AwsStoreProvider;
861
862        // "s3://bucket/中文路径" → url.path() == "/%E4%B8%AD%E6%96%87%E8%B7%AF%E5%BE%84".
863        // The buggy Path::parse(url.path()) stored "%E4%B8%AD..." verbatim; the S3
864        // client then percent-encodes the '%' again, yielding "%25E4%25B8%25AD...".
865        // With Path::from_url_path the Path stores the decoded UTF-8 instead.
866        let url = Url::parse("s3://bucket/中文路径").unwrap();
867        let path = provider.extract_path(&url).unwrap();
868
869        // The Path must hold the decoded UTF-8, not the percent-encoded form.
870        assert_eq!(path.as_ref(), "中文路径");
871    }
872
873    #[test]
874    fn test_is_s3_express() {
875        let cases = [
876            (
877                "s3://bucket/path/to/file",
878                HashMap::from([("s3_express".to_string(), "true".to_string())]),
879                true,
880            ),
881            (
882                "s3://bucket/path/to/file",
883                HashMap::from([("s3_express".to_string(), "false".to_string())]),
884                false,
885            ),
886            ("s3://bucket/path/to/file", HashMap::from([]), false),
887            (
888                "s3://bucket--x-s3/path/to/file",
889                HashMap::from([("s3_express".to_string(), "true".to_string())]),
890                true,
891            ),
892            (
893                "s3://bucket--x-s3/path/to/file",
894                HashMap::from([("s3_express".to_string(), "false".to_string())]),
895                true, // URL takes precedence
896            ),
897            ("s3://bucket--x-s3/path/to/file", HashMap::from([]), true),
898        ];
899
900        for (uri, storage_map, expected) in cases {
901            let url = Url::parse(uri).unwrap();
902            let storage_options = StorageOptions(storage_map);
903            let is_s3_express = check_s3_express(&url, &storage_options);
904            assert_eq!(is_s3_express, expected);
905        }
906    }
907
908    #[tokio::test]
909    async fn test_use_opendal_flag() {
910        use crate::object_store::StorageOptionsAccessor;
911        let provider = AwsStoreProvider;
912        let url = Url::parse("s3://test-bucket/path").unwrap();
913        let params_with_flag = ObjectStoreParams {
914            storage_options_accessor: Some(Arc::new(StorageOptionsAccessor::with_static_options(
915                HashMap::from([
916                    ("use_opendal".to_string(), "true".to_string()),
917                    ("region".to_string(), "us-west-2".to_string()),
918                ]),
919            ))),
920            ..Default::default()
921        };
922
923        let store = provider
924            .new_store(url.clone(), &params_with_flag)
925            .await
926            .unwrap();
927        assert_eq!(store.scheme, "s3");
928    }
929
930    /// S3 Express ignores `start-after` and does not list in key order, but it does hand back
931    /// continuation tokens, which is all the native store resumes from — so an Express bucket
932    /// pages like any other. The OpenDAL arm has no lister to page with at all.
933    #[rstest::rstest]
934    #[case::native("false", true)]
935    #[case::opendal("true", false)]
936    #[tokio::test]
937    async fn test_s3_express_is_paged_by_continuation_token(
938        #[case] use_opendal: &str,
939        #[case] paginated: bool,
940    ) {
941        let provider = AwsStoreProvider;
942        // Express bucket names carry their availability zone, which the S3 client validates.
943        let url = Url::parse("s3://test-bucket--use1-az4--x-s3/path").unwrap();
944        let params = ObjectStoreParams {
945            storage_options_accessor: Some(Arc::new(StorageOptionsAccessor::with_static_options(
946                HashMap::from([
947                    ("use_opendal".to_string(), use_opendal.to_string()),
948                    ("region".to_string(), "us-east-1".to_string()),
949                ]),
950            ))),
951            ..Default::default()
952        };
953
954        let store = provider.new_store(url, &params).await.unwrap();
955
956        assert!(!store.list_is_lexically_ordered);
957        assert_eq!(store.paginated_lister.is_some(), paginated);
958    }
959
960    #[derive(Debug)]
961    struct MockStorageOptionsProvider {
962        call_count: Arc<RwLock<usize>>,
963        expires_in_millis: Option<u64>,
964    }
965
966    impl MockStorageOptionsProvider {
967        fn new(expires_in_millis: Option<u64>) -> Self {
968            Self {
969                call_count: Arc::new(RwLock::new(0)),
970                expires_in_millis,
971            }
972        }
973
974        async fn get_call_count(&self) -> usize {
975            *self.call_count.read().await
976        }
977    }
978
979    #[async_trait::async_trait]
980    impl StorageOptionsProvider for MockStorageOptionsProvider {
981        async fn fetch_storage_options(&self) -> Result<Option<HashMap<String, String>>> {
982            let count = {
983                let mut c = self.call_count.write().await;
984                *c += 1;
985                *c
986            };
987
988            let mut options = HashMap::from([
989                ("aws_access_key_id".to_string(), format!("AKID_{}", count)),
990                (
991                    "aws_secret_access_key".to_string(),
992                    format!("SECRET_{}", count),
993                ),
994                ("aws_session_token".to_string(), format!("TOKEN_{}", count)),
995            ]);
996
997            if let Some(expires_in) = self.expires_in_millis {
998                let now_ms = SystemTime::now()
999                    .duration_since(UNIX_EPOCH)
1000                    .unwrap()
1001                    .as_millis() as u64;
1002                let expires_at = now_ms + expires_in;
1003                options.insert("expires_at_millis".to_string(), expires_at.to_string());
1004            }
1005
1006            Ok(Some(options))
1007        }
1008
1009        fn provider_id(&self) -> String {
1010            let ptr = Arc::as_ptr(&self.call_count) as usize;
1011            format!("MockStorageOptionsProvider {{ id: {} }}", ptr)
1012        }
1013    }
1014
1015    #[tokio::test]
1016    async fn test_dynamic_credential_provider_with_initial_cache() {
1017        MockClock::set_system_time(Duration::from_secs(100_000));
1018
1019        let now_ms = MockClock::system_time().as_millis() as u64;
1020
1021        // Create a mock provider that returns credentials expiring in 10 minutes
1022        let mock = Arc::new(MockStorageOptionsProvider::new(Some(
1023            600_000, // Expires in 10 minutes
1024        )));
1025
1026        // Create initial options with cached credentials that expire in 10 minutes
1027        let expires_at = now_ms + 600_000; // 10 minutes from now
1028        let initial_options = HashMap::from([
1029            ("aws_access_key_id".to_string(), "AKID_CACHED".to_string()),
1030            (
1031                "aws_secret_access_key".to_string(),
1032                "SECRET_CACHED".to_string(),
1033            ),
1034            ("aws_session_token".to_string(), "TOKEN_CACHED".to_string()),
1035            ("expires_at_millis".to_string(), expires_at.to_string()),
1036            ("refresh_offset_millis".to_string(), "300000".to_string()), // 5 minute refresh offset
1037        ]);
1038
1039        let provider = DynamicStorageOptionsCredentialProvider::from_provider_with_initial(
1040            mock.clone(),
1041            initial_options,
1042        );
1043
1044        // First call should use cached credentials (not expired yet)
1045        let cred = provider.get_credential().await.unwrap();
1046        assert_eq!(cred.key_id, "AKID_CACHED");
1047        assert_eq!(cred.secret_key, "SECRET_CACHED");
1048        assert_eq!(cred.token, Some("TOKEN_CACHED".to_string()));
1049
1050        // Should not have called the provider yet
1051        assert_eq!(mock.get_call_count().await, 0);
1052    }
1053
1054    #[tokio::test]
1055    async fn test_dynamic_credential_provider_with_expired_cache() {
1056        MockClock::set_system_time(Duration::from_secs(100_000));
1057
1058        let now_ms = MockClock::system_time().as_millis() as u64;
1059
1060        // Create a mock provider that returns credentials expiring in 10 minutes
1061        let mock = Arc::new(MockStorageOptionsProvider::new(Some(
1062            600_000, // Expires in 10 minutes
1063        )));
1064
1065        // Create initial options with credentials that expired 1 second ago
1066        let expired_time = now_ms - 1_000; // 1 second ago
1067        let initial_options = HashMap::from([
1068            ("aws_access_key_id".to_string(), "AKID_EXPIRED".to_string()),
1069            (
1070                "aws_secret_access_key".to_string(),
1071                "SECRET_EXPIRED".to_string(),
1072            ),
1073            ("expires_at_millis".to_string(), expired_time.to_string()),
1074            ("refresh_offset_millis".to_string(), "300000".to_string()), // 5 minute refresh offset
1075        ]);
1076
1077        let provider = DynamicStorageOptionsCredentialProvider::from_provider_with_initial(
1078            mock.clone(),
1079            initial_options,
1080        );
1081
1082        // First call should fetch new credentials because cached ones are expired
1083        let cred = provider.get_credential().await.unwrap();
1084        assert_eq!(cred.key_id, "AKID_1");
1085        assert_eq!(cred.secret_key, "SECRET_1");
1086        assert_eq!(cred.token, Some("TOKEN_1".to_string()));
1087
1088        // Should have called the provider once
1089        assert_eq!(mock.get_call_count().await, 1);
1090    }
1091
1092    #[tokio::test]
1093    async fn test_dynamic_credential_provider_refresh_lead_time() {
1094        MockClock::set_system_time(Duration::from_secs(100_000));
1095
1096        // Create a mock provider that returns credentials expiring in 30 seconds
1097        let mock = Arc::new(MockStorageOptionsProvider::new(Some(
1098            30_000, // Expires in 30 seconds
1099        )));
1100
1101        // Create credential provider with default 60 second refresh offset
1102        // This means credentials should be refreshed when they have less than 60 seconds left
1103        let provider = DynamicStorageOptionsCredentialProvider::from_provider(mock.clone());
1104
1105        // First call should fetch credentials from provider (no initial cache)
1106        // Credentials expire in 30 seconds, which is less than our 60 second refresh offset,
1107        // so they should be considered "needs refresh" immediately
1108        let cred = provider.get_credential().await.unwrap();
1109        assert_eq!(cred.key_id, "AKID_1");
1110        assert_eq!(mock.get_call_count().await, 1);
1111
1112        // Second call should trigger refresh because credentials expire in 30 seconds
1113        // but our refresh lead time is 60 seconds (now + 60sec > expires_at)
1114        // The mock will return new credentials (AKID_2) with the same expiration
1115        let cred = provider.get_credential().await.unwrap();
1116        assert_eq!(cred.key_id, "AKID_2");
1117        assert_eq!(mock.get_call_count().await, 2);
1118    }
1119
1120    #[tokio::test]
1121    async fn test_dynamic_credential_provider_no_initial_cache() {
1122        MockClock::set_system_time(Duration::from_secs(100_000));
1123
1124        // Create a mock provider that returns credentials expiring in 2 minutes
1125        let mock = Arc::new(MockStorageOptionsProvider::new(Some(
1126            120_000, // Expires in 2 minutes
1127        )));
1128
1129        // Create credential provider without initial cache, using default 60 second refresh offset
1130        let provider = DynamicStorageOptionsCredentialProvider::from_provider(mock.clone());
1131
1132        // First call should fetch from provider (call count = 1)
1133        let cred = provider.get_credential().await.unwrap();
1134        assert_eq!(cred.key_id, "AKID_1");
1135        assert_eq!(cred.secret_key, "SECRET_1");
1136        assert_eq!(cred.token, Some("TOKEN_1".to_string()));
1137        assert_eq!(mock.get_call_count().await, 1);
1138
1139        // Second call should use cached credentials (not expired yet, still > 60 seconds remaining)
1140        let cred = provider.get_credential().await.unwrap();
1141        assert_eq!(cred.key_id, "AKID_1");
1142        assert_eq!(mock.get_call_count().await, 1); // Still 1, didn't fetch again
1143
1144        // Advance time to 90 seconds - should trigger refresh (within 60 sec refresh offset)
1145        // At this point, credentials expire in 30 seconds (< 60 sec offset)
1146        MockClock::set_system_time(Duration::from_secs(100_000 + 90));
1147        let cred = provider.get_credential().await.unwrap();
1148        assert_eq!(cred.key_id, "AKID_2");
1149        assert_eq!(cred.secret_key, "SECRET_2");
1150        assert_eq!(cred.token, Some("TOKEN_2".to_string()));
1151        assert_eq!(mock.get_call_count().await, 2);
1152
1153        // Advance time to 210 seconds total (90 + 120) - should trigger another refresh
1154        MockClock::set_system_time(Duration::from_secs(100_000 + 210));
1155        let cred = provider.get_credential().await.unwrap();
1156        assert_eq!(cred.key_id, "AKID_3");
1157        assert_eq!(cred.secret_key, "SECRET_3");
1158        assert_eq!(mock.get_call_count().await, 3);
1159    }
1160
1161    #[tokio::test]
1162    async fn test_dynamic_credential_provider_with_initial_options() {
1163        MockClock::set_system_time(Duration::from_secs(100_000));
1164
1165        let now_ms = MockClock::system_time().as_millis() as u64;
1166
1167        // Create a mock provider that returns credentials expiring in 10 minutes
1168        let mock = Arc::new(MockStorageOptionsProvider::new(Some(
1169            600_000, // Expires in 10 minutes
1170        )));
1171
1172        // Create initial options with expiration in 10 minutes
1173        let expires_at = now_ms + 600_000; // 10 minutes from now
1174        let initial_options = HashMap::from([
1175            ("aws_access_key_id".to_string(), "AKID_INITIAL".to_string()),
1176            (
1177                "aws_secret_access_key".to_string(),
1178                "SECRET_INITIAL".to_string(),
1179            ),
1180            ("aws_session_token".to_string(), "TOKEN_INITIAL".to_string()),
1181            ("expires_at_millis".to_string(), expires_at.to_string()),
1182            ("refresh_offset_millis".to_string(), "300000".to_string()), // 5 minute refresh offset
1183        ]);
1184
1185        // Create credential provider with initial options
1186        let provider = DynamicStorageOptionsCredentialProvider::from_provider_with_initial(
1187            mock.clone(),
1188            initial_options,
1189        );
1190
1191        // First call should use the initial credential (not expired yet)
1192        let cred = provider.get_credential().await.unwrap();
1193        assert_eq!(cred.key_id, "AKID_INITIAL");
1194        assert_eq!(cred.secret_key, "SECRET_INITIAL");
1195        assert_eq!(cred.token, Some("TOKEN_INITIAL".to_string()));
1196
1197        // Should not have called the provider yet
1198        assert_eq!(mock.get_call_count().await, 0);
1199
1200        // Advance time to 6 minutes - this should trigger a refresh
1201        // (5 minute refresh offset means we refresh 5 minutes before expiration)
1202        MockClock::set_system_time(Duration::from_secs(100_000 + 360));
1203        let cred = provider.get_credential().await.unwrap();
1204        assert_eq!(cred.key_id, "AKID_1");
1205        assert_eq!(cred.secret_key, "SECRET_1");
1206        assert_eq!(cred.token, Some("TOKEN_1".to_string()));
1207
1208        // Should have called the provider once
1209        assert_eq!(mock.get_call_count().await, 1);
1210
1211        // Advance time to 11 minutes total - this should trigger another refresh
1212        MockClock::set_system_time(Duration::from_secs(100_000 + 660));
1213        let cred = provider.get_credential().await.unwrap();
1214        assert_eq!(cred.key_id, "AKID_2");
1215        assert_eq!(cred.secret_key, "SECRET_2");
1216        assert_eq!(cred.token, Some("TOKEN_2".to_string()));
1217
1218        // Should have called the provider twice
1219        assert_eq!(mock.get_call_count().await, 2);
1220
1221        // Advance time to 16 minutes total - this should trigger yet another refresh
1222        MockClock::set_system_time(Duration::from_secs(100_000 + 960));
1223        let cred = provider.get_credential().await.unwrap();
1224        assert_eq!(cred.key_id, "AKID_3");
1225        assert_eq!(cred.secret_key, "SECRET_3");
1226        assert_eq!(cred.token, Some("TOKEN_3".to_string()));
1227
1228        // Should have called the provider three times
1229        assert_eq!(mock.get_call_count().await, 3);
1230    }
1231
1232    #[tokio::test]
1233    async fn test_dynamic_credential_provider_concurrent_access() {
1234        // Create a mock provider with far future expiration
1235        let mock = Arc::new(MockStorageOptionsProvider::new(Some(9999999999999)));
1236
1237        let provider = Arc::new(DynamicStorageOptionsCredentialProvider::from_provider(
1238            mock.clone(),
1239        ));
1240
1241        // Spawn 10 concurrent tasks that all try to get credentials at the same time
1242        let mut handles = vec![];
1243        for i in 0..10 {
1244            let provider = provider.clone();
1245            let handle = tokio::spawn(async move {
1246                let cred = provider.get_credential().await.unwrap();
1247                // Verify we got the correct credentials (should all be AKID_1 from first fetch)
1248                assert_eq!(cred.key_id, "AKID_1");
1249                assert_eq!(cred.secret_key, "SECRET_1");
1250                assert_eq!(cred.token, Some("TOKEN_1".to_string()));
1251                i // Return task number for verification
1252            });
1253            handles.push(handle);
1254        }
1255
1256        // Wait for all tasks to complete
1257        let results: Vec<_> = futures::future::join_all(handles)
1258            .await
1259            .into_iter()
1260            .map(|r| r.unwrap())
1261            .collect();
1262
1263        // Verify all 10 tasks completed successfully
1264        assert_eq!(results.len(), 10);
1265        for i in 0..10 {
1266            assert!(results.contains(&i));
1267        }
1268
1269        // The provider should have been called exactly once (first request triggers fetch,
1270        // subsequent requests use cache)
1271        let call_count = mock.get_call_count().await;
1272        assert_eq!(
1273            call_count, 1,
1274            "Provider should be called exactly once despite concurrent access"
1275        );
1276    }
1277
1278    #[tokio::test]
1279    async fn test_dynamic_credential_provider_concurrent_refresh() {
1280        MockClock::set_system_time(Duration::from_secs(100_000));
1281
1282        let now_ms = MockClock::system_time().as_millis() as u64;
1283
1284        // Create initial options with credentials that expired in the past (1000 seconds ago)
1285        let expires_at = now_ms - 1_000_000;
1286        let initial_options = HashMap::from([
1287            ("aws_access_key_id".to_string(), "AKID_OLD".to_string()),
1288            (
1289                "aws_secret_access_key".to_string(),
1290                "SECRET_OLD".to_string(),
1291            ),
1292            ("aws_session_token".to_string(), "TOKEN_OLD".to_string()),
1293            ("expires_at_millis".to_string(), expires_at.to_string()),
1294            ("refresh_offset_millis".to_string(), "300000".to_string()), // 5 minute refresh offset
1295        ]);
1296
1297        // Mock will return credentials expiring in 1 hour
1298        let mock = Arc::new(MockStorageOptionsProvider::new(Some(
1299            3_600_000, // Expires in 1 hour
1300        )));
1301
1302        let provider = Arc::new(
1303            DynamicStorageOptionsCredentialProvider::from_provider_with_initial(
1304                mock.clone(),
1305                initial_options,
1306            ),
1307        );
1308
1309        // Spawn 20 concurrent tasks that all try to get credentials at the same time
1310        // Since the initial credential is expired, they'll all try to refresh
1311        let mut handles = vec![];
1312        for i in 0..20 {
1313            let provider = provider.clone();
1314            let handle = tokio::spawn(async move {
1315                let cred = provider.get_credential().await.unwrap();
1316                // All should get the new credentials (AKID_1 from first fetch)
1317                assert_eq!(cred.key_id, "AKID_1");
1318                assert_eq!(cred.secret_key, "SECRET_1");
1319                assert_eq!(cred.token, Some("TOKEN_1".to_string()));
1320                i
1321            });
1322            handles.push(handle);
1323        }
1324
1325        // Wait for all tasks to complete
1326        let results: Vec<_> = futures::future::join_all(handles)
1327            .await
1328            .into_iter()
1329            .map(|r| r.unwrap())
1330            .collect();
1331
1332        // Verify all 20 tasks completed successfully
1333        assert_eq!(results.len(), 20);
1334
1335        // The provider should have been called at least once, but possibly more times
1336        // due to the try_write mechanism and race conditions
1337        let call_count = mock.get_call_count().await;
1338        assert!(
1339            call_count >= 1,
1340            "Provider should be called at least once, was called {} times",
1341            call_count
1342        );
1343
1344        // It shouldn't be called 20 times though - the lock should prevent most concurrent fetches
1345        assert!(
1346            call_count < 10,
1347            "Provider should not be called too many times due to lock contention, was called {} times",
1348            call_count
1349        );
1350    }
1351
1352    #[tokio::test]
1353    async fn test_explicit_aws_credentials_takes_precedence_over_accessor() {
1354        // Create a mock storage options provider that should NOT be called
1355        let mock_storage_provider = Arc::new(MockStorageOptionsProvider::new(Some(600_000)));
1356
1357        // Create an accessor with the mock provider
1358        let accessor = Arc::new(StorageOptionsAccessor::with_provider(
1359            mock_storage_provider.clone(),
1360        ));
1361
1362        // Create an explicit AWS credentials provider
1363        let explicit_cred_provider = Arc::new(MockAwsCredentialsProvider::default());
1364
1365        // Build credentials with both aws_credentials AND accessor
1366        // The explicit aws_credentials should take precedence
1367        let (result, _region) = build_aws_credential(
1368            Duration::from_secs(300),
1369            Some(explicit_cred_provider.clone() as AwsCredentialProvider),
1370            None, // no storage_options
1371            Some("us-west-2".to_string()),
1372            Some(accessor),
1373            None,
1374        )
1375        .await
1376        .unwrap();
1377
1378        // Get credential from the result
1379        let cred = result.get_credential().await.unwrap();
1380
1381        // The explicit provider should have been called (it returns empty strings)
1382        assert!(explicit_cred_provider.called.load(Ordering::Relaxed));
1383
1384        // The storage options provider should NOT have been called
1385        assert_eq!(
1386            mock_storage_provider.get_call_count().await,
1387            0,
1388            "Storage options provider should not be called when explicit aws_credentials is provided"
1389        );
1390
1391        // Verify we got credentials from the explicit provider (empty strings)
1392        assert_eq!(cred.key_id, "");
1393        assert_eq!(cred.secret_key, "");
1394    }
1395
1396    #[tokio::test]
1397    async fn test_accessor_used_when_no_explicit_aws_credentials() {
1398        MockClock::set_system_time(Duration::from_secs(100_000));
1399
1400        let now_ms = MockClock::system_time().as_millis() as u64;
1401
1402        // Create a mock storage options provider
1403        let mock_storage_provider = Arc::new(MockStorageOptionsProvider::new(Some(600_000)));
1404
1405        // Create initial options
1406        let expires_at = now_ms + 600_000; // 10 minutes from now
1407        let initial_options = HashMap::from([
1408            (
1409                "aws_access_key_id".to_string(),
1410                "AKID_FROM_ACCESSOR".to_string(),
1411            ),
1412            (
1413                "aws_secret_access_key".to_string(),
1414                "SECRET_FROM_ACCESSOR".to_string(),
1415            ),
1416            (
1417                "aws_session_token".to_string(),
1418                "TOKEN_FROM_ACCESSOR".to_string(),
1419            ),
1420            ("expires_at_millis".to_string(), expires_at.to_string()),
1421            ("refresh_offset_millis".to_string(), "300000".to_string()), // 5 minute refresh offset
1422        ]);
1423
1424        // Create an accessor with initial options and provider
1425        let accessor = Arc::new(StorageOptionsAccessor::with_initial_and_provider(
1426            initial_options,
1427            mock_storage_provider.clone(),
1428        ));
1429
1430        // Build credentials with accessor but NO explicit aws_credentials
1431        let (result, _region) = build_aws_credential(
1432            Duration::from_secs(300),
1433            None, // no explicit aws_credentials
1434            None, // no storage_options
1435            Some("us-west-2".to_string()),
1436            Some(accessor),
1437            None,
1438        )
1439        .await
1440        .unwrap();
1441
1442        // Get credential - should use the initial accessor credentials
1443        let cred = result.get_credential().await.unwrap();
1444        assert_eq!(cred.key_id, "AKID_FROM_ACCESSOR");
1445        assert_eq!(cred.secret_key, "SECRET_FROM_ACCESSOR");
1446
1447        // Storage options provider should NOT have been called yet (using cached initial creds)
1448        assert_eq!(mock_storage_provider.get_call_count().await, 0);
1449
1450        // Advance time to trigger refresh (past the 5 minute refresh offset)
1451        MockClock::set_system_time(Duration::from_secs(100_000 + 360));
1452
1453        // Get credential again - should now fetch from provider
1454        let cred = result.get_credential().await.unwrap();
1455        assert_eq!(cred.key_id, "AKID_1");
1456        assert_eq!(cred.secret_key, "SECRET_1");
1457
1458        // Storage options provider should have been called once
1459        assert_eq!(mock_storage_provider.get_call_count().await, 1);
1460    }
1461
1462    // Test that aws_provider_scheme=token selects static credentials.
1463    #[tokio::test]
1464    async fn test_provider_scheme_token() {
1465        let opts = HashMap::from([
1466            (AmazonS3ConfigKey::AccessKeyId, "AKID".to_string()),
1467            (AmazonS3ConfigKey::SecretAccessKey, "SECRET".to_string()),
1468        ]);
1469
1470        let (provider, _) = build_aws_credential(
1471            Duration::from_secs(300),
1472            None,
1473            Some(&opts),
1474            Some("us-east-1".to_string()),
1475            None,
1476            Some(AwsProviderScheme::Token),
1477        )
1478        .await
1479        .unwrap();
1480
1481        let cred = provider.get_credential().await.unwrap();
1482        assert_eq!(cred.key_id, "AKID");
1483        assert_eq!(cred.secret_key, "SECRET");
1484    }
1485
1486    // Test that aws_provider_scheme=token errors when no static credentials are present.
1487    #[tokio::test]
1488    async fn test_provider_scheme_token_errors_without_credentials() {
1489        let opts: HashMap<AmazonS3ConfigKey, String> = HashMap::new();
1490
1491        let result = build_aws_credential(
1492            Duration::from_secs(300),
1493            None,
1494            Some(&opts),
1495            Some("us-east-1".to_string()),
1496            None,
1497            Some(AwsProviderScheme::Token),
1498        )
1499        .await;
1500        assert!(result.is_err());
1501        assert!(
1502            result
1503                .unwrap_err()
1504                .to_string()
1505                .contains("aws_provider_scheme=token"),
1506            "error should mention aws_provider_scheme=token"
1507        );
1508    }
1509
1510    // Test that aws_provider_scheme=ecs builds a provider without error.
1511    // The ECS provider itself reads from env vars lazily; construction always succeeds.
1512    #[tokio::test]
1513    async fn test_provider_scheme_ecs() {
1514        let opts: HashMap<AmazonS3ConfigKey, String> = HashMap::new();
1515
1516        let result = build_aws_credential(
1517            Duration::from_secs(300),
1518            None,
1519            Some(&opts),
1520            Some("us-east-1".to_string()),
1521            None,
1522            Some(AwsProviderScheme::Ecs),
1523        )
1524        .await;
1525        assert!(result.is_ok(), "ECS provider should build without error");
1526    }
1527
1528    // Test that aws_provider_scheme=irsa builds a provider and attempts credential
1529    // retrieval (which fails with a provider error, not a config error like
1530    // "Missing Region" — confirming the region is wired through to the STS client).
1531    #[tokio::test]
1532    async fn test_provider_scheme_irsa() {
1533        let opts: HashMap<AmazonS3ConfigKey, String> = HashMap::new();
1534
1535        let (provider, _) = build_aws_credential(
1536            Duration::from_secs(300),
1537            None,
1538            Some(&opts),
1539            Some("us-east-1".to_string()),
1540            None,
1541            Some(AwsProviderScheme::Irsa),
1542        )
1543        .await
1544        .unwrap();
1545
1546        // Credential retrieval must fail with a provider error (missing env vars or
1547        // network), NOT a configuration error like "Invalid Configuration: Missing Region".
1548        let err = provider.get_credential().await.unwrap_err();
1549        assert!(
1550            !err.to_string().contains("Missing Region"),
1551            "should not fail with Missing Region; region was provided. got: {err}"
1552        );
1553    }
1554
1555    // Test that an invalid aws_provider_scheme value produces a clear error.
1556    #[test]
1557    fn test_provider_scheme_invalid_value() {
1558        let opts = StorageOptions::new(HashMap::from([(
1559            "aws_provider_scheme".to_string(),
1560            "magic".to_string(),
1561        )]));
1562        let result = opts.aws_provider_scheme();
1563        assert!(result.is_err());
1564        assert!(result.unwrap_err().to_string().contains("magic"));
1565    }
1566
1567    // Test that no aws_provider_scheme falls through to DefaultCredentialsChain without error.
1568    #[tokio::test]
1569    async fn test_no_provider_scheme_uses_default_chain() {
1570        let opts: HashMap<AmazonS3ConfigKey, String> = HashMap::new();
1571
1572        let result = build_aws_credential(
1573            Duration::from_secs(300),
1574            None,
1575            Some(&opts),
1576            Some("us-east-1".to_string()),
1577            None,
1578            None,
1579        )
1580        .await;
1581        assert!(result.is_ok());
1582    }
1583}