pyo3_object_store/
store.rs

1use std::sync::Arc;
2
3use object_store::ObjectStore;
4use pyo3::exceptions::{PyRuntimeWarning, PyValueError};
5use pyo3::prelude::*;
6use pyo3::pybacked::PyBackedStr;
7use pyo3::types::{PyDict, PyTuple};
8use pyo3::{intern, PyTypeInfo};
9
10use crate::{PyAzureStore, PyGCSStore, PyHttpStore, PyLocalStore, PyMemoryStore, PyS3Store};
11
12/// A wrapper around a Rust ObjectStore instance that allows any rust-native implementation of
13/// ObjectStore.
14///
15/// This will only accept ObjectStore instances created from the same library. See
16/// [register_store_module][crate::register_store_module].
17#[derive(Debug, Clone)]
18pub struct PyObjectStore(Arc<dyn ObjectStore>);
19
20impl<'py> FromPyObject<'_, 'py> for PyObjectStore {
21    type Error = PyErr;
22
23    fn extract(obj: Borrowed<'_, 'py, pyo3::PyAny>) -> PyResult<Self> {
24        if let Ok(store) = obj.cast::<PyS3Store>() {
25            Ok(Self(store.get().as_ref().clone()))
26        } else if let Ok(store) = obj.cast::<PyAzureStore>() {
27            Ok(Self(store.get().as_ref().clone()))
28        } else if let Ok(store) = obj.cast::<PyGCSStore>() {
29            Ok(Self(store.get().as_ref().clone()))
30        } else if let Ok(store) = obj.cast::<PyHttpStore>() {
31            Ok(Self(store.get().as_ref().clone()))
32        } else if let Ok(store) = obj.cast::<PyLocalStore>() {
33            Ok(Self(store.get().as_ref().clone()))
34        } else if let Ok(store) = obj.cast::<PyMemoryStore>() {
35            Ok(Self(store.get().as_ref().clone()))
36        } else {
37            let py = obj.py();
38            // Check for object-store instance from other library
39            let cls_name = obj
40                .getattr(intern!(py, "__class__"))?
41                .getattr(intern!(py, "__name__"))?
42                .extract::<PyBackedStr>()?;
43            if [
44                PyAzureStore::type_object(py).name()?.to_str()?,
45                PyGCSStore::type_object(py).name()?.to_str()?,
46                PyHttpStore::type_object(py).name()?.to_str()?,
47                PyLocalStore::type_object(py).name()?.to_str()?,
48                PyMemoryStore::type_object(py).name()?.to_str()?,
49                PyS3Store::type_object(py).name()?.to_str()?,
50            ]
51            .contains(&cls_name.as_str())
52            {
53                return Err(PyValueError::new_err("You must use an object store instance exported from **the same library** as this function. They cannot be used across libraries.\nThis is because object store instances are compiled with a specific version of Rust and Python." ));
54            }
55
56            Err(PyValueError::new_err(format!(
57                "Expected an object store instance, got {}",
58                obj.repr()?
59            )))
60        }
61    }
62}
63
64impl AsRef<Arc<dyn ObjectStore>> for PyObjectStore {
65    fn as_ref(&self) -> &Arc<dyn ObjectStore> {
66        &self.0
67    }
68}
69
70impl From<PyObjectStore> for Arc<dyn ObjectStore> {
71    fn from(value: PyObjectStore) -> Self {
72        value.0
73    }
74}
75
76impl PyObjectStore {
77    /// Consume self and return the underlying [`ObjectStore`].
78    pub fn into_inner(self) -> Arc<dyn ObjectStore> {
79        self.0
80    }
81
82    /// Consume self and return a reference-counted [`ObjectStore`].
83    pub fn into_dyn(self) -> Arc<dyn ObjectStore> {
84        self.0
85    }
86}
87
88#[derive(Debug, Clone)]
89struct PyExternalObjectStoreInner(Arc<dyn ObjectStore>);
90
91impl<'py> FromPyObject<'_, 'py> for PyExternalObjectStoreInner {
92    type Error = PyErr;
93
94    fn extract(obj: Borrowed<'_, 'py, pyo3::PyAny>) -> PyResult<Self> {
95        let py = obj.py();
96        // Check for object-store instance from other library
97        let cls_name = obj
98            .getattr(intern!(py, "__class__"))?
99            .getattr(intern!(py, "__name__"))?
100            .extract::<PyBackedStr>()?;
101
102        if cls_name.as_str() == PyAzureStore::type_object(py).name()? {
103            let (args, kwargs): (Bound<PyTuple>, Bound<PyDict>) = obj
104                .call_method0(intern!(py, "__getnewargs_ex__"))?
105                .extract()?;
106            let store = PyAzureStore::type_object(py)
107                .call(args, Some(&kwargs))?
108                .cast::<PyAzureStore>()?
109                .get()
110                .clone();
111            return Ok(Self(store.into_inner()));
112        }
113
114        if cls_name.as_str() == PyGCSStore::type_object(py).name()? {
115            let (args, kwargs): (Bound<PyTuple>, Bound<PyDict>) = obj
116                .call_method0(intern!(py, "__getnewargs_ex__"))?
117                .extract()?;
118            let store = PyGCSStore::type_object(py)
119                .call(args, Some(&kwargs))?
120                .cast::<PyGCSStore>()?
121                .get()
122                .clone();
123            return Ok(Self(store.into_inner()));
124        }
125
126        if cls_name.as_str() == PyHttpStore::type_object(py).name()? {
127            let (args, kwargs): (Bound<PyTuple>, Bound<PyDict>) = obj
128                .call_method0(intern!(py, "__getnewargs_ex__"))?
129                .extract()?;
130            let store = PyHttpStore::type_object(py)
131                .call(args, Some(&kwargs))?
132                .cast::<PyHttpStore>()?
133                .get()
134                .clone();
135            return Ok(Self(store.into_inner()));
136        }
137
138        if cls_name.as_str() == PyLocalStore::type_object(py).name()? {
139            let (args, kwargs): (Bound<PyTuple>, Bound<PyDict>) = obj
140                .call_method0(intern!(py, "__getnewargs_ex__"))?
141                .extract()?;
142            let store = PyLocalStore::type_object(py)
143                .call(args, Some(&kwargs))?
144                .cast::<PyLocalStore>()?
145                .get()
146                .clone();
147            return Ok(Self(store.into_inner()));
148        }
149
150        if cls_name.as_str() == PyS3Store::type_object(py).name()? {
151            let (args, kwargs): (Bound<PyTuple>, Bound<PyDict>) = obj
152                .call_method0(intern!(py, "__getnewargs_ex__"))?
153                .extract()?;
154            let store = PyS3Store::type_object(py)
155                .call(args, Some(&kwargs))?
156                .cast::<PyS3Store>()?
157                .get()
158                .clone();
159            return Ok(Self(store.into_inner()));
160        }
161
162        Err(PyValueError::new_err(format!(
163            "Expected an object store-compatible instance, got {}",
164            obj.repr()?
165        )))
166    }
167}
168
169/// A wrapper around a Rust [ObjectStore] instance that will extract and recreate an ObjectStore
170/// instance out of a Python object.
171///
172/// This will accept [ObjectStore] instances from **any** Python library exporting store classes
173/// from `pyo3-object_store`.
174///
175/// ## Caveats
176///
177/// - This will extract the configuration of the store and **recreate** the store instance in the
178///   current module. This means that no connection pooling will be reused from the original
179///   library. Also, there is a slight overhead to this as configuration parsing will need to
180///   happen from scratch.
181///
182///   This will work best when the store is created once and used multiple times.
183///
184/// - This relies on the public Python API (`__getnewargs_ex__` and `__init__`) of the store
185///   classes to extract the configuration. If the public API changes in a non-backwards compatible
186///   way, this store conversion may fail.
187///
188/// - While this reuses `__getnewargs_ex__` (from the pickle implementation) to extract arguments
189///   to pass into `__init__`, it does not actually _use_ pickle, and so even non-pickleable
190///   credential providers should work.
191///
192/// - This will not work for `PyMemoryStore` because we can't clone the internal state of the
193///   store.
194#[derive(Debug, Clone)]
195pub struct PyExternalObjectStore(PyExternalObjectStoreInner);
196
197impl From<PyExternalObjectStore> for Arc<dyn ObjectStore> {
198    fn from(value: PyExternalObjectStore) -> Self {
199        value.0 .0
200    }
201}
202
203impl PyExternalObjectStore {
204    /// Consume self and return a reference-counted [`ObjectStore`].
205    pub fn into_dyn(self) -> Arc<dyn ObjectStore> {
206        self.into()
207    }
208}
209
210impl<'py> FromPyObject<'_, 'py> for PyExternalObjectStore {
211    type Error = PyErr;
212
213    fn extract(obj: Borrowed<'_, 'py, pyo3::PyAny>) -> PyResult<Self> {
214        match obj.extract() {
215            Ok(inner) => {
216                #[cfg(feature = "external-store-warning")]
217                {
218                    let py = obj.py();
219
220                    let warnings_mod = py.import(intern!(py, "warnings"))?;
221                    let warning = PyRuntimeWarning::new_err(
222                    "Successfully reconstructed a store defined in another Python module. Connection pooling will not be shared across store instances.",
223                );
224                    let args = PyTuple::new(py, vec![warning])?;
225                    warnings_mod.call_method1(intern!(py, "warn"), args)?;
226                }
227                Ok(Self(inner))
228            }
229            Err(err) => Err(err),
230        }
231    }
232}
233
234/// A convenience wrapper around native and external ObjectStore instances.
235///
236/// Note that there may be performance differences between accepted variants here. If you wish to
237/// only permit the highest-performance stores, use [`PyObjectStore`] directly as the parameter in
238/// your signature.
239#[derive(FromPyObject)]
240pub enum AnyObjectStore {
241    /// A wrapper around a [`PyObjectStore`].
242    PyObjectStore(PyObjectStore),
243    /// A wrapper around a [`PyExternalObjectStore`].
244    PyExternalObjectStore(PyExternalObjectStore),
245}
246
247impl From<AnyObjectStore> for Arc<dyn ObjectStore> {
248    fn from(value: AnyObjectStore) -> Self {
249        match value {
250            AnyObjectStore::PyObjectStore(store) => store.into(),
251            AnyObjectStore::PyExternalObjectStore(store) => store.into(),
252        }
253    }
254}
255
256impl AnyObjectStore {
257    /// Consume self and return a reference-counted [`ObjectStore`].
258    pub fn into_dyn(self) -> Arc<dyn ObjectStore> {
259        self.into()
260    }
261}
pyo3_object_store/store.rs

pyo3_object_store/
store.rs
