pyo3_object_store/
client.rs

1use std::collections::HashMap;
2use std::str::FromStr;
3
4use http::{HeaderMap, HeaderName, HeaderValue};
5use object_store::{Certificate, ClientConfigKey, ClientOptions};
6use pyo3::exceptions::PyValueError;
7use pyo3::prelude::*;
8use pyo3::pybacked::{PyBackedBytes, PyBackedStr};
9use pyo3::types::{PyBytes, PyDict, PyString};
10
11use crate::config::PyConfigValue;
12use crate::error::PyObjectStoreError;
13use crate::PyObjectStoreResult;
14
15/// A wrapper around one or more `Certificate`s parsed from PEM input.
16///
17/// The original PEM is retained so the value round-trips through
18/// [`IntoPyObject`]; parsing happens once, on extraction.
19#[derive(Clone, Debug)]
20struct PyCertificate {
21    pem: Vec<u8>,
22    certificates: Vec<Certificate>,
23}
24
25impl PyCertificate {
26    fn new(pem: Vec<u8>) -> PyObjectStoreResult<Self> {
27        let certificates = Certificate::from_pem_bundle(&pem)?;
28        if certificates.is_empty() {
29            return Err(PyValueError::new_err(
30                "No certificates found in `root_certificate` input; expected one or more PEM-encoded certificates.",
31            )
32            .into());
33        }
34        Ok(Self { pem, certificates })
35    }
36}
37
38impl PartialEq for PyCertificate {
39    fn eq(&self, other: &Self) -> bool {
40        self.pem == other.pem
41    }
42}
43
44impl<'py> FromPyObject<'_, 'py> for PyCertificate {
45    type Error = PyErr;
46
47    fn extract(obj: Borrowed<'_, 'py, pyo3::PyAny>) -> PyResult<Self> {
48        let pem = if let Ok(bytes) = obj.extract::<Vec<u8>>() {
49            bytes
50        } else {
51            obj.extract::<String>()?.into_bytes()
52        };
53        Ok(Self::new(pem)?)
54    }
55}
56
57impl<'py> IntoPyObject<'py> for &PyCertificate {
58    type Target = PyBytes;
59    type Output = Bound<'py, PyBytes>;
60    type Error = PyErr;
61
62    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
63        Ok(PyBytes::new(py, &self.pem))
64    }
65}
66
67/// A wrapper around `ClientConfigKey` that implements [`FromPyObject`].
68#[derive(Clone, Debug, PartialEq, Eq, Hash)]
69pub struct PyClientConfigKey(ClientConfigKey);
70
71impl<'py> FromPyObject<'_, 'py> for PyClientConfigKey {
72    type Error = PyErr;
73
74    fn extract(obj: Borrowed<'_, 'py, pyo3::PyAny>) -> PyResult<Self> {
75        let s = obj.extract::<PyBackedStr>()?.to_lowercase();
76        let key = s.parse().map_err(PyObjectStoreError::ObjectStoreError)?;
77        Ok(Self(key))
78    }
79}
80
81impl<'py> IntoPyObject<'py> for PyClientConfigKey {
82    type Target = PyString;
83    type Output = Bound<'py, PyString>;
84    type Error = PyErr;
85
86    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
87        Ok(PyString::new(py, self.0.as_ref()))
88    }
89}
90
91impl<'py> IntoPyObject<'py> for &PyClientConfigKey {
92    type Target = PyString;
93    type Output = Bound<'py, PyString>;
94    type Error = PyErr;
95
96    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
97        Ok(PyString::new(py, self.0.as_ref()))
98    }
99}
100
101/// A wrapper around `ClientOptions` that implements [`FromPyObject`].
102#[derive(Clone, Debug, PartialEq)]
103pub struct PyClientOptions {
104    string_options: HashMap<PyClientConfigKey, PyConfigValue>,
105    default_headers: Option<PyHeaderMap>,
106    root_certificate: Option<PyCertificate>,
107}
108
109impl<'py> FromPyObject<'_, 'py> for PyClientOptions {
110    type Error = PyErr;
111
112    // Need custom extraction because default headers is not a string value
113    fn extract(obj: Borrowed<'_, 'py, pyo3::PyAny>) -> PyResult<Self> {
114        let dict = obj.extract::<Bound<PyDict>>()?;
115        let mut string_options = HashMap::new();
116        let mut default_headers = None;
117        let mut root_certificate = None;
118
119        for (key, value) in dict.iter() {
120            if let Ok(key) = key.extract::<PyClientConfigKey>() {
121                string_options.insert(key, value.extract::<PyConfigValue>()?);
122            } else {
123                let key = key.extract::<PyBackedStr>()?;
124                match &*key {
125                    "default_headers" => default_headers = Some(value.extract::<PyHeaderMap>()?),
126                    "root_certificate" => {
127                        root_certificate = Some(value.extract::<PyCertificate>()?)
128                    }
129                    _ => return Err(PyValueError::new_err(format!("Invalid key: {key}."))),
130                }
131            }
132        }
133
134        Ok(Self {
135            string_options,
136            default_headers,
137            root_certificate,
138        })
139    }
140}
141
142impl<'py> IntoPyObject<'py> for PyClientOptions {
143    type Target = PyDict;
144    type Output = Bound<'py, PyDict>;
145    type Error = PyErr;
146
147    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
148        let dict = self.string_options.into_pyobject(py)?;
149        if let Some(headers) = self.default_headers {
150            dict.set_item("default_headers", headers)?;
151        }
152        if let Some(certificate) = &self.root_certificate {
153            dict.set_item("root_certificate", certificate)?;
154        }
155        Ok(dict)
156    }
157}
158
159impl<'py> IntoPyObject<'py> for &PyClientOptions {
160    type Target = PyDict;
161    type Output = Bound<'py, PyDict>;
162    type Error = PyErr;
163
164    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
165        let dict = (&self.string_options).into_pyobject(py)?;
166        if let Some(headers) = &self.default_headers {
167            dict.set_item("default_headers", headers)?;
168        }
169        if let Some(certificate) = &self.root_certificate {
170            dict.set_item("root_certificate", certificate)?;
171        }
172        Ok(dict.clone())
173    }
174}
175
176impl From<PyClientOptions> for ClientOptions {
177    fn from(value: PyClientOptions) -> Self {
178        let mut options = ClientOptions::new();
179        for (key, value) in value.string_options.into_iter() {
180            options = options.with_config(key.0, value.0);
181        }
182
183        if let Some(headers) = value.default_headers {
184            options = options.with_default_headers(headers.0);
185        }
186
187        if let Some(certificate) = value.root_certificate {
188            for certificate in certificate.certificates {
189                options = options.with_root_certificate(certificate);
190            }
191        }
192
193        options
194    }
195}
196
197#[derive(Clone, Debug, PartialEq)]
198struct PyHeaderMap(HeaderMap);
199
200impl<'py> FromPyObject<'_, 'py> for PyHeaderMap {
201    type Error = PyErr;
202
203    fn extract(obj: Borrowed<'_, 'py, pyo3::PyAny>) -> PyResult<Self> {
204        let dict = obj.extract::<Bound<PyDict>>()?;
205        let mut header_map = HeaderMap::with_capacity(dict.len());
206        for (key, value) in dict.iter() {
207            let key = HeaderName::from_str(&key.extract::<PyBackedStr>()?)
208                .map_err(|err| PyValueError::new_err(err.to_string()))?;
209
210            // HTTP Header values can have non-ascii bytes, so first try to extract as bytes.
211            let value = if let Ok(value_bytes) = value.extract::<PyBackedBytes>() {
212                HeaderValue::from_bytes(&value_bytes)
213            } else {
214                HeaderValue::from_str(&value.extract::<PyBackedStr>()?)
215            }
216            .map_err(|err| PyValueError::new_err(err.to_string()))?;
217
218            header_map.insert(key, value);
219        }
220        Ok(Self(header_map))
221    }
222}
223
224impl<'py> IntoPyObject<'py> for PyHeaderMap {
225    type Target = PyDict;
226    type Output = Bound<'py, PyDict>;
227    type Error = PyErr;
228
229    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
230        let dict = PyDict::new(py);
231        for (key, value) in self.0.iter() {
232            dict.set_item(key.as_str(), value.as_bytes())?;
233        }
234        Ok(dict)
235    }
236}
237
238impl<'py> IntoPyObject<'py> for &PyHeaderMap {
239    type Target = PyDict;
240    type Output = Bound<'py, PyDict>;
241    type Error = PyErr;
242
243    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
244        let dict = PyDict::new(py);
245        for (key, value) in self.0.iter() {
246            dict.set_item(key.as_str(), value.as_bytes())?;
247        }
248        Ok(dict)
249    }
250}
pyo3_object_store/client.rs

pyo3_object_store/
client.rs
