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