diff --git a/Cargo.lock b/Cargo.lock index 10e2a283..2150dfc9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1643,6 +1643,7 @@ dependencies = [ [[package]] name = "pyo3-object_store" version = "0.9.0" +source = "git+https://github.com/developmentseed/obstore?rev=20010a52a1eaa95e5873c181f08a1cf768a97df4#20010a52a1eaa95e5873c181f08a1cf768a97df4" dependencies = [ "async-trait", "bytes", diff --git a/obstore/Cargo.toml b/obstore/Cargo.toml index 74404fb1..8c9c612d 100644 --- a/obstore/Cargo.toml +++ b/obstore/Cargo.toml @@ -35,7 +35,8 @@ pyo3-arrow = "0.16" pyo3-async-runtimes = { workspace = true, features = ["tokio-runtime"] } pyo3-bytes = "0.6" pyo3-file = { workspace = true } -pyo3-object_store = { path = "../pyo3-object_store" } +# pyo3-object_store = { path = "../pyo3-object_store" } +pyo3-object_store = { git = "https://github.com/developmentseed/obstore", rev = "20010a52a1eaa95e5873c181f08a1cf768a97df4" } tokio = { workspace = true, features = [ "macros", "rt", diff --git a/pyo3-object_store/Cargo.toml b/pyo3-object_store/Cargo.toml index 826e257b..4da01ca1 100644 --- a/pyo3-object_store/Cargo.toml +++ b/pyo3-object_store/Cargo.toml @@ -36,8 +36,8 @@ object_store = { version = "0.13.0", features = [ ] } # This is already an object_store dependency percent-encoding = "2.1" -pyo3 = { version = "0.28", features = ["chrono", "indexmap"] } -pyo3-async-runtimes = { version = "0.28", features = ["tokio-runtime"] } +pyo3 = { version = "0.27", features = ["chrono", "indexmap"] } +pyo3-async-runtimes = { version = "0.27", features = ["tokio-runtime"] } serde = "1" thiserror = "1" tokio = { version = "1.40", features = ["rt-multi-thread"] } diff --git a/pyo3-object_store/src/aws/credentials.rs b/pyo3-object_store/src/aws/credentials.rs index 11dcc0b4..a4aff1e9 100644 --- a/pyo3-object_store/src/aws/credentials.rs +++ b/pyo3-object_store/src/aws/credentials.rs @@ -53,6 +53,7 @@ impl<'py> FromPyObject<'_, 'py> for PyAwsCredential { } } +/// A Python-facing wrapper around a user-provided callback for AWS credential management. // TODO: don't use a cache for static credentials where `expires_at` is `None` // (so you don't need to access a mutex) #[derive(Debug)] diff --git a/pyo3-object_store/src/aws/mod.rs b/pyo3-object_store/src/aws/mod.rs index a4871db1..98c15b1c 100644 --- a/pyo3-object_store/src/aws/mod.rs +++ b/pyo3-object_store/src/aws/mod.rs @@ -1,4 +1,7 @@ +//! AWS S3 object store binding. + mod credentials; mod store; -pub use store::PyS3Store; +pub use credentials::PyAWSCredentialProvider; +pub use store::{PyAmazonS3Config, PyAmazonS3ConfigKey, PyS3Store}; diff --git a/pyo3-object_store/src/aws/store.rs b/pyo3-object_store/src/aws/store.rs index f8041c4a..ecd8b701 100644 --- a/pyo3-object_store/src/aws/store.rs +++ b/pyo3-object_store/src/aws/store.rs @@ -20,7 +20,7 @@ use crate::retry::PyRetryConfig; use crate::PyUrl; #[derive(Debug, Clone, PartialEq)] -struct S3Config { +pub struct S3Config { prefix: Option, config: PyAmazonS3Config, client_options: Option, @@ -29,7 +29,8 @@ struct S3Config { } impl S3Config { - fn bucket(&self) -> &str { + /// Access the bucket name for this config. + pub fn bucket(&self) -> &str { self.config .0 .get(&PyAmazonS3ConfigKey(AmazonS3ConfigKey::Bucket)) @@ -37,6 +38,33 @@ impl S3Config { .as_ref() } + /// Access the prefix for this config, if it exists. + pub fn prefix(&self) -> Option<&PyPath> { + self.prefix.as_ref() + } + + /// Access the config key-value pairs for this config. + /// + /// Note that the bucket **is included** in the returned config. + pub fn config(&self) -> &PyAmazonS3Config { + &self.config + } + + /// Access the client options for this config, if they exist. + pub fn client_options(&self) -> Option<&PyClientOptions> { + self.client_options.as_ref() + } + + /// Access the retry config for this config, if it exists. + pub fn retry_config(&self) -> Option<&PyRetryConfig> { + self.retry_config.as_ref() + } + + /// Access the credential provider for this config, if it exists. + pub fn credential_provider(&self) -> Option<&PyAWSCredentialProvider> { + self.credential_provider.as_ref() + } + fn __getnewargs_ex__<'py>(&'py self, py: Python<'py>) -> PyResult> { let args = PyTuple::empty(py).into_bound_py_any(py)?; let kwargs = PyDict::new(py); @@ -79,6 +107,11 @@ impl PyS3Store { pub fn into_inner(self) -> Arc> { self.store } + + /// Access the config for this store. + pub fn config(&self) -> &S3Config { + &self.config + } } #[pymethods] @@ -201,8 +234,9 @@ impl PyS3Store { self.config.prefix.as_ref() } + #[pyo3(name = "config")] #[getter] - fn config(&self) -> &PyAmazonS3Config { + fn py_config(&self) -> &PyAmazonS3Config { &self.config.config } @@ -222,6 +256,7 @@ impl PyS3Store { } } +/// A Python-facing wrapper around a config key for S3 configuration. #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub struct PyAmazonS3ConfigKey(AmazonS3ConfigKey); @@ -235,6 +270,12 @@ impl<'py> FromPyObject<'_, 'py> for PyAmazonS3ConfigKey { } } +impl AsRef for PyAmazonS3ConfigKey { + fn as_ref(&self) -> &AmazonS3ConfigKey { + &self.0 + } +} + impl AsRef for PyAmazonS3ConfigKey { fn as_ref(&self) -> &str { self.0.as_ref() @@ -278,9 +319,16 @@ impl From for AmazonS3ConfigKey { } } +/// A Python-facing wrapper around a set of S3 configuration key-value pairs. #[derive(Clone, Debug, Default, PartialEq, Eq, IntoPyObject, IntoPyObjectRef)] pub struct PyAmazonS3Config(HashMap); +impl AsRef> for PyAmazonS3Config { + fn as_ref(&self) -> &HashMap { + &self.0 + } +} + // Note: we manually impl FromPyObject instead of deriving it so that we can raise an // UnknownConfigurationKeyError instead of a `TypeError` on invalid config keys. // diff --git a/pyo3-object_store/src/azure/credentials.rs b/pyo3-object_store/src/azure/credentials.rs index 48890fab..471d07db 100644 --- a/pyo3-object_store/src/azure/credentials.rs +++ b/pyo3-object_store/src/azure/credentials.rs @@ -17,7 +17,8 @@ use crate::credentials::{is_awaitable, TemporaryToken, TokenCache}; use crate::path::PyPath; use crate::PyObjectStoreError; -struct PyAzureAccessKey { +/// A wrapper around an [AzureAccessKey] +pub struct PyAzureAccessKey { access_key: AzureAccessKey, expires_at: Option>, } @@ -40,7 +41,8 @@ impl<'py> FromPyObject<'_, 'py> for PyAzureAccessKey { } } -struct PyAzureSASToken { +/// A wrapper around a SAS token, which is a list of key-value pairs, and an optional expiry timestamp. +pub struct PyAzureSASToken { sas_token: Vec<(String, String)>, expires_at: Option>, } @@ -70,7 +72,8 @@ impl<'py> FromPyObject<'_, 'py> for PyAzureSASToken { } } -struct PyBearerToken { +/// A wrapper around a bearer token +pub struct PyBearerToken { token: String, expires_at: Option>, } @@ -86,10 +89,16 @@ impl<'py> FromPyObject<'_, 'py> for PyBearerToken { } } +/// A Python-facing enum wrapper around different Azure credential types. #[derive(FromPyObject)] -enum PyAzureCredential { +pub enum PyAzureCredential { + /// An access key credential AccessKey(PyAzureAccessKey), + + /// A SAS token credential SASToken(PyAzureSASToken), + + /// A bearer token credential BearerToken(PyBearerToken), } @@ -141,6 +150,7 @@ fn split_sas(sas: &str) -> Result, object_store::Error> { Ok(pairs) } +/// A Python-facing wrapper around a user-provided credential provider callback #[derive(Debug)] pub struct PyAzureCredentialProvider { /// The provided user callback to manage credential refresh diff --git a/pyo3-object_store/src/azure/mod.rs b/pyo3-object_store/src/azure/mod.rs index b413a914..689dcd88 100644 --- a/pyo3-object_store/src/azure/mod.rs +++ b/pyo3-object_store/src/azure/mod.rs @@ -1,5 +1,10 @@ +//! Azure Object Store Python bindings. + mod credentials; mod error; mod store; -pub use store::PyAzureStore; +pub use credentials::{ + PyAzureAccessKey, PyAzureCredential, PyAzureCredentialProvider, PyAzureSASToken, PyBearerToken, +}; +pub use store::{PyAzureConfig, PyAzureConfigKey, PyAzureStore}; diff --git a/pyo3-object_store/src/azure/store.rs b/pyo3-object_store/src/azure/store.rs index 97420416..2069cdab 100644 --- a/pyo3-object_store/src/azure/store.rs +++ b/pyo3-object_store/src/azure/store.rs @@ -18,7 +18,7 @@ use crate::retry::PyRetryConfig; use crate::{MaybePrefixedStore, PyUrl}; #[derive(Debug, Clone, PartialEq)] -struct AzureConfig { +pub struct AzureConfig { prefix: Option, config: PyAzureConfig, client_options: Option, @@ -43,6 +43,33 @@ impl AzureConfig { .as_ref() } + /// Access the prefix for this config, if it exists. + pub fn prefix(&self) -> Option<&PyPath> { + self.prefix.as_ref() + } + + /// Access the config key-value pairs for this config. + /// + /// Note that the account name and container name **are included** in the returned config. + pub fn config(&self) -> &PyAzureConfig { + &self.config + } + + /// Access the client options for this config, if they exist. + pub fn client_options(&self) -> Option<&PyClientOptions> { + self.client_options.as_ref() + } + + /// Access the retry config for this config, if it exists. + pub fn retry_config(&self) -> Option<&PyRetryConfig> { + self.retry_config.as_ref() + } + + /// Access the credential provider for this config, if it exists. + pub fn credential_provider(&self) -> Option<&PyAzureCredentialProvider> { + self.credential_provider.as_ref() + } + fn __getnewargs_ex__<'py>(&'py self, py: Python<'py>) -> PyResult> { let args = PyTuple::empty(py).into_bound_py_any(py)?; let kwargs = PyDict::new(py); @@ -85,6 +112,11 @@ impl PyAzureStore { pub fn into_inner(self) -> Arc> { self.store } + + /// Access the config for this store. + pub fn config(&self) -> &AzureConfig { + &self.config + } } #[pymethods] @@ -223,8 +255,9 @@ impl PyAzureStore { self.config.prefix.as_ref() } + #[pyo3(name = "config")] #[getter] - fn config(&self) -> &PyAzureConfig { + fn py_config(&self) -> &PyAzureConfig { &self.config.config } @@ -244,6 +277,7 @@ impl PyAzureStore { } } +/// A Python-facing wrapper around a config key for Azure configuration. #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub struct PyAzureConfigKey(AzureConfigKey); @@ -304,6 +338,7 @@ impl From for AzureConfigKey { } } +/// A Python-facing wrapper around a config for Azure configuration. #[derive(Clone, Debug, Default, PartialEq, Eq, IntoPyObject, IntoPyObjectRef)] pub struct PyAzureConfig(HashMap); diff --git a/pyo3-object_store/src/gcp/credentials.rs b/pyo3-object_store/src/gcp/credentials.rs index 06c8c1cb..4f86a333 100644 --- a/pyo3-object_store/src/gcp/credentials.rs +++ b/pyo3-object_store/src/gcp/credentials.rs @@ -43,6 +43,7 @@ impl<'py> FromPyObject<'_, 'py> for PyGcpCredential { } } +/// A Python-facing wrapper around a user-provided callback for GCP credential management. // TODO: don't use a cache for static credentials where `expires_at` is `None` // (so you don't need to access a mutex) #[derive(Debug)] diff --git a/pyo3-object_store/src/gcp/mod.rs b/pyo3-object_store/src/gcp/mod.rs index 78774709..64283959 100644 --- a/pyo3-object_store/src/gcp/mod.rs +++ b/pyo3-object_store/src/gcp/mod.rs @@ -1,4 +1,7 @@ +//! Google Cloud Storage Object Store Python bindings. + mod credentials; mod store; -pub use store::PyGCSStore; +pub use credentials::PyGcpCredentialProvider; +pub use store::{PyGCSStore, PyGoogleConfig, PyGoogleConfigKey}; diff --git a/pyo3-object_store/src/gcp/store.rs b/pyo3-object_store/src/gcp/store.rs index 260708df..fd386a4b 100644 --- a/pyo3-object_store/src/gcp/store.rs +++ b/pyo3-object_store/src/gcp/store.rs @@ -18,7 +18,7 @@ use crate::retry::PyRetryConfig; use crate::{MaybePrefixedStore, PyUrl}; #[derive(Debug, Clone, PartialEq)] -struct GCSConfig { +pub struct GCSConfig { prefix: Option, config: PyGoogleConfig, client_options: Option, @@ -27,7 +27,8 @@ struct GCSConfig { } impl GCSConfig { - fn bucket(&self) -> &str { + /// Access the bucket name for this config. + pub fn bucket(&self) -> &str { self.config .0 .get(&PyGoogleConfigKey(GoogleConfigKey::Bucket)) @@ -35,6 +36,33 @@ impl GCSConfig { .as_ref() } + /// Access the prefix for this config, if it exists. + pub fn prefix(&self) -> Option<&PyPath> { + self.prefix.as_ref() + } + + /// Access the config key-value pairs for this config. + /// + /// Note that the bucket **is included** in the returned config. + pub fn config(&self) -> &PyGoogleConfig { + &self.config + } + + /// Access the client options for this config, if they exist. + pub fn client_options(&self) -> Option<&PyClientOptions> { + self.client_options.as_ref() + } + + /// Access the retry config for this config, if it exists. + pub fn retry_config(&self) -> Option<&PyRetryConfig> { + self.retry_config.as_ref() + } + + /// Access the credential provider for this config, if it exists. + pub fn credential_provider(&self) -> Option<&PyGcpCredentialProvider> { + self.credential_provider.as_ref() + } + fn __getnewargs_ex__<'py>(&'py self, py: Python<'py>) -> PyResult> { let args = PyTuple::empty(py).into_bound_py_any(py)?; let kwargs = PyDict::new(py); @@ -77,6 +105,11 @@ impl PyGCSStore { pub fn into_inner(self) -> Arc> { self.store } + + /// Access the config for this store. + pub fn config(&self) -> &GCSConfig { + &self.config + } } #[pymethods] @@ -187,8 +220,9 @@ impl PyGCSStore { self.config.prefix.as_ref() } + #[pyo3(name = "config")] #[getter] - fn config(&self) -> &PyGoogleConfig { + fn py_config(&self) -> &PyGoogleConfig { &self.config.config } @@ -208,6 +242,7 @@ impl PyGCSStore { } } +/// A Python-facing wrapper around a config key for Google Cloud Storage. #[derive(Clone, Debug, PartialEq, Eq, Hash)] pub struct PyGoogleConfigKey(GoogleConfigKey); @@ -269,6 +304,7 @@ impl From for GoogleConfigKey { } } +/// A Python-facing wrapper around a config for Google Cloud Storage #[derive(Clone, Debug, Default, PartialEq, Eq, IntoPyObject, IntoPyObjectRef)] pub struct PyGoogleConfig(HashMap); diff --git a/pyo3-object_store/src/lib.rs b/pyo3-object_store/src/lib.rs index 0d7d7fdf..1a2bfdff 100644 --- a/pyo3-object_store/src/lib.rs +++ b/pyo3-object_store/src/lib.rs @@ -2,13 +2,13 @@ #![warn(missing_docs)] mod api; -mod aws; -mod azure; +pub mod aws; +pub mod azure; mod client; mod config; mod credentials; pub(crate) mod error; -mod gcp; +pub mod gcp; mod http; mod local; mod memory; @@ -31,5 +31,5 @@ pub use memory::PyMemoryStore; pub use path::PyPath; pub use prefix::MaybePrefixedStore; pub use simple::from_url; -pub use store::{AnyObjectStore, PyExternalObjectStore, PyObjectStore}; +pub use store::{AnyObjectStore, PyExternalObjectStore, PyObjectStore, PyTypedObjectStore}; pub use url::PyUrl; diff --git a/pyo3-object_store/src/store.rs b/pyo3-object_store/src/store.rs index 2b802f1e..d2b15796 100644 --- a/pyo3-object_store/src/store.rs +++ b/pyo3-object_store/src/store.rs @@ -48,7 +48,7 @@ impl<'py> FromPyObject<'_, 'py> for PyObjectStore { PyMemoryStore::type_object(py).name()?.to_str()?, PyS3Store::type_object(py).name()?.to_str()?, ] - .contains(&cls_name.as_str()) + .contains(&cls_name.as_ref()) { 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." )); } @@ -85,85 +85,158 @@ impl PyObjectStore { } } +/// A typed version of [`PyObjectStore`] that avoids type erasure of the underlying store #[derive(Debug, Clone)] -struct PyExternalObjectStoreInner(Arc); +pub enum PyTypedObjectStore { + /// A wrapper around a [`PyAzureStore`]. + Azure(PyAzureStore), + /// A wrapper around a [`PyGCSStore`]. + Gcs(PyGCSStore), + /// A wrapper around a [`PyHttpStore`]. + Http(PyHttpStore), + /// A wrapper around a [`PyLocalStore`]. + Local(PyLocalStore), + /// A wrapper around a [`PyS3Store`]. + S3(PyS3Store), +} -impl<'py> FromPyObject<'_, 'py> for PyExternalObjectStoreInner { - type Error = PyErr; +impl PyTypedObjectStore { + /// Attempt to extract an object store instance from an externally-defined (not statically + /// linked) ObjectStore instance. + pub fn from_external_store<'py>(obj: Borrowed<'_, 'py, pyo3::PyAny>) -> PyResult { + import_external_object_store(obj) + } +} - fn extract(obj: Borrowed<'_, 'py, pyo3::PyAny>) -> PyResult { - let py = obj.py(); - // Check for object-store instance from other library - let cls_name = obj - .getattr(intern!(py, "__class__"))? - .getattr(intern!(py, "__name__"))? - .extract::()?; - - if cls_name.as_str() == PyAzureStore::type_object(py).name()? { - let (args, kwargs): (Bound, Bound) = obj - .call_method0(intern!(py, "__getnewargs_ex__"))? - .extract()?; - let store = PyAzureStore::type_object(py) - .call(args, Some(&kwargs))? - .cast::()? - .get() - .clone(); - return Ok(Self(store.into_inner())); +impl From for PyObjectStore { + fn from(value: PyTypedObjectStore) -> Self { + match value { + PyTypedObjectStore::Azure(store) => PyObjectStore(store.into_inner()), + PyTypedObjectStore::Gcs(store) => PyObjectStore(store.into_inner()), + PyTypedObjectStore::Http(store) => PyObjectStore(store.into_inner()), + PyTypedObjectStore::Local(store) => PyObjectStore(store.into_inner()), + PyTypedObjectStore::S3(store) => PyObjectStore(store.into_inner()), } + } +} - if cls_name.as_str() == PyGCSStore::type_object(py).name()? { - let (args, kwargs): (Bound, Bound) = obj - .call_method0(intern!(py, "__getnewargs_ex__"))? - .extract()?; - let store = PyGCSStore::type_object(py) - .call(args, Some(&kwargs))? - .cast::()? - .get() - .clone(); - return Ok(Self(store.into_inner())); +impl From for Arc { + fn from(value: PyTypedObjectStore) -> Self { + match value { + PyTypedObjectStore::Azure(store) => store.into_inner(), + PyTypedObjectStore::Gcs(store) => store.into_inner(), + PyTypedObjectStore::Http(store) => store.into_inner(), + PyTypedObjectStore::Local(store) => store.into_inner(), + PyTypedObjectStore::S3(store) => store.into_inner(), } + } +} - if cls_name.as_str() == PyHttpStore::type_object(py).name()? { - let (args, kwargs): (Bound, Bound) = obj - .call_method0(intern!(py, "__getnewargs_ex__"))? - .extract()?; - let store = PyHttpStore::type_object(py) - .call(args, Some(&kwargs))? - .cast::()? - .get() - .clone(); - return Ok(Self(store.into_inner())); - } +fn raise_external_store_warning<'py>(py: Python<'py>) -> PyResult<()> { + let warnings_mod = py.import(intern!(py, "warnings"))?; + let warning = PyRuntimeWarning::new_err( + "Successfully reconstructed a store defined in another Python module. Connection pooling will not be shared across store instances.", + ); + let args = PyTuple::new(py, vec![warning])?; + warnings_mod.call_method1(intern!(py, "warn"), args)?; + Ok(()) +} - if cls_name.as_str() == PyLocalStore::type_object(py).name()? { - let (args, kwargs): (Bound, Bound) = obj - .call_method0(intern!(py, "__getnewargs_ex__"))? - .extract()?; - let store = PyLocalStore::type_object(py) - .call(args, Some(&kwargs))? - .cast::()? - .get() - .clone(); - return Ok(Self(store.into_inner())); - } +fn import_external_object_store<'py>( + obj: Borrowed<'_, 'py, pyo3::PyAny>, +) -> PyResult { + let py = obj.py(); - if cls_name.as_str() == PyS3Store::type_object(py).name()? { - let (args, kwargs): (Bound, Bound) = obj - .call_method0(intern!(py, "__getnewargs_ex__"))? - .extract()?; - let store = PyS3Store::type_object(py) - .call(args, Some(&kwargs))? - .cast::()? - .get() - .clone(); - return Ok(Self(store.into_inner())); - } + // Check for object-store instance from other library + let cls_name = obj + .getattr(intern!(py, "__class__"))? + .getattr(intern!(py, "__name__"))? + .extract::()?; - Err(PyValueError::new_err(format!( - "Expected an object store-compatible instance, got {}", - obj.repr()? - ))) + if cls_name.as_ref() == PyAzureStore::type_object(py).name()? { + let (args, kwargs): (Bound, Bound) = obj + .call_method0(intern!(py, "__getnewargs_ex__"))? + .extract()?; + let store = PyAzureStore::type_object(py) + .call(args, Some(&kwargs))? + .cast::()? + .get() + .clone(); + + #[cfg(feature = "external-store-warning")] + raise_external_store_warning(py)?; + + return Ok(PyTypedObjectStore::Azure(store)); } + + if cls_name.as_ref() == PyGCSStore::type_object(py).name()? { + let (args, kwargs): (Bound, Bound) = obj + .call_method0(intern!(py, "__getnewargs_ex__"))? + .extract()?; + let store = PyGCSStore::type_object(py) + .call(args, Some(&kwargs))? + .cast::()? + .get() + .clone(); + + #[cfg(feature = "external-store-warning")] + raise_external_store_warning(py)?; + + return Ok(PyTypedObjectStore::Gcs(store)); + } + + if cls_name.as_ref() == PyHttpStore::type_object(py).name()? { + let (args, kwargs): (Bound, Bound) = obj + .call_method0(intern!(py, "__getnewargs_ex__"))? + .extract()?; + let store = PyHttpStore::type_object(py) + .call(args, Some(&kwargs))? + .cast::()? + .get() + .clone(); + + #[cfg(feature = "external-store-warning")] + raise_external_store_warning(py)?; + + return Ok(PyTypedObjectStore::Http(store)); + } + + if cls_name.as_ref() == PyLocalStore::type_object(py).name()? { + let (args, kwargs): (Bound, Bound) = obj + .call_method0(intern!(py, "__getnewargs_ex__"))? + .extract()?; + let store = PyLocalStore::type_object(py) + .call(args, Some(&kwargs))? + .cast::()? + .get() + .clone(); + + #[cfg(feature = "external-store-warning")] + raise_external_store_warning(py)?; + + return Ok(PyTypedObjectStore::Local(store)); + } + + if cls_name.as_ref() == PyS3Store::type_object(py).name()? { + let (args, kwargs): (Bound, Bound) = obj + .call_method0(intern!(py, "__getnewargs_ex__"))? + .extract()?; + let store = PyS3Store::type_object(py) + .call(args, Some(&kwargs))? + .cast::()? + .get() + .clone(); + + #[cfg(feature = "external-store-warning")] + raise_external_store_warning(py)?; + + return Ok(PyTypedObjectStore::S3(store)); + } + + Err(PyValueError::new_err(format!( + "Expected an object store-compatible instance, got {}", + obj.repr()? + ))) } /// A wrapper around a Rust [ObjectStore] instance that will extract and recreate an ObjectStore @@ -192,11 +265,11 @@ impl<'py> FromPyObject<'_, 'py> for PyExternalObjectStoreInner { /// - This will not work for `PyMemoryStore` because we can't clone the internal state of the /// store. #[derive(Debug, Clone)] -pub struct PyExternalObjectStore(PyExternalObjectStoreInner); +pub struct PyExternalObjectStore(Arc); impl From for Arc { fn from(value: PyExternalObjectStore) -> Self { - value.0 .0 + value.0 } } @@ -211,23 +284,8 @@ impl<'py> FromPyObject<'_, 'py> for PyExternalObjectStore { type Error = PyErr; fn extract(obj: Borrowed<'_, 'py, pyo3::PyAny>) -> PyResult { - match obj.extract() { - Ok(inner) => { - #[cfg(feature = "external-store-warning")] - { - let py = obj.py(); - - let warnings_mod = py.import(intern!(py, "warnings"))?; - let warning = PyRuntimeWarning::new_err( - "Successfully reconstructed a store defined in another Python module. Connection pooling will not be shared across store instances.", - ); - let args = PyTuple::new(py, vec![warning])?; - warnings_mod.call_method1(intern!(py, "warn"), args)?; - } - Ok(Self(inner)) - } - Err(err) => Err(err), - } + let typed_store = import_external_object_store(obj)?; + Ok(Self(typed_store.into())) } }