fix(storage): open ZipStore archive on first use from every method and reopen it without truncating

changes/4450.bugfix.md

@@ -0,0 +1 @@
+`ZipStore` now opens its archive on first use from every method, so calling `get`, `get_partial_values`, `set_if_not_exists`, or `clear` on a store that was never opened no longer raises `AttributeError`. A store created with `mode="w"` or `mode="x"` now keeps its entries when it is reopened by `move()` or by unpickling, instead of truncating the archive or refusing to open it.

src/zarr/storage/_zip.py

@@ -162,22 +162,33 @@ def __init__(
         self._zmode = mode
         self.compression = compression
         self.allowZip64 = allowZip64
+        self._lock = threading.RLock()
 
     def _sync_open(self) -> None:
         if self._is_open:
             raise ValueError("store is already open")
 
-        self._lock = threading.RLock()
-
         self._zf = zipfile.ZipFile(
             self.path if self.path is not None else self._fileobj,  # type: ignore[arg-type]
             mode=self._zmode,
             compression=self.compression,
             allowZip64=self.allowZip64,
         )
+        # "w" truncates and "x" refuses an existing file. Both apply only to the
+        # first open: reopening after close(), move(), or unpickling must keep
+        # the entries already written.
+        if self._zmode in ("w", "x"):
+            self._zmode = "a"
 
         self._is_open = True
 
+    def _zipfile(self) -> zipfile.ZipFile:
+        """Return the archive, opening it on first use."""
+        with self._lock:
+            if not self._is_open:
+                self._sync_open()
+            return self._zf
+
     async def _open(self) -> None:
         self._sync_open()
 
@@ -197,6 +208,7 @@ def __getstate__(self) -> dict[str, Any]:
 
     def __setstate__(self, state: dict[str, Any]) -> None:
         self.__dict__ = state
+        self._lock = threading.RLock()
         self._is_open = False
         self._sync_open()
 
@@ -216,7 +228,7 @@ async def clear(self) -> None:
                 raise NotImplementedError(
                     "clear() is not supported for a ZipStore backed by a file-like object"
                 )
-            self._zf.close()
+            self._zipfile().close()
             os.remove(self.path)
             self._zf = zipfile.ZipFile(
                 self.path, mode="w", compression=self.compression, allowZip64=self.allowZip64
@@ -243,11 +255,9 @@ def _get(
         prototype: BufferPrototype,
         byte_range: ByteRequest | None = None,
     ) -> Buffer | None:
-        if not self._is_open:
-            self._sync_open()
         # docstring inherited
         try:
-            with self._zf.open(key) as f:  # will raise KeyError
+            with self._zipfile().open(key) as f:  # will raise KeyError
                 if byte_range is None:
                     return prototype.buffer.from_bytes(f.read())
                 elif isinstance(byte_range, RangeByteRequest):
@@ -288,8 +298,6 @@ async def get_partial_values(
         return out
 
     def _set(self, key: str, value: Buffer) -> None:
-        if not self._is_open:
-            self._sync_open()
         # generally, this should be called inside a lock
         keyinfo = zipfile.ZipInfo(filename=key, date_time=time.localtime(time.time())[:6])
         keyinfo.compress_type = self.compression
@@ -298,13 +306,11 @@ def _set(self, key: str, value: Buffer) -> None:
             keyinfo.external_attr |= 0x10  # MS-DOS directory flag
         else:
             keyinfo.external_attr = 0o644 << 16  # ?rw-r--r--
-        self._zf.writestr(keyinfo, value.to_bytes())
+        self._zipfile().writestr(keyinfo, value.to_bytes())
 
     async def set(self, key: str, value: Buffer) -> None:
         # docstring inherited
         self._check_writable()
-        if not self._is_open:
-            self._sync_open()
         if not isinstance(value, Buffer):
             raise TypeError(
                 f"ZipStore.set(): `value` must be a Buffer instance. Got an instance of {type(value)} instead."
@@ -315,7 +321,7 @@ async def set(self, key: str, value: Buffer) -> None:
     async def set_if_not_exists(self, key: str, value: Buffer) -> None:
         self._check_writable()
         with self._lock:
-            members = self._zf.namelist()
+            members = self._zipfile().namelist()
             if key not in members:
                 self._set(key, value)
 
@@ -337,22 +343,18 @@ async def delete(self, key: str) -> None:
 
     async def exists(self, key: str) -> bool:
         # docstring inherited
-        if not self._is_open:
-            self._sync_open()
         with self._lock:
             try:
-                self._zf.getinfo(key)
+                self._zipfile().getinfo(key)
             except KeyError:
                 return False
             else:
                 return True
 
     async def list(self) -> AsyncIterator[str]:
         # docstring inherited
-        if not self._is_open:
-            self._sync_open()
         with self._lock:
-            for key in self._zf.namelist():
+            for key in self._zipfile().namelist():
                 yield key
 
     async def list_prefix(self, prefix: str) -> AsyncIterator[str]:
@@ -363,11 +365,9 @@ async def list_prefix(self, prefix: str) -> AsyncIterator[str]:
 
     async def list_dir(self, prefix: str) -> AsyncIterator[str]:
         # docstring inherited
-        if not self._is_open:
-            self._sync_open()
         prefix = prefix.rstrip("/")
 
-        keys = self._zf.namelist()
+        keys = self._zipfile().namelist()
         seen = set()
         if prefix == "":
             keys_unique = {k.split("/")[0] for k in keys}

tests/test_store/test_zip.py

@@ -188,6 +188,76 @@ async def test_move(self, tmp_path: Path) -> None:
         assert not origin.exists()
         assert np.array_equal(array[...], np.arange(10))
 
+    @pytest.mark.parametrize(
+        "call",
+        [
+            lambda s: s.get("foo", default_buffer_prototype()),
+            lambda s: s.get_partial_values(default_buffer_prototype(), [("foo", None)]),
+            lambda s: s.exists("foo"),
+            lambda s: s.set("bar", cpu.Buffer.from_bytes(b"x")),
+            lambda s: s.set_if_not_exists("bar", cpu.Buffer.from_bytes(b"x")),
+            lambda s: s.clear(),
+            lambda s: s.delete("bar"),
+            lambda s: s.delete_dir("bar"),
+            lambda s: s.is_empty(""),
+        ],
+        ids=[
+            "get",
+            "get_partial_values",
+            "exists",
+            "set",
+            "set_if_not_exists",
+            "clear",
+            "delete",
+            "delete_dir",
+            "is_empty",
+        ],
+    )
+    async def test_methods_open_store_on_first_use(
+        self, store_kwargs: dict[str, Any], call: Any
+    ) -> None:
+        # every method works on a store that was constructed but never opened
+        seed = await self.store_cls.open(**store_kwargs)
+        await seed.set("foo", cpu.Buffer.from_bytes(b"bar"))
+        seed.close()
+
+        store = self.store_cls(**{**store_kwargs, "mode": "a"})
+        assert not store._is_open
+        await call(store)
+        store.close()
+
+    @pytest.mark.parametrize("method", ["list", "list_dir"])
+    async def test_listing_opens_store_on_first_use(
+        self, store_kwargs: dict[str, Any], method: str
+    ) -> None:
+        seed = await self.store_cls.open(**store_kwargs)
+        await seed.set("foo", cpu.Buffer.from_bytes(b"bar"))
+        seed.close()
+
+        store = self.store_cls(**{**store_kwargs, "mode": "r"})
+        args = () if method == "list" else ("",)
+        assert [k async for k in getattr(store, method)(*args)] == ["foo"]
+        store.close()
+
+    @pytest.mark.parametrize("mode", ["w", "x"])
+    @pytest.mark.parametrize("reopen", ["close_twice", "move", "pickle"])
+    async def test_reopen_keeps_entries(self, tmp_path: Path, mode: str, reopen: str) -> None:
+        # "w" truncates and "x" refuses an existing file; neither may apply
+        # when a store that already wrote entries is opened again
+        store = ZipStore(tmp_path / "data.zip", mode=mode)  # type: ignore[arg-type]
+        await store.set("foo", cpu.Buffer.from_bytes(b"bar"))
+        if reopen == "close_twice":
+            store.close()
+            store.close()
+        elif reopen == "move":
+            await store.move(tmp_path / "moved" / "data.zip")
+            store.close()
+        else:
+            store.close()
+            pickle.loads(pickle.dumps(store)).close()
+        with zipfile.ZipFile(store.path) as zf:  # type: ignore[arg-type]
+            assert zf.namelist() == ["foo"]
+
 
 class TestZipStoreFileObj:
     """ZipStore backed by an open binary file-like object instead of a path."""
@@ -352,8 +422,7 @@ class ZipStoreLifecycleMachine(RuleBasedStateMachine):
     Invariant under test: a constructed ZipStore can always be closed without
     raising, regardless of whether it was ever opened or did any I/O. This is a
     property-based generalization of the former example-based regression tests
-    for ZipStore.close() being called on a never-opened store (which raised
-    AttributeError because ``_lock`` is created lazily in ``_sync_open``).
+    for ZipStore.close() being called on a never-opened store.
     """
 
     def __init__(self, tmp_path: Path) -> None: