Skip to main content

polars_io/cloud/
options.rs

1#[cfg(feature = "aws")]
2use std::io::Read;
3#[cfg(feature = "aws")]
4use std::path::Path;
5use std::str::FromStr;
6#[cfg(any(feature = "aws", feature = "gcp", feature = "azure", feature = "http"))]
7use std::sync::Arc;
8use std::sync::LazyLock;
9
10#[cfg(any(feature = "aws", feature = "gcp", feature = "azure", feature = "http"))]
11use object_store::ClientOptions;
12#[cfg(feature = "aws")]
13use object_store::aws::AmazonS3Builder;
14#[cfg(feature = "aws")]
15pub use object_store::aws::AmazonS3ConfigKey;
16#[cfg(feature = "azure")]
17pub use object_store::azure::AzureConfigKey;
18#[cfg(feature = "azure")]
19use object_store::azure::MicrosoftAzureBuilder;
20#[cfg(feature = "gcp")]
21use object_store::gcp::GoogleCloudStorageBuilder;
22#[cfg(feature = "gcp")]
23pub use object_store::gcp::GoogleConfigKey;
24use polars_error::*;
25#[cfg(feature = "aws")]
26use polars_utils::cache::LruCache;
27use polars_utils::pl_path::{CloudScheme, PlRefPath};
28use polars_utils::total_ord::TotalOrdWrap;
29#[cfg(feature = "http")]
30use reqwest::header::HeaderMap;
31#[cfg(feature = "serde")]
32use serde::{Deserialize, Serialize};
33
34#[cfg(feature = "cloud")]
35use super::credential_provider::PlCredentialProvider;
36#[cfg(feature = "cloud")]
37use crate::cloud::ObjectStoreErrorContext;
38#[cfg(any(feature = "aws", feature = "gcp", feature = "azure", feature = "http"))]
39use crate::cloud::dns::{CachingResolver, DnsResolverConfig};
40#[cfg(feature = "file_cache")]
41use crate::file_cache::get_env_file_cache_ttl;
42#[cfg(feature = "aws")]
43use crate::pl_async::with_concurrency_budget;
44
45#[cfg(feature = "aws")]
46fn to_io_err(err: reqwest::Error) -> PolarsError {
47    PolarsError::IO {
48        error: Arc::new(std::io::Error::other(err)),
49        msg: None,
50    }
51}
52
53#[cfg(feature = "aws")]
54static BUCKET_REGION: LazyLock<
55    std::sync::Mutex<LruCache<polars_utils::pl_str::PlSmallStr, polars_utils::pl_str::PlSmallStr>>,
56> = LazyLock::new(|| std::sync::Mutex::new(LruCache::with_capacity(32)));
57
58/// The type of the config keys must satisfy the following requirements:
59/// 1. must be easily collected into a HashMap, the type required by the object_crate API.
60/// 2. be Serializable, required when the serde-lazy feature is defined.
61/// 3. not actually use HashMap since that type is disallowed in Polars for performance reasons.
62///
63/// Currently this type is a vector of pairs config key - config value.
64#[allow(dead_code)]
65type Configs<T> = Vec<(T, String)>;
66
67#[derive(Clone, Debug, PartialEq, Hash, Eq)]
68#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
69#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
70pub enum CloudConfig {
71    #[cfg(feature = "aws")]
72    Aws(
73        #[cfg_attr(feature = "dsl-schema", schemars(with = "Vec<(String, String)>"))]
74        Configs<AmazonS3ConfigKey>,
75    ),
76    #[cfg(feature = "azure")]
77    Azure(
78        #[cfg_attr(feature = "dsl-schema", schemars(with = "Vec<(String, String)>"))]
79        Configs<AzureConfigKey>,
80    ),
81    #[cfg(feature = "gcp")]
82    Gcp(
83        #[cfg_attr(feature = "dsl-schema", schemars(with = "Vec<(String, String)>"))]
84        Configs<GoogleConfigKey>,
85    ),
86    #[cfg(feature = "http")]
87    Http {
88        headers: Vec<(String, String)>,
89    },
90    Ext {
91        options: Vec<(String, String)>,
92    },
93}
94
95#[derive(Clone, Debug, PartialEq, Hash, Eq)]
96#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
97#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
98/// Options to connect to various cloud providers.
99pub struct CloudOptions {
100    #[cfg(feature = "file_cache")]
101    pub file_cache_ttl: u64,
102    pub config: Option<CloudConfig>,
103    #[cfg_attr(feature = "serde", serde(default))]
104    pub retry_config: CloudRetryConfig,
105    #[cfg(feature = "cloud")]
106    /// Note: In most cases you will want to access this via [`CloudOptions::initialized_credential_provider`]
107    /// rather than directly.
108    pub(crate) credential_provider: Option<PlCredentialProvider>,
109}
110
111impl Default for CloudOptions {
112    fn default() -> Self {
113        Self::default_static_ref().clone()
114    }
115}
116
117impl CloudOptions {
118    pub fn default_static_ref() -> &'static Self {
119        static DEFAULT: LazyLock<CloudOptions> = LazyLock::new(|| CloudOptions {
120            #[cfg(feature = "file_cache")]
121            file_cache_ttl: get_env_file_cache_ttl(),
122            config: None,
123            retry_config: CloudRetryConfig::default(),
124            #[cfg(feature = "cloud")]
125            credential_provider: None,
126        });
127
128        &DEFAULT
129    }
130}
131
132#[derive(Clone, Copy, Default, Debug, PartialEq, Hash, Eq)]
133#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
134#[cfg_attr(feature = "dsl-schema", derive(schemars::JsonSchema))]
135pub struct CloudRetryConfig {
136    pub max_retries: Option<usize>,
137    pub retry_timeout: Option<std::time::Duration>,
138    pub retry_init_backoff: Option<std::time::Duration>,
139    pub retry_max_backoff: Option<std::time::Duration>,
140    pub retry_base_multiplier: Option<TotalOrdWrap<f64>>,
141}
142
143#[cfg(any(feature = "aws", feature = "gcp", feature = "azure"))]
144impl From<CloudRetryConfig> for object_store::RetryConfig {
145    fn from(value: CloudRetryConfig) -> Self {
146        use std::time::Duration;
147
148        use polars_core::config::verbose;
149
150        let out = object_store::RetryConfig {
151            backoff: object_store::BackoffConfig {
152                init_backoff: value
153                    .retry_init_backoff
154                    .unwrap_or_else(|| DEFAULTS.backoff.init_backoff),
155                max_backoff: value
156                    .retry_max_backoff
157                    .unwrap_or_else(|| DEFAULTS.backoff.max_backoff),
158                base: value
159                    .retry_base_multiplier
160                    .map_or_else(|| DEFAULTS.backoff.base, |x| x.0),
161            },
162            max_retries: value.max_retries.unwrap_or_else(|| DEFAULTS.max_retries),
163            retry_timeout: value
164                .retry_timeout
165                .unwrap_or_else(|| DEFAULTS.retry_timeout),
166        };
167
168        if verbose() {
169            eprintln!("object-store retry config: {:?}", out)
170        }
171
172        return out;
173
174        static DEFAULTS: LazyLock<object_store::RetryConfig> =
175            LazyLock::new(|| object_store::RetryConfig {
176                backoff: object_store::BackoffConfig {
177                    init_backoff: Duration::from_millis(parse_env_var(
178                        100,
179                        "POLARS_CLOUD_RETRY_INIT_BACKOFF_MS",
180                    )),
181                    max_backoff: Duration::from_millis(parse_env_var(
182                        15 * 1000,
183                        "POLARS_CLOUD_RETRY_MAX_BACKOFF_MS",
184                    )),
185                    base: parse_env_var(2., "POLARS_CLOUD_RETRY_BASE_MULTIPLIER"),
186                },
187                max_retries: parse_env_var(2, "POLARS_CLOUD_MAX_RETRIES"),
188                retry_timeout: Duration::from_millis(parse_env_var(
189                    10 * 1000,
190                    "POLARS_CLOUD_RETRY_TIMEOUT_MS",
191                )),
192            });
193
194        fn parse_env_var<T: FromStr>(default: T, name: &'static str) -> T {
195            std::env::var(name).map_or(default, |x| {
196                x.parse::<T>()
197                    .ok()
198                    .unwrap_or_else(|| panic!("invalid value for {name}: {x}"))
199            })
200        }
201    }
202}
203
204#[cfg(feature = "http")]
205pub(crate) fn try_build_http_header_map_from_items_slice<S: AsRef<str>>(
206    headers: &[(S, S)],
207) -> PolarsResult<HeaderMap> {
208    use reqwest::header::{HeaderName, HeaderValue};
209
210    let mut map = HeaderMap::with_capacity(headers.len());
211    for (k, v) in headers {
212        let (k, v) = (k.as_ref(), v.as_ref());
213        map.insert(
214            HeaderName::from_str(k).map_err(to_compute_err)?,
215            HeaderValue::from_str(v).map_err(to_compute_err)?,
216        );
217    }
218
219    Ok(map)
220}
221
222#[allow(dead_code)]
223/// Parse an untype configuration hashmap to a typed configuration for the given configuration key type.
224fn parse_untyped_config<T, I: IntoIterator<Item = (impl AsRef<str>, impl Into<String>)>>(
225    config: I,
226) -> PolarsResult<Configs<T>>
227where
228    T: FromStr + Eq + std::hash::Hash,
229{
230    Ok(config
231        .into_iter()
232        // Silently ignores custom upstream storage_options
233        .filter_map(|(key, val)| {
234            T::from_str(key.as_ref().to_ascii_lowercase().as_str())
235                .ok()
236                .map(|typed_key| (typed_key, val.into()))
237        })
238        .collect::<Configs<T>>())
239}
240
241#[derive(Debug, Copy, Clone, PartialEq)]
242pub enum CloudType {
243    Aws,
244    Azure,
245    /// URI with 'file:' scheme
246    File,
247    /// Google cloud platform
248    Gcp,
249    Http,
250    /// HuggingFace
251    Hf,
252    /// Externally registered scheme (e.g. hdfs:// as "hdfs")
253    Ext(&'static str),
254}
255
256impl CloudType {
257    pub fn from_cloud_scheme(scheme: CloudScheme) -> Self {
258        match scheme {
259            CloudScheme::Abfs
260            | CloudScheme::Abfss
261            | CloudScheme::Adl
262            | CloudScheme::Az
263            | CloudScheme::Azure => Self::Azure,
264
265            CloudScheme::File | CloudScheme::FileNoHostname => Self::File,
266
267            CloudScheme::Gcs | CloudScheme::Gs => Self::Gcp,
268
269            CloudScheme::Hf => Self::Hf,
270
271            CloudScheme::Http | CloudScheme::Https => Self::Http,
272
273            CloudScheme::S3 | CloudScheme::S3a => Self::Aws,
274
275            CloudScheme::Ext(scheme) => Self::Ext(scheme),
276        }
277    }
278}
279
280pub static USER_AGENT: &str = concat!("polars", "/", env!("CARGO_PKG_VERSION"),);
281
282#[cfg(any(feature = "aws", feature = "gcp", feature = "azure", feature = "http"))]
283pub(super) fn get_client_options() -> ClientOptions {
284    use std::num::NonZeroU64;
285
286    use reqwest::header::HeaderValue;
287
288    ClientOptions::new()
289        // Disables the time limit for downloading the response body.
290        .with_timeout_disabled()
291        // Set the time limit for establishing the connection.
292        .with_connect_timeout(std::time::Duration::from_secs(
293            std::env::var("POLARS_HTTP_CONNECT_TIMEOUT_SECONDS")
294                .map(|x| {
295                    x.parse::<NonZeroU64>()
296                        .ok()
297                        .unwrap_or_else(|| {
298                            panic!("invalid value for POLARS_HTTP_CONNECT_TIMEOUT_SECONDS: {x}")
299                        })
300                        .get()
301                })
302                .unwrap_or(5 * 60),
303        ))
304        .with_user_agent(HeaderValue::from_static(USER_AGENT))
305        .with_allow_http(true)
306        .with_dns_resolver(Arc::new(
307            CachingResolver::new(DnsResolverConfig::from_env()),
308        ))
309}
310
311#[cfg(feature = "aws")]
312fn read_config(
313    builder: &mut AmazonS3Builder,
314    items: &[(&Path, &[(&str, AmazonS3ConfigKey)])],
315) -> Option<()> {
316    use crate::path_utils::resolve_homedir;
317
318    for (path, keys) in items {
319        if keys
320            .iter()
321            .all(|(_, key)| builder.get_config_value(key).is_some())
322        {
323            continue;
324        }
325
326        let mut config = std::fs::File::open(resolve_homedir(path)).ok()?;
327        let mut buf = vec![];
328        config.read_to_end(&mut buf).ok()?;
329        let content = std::str::from_utf8(buf.as_ref()).ok()?;
330
331        for (pattern, key) in keys.iter() {
332            if builder.get_config_value(key).is_none() {
333                let reg = polars_utils::regex_cache::compile_regex(pattern).unwrap();
334                let cap = reg.captures(content)?;
335                let m = cap.get(1)?;
336                let parsed = m.as_str();
337                *builder = std::mem::take(builder).with_config(*key, parsed);
338            }
339        }
340    }
341    Some(())
342}
343
344impl CloudOptions {
345    pub fn with_retry_config(mut self, retry_config: CloudRetryConfig) -> Self {
346        self.retry_config = retry_config;
347        self
348    }
349
350    #[cfg(feature = "cloud")]
351    pub fn with_credential_provider(
352        mut self,
353        credential_provider: Option<PlCredentialProvider>,
354    ) -> Self {
355        self.credential_provider = credential_provider;
356        self
357    }
358
359    /// Set the configuration for AWS connections. This is the preferred API from rust.
360    #[cfg(feature = "aws")]
361    pub fn with_aws<I: IntoIterator<Item = (AmazonS3ConfigKey, impl Into<String>)>>(
362        mut self,
363        configs: I,
364    ) -> Self {
365        self.config = Some(CloudConfig::Aws(
366            configs.into_iter().map(|(k, v)| (k, v.into())).collect(),
367        ));
368        self
369    }
370
371    /// Build the [`object_store::ObjectStore`] implementation for AWS.
372    #[cfg(feature = "aws")]
373    pub async fn build_aws(
374        &self,
375        url: PlRefPath,
376        clear_cached_credentials: bool,
377    ) -> PolarsResult<impl object_store::ObjectStore> {
378        use super::credential_provider::IntoCredentialProvider;
379
380        let opt_credential_provider =
381            self.initialized_credential_provider(clear_cached_credentials)?;
382
383        let mut builder = AmazonS3Builder::from_env()
384            .with_client_options(get_client_options())
385            .with_url(url.clone().to_string());
386
387        if let Some(credential_provider) = &opt_credential_provider {
388            let storage_update_options = parse_untyped_config::<AmazonS3ConfigKey, _>(
389                credential_provider
390                    .storage_update_options()?
391                    .into_iter()
392                    .map(|(k, v)| (k, v.to_string())),
393            )?;
394
395            for (key, value) in storage_update_options {
396                builder = builder.with_config(key, value);
397            }
398        }
399
400        read_config(
401            &mut builder,
402            &[(
403                Path::new("~/.aws/config"),
404                &[("region\\s*=\\s*([^\r\n]*)", AmazonS3ConfigKey::Region)],
405            )],
406        );
407
408        read_config(
409            &mut builder,
410            &[(
411                Path::new("~/.aws/credentials"),
412                &[
413                    (
414                        "aws_access_key_id\\s*=\\s*([^\\r\\n]*)",
415                        AmazonS3ConfigKey::AccessKeyId,
416                    ),
417                    (
418                        "aws_secret_access_key\\s*=\\s*([^\\r\\n]*)",
419                        AmazonS3ConfigKey::SecretAccessKey,
420                    ),
421                    (
422                        "aws_session_token\\s*=\\s*([^\\r\\n]*)",
423                        AmazonS3ConfigKey::Token,
424                    ),
425                ],
426            )],
427        );
428
429        if let Some(options) = &self.config {
430            let CloudConfig::Aws(options) = options else {
431                panic!("impl error: cloud type mismatch")
432            };
433            for (key, value) in options {
434                builder = builder.with_config(*key, value);
435            }
436        }
437
438        if builder
439            .get_config_value(&AmazonS3ConfigKey::DefaultRegion)
440            .is_none()
441            && builder
442                .get_config_value(&AmazonS3ConfigKey::Region)
443                .is_none()
444        {
445            let bucket = crate::cloud::CloudLocation::new(url.clone(), false)?.bucket;
446            let region = {
447                let mut bucket_region = BUCKET_REGION.lock().unwrap();
448                bucket_region.get(bucket.as_str()).cloned()
449            };
450
451            match region {
452                Some(region) => {
453                    builder = builder.with_config(AmazonS3ConfigKey::Region, region.as_str())
454                },
455                None => {
456                    if builder
457                        .get_config_value(&AmazonS3ConfigKey::Endpoint)
458                        .is_some()
459                    {
460                        // Set a default value if the endpoint is not aws.
461                        // See: #13042
462                        builder = builder.with_config(AmazonS3ConfigKey::Region, "us-east-1");
463                    } else {
464                        polars_warn!(
465                            "'(default_)region' not set; polars will try to get it from bucket\n\nSet the region manually to silence this warning."
466                        );
467                        let result = with_concurrency_budget(1, || async {
468                            reqwest::Client::builder()
469                                .user_agent(USER_AGENT)
470                                .build()
471                                .unwrap()
472                                .head(format!("https://{bucket}.s3.amazonaws.com"))
473                                .send()
474                                .await
475                                .map_err(to_io_err)
476                        })
477                        .await?;
478                        if let Some(region) = result.headers().get("x-amz-bucket-region") {
479                            let region =
480                                std::str::from_utf8(region.as_bytes()).map_err(to_compute_err)?;
481                            let mut bucket_region = BUCKET_REGION.lock().unwrap();
482                            bucket_region.insert(bucket, region.into());
483                            builder = builder.with_config(AmazonS3ConfigKey::Region, region)
484                        }
485                    }
486                },
487            };
488        };
489
490        let builder = builder.with_retry(self.retry_config.into());
491
492        let opt_credential_provider = match opt_credential_provider {
493            #[cfg(feature = "python")]
494            Some(PlCredentialProvider::Python(object)) => {
495                if pyo3::Python::attach(|py| {
496                    let Ok(func_object) = object
497                        .unwrap_as_provider_ref()
498                        .getattr(py, "_can_use_as_provider")
499                    else {
500                        return PolarsResult::Ok(true);
501                    };
502
503                    Ok(func_object.call0(py)?.extract::<bool>(py).unwrap())
504                })? {
505                    Some(PlCredentialProvider::Python(object))
506                } else {
507                    None
508                }
509            },
510
511            v => v,
512        };
513
514        let builder = if let Some(credential_provider) = opt_credential_provider {
515            builder.with_credentials(credential_provider.into_aws_provider())
516        } else {
517            builder
518        };
519
520        let builder = if builder
521            .get_config_value(&AmazonS3ConfigKey::Checksum)
522            .is_none()
523        {
524            // AWS default checksum, which is also more efficient than SHA256.
525            builder.with_checksum_algorithm(object_store::aws::Checksum::CRC64NVME)
526        } else {
527            builder
528        };
529
530        let out = builder
531            .with_unsigned_payload(true)
532            .build()
533            .map_err(|e| ObjectStoreErrorContext::new(url).attach_err_info(e))?;
534
535        Ok(out)
536    }
537
538    /// Set the configuration for Azure connections. This is the preferred API from rust.
539    #[cfg(feature = "azure")]
540    pub fn with_azure<I: IntoIterator<Item = (AzureConfigKey, impl Into<String>)>>(
541        mut self,
542        configs: I,
543    ) -> Self {
544        self.config = Some(CloudConfig::Azure(
545            configs.into_iter().map(|(k, v)| (k, v.into())).collect(),
546        ));
547        self
548    }
549
550    /// Build the [`object_store::ObjectStore`] implementation for Azure.
551    #[cfg(feature = "azure")]
552    pub fn build_azure(
553        &self,
554        url: PlRefPath,
555        clear_cached_credentials: bool,
556    ) -> PolarsResult<impl object_store::ObjectStore> {
557        use super::credential_provider::IntoCredentialProvider;
558        use crate::cloud::ObjectStoreErrorContext;
559
560        let verbose = polars_core::config::verbose();
561
562        // The credential provider `self.credentials` is prioritized if it is set. We also need
563        // `from_env()` as it may source environment configured storage account name.
564        let mut builder =
565            MicrosoftAzureBuilder::from_env().with_client_options(get_client_options());
566
567        if let Some(options) = &self.config {
568            let CloudConfig::Azure(options) = options else {
569                panic!("impl error: cloud type mismatch")
570            };
571            for (key, value) in options.iter() {
572                builder = builder.with_config(*key, value);
573            }
574        }
575
576        let builder = builder
577            .with_url(url.to_string())
578            .with_retry(self.retry_config.into());
579
580        let builder =
581            if let Some(v) = self.initialized_credential_provider(clear_cached_credentials)? {
582                if verbose {
583                    eprintln!(
584                        "[CloudOptions::build_azure]: Using credential provider {:?}",
585                        v
586                    );
587                }
588                builder.with_credentials(v.into_azure_provider())
589            } else {
590                builder
591            };
592
593        let out = builder
594            .build()
595            .map_err(|e| ObjectStoreErrorContext::new(url).attach_err_info(e))?;
596
597        Ok(out)
598    }
599
600    /// Set the configuration for GCP connections. This is the preferred API from rust.
601    #[cfg(feature = "gcp")]
602    pub fn with_gcp<I: IntoIterator<Item = (GoogleConfigKey, impl Into<String>)>>(
603        mut self,
604        configs: I,
605    ) -> Self {
606        self.config = Some(CloudConfig::Gcp(
607            configs.into_iter().map(|(k, v)| (k, v.into())).collect(),
608        ));
609        self
610    }
611
612    /// Build the [`object_store::ObjectStore`] implementation for GCP.
613    #[cfg(feature = "gcp")]
614    pub fn build_gcp(
615        &self,
616        url: PlRefPath,
617        clear_cached_credentials: bool,
618    ) -> PolarsResult<impl object_store::ObjectStore> {
619        use super::credential_provider::IntoCredentialProvider;
620
621        let credential_provider = self.initialized_credential_provider(clear_cached_credentials)?;
622
623        let builder = if credential_provider.is_none() {
624            GoogleCloudStorageBuilder::from_env()
625        } else {
626            GoogleCloudStorageBuilder::new()
627        };
628
629        let mut builder = builder.with_client_options(get_client_options());
630
631        if let Some(options) = &self.config {
632            let CloudConfig::Gcp(options) = options else {
633                panic!("impl error: cloud type mismatch")
634            };
635            for (key, value) in options.iter() {
636                builder = builder.with_config(*key, value);
637            }
638        }
639
640        let builder = builder
641            .with_url(url.to_string())
642            .with_retry(self.retry_config.into());
643
644        let builder = if let Some(v) = credential_provider {
645            builder.with_credentials(v.into_gcp_provider())
646        } else {
647            builder
648        };
649
650        let out = builder
651            .build()
652            .map_err(|e| ObjectStoreErrorContext::new(url).attach_err_info(e))?;
653
654        Ok(out)
655    }
656
657    #[cfg(feature = "http")]
658    pub fn build_http(&self, url: PlRefPath) -> PolarsResult<impl object_store::ObjectStore> {
659        let out = object_store::http::HttpBuilder::new()
660            .with_url(url.to_string())
661            .with_client_options({
662                let mut opts = super::get_client_options();
663                if let Some(CloudConfig::Http { headers }) = &self.config {
664                    opts = opts.with_default_headers(try_build_http_header_map_from_items_slice(
665                        headers.as_slice(),
666                    )?);
667                }
668                opts
669            })
670            .build()
671            .map_err(|e| ObjectStoreErrorContext::new(url).attach_err_info(e))?;
672
673        Ok(out)
674    }
675
676    /// Parse a configuration from a Hashmap. This is the interface from Python.
677    #[allow(unused_variables)]
678    pub fn from_untyped_config<I: IntoIterator<Item = (impl AsRef<str>, impl Into<String>)>>(
679        scheme: Option<CloudScheme>,
680        config: I,
681    ) -> PolarsResult<Self> {
682        match scheme.map_or(CloudType::File, CloudType::from_cloud_scheme) {
683            CloudType::Aws => {
684                #[cfg(feature = "aws")]
685                {
686                    parse_untyped_config::<AmazonS3ConfigKey, _>(config)
687                        .map(|aws| Self::default().with_aws(aws))
688                }
689                #[cfg(not(feature = "aws"))]
690                {
691                    polars_bail!(ComputeError: "'aws' feature is not enabled");
692                }
693            },
694            CloudType::Azure => {
695                #[cfg(feature = "azure")]
696                {
697                    parse_untyped_config::<AzureConfigKey, _>(config)
698                        .map(|azure| Self::default().with_azure(azure))
699                }
700                #[cfg(not(feature = "azure"))]
701                {
702                    polars_bail!(ComputeError: "'azure' feature is not enabled");
703                }
704            },
705            CloudType::File => Ok(Self::default()),
706            CloudType::Http => Ok(Self::default()),
707            CloudType::Gcp => {
708                #[cfg(feature = "gcp")]
709                {
710                    parse_untyped_config::<GoogleConfigKey, _>(config)
711                        .map(|gcp| Self::default().with_gcp(gcp))
712                }
713                #[cfg(not(feature = "gcp"))]
714                {
715                    polars_bail!(ComputeError: "'gcp' feature is not enabled");
716                }
717            },
718            CloudType::Hf => {
719                #[cfg(feature = "http")]
720                {
721                    use polars_core::config;
722
723                    use crate::path_utils::resolve_homedir;
724
725                    let mut this = Self::default();
726                    let mut token = None;
727                    let verbose = config::verbose();
728
729                    for (i, (k, v)) in config.into_iter().enumerate() {
730                        let (k, v) = (k.as_ref(), v.into());
731
732                        if i == 0 && k == "token" {
733                            if verbose {
734                                eprintln!("HF token sourced from storage_options");
735                            }
736                            token = Some(v);
737                        } else {
738                            polars_bail!(ComputeError: "unknown configuration key for HF: {}", k)
739                        }
740                    }
741
742                    token = token
743                        .or_else(|| {
744                            let v = std::env::var("HF_TOKEN").ok();
745                            if v.is_some() && verbose {
746                                eprintln!("HF token sourced from HF_TOKEN env var");
747                            }
748                            v
749                        })
750                        .or_else(|| {
751                            let hf_home = std::env::var("HF_HOME");
752                            let hf_home = hf_home.as_deref();
753                            let hf_home = hf_home.unwrap_or("~/.cache/huggingface");
754                            let hf_home = resolve_homedir(hf_home);
755                            let cached_token_path = hf_home.join("token");
756
757                            let v = std::string::String::from_utf8(
758                                std::fs::read(&cached_token_path).ok()?,
759                            )
760                            .ok()
761                            .filter(|x| !x.is_empty());
762
763                            if v.is_some() && verbose {
764                                eprintln!("HF token sourced from {:?}", cached_token_path);
765                            }
766
767                            v
768                        });
769
770                    if let Some(v) = token {
771                        this.config = Some(CloudConfig::Http {
772                            headers: vec![("Authorization".into(), format!("Bearer {v}"))],
773                        })
774                    }
775
776                    Ok(this)
777                }
778                #[cfg(not(feature = "http"))]
779                {
780                    polars_bail!(ComputeError: "'http' feature is not enabled");
781                }
782            },
783            CloudType::Ext(_) => {
784                let pairs: Vec<(String, String)> = config
785                    .into_iter()
786                    .map(|(k, v)| (k.as_ref().to_string(), v.into()))
787                    .collect();
788
789                Ok(Self {
790                    config: if pairs.is_empty() {
791                        None
792                    } else {
793                        Some(CloudConfig::Ext { options: pairs })
794                    },
795                    ..Self::default()
796                })
797            },
798        }
799    }
800
801    /// Python passes a credential provider builder that needs to be called to get the actual credential
802    /// provider.
803    #[cfg(feature = "cloud")]
804    fn initialized_credential_provider(
805        &self,
806        clear_cached_credentials: bool,
807    ) -> PolarsResult<Option<PlCredentialProvider>> {
808        if let Some(v) = self.credential_provider.clone() {
809            v.try_into_initialized(clear_cached_credentials)
810        } else {
811            Ok(None)
812        }
813    }
814}
815
816#[cfg(feature = "cloud")]
817#[cfg(test)]
818mod tests {
819    use hashbrown::HashMap;
820
821    use super::parse_untyped_config;
822
823    #[cfg(feature = "aws")]
824    #[test]
825    fn test_parse_untyped_config() {
826        use object_store::aws::AmazonS3ConfigKey;
827
828        let aws_config = [
829            ("aws_secret_access_key", "a_key"),
830            ("aws_s3_allow_unsafe_rename", "true"),
831        ]
832        .into_iter()
833        .collect::<HashMap<_, _>>();
834        let aws_keys = parse_untyped_config::<AmazonS3ConfigKey, _>(aws_config)
835            .expect("Parsing keys shouldn't have thrown an error");
836
837        assert_eq!(
838            aws_keys.first().unwrap().0,
839            AmazonS3ConfigKey::SecretAccessKey
840        );
841        assert_eq!(aws_keys.len(), 1);
842
843        let aws_config = [
844            ("AWS_SECRET_ACCESS_KEY", "a_key"),
845            ("aws_s3_allow_unsafe_rename", "true"),
846        ]
847        .into_iter()
848        .collect::<HashMap<_, _>>();
849        let aws_keys = parse_untyped_config::<AmazonS3ConfigKey, _>(aws_config)
850            .expect("Parsing keys shouldn't have thrown an error");
851
852        assert_eq!(
853            aws_keys.first().unwrap().0,
854            AmazonS3ConfigKey::SecretAccessKey
855        );
856        assert_eq!(aws_keys.len(), 1);
857    }
858}