Support pickling & always manage store prefix

docs/api/store/middleware.md

@@ -1,5 +1,3 @@
 # Middleware
 
 Wrappers around other `ObjectStore` instances to provide monitoring or other modifications.
-
-::: obstore.store.PrefixStore

docs/getting-started.md

@@ -16,10 +16,6 @@ Classes to construct a store are exported from the `obstore.store` submodule:
 - [`LocalStore`][obstore.store.LocalStore]: Local filesystem storage providing the same object store interface.
 - [`MemoryStore`][obstore.store.MemoryStore]: A fully in-memory implementation of ObjectStore.
 
-Additionally, some middlewares exist:
-
-- [`PrefixStore`][obstore.store.PrefixStore]: Store wrapper that applies a constant prefix to all paths handled by the store.
-
 Each store concept has a variety of constructors, and a host of configuration options.
 
 **Example:**

mkdocs.yml

@@ -35,7 +35,7 @@ nav:
           - api/store/local.md
           - api/store/memory.md
           - api/store/config.md
-          - api/store/middleware.md
+          # - api/store/middleware.md
       - api/copy.md
       - api/delete.md
       - api/get.md

obstore/python/obstore/store/__init__.pyi

@@ -9,7 +9,6 @@ from ._client import ClientConfig as ClientConfig
 from ._gcs import GCSConfig as GCSConfig
 from ._gcs import GCSStore as GCSStore
 from ._http import HTTPStore as HTTPStore
-from ._prefix import PrefixStore as PrefixStore
 from ._retry import BackoffConfig as BackoffConfig
 from ._retry import RetryConfig as RetryConfig
 
@@ -73,7 +72,5 @@ class MemoryStore:
     def __init__(self) -> None: ...
     def __repr__(self) -> str: ...
 
-ObjectStore = (
-    AzureStore | GCSStore | HTTPStore | S3Store | LocalStore | MemoryStore | PrefixStore
-)
+ObjectStore = AzureStore | GCSStore | HTTPStore | S3Store | LocalStore | MemoryStore
 """All supported ObjectStore implementations."""

obstore/python/obstore/store/_aws.pyi

@@ -521,6 +521,7 @@ class S3Store:
         self,
         bucket: str | None = None,
         *,
+        prefix: str | None = None,
         config: S3Config | None = None,
         client_options: ClientConfig | None = None,
         retry_config: RetryConfig | None = None,
@@ -532,6 +533,7 @@ class S3Store:
             bucket: The AWS bucket to use.
 
         Keyword Args:
+            prefix: A prefix within the bucket to use for all operations.
             config: AWS Configuration. Values in this config will override values inferred from the environment. Defaults to None.
             client_options: HTTP Client options. Defaults to None.
             retry_config: Retry configuration. Defaults to None.
