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

pyo3_object_store/
store.rs
