pyo3_object_store/
http.rs

1use std::sync::Arc;
2
3use object_store::http::{HttpBuilder, HttpStore};
4use pyo3::prelude::*;
5use pyo3::types::{PyDict, PyTuple, PyType};
6use pyo3::{intern, IntoPyObjectExt};
7
8use crate::error::PyObjectStoreResult;
9use crate::retry::PyRetryConfig;
10use crate::{PyClientOptions, PyUrl};
11
12#[derive(Debug, Clone, PartialEq)]
13struct HTTPConfig {
14    url: PyUrl,
15    client_options: Option<PyClientOptions>,
16    retry_config: Option<PyRetryConfig>,
17}
18
19impl HTTPConfig {
20    fn __getnewargs_ex__<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyTuple>> {
21        let args = PyTuple::new(py, vec![self.url.clone()])?.into_bound_py_any(py)?;
22        let kwargs = PyDict::new(py);
23
24        if let Some(client_options) = &self.client_options {
25            kwargs.set_item(intern!(py, "client_options"), client_options.clone())?;
26        }
27        if let Some(retry_config) = &self.retry_config {
28            kwargs.set_item(intern!(py, "retry_config"), retry_config.clone())?;
29        }
30
31        PyTuple::new(py, [args, kwargs.into_bound_py_any(py)?])
32    }
33}
34
35/// A Python-facing wrapper around a [`HttpStore`].
36#[derive(Debug, Clone)]
37#[pyclass(name = "HTTPStore", frozen, subclass)]
38pub struct PyHttpStore {
39    // Note: we don't need to wrap this in a MaybePrefixedStore because the HttpStore manages its
40    // own prefix.
41    store: Arc<HttpStore>,
42    /// A config used for pickling. This must stay in sync with the underlying store's config.
43    config: HTTPConfig,
44}
45
46impl AsRef<Arc<HttpStore>> for PyHttpStore {
47    fn as_ref(&self) -> &Arc<HttpStore> {
48        &self.store
49    }
50}
51
52impl PyHttpStore {
53    /// Consume self and return the underlying [`HttpStore`].
54    pub fn into_inner(self) -> Arc<HttpStore> {
55        self.store
56    }
57}
58
59#[pymethods]
60impl PyHttpStore {
61    #[new]
62    #[pyo3(signature = (url, *, client_options=None, retry_config=None))]
63    fn new(
64        url: PyUrl,
65        client_options: Option<PyClientOptions>,
66        retry_config: Option<PyRetryConfig>,
67    ) -> PyObjectStoreResult<Self> {
68        let mut builder = HttpBuilder::new().with_url(url.clone());
69        if let Some(client_options) = client_options.clone() {
70            builder = builder.with_client_options(client_options.into())
71        }
72        if let Some(retry_config) = retry_config.clone() {
73            builder = builder.with_retry(retry_config.into())
74        }
75        Ok(Self {
76            store: Arc::new(builder.build()?),
77            config: HTTPConfig {
78                url,
79                client_options,
80                retry_config,
81            },
82        })
83    }
84
85    #[classmethod]
86    #[pyo3(signature = (url, *, client_options=None, retry_config=None))]
87    pub(crate) fn from_url<'py>(
88        cls: &Bound<'py, PyType>,
89        py: Python<'py>,
90        url: PyUrl,
91        client_options: Option<PyClientOptions>,
92        retry_config: Option<PyRetryConfig>,
93    ) -> PyObjectStoreResult<Bound<'py, PyAny>> {
94        // Note: we pass **back** through Python so that if cls is a subclass, we instantiate the
95        // subclass
96        let kwargs = PyDict::new(py);
97        kwargs.set_item("url", url)?;
98        kwargs.set_item("client_options", client_options)?;
99        kwargs.set_item("retry_config", retry_config)?;
100        Ok(cls.call((), Some(&kwargs))?)
101    }
102
103    fn __eq__(&self, other: &Bound<PyAny>) -> bool {
104        // Ensure we never error on __eq__ by returning false if the other object is not the same
105        // type
106        other
107            .cast::<PyHttpStore>()
108            .map(|other| self.config == other.get().config)
109            .unwrap_or(false)
110    }
111
112    fn __getnewargs_ex__<'py>(&'py self, py: Python<'py>) -> PyResult<Bound<'py, PyTuple>> {
113        self.config.__getnewargs_ex__(py)
114    }
115
116    fn __repr__(&self) -> String {
117        format!("HTTPStore(\"{}\")", &self.config.url.as_ref())
118    }
119
120    #[getter]
121    fn url(&self) -> &PyUrl {
122        &self.config.url
123    }
124
125    #[getter]
126    fn client_options(&self) -> Option<PyClientOptions> {
127        self.config.client_options.clone()
128    }
129
130    #[getter]
131    fn retry_config(&self) -> Option<PyRetryConfig> {
132        self.config.retry_config.clone()
133    }
134}