diff --git a/src/lib.rs b/src/lib.rs index 9ef05f2d..c53e73de 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -26,7 +26,7 @@ use zarrs::array::{ use zarrs::config::global_config; use zarrs::convert::array_metadata_v2_to_v3; use zarrs::plugin::ZarrVersion; -use zarrs::storage::{ReadableWritableListableStorage, StorageHandle, StoreKey}; +use zarrs::storage::{ReadableStorage, ReadableWritableListableStorage, StorageHandle, StoreKey}; mod chunk_item; mod concurrency; @@ -44,7 +44,12 @@ use crate::utils::{PyCodecErrExt, PyErrExt as _}; #[gen_stub_pyclass] #[pyclass] pub(crate) struct CodecPipelineImpl { - pub(crate) store: ReadableWritableListableStorage, + /// Read handle, deliberately the READ-ONLY type: `set`/`erase` are not in its interface, + /// so the only way to reach a write is `writable()`. + pub(crate) readable_store: ReadableStorage, + /// The writable handle -- `None` when zarr-python opened the store read-only. Same object + /// as `readable_store`; the read-only case simply never keeps a writable view of it. + pub(crate) writable_store: Option, pub(crate) codec_chain: Arc, pub(crate) codec_options: CodecOptions, pub(crate) chunk_concurrent_minimum: usize, @@ -61,7 +66,10 @@ impl CodecPipelineImpl { codec_chain: &CodecChain, codec_options: &CodecOptions, ) -> PyResult> { - let value_encoded = self.store.get(&item.key).map_py_err::()?; + let value_encoded = self + .readable_store + .get(&item.key) + .map_py_err::()?; let value_decoded = if let Some(value_encoded) = value_encoded { let value_encoded: Vec = value_encoded.into(); // zero-copy in this case codec_chain @@ -80,6 +88,13 @@ impl CodecPipelineImpl { Ok(value_decoded) } + /// The writable store, or zarr-python's own refusal, verbatim. + fn writable(&self) -> PyResult<&ReadableWritableListableStorage> { + self.writable_store.as_ref().ok_or_else(|| { + PyValueError::new_err("store was opened in read-only mode and does not support writing") + }) + } + fn store_chunk_bytes( &self, item: &ChunkItem, @@ -91,8 +106,10 @@ impl CodecPipelineImpl { .validate(item.num_elements, &self.data_type) .map_codec_err()?; + let store = self.writable()?; + if value_decoded.is_fill_value(&self.fill_value) { - self.store.erase(&item.key).map_py_err::() + store.erase(&item.key).map_py_err::() } else { let value_encoded = codec_chain .encode( @@ -106,7 +123,7 @@ impl CodecPipelineImpl { .map_codec_err()?; // Store the encoded chunk - self.store + store .set(&item.key, value_encoded.into()) .map_py_err::() } @@ -252,6 +269,8 @@ impl CodecPipelineImpl { let store: ReadableWritableListableStorage = (&store_config).try_into().map_py_err::()?; + let writable_store = (!store_config.read_only).then(|| store.clone()); + let readable_store: ReadableStorage = store.readable(); let data_type = DataType::from_metadata(&metadata_v3.data_type).map_py_err::()?; @@ -270,7 +289,7 @@ impl CodecPipelineImpl { .map_py_err::()?; Ok(Self { - store, + readable_store, codec_chain, codec_options, chunk_concurrent_minimum, @@ -278,6 +297,7 @@ impl CodecPipelineImpl { num_threads, fill_value, data_type, + writable_store, }) } @@ -308,7 +328,7 @@ impl CodecPipelineImpl { if !partial_chunk_items.is_empty() { let key_decoder_pairs = iter_concurrent_limit!(chunk_concurrent_limit, partial_chunk_items, map, |item| { - let storage_handle = Arc::new(StorageHandle::new(self.store.clone())); + let storage_handle = Arc::new(StorageHandle::new(self.readable_store.clone())); let input_handle = StoragePartialDecoder::new(storage_handle, item.key.clone()); let partial_decoder = self .codec_chain @@ -352,8 +372,10 @@ impl CodecPipelineImpl { // See zarrs::array::Array::retrieve_chunk_subset_into if is_whole_chunk(&item) { // See zarrs::array::Array::retrieve_chunk_into - if let Some(chunk_encoded) = - self.store.get(&item.key).map_py_err::()? + if let Some(chunk_encoded) = self + .readable_store + .get(&item.key) + .map_py_err::()? { // Decode the encoded data into the output buffer let chunk_encoded: Vec = chunk_encoded.into(); @@ -397,6 +419,9 @@ impl CodecPipelineImpl { value: &Bound<'_, PyUntypedArray>, write_empty_chunks: bool, ) -> PyResult<()> { + // Fail before decoding anything; the write site checks again by construction. + self.writable()?; + enum InputValue<'a> { Array(ArrayBytes<'a>), Constant(FillValue), diff --git a/src/store.rs b/src/store.rs index 0176fa85..b0ad2ff5 100644 --- a/src/store.rs +++ b/src/store.rs @@ -22,7 +22,17 @@ pub use self::http::HttpStoreConfig; pub use self::obstore::ObStoreConfig; #[derive(Debug, Clone)] -pub enum StoreConfig { +pub struct StoreConfig { + pub kind: StoreKind, + /// Whether zarr-python opened the store read-only, i.e. `mode="r"`. + /// + /// zarr-python resolves the mode into the store itself (`Store.read_only`), so it comes + /// off the same object every other field here does. + pub read_only: bool, +} + +#[derive(Debug, Clone)] +pub enum StoreKind { Filesystem(FilesystemStoreConfig), Http(HttpStoreConfig), ObStore(ObStoreConfig), @@ -35,10 +45,10 @@ impl<'py> FromPyObject<'_, 'py> for StoreConfig { fn extract(store: Borrowed<'_, 'py, PyAny>) -> PyResult { let name = store.get_type().name()?; let name = name.to_str()?; - match name { + let kind = match name { "LocalStore" => { let root: String = store.getattr("root")?.call_method0("__str__")?.extract()?; - Ok(StoreConfig::Filesystem(FilesystemStoreConfig::new(root))) + StoreKind::Filesystem(FilesystemStoreConfig::new(root)) } "FsspecStore" => { let fs = store.getattr("fs")?; @@ -48,13 +58,14 @@ impl<'py> FromPyObject<'_, 'py> for StoreConfig { let storage_options: HashMap> = fs.getattr("storage_options")?.extract()?; match fs_name { - "HTTPFileSystem" => Ok(StoreConfig::Http(HttpStoreConfig::new( - &path, - &storage_options, - )?)), - _ => Err(PyErr::new::(format!( - "zarrs-python does not support {fs_name} (FsspecStore) stores" - ))), + "HTTPFileSystem" => { + StoreKind::Http(HttpStoreConfig::new(&path, &storage_options)?) + } + _ => { + return Err(PyErr::new::(format!( + "zarrs-python does not support {fs_name} (FsspecStore) stores" + ))); + } } } "ObjectStore" => { @@ -62,29 +73,35 @@ impl<'py> FromPyObject<'_, 'py> for StoreConfig { let external_object_store: PyExternalObjectStore = underlying_store.extract()?; let object_store: Arc = external_object_store.into_dyn(); - Ok(StoreConfig::ObStore(ObStoreConfig::new(object_store))) + StoreKind::ObStore(ObStoreConfig::new(object_store)) } - _ => Err(PyErr::new::(format!( - "zarrs-python does not support {name} stores" - ))), - } + _ => { + return Err(PyErr::new::(format!( + "zarrs-python does not support {name} stores" + ))); + } + }; + Ok(StoreConfig { + kind, + read_only: store.getattr("read_only")?.extract()?, + }) } } impl StoreConfig { pub fn direct_io(&mut self, flag: bool) { - match self { - StoreConfig::Filesystem(config) => config.direct_io(flag), - StoreConfig::Http(_config) => (), - StoreConfig::ObStore(_config) => (), + match &mut self.kind { + StoreKind::Filesystem(config) => config.direct_io(flag), + StoreKind::Http(_config) => (), + StoreKind::ObStore(_config) => (), } } pub fn file_handle_cache_size(&mut self, size: usize) { - match self { - StoreConfig::Filesystem(config) => config.file_handle_cache_size(size), - StoreConfig::Http(_config) => (), - StoreConfig::ObStore(_config) => (), + match &mut self.kind { + StoreKind::Filesystem(config) => config.file_handle_cache_size(size), + StoreKind::Http(_config) => (), + StoreKind::ObStore(_config) => (), } } } @@ -99,10 +116,10 @@ impl TryFrom<&StoreConfig> for ReadableWritableListableStorage { type Error = PyErr; fn try_from(value: &StoreConfig) -> Result { - match value { - StoreConfig::Filesystem(config) => config.try_into(), - StoreConfig::Http(config) => config.try_into(), - StoreConfig::ObStore(config) => config.try_into(), + match &value.kind { + StoreKind::Filesystem(config) => config.try_into(), + StoreKind::Http(config) => config.try_into(), + StoreKind::ObStore(config) => config.try_into(), } } } diff --git a/tests/test_read_only_store.py b/tests/test_read_only_store.py new file mode 100644 index 00000000..bb58c0c6 --- /dev/null +++ b/tests/test_read_only_store.py @@ -0,0 +1,57 @@ +"""A store opened read-only must refuse writes, as zarr-python's own pipeline does. + +zarr-python enforces this in the store itself -- `Store._check_writable`, reached from the +concrete store's `_set`. This pipeline never gets there: it is handed a `StoreConfig` and +builds its own Rust store, so nothing zarr-python checks is in the path -- a `mode="r"` write +SUCCEEDS unless this side refuses it. It refuses by keeping a writable handle only when the +store is writable; the read handle has no `set`/`erase` in its interface at all. + +Opened STRICT throughout, and that is what makes the assertion mean anything: zarr's own +refusal message is byte-identical to the Rust one, so with a fallback available these +tests would pass whether the guard fired or zarr-python served the write. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np +import pytest +import zarr + +if TYPE_CHECKING: + from pathlib import Path + + +@pytest.fixture +def array(tmp_path: Path) -> tuple[Path, np.ndarray]: + values = np.arange(64, dtype=np.float32) + path = tmp_path / "a" + zarr.create_array(path, dtype=values.dtype, shape=values.shape, chunks=(16,))[:] = ( + values + ) + return path, values + + +def open_strict(path: Path, mode: str) -> zarr.Array: + """Open with no fallback. `strict` has to be set BEFORE the open: that is when the + pipeline decides whether it has one.""" + with zarr.config.set({"codec_pipeline.strict": True}): + return zarr.open_array(path, mode=mode) + + +def test_write_to_a_read_only_array_raises(array: tuple[Path, np.ndarray]) -> None: + path, values = array + z = open_strict(path, "r") + with pytest.raises(ValueError, match="read-only"): + z[0:16] = -1.0 + np.testing.assert_array_equal(zarr.open_array(path, mode="r")[:], values) + + +def test_a_writable_array_still_writes(array: tuple[Path, np.ndarray]) -> None: + path, values = array + z = open_strict(path, "r+") + z[0:16] = -1.0 + expected = values.copy() + expected[0:16] = -1.0 + np.testing.assert_array_equal(zarr.open_array(path, mode="r")[:], expected)