diff --git a/Cargo.lock b/Cargo.lock index 38f372e4..be94fddf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1422,12 +1422,11 @@ dependencies = [ [[package]] name = "object_store" version = "0.14.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d354792e39fa5f0009e47623cf8b15b099bf9a652fa55c6f817fe28ac84fea50" +source = "git+https://github.com/apache/arrow-rs-object-store?rev=978fb48b9ea6ebe3f8f4833f656e793cde79ab15#978fb48b9ea6ebe3f8f4833f656e793cde79ab15" dependencies = [ "async-trait", "aws-lc-rs", - "base64 0.22.1", + "base64 0.23.1", "bytes", "chrono", "crc-fast", diff --git a/Cargo.toml b/Cargo.toml index c1ff73ba..292011d5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -26,7 +26,7 @@ chrono = "0.4.44" futures = "0.3.31" http = "1.4" indexmap = "2" -object_store = { version = "0.14.0", features = [ +object_store = { git = "https://github.com/apache/arrow-rs-object-store", rev = "978fb48b9ea6ebe3f8f4833f656e793cde79ab15", features = [ "aws", "azure", "gcp", diff --git a/obstore/python/obstore/_store/__init__.pyi b/obstore/python/obstore/_store/__init__.pyi index 39eb60ec..620a62f9 100644 --- a/obstore/python/obstore/_store/__init__.pyi +++ b/obstore/python/obstore/_store/__init__.pyi @@ -185,6 +185,33 @@ class LocalStore: def __eq__(self, value: object) -> bool: ... def __getnewargs_ex__(self): ... + def replace_prefix(self, prefix: str | Path | None) -> Self: + """Construct a new store with a different prefix. + + All other configuration is inherited from this store. Note that unlike the + remote stores, there's no underlying connection pool to reuse here, so this is + equivalent to constructing a new [`LocalStore`][obstore.store.LocalStore] + directly. + + The new prefix fully replaces the existing one; it is not appended to it. + + **Example:** + + ```py + store = LocalStore(prefix="/data/2024") + store_2025 = store.replace_prefix("/data/2025") + ``` + + Args: + prefix: The prefix to apply to all operations in the new store. Pass `None` + to construct a store with no prefix. If `mkdir` is `True` on this store, + the directory at `prefix` will attempt to be created. + + Returns: + LocalStore + + """ + @property def prefix(self) -> Path | None: """Get the prefix applied to all operations in this store, if any.""" diff --git a/obstore/python/obstore/_store/_aws.pyi b/obstore/python/obstore/_store/_aws.pyi index 365b6d00..32eea198 100644 --- a/obstore/python/obstore/_store/_aws.pyi +++ b/obstore/python/obstore/_store/_aws.pyi @@ -620,6 +620,35 @@ class S3Store: def __eq__(self, value: object) -> bool: ... def __getnewargs_ex__(self): ... + def replace_prefix(self, prefix: str | None) -> Self: + """Construct a new store with a different prefix. + + The new store **shares the underlying HTTP client and connection pool** with + this store, so no new connections need to be established, and any already-open + connections stay warm. This makes it cheap to create many stores pointing at + different prefixes of the same bucket. + + All other configuration is inherited from this store. + + The new prefix fully replaces the existing one; it is not appended to it. It is + always interpreted relative to the root of the bucket. + + **Example:** + + ```py + store = S3Store("bucket", prefix="data/2024") + store_2025 = store.replace_prefix("data/2025") + ``` + + Args: + prefix: The prefix to apply to all operations in the new store. Pass `None` + to construct a store with no prefix. + + Returns: + S3Store + + """ + @property def prefix(self) -> str | None: """Get the prefix applied to all operations in this store, if any.""" diff --git a/obstore/python/obstore/_store/_azure.pyi b/obstore/python/obstore/_store/_azure.pyi index d0084172..46f00e6c 100644 --- a/obstore/python/obstore/_store/_azure.pyi +++ b/obstore/python/obstore/_store/_azure.pyi @@ -448,6 +448,35 @@ class AzureStore: def __eq__(self, value: object) -> bool: ... def __getnewargs_ex__(self): ... + def replace_prefix(self, prefix: str | None) -> Self: + """Construct a new store with a different prefix. + + The new store **shares the underlying HTTP client and connection pool** with + this store, so no new connections need to be established, and any already-open + connections stay warm. This makes it cheap to create many stores pointing at + different prefixes of the same container. + + All other configuration is inherited from this store. + + The new prefix fully replaces the existing one; it is not appended to it. It is + always interpreted relative to the root of the container. + + **Example:** + + ```py + store = AzureStore("container", prefix="data/2024") + store_2025 = store.replace_prefix("data/2025") + ``` + + Args: + prefix: The prefix to apply to all operations in the new store. Pass `None` + to construct a store with no prefix. + + Returns: + AzureStore + + """ + @property def prefix(self) -> str | None: """Get the prefix applied to all operations in this store, if any.""" diff --git a/obstore/python/obstore/_store/_gcs.pyi b/obstore/python/obstore/_store/_gcs.pyi index 2d8d9477..fa7a2b9f 100644 --- a/obstore/python/obstore/_store/_gcs.pyi +++ b/obstore/python/obstore/_store/_gcs.pyi @@ -226,6 +226,35 @@ class GCSStore: def __eq__(self, value: object) -> bool: ... def __getnewargs_ex__(self): ... + def replace_prefix(self, prefix: str | None) -> Self: + """Construct a new store with a different prefix. + + The new store **shares the underlying HTTP client and connection pool** with + this store, so no new connections need to be established, and any already-open + connections stay warm. This makes it cheap to create many stores pointing at + different prefixes of the same bucket. + + All other configuration is inherited from this store. + + The new prefix fully replaces the existing one; it is not appended to it. It is + always interpreted relative to the root of the bucket. + + **Example:** + + ```py + store = GCSStore("bucket", prefix="data/2024") + store_2025 = store.replace_prefix("data/2025") + ``` + + Args: + prefix: The prefix to apply to all operations in the new store. Pass `None` + to construct a store with no prefix. + + Returns: + GCSStore + + """ + @property def prefix(self) -> str | None: """Get the prefix applied to all operations in this store, if any.""" diff --git a/pyo3-object_store/Cargo.toml b/pyo3-object_store/Cargo.toml index 9c6577da..0205e12d 100644 --- a/pyo3-object_store/Cargo.toml +++ b/pyo3-object_store/Cargo.toml @@ -28,7 +28,7 @@ humantime = "2.1" http = "1" # This is already an object_store dependency itertools = "0.15.0" -object_store = { version = "0.14.0", features = [ +object_store = { git = "https://github.com/apache/arrow-rs-object-store", rev = "978fb48b9ea6ebe3f8f4833f656e793cde79ab15", features = [ "aws", "azure", "gcp", diff --git a/pyo3-object_store/src/aws/store.rs b/pyo3-object_store/src/aws/store.rs index f8041c4a..ef949efd 100644 --- a/pyo3-object_store/src/aws/store.rs +++ b/pyo3-object_store/src/aws/store.rs @@ -57,6 +57,16 @@ impl S3Config { PyTuple::new(py, [args, kwargs.into_bound_py_any(py)?]) } + + fn replace_prefix(&self, new_prefix: Option) -> Self { + Self { + prefix: new_prefix, + config: self.config.clone(), + client_options: self.client_options.clone(), + retry_config: self.retry_config.clone(), + credential_provider: self.credential_provider.clone(), + } + } } /// A Python-facing wrapper around an [`AmazonS3`]. @@ -216,6 +226,13 @@ impl PyS3Store { self.config.credential_provider.as_ref() } + fn replace_prefix(&self, prefix: Option) -> PyObjectStoreResult { + Ok(Self { + store: Arc::new(self.store.replace_prefix(prefix.clone())), + config: self.config.replace_prefix(prefix), + }) + } + #[getter] fn retry_config(&self) -> Option<&PyRetryConfig> { self.config.retry_config.as_ref() diff --git a/pyo3-object_store/src/azure/store.rs b/pyo3-object_store/src/azure/store.rs index d572b53c..0df27a49 100644 --- a/pyo3-object_store/src/azure/store.rs +++ b/pyo3-object_store/src/azure/store.rs @@ -63,6 +63,16 @@ impl AzureConfig { PyTuple::new(py, [args, kwargs.into_bound_py_any(py)?]) } + + fn replace_prefix(&self, new_prefix: Option) -> Self { + Self { + prefix: new_prefix, + config: self.config.clone(), + client_options: self.client_options.clone(), + retry_config: self.retry_config.clone(), + credential_provider: self.credential_provider.clone(), + } + } } /// A Python-facing wrapper around a [`MicrosoftAzure`]. @@ -238,6 +248,13 @@ impl PyAzureStore { self.config.credential_provider.as_ref() } + fn replace_prefix(&self, prefix: Option) -> PyObjectStoreResult { + Ok(Self { + store: Arc::new(self.store.replace_prefix(prefix.clone())), + config: self.config.replace_prefix(prefix), + }) + } + #[getter] fn retry_config(&self) -> Option<&PyRetryConfig> { self.config.retry_config.as_ref() diff --git a/pyo3-object_store/src/gcp/store.rs b/pyo3-object_store/src/gcp/store.rs index 260708df..99a38a1a 100644 --- a/pyo3-object_store/src/gcp/store.rs +++ b/pyo3-object_store/src/gcp/store.rs @@ -55,6 +55,16 @@ impl GCSConfig { PyTuple::new(py, [args, kwargs.into_bound_py_any(py)?]) } + + fn replace_prefix(&self, new_prefix: Option) -> Self { + Self { + prefix: new_prefix, + config: self.config.clone(), + client_options: self.client_options.clone(), + retry_config: self.retry_config.clone(), + credential_provider: self.credential_provider.clone(), + } + } } /// A Python-facing wrapper around a [`GoogleCloudStorage`]. @@ -202,6 +212,13 @@ impl PyGCSStore { self.config.credential_provider.as_ref() } + fn replace_prefix(&self, prefix: Option) -> PyObjectStoreResult { + Ok(Self { + store: Arc::new(self.store.replace_prefix(prefix.clone())), + config: self.config.replace_prefix(prefix), + }) + } + #[getter] fn retry_config(&self) -> Option<&PyRetryConfig> { self.config.retry_config.as_ref() diff --git a/pyo3-object_store/src/local.rs b/pyo3-object_store/src/local.rs index fc36df91..61d40c72 100644 --- a/pyo3-object_store/src/local.rs +++ b/pyo3-object_store/src/local.rs @@ -140,4 +140,11 @@ impl PyLocalStore { py.None().into_bound_py_any(py) } } + + fn replace_prefix(&self, prefix: Option) -> PyObjectStoreResult { + // Here we use Self::new instead of `replace_prefix` as on the other stores because 1) this + // doesn't use a MaybePrefixedStore wrapper and 2) there's no underlying connection pool we + // need to reuse. + Self::new(prefix, self.config.automatic_cleanup, self.config.mkdir) + } } diff --git a/pyo3-object_store/src/prefix.rs b/pyo3-object_store/src/prefix.rs index 78f3fc1a..2e6739c7 100644 --- a/pyo3-object_store/src/prefix.rs +++ b/pyo3-object_store/src/prefix.rs @@ -88,6 +88,16 @@ impl MaybePrefixedStore { } } +impl MaybePrefixedStore { + /// Create a new instance of [`MaybePrefixedStore`] with a new prefix + pub fn replace_prefix(&self, new_prefix: Option>) -> Self { + Self { + prefix: new_prefix.map(|x| x.into()), + inner: self.inner.clone(), + } + } +} + // Note: This is a relative hack to move these two functions to pure functions so they don't rely // on the `self` lifetime. Expected to be cleaned up before merge. // diff --git a/tests/store/test_replace_prefix.py b/tests/store/test_replace_prefix.py new file mode 100644 index 00000000..54d3f3f5 --- /dev/null +++ b/tests/store/test_replace_prefix.py @@ -0,0 +1,198 @@ +"""Tests for `Store.replace_prefix` across each store backend.""" + +from __future__ import annotations + +import pickle +from pathlib import Path +from typing import TYPE_CHECKING + +import pytest + +from obstore.store import AzureStore, GCSStore, LocalStore, S3Store + +if TYPE_CHECKING: + from obstore.store import ClientConfig, S3Config + + +def remote_stores() -> list[S3Store | AzureStore | GCSStore]: + """One store per remote backend, each configured with a prefix.""" + return [ + S3Store( + "bucket", + prefix="data/2024", + region="us-east-1", + skip_signature=True, + client_options={"timeout": "10s"}, + retry_config={"max_retries": 5}, + ), + AzureStore( + "container", + prefix="data/2024", + account_name="account", + skip_signature=True, + client_options={"timeout": "10s"}, + retry_config={"max_retries": 5}, + ), + GCSStore( + "bucket", + prefix="data/2024", + skip_signature=True, + client_options={"timeout": "10s"}, + retry_config={"max_retries": 5}, + ), + ] + + +@pytest.fixture(params=remote_stores(), ids=lambda store: type(store).__name__) +def remote_store(request: pytest.FixtureRequest) -> S3Store | AzureStore | GCSStore: + return request.param + + +def test_replaces_prefix(remote_store: S3Store | AzureStore | GCSStore): + assert remote_store.replace_prefix("data/2025").prefix == "data/2025" + + +def test_replaces_prefix_rather_than_appending( + remote_store: S3Store | AzureStore | GCSStore, +): + twice = remote_store.replace_prefix("data/2025").replace_prefix("data/2026") + assert twice.prefix == "data/2026" + + +def test_none_clears_prefix(remote_store: S3Store | AzureStore | GCSStore): + assert remote_store.replace_prefix(None).prefix is None + + +def test_original_store_is_unchanged(remote_store: S3Store | AzureStore | GCSStore): + remote_store.replace_prefix("data/2025") + assert remote_store.prefix == "data/2024" + + +def test_other_config_is_inherited(remote_store: S3Store | AzureStore | GCSStore): + new_store = remote_store.replace_prefix("data/2025") + assert new_store.config == remote_store.config + assert new_store.client_options == remote_store.client_options + assert new_store.retry_config == remote_store.retry_config + + +def test_eq_matches_a_directly_constructed_store( + remote_store: S3Store | AzureStore | GCSStore, +): + new_store = remote_store.replace_prefix("data/2025") + directly = type(remote_store)( + prefix="data/2025", + config=remote_store.config, # type: ignore[arg-type] + client_options=remote_store.client_options, + retry_config=remote_store.retry_config, + ) + assert new_store == directly + + +def test_pickle_round_trip(remote_store: S3Store | AzureStore | GCSStore): + """The pickling config must stay in sync with the underlying store's prefix.""" + new_store = remote_store.replace_prefix("data/2025") + restored = pickle.loads(pickle.dumps(new_store)) + assert restored.prefix == "data/2025" + assert restored == new_store + + +def test_preserves_subclass(remote_store: S3Store | AzureStore | GCSStore): + init_calls = [] + + class Subclass(type(remote_store)): # type: ignore[misc] + def __init__(self, *_args: object, **_kwargs: object) -> None: + init_calls.append(1) + + store = Subclass( + prefix="data/2024", + config=remote_store.config, # type: ignore[arg-type] + client_options=remote_store.client_options, + retry_config=remote_store.retry_config, + ) + assert len(init_calls) == 1 + + new_store = store.replace_prefix("data/2025") + assert type(new_store) is Subclass + assert new_store.prefix == "data/2025" + assert new_store.config == store.config + # `__init__` is not re-run, just as unpickling a store does not call it. + assert len(init_calls) == 1 + + +def test_local_replaces_prefix(tmp_path: Path): + (tmp_path / "2024").mkdir() + (tmp_path / "2025").mkdir() + + store = LocalStore(tmp_path / "2024", automatic_cleanup=True) + new_store = store.replace_prefix(tmp_path / "2025") + + assert new_store.prefix == tmp_path / "2025" + assert isinstance(new_store.prefix, Path) + # The original is untouched + assert store.prefix == tmp_path / "2024" + # And the rest of the config is inherited + assert new_store == LocalStore(tmp_path / "2025", automatic_cleanup=True) + + +def test_local_none_clears_prefix(tmp_path: Path): + assert LocalStore(tmp_path).replace_prefix(None).prefix is None + + +def test_local_mkdir_is_inherited(tmp_path: Path): + store = LocalStore(tmp_path / "2024", mkdir=True) + new_dir = tmp_path / "2025" + assert not new_dir.exists() + + store.replace_prefix(new_dir) + assert new_dir.exists() + + +def test_local_writes_to_the_new_prefix(tmp_path: Path): + store = LocalStore(tmp_path / "2024", mkdir=True) + new_store = store.replace_prefix(tmp_path / "2025") + new_store.put("afile.txt", b"hello world") + + assert (tmp_path / "2025" / "afile.txt").read_bytes() == b"hello world" + assert not (tmp_path / "2024" / "afile.txt").exists() + + +def test_local_pickle_round_trip(tmp_path: Path): + store = LocalStore(tmp_path / "2024", mkdir=True) + new_store = store.replace_prefix(tmp_path / "2025") + restored: LocalStore = pickle.loads(pickle.dumps(new_store)) + assert restored.prefix == tmp_path / "2025" + assert restored == new_store + + +def test_local_preserves_subclass(tmp_path: Path): + init_calls = [] + + class MyLocalStore(LocalStore): + def __init__(self, *_args: object, **_kwargs: object) -> None: + init_calls.append(1) + + store = MyLocalStore(tmp_path / "2024", mkdir=True) + assert len(init_calls) == 1 + + new_store = store.replace_prefix(tmp_path / "2025") + assert type(new_store) is MyLocalStore + assert new_store.prefix == tmp_path / "2025" + # `__init__` is not re-run, just as unpickling a store does not call it. + assert len(init_calls) == 1 + + +def test_writes_to_the_new_prefix(minio_bucket: tuple[S3Config, ClientConfig]): + """End-to-end: the underlying store, not just the config, is re-prefixed.""" + config, client_options = minio_bucket + store = S3Store(prefix="data/2024", config=config, client_options=client_options) + new_store = store.replace_prefix("data/2025") + + store.put("afile.txt", b"2024") + new_store.put("afile.txt", b"2025") + + unprefixed = S3Store(config=config, client_options=client_options) + assert unprefixed.get("data/2024/afile.txt").bytes() == b"2024" + assert unprefixed.get("data/2025/afile.txt").bytes() == b"2025" + + # Listing through the new store only sees the new prefix + assert [obj["path"] for obj in new_store.list().collect()] == ["afile.txt"]