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