@@ -544,8 +546,9 @@ class S3Store:
     def from_session(
         cls,
         session: boto3.session.Session | botocore.session.Session,
-        bucket: str,
+        bucket: str | None = None,
         *,
+        prefix: str | None = None,
         config: S3Config | None = None,
         client_options: ClientConfig | None = None,
         retry_config: RetryConfig | None = None,
@@ -574,6 +577,7 @@ class S3Store:
             bucket: The AWS bucket to use.
 
         Keyword Args:
+            prefix: A prefix within the bucket to use for all operations.
             config: AWS Configuration. Values in this config will override values inferred from the session. Defaults to None.
             client_options: HTTP Client options. Defaults to None.
             retry_config: Retry configuration. Defaults to None.
@@ -601,11 +605,6 @@ class S3Store:
         - `https://<bucket>.s3.<region>.amazonaws.com`
         - `https://ACCOUNT_ID.r2.cloudflarestorage.com/bucket`
 
-        !!! note
-            Note that `from_url` will not use any additional parts of the path as a
-            bucket prefix. It will only extract the bucket, region, and endpoint. If you
-            wish to use a path prefix, consider wrapping this with `PrefixStore`.
-
         Args:
             url: well-known storage URL.
 
@@ -619,4 +618,5 @@ class S3Store:
             S3Store
         """
 
+    def __getnewargs_ex__(self): ...
     def __repr__(self) -> str: ...

obstore/python/obstore/store/_azure.pyi

@@ -328,12 +328,6 @@ class AzureStore:
         - `https://<account>.blob.fabric.microsoft.com`
         - `https://<account>.blob.fabric.microsoft.com/<container>`
 
-        !!! note
-            Note that `from_url` will not use any additional parts of the path as a
-            bucket prefix. It will only extract the container name, account name, and
-            whether it's a fabric endpoint. If you wish to use a path prefix, consider
-            wrapping this with `PrefixStore`.
-
         Args:
             url: well-known storage URL.
 
@@ -346,4 +340,5 @@ class AzureStore:
             AzureStore
         """
 
+    def __getnewargs_ex__(self): ...
     def __repr__(self) -> str: ...

obstore/python/obstore/store/_gcs.pyi

@@ -120,11 +120,6 @@ class GCSStore:
 
         - `gs://<bucket>/<path>`
 
-        !!! note
-            Note that `from_url` will not use any additional parts of the path as a
-            bucket prefix. It will only extract the bucket name. If you wish to use a
-            path prefix, consider wrapping this with `PrefixStore`.
-
         Args:
             url: well-known storage URL.
 
@@ -137,4 +132,5 @@ class GCSStore:
             GCSStore
         """
 
+    def __getnewargs_ex__(self): ...
     def __repr__(self) -> str: ...

obstore/src/signer.rs

@@ -15,7 +15,8 @@ use pyo3::prelude::*;
 use pyo3::pybacked::PyBackedStr;
 use pyo3::types::PyString;
 use pyo3_object_store::{
-    PyAzureStore, PyGCSStore, PyObjectStoreError, PyObjectStoreResult, PyS3Store,
+    MaybePrefixedStore, PyAzureStore, PyGCSStore, PyObjectStoreError, PyObjectStoreResult,
+    PyS3Store,
 };
 use url::Url;
 
@@ -24,19 +25,19 @@ use crate::runtime::get_runtime;
 
 #[derive(Debug)]
 pub(crate) enum SignCapableStore {
-    S3(Arc<AmazonS3>),
-    Gcs(Arc<GoogleCloudStorage>),
-    Azure(Arc<MicrosoftAzure>),
+    S3(Arc<MaybePrefixedStore<AmazonS3>>),
+    Gcs(Arc<MaybePrefixedStore<GoogleCloudStorage>>),
+    Azure(Arc<MaybePrefixedStore<MicrosoftAzure>>),
 }
 
 impl<'py> FromPyObject<'py> for SignCapableStore {
     fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult<Self> {
         if let Ok(store) = ob.downcast::<PyS3Store>() {
-            Ok(Self::S3(store.borrow().as_ref().clone()))
+            Ok(Self::S3(store.get().as_ref().clone()))
         } else if let Ok(store) = ob.downcast::<PyGCSStore>() {
-            Ok(Self::Gcs(store.borrow().as_ref().clone()))
+            Ok(Self::Gcs(store.get().as_ref().clone()))
         } else if let Ok(store) = ob.downcast::<PyAzureStore>() {
-            Ok(Self::Azure(store.borrow().as_ref().clone()))
+            Ok(Self::Azure(store.get().as_ref().clone()))
         } else {
             let py = ob.py();
             // Check for object-store instance from other library
@@ -78,9 +79,9 @@ impl Signer for SignCapableStore {
         Self: 'async_trait,
     {
         match self {
-            Self::S3(inner) => inner.signed_url(method, path, expires_in),
-            Self::Gcs(inner) => inner.signed_url(method, path, expires_in),
-            Self::Azure(inner) => inner.signed_url(method, path, expires_in),
+            Self::S3(inner) => inner.as_ref().inner().signed_url(method, path, expires_in),
+            Self::Gcs(inner) => inner.as_ref().inner().signed_url(method, path, expires_in),
+            Self::Azure(inner) => inner.as_ref().inner().signed_url(method, path, expires_in),
         }
     }
 
@@ -96,9 +97,18 @@ impl Signer for SignCapableStore {
         Self: 'async_trait,
     {
         match self {
-            Self::S3(inner) => inner.signed_urls(method, paths, expires_in),
-            Self::Gcs(inner) => inner.signed_urls(method, paths, expires_in),
-            Self::Azure(inner) => inner.signed_urls(method, paths, expires_in),
+            Self::S3(inner) => inner
+                .as_ref()
+                .inner()
+                .signed_urls(method, paths, expires_in),
+            Self::Gcs(inner) => inner
+                .as_ref()
+                .inner()
+                .signed_urls(method, paths, expires_in),
+            Self::Azure(inner) => inner
+                .as_ref()
+                .inner()
+                .signed_urls(method, paths, expires_in),
         }
     }
 }

pyo3-object_store/Cargo.toml

@@ -14,6 +14,8 @@ rust-version = "1.75"
 include = ["src", "type-hints", "README.md", "LICENSE"]
 
 [dependencies]
+async-trait = "0.1.85"
+bytes = "1"
 futures = "0.3"
 # This is already an object_store dependency
 humantime = "2.1"
@@ -25,6 +27,7 @@ object_store = { version = "0.11.2", features = [
 ] }
 pyo3 = { version = "0.23", features = ["chrono", "indexmap"] }
 pyo3-async-runtimes = { version = "0.23", features = ["tokio-runtime"] }
+serde = "1"
 thiserror = "1"
 url = "2"
 

pyo3-object_store/src/api.rs

@@ -2,9 +2,7 @@ use pyo3::intern;
 use pyo3::prelude::*;
 
 use crate::error::*;
-use crate::{
-    PyAzureStore, PyGCSStore, PyHttpStore, PyLocalStore, PyMemoryStore, PyPrefixStore, PyS3Store,
-};
+use crate::{PyAzureStore, PyGCSStore, PyHttpStore, PyLocalStore, PyMemoryStore, PyS3Store};
 
 /// Export the default Python API as a submodule named `store` within the given parent module
 ///
@@ -52,7 +50,6 @@ pub fn register_store_module(
     child_module.add_class::<PyLocalStore>()?;
     child_module.add_class::<PyMemoryStore>()?;
     child_module.add_class::<PyS3Store>()?;
-    child_module.add_class::<PyPrefixStore>()?;
 
     parent_module.add_submodule(&child_module)?;