Prefix store (#117)

* Prefix store

* Add tests for prefix

* Bump python beta version

Author: kylebarron

Cargo.lock

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

docs/api/store/middleware.md

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

mkdocs.yml

@@ -1,6 +1,6 @@
 [package]
 name = "obstore"
-version = "0.3.0-beta.9"
+version = "0.3.0-beta.10"
 authors = { workspace = true }
 edition = { workspace = true }
 description = "A Python interface to the Rust object_store crate, providing a uniform API for interacting with object storage services and local files."

obstore/Cargo.toml

@@ -9,6 +9,7 @@ from ._client import ClientConfigKey as ClientConfigKey
 from ._gcs import GCSConfigKey as GCSConfigKey
 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
 
@@ -57,5 +58,7 @@ class MemoryStore:
     def __init__(self) -> None: ...
     def __repr__(self) -> str: ...
 
-ObjectStore = AzureStore | GCSStore | HTTPStore | S3Store | LocalStore | MemoryStore
+ObjectStore = (
+    AzureStore | GCSStore | HTTPStore | S3Store | LocalStore | MemoryStore | PrefixStore
+)
 """All supported ObjectStore implementations."""

obstore/python/obstore/store/__init__.pyi

@@ -0,0 +1,42 @@
+from obstore.store import ObjectStore
+
+class PrefixStore:
+    """Store wrapper that applies a constant prefix to all paths handled by the store.
+
+    **Example**:
+
+    ```py
+    import obstore as obs
+    from obstore.store import MemoryStore, PrefixStore
+
+    store = MemoryStore()
+
+    data = b"the quick brown fox jumps over the lazy dog"
+    path = "a/b/c/data.txt"
+
+    obs.put(store, path, data)
+
+    prefix_store = PrefixStore(store, "a/")
+    assert obs.get(prefix_store, "b/c/data.txt").bytes() == data
+
+    # The / after the passed-in prefix is inferred
+    prefix_store2 = PrefixStore(store, "a")
+    assert obs.get(prefix_store2, "b/c/data.txt").bytes() == data
+
+    # The prefix is removed from list results
+    assert obs.list(prefix_store).collect()[0]["path"] == "b/c/data.txt"
+
+    # More deeply nested prefix
+    prefix_store3 = PrefixStore(store, "a/b/c")
+    assert obs.get(prefix_store3, "data.txt").bytes() == data
+    ```
+    """
+    def __init__(self, store: ObjectStore, prefix: str) -> None:
+        """Create a new PrefixStore with the provided prefix.
+
+        Args:
+            store: The underlying store to wrap.
+            prefix: If the prefix does not end with `/`, one will be added.
+        """
+
+    def __repr__(self) -> str: ...

obstore/python/obstore/store/_prefix.pyi

@@ -2,7 +2,9 @@ use pyo3::intern;
 use pyo3::prelude::*;
 
 use crate::error::*;
-use crate::{PyAzureStore, PyGCSStore, PyHttpStore, PyLocalStore, PyMemoryStore, PyS3Store};
+use crate::{
+    PyAzureStore, PyGCSStore, PyHttpStore, PyLocalStore, PyMemoryStore, PyPrefixStore, PyS3Store,
+};
 
 /// Export the default Python API as a submodule named `store` within the given parent module
 ///
@@ -50,6 +52,7 @@ 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)?;
 

pyo3-object_store/src/api.rs

@@ -10,6 +10,7 @@ mod gcp;
 mod http;
 mod local;
 mod memory;
+mod prefix;
 mod retry;
 mod store;
 
@@ -22,4 +23,5 @@ pub use gcp::PyGCSStore;
 pub use http::PyHttpStore;
 pub use local::PyLocalStore;
 pub use memory::PyMemoryStore;
+pub use prefix::PyPrefixStore;
 pub use store::PyObjectStore;

pyo3-object_store/src/lib.rs

@@ -0,0 +1,30 @@
+use std::sync::Arc;
+
+use pyo3::prelude::*;
+
+use object_store::prefix::PrefixStore;
+use object_store::ObjectStore;
+
+use crate::PyObjectStore;
+
+/// A Python-facing wrapper around a [`PrefixStore`].
+#[pyclass(name = "PrefixStore", frozen)]
+pub struct PyPrefixStore(Arc<PrefixStore<Arc<dyn ObjectStore>>>);
+
+impl AsRef<Arc<PrefixStore<Arc<dyn ObjectStore>>>> for PyPrefixStore {
+    fn as_ref(&self) -> &Arc<PrefixStore<Arc<dyn ObjectStore>>> {
+        &self.0
+    }
+}
+
+#[pymethods]
+impl PyPrefixStore {
+    #[new]
+    fn new(store: PyObjectStore, prefix: String) -> Self {
+        Self(Arc::new(PrefixStore::new(store.into_inner(), prefix)))
+    }
+
+    fn __repr__(&self) -> String {
+        self.0.to_string()
+    }
+}

pyo3-object_store/src/prefix.rs

@@ -6,8 +6,9 @@ use pyo3::intern;
 use pyo3::prelude::*;
 use pyo3::pybacked::PyBackedStr;
 
-use crate::http::PyHttpStore;
-use crate::{PyAzureStore, PyGCSStore, PyLocalStore, PyMemoryStore, PyS3Store};
+use crate::{
+    PyAzureStore, PyGCSStore, PyHttpStore, PyLocalStore, PyMemoryStore, PyPrefixStore, PyS3Store,
+};
 
 /// A wrapper around a Rust ObjectStore instance that allows any rust-native implementation of
 /// ObjectStore.
@@ -29,6 +30,8 @@ impl<'py> FromPyObject<'py> for PyObjectStore {
             Ok(Self(store.get().as_ref().clone()))
         } else if let Ok(store) = ob.downcast::<PyMemoryStore>() {
             Ok(Self(store.get().as_ref().clone()))
+        } else if let Ok(store) = ob.downcast::<PyPrefixStore>() {
+            Ok(Self(store.get().as_ref().clone()))
         } else {
             let py = ob.py();
             // Check for object-store instance from other library
@@ -43,6 +46,7 @@ impl<'py> FromPyObject<'py> for PyObjectStore {
                 "LocalStore",
                 "MemoryStore",
                 "S3Store",
+                "PrefixStore",
             ]
             .contains(&cls_name.as_ref())
             {