diff --git a/changes/4168.bugfix.md b/changes/4168.bugfix.md new file mode 100644 index 0000000000..5415a07751 --- /dev/null +++ b/changes/4168.bugfix.md @@ -0,0 +1,5 @@ +Path-backed `ZipStore` now rejects pickling unless it was opened in read-only +mode. Pickling a writable store previously allowed multiple processes to reopen +the same ZIP archive for writing, so each writer could replace the archive's +central directory when it closed and corrupt the result. Read-only, path-backed +`ZipStore` instances remain pickleable. diff --git a/docs/user-guide/performance.md b/docs/user-guide/performance.md index 52c1cf0d71..e4c9c8e544 100644 --- a/docs/user-guide/performance.md +++ b/docs/user-guide/performance.md @@ -307,8 +307,13 @@ When writing to the same chunks from multiple processes, you should use external ## Pickle support -Zarr arrays and groups can be pickled, as long as the underlying store object can be -pickled. All of the storage classes provided in the `zarr.storage` module can be pickled. +Zarr arrays and groups can be pickled as long as the underlying store object can be +pickled. All path-backed storage classes provided in `zarr.storage` support pickling +for read-only workloads. A `ZipStore` opened in an archive-writing mode rejects +pickling because independent ZIP writers can corrupt the archive. A `ZipStore` backed +by an open file object also rejects pickling because it cannot reopen that object from +a path. For parallel writes, use a store such as `LocalStore` and create the ZIP archive +only after writing is complete. If an array or group is backed by a persistent store such as a `zarr.storage.LocalStore`, `zarr.storage.ZipStore` or `zarr.storage.FsspecStore` then the store data diff --git a/src/zarr/storage/_zip.py b/src/zarr/storage/_zip.py index 69ae18bc2c..b3daa15222 100644 --- a/src/zarr/storage/_zip.py +++ b/src/zarr/storage/_zip.py @@ -97,6 +97,14 @@ class ZipStore(Store): will raise an exception when the ZIP file would require ZIP64 extensions. + Notes + ----- + A ``ZipStore`` opened in an archive-writing mode cannot be pickled. ZIP + archives do not support independent concurrent writers, and reopening a + serialized writer can corrupt the archive. For parallel writes, use a + store that supports them and create the ZIP archive after writing is + complete. + Attributes ---------- allowed_exceptions @@ -182,6 +190,13 @@ async def _open(self) -> None: self._sync_open() def __getstate__(self) -> dict[str, Any]: + if self._zmode != "r": + raise TypeError( + "ZipStore instances opened in an archive-writing mode cannot be pickled, " + "because independent ZIP writers can corrupt the archive. Use a LocalStore " + "for parallel writes and create the ZIP archive after writing is complete, " + "or reopen the ZipStore with mode='r' before pickling." + ) if self.path is None: # A path-backed store pickles its path and reopens the file on # unpickling; an open file object cannot be serialized that way. diff --git a/src/zarr/testing/store.py b/src/zarr/testing/store.py index f64d8e9364..0fcb38cd0b 100644 --- a/src/zarr/testing/store.py +++ b/src/zarr/testing/store.py @@ -129,8 +129,8 @@ async def test_serializable_store(self, store: S) -> None: # quickly roundtrip data to a key to test that new store works data_buf = self.buffer_cls.from_bytes(b"\x01\x02\x03\x04") key = "foo" - await store.set(key, data_buf) - observed = await store.get(key, prototype=default_buffer_prototype()) + await new_store.set(key, data_buf) + observed = await new_store.get(key, prototype=default_buffer_prototype()) assert_bytes_equal(observed, data_buf) def test_store_read_only(self, store: S) -> None: diff --git a/tests/test_store/test_zip.py b/tests/test_store/test_zip.py index 32b18c5273..ba7508212b 100644 --- a/tests/test_store/test_zip.py +++ b/tests/test_store/test_zip.py @@ -85,6 +85,35 @@ def test_store_supports_writes(self, store: ZipStore) -> None: def test_store_supports_listing(self, store: ZipStore) -> None: assert store.supports_listing + async def test_serializable_store(self, store: ZipStore) -> None: + data = cpu.Buffer.from_bytes(b"preserve me") + await store.set("sentinel", data) + + with pytest.raises(TypeError, match="archive-writing mode cannot be pickled"): + pickle.dumps(store) + + store.close() + assert store.path is not None + with zipfile.ZipFile(store.path, mode="r") as archive: + assert archive.read("sentinel") == data.to_bytes() + + async def test_read_only_store_is_serializable(self, tmp_path: Path) -> None: + path = tmp_path / "data.zip" + data = cpu.Buffer.from_bytes(b"preserve me") + + writable = await ZipStore.open(path, mode="w") + await writable.set("sentinel", data) + writable.close() + + read_only = await ZipStore.open(path, mode="r") + restored = pickle.loads(pickle.dumps(read_only)) + observed = await restored.get("sentinel", prototype=default_buffer_prototype()) + + assert observed is not None + assert observed.to_bytes() == data.to_bytes() + read_only.close() + restored.close() + # TODO: fix this warning @pytest.mark.filterwarnings("ignore:Unclosed client session:ResourceWarning") def test_api_integration(self, store: ZipStore) -> None: