polars_python/io/
cloud_options.rs

1use std::time::Duration;
2
3use polars::prelude::CloudScheme;
4use polars_core::config::verbose_print_sensitive;
5use polars_io::cloud::{CloudOptions, CloudRetryConfig};
6use polars_utils::total_ord::TotalOrdWrap;
7use pyo3::exceptions::PyValueError;
8use pyo3::intern;
9use pyo3::prelude::*;
10use pyo3::pybacked::PyBackedStr;
11use pyo3::types::PyDict;
12
13use crate::utils::to_py_err;
14
15/// Interface to `StorageOptionsDict | None` on the Python side
16pub struct OptPyCloudOptions<'py>(Bound<'py, PyAny>);
17
18impl<'a, 'py> FromPyObject<'a, 'py> for OptPyCloudOptions<'py> {
19    type Error = PyErr;
20
21    fn extract(ob: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
22        Ok(Self(ob.to_owned()))
23    }
24}
25
26impl OptPyCloudOptions<'_> {
27    pub fn extract_opt_cloud_options(
28        &self,
29        cloud_scheme: Option<CloudScheme>,
30        credential_provider: Option<Py<PyAny>>,
31    ) -> PyResult<Option<CloudOptions>> {
32        let py = self.0.py();
33
34        let mut storage_options: Vec<(PyBackedStr, String)> = vec![];
35        let mut file_cache_ttl: u64 = 2;
36        let mut retry_config = CloudRetryConfig::default();
37
38        let storage_options_dict: Option<Bound<'_, PyDict>> = self.0.extract()?;
39
40        if let Some(storage_options_dict) = storage_options_dict {
41            storage_options.reserve(
42                storage_options_dict
43                    .call_method0(intern!(py, "__len__"))?
44                    .extract()?,
45            );
46
47            for v in storage_options_dict
48                .call_method0(intern!(py, "items"))?
49                .try_iter()?
50            {
51                let (key, value): (PyBackedStr, Bound<'_, PyAny>) = v?.extract()?;
52
53                macro_rules! expected_type {
54                    ($key_name:expr, $type_name:expr) => {{
55                        |_| {
56                            let key_name = $key_name;
57                            let type_name = $type_name;
58                            PyValueError::new_err(format!(
59                                "invalid value for '{key_name}': '{value}' (expected {type_name})"
60                            ))
61                        }
62                    }};
63                }
64
65                match &*key {
66                    "file_cache_ttl" => {
67                        file_cache_ttl = value
68                            .extract()
69                            .map_err(expected_type!("file_cache_ttl", "int"))?;
70                    },
71                    "max_retries" => {
72                        retry_config.max_retries = value
73                            .extract()
74                            .map_err(expected_type!("max_retries", "int"))?;
75                    },
76                    "retry_timeout_ms" => {
77                        retry_config.retry_timeout = Some(Duration::from_millis(
78                            value
79                                .extract()
80                                .map_err(expected_type!("retry_timeout", "int"))?,
81                        ));
82                    },
83                    "retry_init_backoff_ms" => {
84                        retry_config.retry_init_backoff = Some(Duration::from_millis(
85                            value
86                                .extract()
87                                .map_err(expected_type!("retry_init_backoff", "int"))?,
88                        ));
89                    },
90                    "retry_max_backoff_ms" => {
91                        retry_config.retry_max_backoff = Some(Duration::from_millis(
92                            value
93                                .extract()
94                                .map_err(expected_type!("retry_max_backoff", "int"))?,
95                        ));
96                    },
97                    "retry_base_multiplier" => {
98                        retry_config.retry_base_multiplier = Some(TotalOrdWrap(
99                            value
100                                .extract()
101                                .map_err(expected_type!("retry_base_multiplier", "float"))?,
102                        ));
103                    },
104                    _ => {
105                        let value: String = value.extract().map_err(expected_type!(&key, "str"))?;
106                        storage_options.push((key, value))
107                    },
108                }
109            }
110        }
111
112        let cloud_options = CloudOptions::from_untyped_config(cloud_scheme, storage_options)
113            .map_err(to_py_err)?
114            .with_retry_config(retry_config);
115
116        #[cfg(feature = "cloud")]
117        let mut cloud_options =
118            cloud_options.with_credential_provider(credential_provider.map(
119                polars_io::cloud::credential_provider::PlCredentialProvider::from_python_builder,
120            ));
121
122        #[cfg(feature = "cloud")]
123        if file_cache_ttl > 0 {
124            cloud_options.file_cache_ttl = file_cache_ttl;
125        }
126        verbose_print_sensitive(|| format!("extracted cloud_options: {:?}", cloud_options));
127
128        Ok(Some(cloud_options))
129    }
130}
polars_python/io/cloud_options.rs

polars_python/io/
cloud_options.rs
