diff --git a/.github/workflows/zarr-storage-extraction.yml b/.github/workflows/zarr-storage-extraction.yml new file mode 100644 index 0000000000..9f8d9121dc --- /dev/null +++ b/.github/workflows/zarr-storage-extraction.yml @@ -0,0 +1,75 @@ +name: zarr-storage extraction + +on: + pull_request: + paths: + - 'packages/zarr-storage/**' + - 'src/zarr/abc/store.py' + - 'src/zarr/core/_coalesce.py' + - 'src/zarr/storage/**' + - 'src/zarr/testing/**' + - 'src/zarr/experimental/cache_store.py' + - 'tests/test_store/**' + - 'tests/test_experimental/test_cache_store.py' + - '.github/workflows/zarr-storage-extraction.yml' + push: + branches: [main] + paths: + - 'packages/zarr-storage/**' + - 'src/zarr/abc/store.py' + - 'src/zarr/core/_coalesce.py' + - 'src/zarr/storage/**' + - 'src/zarr/testing/**' + - 'src/zarr/experimental/cache_store.py' + - 'tests/test_store/**' + - 'tests/test_experimental/test_cache_store.py' + - '.github/workflows/zarr-storage-extraction.yml' + workflow_dispatch: + +permissions: + contents: read + +jobs: + test: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + include: + - python-version: '3.12' + dependencies: minimal + - python-version: '3.14' + dependencies: minimal + - python-version: '3.12' + dependencies: optional + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + with: + persist-credentials: false + - uses: astral-sh/setup-uv@20cfd1bf945f4377ade1205e4dbc17946fc9a30d # v10.0.1 + - name: Install Python and Hatch + run: | + uv python install ${{ matrix.python-version }} + uv tool install hatch==1.18.0 + - name: Test extracted interfaces against this checkout + env: + PYTHONPATH: packages/zarr-storage/src + run: >- + hatch run test.py${{ matrix.python-version }}-${{ matrix.dependencies }}:pytest + packages/zarr-storage/tests --import-mode=importlib + - name: Build storage distributions + working-directory: packages/zarr-storage + run: hatch build + - name: Test built wheels + env: + PYTHONPATH: packages/zarr-storage/dist/zarr_storage-0.1.0-py3-none-any.whl + run: >- + hatch run test.py${{ matrix.python-version }}-${{ matrix.dependencies }}:pytest + packages/zarr-storage/tests --import-mode=importlib + - name: Run stateful store tests + if: matrix.dependencies == 'optional' + env: + PYTHONPATH: packages/zarr-storage/dist/zarr_storage-0.1.0-py3-none-any.whl + run: >- + hatch run test.py${{ matrix.python-version }}-${{ matrix.dependencies }}:pytest + packages/zarr-storage/tests/test_store/test_stateful.py --run-slow-hypothesis diff --git a/packages/zarr-storage/CHANGELOG.md b/packages/zarr-storage/CHANGELOG.md new file mode 100644 index 0000000000..a6c45e7cff --- /dev/null +++ b/packages/zarr-storage/CHANGELOG.md @@ -0,0 +1,9 @@ +# zarr-storage changelog + +## Unreleased + +- Extract the existing storage interfaces, concrete implementations, and accessory + stores into `zarr_storage.legacy`, preserving Zarr's current runtime imports. +- Include the store and experimental cache-store suites and distribute reusable + conformance tests and stateful testing utilities in `zarr_storage.testing`. +- Make `LatencyStore` usable without importing pytest. diff --git a/packages/zarr-storage/LICENSE.txt b/packages/zarr-storage/LICENSE.txt new file mode 100644 index 0000000000..1e8da4d242 --- /dev/null +++ b/packages/zarr-storage/LICENSE.txt @@ -0,0 +1,21 @@ +The MIT License (MIT) + +Copyright (c) 2015-2025 Zarr Developers + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/packages/zarr-storage/README.md b/packages/zarr-storage/README.md new file mode 100644 index 0000000000..a3ead03a95 --- /dev/null +++ b/packages/zarr-storage/README.md @@ -0,0 +1,80 @@ +# zarr-storage + +An extraction of Zarr-Python's existing storage layer into the `legacy` +namespace, with concrete stores, wrappers, and reusable conformance tests. + +```python +from zarr.core.buffer import default_buffer_prototype +from zarr_storage.legacy import LatencyStore, MemoryStore, WrapperStore + +store = LatencyStore(MemoryStore(), get_latency=0.01) +``` + +## Included APIs + +- The `Store` ABC, byte requests, byte getter/setter protocols, and sync capabilities. +- `MemoryStore`, `ManagedMemoryStore`, `GpuMemoryStore`, `LocalStore`, `ZipStore`, + `FsspecStore`, and `ObjectStore`. +- `WrapperStore`, `LoggingStore`, and `LatencyStore`. +- `StorePath`, store construction/path utilities, and byte-range coalescing. +- Experimental `CacheStore` at `zarr_storage.legacy.experimental.cache_store`. +- `zarr_storage.testing.StoreTests`, state machines, strategies, buffer fixtures, + and assertion helpers for third-party implementations. + +`LatencyStore` is available without pytest. The conformance utilities are optional: +install `zarr-storage[testing]` to use them. Install `zarr-storage[remote]` for +fsspec and obstore backends. GPU execution additionally requires a suitable CuPy +installation and hardware, as it does in Zarr-Python. + +## Status and compatibility + +This is an experimental extraction draft. The source comes from Zarr-Python's +`src/zarr/abc/store.py`, `src/zarr/storage`, and storage-related testing utilities +at commit `9c29a0da9`. Imports are redirected to the extracted implementations; +legacy signatures, inherited behavior, and async execution are preserved. +Future APIs can coexist under a different namespace. No new storage contract or +deprecation is introduced here. + +**There is still a runtime dependency on `zarr`.** Shared buffers, configuration, +concurrency, metadata-aware IO helpers, and sync utilities remain there. This +package must not become a dependency of `zarr` until that dependency is removed. +The tested compatibility baseline is this source checkout, not every published +release in the declared dependency range. + +The extracted classes have distinct identities. Importing this package does not +change Zarr's imports or make its array entry points accept these stores. At +runtime adoption, the old Zarr import paths should re-export one canonical +implementation. Until then, use these stores through their storage APIs. + +## Tests + +The package includes the full `tests/test_store` suite and the experimental +cache-store suite, redirected to the extracted implementations. The original +Zarr tests remain in place. Tests cover sync/async IO, ranges, listing, lifecycle, +read-only behavior, pickling, wrappers, caching, and array/group integration. + +The integration tests have an explicit, test-only fixture that rebinds Zarr's +storage references to the extracted implementations, simulating future +re-exports. They use real stores and real arrays; no IO methods are mocked by +that fixture. Bindings are restored after each test. Import-isolation tests +separately verify that normal package imports leave Zarr unchanged. + +From the repository root: + +```sh +# Core stores; optional-backend and GPU tests skip when dependencies are absent. +PYTHONPATH=packages/zarr-storage/src hatch run test.py3.12-minimal:pytest packages/zarr-storage/tests + +# Includes fsspec/obstore tests and a local moto S3 server. +PYTHONPATH=packages/zarr-storage/src hatch run test.py3.12-optional:pytest packages/zarr-storage/tests + +# Stateful tests retain the upstream opt-in flag. +PYTHONPATH=packages/zarr-storage/src hatch run test.py3.12-minimal:pytest packages/zarr-storage/tests/test_store/test_stateful.py --run-slow-hypothesis +``` + +Unsupported backend operations retain their upstream skips/xfails. The standalone +package suite uses a 50-example Hypothesis profile with no deadline; select a +registered profile via `HYPOTHESIS_PROFILE` to override it. + +Build from this directory with `hatch build`. CI runs source and wheel tests on +Python 3.12/3.14, with an additional optional-backend job on Python 3.12. diff --git a/packages/zarr-storage/docs/api/index.md b/packages/zarr-storage/docs/api/index.md new file mode 100644 index 0000000000..cdca800f84 --- /dev/null +++ b/packages/zarr-storage/docs/api/index.md @@ -0,0 +1,17 @@ +# Legacy storage API + +::: zarr_storage.legacy + +## Experimental cache + +::: zarr_storage.legacy.experimental.cache_store.CacheStore + +## Conformance utilities + +Install `zarr-storage[testing]` to use these utilities. + +::: zarr_storage.testing.StoreTests + +::: zarr_storage.testing.stateful.ZarrStoreStateMachine + +::: zarr_storage.testing.stateful.ZarrHierarchyStateMachine diff --git a/packages/zarr-storage/pyproject.toml b/packages/zarr-storage/pyproject.toml new file mode 100644 index 0000000000..6c2aaa57e4 --- /dev/null +++ b/packages/zarr-storage/pyproject.toml @@ -0,0 +1,75 @@ +[build-system] +requires = ["hatchling>=1.29.0"] +build-backend = "hatchling.build" + +[project] +name = "zarr-storage" +version = "0.1.0" +description = "Legacy Zarr stores, interfaces, and reusable conformance tests." +readme = "README.md" +requires-python = ">=3.12" +license = "MIT" +license-files = ["LICENSE.txt"] +authors = [{ name = "Davis Bennett", email = "davis.v.bennett@gmail.com" }] +# Temporary: legacy interfaces still consume Zarr's shared foundations. +# Remove this dependency before making zarr depend on this distribution. +dependencies = ["zarr>=3.3,<3.4", "numpy>=2", "typing_extensions>=4.14"] + +[project.optional-dependencies] +remote = ["fsspec>=2023.10.0", "obstore>=0.5.1"] +testing = ["pytest>=9.1", "pytest-asyncio>=1.0", "hypothesis>=6.160.0"] + +[project.urls] +Source = "https://github.com/zarr-developers/zarr-python/tree/main/packages/zarr-storage" +Issues = "https://github.com/zarr-developers/zarr-python/issues" + +[tool.hatch.build.targets.wheel] +packages = ["src/zarr_storage"] + +[tool.hatch.build.targets.sdist] +include = ["/src", "/tests", "/docs", "/CHANGELOG.md"] + +[tool.pytest.ini_options] +addopts = ["--import-mode=importlib"] +strict = true +faulthandler_timeout = 600 +faulthandler_exit_on_timeout = true +asyncio_mode = "auto" +asyncio_default_fixture_loop_scope = "function" +filterwarnings = [ + "error", + "ignore:Unclosed client session Sequence[tuple[int, Buffer | None]]: + """Fetch one byte range. Raises FileNotFoundError if the key is absent.""" + async with ctx.semaphore: + buf = await ctx.fetch(req) + if buf is None: + raise FileNotFoundError + return ((idx, buf),) + + +async def _fetch_group( + ctx: _WorkerCtx, members: list[tuple[int, RangeByteRequest]] +) -> Sequence[tuple[int, Buffer | None]]: + """Fetch one merged byte range and slice it back into per-input buffers. + + `members` must already be sorted by `start`; callers in this module + build it from the sorted mergeable list. Raises `FileNotFoundError` + if the key is absent. + """ + if len(members) == 1: + solo_idx, solo_req = members[0] + return await _fetch_single(ctx, solo_idx, solo_req) + + start = members[0][1].start + end = max(r.end for _, r in members) + async with ctx.semaphore: + big = await ctx.fetch(RangeByteRequest(start, end)) + if big is None: + raise FileNotFoundError + sliced = [(idx, big[r.start - start : r.end - start]) for idx, r in members] + return tuple(sliced) + + +def coalesce_ranges( + byte_ranges: Sequence[ByteRequest | None], + *, + max_gap_bytes: int, + max_coalesced_bytes: int, +) -> tuple[ + list[list[tuple[int, RangeByteRequest]]], + list[tuple[int, ByteRequest | None]], +]: + """Plan a set of byte-range fetches: which inputs merge, which stand alone. + + Pure (no I/O). The result is the I/O plan a caller would execute: each + group corresponds to one fetch of a coalesced byte range, and each + uncoalescable item corresponds to one fetch of the original request. + + All tuning knobs are required keyword arguments. `Store.get_ranges` is + the public entry point and owns the canonical default values; this + function takes them explicitly to avoid duplicating policy. + + Parameters + ---------- + byte_ranges + Input ranges. `None` means "the whole value". + max_gap_bytes + Two `RangeByteRequest`s separated by at most this many bytes may be + merged into one fetch. + max_coalesced_bytes + Upper bound on the size of a single merged fetch. + + Returns + ------- + groups + List of merged groups. Each group is a list of + `(input_index, RangeByteRequest)` pairs sorted by `start`. A + single-element group represents a `RangeByteRequest` that did not + merge with any neighbor. + uncoalescable + List of `(input_index, request)` pairs for inputs that are not + `RangeByteRequest` (`OffsetByteRequest`, `SuffixByteRequest`, + `None`). Indices are preserved from the input order. + + Notes + ----- + Only `RangeByteRequest` inputs participate in coalescing. Two ranges + merge when both: their gap (next `start` minus current group's running + `end`) is `<= max_gap_bytes`, and the resulting merged span is + `<= max_coalesced_bytes`. + """ + indexed = list(enumerate(byte_ranges)) + mergeable = [(i, r) for i, r in indexed if isinstance(r, RangeByteRequest)] + uncoalescable: list[tuple[int, ByteRequest | None]] = [ + (i, r) for i, r in indexed if not isinstance(r, RangeByteRequest) + ] + + # Sort mergeables by start offset, then merge. Track running start/end of the + # current group so each merge step is O(1) instead of O(group size). + mergeable.sort(key=lambda pair: pair[1].start) + groups: list[list[tuple[int, RangeByteRequest]]] = [] + group_start = 0 + group_end = 0 + for pair in mergeable: + _i, r = pair + if groups and r.start - group_end <= max_gap_bytes: + prospective_end = max(group_end, r.end) + if prospective_end - group_start <= max_coalesced_bytes: + groups[-1].append(pair) + group_end = prospective_end + continue + groups.append([pair]) + group_start = r.start + group_end = r.end + + return groups, uncoalescable + + +async def coalesced_get( + fetch: Callable[[ByteRequest | None], Awaitable[Buffer | None]], + byte_ranges: Sequence[ByteRequest | None], + *, + max_concurrency: int, + max_gap_bytes: int, + max_coalesced_bytes: int, +) -> AsyncGenerator[Sequence[tuple[int, Buffer | None]]]: + """Read many byte ranges through `fetch` with coalescing and concurrency. + + Nearby ranges are merged into a single underlying I/O, and merged fetches + are run concurrently. Each yield corresponds to exactly one underlying I/O + operation: a sequence of `(input_index, result)` tuples for all input + ranges served by that I/O. Tuples within a yielded sequence are ordered by + start offset. Yields across groups are in completion order, not input + order. + + All tuning knobs are required keyword arguments. `Store.get_ranges` is + the public entry point and owns the canonical default values; this + function takes them explicitly to avoid duplicating policy. + + Parameters + ---------- + fetch + Callable that reads one byte range and returns a `Buffer` (or `None` + if the underlying key does not exist). Typically constructed via + `functools.partial(store.get, key, prototype)`. + byte_ranges + Input ranges. `None` means "the whole value". + max_concurrency + Maximum number of merged fetches in flight at once. + max_gap_bytes + Forwarded to `coalesce_ranges`. + max_coalesced_bytes + Forwarded to `coalesce_ranges`. + + Yields + ------ + Sequence[tuple[int, Buffer | None]] + Per-I/O batch of `(input_index, result)` tuples. + + Notes + ----- + - Only `RangeByteRequest` inputs are coalesced. `OffsetByteRequest`, + `SuffixByteRequest`, and `None` are each treated as uncoalescable + (one fetch, one single-tuple yield per input). + - Failures from underlying fetches surface as a `BaseExceptionGroup` + (PEP 654). Inner exceptions include `FileNotFoundError` if a fetch + returns `None`, plus any exception `fetch` raises. Pending fetches are + cancelled as soon as one task fails, so the group typically contains a + single non-`CancelledError` exception even under high concurrency. + - Groups completed before the failure remain observable on the yields + preceding the raise. + - `GeneratorExit` raised by `aclose()` is filtered out so the iterator + closes cleanly; callers don't see a group containing only it. + """ + if not byte_ranges: + return + + groups, singles = coalesce_ranges( + byte_ranges, + max_gap_bytes=max_gap_bytes, + max_coalesced_bytes=max_coalesced_bytes, + ) + + ctx = _WorkerCtx(fetch=fetch, semaphore=asyncio.Semaphore(max_concurrency)) + + # Launch all work as tasks. The semaphore bounds actual I/O concurrency. + # TaskGroup wraps task exceptions in BaseExceptionGroup; we propagate the + # group unchanged as part of the public contract (callers handle batch + # failures via `except*` / PEP 654). GeneratorExit (raised when the + # consumer calls aclose()) is filtered out so close completes cleanly. + try: + async with asyncio.TaskGroup() as tg: + tasks = [ + *(tg.create_task(_fetch_group(ctx, group)) for group in groups), + *(tg.create_task(_fetch_single(ctx, i, single)) for i, single in singles), + ] + + for fut in asyncio.as_completed(tasks): + yield await fut + except BaseExceptionGroup as eg: + # Strip GeneratorExits (consumer aclose()) and propagate whatever remains. + _, other_errors = eg.split(GeneratorExit) + + if other_errors is not None: + raise other_errors from None diff --git a/packages/zarr-storage/src/zarr_storage/legacy/__init__.py b/packages/zarr-storage/src/zarr_storage/legacy/__init__.py new file mode 100644 index 0000000000..0e371fc0d3 --- /dev/null +++ b/packages/zarr-storage/src/zarr_storage/legacy/__init__.py @@ -0,0 +1,107 @@ +import sys +import warnings +from types import ModuleType +from typing import Any + +from zarr.errors import ZarrDeprecationWarning + +from zarr_storage.legacy._abc import ( + ByteGetter as ByteGetter, +) +from zarr_storage.legacy._abc import ( + ByteRequest as ByteRequest, +) +from zarr_storage.legacy._abc import ( + ByteSetter as ByteSetter, +) +from zarr_storage.legacy._abc import ( + OffsetByteRequest as OffsetByteRequest, +) +from zarr_storage.legacy._abc import ( + RangeByteRequest as RangeByteRequest, +) +from zarr_storage.legacy._abc import ( + Store as Store, +) +from zarr_storage.legacy._abc import ( + SuffixByteRequest as SuffixByteRequest, +) +from zarr_storage.legacy._abc import ( + SupportsDeleteSync as SupportsDeleteSync, +) +from zarr_storage.legacy._abc import ( + SupportsGetSync as SupportsGetSync, +) +from zarr_storage.legacy._abc import ( + SupportsSetSync as SupportsSetSync, +) +from zarr_storage.legacy._abc import ( + SupportsSyncStore as SupportsSyncStore, +) +from zarr_storage.legacy._abc import ( + SyncByteGetter as SyncByteGetter, +) +from zarr_storage.legacy._abc import ( + SyncByteSetter as SyncByteSetter, +) +from zarr_storage.legacy._abc import ( + _store_supports_sync_io as _store_supports_sync_io, +) +from zarr_storage.legacy._abc import ( + set_or_delete as set_or_delete, +) +from zarr_storage.legacy._common import StoreLike, StorePath +from zarr_storage.legacy._fsspec import FsspecStore +from zarr_storage.legacy._latency import LatencyStore +from zarr_storage.legacy._local import LocalStore +from zarr_storage.legacy._logging import LoggingStore +from zarr_storage.legacy._memory import GpuMemoryStore, ManagedMemoryStore, MemoryStore +from zarr_storage.legacy._obstore import ObjectStore +from zarr_storage.legacy._wrapper import WrapperStore +from zarr_storage.legacy._zip import ZipStore + +__all__ = [ + "ByteGetter", + "ByteRequest", + "ByteSetter", + "FsspecStore", + "GpuMemoryStore", + "LatencyStore", + "LocalStore", + "LoggingStore", + "ManagedMemoryStore", + "MemoryStore", + "ObjectStore", + "OffsetByteRequest", + "RangeByteRequest", + "Store", + "StoreLike", + "StorePath", + "SuffixByteRequest", + "SupportsDeleteSync", + "SupportsGetSync", + "SupportsSetSync", + "SupportsSyncStore", + "SyncByteGetter", + "SyncByteSetter", + "WrapperStore", + "ZipStore", + "set_or_delete", +] + + +class VerboseModule(ModuleType): + def __setattr__(self, attr: str, value: Any) -> None: + if attr == "default_compressor": + warnings.warn( + "setting zarr_storage.legacy.default_compressor is deprecated, use " + "zarr.config to configure array.v2_default_compressor " + "e.g. config.set({'codecs.zstd':'numcodecs.Zstd', 'array.v2_default_compressor.numeric': 'zstd'})", + ZarrDeprecationWarning, + stacklevel=1, + ) + else: + super().__setattr__(attr, value) + + +sys.modules[__name__].__class__ = VerboseModule diff --git a/packages/zarr-storage/src/zarr_storage/legacy/_abc.py b/packages/zarr-storage/src/zarr_storage/legacy/_abc.py new file mode 100644 index 0000000000..4183c59888 --- /dev/null +++ b/packages/zarr-storage/src/zarr_storage/legacy/_abc.py @@ -0,0 +1,748 @@ +from __future__ import annotations + +import asyncio +from abc import ABC, abstractmethod +from dataclasses import dataclass +from functools import partial +from itertools import starmap +from typing import TYPE_CHECKING, Literal, Protocol, runtime_checkable + +if TYPE_CHECKING: + from collections.abc import AsyncGenerator, AsyncIterator, Iterable, Sequence + from types import TracebackType + from typing import Any, Self + + from zarr.core.buffer import Buffer, BufferPrototype + +__all__ = [ + "ByteGetter", + "ByteSetter", + "Store", + "SupportsDeleteSync", + "SupportsGetSync", + "SupportsSetSync", + "SupportsSyncStore", + "SyncByteGetter", + "SyncByteSetter", + "set_or_delete", +] + + +@dataclass(frozen=True, slots=True) +class RangeByteRequest: + """Request a specific byte range""" + + start: int + """The start of the byte range request (inclusive).""" + end: int + """The end of the byte range request (exclusive).""" + + +@dataclass(frozen=True, slots=True) +class OffsetByteRequest: + """Request all bytes starting from a given byte offset""" + + offset: int + """The byte offset for the offset range request.""" + + +@dataclass(frozen=True, slots=True) +class SuffixByteRequest: + """Request up to the last `n` bytes""" + + suffix: int + """The number of bytes from the suffix to request.""" + + +type ByteRequest = RangeByteRequest | OffsetByteRequest | SuffixByteRequest + + +class Store(ABC): + """ + Abstract base class for Zarr stores. + """ + + _read_only: bool + _is_open: bool + + def __init__(self, *, read_only: bool = False) -> None: + self._is_open = False + self._read_only = read_only + + @classmethod + async def open(cls, *args: Any, **kwargs: Any) -> Self: + """ + Create and open the store. + + Parameters + ---------- + *args : Any + Positional arguments to pass to the store constructor. + **kwargs : Any + Keyword arguments to pass to the store constructor. + + Returns + ------- + Store + The opened store instance. + """ + store = cls(*args, **kwargs) + await store._open() + return store + + def with_read_only(self, read_only: bool = False) -> Store: + """ + Return a new store with a new read_only setting. + + The new store points to the same location with the specified new read_only state. + The returned Store is not automatically opened, and this store is + not automatically closed. + + Parameters + ---------- + read_only + If True, the store will be created in read-only mode. Defaults to False. + + Returns + ------- + A new store of the same type with the new read only attribute. + """ + raise NotImplementedError( + f"with_read_only is not implemented for the {type(self)} store type." + ) + + def __enter__(self) -> Self: + """Enter a context manager that will close the store upon exiting.""" + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> None: + """Close the store.""" + self.close() + + async def _open(self) -> None: + """ + Open the store. + + Raises + ------ + ValueError + If the store is already open. + """ + if self._is_open: + raise ValueError("store is already open") + self._is_open = True + + async def _ensure_open(self) -> None: + """Open the store if it is not already open.""" + if not self._is_open: + await self._open() + + async def is_empty(self, prefix: str) -> bool: + """ + Check if the directory is empty. + + Parameters + ---------- + prefix : str + Prefix of keys to check. + + Returns + ------- + bool + True if the store is empty, False otherwise. + """ + if not self.supports_listing: + raise NotImplementedError + if prefix != "" and not prefix.endswith("/"): + prefix += "/" + async for _ in self.list_prefix(prefix): + return False + return True + + async def clear(self) -> None: + """ + Clear the store. + + Remove all keys and values from the store. + """ + if not self.supports_deletes: + raise NotImplementedError + if not self.supports_listing: + raise NotImplementedError + self._check_writable() + await self.delete_dir("") + + @property + def read_only(self) -> bool: + """Is the store read-only?""" + return self._read_only + + def _check_writable(self) -> None: + """Raise an exception if the store is not writable.""" + if self.read_only: + raise ValueError("store was opened in read-only mode and does not support writing") + + @abstractmethod + def __eq__(self, value: object) -> bool: + """Equality comparison.""" + ... + + @abstractmethod + async def get( + self, + key: str, + prototype: BufferPrototype, + byte_range: ByteRequest | None = None, + ) -> Buffer | None: + """Retrieve the value associated with a given key. + + Parameters + ---------- + key : str + prototype : BufferPrototype + The prototype of the output buffer. Stores may support a default buffer prototype. + byte_range : ByteRequest, optional + ByteRequest may be one of the following. If not provided, all data associated with the key is retrieved. + - RangeByteRequest(int, int): Request a specific range of bytes in the form (start, end). The end is exclusive. If the given range is zero-length or starts after the end of the object, an error will be returned. Additionally, if the range ends after the end of the object, the entire remainder of the object will be returned. Otherwise, the exact requested range will be returned. + - OffsetByteRequest(int): Request all bytes starting from a given byte offset. This is equivalent to bytes={int}- as an HTTP header. + - SuffixByteRequest(int): Request the last int bytes. Note that here, int is the size of the request, not the byte offset. This is equivalent to bytes=-{int} as an HTTP header. + + Returns + ------- + Buffer + """ + ... + + @abstractmethod + async def get_partial_values( + self, + prototype: BufferPrototype, + key_ranges: Iterable[tuple[str, ByteRequest | None]], + ) -> list[Buffer | None]: + """Retrieve possibly partial values from given key_ranges. + + Parameters + ---------- + prototype : BufferPrototype + The prototype of the output buffer. Stores may support a default buffer prototype. + key_ranges : Iterable[tuple[str, tuple[int | None, int | None]]] + Ordered set of key, range pairs, a key may occur multiple times with different ranges + + Returns + ------- + list of values, in the order of the key_ranges, may contain null/none for missing keys + """ + ... + + @abstractmethod + async def exists(self, key: str) -> bool: + """Check if a key exists in the store. + + Parameters + ---------- + key : str + + Returns + ------- + bool + """ + ... + + @property + @abstractmethod + def supports_writes(self) -> bool: + """Does the store support writes?""" + ... + + @abstractmethod + async def set(self, key: str, value: Buffer) -> None: + """Store a (key, value) pair. + + Parameters + ---------- + key : str + value : Buffer + """ + ... + + async def set_if_not_exists(self, key: str, value: Buffer) -> None: + """ + Store a key to ``value`` if the key is not already present. + + Parameters + ---------- + key : str + value : Buffer + """ + # Note for implementers: the default implementation provided here + # is not safe for concurrent writers. There's a race condition between + # the `exists` check and the `set` where another writer could set some + # value at `key` or delete `key`. + if not await self.exists(key): + await self.set(key, value) + + async def _set_many(self, values: Iterable[tuple[str, Buffer]]) -> None: + """ + Insert multiple (key, value) pairs into storage. + """ + await asyncio.gather(*starmap(self.set, values)) + + @property + def supports_consolidated_metadata(self) -> bool: + """ + Does the store support consolidated metadata?. + + If it doesn't an error will be raised on requests to consolidate the metadata. + Returning `False` can be useful for stores which implement their own + consolidation mechanism outside of the zarr-python implementation. + """ + + return True + + @property + @abstractmethod + def supports_deletes(self) -> bool: + """Does the store support deletes?""" + ... + + @abstractmethod + async def delete(self, key: str) -> None: + """Remove a key from the store + + Parameters + ---------- + key : str + """ + ... + + @property + def supports_partial_writes(self) -> Literal[False]: + """Does the store support partial writes? + + Partial writes are no longer used by Zarr, so this is always false. + """ + return False + + @property + @abstractmethod + def supports_listing(self) -> bool: + """Does the store support listing?""" + ... + + @abstractmethod + def list(self) -> AsyncIterator[str]: + """Retrieve all keys in the store. + + Returns + ------- + AsyncIterator[str] + """ + # This method should be async, like overridden methods in child classes. + # However, that's not straightforward: + # https://stackoverflow.com/questions/68905848 + + @abstractmethod + def list_prefix(self, prefix: str) -> AsyncIterator[str]: + """ + Retrieve all keys in the store that begin with a given prefix. Keys are returned relative + to the root of the store. + + Parameters + ---------- + prefix : str + + Returns + ------- + AsyncIterator[str] + """ + # This method should be async, like overridden methods in child classes. + # However, that's not straightforward: + # https://stackoverflow.com/questions/68905848 + + @abstractmethod + def list_dir(self, prefix: str) -> AsyncIterator[str]: + """ + Retrieve all keys and prefixes with a given prefix and which do not contain the character + “/” after the given prefix. + + Parameters + ---------- + prefix : str + + Returns + ------- + AsyncIterator[str] + """ + # This method should be async, like overridden methods in child classes. + # However, that's not straightforward: + # https://stackoverflow.com/questions/68905848 + + async def delete_dir(self, prefix: str) -> None: + """ + Remove all keys and prefixes in the store that begin with a given prefix. + """ + if not self.supports_deletes: + raise NotImplementedError + if not self.supports_listing: + raise NotImplementedError + self._check_writable() + if prefix != "" and not prefix.endswith("/"): + prefix += "/" + async for key in self.list_prefix(prefix): + await self.delete(key) + + def close(self) -> None: + """Close the store.""" + self._is_open = False + + async def _get_many( + self, requests: Iterable[tuple[str, BufferPrototype, ByteRequest | None]] + ) -> AsyncGenerator[tuple[str, Buffer | None], None]: + """ + Retrieve a collection of objects from storage. In general this method does not guarantee + that objects will be retrieved in the order in which they were requested, so this method + yields tuple[str, Buffer | None] instead of just Buffer | None + """ + for req in requests: + yield (req[0], await self.get(*req)) + + async def get_ranges( + self, + key: str, + byte_ranges: Sequence[ByteRequest | None], + *, + prototype: BufferPrototype, + max_concurrency: int = 10, + max_gap_bytes: int = 1 << 20, # 1 MiB + max_coalesced_bytes: int = 16 << 20, # 16 MiB + ) -> AsyncIterator[Sequence[tuple[int, Buffer | None]]]: + """Read many byte ranges from `key`. + + Yields one batch per underlying I/O operation, each a sequence of + `(input_index, Buffer | None)` tuples. Batches across yields arrive in + completion order, not input order. The default implementation built + into `Store` runs the coalescer over `self.get`, so subclasses get a + working implementation for free; stores that have a more efficient + backend (e.g. ranged HTTP, S3 byte-range fetches) should override. + + Parameters + ---------- + key + Storage key to read from. + byte_ranges + Input ranges. `None` means "the whole value". + prototype + Buffer prototype, forwarded to `self.get`. + max_concurrency + Maximum number of merged fetches in flight at once. + max_gap_bytes + Two `RangeByteRequest`s separated by at most this many bytes may + be merged into one fetch. + max_coalesced_bytes + Upper bound on the size of a single merged fetch. + + Raises + ------ + BaseExceptionGroup + Failures from underlying fetches are reported as a + `BaseExceptionGroup` (PEP 654) and should be handled with + `except*`. Inner exceptions include `FileNotFoundError` if any + fetch returns `None` (i.e. `key` is absent), and any exception + raised by `self.get` for the corresponding range. Pending + fetches are cancelled as soon as one task fails, so the group + typically contains a single non-`CancelledError` exception even + under high concurrency. + """ + # Local import: zarr.core._coalesce imports symbols from this module. + from zarr_storage._coalesce import coalesced_get + + fetch = partial(self.get, key, prototype) + async for group in coalesced_get( + fetch, + byte_ranges, + max_concurrency=max_concurrency, + max_gap_bytes=max_gap_bytes, + max_coalesced_bytes=max_coalesced_bytes, + ): + yield group + + def get_ranges_sync( + self, + key: str, + byte_ranges: Sequence[ByteRequest | None], + *, + prototype: BufferPrototype, + max_gap_bytes: int = 1 << 20, # 1 MiB + max_coalesced_bytes: int = 16 << 20, # 16 MiB + ) -> Sequence[tuple[int, Buffer | None]]: + """Synchronous, coalescing counterpart of `get_ranges`. + + Plans merged fetches with the same `coalesce_ranges` policy as the async + path, then issues one synchronous `get_sync` per merged group (or per + uncoalescable request) and slices results back into per-input buffers. + Used by the sync codec pipeline's partial-shard reads so they get the + same byte-range coalescing as the async path, without an event loop. + + Returns a list of `(input_index, Buffer | None)`. Raises + `BaseExceptionGroup` containing a `FileNotFoundError` if the key is + absent (matching `get_ranges`), so callers can handle a deleted shard + uniformly across the sync and async paths. + + Requires the store to implement `get_sync` (`SupportsGetSync`). + """ + from zarr_storage._coalesce import coalesce_ranges + + if not isinstance(self, SupportsGetSync): + raise TypeError(f"{type(self).__name__} does not support synchronous reads") + + groups, uncoalescable = coalesce_ranges( + byte_ranges, max_gap_bytes=max_gap_bytes, max_coalesced_bytes=max_coalesced_bytes + ) + results: list[tuple[int, Buffer | None]] = [] + errors: list[BaseException] = [] + + def _get(req: ByteRequest | None) -> Buffer | None: + return self.get_sync(key, prototype=prototype, byte_range=req) + + for idx, req in uncoalescable: + buf = _get(req) + if buf is None: + errors.append(FileNotFoundError(key)) + else: + results.append((idx, buf)) + + for members in groups: + if len(members) == 1: + solo_idx, solo_req = members[0] + buf = _get(solo_req) + if buf is None: + errors.append(FileNotFoundError(key)) + else: + results.append((solo_idx, buf)) + continue + start = members[0][1].start + end = max(r.end for _, r in members) + big = _get(RangeByteRequest(start, end)) + if big is None: + errors.append(FileNotFoundError(key)) + continue + for member_idx, r in members: + results.append((member_idx, big[r.start - start : r.end - start])) + + if errors: + raise BaseExceptionGroup("chunk read failed", errors) + return results + + async def getsize(self, key: str) -> int: + """ + Return the size, in bytes, of a value in a Store. + + Parameters + ---------- + key : str + + Returns + ------- + nbytes : int + The size of the value (in bytes). + + Raises + ------ + FileNotFoundError + When the given key does not exist in the store. + """ + # Note to implementers: this default implementation is very inefficient since + # it requires reading the entire object. Many systems will have ways to get the + # size of an object without reading it. + # avoid circular import + from zarr.core.buffer.core import default_buffer_prototype + + value = await self.get(key, prototype=default_buffer_prototype()) + if value is None: + raise FileNotFoundError(key) + return len(value) + + async def getsize_prefix(self, prefix: str) -> int: + """ + Return the size, in bytes, of all values under a prefix. + + Parameters + ---------- + prefix : str + The prefix of the directory to measure. + + Returns + ------- + nbytes : int + The sum of the sizes of the values in the directory (in bytes). + + See Also + -------- + zarr.Array.nbytes_stored + Store.getsize + + Notes + ----- + ``getsize_prefix`` is just provided as a potentially faster alternative to + listing all the keys under a prefix calling [`Store.getsize`][zarr.abc.store.Store.getsize] on each. + + In general, ``prefix`` should be the path of an Array or Group in the Store. + Implementations may differ on the behavior when some other ``prefix`` + is provided. + """ + # TODO: Overlap listing keys with getsize calls. + # Currently, we load the list of keys into memory and only then move + # on to getting sizes. Ideally we would overlap those two, which should + # improve tail latency and might reduce memory pressure (since not all keys + # would be in memory at once). + + # avoid circular import + from zarr.core.common import concurrent_map + from zarr.core.config import config + + if prefix != "" and not prefix.endswith("/"): + prefix += "/" + keys = [(x,) async for x in self.list_prefix(prefix)] + limit = config.get("async.concurrency") + sizes = await concurrent_map(keys, self.getsize, limit=limit) + return sum(sizes) + + +@runtime_checkable +class ByteGetter(Protocol): + async def get( + self, prototype: BufferPrototype, byte_range: ByteRequest | None = None + ) -> Buffer | None: ... + + +@runtime_checkable +class ByteSetter(Protocol): + async def get( + self, prototype: BufferPrototype, byte_range: ByteRequest | None = None + ) -> Buffer | None: ... + + async def set(self, value: Buffer) -> None: ... + + async def delete(self) -> None: ... + + async def set_if_not_exists(self, default: Buffer) -> None: ... + + +@runtime_checkable +class SyncByteGetter(Protocol): + """A `ByteGetter` that can also fetch synchronously, without an event loop. + + Non-StorePath byte getters (e.g. the sharding codec's in-memory + `_ShardingByteGetter`) implement this so a synchronous codec pipeline can + take its sync fast path on them instead of scheduling one coroutine per + chunk. Note that `StorePath` also *has* a `get_sync` method (so it matches + this protocol structurally) but it only works when its store supports + synchronous IO — callers gate `StorePath` on the store's `SupportsGetSync` + instead of on this protocol. + """ + + def get_sync( + self, prototype: BufferPrototype | None = None, byte_range: ByteRequest | None = None + ) -> Buffer | None: ... + + +@runtime_checkable +class SyncByteSetter(SyncByteGetter, Protocol): + """A `ByteSetter` that can also write synchronously. See `SyncByteGetter`.""" + + def set_sync(self, value: Buffer) -> None: ... + + def delete_sync(self) -> None: ... + + +@runtime_checkable +class SupportsGetSync(Protocol): + """Store protocol for synchronous reads (`get_sync`). + + The store sync surface is all-or-nothing: a store implementing any of the + `*_sync` methods must implement all of them (`SupportsSyncStore`), because + consumers mix sync reads, writes, and deletes within one operation. + Capability-gated callers consult `_store_supports_sync_io` rather than the + individual protocols. + """ + + def get_sync( + self, + key: str, + *, + prototype: BufferPrototype | None = None, + byte_range: ByteRequest | None = None, + ) -> Buffer | None: ... + + +@runtime_checkable +class SupportsSetSync(Protocol): + """Store protocol for synchronous writes (`set_sync`). + + See `SupportsGetSync` for the all-or-nothing contract on the store sync + surface. + """ + + def set_sync(self, key: str, value: Buffer) -> None: ... + + +@runtime_checkable +class SupportsDeleteSync(Protocol): + """Store protocol for synchronous deletes (`delete_sync`). + + See `SupportsGetSync` for the all-or-nothing contract on the store sync + surface. + """ + + def delete_sync(self, key: str) -> None: ... + + +@runtime_checkable +class SupportsSyncStore(SupportsGetSync, SupportsSetSync, SupportsDeleteSync, Protocol): + """The full store sync surface: `get_sync`, `set_sync`, and `delete_sync`.""" + + +def _store_supports_sync_io(store: object) -> bool: + """Whether `store` can serve the full synchronous IO surface right now. + + Structural membership in `SupportsSyncStore` is necessary but not always + sufficient: a store can present the `*_sync` methods while its ability to + run them depends on runtime state the type system cannot see. Wrapper + stores are the canonical case — `WrapperStore` delegates the sync methods + to the store it wraps, so they only work when the wrapped store is itself + sync-capable. Such stores opt out dynamically via a `_supports_sync_io` + attribute/property (absent means capable). + + This is an interim, private convention pending a formal sync/async store + architecture — the store-side twin of the codec-side `_sync_capable` + convention consulted by `zarr.abc.codec._codec_supports_sync`. + + Synchronous IO is all-or-nothing: consumers such as the fused codec + pipeline mix synchronous reads, writes, and deletes within one batch + (e.g. a partial-chunk write reads existing bytes and an all-fill chunk is + deleted), so a partial sync surface never satisfies this predicate. + """ + return isinstance(store, SupportsSyncStore) and getattr(store, "_supports_sync_io", True) + + +async def set_or_delete(byte_setter: ByteSetter, value: Buffer | None) -> None: + """Set or delete a value in a byte setter + + Parameters + ---------- + byte_setter : ByteSetter + value : Buffer | None + + Notes + ----- + If value is None, the key will be deleted. + """ + if value is None: + await byte_setter.delete() + else: + await byte_setter.set(value) diff --git a/packages/zarr-storage/src/zarr_storage/legacy/_common.py b/packages/zarr-storage/src/zarr_storage/legacy/_common.py new file mode 100644 index 0000000000..63553256a8 --- /dev/null +++ b/packages/zarr-storage/src/zarr_storage/legacy/_common.py @@ -0,0 +1,663 @@ +from __future__ import annotations + +import importlib.util +from pathlib import Path +from typing import TYPE_CHECKING, Any, Literal, Self + +from zarr.core._json import buffer_to_json_object, get_json +from zarr.core.buffer import Buffer, default_buffer_prototype +from zarr.core.common import ( + ANY_ACCESS_MODE, + ZARR_JSON, + ZARRAY_JSON, + ZGROUP_JSON, + AccessModeLiteral, + ZarrFormat, +) +from zarr.errors import ContainsArrayAndGroupError, ContainsArrayError, ContainsGroupError + +from zarr_storage.legacy._abc import ( + ByteRequest, + Store, + SupportsDeleteSync, + SupportsGetSync, + SupportsSetSync, +) +from zarr_storage.legacy._local import LocalStore +from zarr_storage.legacy._memory import ManagedMemoryStore, MemoryStore +from zarr_storage.legacy._utils import UPath, _join_paths, normalize_path, parse_store_url + +_has_fsspec = importlib.util.find_spec("fsspec") +if _has_fsspec: + from fsspec.mapping import FSMap +else: + FSMap = None + +if TYPE_CHECKING: + from zarr.core.buffer import BufferPrototype + from zarr.core.common import JSON + + +class StorePath: + """ + Path-like interface for a Store. + + Parameters + ---------- + store : Store + The store to use. + path : str + The path within the store. + """ + + store: Store + path: str + + def __init__(self, store: Store, path: str = "") -> None: + self.store = store + self.path = normalize_path(path) + + @property + def read_only(self) -> bool: + return self.store.read_only + + @classmethod + async def _create_open_instance(cls, store: Store, path: str) -> Self: + """Helper to create and return a StorePath instance.""" + await store._ensure_open() + return cls(store, path) + + @classmethod + async def open(cls, store: Store, path: str, mode: AccessModeLiteral | None = None) -> Self: + """ + Open StorePath based on the provided mode. + + * If the mode is None, return an opened version of the store with no changes. + * If the mode is 'r+', 'w-', 'w', or 'a' and the store is read-only, raise a ValueError. + * If the mode is 'r' and the store is not read-only, return a copy of the store with read_only set to True. + * If the mode is 'w-' and the store is not read-only and the StorePath contains keys, raise a FileExistsError. + * If the mode is 'w' and the store is not read-only, delete all keys nested within the StorePath. + + Parameters + ---------- + mode : AccessModeLiteral + The mode to use when initializing the store path. + + The accepted values are: + + - `'r'`: read only (must exist) + - `'r+'`: read/write (must exist) + - `'a'`: read/write (create if doesn't exist) + - `'w'`: read/write (overwrite if exists) + - `'w-'`: read/write (create if doesn't exist). + + Raises + ------ + FileExistsError + If the mode is 'w-' and the store path already exists. + ValueError + If the mode is not "r" and the store is read-only, or + """ + + # fastpath if mode is None + if mode is None: + return await cls._create_open_instance(store, path) + + if mode not in ANY_ACCESS_MODE: + raise ValueError(f"Invalid mode: {mode}, expected one of {ANY_ACCESS_MODE}") + + if store.read_only: + # Don't allow write operations on a read-only store + if mode != "r": + raise ValueError( + f"Store is read-only but mode is {mode!r}. Create a writable store or use 'r' mode." + ) + self = await cls._create_open_instance(store, path) + elif mode == "r": + # Create read-only copy for read mode on writable store + try: + read_only_store = store.with_read_only(True) + except NotImplementedError as e: + raise ValueError( + "Store is not read-only but mode is 'r'. Unable to create a read-only copy of the store. " + "Please use a read-only store or a storage class that implements .with_read_only()." + ) from e + self = await cls._create_open_instance(read_only_store, path) + else: + # writable store and writable mode + self = await cls._create_open_instance(store, path) + + # Handle mode-specific operations + match mode: + case "w-": + if not await self.is_empty(): + raise FileExistsError( + f"Cannot create '{path}' with mode 'w-' because it already contains data. " + f"Use mode 'w' to overwrite or 'a' to append." + ) + case "w": + await self.delete_dir() + return self + + async def get( + self, + prototype: BufferPrototype | None = None, + byte_range: ByteRequest | None = None, + ) -> Buffer | None: + """ + Read bytes from the store. + + Parameters + ---------- + prototype : BufferPrototype, optional + The buffer prototype to use when reading the bytes. + byte_range : ByteRequest, optional + The range of bytes to read. + + Returns + ------- + buffer : Buffer or None + The read bytes, or None if the key does not exist. + """ + if prototype is None: + prototype = default_buffer_prototype() + return await self.store.get(self.path, prototype=prototype, byte_range=byte_range) + + async def get_json(self, *, byte_range: ByteRequest | None = None) -> JSON | None: + """ + Read and parse the JSON document at this path, or None if it is absent. + + Parameters + ---------- + byte_range : ByteRequest, optional + If given, read only this portion of the value. Note that a partial + read of a JSON document may not be valid JSON. + + Returns + ------- + JSON or None + The parsed JSON value, or None if this path does not exist. + """ + return await get_json(self.store, self.path, byte_range=byte_range) + + async def set(self, value: Buffer) -> None: + """ + Write bytes to the store. + + Parameters + ---------- + value : Buffer + The buffer to write. + """ + await self.store.set(self.path, value) + + async def delete(self) -> None: + """ + Delete the key from the store. + + Raises + ------ + NotImplementedError + If the store does not support deletion. + """ + await self.store.delete(self.path) + + async def delete_dir(self) -> None: + """ + Delete all keys with the given prefix from the store. + """ + await self.store.delete_dir(self.path) + + async def set_if_not_exists(self, default: Buffer) -> None: + """ + Store a key to `value` if the key is not already present. + + Parameters + ---------- + default : Buffer + The buffer to store if the key is not already present. + """ + await self.store.set_if_not_exists(self.path, default) + + async def exists(self) -> bool: + """ + Check if the key exists in the store. + + Returns + ------- + bool + True if the key exists in the store, False otherwise. + """ + return await self.store.exists(self.path) + + async def is_empty(self) -> bool: + """ + Check if any keys exist in the store with the given prefix. + + Returns + ------- + bool + True if no keys exist in the store with the given prefix, False otherwise. + """ + return await self.store.is_empty(self.path) + + # ------------------------------------------------------------------- + # Synchronous IO delegation + # ------------------------------------------------------------------- + + def get_sync( + self, + *, + prototype: BufferPrototype | None = None, + byte_range: ByteRequest | None = None, + ) -> Buffer | None: + """Synchronous read — delegates to `self.store.get_sync(self.path, ...)`.""" + if not isinstance(self.store, SupportsGetSync): + raise TypeError(f"Store {type(self.store).__name__} does not support synchronous get.") + if prototype is None: + prototype = default_buffer_prototype() + return self.store.get_sync(self.path, prototype=prototype, byte_range=byte_range) + + def set_sync(self, value: Buffer) -> None: + """Synchronous write — delegates to `self.store.set_sync(self.path, value)`.""" + if not isinstance(self.store, SupportsSetSync): + raise TypeError(f"Store {type(self.store).__name__} does not support synchronous set.") + self.store.set_sync(self.path, value) + + def delete_sync(self) -> None: + """Synchronous delete — delegates to `self.store.delete_sync(self.path)`.""" + if not isinstance(self.store, SupportsDeleteSync): + raise TypeError( + f"Store {type(self.store).__name__} does not support synchronous delete." + ) + self.store.delete_sync(self.path) + + def __truediv__(self, other: str) -> StorePath: + """Combine this store path with another path""" + return self.__class__(self.store, _join_paths([self.path, other])) + + def __str__(self) -> str: + return _join_paths([str(self.store), self.path]) + + def __repr__(self) -> str: + return f"StorePath({self.store.__class__.__name__}, '{self}')" + + def __eq__(self, other: object) -> bool: + """ + Check if two StorePath objects are equal. + + Returns + ------- + bool + True if the two objects are equal, False otherwise. + + Notes + ----- + Two StorePath objects are considered equal if their stores are equal + and their paths are equal. + """ + try: + return self.store == other.store and self.path == other.path # type: ignore[attr-defined, no-any-return] + except AttributeError: + return False + + +type StoreLike = Store | StorePath | FSMap | Path | UPath | str | dict[str, Buffer] + + +async def make_store( + store_like: StoreLike | None, + *, + mode: AccessModeLiteral | None = None, + storage_options: dict[str, Any] | None = None, +) -> Store: + """ + Convert a `StoreLike` object into a Store object. + + `StoreLike` objects are converted to `Store` as follows: + + - `Store` or `StorePath` = `Store` object. + - `Path` or `str` = `LocalStore` object. + - `str` that starts with a protocol = `FsspecStore` object. + - `dict[str, Buffer]` = `MemoryStore` object. + - `None` = `MemoryStore` object. + - `FSMap` = `FsspecStore` object. + - `UPath` = `FsspecStore` object, or `LocalStore` for a local `UPath`. + + Parameters + ---------- + store_like : StoreLike | None + The `StoreLike` object to convert to a `Store` object. See the + [storage documentation in the user guide][user-guide-store-like] + for a description of all valid StoreLike values. + mode : StoreAccessMode | None, optional + The mode to use when creating the `Store` object. If None, the + default mode is 'r'. + storage_options : dict[str, Any] | None, optional + The storage options to use when creating the `RemoteStore` object. If + None, the default storage options are used. + + Returns + ------- + Store + The converted Store object. + + Raises + ------ + TypeError + If the StoreLike object is not one of the supported types, or if storage_options is provided but not used. + """ + from zarr_storage.legacy._fsspec import FsspecStore # circular import + + # Parse URL early so we can reuse the result for both validation and routing + parsed = parse_store_url(store_like) if isinstance(store_like, str) else None + + # Check if storage_options is valid for this store_like + if storage_options is not None: + is_fsspec_uri = parsed is not None and parsed.scheme not in ("", "memory", "file") + if not is_fsspec_uri: + raise TypeError( + "'storage_options' was provided but unused. " + "'storage_options' is only used when the store is passed as an FSSpec URI string.", + ) + + assert mode in (None, "r", "r+", "a", "w", "w-") + _read_only = mode == "r" + + if isinstance(store_like, StorePath): + # Get underlying store + return store_like.store + + elif isinstance(store_like, Store): + # Already a Store + return store_like + + elif isinstance(store_like, dict): + # Already a dictionary that can be a MemoryStore + # + # We deliberate only consider dict[str, Buffer] here, and not arbitrary mutable mappings. + # By only allowing dictionaries, which are in-memory, we know that MemoryStore appropriate. + return await MemoryStore.open(store_dict=store_like, read_only=_read_only) + + elif store_like is None: + # Create a new in-memory store + return await make_store({}, mode=mode, storage_options=storage_options) + + elif isinstance(store_like, UPath): + # This must be checked before Path: in universal-pathlib < 0.3 every UPath, including + # remote ones like S3Path, subclasses pathlib.Path, and would otherwise be misrouted to a + # LocalStore. Local UPaths get a LocalStore so that UPath("/data") and Path("/data") agree, + # mirroring how the equivalent strings are routed below. + if store_like.protocol in ("", "file"): + return await make_store( + Path(store_like.path), mode=mode, storage_options=storage_options + ) + return FsspecStore.from_upath(store_like, read_only=_read_only) + + elif isinstance(store_like, Path): + # Create a new LocalStore + return await LocalStore.open(root=store_like, mode=mode, read_only=_read_only) + + elif isinstance(store_like, str) and parsed is not None: + if parsed.scheme == "memory" and not _has_fsspec: + # Create or get a ManagedMemoryStore + return ManagedMemoryStore(name=parsed.name, path=parsed.path, read_only=_read_only) + elif parsed.scheme == "file" or not parsed.scheme: + # Local filesystem path — use parsed.path to strip the file:// scheme + return await make_store(Path(parsed.path), mode=mode, storage_options=storage_options) + else: + # Assume fsspec can handle it (s3://, gs://, http://, etc.) + return FsspecStore.from_url( + store_like, storage_options=storage_options, read_only=_read_only + ) + + elif _has_fsspec and isinstance(store_like, FSMap): + return FsspecStore.from_mapper(store_like, read_only=_read_only) + + else: + raise TypeError(f"Unsupported type for store_like: '{type(store_like).__name__}'") + + +async def make_store_path( + store_like: StoreLike | None, + *, + path: str | None = "", + mode: AccessModeLiteral | None = None, + storage_options: dict[str, Any] | None = None, +) -> StorePath: + """ + Convert a `StoreLike` object into a StorePath object. + + This function takes a `StoreLike` object and returns a `StorePath` object. See `make_store` for details + of which `Store` is used for each type of `store_like` object. + + Parameters + ---------- + store_like : StoreLike or None, default=None + The `StoreLike` object to convert to a `StorePath` object. See the + [storage documentation in the user guide][user-guide-store-like] + for a description of all valid StoreLike values. + path : str | None, optional + The path to use when creating the `StorePath` object. If None, the + default path is the empty string. + mode : StoreAccessMode | None, optional + The mode to use when creating the `StorePath` object. If None, the + default mode is 'r'. + storage_options : dict[str, Any] | None, optional + The storage options to use when creating the `RemoteStore` object. If + None, the default storage options are used. + + Returns + ------- + StorePath + The converted StorePath object. + + Raises + ------ + TypeError + If the StoreLike object is not one of the supported types, or if storage_options is provided but not used. + ValueError + If path is provided for a store that does not support it. + + See Also + -------- + make_store + """ + path_normalized = normalize_path(path) + + if isinstance(store_like, StorePath): + # Already a StorePath + if storage_options: + raise TypeError( + "'storage_options' was provided but unused. " + "'storage_options' is only used when the store is passed as an FSSpec URI string.", + ) + return store_like / path_normalized + + elif _has_fsspec and isinstance(store_like, FSMap) and path: + raise ValueError( + "'path' was provided but is not used for FSMap store_like objects. Specify the path when creating the FSMap instance instead." + ) + + else: + store = await make_store(store_like, mode=mode, storage_options=storage_options) + return await StorePath.open(store, path=path_normalized, mode=mode) + + +async def ensure_no_existing_node( + store_path: StorePath, + zarr_format: ZarrFormat, + node_type: Literal["array", "group"] | None = None, +) -> None: + """ + Check if a store_path is safe for array / group creation. + Returns `None` or raises an exception. + + Parameters + ---------- + store_path : StorePath + The storage location to check. + zarr_format : ZarrFormat + The Zarr format to check. + node_type : str | None, optional + Raise an error if an "array", or "group" exists. By default (when None), raises an error for either. + + Raises + ------ + ContainsArrayError, ContainsGroupError, ContainsArrayAndGroupError + """ + if zarr_format == 2: + extant_node = await _contains_node_v2(store_path) + elif zarr_format == 3: + extant_node = await _contains_node_v3(store_path) + + match extant_node: + case "array": + if node_type != "group": + msg = f"An array exists in store {store_path.store!r} at path {store_path.path!r}." + raise ContainsArrayError(msg) + + case "group": + if node_type != "array": + msg = f"A group exists in store {store_path.store!r} at path {store_path.path!r}." + raise ContainsGroupError(msg) + + case "nothing": + return + + case _: + msg = f"Invalid value for extant_node: {extant_node}" # type: ignore[unreachable] + raise ValueError(msg) + + +async def _contains_node_v3(store_path: StorePath) -> Literal["array", "group", "nothing"]: + """ + Check if a store_path contains nothing, an array, or a group. This function + returns the string "array", "group", or "nothing" to denote containing an array, a group, or + nothing. + + Parameters + ---------- + store_path : StorePath + The location in storage to check. + + Returns + ------- + Literal["array", "group", "nothing"] + A string representing the zarr node found at store_path. + """ + result: Literal["array", "group", "nothing"] = "nothing" + extant_meta_bytes = await (store_path / ZARR_JSON).get() + # if no metadata document could be loaded, then we just return "nothing" + if extant_meta_bytes is not None: + try: + extant_meta_json = buffer_to_json_object(extant_meta_bytes) + # avoid constructing a full metadata document here in the name of speed. + if extant_meta_json["node_type"] == "array": + result = "array" + elif extant_meta_json["node_type"] == "group": + result = "group" + except (KeyError, TypeError, ValueError): + # any of these errors is consistent with no array or group present. `ValueError` + # covers both malformed JSON (`json.JSONDecodeError`) and non-UTF-8 bytes + # (`UnicodeDecodeError`), each a `ValueError` subclass. + pass + return result + + +async def _contains_node_v2(store_path: StorePath) -> Literal["array", "group", "nothing"]: + """ + Check if a store_path contains nothing, an array, a group, or both. If both an array and a + group are detected, a `ContainsArrayAndGroup` exception is raised. Otherwise, this function + returns the string "array", "group", or "nothing" to denote containing an array, a group, or + nothing. + + Parameters + ---------- + store_path : StorePath + The location in storage to check. + + Returns + ------- + Literal["array", "group", "nothing"] + A string representing the zarr node found at store_path. + """ + _array = await contains_array(store_path=store_path, zarr_format=2) + _group = await contains_group(store_path=store_path, zarr_format=2) + + if _array and _group: + msg = ( + "Array and group metadata documents (.zarray and .zgroup) were both found in store " + f"{store_path.store!r} at path {store_path.path!r}. " + "Only one of these files may be present in a given directory / prefix. " + "Remove the .zarray file, or the .zgroup file, or both." + ) + raise ContainsArrayAndGroupError(msg) + elif _array: + return "array" + elif _group: + return "group" + else: + return "nothing" + + +async def contains_array(store_path: StorePath, zarr_format: ZarrFormat) -> bool: + """ + Check if an array exists at a given StorePath. + + Parameters + ---------- + store_path : StorePath + The StorePath to check for an existing group. + zarr_format : + The zarr format to check for. + + Returns + ------- + bool + True if the StorePath contains a group, False otherwise. + + """ + if zarr_format == 3: + extant_meta_bytes = await (store_path / ZARR_JSON).get() + if extant_meta_bytes is None: + return False + else: + try: + extant_meta_json = buffer_to_json_object(extant_meta_bytes) + # we avoid constructing a full metadata document here in the name of speed. + if extant_meta_json["node_type"] == "array": + return True + except (ValueError, KeyError, TypeError): + return False + elif zarr_format == 2: + return await (store_path / ZARRAY_JSON).exists() + msg = f"Invalid zarr_format provided. Got {zarr_format}, expected 2 or 3" + raise ValueError(msg) + + +async def contains_group(store_path: StorePath, zarr_format: ZarrFormat) -> bool: + """ + Check if a group exists at a given StorePath. + + Parameters + ---------- + + store_path : StorePath + The StorePath to check for an existing group. + zarr_format : + The zarr format to check for. + + Returns + ------- + + bool + True if the StorePath contains a group, False otherwise + + """ + if zarr_format == 3: + return (await _contains_node_v3(store_path)) == "group" + elif zarr_format == 2: + return await (store_path / ZGROUP_JSON).exists() + msg = f"Invalid zarr_format provided. Got {zarr_format}, expected 2 or 3" # type: ignore[unreachable] + raise ValueError(msg) diff --git a/packages/zarr-storage/src/zarr_storage/legacy/_fsspec.py b/packages/zarr-storage/src/zarr_storage/legacy/_fsspec.py new file mode 100644 index 0000000000..930190059b --- /dev/null +++ b/packages/zarr-storage/src/zarr_storage/legacy/_fsspec.py @@ -0,0 +1,454 @@ +from __future__ import annotations + +import warnings +from contextlib import suppress +from typing import TYPE_CHECKING, Any + +from packaging.version import parse as parse_version +from zarr.core.buffer import Buffer +from zarr.errors import ZarrUserWarning + +from zarr_storage.legacy._abc import ( + ByteRequest, + OffsetByteRequest, + RangeByteRequest, + Store, + SuffixByteRequest, +) +from zarr_storage.legacy._utils import _dereference_path + +if TYPE_CHECKING: + from collections.abc import AsyncIterator, Iterable + + from fsspec import AbstractFileSystem + from fsspec.asyn import AsyncFileSystem + from fsspec.mapping import FSMap + from zarr.core.buffer import BufferPrototype + + +ALLOWED_EXCEPTIONS: tuple[type[Exception], ...] = ( + FileNotFoundError, + IsADirectoryError, + NotADirectoryError, +) + + +def _make_async(fs: AbstractFileSystem) -> AsyncFileSystem: + """Convert a sync FSSpec filesystem to an async FFSpec filesystem + + If the filesystem class supports async operations, a new async instance is created + from the existing instance. + + If the filesystem class does not support async operations, the existing instance + is wrapped with AsyncFileSystemWrapper. + """ + import fsspec + + fsspec_version = parse_version(fsspec.__version__) + if fs.async_impl and fs.asynchronous: + # Already an async instance of an async filesystem, nothing to do + return fs + if fs.async_impl: + # Convert sync instance of an async fs to an async instance. Reuse the original + # constructor arguments rather than round-tripping through JSON, since storage + # options may hold objects that are not JSON-serializable (e.g. credentials). + return type(fs)(*fs.storage_args, **{**fs.storage_options, "asynchronous": True}) + + if fsspec_version < parse_version("2024.12.0"): + raise ImportError( + f"The filesystem '{fs}' is synchronous, and the required " + "AsyncFileSystemWrapper is not available. Upgrade fsspec to version " + "2024.12.0 or later to enable this functionality." + ) + from fsspec.implementations.asyn_wrapper import AsyncFileSystemWrapper + + return AsyncFileSystemWrapper(fs, asynchronous=True) + + +class FsspecStore(Store): + """ + Store for remote data based on FSSpec. + + Parameters + ---------- + fs : AsyncFileSystem + The Async FSSpec filesystem to use with this store. + read_only : bool + Whether the store is read-only + path : str + The root path of the store. This should be a relative path and must not include the + filesystem scheme. + allowed_exceptions : tuple[type[Exception], ...] + When fetching data, these cases will be deemed to correspond to missing keys. + + Attributes + ---------- + fs + allowed_exceptions + supports_writes + supports_deletes + supports_listing + + Raises + ------ + TypeError + If the Filesystem does not support async operations. + ValueError + If the path argument includes a scheme. + + Warns + ----- + ZarrUserWarning + If the file system (fs) was not created with `asynchronous=True`. + + Notes + ----- + Closing the store does not close the underlying filesystem or its network + session. fsspec caches and shares filesystem instances across callers, so + the store cannot know whether it is the only user, and closing a shared + session would break other stores. The filesystem's lifecycle belongs to + whoever created it; use fsspec's own tools (e.g. `clear_instance_cache`) + to release it. + + See Also + -------- + FsspecStore.from_upath + FsspecStore.from_url + """ + + # based on FSSpec + supports_writes: bool = True + supports_deletes: bool = True + supports_listing: bool = True + + fs: AsyncFileSystem + allowed_exceptions: tuple[type[Exception], ...] + path: str + + def __init__( + self, + fs: AsyncFileSystem, + read_only: bool = False, + path: str = "/", + allowed_exceptions: tuple[type[Exception], ...] = ALLOWED_EXCEPTIONS, + ) -> None: + super().__init__(read_only=read_only) + self.fs = fs + self.path = path + self.allowed_exceptions = allowed_exceptions + + if not self.fs.async_impl: + raise TypeError("Filesystem needs to support async operations.") + if not self.fs.asynchronous: + warnings.warn( + f"fs ({fs}) was not created with `asynchronous=True`, this may lead to surprising behavior", + category=ZarrUserWarning, + stacklevel=2, + ) + + @classmethod + def from_upath( + cls, + upath: Any, + read_only: bool = False, + allowed_exceptions: tuple[type[Exception], ...] = ALLOWED_EXCEPTIONS, + ) -> FsspecStore: + """ + Create an FsspecStore from a upath object. + + Parameters + ---------- + upath : UPath + The upath to the root of the store. + read_only : bool + Whether the store is read-only, defaults to False. + allowed_exceptions : tuple, optional + The exceptions that are allowed to be raised when accessing the + store. Defaults to ALLOWED_EXCEPTIONS. + + Returns + ------- + FsspecStore + """ + # A UPath hands back a filesystem in whatever mode it was constructed with, which is + # synchronous unless the caller passed asynchronous=True. Route it through _make_async so + # that sync-mode instances of async filesystems are re-created in async mode, and + # genuinely synchronous filesystems are wrapped. + return cls( + fs=_make_async(upath.fs), + path=upath.path.rstrip("/"), + read_only=read_only, + allowed_exceptions=allowed_exceptions, + ) + + @classmethod + def from_mapper( + cls, + fs_map: FSMap, + read_only: bool = False, + allowed_exceptions: tuple[type[Exception], ...] = ALLOWED_EXCEPTIONS, + ) -> FsspecStore: + """ + Create an FsspecStore from an FSMap object. + + Parameters + ---------- + fs_map : FSMap + Fsspec mutable mapping object. + read_only : bool + Whether the store is read-only, defaults to False. + allowed_exceptions : tuple, optional + The exceptions that are allowed to be raised when accessing the + store. Defaults to ALLOWED_EXCEPTIONS. + + Returns + ------- + FsspecStore + """ + fs = _make_async(fs_map.fs) + return cls( + fs=fs, + path=fs_map.root, + read_only=read_only, + allowed_exceptions=allowed_exceptions, + ) + + @classmethod + def from_url( + cls, + url: str, + storage_options: dict[str, Any] | None = None, + read_only: bool = False, + allowed_exceptions: tuple[type[Exception], ...] = ALLOWED_EXCEPTIONS, + ) -> FsspecStore: + """ + Create an FsspecStore from a URL. The type of store is determined from the URL scheme. + + Parameters + ---------- + url : str + The URL to the root of the store. + storage_options : dict, optional + The options to pass to fsspec when creating the filesystem. + read_only : bool + Whether the store is read-only, defaults to False. + allowed_exceptions : tuple, optional + The exceptions that are allowed to be raised when accessing the + store. Defaults to ALLOWED_EXCEPTIONS. + + Returns + ------- + FsspecStore + """ + try: + from fsspec import url_to_fs + except ImportError: + # before fsspec==2024.3.1 + from fsspec.core import url_to_fs + + opts = storage_options or {} + opts = {"asynchronous": True, **opts} + + fs, path = url_to_fs(url, **opts) + if not fs.async_impl: + fs = _make_async(fs) + + return cls(fs=fs, path=path, read_only=read_only, allowed_exceptions=allowed_exceptions) + + def with_read_only(self, read_only: bool = False) -> FsspecStore: + # docstring inherited + return type(self)( + fs=self.fs, + path=self.path, + allowed_exceptions=self.allowed_exceptions, + read_only=read_only, + ) + + async def clear(self) -> None: + # docstring inherited + try: + for subpath in await self.fs._find(self.path, withdirs=True): + if subpath != self.path: + await self.fs._rm(subpath, recursive=True) + except FileNotFoundError: + pass + + def __repr__(self) -> str: + return f"" + + def __eq__(self, other: object) -> bool: + return ( + isinstance(other, type(self)) + and self.path == other.path + and self.read_only == other.read_only + and self.fs == other.fs + ) + + async def get( + self, + key: str, + prototype: BufferPrototype, + byte_range: ByteRequest | None = None, + ) -> Buffer | None: + # docstring inherited + if not self._is_open: + await self._open() + path = _dereference_path(self.path, key) + + try: + if byte_range is None: + value = prototype.buffer.from_bytes(await self.fs._cat_file(path)) + elif isinstance(byte_range, RangeByteRequest): + value = prototype.buffer.from_bytes( + await self.fs._cat_file( + path, + start=byte_range.start, + end=byte_range.end, + ) + ) + elif isinstance(byte_range, OffsetByteRequest): + value = prototype.buffer.from_bytes( + await self.fs._cat_file(path, start=byte_range.offset, end=None) + ) + elif isinstance(byte_range, SuffixByteRequest): + value = prototype.buffer.from_bytes( + await self.fs._cat_file(path, start=-byte_range.suffix, end=None) + ) + else: + raise ValueError(f"Unexpected byte_range, got {byte_range}.") + except self.allowed_exceptions: + return None + except OSError as e: + if "not satisfiable" in str(e): + # this is an s3-specific condition we probably don't want to leak + return prototype.buffer.from_bytes(b"") + raise + else: + return value + + async def set( + self, + key: str, + value: Buffer, + byte_range: tuple[int, int] | None = None, + ) -> None: + # docstring inherited + if not self._is_open: + await self._open() + self._check_writable() + if not isinstance(value, Buffer): + raise TypeError( + f"FsspecStore.set(): `value` must be a Buffer instance. Got an instance of {type(value)} instead." + ) + path = _dereference_path(self.path, key) + # write data + if byte_range: + raise NotImplementedError + await self.fs._pipe_file(path, value.to_bytes()) + + async def delete(self, key: str) -> None: + # docstring inherited + self._check_writable() + path = _dereference_path(self.path, key) + try: + await self.fs._rm(path) + except FileNotFoundError: + pass + except self.allowed_exceptions: + pass + + async def delete_dir(self, prefix: str) -> None: + # docstring inherited + if not self.supports_deletes: + raise NotImplementedError( + "This method is only available for stores that support deletes." + ) + self._check_writable() + + path_to_delete = _dereference_path(self.path, prefix) + + with suppress(*self.allowed_exceptions): + await self.fs._rm(path_to_delete, recursive=True) + + async def exists(self, key: str) -> bool: + # docstring inherited + path = _dereference_path(self.path, key) + exists: bool = await self.fs._exists(path) + return exists + + async def get_partial_values( + self, + prototype: BufferPrototype, + key_ranges: Iterable[tuple[str, ByteRequest | None]], + ) -> list[Buffer | None]: + # docstring inherited + # Materialise first: key_ranges may be a one-shot iterable, so a bare + # truthiness check (e.g. `if key_ranges`) would be unreliable for an + # empty generator. _cat_ranges also expects lists of paths/starts/stops. + key_ranges = list(key_ranges) + if not key_ranges: + return [] + paths: list[str] = [] + starts: list[int | None] = [] + stops: list[int | None] = [] + for key, byte_range in key_ranges: + paths.append(_dereference_path(self.path, key)) + if byte_range is None: + starts.append(None) + stops.append(None) + elif isinstance(byte_range, RangeByteRequest): + starts.append(byte_range.start) + stops.append(byte_range.end) + elif isinstance(byte_range, OffsetByteRequest): + starts.append(byte_range.offset) + stops.append(None) + elif isinstance(byte_range, SuffixByteRequest): + starts.append(-byte_range.suffix) + stops.append(None) + else: + raise ValueError(f"Unexpected byte_range, got {byte_range}.") + # TODO: expectations for exceptions or missing keys? + res = await self.fs._cat_ranges(paths, starts, stops, on_error="return") + # the following is an s3-specific condition we probably don't want to leak + res = [b"" if (isinstance(r, OSError) and "not satisfiable" in str(r)) else r for r in res] + for r in res: + if isinstance(r, Exception) and not isinstance(r, self.allowed_exceptions): + raise r + + return [None if isinstance(r, Exception) else prototype.buffer.from_bytes(r) for r in res] + + async def list(self) -> AsyncIterator[str]: + # docstring inherited + allfiles = await self.fs._find(self.path, detail=False, withdirs=False) + for onefile in (a.removeprefix(f"{self.path}/") for a in allfiles): + yield onefile + + async def list_dir(self, prefix: str) -> AsyncIterator[str]: + # docstring inherited + prefix = f"{self.path}/{prefix.rstrip('/')}" + try: + allfiles = await self.fs._ls(prefix, detail=False) + except FileNotFoundError: + return + for onefile in (a.replace(f"{prefix}/", "") for a in allfiles): + yield onefile.removeprefix(self.path).removeprefix("/") + + async def list_prefix(self, prefix: str) -> AsyncIterator[str]: + # docstring inherited + for onefile in await self.fs._find( + f"{self.path}/{prefix}", detail=False, maxdepth=None, withdirs=False + ): + yield onefile.removeprefix(f"{self.path}/") + + async def getsize(self, key: str) -> int: + path = _dereference_path(self.path, key) + info = await self.fs._info(path) + + size = info.get("size") + + if size is None: + # Not all filesystems support size. Fall back to reading the entire object + return await super().getsize(key) + else: + # fsspec doesn't have typing. We'll need to assume or verify this is true + return int(size) diff --git a/packages/zarr-storage/src/zarr_storage/legacy/_latency.py b/packages/zarr-storage/src/zarr_storage/legacy/_latency.py new file mode 100644 index 0000000000..4fabd10276 --- /dev/null +++ b/packages/zarr-storage/src/zarr_storage/legacy/_latency.py @@ -0,0 +1,173 @@ +from __future__ import annotations + +import asyncio +import time +from typing import TYPE_CHECKING, Self + +import numpy as np + +from zarr_storage.legacy._abc import ByteRequest, Store +from zarr_storage.legacy._wrapper import WrapperStore + +if TYPE_CHECKING: + from collections.abc import AsyncIterator, Iterable, Sequence + + from zarr.core.buffer import Buffer + from zarr.core.buffer.core import BufferPrototype + + +class LatencyStore(WrapperStore[Store]): + """ + A wrapper class that takes any store class in its constructor and + adds latency to the `set` and `get` methods. This can be used for + performance testing. + """ + + _get_latency: float | tuple[float, float] + _set_latency: float | tuple[float, float] + + def __init__( + self, + store: Store, + *, + get_latency: float | tuple[float, float] = 0, + set_latency: float | tuple[float, float] = 0, + ) -> None: + super().__init__(store) + self._get_latency = get_latency if isinstance(get_latency, tuple) else float(get_latency) + self._set_latency = set_latency if isinstance(set_latency, tuple) else float(set_latency) + + @property + def get_latency(self) -> float: + if isinstance(self._get_latency, float): + return self._get_latency + return max(0.0, np.random.normal(loc=self._get_latency[0], scale=self._get_latency[1])) + + @property + def set_latency(self) -> float: + if isinstance(self._set_latency, float): + return self._set_latency + return max(0.0, np.random.normal(loc=self._set_latency[0], scale=self._set_latency[1])) + + def _with_store(self, store: Store) -> Self: + # Pass the raw latency config, not the sampled `get_latency`/`set_latency` + # properties — sampling would freeze a `(loc, scale)` distribution into + # one fixed float on derived stores (e.g. via `with_read_only`). + return type(self)(store, get_latency=self._get_latency, set_latency=self._set_latency) + + async def set(self, key: str, value: Buffer) -> None: + """ + Add latency to the ``set`` method. + + Calls ``asyncio.sleep(self.set_latency)`` before invoking the wrapped ``set`` method. + + Parameters + ---------- + key : str + The key to set + value : Buffer + The value to set + + Returns + ------- + None + """ + await asyncio.sleep(self.set_latency) + await self._store.set(key, value) + + async def get( + self, key: str, prototype: BufferPrototype, byte_range: ByteRequest | None = None + ) -> Buffer | None: + """ + Add latency to the ``get`` method. + + Calls ``asyncio.sleep(self.get_latency)`` before invoking the wrapped ``get`` method. + + Parameters + ---------- + key : str + The key to get + prototype : BufferPrototype + The BufferPrototype to use. + byte_range : ByteRequest, optional + An optional byte range. + + Returns + ------- + buffer : Buffer or None + """ + await asyncio.sleep(self.get_latency) + return await self._store.get(key, prototype=prototype, byte_range=byte_range) + + def get_sync( + self, + key: str, + *, + prototype: BufferPrototype | None = None, + byte_range: ByteRequest | None = None, + ) -> Buffer | None: + """Add latency to `get_sync`. + + Sleeps `self.get_latency` on the calling thread (the sync path runs on + worker threads, not the event loop) before delegating to the wrapped + store. + """ + time.sleep(self.get_latency) + return super().get_sync(key, prototype=prototype, byte_range=byte_range) + + def set_sync(self, key: str, value: Buffer) -> None: + """Add latency to `set_sync`. + + Sleeps `self.set_latency` on the calling thread (the sync path runs on + worker threads, not the event loop) before delegating to the wrapped + store. + """ + time.sleep(self.set_latency) + super().set_sync(key, value) + + async def get_ranges( + self, + key: str, + byte_ranges: Sequence[ByteRequest | None], + *, + prototype: BufferPrototype, + max_concurrency: int | None = None, + max_gap_bytes: int | None = None, + max_coalesced_bytes: int | None = None, + ) -> AsyncIterator[Sequence[tuple[int, Buffer | None]]]: + """Byte-range reads built on `self.get`, so each fetch pays latency. + + Routes through the coalescing `Store.get_ranges` default instead of the + `WrapperStore` delegation, which would bypass this wrapper's `get` and + therefore the synthetic latency. `None` for a coalescing kwarg means + "use the `Store` default". + """ + kwargs: dict[str, int] = {} + if max_concurrency is not None: + kwargs["max_concurrency"] = max_concurrency + if max_gap_bytes is not None: + kwargs["max_gap_bytes"] = max_gap_bytes + if max_coalesced_bytes is not None: + kwargs["max_coalesced_bytes"] = max_coalesced_bytes + async for group in Store.get_ranges(self, key, byte_ranges, prototype=prototype, **kwargs): + yield group + + async def get_partial_values( + self, + prototype: BufferPrototype, + key_ranges: Iterable[tuple[str, ByteRequest | None]], + ) -> list[Buffer | None]: + """Partial-value reads built on `self.get`, so each fetch pays latency. + + Issues one `self.get` per `(key, byte_range)` pair instead of the + `WrapperStore` delegation, which would bypass this wrapper's `get` and + therefore the synthetic latency. + """ + return list( + await asyncio.gather( + *( + self.get(key, prototype=prototype, byte_range=byte_range) + for key, byte_range in key_ranges + ) + ) + ) diff --git a/packages/zarr-storage/src/zarr_storage/legacy/_local.py b/packages/zarr-storage/src/zarr_storage/legacy/_local.py new file mode 100644 index 0000000000..88638928a0 --- /dev/null +++ b/packages/zarr-storage/src/zarr_storage/legacy/_local.py @@ -0,0 +1,373 @@ +from __future__ import annotations + +import asyncio +import contextlib +import io +import os +import shutil +import sys +import uuid +from pathlib import Path +from typing import TYPE_CHECKING, BinaryIO, Literal, Self + +from zarr.core.buffer import Buffer +from zarr.core.buffer.core import default_buffer_prototype +from zarr.core.common import AccessModeLiteral, concurrent_map + +from zarr_storage.legacy._abc import ( + ByteRequest, + OffsetByteRequest, + RangeByteRequest, + Store, + SuffixByteRequest, +) + +if TYPE_CHECKING: + from collections.abc import AsyncIterator, Iterable, Iterator + + from zarr.core.buffer import BufferPrototype + + +def _get(path: Path, prototype: BufferPrototype, byte_range: ByteRequest | None) -> Buffer: + if byte_range is None: + return prototype.buffer.from_bytes(path.read_bytes()) + with path.open("rb") as f: + size = f.seek(0, io.SEEK_END) + if isinstance(byte_range, RangeByteRequest): + f.seek(byte_range.start) + return prototype.buffer.from_bytes(f.read(byte_range.end - f.tell())) + elif isinstance(byte_range, OffsetByteRequest): + f.seek(byte_range.offset) + elif isinstance(byte_range, SuffixByteRequest): + f.seek(max(0, size - byte_range.suffix)) + else: + raise TypeError(f"Unexpected byte_range, got {byte_range}.") + return prototype.buffer.from_bytes(f.read()) + + +if sys.platform == "win32": + # Per the os.rename docs: + # On Windows, if dst exists a FileExistsError is always raised. + _safe_move = os.rename +else: + # On Unix, os.rename silently replace files, so instead we use os.link like + # atomicwrites: + # https://github.com/untitaker/python-atomicwrites/blob/1.4.1/atomicwrites/__init__.py#L59-L60 + # This also raises FileExistsError if dst exists. + def _safe_move(src: Path, dst: Path) -> None: + os.link(src, dst) + os.unlink(src) + + +@contextlib.contextmanager +def _atomic_write( + path: Path, + mode: Literal["r+b", "wb"], + exclusive: bool = False, +) -> Iterator[BinaryIO]: + tmp_path = path.with_suffix(f".{uuid.uuid4().hex}.partial") + try: + with tmp_path.open(mode) as f: + yield f + if exclusive: + _safe_move(tmp_path, path) + else: + tmp_path.replace(path) + except Exception: + tmp_path.unlink(missing_ok=True) + raise + + +def _put(path: Path, value: Buffer, exclusive: bool = False) -> int: + path.parent.mkdir(parents=True, exist_ok=True) + # write takes any object supporting the buffer protocol + view = value.as_buffer_like() + with _atomic_write(path, "wb", exclusive=exclusive) as f: + return f.write(view) + + +class LocalStore(Store): + """ + Store for the local file system. + + Parameters + ---------- + root : str or Path + Directory to use as root of store. + read_only : bool + Whether the store is read-only + + Attributes + ---------- + supports_writes + supports_deletes + supports_listing + root + """ + + supports_writes: bool = True + supports_deletes: bool = True + supports_listing: bool = True + + root: Path + + def __init__(self, root: Path | str, *, read_only: bool = False) -> None: + super().__init__(read_only=read_only) + if isinstance(root, str): + root = Path(root) + if not isinstance(root, Path): + raise TypeError( + f"'root' must be a string or Path instance. Got an instance of {type(root)} instead." + ) + self.root = root + + def with_read_only(self, read_only: bool = False) -> Self: + # docstring inherited + return type(self)( + root=self.root, + read_only=read_only, + ) + + @classmethod + async def open( + cls, root: Path | str, *, read_only: bool = False, mode: AccessModeLiteral | None = None + ) -> Self: + """ + Create and open the store. + + Parameters + ---------- + root : str or Path + Directory to use as root of store. + read_only : bool + Whether the store is read-only + mode : + Mode in which to create the store. This only affects opening the store, + and the final read-only state of the store is controlled through the + read_only parameter. + + Returns + ------- + Store + The opened store instance. + """ + # If mode = 'r+', want to open in read only mode (fail if exists), + # but return a writeable store + if mode is not None: + read_only_creation = mode in ["r", "r+"] + else: + read_only_creation = read_only + store = cls(root, read_only=read_only_creation) + await store._open() + + # Set read_only state + store = store.with_read_only(read_only) + await store._open() + return store + + async def _open(self, *, mode: AccessModeLiteral | None = None) -> None: + if not self.read_only: + self.root.mkdir(parents=True, exist_ok=True) + + if not self.root.exists(): + raise FileNotFoundError(f"{self.root} does not exist") + return await super()._open() + + async def clear(self) -> None: + # docstring inherited + self._check_writable() + shutil.rmtree(self.root) + self.root.mkdir() + + def __str__(self) -> str: + return f"file://{self.root.as_posix()}" + + def __repr__(self) -> str: + return f"LocalStore('{self}')" + + def __eq__(self, other: object) -> bool: + return isinstance(other, type(self)) and self.root == other.root + + # ------------------------------------------------------------------- + # Synchronous store methods + # ------------------------------------------------------------------- + + def _ensure_open_sync(self) -> None: + if not self._is_open: + if not self.read_only: + self.root.mkdir(parents=True, exist_ok=True) + if not self.root.exists(): + raise FileNotFoundError(f"{self.root} does not exist") + self._is_open = True + + def get_sync( + self, + key: str, + *, + prototype: BufferPrototype | None = None, + byte_range: ByteRequest | None = None, + ) -> Buffer | None: + if prototype is None: + prototype = default_buffer_prototype() + self._ensure_open_sync() + assert isinstance(key, str) + path = self.root / key + try: + return _get(path, prototype, byte_range) + except (FileNotFoundError, IsADirectoryError, NotADirectoryError): + return None + + def set_sync(self, key: str, value: Buffer) -> None: + self._ensure_open_sync() + self._check_writable() + assert isinstance(key, str) + if not isinstance(value, Buffer): + raise TypeError( + f"LocalStore.set(): `value` must be a Buffer instance. " + f"Got an instance of {type(value)} instead." + ) + path = self.root / key + _put(path, value) + + def delete_sync(self, key: str) -> None: + self._ensure_open_sync() + self._check_writable() + path = self.root / key + if path.is_dir(): + shutil.rmtree(path) + else: + path.unlink(missing_ok=True) + + async def get( + self, + key: str, + prototype: BufferPrototype | None = None, + byte_range: ByteRequest | None = None, + ) -> Buffer | None: + # docstring inherited + if prototype is None: + prototype = default_buffer_prototype() + if not self._is_open: + await self._open() + assert isinstance(key, str) + path = self.root / key + + try: + return await asyncio.to_thread(_get, path, prototype, byte_range) + except (FileNotFoundError, IsADirectoryError, NotADirectoryError): + return None + + async def get_partial_values( + self, + prototype: BufferPrototype, + key_ranges: Iterable[tuple[str, ByteRequest | None]], + ) -> list[Buffer | None]: + # docstring inherited + args = [] + for key, byte_range in key_ranges: + assert isinstance(key, str) + path = self.root / key + args.append((_get, path, prototype, byte_range)) + return await concurrent_map(args, asyncio.to_thread, limit=None) # TODO: fix limit + + async def set(self, key: str, value: Buffer) -> None: + # docstring inherited + return await self._set(key, value) + + async def set_if_not_exists(self, key: str, value: Buffer) -> None: + # docstring inherited + try: + return await self._set(key, value, exclusive=True) + except FileExistsError: + pass + + async def _set(self, key: str, value: Buffer, exclusive: bool = False) -> None: + if not self._is_open: + await self._open() + self._check_writable() + assert isinstance(key, str) + if not isinstance(value, Buffer): + raise TypeError( + f"LocalStore.set(): `value` must be a Buffer instance. Got an instance of {type(value)} instead." + ) + path = self.root / key + await asyncio.to_thread(_put, path, value, exclusive=exclusive) + + async def delete(self, key: str) -> None: + """ + Remove a key from the store. + + Parameters + ---------- + key : str + + Notes + ----- + If ``key`` is a directory within this store, the entire directory + at ``store.root / key`` is deleted. + """ + # docstring inherited + self._check_writable() + path = self.root / key + if path.is_dir(): # TODO: support deleting directories? shutil.rmtree? + shutil.rmtree(path) + else: + await asyncio.to_thread(path.unlink, True) # Q: we may want to raise if path is missing + + async def delete_dir(self, prefix: str) -> None: + # docstring inherited + self._check_writable() + path = self.root / prefix + if path.is_dir(): + shutil.rmtree(path) + elif path.is_file(): + raise ValueError(f"delete_dir was passed a {prefix=!r} that is a file") + else: + # Non-existent directory + # This path is tested by test_group:test_create_creates_parents for one + pass + + async def exists(self, key: str) -> bool: + # docstring inherited + path = self.root / key + return await asyncio.to_thread(path.is_file) + + async def list(self) -> AsyncIterator[str]: + # docstring inherited + to_strip = self.root.as_posix() + "/" + for p in list(self.root.rglob("*")): + if p.is_file(): + yield p.as_posix().replace(to_strip, "") + + async def list_prefix(self, prefix: str) -> AsyncIterator[str]: + # docstring inherited + to_strip = self.root.as_posix() + "/" + prefix = prefix.rstrip("/") + for p in (self.root / prefix).rglob("*"): + if p.is_file(): + yield p.as_posix().replace(to_strip, "") + + async def list_dir(self, prefix: str) -> AsyncIterator[str]: + # docstring inherited + base = self.root / prefix + try: + key_iter = base.iterdir() + for key in key_iter: + yield key.relative_to(base).as_posix() + except (FileNotFoundError, NotADirectoryError): + pass + + async def move(self, dest_root: Path | str) -> None: + """ + Move the store to another path. The old root directory is deleted. + """ + if isinstance(dest_root, str): + dest_root = Path(dest_root) + os.makedirs(dest_root.parent, exist_ok=True) + if dest_root.exists(): + raise FileExistsError(f"Destination root {dest_root} already exists.") + shutil.move(self.root, dest_root) + self.root = dest_root + + async def getsize(self, key: str) -> int: + return (self.root / key).stat().st_size diff --git a/packages/zarr-storage/src/zarr_storage/legacy/_logging.py b/packages/zarr-storage/src/zarr_storage/legacy/_logging.py new file mode 100644 index 0000000000..71ecbe430f --- /dev/null +++ b/packages/zarr-storage/src/zarr_storage/legacy/_logging.py @@ -0,0 +1,258 @@ +from __future__ import annotations + +import inspect +import logging +import sys +import time +from collections import defaultdict +from contextlib import contextmanager +from typing import TYPE_CHECKING, Any, Self + +from zarr_storage.legacy._abc import Store +from zarr_storage.legacy._wrapper import WrapperStore + +if TYPE_CHECKING: + from collections.abc import AsyncGenerator, Generator, Iterable + + from zarr.core.buffer import Buffer, BufferPrototype + + from zarr_storage.legacy._abc import ByteRequest + + counter: defaultdict[str, int] + + +class LoggingStore[T_Store: Store](WrapperStore[T_Store]): + """ + Store that logs all calls to another wrapped store. + + Parameters + ---------- + store : Store + Store to wrap + log_level : str + Log level + log_handler : logging.Handler + Log handler + + Attributes + ---------- + counter : dict + Counter of number of times each method has been called + """ + + counter: defaultdict[str, int] + + def __init__( + self, + store: T_Store, + log_level: str = "DEBUG", + log_handler: logging.Handler | None = None, + ) -> None: + super().__init__(store) + self.counter = defaultdict(int) + self.log_level = log_level + self.log_handler = log_handler + self._configure_logger(log_level, log_handler) + + def _configure_logger( + self, log_level: str = "DEBUG", log_handler: logging.Handler | None = None + ) -> None: + self.log_level = log_level + self.logger = logging.getLogger(f"LoggingStore({self._store})") + self.logger.setLevel(log_level) + + if not self.logger.hasHandlers(): + if not log_handler: + log_handler = self._default_handler() + # Add handler to logger + self.logger.addHandler(log_handler) + + def _default_handler(self) -> logging.Handler: + """Define a default log handler""" + handler = logging.StreamHandler(stream=sys.stdout) + handler.setLevel(self.log_level) + handler.setFormatter( + logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s") + ) + return handler + + def _with_store(self, store: T_Store) -> Self: + return type(self)(store=store, log_level=self.log_level, log_handler=self.log_handler) + + @contextmanager + def log(self, hint: Any = "") -> Generator[None, None, None]: + """Context manager to log method calls + + Each call to the wrapped store is logged to the configured logger and added to + the counter dict. + """ + method = inspect.stack()[2].function + op = f"{type(self._store).__name__}.{method}" + if hint: + op = f"{op}({hint})" + self.logger.info(" Calling %s", op) + start_time = time.time() + try: + self.counter[method] += 1 + yield + finally: + end_time = time.time() + self.logger.info("Finished %s [%.2f s]", op, end_time - start_time) + + @classmethod + async def open(cls: type[Self], store_cls: type[T_Store], *args: Any, **kwargs: Any) -> Self: + log_level = kwargs.pop("log_level", "DEBUG") + log_handler = kwargs.pop("log_handler", None) + store = store_cls(*args, **kwargs) + await store._open() + return cls(store=store, log_level=log_level, log_handler=log_handler) + + @property + def supports_writes(self) -> bool: + with self.log(): + return self._store.supports_writes + + @property + def supports_deletes(self) -> bool: + with self.log(): + return self._store.supports_deletes + + @property + def supports_listing(self) -> bool: + with self.log(): + return self._store.supports_listing + + @property + def read_only(self) -> bool: + with self.log(): + return self._store.read_only + + @property + def _is_open(self) -> bool: + with self.log(): + return self._store._is_open + + @_is_open.setter + def _is_open(self, value: bool) -> None: + raise NotImplementedError("LoggingStore must be opened via the `_open` method") + + async def _open(self) -> None: + with self.log(): + return await self._store._open() + + async def _ensure_open(self) -> None: + with self.log(): + return await self._store._ensure_open() + + async def is_empty(self, prefix: str = "") -> bool: + # docstring inherited + with self.log(): + return await self._store.is_empty(prefix=prefix) + + async def clear(self) -> None: + # docstring inherited + with self.log(): + return await self._store.clear() + + def __str__(self) -> str: + return f"logging-{self._store}" + + def __repr__(self) -> str: + return f"LoggingStore({self._store.__class__.__name__}, '{self._store}')" + + def __eq__(self, other: object) -> bool: + with self.log(other): + return type(self) is type(other) and self._store.__eq__(other._store) + + async def get( + self, + key: str, + prototype: BufferPrototype, + byte_range: ByteRequest | None = None, + ) -> Buffer | None: + # docstring inherited + with self.log(key): + return await self._store.get(key=key, prototype=prototype, byte_range=byte_range) + + async def get_partial_values( + self, + prototype: BufferPrototype, + key_ranges: Iterable[tuple[str, ByteRequest | None]], + ) -> list[Buffer | None]: + # docstring inherited + key_ranges = list(key_ranges) + keys = ",".join([k[0] for k in key_ranges]) + with self.log(keys): + return await self._store.get_partial_values(prototype=prototype, key_ranges=key_ranges) + + async def exists(self, key: str) -> bool: + # docstring inherited + with self.log(key): + return await self._store.exists(key) + + async def set(self, key: str, value: Buffer) -> None: + # docstring inherited + with self.log(key): + return await self._store.set(key=key, value=value) + + async def set_if_not_exists(self, key: str, value: Buffer) -> None: + # docstring inherited + with self.log(key): + return await self._store.set_if_not_exists(key=key, value=value) + + async def delete(self, key: str) -> None: + # docstring inherited + with self.log(key): + return await self._store.delete(key=key) + + def get_sync( + self, + key: str, + *, + prototype: BufferPrototype | None = None, + byte_range: ByteRequest | None = None, + ) -> Buffer | None: + # docstring inherited + with self.log(key): + return super().get_sync(key, prototype=prototype, byte_range=byte_range) + + def set_sync(self, key: str, value: Buffer) -> None: + # docstring inherited + with self.log(key): + return super().set_sync(key, value) + + def delete_sync(self, key: str) -> None: + # docstring inherited + with self.log(key): + return super().delete_sync(key) + + async def list(self) -> AsyncGenerator[str, None]: + # docstring inherited + with self.log(): + async for key in self._store.list(): + yield key + + async def list_prefix(self, prefix: str) -> AsyncGenerator[str, None]: + # docstring inherited + with self.log(prefix): + async for key in self._store.list_prefix(prefix=prefix): + yield key + + async def list_dir(self, prefix: str) -> AsyncGenerator[str, None]: + # docstring inherited + with self.log(prefix): + async for key in self._store.list_dir(prefix=prefix): + yield key + + async def delete_dir(self, prefix: str) -> None: + # docstring inherited + with self.log(prefix): + await self._store.delete_dir(prefix=prefix) + + async def getsize(self, key: str) -> int: + with self.log(key): + return await self._store.getsize(key) + + async def getsize_prefix(self, prefix: str) -> int: + with self.log(prefix): + return await self._store.getsize_prefix(prefix) diff --git a/packages/zarr-storage/src/zarr_storage/legacy/_memory.py b/packages/zarr-storage/src/zarr_storage/legacy/_memory.py new file mode 100644 index 0000000000..085db0a41d --- /dev/null +++ b/packages/zarr-storage/src/zarr_storage/legacy/_memory.py @@ -0,0 +1,710 @@ +from __future__ import annotations + +import os +import threading +import weakref +from logging import getLogger +from typing import TYPE_CHECKING, Any, Self + +from zarr.core.buffer import Buffer, gpu +from zarr.core.buffer.core import default_buffer_prototype +from zarr.core.common import concurrent_map + +from zarr_storage.legacy._abc import ByteRequest, Store +from zarr_storage.legacy._utils import ( + _join_paths, + _normalize_byte_range_index, + normalize_path, + parse_store_url, +) + +if TYPE_CHECKING: + from collections.abc import AsyncIterator, Iterable, MutableMapping + + from zarr.core.buffer import BufferPrototype + + +logger = getLogger(__name__) + + +def _copy_buffer(value: Buffer) -> Buffer: + """Copy `value` so the store does not retain the caller's memory. + + Encoding a chunk can hand the store a zero-copy view of the user's array + (an uncompressed write is the common case), and unlike stores that + serialize on write, this one keeps whatever it is given alive in a dict. + Without this copy a later mutation of the user's array would rewrite + chunks already committed to the store. + """ + return type(value).from_array_like(value.as_array_like().copy()) + + +class MemoryStore(Store): + """ + Store for local memory. + + Parameters + ---------- + store_dict : dict + Initial data + read_only : bool + Whether the store is read-only + + Attributes + ---------- + supports_writes + supports_deletes + supports_listing + + Notes + ----- + Writes copy the buffer they are given, so the store never aliases the + caller's memory. Buffers passed via `store_dict` are the caller's + responsibility and are stored as-is. + """ + + supports_writes: bool = True + supports_deletes: bool = True + supports_listing: bool = True + + _store_dict: MutableMapping[str, Buffer] + + def __init__( + self, + store_dict: MutableMapping[str, Buffer] | None = None, + *, + read_only: bool = False, + ) -> None: + super().__init__(read_only=read_only) + if store_dict is None: + store_dict = {} + self._store_dict = store_dict + + def with_read_only(self, read_only: bool = False) -> MemoryStore: + # docstring inherited + return type(self)( + store_dict=self._store_dict, + read_only=read_only, + ) + + async def clear(self) -> None: + # docstring inherited + self._store_dict.clear() + + def __str__(self) -> str: + return f"memory://{id(self._store_dict)}" + + def __repr__(self) -> str: + return f"MemoryStore('{self}')" + + def __eq__(self, other: object) -> bool: + return ( + isinstance(other, type(self)) + and self._store_dict == other._store_dict + and self.read_only == other.read_only + ) + + # ------------------------------------------------------------------- + # Synchronous store methods + # ------------------------------------------------------------------- + + def get_sync( + self, + key: str, + *, + prototype: BufferPrototype | None = None, + byte_range: ByteRequest | None = None, + ) -> Buffer | None: + if prototype is None: + prototype = default_buffer_prototype() + if not self._is_open: + self._is_open = True + assert isinstance(key, str) + try: + value = self._store_dict[key] + start, stop = _normalize_byte_range_index(value, byte_range) + return prototype.buffer.from_buffer(value[start:stop]) + except KeyError: + return None + + def set_sync(self, key: str, value: Buffer) -> None: + self._check_writable() + if not self._is_open: + self._is_open = True + assert isinstance(key, str) + if not isinstance(value, Buffer): + raise TypeError( + f"MemoryStore.set(): `value` must be a Buffer instance. Got an instance of {type(value)} instead." + ) + self._store_dict[key] = _copy_buffer(value) + + def delete_sync(self, key: str) -> None: + self._check_writable() + if not self._is_open: + self._is_open = True + try: + del self._store_dict[key] + except KeyError: + logger.debug("Key %s does not exist.", key) + + async def get( + self, + key: str, + prototype: BufferPrototype | None = None, + byte_range: ByteRequest | None = None, + ) -> Buffer | None: + # docstring inherited + if prototype is None: + prototype = default_buffer_prototype() + if not self._is_open: + await self._open() + assert isinstance(key, str) + try: + value = self._store_dict[key] + start, stop = _normalize_byte_range_index(value, byte_range) + return prototype.buffer.from_buffer(value[start:stop]) + except KeyError: + return None + + async def get_partial_values( + self, + prototype: BufferPrototype, + key_ranges: Iterable[tuple[str, ByteRequest | None]], + ) -> list[Buffer | None]: + # docstring inherited + + # All the key-ranges arguments goes with the same prototype + async def _get(key: str, byte_range: ByteRequest | None) -> Buffer | None: + return await self.get(key, prototype=prototype, byte_range=byte_range) + + return await concurrent_map(key_ranges, _get, limit=None) + + async def exists(self, key: str) -> bool: + # docstring inherited + return key in self._store_dict + + async def set(self, key: str, value: Buffer, byte_range: tuple[int, int] | None = None) -> None: + # docstring inherited + self._check_writable() + await self._ensure_open() + assert isinstance(key, str) + if not isinstance(value, Buffer): + raise TypeError( + f"MemoryStore.set(): `value` must be a Buffer instance. Got an instance of {type(value)} instead." + ) + if byte_range is not None: + buf = self._store_dict[key] + buf[byte_range[0] : byte_range[1]] = value + self._store_dict[key] = buf + else: + self._store_dict[key] = _copy_buffer(value) + + async def set_if_not_exists(self, key: str, value: Buffer) -> None: + # docstring inherited + self._check_writable() + await self._ensure_open() + self._store_dict.setdefault(key, _copy_buffer(value)) + + async def delete(self, key: str) -> None: + # docstring inherited + self._check_writable() + try: + del self._store_dict[key] + except KeyError: + logger.debug("Key %s does not exist.", key) + + async def list(self) -> AsyncIterator[str]: + # docstring inherited + for key in self._store_dict: + yield key + + async def list_prefix(self, prefix: str) -> AsyncIterator[str]: + # docstring inherited + # note: we materialize all dict keys into a list here so we can mutate the dict in-place (e.g. in delete_prefix) + for key in list(self._store_dict): + if key.startswith(prefix): + yield key + + async def list_dir(self, prefix: str) -> AsyncIterator[str]: + # docstring inherited + prefix = prefix.rstrip("/") + + if prefix == "": + keys_unique = {k.split("/")[0] for k in self._store_dict} + else: + # Our dictionary doesn't contain directory markers, but we want to include + # a pseudo directory when there's a nested item and we're listing an + # intermediate level. + keys_unique = { + key.removeprefix(f"{prefix}/").split("/")[0] + for key in self._store_dict + if key.startswith(f"{prefix}/") and key not in {prefix, f"{prefix}/"} + } + + for key in keys_unique: + yield key + + +class GpuMemoryStore(MemoryStore): + """ + Store for GPU memory. + + Stores every chunk in GPU memory irrespective of the original location. + + The dictionary of buffers to initialize this memory store with *must* be + GPU Buffers. + + Writing data to this store through ``.set`` will move the buffer to the GPU + if necessary. + + Parameters + ---------- + store_dict : MutableMapping, optional + A mutable mapping with string keys and [zarr.core.buffer.gpu.Buffer][] + values. + read_only : bool + Whether to open the store in read-only mode. + """ + + _store_dict: MutableMapping[str, gpu.Buffer] # type: ignore[assignment] + + def __init__( + self, + store_dict: MutableMapping[str, gpu.Buffer] | None = None, + *, + read_only: bool = False, + ) -> None: + super().__init__(store_dict=store_dict, read_only=read_only) # type: ignore[arg-type] + + def __str__(self) -> str: + return f"gpumemory://{id(self._store_dict)}" + + def __repr__(self) -> str: + return f"GpuMemoryStore('{self}')" + + @classmethod + def from_dict(cls, store_dict: MutableMapping[str, Buffer]) -> Self: + """ + Create a GpuMemoryStore from a dictionary of buffers at any location. + + The dictionary backing the newly created ``GpuMemoryStore`` will not be + the same as ``store_dict``. + + Parameters + ---------- + store_dict : mapping + A mapping of strings keys to arbitrary Buffers. The buffer data + will be moved into a [`gpu.Buffer`][zarr.core.buffer.gpu.Buffer]. + + Returns + ------- + GpuMemoryStore + """ + gpu_store_dict = {k: gpu.Buffer.from_buffer(v) for k, v in store_dict.items()} + return cls(gpu_store_dict) + + async def set(self, key: str, value: Buffer, byte_range: tuple[int, int] | None = None) -> None: + # docstring inherited + self._check_writable() + assert isinstance(key, str) + if not isinstance(value, Buffer): + raise TypeError( + f"GpuMemoryStore.set(): `value` must be a Buffer instance. Got an instance of {type(value)} instead." + ) + # Convert to gpu.Buffer + gpu_value = value if isinstance(value, gpu.Buffer) else gpu.Buffer.from_buffer(value) + await super().set(key, gpu_value, byte_range=byte_range) + + def set_sync(self, key: str, value: Buffer) -> None: + # docstring inherited + self._check_writable() + assert isinstance(key, str) + if not isinstance(value, Buffer): + raise TypeError( + f"GpuMemoryStore.set(): `value` must be a Buffer instance. Got an instance of {type(value)} instead." + ) + # Convert to gpu.Buffer, mirroring `set` above: every value in this store's + # backing dict must be a gpu.Buffer, regardless of which API wrote it. + gpu_value = value if isinstance(value, gpu.Buffer) else gpu.Buffer.from_buffer(value) + super().set_sync(key, gpu_value) + + +# ----------------------------------------------------------------------------- +# ManagedMemoryStore and its registry +# ----------------------------------------------------------------------------- +# ManagedMemoryStore owns the lifecycle of its backing dict, enabling proper +# weakref-based tracking. This allows memory:// URLs to be resolved back to +# the store's dict within the same process. + + +class _ManagedStoreDict(dict[str, Buffer]): + """ + A dict subclass that supports weak references. + + Regular dicts don't support weakrefs, but we need to track managed store dicts + in a WeakValueDictionary so they can be garbage collected when no longer + referenced. This subclass adds the necessary __weakref__ slot. + """ + + __slots__ = ("__weakref__",) + + +class _ManagedStoreDictRegistry: + """ + Registry for managed store dicts. + + This registry is the source of truth for managed store dicts. It creates + new dicts, tracks them via weak references, and looks them up by name. + """ + + def __init__(self) -> None: + self._registry: weakref.WeakValueDictionary[str, _ManagedStoreDict] = ( + weakref.WeakValueDictionary() + ) + self._counter = 0 + self._lock = threading.Lock() + + def _generate_name(self) -> str: + """Generate a unique name for a store. + + Must be called while holding `self._lock`. + """ + name = str(self._counter) + self._counter += 1 + return name + + def get_or_create(self, name: str | None = None) -> tuple[_ManagedStoreDict, str]: + """ + Get an existing managed dict by name, or create a new one. + + Thread-safe: uses a lock to prevent TOCTOU races between + checking for an existing entry and inserting a new one. + + Parameters + ---------- + name : str | None + The name for the store. If None, a unique name is auto-generated. + If a store with this name already exists, returns the existing store. + Names cannot contain '/' characters. + + Returns + ------- + tuple[_ManagedStoreDict, str] + The store dict and its name. + + Raises + ------ + ValueError + If the name contains '/' characters. + """ + with self._lock: + if name is None: + name = self._generate_name() + elif "/" in name: + raise ValueError( + f"Store name cannot contain '/': {name!r}. " + "Use the 'path' parameter to specify a path within the store." + ) + + existing = self._registry.get(name) + if existing is not None: + return existing, name + + store_dict = _ManagedStoreDict() + self._registry[name] = store_dict + return store_dict, name + + def get(self, name: str) -> _ManagedStoreDict | None: + """ + Look up a managed store dict by name. + + Parameters + ---------- + name : str + The name of the store. + + Returns + ------- + _ManagedStoreDict | None + The store dict if found, None otherwise. + """ + return self._registry.get(name) + + +_managed_store_dict_registry = _ManagedStoreDictRegistry() + + +class ManagedMemoryStore(MemoryStore): + """ + A memory store that owns and manages the lifecycle of its backing dict. + + Unlike ``MemoryStore`` which accepts any ``MutableMapping``, this store + creates and owns its backing dict internally. This enables proper lifecycle + management and allows the store to be looked up by its ``memory://`` URL + within the same process. + + Parameters + ---------- + name : str | None + The name for this store, used in the ``memory://`` URL. If None, a unique + name is auto-generated. If a store with this name already exists, the + new store will share the same backing dict. + path : str + The root path for this store. All keys will be prefixed with this path. + read_only : bool + Whether the store is read-only. + + Attributes + ---------- + name : str + The name of this store. + path : str + The root path of this store. + + Notes + ----- + The backing dict is tracked via weak references and will be garbage collected + when no ``ManagedMemoryStore`` instances reference it. URLs pointing to a + garbage-collected store will fail to resolve. + + See Also + -------- + MemoryStore : A memory store that accepts any MutableMapping. + + Examples + -------- + >>> store = ManagedMemoryStore(name="my-data") + >>> str(store) + 'memory://my-data' + >>> # Later, resolve the URL back to the store's dict + >>> store2 = ManagedMemoryStore.from_url("memory://my-data") + >>> store2._store_dict is store._store_dict + True + >>> # Create a store with a path prefix + >>> store3 = ManagedMemoryStore.from_url("memory://my-data/subdir") + >>> store3.path + 'subdir' + """ + + _store_dict: _ManagedStoreDict + _name: str + path: str + + def __init__(self, name: str | None = None, *, path: str = "", read_only: bool = False) -> None: + # Skip MemoryStore.__init__ and call Store.__init__ directly + # because we manage _store_dict via the registry, not via a user-supplied + # MutableMapping. If MemoryStore.__init__ ever adds logic beyond setting + # _store_dict, that logic must be replicated here. + Store.__init__(self, read_only=read_only) + + # Get or create a managed dict from the registry + self._store_dict, self._name = _managed_store_dict_registry.get_or_create(name) + self.path = normalize_path(path) + + def __str__(self) -> str: + return _join_paths([f"memory://{self._name}", self.path]) + + def __repr__(self) -> str: + return f"ManagedMemoryStore('{self}')" + + def __eq__(self, other: object) -> bool: + return ( + isinstance(other, type(self)) + and self._store_dict is other._store_dict + and self.path == other.path + and self.read_only == other.read_only + ) + + @property + def name(self) -> str: + """The name of this store, used in the memory:// URL.""" + return self._name + + @classmethod + def _from_managed_dict( + cls, + managed_dict: _ManagedStoreDict, + name: str, + *, + path: str = "", + read_only: bool = False, + ) -> ManagedMemoryStore: + """Internal: create a store from an existing managed dict.""" + store = object.__new__(cls) + Store.__init__(store, read_only=read_only) + store._store_dict = managed_dict + store._name = name + store.path = normalize_path(path) + return store + + def with_read_only(self, read_only: bool = False) -> ManagedMemoryStore: + # docstring inherited + return type(self)._from_managed_dict( + self._store_dict, self._name, path=self.path, read_only=read_only + ) + + @classmethod + def from_url(cls, url: str, *, read_only: bool = False) -> ManagedMemoryStore: + """ + Create a ManagedMemoryStore from a memory:// URL. + + This looks up the backing dict in the process-wide registry and creates + a new store instance that shares the same dict. + + Parameters + ---------- + url : str + A URL like "memory://my-store" or "memory://my-store/path/to/data" + identifying the store and optional path prefix. + read_only : bool + Whether the new store should be read-only. + + Returns + ------- + ManagedMemoryStore + A store sharing the same backing dict as the original. + + Raises + ------ + ValueError + If the URL is not a valid memory:// URL or the store has been + garbage collected. + """ + parsed = parse_store_url(url) + if parsed.scheme != "memory": + raise ValueError( + f"Expected a 'memory://' URL, got scheme {parsed.scheme!r} in '{url}'." + ) + name = parsed.name or "" + managed_dict = _managed_store_dict_registry.get(name) + if managed_dict is None: + raise ValueError( + f"Memory store not found for URL '{url}'. " + "The store may have been garbage collected." + ) + return cls._from_managed_dict(managed_dict, name, path=parsed.path, read_only=read_only) + + # Override MemoryStore methods to use path prefix and check process + + def get_sync( + self, + key: str, + *, + prototype: BufferPrototype | None = None, + byte_range: ByteRequest | None = None, + ) -> Buffer | None: + # docstring inherited + return super().get_sync( + _join_paths([self.path, key]), prototype=prototype, byte_range=byte_range + ) + + def set_sync(self, key: str, value: Buffer) -> None: + # docstring inherited + super().set_sync(_join_paths([self.path, key]), value) + + def delete_sync(self, key: str) -> None: + # docstring inherited + super().delete_sync(_join_paths([self.path, key])) + + async def get( + self, + key: str, + prototype: BufferPrototype | None = None, + byte_range: ByteRequest | None = None, + ) -> Buffer | None: + # docstring inherited + return await super().get( + _join_paths([self.path, key]), prototype=prototype, byte_range=byte_range + ) + + # get_partial_values is intentionally NOT overridden here: MemoryStore.get_partial_values + # dispatches per-key through `self.get`, which already resolves to the override above. + # Re-prefixing the keys here as well would apply `self.path` twice. + + async def exists(self, key: str) -> bool: + # docstring inherited + return await super().exists(_join_paths([self.path, key])) + + async def set(self, key: str, value: Buffer, byte_range: tuple[int, int] | None = None) -> None: + # docstring inherited + return await super().set(_join_paths([self.path, key]), value, byte_range=byte_range) + + async def set_if_not_exists(self, key: str, value: Buffer) -> None: + # docstring inherited + return await super().set_if_not_exists(_join_paths([self.path, key]), value) + + async def delete(self, key: str) -> None: + # docstring inherited + return await super().delete(_join_paths([self.path, key])) + + async def list(self) -> AsyncIterator[str]: + # docstring inherited + prefix = f"{self.path}/" if self.path else "" + async for key in super().list(): + if key.startswith(prefix): + yield key.removeprefix(prefix) + + async def list_prefix(self, prefix: str) -> AsyncIterator[str]: + # docstring inherited + # Manual concatenation instead of _join_paths because we need "path/" + # as the prefix when prefix is empty (to list all keys under self.path) + full_prefix = f"{self.path}/{prefix}" if self.path else prefix + path_prefix = f"{self.path}/" if self.path else "" + async for key in super().list_prefix(full_prefix): + yield key.removeprefix(path_prefix) + + async def list_dir(self, prefix: str) -> AsyncIterator[str]: + # docstring inherited + full_prefix = _join_paths([self.path, prefix]) + async for key in super().list_dir(full_prefix): + yield key + + def __reduce__( + self, + ) -> tuple[type[ManagedMemoryStore], tuple[str | None], dict[str, Any]]: + """ + Support pickling of ManagedMemoryStore. + + On unpickle, the store will reconnect to an existing store with the same + name if one exists in the registry, or create a new empty store otherwise. + + Note that the backing dict data is NOT serialized - only the store's + identity (name, path, read_only) is preserved. If the original store has + been garbage collected, the unpickled store will have an empty dict. + + The current process ID is preserved so that cross-process unpickling can be + detected and will raise an error at unpickle time. + """ + return ( + self.__class__, + (self._name,), + { + "path": self.path, + "read_only": self.read_only, + "created_pid": os.getpid(), + }, + ) + + def __setstate__(self, state: dict[str, Any]) -> None: + """Restore state after unpickling. + + The pickle protocol calls ``cls(name)`` (via ``__reduce__``'s args) + then ``__setstate__(state)``. ``__init__`` already set up + ``_store_dict`` and ``_name`` from the registry — we just restore + path and read_only here. + """ + # Check for cross-process usage first, before mutating state + created_pid = state.get("created_pid") + if created_pid is not None and created_pid != os.getpid(): + raise RuntimeError( + f"ManagedMemoryStore '{self._name}' was created in process {created_pid} " + f"but is being unpickled in process {os.getpid()}. " + "ManagedMemoryStore instances cannot be shared across processes because " + "their backing dict is not serialized. Use a persistent store (e.g., " + "LocalStore, ZipStore) for cross-process data sharing." + ) + + self.path = normalize_path(state.get("path", "")) + # Use the Store-level _read_only attribute directly because + # Store.__init__ was already called by __init__ during unpickling + self._read_only = state.get("read_only", False) diff --git a/packages/zarr-storage/src/zarr_storage/legacy/_obstore.py b/packages/zarr-storage/src/zarr_storage/legacy/_obstore.py new file mode 100644 index 0000000000..ae29ec3336 --- /dev/null +++ b/packages/zarr-storage/src/zarr_storage/legacy/_obstore.py @@ -0,0 +1,513 @@ +from __future__ import annotations + +import asyncio +import contextlib +import pickle +from collections import defaultdict +from itertools import chain +from operator import itemgetter +from typing import TYPE_CHECKING, Self, TypedDict + +from zarr.core.common import concurrent_map +from zarr.core.config import config + +from zarr_storage.legacy._abc import ( + ByteRequest, + OffsetByteRequest, + RangeByteRequest, + Store, + SuffixByteRequest, +) +from zarr_storage.legacy._utils import _relativize_path + +if TYPE_CHECKING: + from collections.abc import AsyncGenerator, Coroutine, Iterable, Sequence + from typing import Any + + from obstore import ListResult, ListStream, ObjectMeta, OffsetRange, SuffixRange + from obstore.store import ObjectStore as _UpstreamObjectStore + from zarr.core.buffer import Buffer, BufferPrototype + +__all__ = ["ObjectStore"] + +_ALLOWED_EXCEPTIONS: tuple[type[Exception], ...] = ( + FileNotFoundError, + IsADirectoryError, + NotADirectoryError, +) + + +class ObjectStore[T_Store: "_UpstreamObjectStore"](Store): + """ + Store that uses obstore for fast read/write from AWS, GCP, Azure. + + Parameters + ---------- + store : obstore.store.ObjectStore + An obstore store instance that is set up with the proper credentials. + read_only : bool + Whether to open the store in read-only mode. + + Warnings + -------- + ObjectStore is experimental and subject to API changes without notice. Please + raise an issue with any comments/concerns about the store. + """ + + store: T_Store + """The underlying obstore instance.""" + + def __eq__(self, value: object) -> bool: + if not isinstance(value, ObjectStore): + return False + + if not self.read_only == value.read_only: + return False + + return self.store == value.store # type: ignore[no-any-return] + + def __init__(self, store: T_Store, *, read_only: bool = False) -> None: + if not store.__class__.__module__.startswith("obstore"): + raise TypeError(f"expected ObjectStore class, got {store!r}") + super().__init__(read_only=read_only) + self.store = store + + def with_read_only(self, read_only: bool = False) -> Self: + # docstring inherited + return type(self)( + store=self.store, + read_only=read_only, + ) + + def __str__(self) -> str: + return f"object_store://{self.store}" + + def __repr__(self) -> str: + return f"{type(self).__name__}({self})" + + def __getstate__(self) -> dict[Any, Any]: + state = self.__dict__.copy() + state["store"] = pickle.dumps(self.store) + return state + + def __setstate__(self, state: dict[Any, Any]) -> None: + state["store"] = pickle.loads(state["store"]) + self.__dict__.update(state) + + async def get( + self, key: str, prototype: BufferPrototype, byte_range: ByteRequest | None = None + ) -> Buffer | None: + # docstring inherited + import obstore as obs + + try: + if byte_range is None: + resp = await obs.get_async(self.store, key) + return prototype.buffer.from_bytes(await resp.bytes_async()) # type: ignore[arg-type] + elif isinstance(byte_range, RangeByteRequest): + bytes = await obs.get_range_async( + self.store, key, start=byte_range.start, end=byte_range.end + ) + return prototype.buffer.from_bytes(bytes) # type: ignore[arg-type] + elif isinstance(byte_range, OffsetByteRequest): + resp = await obs.get_async( + self.store, key, options={"range": {"offset": byte_range.offset}} + ) + return prototype.buffer.from_bytes(await resp.bytes_async()) # type: ignore[arg-type] + elif isinstance(byte_range, SuffixByteRequest): + # some object stores (Azure) don't support suffix requests. In this + # case, our workaround is to first get the length of the object and then + # manually request the byte range at the end. + try: + resp = await obs.get_async( + self.store, key, options={"range": {"suffix": byte_range.suffix}} + ) + return prototype.buffer.from_bytes(await resp.bytes_async()) # type: ignore[arg-type] + except obs.exceptions.NotSupportedError: + head_resp = await obs.head_async(self.store, key) + file_size = head_resp["size"] + suffix_len = byte_range.suffix + buffer = await obs.get_range_async( + self.store, + key, + start=file_size - suffix_len, + length=suffix_len, + ) + return prototype.buffer.from_bytes(buffer) # type: ignore[arg-type] + else: + raise ValueError(f"Unexpected byte_range, got {byte_range}") + except _ALLOWED_EXCEPTIONS: + return None + + async def get_partial_values( + self, + prototype: BufferPrototype, + key_ranges: Iterable[tuple[str, ByteRequest | None]], + ) -> list[Buffer | None]: + # docstring inherited + return await _get_partial_values(self.store, prototype=prototype, key_ranges=key_ranges) + + async def exists(self, key: str) -> bool: + # docstring inherited + import obstore as obs + + try: + await obs.head_async(self.store, key) + except FileNotFoundError: + return False + else: + return True + + @property + def supports_writes(self) -> bool: + # docstring inherited + return True + + async def set(self, key: str, value: Buffer) -> None: + # docstring inherited + import obstore as obs + + self._check_writable() + + buf = value.as_buffer_like() + await obs.put_async(self.store, key, buf) + + async def set_if_not_exists(self, key: str, value: Buffer) -> None: + # docstring inherited + import obstore as obs + + self._check_writable() + buf = value.as_buffer_like() + with contextlib.suppress(obs.exceptions.AlreadyExistsError): + await obs.put_async(self.store, key, buf, mode="create") + + @property + def supports_deletes(self) -> bool: + # docstring inherited + return True + + async def delete(self, key: str) -> None: + # docstring inherited + import obstore as obs + + self._check_writable() + + # Some obstore stores such as local filesystems, GCP and Azure raise an error + # when deleting a non-existent key, while others such as S3 and in-memory do + # not. We suppress the error to make the behavior consistent across all obstore + # stores. This is also in line with the behavior of the other Zarr store adapters. + with contextlib.suppress(FileNotFoundError): + await obs.delete_async(self.store, key) + + async def delete_dir(self, prefix: str) -> None: + # docstring inherited + import obstore as obs + + self._check_writable() + if prefix != "" and not prefix.endswith("/"): + prefix += "/" + + metas = await obs.list(self.store, prefix).collect_async() + keys = [(m["path"],) for m in metas] + await concurrent_map(keys, self.delete, limit=config.get("async.concurrency")) + + @property + def supports_listing(self) -> bool: + # docstring inherited + return True + + async def _list(self, prefix: str | None = None) -> AsyncGenerator[ObjectMeta, None]: + import obstore as obs + + objects: ListStream[Sequence[ObjectMeta]] = obs.list(self.store, prefix=prefix) + async for batch in objects: + for item in batch: + yield item + + def list(self) -> AsyncGenerator[str, None]: + # docstring inherited + return (obj["path"] async for obj in self._list()) + + def list_prefix(self, prefix: str) -> AsyncGenerator[str, None]: + # docstring inherited + return (obj["path"] async for obj in self._list(prefix)) + + def list_dir(self, prefix: str) -> AsyncGenerator[str, None]: + # docstring inherited + import obstore as obs + + coroutine = obs.list_with_delimiter_async(self.store, prefix=prefix) + return _transform_list_dir(coroutine, prefix) + + async def getsize(self, key: str) -> int: + # docstring inherited + import obstore as obs + + resp = await obs.head_async(self.store, key) + return resp["size"] + + async def getsize_prefix(self, prefix: str) -> int: + # docstring inherited + sizes = [obj["size"] async for obj in self._list(prefix=prefix)] + return sum(sizes) + + +async def _transform_list_dir( + list_result_coroutine: Coroutine[Any, Any, ListResult[Sequence[ObjectMeta]]], prefix: str +) -> AsyncGenerator[str, None]: + """ + Transform the result of list_with_delimiter into an async generator of paths. + """ + list_result = await list_result_coroutine + + # We assume that the underlying object-store implementation correctly handles the + # prefix, so we don't double-check that the returned results actually start with the + # given prefix. + prefix = prefix.rstrip("/") + for path in chain( + list_result["common_prefixes"], map(itemgetter("path"), list_result["objects"]) + ): + if prefix != "" and path == prefix: + continue + relpath = _relativize_path(path=path, prefix=prefix) + if relpath: + yield relpath + + +class _BoundedRequest(TypedDict): + """Range request with a known start and end byte. + + These requests can be multiplexed natively on the Rust side with + `obstore.get_ranges_async`. + """ + + original_request_index: int + """The positional index in the original key_ranges input""" + + start: int + """Start byte offset.""" + + end: int + """End byte offset.""" + + +class _OtherRequest(TypedDict): + """Offset or suffix range requests. + + These requests cannot be concurrent on the Rust side, and each need their own call + to `obstore.get_async`, passing in the `range` parameter. + """ + + original_request_index: int + """The positional index in the original key_ranges input""" + + path: str + """The path to request from.""" + + range: OffsetRange | None + # Note: suffix requests are handled separately because some object stores (Azure) + # don't support them + """The range request type.""" + + +class _SuffixRequest(TypedDict): + """Offset or suffix range requests. + + These requests cannot be concurrent on the Rust side, and each need their own call + to `obstore.get_async`, passing in the `range` parameter. + """ + + original_request_index: int + """The positional index in the original key_ranges input""" + + path: str + """The path to request from.""" + + range: SuffixRange + """The suffix range.""" + + +class _Response(TypedDict): + """A response buffer associated with the original index that it should be restored to.""" + + original_request_index: int + """The positional index in the original key_ranges input""" + + buffer: Buffer + """The buffer returned from obstore's range request.""" + + +async def _make_bounded_requests( + store: _UpstreamObjectStore, + path: str, + requests: list[_BoundedRequest], + prototype: BufferPrototype, + semaphore: asyncio.Semaphore, +) -> list[_Response]: + """Make all bounded requests for a specific file. + + `obstore.get_ranges_async` allows for making concurrent requests for multiple ranges + within a single file, and will e.g. merge concurrent requests. This only uses one + single Python coroutine. + """ + import obstore as obs + + starts = [r["start"] for r in requests] + ends = [r["end"] for r in requests] + async with semaphore: + responses = await obs.get_ranges_async(store, path=path, starts=starts, ends=ends) + + buffer_responses: list[_Response] = [] + for request, response in zip(requests, responses, strict=True): + buffer_responses.append( + { + "original_request_index": request["original_request_index"], + "buffer": prototype.buffer.from_bytes(response), # type: ignore[arg-type] + } + ) + + return buffer_responses + + +async def _make_other_request( + store: _UpstreamObjectStore, + request: _OtherRequest, + prototype: BufferPrototype, + semaphore: asyncio.Semaphore, +) -> list[_Response]: + """Make offset or full-file requests. + + We return a `list[_Response]` for symmetry with `_make_bounded_requests` so that all + futures can be gathered together. + """ + import obstore as obs + + async with semaphore: + if request["range"] is None: + resp = await obs.get_async(store, request["path"]) + else: + resp = await obs.get_async(store, request["path"], options={"range": request["range"]}) + buffer = await resp.bytes_async() + + return [ + { + "original_request_index": request["original_request_index"], + "buffer": prototype.buffer.from_bytes(buffer), # type: ignore[arg-type] + } + ] + + +async def _make_suffix_request( + store: _UpstreamObjectStore, + request: _SuffixRequest, + prototype: BufferPrototype, + semaphore: asyncio.Semaphore, +) -> list[_Response]: + """Make suffix requests. + + This is separated out from `_make_other_request` because some object stores (Azure) + don't support suffix requests. In this case, our workaround is to first get the + length of the object and then manually request the byte range at the end. + + We return a `list[_Response]` for symmetry with `_make_bounded_requests` so that all + futures can be gathered together. + """ + import obstore as obs + + async with semaphore: + try: + resp = await obs.get_async(store, request["path"], options={"range": request["range"]}) + buffer = await resp.bytes_async() + except obs.exceptions.NotSupportedError: + head_resp = await obs.head_async(store, request["path"]) + file_size = head_resp["size"] + suffix_len = request["range"]["suffix"] + buffer = await obs.get_range_async( + store, + request["path"], + start=file_size - suffix_len, + length=suffix_len, + ) + + return [ + { + "original_request_index": request["original_request_index"], + "buffer": prototype.buffer.from_bytes(buffer), # type: ignore[arg-type] + } + ] + + +async def _get_partial_values( + store: _UpstreamObjectStore, + prototype: BufferPrototype, + key_ranges: Iterable[tuple[str, ByteRequest | None]], +) -> list[Buffer | None]: + """Make multiple range requests. + + ObjectStore has a `get_ranges` method that will additionally merge nearby ranges, + but it's _per_ file. So we need to split these key_ranges into **per-file** key + ranges, and then reassemble the results in the original order. + + We separate into different requests: + + - One call to `obstore.get_ranges_async` **per target file** + - One call to `obstore.get_async` for each other request. + """ + key_ranges = list(key_ranges) + per_file_bounded_requests: dict[str, list[_BoundedRequest]] = defaultdict(list) + other_requests: list[_OtherRequest] = [] + suffix_requests: list[_SuffixRequest] = [] + + for idx, (path, byte_range) in enumerate(key_ranges): + if byte_range is None: + other_requests.append( + { + "original_request_index": idx, + "path": path, + "range": None, + } + ) + elif isinstance(byte_range, RangeByteRequest): + per_file_bounded_requests[path].append( + {"original_request_index": idx, "start": byte_range.start, "end": byte_range.end} + ) + elif isinstance(byte_range, OffsetByteRequest): + other_requests.append( + { + "original_request_index": idx, + "path": path, + "range": {"offset": byte_range.offset}, + } + ) + elif isinstance(byte_range, SuffixByteRequest): + suffix_requests.append( + { + "original_request_index": idx, + "path": path, + "range": {"suffix": byte_range.suffix}, + } + ) + else: + raise ValueError(f"Unsupported range input: {byte_range}") + + semaphore = asyncio.Semaphore(config.get("async.concurrency")) + + futs: list[Coroutine[Any, Any, list[_Response]]] = [] + for path, bounded_ranges in per_file_bounded_requests.items(): + futs.append( + _make_bounded_requests(store, path, bounded_ranges, prototype, semaphore=semaphore) + ) + + for request in other_requests: + futs.append(_make_other_request(store, request, prototype, semaphore=semaphore)) # noqa: PERF401 + + for suffix_request in suffix_requests: + futs.append(_make_suffix_request(store, suffix_request, prototype, semaphore=semaphore)) # noqa: PERF401 + + buffers: list[Buffer | None] = [None] * len(key_ranges) + + for responses in await asyncio.gather(*futs): + for resp in responses: + buffers[resp["original_request_index"]] = resp["buffer"] + + return buffers diff --git a/packages/zarr-storage/src/zarr_storage/legacy/_utils.py b/packages/zarr-storage/src/zarr_storage/legacy/_utils.py new file mode 100644 index 0000000000..847fa14279 --- /dev/null +++ b/packages/zarr-storage/src/zarr_storage/legacy/_utils.py @@ -0,0 +1,306 @@ +from __future__ import annotations + +import importlib +import re +from pathlib import Path, PureWindowsPath +from urllib.parse import urlparse + +if importlib.util.find_spec("upath"): + # Re-exported for zarr_storage.legacy._common, which needs it to recognize UPath store_like values. + # The redundant-looking alias is the explicit re-export mypy requires under strict mode. + from upath.core import UPath as UPath # noqa: PLC0414 +else: + + class UPath: # type: ignore[no-redef] + pass + + +import sys +from typing import TYPE_CHECKING, NamedTuple + +from zarr_storage.legacy._abc import OffsetByteRequest, RangeByteRequest, SuffixByteRequest + +if TYPE_CHECKING: + from collections.abc import Iterable, Mapping + + from zarr.core.buffer import Buffer + + from zarr_storage.legacy._abc import ByteRequest + + +class ParsedStoreUrl(NamedTuple): + """ + Parsed components of a store URL. + + Attributes + ---------- + scheme : str + The URL scheme (e.g., "memory", "file", "s3"). Empty string for local paths. + name : str | None + The store name/host component. For memory:// URLs this is the store name. + None if empty. + path : str + The path component within the store. + raw : str + The original URL string. + """ + + scheme: str + name: str | None + path: str + raw: str + + +def parse_store_url(url: str) -> ParsedStoreUrl: + """ + Parse a store URL into its components. + + Parameters + ---------- + url : str + A URL like "memory://store-name/path" or "s3://bucket/key" or a local path. + + Returns + ------- + ParsedStoreUrl + Named tuple with scheme, name, path, and raw URL. + + Examples + -------- + >>> parse_store_url("memory://mystore") + ParsedStoreUrl(scheme='memory', name='mystore', path='', raw='memory://mystore') + + >>> parse_store_url("memory://mystore/path/to/data") + ParsedStoreUrl(scheme='memory', name='mystore', path='path/to/data', raw='memory://mystore/path/to/data') + + >>> parse_store_url("s3://bucket/key") + ParsedStoreUrl(scheme='s3', name='bucket', path='key', raw='s3://bucket/key') + + >>> parse_store_url("/local/path") + ParsedStoreUrl(scheme='', name=None, path='/local/path', raw='/local/path') + + Note that ``memory://name/path`` and ``memory:///path`` are different: + the first has ``name="name"`` and ``path="path"``, while the second has + ``name=None`` and ``path="/path"`` (no host component between ``//`` and ``/``). + """ + # On Windows, bare paths like "C:\foo" or "C:/foo" cause urlparse to + # misinterpret the drive letter as a URL scheme. Detect this early and + # return a local-path result without going through urlparse. + if sys.platform == "win32" and PureWindowsPath(url).drive: + return ParsedStoreUrl(scheme="", name=None, path=url, raw=url) + + parsed = urlparse(url) + + # netloc is the "host" part (store name for memory://, bucket for s3://, etc.) + name = parsed.netloc or None + + # For URLs with a scheme and netloc (like memory://store/path or s3://bucket/key), + # strip the leading slash from the path component. + # For local paths (no scheme), preserve the path as-is. + if parsed.scheme and parsed.netloc: + path = parsed.path.lstrip("/") + else: + path = parsed.path + + return ParsedStoreUrl(scheme=parsed.scheme, name=name, path=path, raw=url) + + +def normalize_path(path: str | bytes | Path | None) -> str: + if path is None: + result = "" + elif isinstance(path, bytes): + result = str(path, "ascii") + + # handle pathlib.Path + + elif isinstance(path, Path | UPath): + result = str(path) + + elif isinstance(path, str): + result = path + + else: + raise TypeError(f'Object {path} has an invalid type for "path": {type(path).__name__}') + + # convert backslash to forward slash + result = result.replace("\\", "/") + + # remove leading and trailing slashes + result = result.strip("/") + + # collapse any repeated slashes + pat = re.compile(r"//+") + result = pat.sub("/", result) + + # disallow path segments with just '.' or '..' + segments = result.split("/") + if any(s in {".", ".."} for s in segments): + raise ValueError( + f"The path {path!r} is invalid because its string representation contains '.' or '..' segments." + ) + + return result + + +def _normalize_byte_range_index(data: Buffer, byte_range: ByteRequest | None) -> tuple[int, int]: + """ + Convert a ByteRequest into an explicit start and stop + """ + if byte_range is None: + start = 0 + stop = len(data) + 1 + elif isinstance(byte_range, RangeByteRequest): + start = byte_range.start + stop = byte_range.end + elif isinstance(byte_range, OffsetByteRequest): + start = byte_range.offset + stop = len(data) + 1 + elif isinstance(byte_range, SuffixByteRequest): + start = max(0, len(data) - byte_range.suffix) + stop = len(data) + 1 + else: + raise ValueError(f"Unexpected byte_range, got {byte_range}.") + return (start, stop) + + +def _join_paths(paths: Iterable[str]) -> str: + """ + Filter out instances of '' and join the remaining strings with '/'. + + Parameters + ---------- + paths : Iterable[str] + + Returns + ------- + str + + Examples + -------- + ```python + from zarr_storage.legacy._utils import _join_paths + _join_paths(["", "a", "b"]) + # 'a/b' + _join_paths(["a", "b", "c"]) + # 'a/b/c' + ``` + """ + return "/".join(filter(lambda v: v != "", paths)) + + +def _dereference_path(root: str, path: str) -> str: + """ + Combine a store-side root with a key into a single fully-qualified path. + + Unlike `_join_paths`, this is purpose-built for the case where `root` is + an opaque backend-side prefix that may use `"/"` as a sentinel for "root + of the filesystem" (notably for fsspec's `ReferenceFileSystem`). A + trailing `"/"` is stripped from `root` before joining; if `root` is then + empty, the bare `path` is returned so that joining `"/"` with `"key"` + yields `"key"` rather than `"//key"`. A trailing `"/"` on the result is + also stripped. + + Leading slashes on `root` are preserved -- a backend-side path like + `"/home/foo/data.zarr"` is an absolute filesystem path for + `LocalFileSystem` and must not lose its leading separator. + + Parameters + ---------- + root : str + The backend-side root of a store. May be `""`, `"/"`, an absolute + filesystem path, or a backend-specific prefix. + path : str + The key within the store, typically a zarr key like `"zarr.json"` + or `"a/b/c/zarr.json"`. + + Returns + ------- + str + `root` and `path` joined by a single `"/"`, with the `"/"` sentinel + collapsed and trailing slashes removed. + + Examples + -------- + ```python + from zarr_storage.legacy._utils import _dereference_path + _dereference_path("/", "zarr.json") # 'zarr.json' + _dereference_path("", "zarr.json") # 'zarr.json' + _dereference_path("/home/foo", "zarr.json") # '/home/foo/zarr.json' + _dereference_path("/home/foo/", "zarr.json") # '/home/foo/zarr.json' + _dereference_path("bucket/p", "zarr.json") # 'bucket/p/zarr.json' + ``` + """ + root = root.rstrip("/") + path = f"{root}/{path}" if root else path + return path.rstrip("/") + + +def _relativize_path(*, path: str, prefix: str) -> str: + """ + Make a "/"-delimited path relative to some prefix. If the prefix is '', then the path is + returned as-is. Otherwise, the prefix is removed from the path as well as the separator + string "/". + + If ``prefix`` is not the empty string and ``path`` does not start with ``prefix`` + followed by a "/" character, then an error is raised. + + This function assumes that the prefix does not end with "/". + + Parameters + ---------- + path : str + The path to make relative to the prefix. + prefix : str + The prefix to make the path relative to. + + Returns + ------- + str + + Examples + -------- + ```python + from zarr_storage.legacy._utils import _relativize_path + _relativize_path(path="a/b", prefix="") + # 'a/b' + _relativize_path(path="a/b/c", prefix="a/b") + # 'c' + ``` + """ + if prefix == "": + return path + else: + _prefix = f"{prefix}/" + if not path.startswith(_prefix): + raise ValueError(f"The first component of {path} does not start with {prefix}.") + return path.removeprefix(_prefix) + + +def _normalize_paths(paths: Iterable[str]) -> tuple[str, ...]: + """ + Normalize the input paths according to the normalization scheme used for zarr node paths. + If any two paths normalize to the same value, raise a ValueError. + """ + path_map: dict[str, str] = {} + for path in paths: + parsed = normalize_path(path) + if parsed in path_map: + msg = ( + f"After normalization, the value '{path}' collides with '{path_map[parsed]}'. " + f"Both '{path}' and '{path_map[parsed]}' normalize to the same value: '{parsed}'. " + f"You should use either '{path}' or '{path_map[parsed]}', but not both." + ) + raise ValueError(msg) + path_map[parsed] = path + return tuple(path_map.keys()) + + +def _normalize_path_keys[T](data: Mapping[str, T]) -> dict[str, T]: + """ + Normalize the keys of the input dict according to the normalization scheme used for zarr node + paths. If any two keys in the input normalize to the same value, raise a ValueError. + Returns a dict where the keys are the elements of the input and the values are the + normalized form of each key. + """ + parsed_keys = _normalize_paths(data.keys()) + return dict(zip(parsed_keys, data.values(), strict=True)) diff --git a/packages/zarr-storage/src/zarr_storage/legacy/_wrapper.py b/packages/zarr-storage/src/zarr_storage/legacy/_wrapper.py new file mode 100644 index 0000000000..cd781e400b --- /dev/null +++ b/packages/zarr-storage/src/zarr_storage/legacy/_wrapper.py @@ -0,0 +1,219 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, cast + +if TYPE_CHECKING: + from collections.abc import AsyncGenerator, AsyncIterator, Iterable, Sequence + from types import TracebackType + from typing import Any, Self + + from zarr.abc.buffer import Buffer + from zarr.core.buffer import BufferPrototype + + from zarr_storage.legacy._abc import ByteRequest + +from zarr_storage.legacy._abc import ( + Store, + SupportsDeleteSync, + SupportsGetSync, + SupportsSetSync, + _store_supports_sync_io, +) + + +class WrapperStore[T_Store: Store](Store): + """ + Store that wraps an existing Store. + + By default all of the store methods are delegated to the wrapped store instance, which is + accessible via the ``._store`` attribute of this class. + + Use this class to modify or extend the behavior of the other store classes. + """ + + _store: T_Store + + def __init__(self, store: T_Store) -> None: + self._store = store + + def _with_store(self, store: T_Store) -> Self: + """ + Constructs a new instance of the wrapper store with the same details but a new store. + """ + return type(self)(store=store) + + @classmethod + async def open(cls: type[Self], store_cls: type[T_Store], *args: Any, **kwargs: Any) -> Self: + store = store_cls(*args, **kwargs) + await store._open() + return cls(store=store) + + def with_read_only(self, read_only: bool = False) -> Self: + return self._with_store(cast(T_Store, self._store.with_read_only(read_only))) + + def __enter__(self) -> Self: + return self._with_store(self._store.__enter__()) + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> None: + return self._store.__exit__(exc_type, exc_value, traceback) + + async def _open(self) -> None: + await self._store._open() + + async def _ensure_open(self) -> None: + await self._store._ensure_open() + + async def is_empty(self, prefix: str) -> bool: + return await self._store.is_empty(prefix) + + @property + def _is_open(self) -> bool: + return self._store._is_open + + @_is_open.setter + def _is_open(self, value: bool) -> None: + raise NotImplementedError("WrapperStore must be opened via the `_open` method") + + async def clear(self) -> None: + return await self._store.clear() + + @property + def read_only(self) -> bool: + return self._store.read_only + + def _check_writable(self) -> None: + return self._store._check_writable() + + def __eq__(self, value: object) -> bool: + return type(self) is type(value) and self._store.__eq__(value._store) + + def __str__(self) -> str: + return f"wrapping-{self._store}" + + def __repr__(self) -> str: + return f"WrapperStore({self._store.__class__.__name__}, '{self._store}')" + + async def get( + self, key: str, prototype: BufferPrototype, byte_range: ByteRequest | None = None + ) -> Buffer | None: + return await self._store.get(key, prototype, byte_range) + + async def get_partial_values( + self, + prototype: BufferPrototype, + key_ranges: Iterable[tuple[str, ByteRequest | None]], + ) -> list[Buffer | None]: + return await self._store.get_partial_values(prototype, key_ranges) + + async def get_ranges( + self, + key: str, + byte_ranges: Sequence[ByteRequest | None], + *, + prototype: BufferPrototype, + max_concurrency: int | None = None, + max_gap_bytes: int | None = None, + max_coalesced_bytes: int | None = None, + ) -> AsyncIterator[Sequence[tuple[int, Buffer | None]]]: + """Forward `get_ranges` to the wrapped store. + + Default values for the coalescing kwargs are not declared here; the + wrapped store decides them. `None` means "don't override the wrapped + store's default". + """ + kwargs: dict[str, int] = {} + if max_concurrency is not None: + kwargs["max_concurrency"] = max_concurrency + if max_gap_bytes is not None: + kwargs["max_gap_bytes"] = max_gap_bytes + if max_coalesced_bytes is not None: + kwargs["max_coalesced_bytes"] = max_coalesced_bytes + async for group in self._store.get_ranges(key, byte_ranges, prototype=prototype, **kwargs): + yield group + + async def exists(self, key: str) -> bool: + return await self._store.exists(key) + + async def set(self, key: str, value: Buffer) -> None: + await self._store.set(key, value) + + async def set_if_not_exists(self, key: str, value: Buffer) -> None: + return await self._store.set_if_not_exists(key, value) + + async def _set_many(self, values: Iterable[tuple[str, Buffer]]) -> None: + await self._store._set_many(values) + + @property + def supports_writes(self) -> bool: + return self._store.supports_writes + + @property + def supports_deletes(self) -> bool: + return self._store.supports_deletes + + @property + def _supports_sync_io(self) -> bool: + # The delegating `*_sync` methods below make every wrapper structurally + # satisfy `SupportsSyncStore`; whether they can actually run depends on + # the wrapped store, so forward its capability (see + # `zarr_storage.legacy._abc._store_supports_sync_io`). + return _store_supports_sync_io(self._store) + + def get_sync( + self, + key: str, + *, + prototype: BufferPrototype | None = None, + byte_range: ByteRequest | None = None, + ) -> Buffer | None: + """Forward `get_sync` to the wrapped store.""" + if not isinstance(self._store, SupportsGetSync): + raise TypeError(f"Store {type(self._store).__name__} does not support synchronous get.") + return self._store.get_sync(key, prototype=prototype, byte_range=byte_range) # type: ignore[unreachable] + + def set_sync(self, key: str, value: Buffer) -> None: + """Forward `set_sync` to the wrapped store.""" + if not isinstance(self._store, SupportsSetSync): + raise TypeError(f"Store {type(self._store).__name__} does not support synchronous set.") + self._store.set_sync(key, value) # type: ignore[unreachable] + + def delete_sync(self, key: str) -> None: + """Forward `delete_sync` to the wrapped store.""" + if not isinstance(self._store, SupportsDeleteSync): + raise TypeError( + f"Store {type(self._store).__name__} does not support synchronous delete." + ) + self._store.delete_sync(key) # type: ignore[unreachable] + + async def delete(self, key: str) -> None: + await self._store.delete(key) + + @property + def supports_listing(self) -> bool: + return self._store.supports_listing + + def list(self) -> AsyncIterator[str]: + return self._store.list() + + def list_prefix(self, prefix: str) -> AsyncIterator[str]: + return self._store.list_prefix(prefix) + + def list_dir(self, prefix: str) -> AsyncIterator[str]: + return self._store.list_dir(prefix) + + async def delete_dir(self, prefix: str) -> None: + return await self._store.delete_dir(prefix) + + def close(self) -> None: + self._store.close() + + async def _get_many( + self, requests: Iterable[tuple[str, BufferPrototype, ByteRequest | None]] + ) -> AsyncGenerator[tuple[str, Buffer | None], None]: + async for req in self._store._get_many(requests): + yield req diff --git a/packages/zarr-storage/src/zarr_storage/legacy/_zip.py b/packages/zarr-storage/src/zarr_storage/legacy/_zip.py new file mode 100644 index 0000000000..5b1e89541b --- /dev/null +++ b/packages/zarr-storage/src/zarr_storage/legacy/_zip.py @@ -0,0 +1,403 @@ +from __future__ import annotations + +import io +import os +import shutil +import threading +import time +import zipfile +from pathlib import Path +from typing import IO, TYPE_CHECKING, Any, Literal + +from zarr.core.buffer import Buffer, BufferPrototype + +from zarr_storage.legacy._abc import ( + ByteRequest, + OffsetByteRequest, + RangeByteRequest, + Store, + SuffixByteRequest, +) + +if TYPE_CHECKING: + from collections.abc import AsyncIterator, Iterable + +ZipStoreAccessModeLiteral = Literal["r", "w", "a"] + + +class _RawReaderAdapter(io.RawIOBase): + """ + Adapt a minimal seekable reader to the `io` interface `zipfile` needs. + + Some file-like objects (e.g. `obstore.ReadableFile`) implement + `read`/`seek`/`tell` but are not `io.IOBase` instances, and their + `read` may return a buffer-protocol object rather than `bytes`. + Wrapping in this adapter plus `io.BufferedReader` yields real `bytes`. + + Reads are clamped to the bytes remaining before EOF: some readers + (obstore < 0.6) raise on short reads rather than returning fewer bytes. + The size is cached, which is safe because the adapter is only used for + read-only access. + """ + + def __init__(self, fileobj: IO[bytes]) -> None: + self._fileobj = fileobj + self._size: int | None = None + + def _get_size(self) -> int: + if self._size is None: + pos = self._fileobj.tell() + self._size = self._fileobj.seek(0, os.SEEK_END) + self._fileobj.seek(pos) + return self._size + + def readable(self) -> bool: + return True + + def seekable(self) -> bool: + return True + + def seek(self, pos: int, whence: int = 0) -> int: + return self._fileobj.seek(pos, whence) + + def tell(self) -> int: + return self._fileobj.tell() + + def readinto(self, b: Any) -> int: + n_requested = min(len(b), self._get_size() - self._fileobj.tell()) + if n_requested <= 0: + return 0 + data = self._fileobj.read(n_requested) + n = len(data) + b[:n] = memoryview(data) + return n + + +class ZipStore(Store): + """ + Store using a ZIP file. + + Parameters + ---------- + path : str, Path, or IO[bytes] + Location of file, or an open binary file object. A file object must + support `read`, `seek`, and `tell`; objects that are not `io.IOBase` + instances (e.g. an `obstore` reader) are adapted automatically but + can only be used for reading (`mode="r"`). The file object must stay + open for the lifetime of the store, and operations that require a + filesystem location (`clear`, `move`, pickling) are not supported. + mode : str, optional + One of 'r' to read an existing file, 'w' to truncate and write a new + file, 'a' to append to an existing file, or 'x' to exclusively create + and write a new file. + compression : int, optional + Compression method to use when writing to the archive. + allowZip64 : bool, optional + If True (the default) will create ZIP files that use the ZIP64 + extensions when the zipfile is larger than 2 GiB. If False + will raise an exception when the ZIP file would require ZIP64 + extensions. + + Attributes + ---------- + allowed_exceptions + supports_writes + supports_deletes + supports_listing + path + compression + allowZip64 + """ + + supports_writes: bool = True + supports_deletes: bool = False + supports_listing: bool = True + + path: Path | None + compression: int + allowZip64: bool + + _zf: zipfile.ZipFile + _lock: threading.RLock + _fileobj: IO[bytes] | None + + def __init__( + self, + path: Path | str | IO[bytes], + *, + mode: ZipStoreAccessModeLiteral = "r", + read_only: bool | None = None, + compression: int = zipfile.ZIP_STORED, + allowZip64: bool = True, + ) -> None: + if read_only is None: + read_only = mode == "r" + + super().__init__(read_only=read_only) + + if isinstance(path, str): + path = Path(path) + if isinstance(path, Path): + self.path = path # root? + self._fileobj = None + else: + self.path = None + if not isinstance(path, io.IOBase): + if not all( + callable(getattr(path, attr, None)) for attr in ("read", "seek", "tell") + ): + raise TypeError( + f"expected a path or an open binary file object supporting " + f"read/seek/tell, got {type(path).__name__}" + ) + if mode != "r": + raise TypeError( + f"a file object that is not an io.IOBase instance can only be " + f"opened for reading (mode='r', got mode={mode!r})" + ) + # e.g. an obstore ReadableFile: readable and seekable, but + # not an io object and reads may not return bytes + path = io.BufferedReader(_RawReaderAdapter(path)) + self._fileobj = path + + self._zmode = mode + self.compression = compression + self.allowZip64 = allowZip64 + + 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, + ) + + self._is_open = True + + async def _open(self) -> None: + self._sync_open() + + def __getstate__(self) -> dict[str, Any]: + 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. + raise TypeError( + "cannot pickle a ZipStore backed by a file-like object; " + "construct the store from a path instead" + ) + # We need a copy to not modify the state of the original store + state = self.__dict__.copy() + for attr in ["_zf", "_lock"]: + state.pop(attr, None) + return state + + def __setstate__(self, state: dict[str, Any]) -> None: + self.__dict__ = state + self._is_open = False + self._sync_open() + + def close(self) -> None: + # docstring inherited + if not self._is_open: + return + super().close() + with self._lock: + self._zf.close() + + async def clear(self) -> None: + # docstring inherited + with self._lock: + self._check_writable() + if self.path is None: + raise NotImplementedError( + "clear() is not supported for a ZipStore backed by a file-like object" + ) + self._zf.close() + os.remove(self.path) + self._zf = zipfile.ZipFile( + self.path, mode="w", compression=self.compression, allowZip64=self.allowZip64 + ) + + def __str__(self) -> str: + if self.path is None: + return f"zip://{self._fileobj!r}" + return f"zip://{self.path}" + + def __repr__(self) -> str: + return f"ZipStore('{self}')" + + def __eq__(self, other: object) -> bool: + return ( + isinstance(other, type(self)) + and self.path == other.path + and self._fileobj is other._fileobj + ) + + def _get( + self, + key: str, + 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 + if byte_range is None: + return prototype.buffer.from_bytes(f.read()) + elif isinstance(byte_range, RangeByteRequest): + f.seek(byte_range.start) + return prototype.buffer.from_bytes(f.read(byte_range.end - f.tell())) + size = f.seek(0, os.SEEK_END) + if isinstance(byte_range, OffsetByteRequest): + f.seek(byte_range.offset) + elif isinstance(byte_range, SuffixByteRequest): + f.seek(max(0, size - byte_range.suffix)) + else: + raise TypeError(f"Unexpected byte_range, got {byte_range}.") + return prototype.buffer.from_bytes(f.read()) + except KeyError: + return None + + async def get( + self, + key: str, + prototype: BufferPrototype, + byte_range: ByteRequest | None = None, + ) -> Buffer | None: + # docstring inherited + assert isinstance(key, str) + + with self._lock: + return self._get(key, prototype=prototype, byte_range=byte_range) + + async def get_partial_values( + self, + prototype: BufferPrototype, + key_ranges: Iterable[tuple[str, ByteRequest | None]], + ) -> list[Buffer | None]: + # docstring inherited + out = [] + with self._lock: + for key, byte_range in key_ranges: + out.append(self._get(key, prototype=prototype, byte_range=byte_range)) + 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 + if keyinfo.filename[-1] == os.sep: + keyinfo.external_attr = 0o40775 << 16 # drwxrwxr-x + keyinfo.external_attr |= 0x10 # MS-DOS directory flag + else: + keyinfo.external_attr = 0o644 << 16 # ?rw-r--r-- + self._zf.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() + assert isinstance(key, str) + if not isinstance(value, Buffer): + raise TypeError( + f"ZipStore.set(): `value` must be a Buffer instance. Got an instance of {type(value)} instead." + ) + with self._lock: + self._set(key, value) + + async def set_if_not_exists(self, key: str, value: Buffer) -> None: + self._check_writable() + with self._lock: + members = self._zf.namelist() + if key not in members: + self._set(key, value) + + async def delete_dir(self, prefix: str) -> None: + # only raise NotImplementedError if any keys are found + self._check_writable() + if prefix != "" and not prefix.endswith("/"): + prefix += "/" + async for _ in self.list_prefix(prefix): + raise NotImplementedError + + async def delete(self, key: str) -> None: + # docstring inherited + # we choose to only raise NotImplementedError here if the key exists + # this allows the array/group APIs to avoid the overhead of existence checks + self._check_writable() + if await self.exists(key): + raise NotImplementedError + + 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) + 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(): + yield key + + async def list_prefix(self, prefix: str) -> AsyncIterator[str]: + # docstring inherited + async for key in self.list(): + if key.startswith(prefix): + yield key + + 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() + seen = set() + if prefix == "": + keys_unique = {k.split("/")[0] for k in keys} + for key in keys_unique: + if key not in seen: + seen.add(key) + yield key + else: + for key in keys: + if key.startswith(f"{prefix}/") and key.strip("/") != prefix: + k = key.removeprefix(f"{prefix}/").split("/")[0] + if k not in seen: + seen.add(k) + yield k + + async def move(self, path: Path | str) -> None: + """ + Move the store to another path. + """ + if self.path is None: + raise NotImplementedError( + "move() is not supported for a ZipStore backed by a file-like object" + ) + if isinstance(path, str): + path = Path(path) + self.close() + os.makedirs(path.parent, exist_ok=True) + shutil.move(self.path, path) + self.path = path + await self._open() diff --git a/packages/zarr-storage/src/zarr_storage/legacy/experimental/__init__.py b/packages/zarr-storage/src/zarr_storage/legacy/experimental/__init__.py new file mode 100644 index 0000000000..8068348f71 --- /dev/null +++ b/packages/zarr-storage/src/zarr_storage/legacy/experimental/__init__.py @@ -0,0 +1 @@ +"""Experimental legacy storage implementations.""" diff --git a/packages/zarr-storage/src/zarr_storage/legacy/experimental/cache_store.py b/packages/zarr-storage/src/zarr_storage/legacy/experimental/cache_store.py new file mode 100644 index 0000000000..f0161f36c8 --- /dev/null +++ b/packages/zarr-storage/src/zarr_storage/legacy/experimental/cache_store.py @@ -0,0 +1,461 @@ +from __future__ import annotations + +import asyncio +import logging +import time +from collections import OrderedDict +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Literal, Self + +from zarr_storage.legacy._abc import ByteRequest, Store +from zarr_storage.legacy._wrapper import WrapperStore + +logger = logging.getLogger(__name__) + +if TYPE_CHECKING: + from zarr.core.buffer.core import Buffer, BufferPrototype + +# A cache entry identifier. Plain ``str`` for full-key entries that live in +# the Store-backed cache; ``(str, ByteRequest)`` for byte-range entries that +# live in the in-memory range cache. +_CacheEntryKey = str | tuple[str, ByteRequest] + + +@dataclass(slots=True) +class _CacheState: + cache_order: OrderedDict[_CacheEntryKey, None] = field(default_factory=OrderedDict) + current_size: int = 0 + key_sizes: dict[_CacheEntryKey, int] = field(default_factory=dict) + lock: asyncio.Lock = field(default_factory=asyncio.Lock) + hits: int = 0 + misses: int = 0 + evictions: int = 0 + key_insert_times: dict[_CacheEntryKey, float] = field(default_factory=dict) + range_cache: dict[str, dict[ByteRequest, Buffer]] = field(default_factory=dict) + + +class CacheStore(WrapperStore[Store]): + """ + A dual-store caching implementation for Zarr stores. + + This cache wraps any Store implementation and uses a separate Store instance + as the cache backend. This provides persistent caching capabilities with + time-based expiration, size-based eviction, and flexible cache storage options. + + Full-key reads are cached in the Store-backed cache. Byte-range reads are + cached in a separate in-memory dictionary so that partial reads never + pollute the filesystem (or other persistent backend). Both caches share + the same ``max_size`` budget and LRU eviction policy. + + Parameters + ---------- + store : Store + The underlying store to wrap with caching + cache_store : Store + The store to use for caching (can be any Store implementation that + supports deletes) + max_age_seconds : int or "infinity", optional + Maximum age of cached entries in seconds. The string "infinity" means + entries never expire. Default is "infinity". + max_size : int | None, optional + Maximum size of the cache in bytes. When exceeded, least recently used + items are evicted. None means unlimited size. Default is None. + Note: Individual values larger than max_size will not be cached. + cache_set_data : bool, optional + Whether to cache data when it's written to the store. Default is True. + + Examples + -------- + ```python + import zarr + from zarr_storage.legacy import MemoryStore + from zarr_storage.legacy.experimental.cache_store import CacheStore + + # Create a cached store + source_store = MemoryStore() + cache_store = MemoryStore() + cached_store = CacheStore( + store=source_store, + cache_store=cache_store, + max_age_seconds=60, + max_size=1024*1024 + ) + + # Use it like any other store + array = zarr.create(shape=(100,), store=cached_store) + array[:] = 42 + ``` + + """ + + _cache: Store + max_age_seconds: int | Literal["infinity"] + max_size: int | None + cache_set_data: bool + _state: _CacheState + + def __init__( + self, + store: Store, + *, + cache_store: Store, + max_age_seconds: int | str = "infinity", + max_size: int | None = None, + cache_set_data: bool = True, + ) -> None: + super().__init__(store) + + if not cache_store.supports_deletes: + msg = ( + f"The provided cache store {cache_store} does not support deletes. " + "The cache_store must support deletes for CacheStore to function properly." + ) + raise ValueError(msg) + + self._cache = cache_store + # Validate and set max_age_seconds + if isinstance(max_age_seconds, str): + if max_age_seconds != "infinity": + raise ValueError("max_age_seconds string value must be 'infinity'") + self.max_age_seconds = "infinity" + else: + self.max_age_seconds = max_age_seconds + self.max_size = max_size + self.cache_set_data = cache_set_data + self._state = _CacheState() + + def _with_store(self, store: Store) -> Self: + # Cannot support this operation because it would share a cache, but have a new store + # So cache keys would conflict + raise NotImplementedError("CacheStore does not support this operation.") + + def with_read_only(self, read_only: bool = False) -> Self: + # Create a new cache store that shares the same cache and mutable state + store = type(self)( + store=self._store.with_read_only(read_only), + cache_store=self._cache, + max_age_seconds=self.max_age_seconds, + max_size=self.max_size, + cache_set_data=self.cache_set_data, + ) + store._state = self._state + return store + + def _is_key_fresh(self, entry_key: _CacheEntryKey) -> bool: + """Check if a cached entry is still fresh based on max_age_seconds. + + Uses monotonic time for accurate elapsed time measurement. + """ + if self.max_age_seconds == "infinity": + return True + now = time.monotonic() + elapsed = now - self._state.key_insert_times.get(entry_key, 0) + return elapsed < self.max_age_seconds + + async def _accommodate_value(self, value_size: int) -> None: + """Ensure there is enough space in the cache for a new value. + + Must be called while holding self._state.lock. + """ + if self.max_size is None: + return + + # Remove least recently used items until we have enough space + while self._state.current_size + value_size > self.max_size and self._state.cache_order: + # Get the least recently used key (first in OrderedDict) + lru_key = next(iter(self._state.cache_order)) + await self._evict_key(lru_key) + + async def _evict_key(self, entry_key: _CacheEntryKey) -> None: + """Evict a cache entry. + + Must be called while holding self._state.lock. + + For ``str`` keys the entry is deleted from the Store-backed cache. + For ``(str, ByteRequest)`` keys the entry is removed from the + in-memory range cache. + """ + key_size = self._state.key_sizes.get(entry_key, 0) + + if isinstance(entry_key, str): + await self._cache.delete(entry_key) + else: + base_key, byte_range = entry_key + per_key = self._state.range_cache.get(base_key) + if per_key is not None: + per_key.pop(byte_range, None) + if not per_key: + del self._state.range_cache[base_key] + + self._state.cache_order.pop(entry_key, None) + self._state.key_insert_times.pop(entry_key, None) + self._state.key_sizes.pop(entry_key, None) + self._state.current_size = max(0, self._state.current_size - key_size) + self._state.evictions += 1 + + async def _track_entry(self, entry_key: _CacheEntryKey, value: Buffer) -> bool: + """Register *entry_key* in the shared size / LRU tracking. + + Returns ``True`` if the entry was tracked, ``False`` if the value + exceeds ``max_size`` and was skipped. Callers should roll back any + data they already stored when this returns ``False``. + + This method holds the lock for the entire operation to ensure atomicity. + """ + value_size = len(value) + + # Check if value exceeds max size + if self.max_size is not None and value_size > self.max_size: + return False + + async with self._state.lock: + # If key already exists, subtract old size first + if entry_key in self._state.key_sizes: + old_size = self._state.key_sizes[entry_key] + self._state.current_size -= old_size + + # Make room for the new value + await self._accommodate_value(value_size) + + # Update tracking atomically + self._state.cache_order[entry_key] = None + self._state.current_size += value_size + self._state.key_sizes[entry_key] = value_size + self._state.key_insert_times[entry_key] = time.monotonic() + + return True + + async def _update_access_order(self, entry_key: _CacheEntryKey) -> None: + """Update the access order for LRU tracking.""" + if entry_key in self._state.cache_order: + async with self._state.lock: + self._state.cache_order.move_to_end(entry_key) + + def _remove_from_tracking(self, entry_key: _CacheEntryKey) -> None: + """Remove an entry from all tracking structures. + + Must be called while holding self._state.lock. + """ + self._state.cache_order.pop(entry_key, None) + self._state.key_insert_times.pop(entry_key, None) + self._state.key_sizes.pop(entry_key, None) + + def _invalidate_range_entries(self, key: str) -> None: + """Remove all byte-range entries for *key* from the range cache and tracking. + + Must be called while holding self._state.lock. + """ + per_key = self._state.range_cache.pop(key, None) + if per_key is not None: + for byte_range in per_key: + entry_key: _CacheEntryKey = (key, byte_range) + entry_size = self._state.key_sizes.pop(entry_key, 0) + self._state.cache_order.pop(entry_key, None) + self._state.key_insert_times.pop(entry_key, None) + self._state.current_size = max(0, self._state.current_size - entry_size) + + # ------------------------------------------------------------------ + # get helpers + # ------------------------------------------------------------------ + + async def _cache_miss( + self, key: str, byte_range: ByteRequest | None, result: Buffer | None + ) -> None: + """Handle a cache miss by storing or cleaning up after a source-store fetch.""" + if result is None: + if byte_range is None: + await self._cache.delete(key) + async with self._state.lock: + self._remove_from_tracking(key) + else: + entry_key: _CacheEntryKey = (key, byte_range) + async with self._state.lock: + per_key = self._state.range_cache.get(key) + if per_key is not None: + per_key.pop(byte_range, None) + if not per_key: + del self._state.range_cache[key] + self._remove_from_tracking(entry_key) + else: + if byte_range is None: + await self._cache.set(key, result) + await self._track_entry(key, result) + else: + entry_key = (key, byte_range) + self._state.range_cache.setdefault(key, {})[byte_range] = result + tracked = await self._track_entry(entry_key, result) + if not tracked: + # Value too large for the cache — roll back the insertion + per_key = self._state.range_cache.get(key) + if per_key is not None: + per_key.pop(byte_range, None) + if not per_key: + del self._state.range_cache[key] + + async def _get_try_cache( + self, key: str, prototype: BufferPrototype, byte_range: ByteRequest | None = None + ) -> Buffer | None: + """Try to get data from cache first, falling back to source store.""" + if byte_range is None: + # Full-key read — use Store-backed cache + maybe_cached = await self._cache.get(key, prototype) + if maybe_cached is not None: + self._state.hits += 1 + await self._update_access_order(key) + return maybe_cached + else: + # Byte-range read — use in-memory range cache + entry_key: _CacheEntryKey = (key, byte_range) + per_key = self._state.range_cache.get(key) + if per_key is not None: + cached_buf = per_key.get(byte_range) + if cached_buf is not None: + self._state.hits += 1 + await self._update_access_order(entry_key) + return cached_buf + + # Cache miss — fetch from source store + self._state.misses += 1 + result = await super().get(key, prototype, byte_range) + await self._cache_miss(key, byte_range, result) + return result + + async def _get_no_cache( + self, key: str, prototype: BufferPrototype, byte_range: ByteRequest | None = None + ) -> Buffer | None: + """Get data directly from source store and update cache.""" + self._state.misses += 1 + result = await super().get(key, prototype, byte_range) + await self._cache_miss(key, byte_range, result) + return result + + @property + def _supports_sync_io(self) -> bool: + # The caching logic lives only in the async get/set/delete overrides; + # the sync methods inherited from `WrapperStore` delegate straight to + # the source store, so a sync-capable consumer (the fused codec + # pipeline) would write and delete around the cache, leaving stale + # entries that later async reads serve as current data. Opt out of + # sync IO until the sync surface is cache-aware. + return False + + async def get( + self, + key: str, + prototype: BufferPrototype, + byte_range: ByteRequest | None = None, + ) -> Buffer | None: + """ + Retrieve data from the store, using cache when appropriate. + + Parameters + ---------- + key : str + The key to retrieve + prototype : BufferPrototype + Buffer prototype for creating the result buffer + byte_range : ByteRequest, optional + Byte range to retrieve + + Returns + ------- + Buffer | None + The retrieved data, or None if not found + """ + entry_key: _CacheEntryKey = (key, byte_range) if byte_range is not None else key + if not self._is_key_fresh(entry_key): + return await self._get_no_cache(key, prototype, byte_range) + else: + return await self._get_try_cache(key, prototype, byte_range) + + async def set(self, key: str, value: Buffer) -> None: + """ + Store data in the underlying store and optionally in cache. + + Parameters + ---------- + key : str + The key to store under + value : Buffer + The data to store + """ + await super().set(key, value) + # Invalidate all cached byte-range entries (source data changed) + async with self._state.lock: + self._invalidate_range_entries(key) + if self.cache_set_data: + await self._cache.set(key, value) + await self._track_entry(key, value) + else: + await self._cache.delete(key) + async with self._state.lock: + self._remove_from_tracking(key) + + async def delete(self, key: str) -> None: + """ + Delete data from both the underlying store and cache. + + Parameters + ---------- + key : str + The key to delete + """ + await super().delete(key) + # Invalidate all cached byte-range entries + async with self._state.lock: + self._invalidate_range_entries(key) + await self._cache.delete(key) + async with self._state.lock: + self._remove_from_tracking(key) + + def cache_info(self) -> dict[str, Any]: + """Return information about the cache state.""" + return { + "cache_store_type": type(self._cache).__name__, + "max_age_seconds": "infinity" + if self.max_age_seconds == "infinity" + else self.max_age_seconds, + "max_size": self.max_size, + "current_size": self._state.current_size, + "cache_set_data": self.cache_set_data, + "tracked_keys": len(self._state.key_insert_times), + "cached_keys": len(self._state.cache_order), + } + + def cache_stats(self) -> dict[str, Any]: + """Return cache performance statistics.""" + total_requests = self._state.hits + self._state.misses + hit_rate = self._state.hits / total_requests if total_requests > 0 else 0.0 + return { + "hits": self._state.hits, + "misses": self._state.misses, + "evictions": self._state.evictions, + "total_requests": total_requests, + "hit_rate": hit_rate, + } + + async def clear_cache(self) -> None: + """Clear all cached data and tracking information.""" + # Clear the cache store if it supports clear + if hasattr(self._cache, "clear"): + await self._cache.clear() + + # Reset tracking + async with self._state.lock: + self._state.key_insert_times.clear() + self._state.cache_order.clear() + self._state.key_sizes.clear() + self._state.range_cache.clear() + self._state.current_size = 0 + + def __repr__(self) -> str: + """Return string representation of the cache store.""" + return ( + f"{self.__class__.__name__}(" + f"store={self._store!r}, " + f"cache_store={self._cache!r}, " + f"max_age_seconds={self.max_age_seconds}, " + f"max_size={self.max_size}, " + f"current_size={self._state.current_size}, " + f"cached_keys={len(self._state.cache_order)})" + ) diff --git a/packages/zarr-storage/src/zarr_storage/py.typed b/packages/zarr-storage/src/zarr_storage/py.typed new file mode 100644 index 0000000000..e69de29bb2 diff --git a/packages/zarr-storage/src/zarr_storage/testing/__init__.py b/packages/zarr-storage/src/zarr_storage/testing/__init__.py new file mode 100644 index 0000000000..62a1545883 --- /dev/null +++ b/packages/zarr-storage/src/zarr_storage/testing/__init__.py @@ -0,0 +1,6 @@ +"""Reusable conformance tests; install zarr-storage[testing].""" + +from zarr_storage.testing.store import StoreTests +from zarr_storage.testing.utils import assert_bytes_equal + +__all__ = ["StoreTests", "assert_bytes_equal"] diff --git a/packages/zarr-storage/src/zarr_storage/testing/buffer.py b/packages/zarr-storage/src/zarr_storage/testing/buffer.py new file mode 100644 index 0000000000..0de5415ec3 --- /dev/null +++ b/packages/zarr-storage/src/zarr_storage/testing/buffer.py @@ -0,0 +1,110 @@ +# mypy: ignore-errors +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, Literal + +import numpy as np +import numpy.typing as npt +from zarr.core.buffer import Buffer, BufferPrototype, cpu + +from zarr_storage.legacy import MemoryStore + +if TYPE_CHECKING: + from collections.abc import Iterable + from typing import Self + + from zarr_storage.legacy._abc import ByteRequest + + +__all__ = [ + "NDBufferUsingTestNDArrayLike", + "StoreExpectingTestBuffer", + "TestBuffer", +] + + +class TestNDArrayLike(np.ndarray): + """An example of an ndarray-like class""" + + __test__ = False + + +class TestBuffer(cpu.Buffer): + """Example of a custom Buffer that handles ArrayLike""" + + __test__ = False + + +class NDBufferUsingTestNDArrayLike(cpu.NDBuffer): + """Example of a custom NDBuffer that handles MyNDArrayLike""" + + @classmethod + def create( + cls, + *, + shape: Iterable[int], + dtype: npt.DTypeLike, + order: Literal["C", "F"] = "C", + fill_value: Any | None = None, + ) -> Self: + """Overwrite `NDBuffer.create` to create a TestNDArrayLike instance""" + ret = cls(TestNDArrayLike(shape=shape, dtype=dtype, order=order)) + if fill_value is not None: + ret.fill(fill_value) + return ret + + @classmethod + def empty( + cls, + shape: tuple[int, ...], + dtype: npt.DTypeLike, + order: Literal["C", "F"] = "C", + ) -> Self: + return super(cpu.NDBuffer, cls).empty(shape=shape, dtype=dtype, order=order) + + +class StoreExpectingTestBuffer(MemoryStore): + """Example of a custom Store that expect MyBuffer for all its non-metadata + + We assume that keys containing "json" is metadata + """ + + async def set(self, key: str, value: Buffer, byte_range: tuple[int, int] | None = None) -> None: + if "json" not in key: + assert isinstance(value, TestBuffer) + await super().set(key, value, byte_range) + + def set_sync(self, key: str, value: Buffer) -> None: + # Synchronous counterpart of `set`, used by FusedCodecPipeline. Mirror the + # same buffer-type guard so the invariant holds whichever pipeline writes. + if "json" not in key: + assert isinstance(value, TestBuffer) + super().set_sync(key, value) + + async def get( + self, + key: str, + prototype: BufferPrototype, + byte_range: tuple[int, int | None] | None = None, + ) -> Buffer | None: + if "json" not in key: + assert prototype.buffer is TestBuffer + ret = await super().get(key=key, prototype=prototype, byte_range=byte_range) + if ret is not None: + assert isinstance(ret, prototype.buffer) + return ret + + def get_sync( + self, + key: str, + *, + prototype: BufferPrototype | None = None, + byte_range: ByteRequest | None = None, + ) -> Buffer | None: + # Synchronous counterpart of `get`, used by FusedCodecPipeline. + if "json" not in key and prototype is not None: + assert prototype.buffer is TestBuffer + ret = super().get_sync(key=key, prototype=prototype, byte_range=byte_range) + if ret is not None and prototype is not None: + assert isinstance(ret, prototype.buffer) + return ret diff --git a/packages/zarr-storage/src/zarr_storage/testing/stateful.py b/packages/zarr-storage/src/zarr_storage/testing/stateful.py new file mode 100644 index 0000000000..8bb7202edc --- /dev/null +++ b/packages/zarr-storage/src/zarr_storage/testing/stateful.py @@ -0,0 +1,671 @@ +import builtins +import functools +from collections.abc import Callable, Iterable +from typing import Any, cast + +import hypothesis.extra.numpy as npst +import hypothesis.strategies as st +import numpy as np +import zarr +from hypothesis import assume, note +from hypothesis.stateful import ( + RuleBasedStateMachine, + initialize, + invariant, + precondition, + rule, +) +from hypothesis.strategies import DataObject +from zarr import Array +from zarr.codecs.bytes import BytesCodec +from zarr.core.buffer import Buffer, BufferPrototype, cpu, default_buffer_prototype +from zarr.core.sync import SyncMixin + +from zarr_storage.legacy import LocalStore, MemoryStore +from zarr_storage.legacy._abc import ( + OffsetByteRequest, + RangeByteRequest, + Store, + SuffixByteRequest, +) +from zarr_storage.testing.strategies import ( + arrays as zarr_arrays, +) +from zarr_storage.testing.strategies import ( + basic_indices, + chunk_paths, + key_ranges, + node_names, + orthogonal_indices, +) +from zarr_storage.testing.strategies import keys as zarr_keys + +MAX_BINARY_SIZE = 100 + + +def with_frequency[F: Callable[..., Any]](frequency: float) -> Callable[[F], F]: + """This needs to be deterministic for hypothesis replaying""" + + def decorator(func: F) -> F: + counter_attr = f"__{func.__name__}_counter" + + @functools.wraps(func) + def wrapper(*args: Any, **kwargs: Any) -> Any: + return func(*args, **kwargs) + + @precondition + def frequency_check(f: Any) -> Any: + if not hasattr(f, counter_attr): + setattr(f, counter_attr, 0) + + current_count = getattr(f, counter_attr) + 1 + setattr(f, counter_attr, current_count) + + return (current_count * frequency) % 1.0 >= (1.0 - frequency) + + return cast(F, frequency_check(wrapper)) + + return decorator + + +def split_prefix_name(path: str) -> tuple[str, str]: + split = path.rsplit("/", maxsplit=1) + if len(split) > 1: + prefix, name = split + else: + prefix = "" + (name,) = split + return prefix, name + + +class ZarrHierarchyStateMachine(SyncMixin, RuleBasedStateMachine): + """ + This state machine models operations that modify a zarr store's + hierarchy. That is, user actions that modify arrays/groups as well + as list operations. It is intended to be used by external stores, and + compares their results to a MemoryStore that is assumed to be perfect. + """ + + def __init__(self, store: Store) -> None: + super().__init__() + + self.store = store + + self.model = MemoryStore() + zarr.group(store=self.model) + + # Track state of the hierarchy, these should contain fully qualified paths + self.all_groups: set[str] = set() + self.all_arrays: set[str] = set() + + @initialize() + def init_store(self) -> None: + # This lets us reuse the fixture provided store. + self._sync(self.store.clear()) + zarr.group(store=self.store) + + def can_add(self, path: str) -> bool: + return path not in self.all_groups and path not in self.all_arrays + + # -------------------- store operations ----------------------- + @rule(name=node_names, data=st.data()) + def add_group(self, name: str, data: DataObject) -> None: + # Handle possible case-insensitive file systems (e.g. MacOS) + if isinstance(self.store, LocalStore): + name = name.lower() + if self.all_groups: + parent = data.draw(st.sampled_from(sorted(self.all_groups)), label="Group parent") + else: + parent = "" + path = f"{parent}/{name}".lstrip("/") + assume(self.can_add(path)) + note(f"Adding group: path='{path}'") + self.all_groups.add(path) + zarr.group(store=self.store, path=path) + zarr.group(store=self.model, path=path) + + @rule(data=st.data(), name=node_names) + def add_array(self, data: DataObject, name: str) -> None: + # Handle possible case-insensitive file systems (e.g. MacOS) + if isinstance(self.store, LocalStore): + name = name.lower() + if self.all_groups: + parent = data.draw(st.sampled_from(sorted(self.all_groups)), label="Array parent") + else: + parent = "" + # TODO: support creating deeper paths + # TODO: support overwriting potentially by just skipping `self.can_add` + path = f"{parent}/{name}".lstrip("/") + assume(self.can_add(path)) + + # Generate array on the model store using the arrays strategy + a = data.draw( + zarr_arrays( + stores=st.just(self.model), + paths=st.just(parent), + array_names=st.just(name), + zarr_formats=st.just(3), + compressors=st.just(BytesCodec()), + open_mode="a", + ), + label="generated array", + ) + note(f"Adding array: path='{path}' shape={a.shape} chunks={a.metadata.chunk_grid}") + + # Recreate the same array in the store under test. + # The data is copied here rather than by `write_data=True`, + # whose shard-wise copy does not support rectilinear chunk grids. + arr = zarr.from_array( + self.store, + data=a, + name=path, + write_data=False, + ) + arr[:] = a[:] + self.all_arrays.add(path) + + @rule() + @with_frequency(0.25) + def clear(self) -> None: + note("clearing") + import zarr + + self._sync(self.store.clear()) + self._sync(self.model.clear()) + + assert self._sync(self.store.is_empty("/")) + assert self._sync(self.model.is_empty("/")) + + self.all_groups.clear() + self.all_arrays.clear() + + zarr.group(store=self.store) + zarr.group(store=self.model) + + # TODO: MemoryStore is broken? + # assert not self._sync(self.store.is_empty("/")) + # assert not self._sync(self.model.is_empty("/")) + + def draw_directory(self, data: DataObject) -> str: + group_st = st.sampled_from(sorted(self.all_groups)) if self.all_groups else st.nothing() + array_st = st.sampled_from(sorted(self.all_arrays)) if self.all_arrays else st.nothing() + array_or_group = data.draw(st.one_of(group_st, array_st)) + if data.draw(st.booleans()) and array_or_group in self.all_arrays: + arr = zarr.open_array(path=array_or_group, store=self.model) + path = data.draw( + st.one_of( + st.sampled_from([array_or_group]), + chunk_paths(ndim=arr.ndim, numblocks=arr.cdata_shape).map( + lambda x: f"{array_or_group}/c/" + ), + ) + ) + else: + path = array_or_group + return path + + @precondition(lambda self: bool(self.all_groups)) + @rule(data=st.data()) + def check_list_dir(self, data: DataObject) -> None: + path = self.draw_directory(data) + note(f"list_dir for {path=!r}") + # Consider .list_dir("path/to/array") for an array with a single chunk. + # The MemoryStore model will return `"c", "zarr.json"` only if the chunk exists + # If that chunk was deleted, then `"c"` is not returned. + # LocalStore will not have this behaviour :/ + # There are similar consistency issues with delete_dir("/path/to/array/c/0/0") + assume(not isinstance(self.store, LocalStore)) + model_ls = sorted(self._sync_iter(self.model.list_dir(path))) + store_ls = sorted(self._sync_iter(self.store.list_dir(path))) + assert model_ls == store_ls, (model_ls, store_ls) + + @precondition(lambda self: bool(self.all_arrays)) + @rule(data=st.data()) + def delete_chunk(self, data: DataObject) -> None: + array = data.draw(st.sampled_from(sorted(self.all_arrays))) + arr = zarr.open_array(path=array, store=self.model) + chunk_path = data.draw(chunk_paths(ndim=arr.ndim, numblocks=arr.cdata_shape, subset=False)) + path = f"{array}/c/{chunk_path}" + note(f"deleting chunk {path=!r}") + self._sync(self.model.delete(path)) + self._sync(self.store.delete(path)) + + @precondition(lambda self: bool(self.all_arrays)) + @rule(data=st.data()) + def check_array(self, data: DataObject) -> None: + path = data.draw(st.sampled_from(sorted(self.all_arrays))) + actual = zarr.open_array(self.store, path=path)[:] + expected = zarr.open_array(self.model, path=path)[:] + np.testing.assert_equal(actual, expected) + + @precondition(lambda self: bool(self.all_arrays)) + @rule(data=st.data()) + def overwrite_array_basic_indexing(self, data: DataObject) -> None: + array = data.draw(st.sampled_from(sorted(self.all_arrays))) + model_array = zarr.open_array(path=array, store=self.model) + store_array = zarr.open_array(path=array, store=self.store) + slicer = data.draw(basic_indices(shape=model_array.shape)) + note(f"overwriting array with basic indexer: {slicer=}") + new_data = data.draw( + npst.arrays(shape=np.shape(model_array[slicer]), dtype=model_array.dtype) + ) + model_array[slicer] = new_data + store_array[slicer] = new_data + + @precondition(lambda self: bool(self.all_arrays)) + @rule(data=st.data()) + def overwrite_array_orthogonal_indexing(self, data: DataObject) -> None: + array = data.draw(st.sampled_from(sorted(self.all_arrays))) + model_array = zarr.open_array(path=array, store=self.model) + store_array = zarr.open_array(path=array, store=self.store) + indexer, _ = data.draw(orthogonal_indices(shape=model_array.shape)) + note(f"overwriting array orthogonal {indexer=}") + new_data = data.draw( + npst.arrays(shape=model_array.oindex[indexer].shape, dtype=model_array.dtype) # type: ignore[union-attr] + ) + model_array.oindex[indexer] = new_data + store_array.oindex[indexer] = new_data + + @precondition(lambda self: bool(self.all_arrays)) + @rule(data=st.data()) + def resize_array(self, data: DataObject) -> None: + array = data.draw(st.sampled_from(sorted(self.all_arrays))) + model_array = zarr.open_array(path=array, store=self.model) + store_array = zarr.open_array(path=array, store=self.store) + ndim = model_array.ndim + new_shape = tuple( + 0 if oldsize == 0 else newsize + for newsize, oldsize in zip( + data.draw(npst.array_shapes(max_dims=ndim, min_dims=ndim, min_side=0)), + model_array.shape, + strict=True, + ) + ) + + note(f"resizing array from {model_array.shape} to {new_shape}") + model_array.resize(new_shape) + store_array.resize(new_shape) + + @precondition(lambda self: bool(self.all_arrays) or bool(self.all_groups)) + @rule(data=st.data()) + def delete_dir(self, data: DataObject) -> None: + path = self.draw_directory(data) + note(f"delete_dir with {path=!r}") + self._sync(self.model.delete_dir(path)) + self._sync(self.store.delete_dir(path)) + + matches = set() + for node in self.all_groups | self.all_arrays: + if node == path or node.startswith(path + "/"): + matches.add(node) + self.all_groups = self.all_groups - matches + self.all_arrays = self.all_arrays - matches + + # @precondition(lambda self: bool(self.all_groups)) + # @precondition(lambda self: bool(self.all_arrays)) + # @rule(data=st.data()) + # def move_array(self, data): + # array_path = data.draw(st.sampled_from(self.all_arrays), label="Array move source") + # to_group = data.draw(st.sampled_from(self.all_groups), label="Array move destination") + + # # fixme renaming to self? + # array_name = os.path.basename(array_path) + # assume(self.model.can_add(to_group, array_name)) + # new_path = f"{to_group}/{array_name}".lstrip("/") + # note(f"moving array '{array_path}' -> '{new_path}'") + # self.model.rename(array_path, new_path) + # self.repo.store.rename(array_path, new_path) + + # @precondition(lambda self: len(self.all_groups) >= 2) + # @rule(data=st.data()) + # def move_group(self, data): + # from_group = data.draw(st.sampled_from(self.all_groups), label="Group move source") + # to_group = data.draw(st.sampled_from(self.all_groups), label="Group move destination") + # assume(not to_group.startswith(from_group)) + + # from_group_name = os.path.basename(from_group) + # assume(self.model.can_add(to_group, from_group_name)) + # # fixme renaming to self? + # new_path = f"{to_group}/{from_group_name}".lstrip("/") + # note(f"moving group '{from_group}' -> '{new_path}'") + # self.model.rename(from_group, new_path) + # self.repo.store.rename(from_group, new_path) + + @precondition(lambda self: self.store.supports_deletes) + @precondition(lambda self: len(self.all_arrays) >= 1) + @rule(data=st.data()) + def delete_array_using_del(self, data: DataObject) -> None: + array_path = data.draw( + st.sampled_from(sorted(self.all_arrays)), label="Array deletion target" + ) + prefix, array_name = split_prefix_name(array_path) + note(f"Deleting array '{array_path}' ({prefix=!r}, {array_name=!r}) using del") + for store in [self.model, self.store]: + group = zarr.open_group(path=prefix, store=store) + group[array_name] # check that it exists + del group[array_name] + self.all_arrays.remove(array_path) + + @precondition(lambda self: self.store.supports_deletes) + @precondition(lambda self: bool(self.all_groups)) + @rule(data=st.data()) + def delete_group_using_del(self, data: DataObject) -> None: + group_path = data.draw( + st.sampled_from(sorted(self.all_groups)), + label="Group deletion target", + ) + prefix, group_name = split_prefix_name(group_path) + note(f"Deleting group '{group_path=!r}', {prefix=!r}, {group_name=!r} using delete") + members = zarr.open_group(store=self.model, path=group_path).members(max_depth=None) + for _, obj in members: + if isinstance(obj, Array): + self.all_arrays.remove(obj.path) + else: + self.all_groups.remove(obj.path) + for store in [self.store, self.model]: + group = zarr.open_group(store=store, path=prefix) + group[group_name] # check that it exists + del group[group_name] + self.all_groups.remove(group_path) + + # # --------------- assertions ----------------- + # def check_group_arrays(self, group): + # # note(f"Checking arrays of '{group}'") + # g1 = self.model.get_group(group) + # g2 = zarr.open_group(path=group, mode="r", store=self.repo.store) + # model_arrays = sorted(g1.arrays(), key=itemgetter(0)) + # our_arrays = sorted(g2.arrays(), key=itemgetter(0)) + # for (n1, a1), (n2, a2) in zip_longest(model_arrays, our_arrays): + # assert n1 == n2 + # assert_array_equal(a1, a2) + + # def check_subgroups(self, group_path): + # g1 = self.model.get_group(group_path) + # g2 = zarr.open_group(path=group_path, mode="r", store=self.repo.store) + # g1_children = [name for (name, _) in g1.groups()] + # g2_children = [name for (name, _) in g2.groups()] + # # note(f"Checking {len(g1_children)} subgroups of group '{group_path}'") + # assert g1_children == g2_children + + # def check_list_prefix_from_group(self, group): + # prefix = f"meta/root/{group}" + # model_list = sorted(self.model.list_prefix(prefix)) + # al_list = sorted(self.repo.store.list_prefix(prefix)) + # # note(f"Checking {len(model_list)} keys under '{prefix}'") + # assert model_list == al_list + + # prefix = f"data/root/{group}" + # model_list = sorted(self.model.list_prefix(prefix)) + # al_list = sorted(self.repo.store.list_prefix(prefix)) + # # note(f"Checking {len(model_list)} keys under '{prefix}'") + # assert model_list == al_list + + # @precondition(lambda self: self.model.is_persistent_session()) + # @rule(data=st.data()) + # def check_group_path(self, data): + # t0 = time.time() + # group = data.draw(st.sampled_from(self.all_groups)) + # self.check_list_prefix_from_group(group) + # self.check_subgroups(group) + # self.check_group_arrays(group) + # t1 = time.time() + # note(f"Checks took {t1 - t0} sec.") + @invariant() + def check_list_prefix_from_root(self) -> None: + model_list = self._sync_iter(self.model.list_prefix("")) + store_list = self._sync_iter(self.store.list_prefix("")) + note(f"Checking {len(model_list)} expected keys vs {len(store_list)} actual keys") + assert sorted(model_list) == sorted(store_list), ( + sorted(model_list), + sorted(store_list), + ) + + # check that our internal state matches that of the store and model + assert all(f"{path}/zarr.json" in model_list for path in self.all_groups | self.all_arrays) + assert all(f"{path}/zarr.json" in store_list for path in self.all_groups | self.all_arrays) + + +class SyncStoreWrapper(zarr.core.sync.SyncMixin): + def __init__(self, store: Store) -> None: + """Synchronous Store wrapper + + This class holds synchronous methods that map to async methods of Store classes. + The synchronous wrapper is needed because hypothesis' stateful testing infra does + not support asyncio so we redefine sync versions of the Store API. + https://github.com/HypothesisWorks/hypothesis/issues/3712#issuecomment-1668999041 + """ + self.store = store + + @property + def read_only(self) -> bool: + return self.store.read_only + + def set(self, key: str, data_buffer: Buffer) -> None: + return self._sync(self.store.set(key, data_buffer)) + + def list(self) -> builtins.list[str]: + return self._sync_iter(self.store.list()) + + def get(self, key: str, prototype: BufferPrototype) -> Buffer | None: + return self._sync(self.store.get(key, prototype=prototype)) + + def get_partial_values( + self, key_ranges: Iterable[Any], prototype: BufferPrototype + ) -> builtins.list[Buffer | None]: + return self._sync(self.store.get_partial_values(prototype=prototype, key_ranges=key_ranges)) + + def delete(self, path: str) -> None: + return self._sync(self.store.delete(path)) + + def is_empty(self, prefix: str) -> bool: + return self._sync(self.store.is_empty(prefix=prefix)) + + def clear(self) -> None: + return self._sync(self.store.clear()) + + def exists(self, key: str) -> bool: + return self._sync(self.store.exists(key)) + + def getsize_prefix(self, prefix: str) -> int: + return self._sync(self.store.getsize_prefix(prefix)) + + def list_dir(self, prefix: str) -> None: + raise NotImplementedError + + def list_prefix(self, prefix: str) -> None: + raise NotImplementedError + + @property + def supports_listing(self) -> bool: + return self.store.supports_listing + + @property + def supports_writes(self) -> bool: + return self.store.supports_writes + + @property + def supports_deletes(self) -> bool: + return self.store.supports_deletes + + +class ZarrStoreStateMachine(RuleBasedStateMachine): + """ " + Zarr store state machine + + This is a subclass of a Hypothesis RuleBasedStateMachine. + It is testing a framework to ensure that the state of a Zarr store matches + an expected state after a set of random operations. It contains a store + (currently, a Zarr MemoryStore) and a model, a simplified version of a + zarr store (in this case, a dict). It also contains rules which represent + actions that can be applied to a zarr store. Rules apply an action to both + the store and the model, and invariants assert that the state of the model + is equal to the state of the store. Hypothesis then generates sequences of + rules, running invariants after each rule. It raises an error if a sequence + produces discontinuity between state of the model and state of the store + (ie. an invariant is violated). + https://hypothesis.readthedocs.io/en/latest/stateful.html + """ + + def __init__(self, store: Store) -> None: + super().__init__() + self.model: dict[str, Buffer] = {} + self.store = SyncStoreWrapper(store) + self.prototype = default_buffer_prototype() + + @initialize() + def init_store(self) -> None: + self.store.clear() + + @rule(key=zarr_keys(), data=st.binary(min_size=0, max_size=MAX_BINARY_SIZE)) + def set(self, key: str, data: bytes) -> None: + note(f"(set) Setting {key!r} with {data!r}") + assert not self.store.read_only + data_buf = cpu.Buffer.from_bytes(data) + self.store.set(key, data_buf) + self.model[key] = data_buf + + @precondition(lambda self: len(self.model.keys()) > 0) + @rule(key=zarr_keys(), data=st.data()) + def get(self, key: str, data: DataObject) -> None: + key = data.draw( + st.sampled_from(sorted(self.model.keys())) + ) # hypothesis wants to sample from sorted list + note("(get)") + store_value = self.store.get(key, self.prototype) + # to bytes here necessary because data_buf set to model in set() + assert self.model[key] == store_value + + @rule(key=zarr_keys(), data=st.data()) + def get_invalid_zarr_keys(self, key: str, data: DataObject) -> None: + note("(get_invalid)") + assume(key not in self.model) + assert self.store.get(key, self.prototype) is None + + @precondition(lambda self: len(self.model.keys()) > 0) + @rule(data=st.data()) + def get_partial_values(self, data: DataObject) -> None: + key_range = data.draw( + key_ranges(keys=st.sampled_from(sorted(self.model.keys())), max_size=MAX_BINARY_SIZE) + ) + note(f"(get partial) {key_range=}") + # Pass a one-shot generator rather than a list: stores (and wrappers such + # as LoggingStore) must not exhaust the iterable before using it. + obs_maybe = self.store.get_partial_values((kr for kr in key_range), self.prototype) + observed = [] + + for obs in obs_maybe: + assert obs is not None + observed.append(obs.to_bytes()) + + model_vals_ls = [] + + for key, byte_range in key_range: + # Independently model each ByteRequest variant (do NOT reuse the + # store's _normalize_byte_range_index helper, so this stays an + # independent oracle). Bounds may exceed the value length. + value = self.model[key] + n = len(value) + if byte_range is None: + expected = value[:] + elif isinstance(byte_range, RangeByteRequest): + expected = value[byte_range.start : byte_range.end] + elif isinstance(byte_range, OffsetByteRequest): + expected = value[byte_range.offset :] + elif isinstance(byte_range, SuffixByteRequest): + # "last suffix bytes"; suffix > n means the whole value. + expected = value[max(0, n - byte_range.suffix) :] + else: + raise AssertionError(f"unexpected byte_range {byte_range!r}") + model_vals_ls.append(expected) + + assert all( + obs == exp.to_bytes() for obs, exp in zip(observed, model_vals_ls, strict=True) + ), ( + observed, + model_vals_ls, + ) + + @precondition(lambda self: self.store.supports_deletes) + @precondition(lambda self: len(self.model.keys()) > 0) + @rule(data=st.data()) + def delete(self, data: DataObject) -> None: + key = data.draw(st.sampled_from(sorted(self.model.keys()))) + note(f"(delete) Deleting {key=}") + + self.store.delete(key) + del self.model[key] + + @rule() + def clear(self) -> None: + assert not self.store.read_only + note("(clear)") + self.store.clear() + self.model.clear() + + assert self.store.is_empty("") + + assert len(self.model.keys()) == len(list(self.store.list())) == 0 + + @rule() + # Local store can be non-empty when there are subdirectories but no files + @precondition(lambda self: not isinstance(self.store.store, LocalStore)) + def is_empty(self) -> None: + note("(is_empty)") + + # make sure they either both are or both aren't empty (same state) + assert self.store.is_empty("") == (not self.model) + + @rule(key=zarr_keys()) + def exists(self, key: str) -> None: + note("(exists)") + + assert self.store.exists(key) == (key in self.model) + + @precondition(lambda self: len(self.model.keys()) > 0) + @rule(data=st.data()) + def getsize_prefix(self, data: DataObject) -> None: + # Measure the size under the first path segment of some existing key. + # getsize_prefix(node) must count only keys under the directory "node/", + # not sibling keys that merely share the string prefix (e.g. measuring + # "a" must not include a sibling key "ab/..."). + key = data.draw(st.sampled_from(sorted(self.model.keys()))) + node = key.split("/")[0] + note(f"(getsize_prefix) {node=}") + + observed = self.store.getsize_prefix(node) + expected = sum(len(value) for k, value in self.model.items() if k.startswith(node + "/")) + assert observed == expected, (observed, expected, node) + + @invariant() + def check_paths_equal(self) -> None: + note("Checking that paths are equal") + paths = sorted(self.store.list()) + + assert sorted(self.model.keys()) == paths + + @invariant() + def check_vals_equal(self) -> None: + note("Checking values equal") + for key, val in self.model.items(): + store_item = self.store.get(key, self.prototype) + assert val == store_item + + @invariant() + def check_num_zarr_keys_equal(self) -> None: + note("check num zarr_keys equal") + + assert len(self.model) == len(list(self.store.list())) + + @invariant() + def check_zarr_keys(self) -> None: + keys = list(self.store.list()) + + if not keys: + assert self.store.is_empty("") is True + + else: + assert self.store.is_empty("") is False + + for key in keys: + assert self.store.exists(key) is True + note("checking keys / exists / empty") diff --git a/packages/zarr-storage/src/zarr_storage/testing/store.py b/packages/zarr-storage/src/zarr_storage/testing/store.py new file mode 100644 index 0000000000..168a83a13e --- /dev/null +++ b/packages/zarr-storage/src/zarr_storage/testing/store.py @@ -0,0 +1,680 @@ +from __future__ import annotations + +import pickle +from abc import abstractmethod +from typing import TYPE_CHECKING + +from zarr_storage.legacy._latency import LatencyStore + +if TYPE_CHECKING: + from typing import Any + + +import pytest +from zarr.core.buffer import Buffer, default_buffer_prototype +from zarr.core.sync import _collect_aiterator, sync + +from zarr_storage.legacy._abc import ( + ByteRequest, + OffsetByteRequest, + RangeByteRequest, + Store, + SuffixByteRequest, + SupportsDeleteSync, + SupportsGetSync, + SupportsSetSync, +) +from zarr_storage.legacy._utils import _normalize_byte_range_index +from zarr_storage.testing.utils import assert_bytes_equal + +__all__ = ["LatencyStore", "StoreTests"] + + +class StoreTests[S: Store, B: Buffer]: + store_cls: type[S] + buffer_cls: type[B] + + @staticmethod + def _require_get_sync(store: S) -> SupportsGetSync: + """Skip unless *store* implements [`SupportsGetSync`][zarr_storage.legacy._abc.SupportsGetSync].""" + if not isinstance(store, SupportsGetSync): + pytest.skip("store does not implement SupportsGetSync") + return store # type: ignore[unreachable] + + @staticmethod + def _require_set_sync(store: S) -> SupportsSetSync: + """Skip unless *store* implements [`SupportsSetSync`][zarr_storage.legacy._abc.SupportsSetSync].""" + if not isinstance(store, SupportsSetSync): + pytest.skip("store does not implement SupportsSetSync") + return store # type: ignore[unreachable] + + @staticmethod + def _require_delete_sync(store: S) -> SupportsDeleteSync: + """Skip unless *store* implements [`SupportsDeleteSync`][zarr_storage.legacy._abc.SupportsDeleteSync].""" + if not isinstance(store, SupportsDeleteSync): + pytest.skip("store does not implement SupportsDeleteSync") + return store # type: ignore[unreachable] + + @abstractmethod + async def set(self, store: S, key: str, value: Buffer) -> None: + """ + Insert a value into a storage backend, with a specific key. + This should not use any store methods. Bypassing the store methods allows them to be + tested. + """ + ... + + @abstractmethod + async def get(self, store: S, key: str) -> Buffer: + """ + Retrieve a value from a storage backend, by key. + This should not use any store methods. Bypassing the store methods allows them to be + tested. + """ + ... + + @abstractmethod + @pytest.fixture + def store_kwargs(self, *args: Any, **kwargs: Any) -> dict[str, Any]: + """Kwargs for instantiating a store""" + ... + + @abstractmethod + def test_store_repr(self, store: S) -> None: ... + + @abstractmethod + def test_store_supports_writes(self, store: S) -> None: ... + + def test_store_supports_partial_writes(self, store: S) -> None: + assert not store.supports_partial_writes + + @abstractmethod + def test_store_supports_listing(self, store: S) -> None: ... + + @pytest.fixture + def open_kwargs(self, store_kwargs: dict[str, Any]) -> dict[str, Any]: + return store_kwargs + + @pytest.fixture + async def store(self, open_kwargs: dict[str, Any]) -> Store: + return await self.store_cls.open(**open_kwargs) + + @pytest.fixture + async def store_not_open(self, store_kwargs: dict[str, Any]) -> Store: + return self.store_cls(**store_kwargs) + + def test_store_type(self, store: S) -> None: + assert isinstance(store, Store) + assert isinstance(store, self.store_cls) + + def test_store_eq(self, store: S, store_kwargs: dict[str, Any]) -> None: + # check self equality + assert store == store # noqa: PLR0124 + + # check store equality with same inputs + # asserting this is important for being able to compare (de)serialized stores + store2 = self.store_cls(**store_kwargs) + assert store == store2 + + async def test_serializable_store(self, store: S) -> None: + new_store: S = pickle.loads(pickle.dumps(store)) + assert new_store == store + assert new_store.read_only == store.read_only + # 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()) + assert_bytes_equal(observed, data_buf) + + def test_store_read_only(self, store: S) -> None: + assert not store.read_only + + with pytest.raises(AttributeError): + store.read_only = False # type: ignore[misc] + + @pytest.mark.parametrize("read_only", [True, False]) + async def test_store_open_read_only(self, open_kwargs: dict[str, Any], read_only: bool) -> None: + open_kwargs["read_only"] = read_only + store = await self.store_cls.open(**open_kwargs) + assert store._is_open + assert store.read_only == read_only + + async def test_store_context_manager(self, open_kwargs: dict[str, Any]) -> None: + # Test that the context manager closes the store + with await self.store_cls.open(**open_kwargs) as store: + assert store._is_open + # Test trying to open an already open store + with pytest.raises(ValueError, match="store is already open"): + await store._open() + assert not store._is_open + + async def test_read_only_store_raises(self, open_kwargs: dict[str, Any]) -> None: + kwargs = {**open_kwargs, "read_only": True} + store = await self.store_cls.open(**kwargs) + assert store.read_only + + # set + with pytest.raises( + ValueError, match="store was opened in read-only mode and does not support writing" + ): + await store.set("foo", self.buffer_cls.from_bytes(b"bar")) + + # delete + with pytest.raises( + ValueError, match="store was opened in read-only mode and does not support writing" + ): + await store.delete("foo") + + async def test_with_read_only_store(self, open_kwargs: dict[str, Any]) -> None: + kwargs = {**open_kwargs, "read_only": True} + store = await self.store_cls.open(**kwargs) + assert store.read_only + + # Test that you cannot write to a read-only store + with pytest.raises( + ValueError, match="store was opened in read-only mode and does not support writing" + ): + await store.set("foo", self.buffer_cls.from_bytes(b"bar")) + + # Check if the store implements with_read_only + try: + writer = store.with_read_only(read_only=False) + except NotImplementedError: + # Test that stores that do not implement with_read_only raise NotImplementedError with the correct message + with pytest.raises( + NotImplementedError, + match=f"with_read_only is not implemented for the {type(store)} store type.", + ): + store.with_read_only(read_only=False) + return + + # Test that you can write to a new store copy + assert not writer._is_open + assert not writer.read_only + await writer.set("foo", self.buffer_cls.from_bytes(b"bar")) + await writer.delete("foo") + + # Test that you cannot write to the original store + assert store.read_only + with pytest.raises( + ValueError, match="store was opened in read-only mode and does not support writing" + ): + await store.set("foo", self.buffer_cls.from_bytes(b"bar")) + with pytest.raises( + ValueError, match="store was opened in read-only mode and does not support writing" + ): + await store.delete("foo") + + # Test that you cannot write to a read-only store copy + reader = store.with_read_only(read_only=True) + assert reader.read_only + with pytest.raises( + ValueError, match="store was opened in read-only mode and does not support writing" + ): + await reader.set("foo", self.buffer_cls.from_bytes(b"bar")) + with pytest.raises( + ValueError, match="store was opened in read-only mode and does not support writing" + ): + await reader.delete("foo") + + @pytest.mark.parametrize("key", ["c/0", "foo/c/0.0", "foo/0/0"]) + @pytest.mark.parametrize( + ("data", "byte_range"), + [ + (b"\x01\x02\x03\x04", None), + (b"\x01\x02\x03\x04", RangeByteRequest(1, 4)), + (b"\x01\x02\x03\x04", OffsetByteRequest(1)), + (b"\x01\x02\x03\x04", SuffixByteRequest(1)), + (b"", None), + ], + ) + async def test_get(self, store: S, key: str, data: bytes, byte_range: ByteRequest) -> None: + """ + Ensure that data can be read from the store using the store.get method. + """ + data_buf = self.buffer_cls.from_bytes(data) + await self.set(store, key, data_buf) + observed = await store.get(key, prototype=default_buffer_prototype(), byte_range=byte_range) + start, stop = _normalize_byte_range_index(data_buf, byte_range=byte_range) + expected = data_buf[start:stop] + assert_bytes_equal(observed, expected) + + async def test_get_not_open(self, store_not_open: S) -> None: + """ + Ensure that data can be read from the store that isn't yet open using the store.get method. + """ + assert not store_not_open._is_open + data_buf = self.buffer_cls.from_bytes(b"\x01\x02\x03\x04") + key = "c/0" + await self.set(store_not_open, key, data_buf) + observed = await store_not_open.get(key, prototype=default_buffer_prototype()) + assert_bytes_equal(observed, data_buf) + + async def test_get_raises(self, store: S) -> None: + """ + Ensure that a ValueError is raise for invalid byte range syntax + """ + data_buf = self.buffer_cls.from_bytes(b"\x01\x02\x03\x04") + await self.set(store, "c/0", data_buf) + with pytest.raises((ValueError, TypeError), match=r"Unexpected byte_range, got.*"): + await store.get("c/0", prototype=default_buffer_prototype(), byte_range=(0, 2)) # type: ignore[arg-type] + + async def test_get_many(self, store: S) -> None: + """ + Ensure that multiple keys can be retrieved at once with the _get_many method. + """ + keys = tuple(map(str, range(10))) + values = tuple(f"{k}".encode() for k in keys) + for k, v in zip(keys, values, strict=False): + await self.set(store, k, self.buffer_cls.from_bytes(v)) + observed_buffers = await _collect_aiterator( + store._get_many( + zip( + keys, + (default_buffer_prototype(),) * len(keys), + (None,) * len(keys), + strict=False, + ) + ) + ) + observed_kvs = sorted(((k, b.to_bytes()) for k, b in observed_buffers)) # type: ignore[union-attr] + expected_kvs = sorted(((k, b) for k, b in zip(keys, values, strict=False))) + assert observed_kvs == expected_kvs + + @pytest.mark.parametrize("key", ["c/0", "foo/c/0.0", "foo/0/0"]) + @pytest.mark.parametrize("data", [b"\x01\x02\x03\x04", b""]) + async def test_getsize(self, store: S, key: str, data: bytes) -> None: + """ + Test the result of store.getsize(). + """ + data_buf = self.buffer_cls.from_bytes(data) + expected = len(data_buf) + await self.set(store, key, data_buf) + observed = await store.getsize(key) + assert observed == expected + + async def test_getsize_prefix(self, store: S) -> None: + """ + Test the result of store.getsize_prefix(). + + Includes a sibling key ("cc/0") that shares the string prefix "c" but + belongs to a different directory: getsize_prefix("c") must not count it, + i.e. the prefix is matched as a directory ("c/...") not a raw substring. + """ + data_buf = self.buffer_cls.from_bytes(b"\x01\x02\x03\x04") + keys = ["c/0/0", "c/0/1", "c/1/0", "c/1/1"] + # Sibling directory sharing the "c" string prefix; must be excluded. + sibling_keys = ["cc/0"] + keys_values = [(k, data_buf) for k in keys + sibling_keys] + await store._set_many(keys_values) + expected = len(data_buf) * len(keys) + observed = await store.getsize_prefix("c") + assert observed == expected + + async def test_getsize_raises(self, store: S) -> None: + """ + Test that getsize() raise a FileNotFoundError if the key doesn't exist. + """ + with pytest.raises(FileNotFoundError): + await store.getsize("c/1000") + + @pytest.mark.parametrize("key", ["zarr.json", "c/0", "foo/c/0.0", "foo/0/0"]) + @pytest.mark.parametrize("data", [b"\x01\x02\x03\x04", b""]) + async def test_set(self, store: S, key: str, data: bytes) -> None: + """ + Ensure that data can be written to the store using the store.set method. + """ + assert not store.read_only + data_buf = self.buffer_cls.from_bytes(data) + await store.set(key, data_buf) + observed = await self.get(store, key) + assert_bytes_equal(observed, data_buf) + + async def test_set_not_open(self, store_not_open: S) -> None: + """ + Ensure that data can be written to the store that's not yet open using the store.set method. + """ + assert not store_not_open._is_open + data_buf = self.buffer_cls.from_bytes(b"\x01\x02\x03\x04") + key = "c/0" + await store_not_open.set(key, data_buf) + observed = await self.get(store_not_open, key) + assert_bytes_equal(observed, data_buf) + + async def test_set_many(self, store: S) -> None: + """ + Test that a dict of key : value pairs can be inserted into the store via the + `_set_many` method. + """ + keys = ["zarr.json", "c/0", "foo/c/0.0", "foo/0/0"] + data_buf = [self.buffer_cls.from_bytes(k.encode()) for k in keys] + store_dict = dict(zip(keys, data_buf, strict=True)) + await store._set_many(store_dict.items()) + for k, v in store_dict.items(): + assert (await self.get(store, k)).to_bytes() == v.to_bytes() + + @pytest.mark.parametrize( + "key_ranges", + [ + [], + [("zarr.json", RangeByteRequest(0, 2))], + [("c/0", RangeByteRequest(0, 2)), ("zarr.json", None)], + [ + ("c/0/0", RangeByteRequest(0, 2)), + ("c/0/1", SuffixByteRequest(2)), + ("c/0/2", OffsetByteRequest(2)), + ], + ], + ) + async def test_get_partial_values( + self, store: S, key_ranges: list[tuple[str, ByteRequest]] + ) -> None: + # put all of the data + for key, _ in key_ranges: + await self.set(store, key, self.buffer_cls.from_bytes(bytes(key, encoding="utf-8"))) + + # read back just part of it. Pass key_ranges as a one-shot generator + # (a valid Iterable per the method signature) to ensure stores and + # wrappers do not exhaust the iterable before handing it to the backend. + observed_maybe = await store.get_partial_values( + prototype=default_buffer_prototype(), + key_ranges=(kr for kr in key_ranges), + ) + + # One result must be returned per requested key range. Checking this + # explicitly guards against a store/wrapper exhausting the key_ranges + # iterable early and silently returning fewer (or no) results. + assert len(observed_maybe) == len(key_ranges) + + observed: list[Buffer] = [] + expected: list[Buffer] = [] + + for obs in observed_maybe: + assert obs is not None + observed.append(obs) + + for key, byte_range in key_ranges: + result = await store.get( + key, prototype=default_buffer_prototype(), byte_range=byte_range + ) + assert result is not None + expected.append(result) + + assert all( + obs.to_bytes() == exp.to_bytes() for obs, exp in zip(observed, expected, strict=True) + ) + + async def test_exists(self, store: S) -> None: + assert not await store.exists("foo") + await store.set("foo/zarr.json", self.buffer_cls.from_bytes(b"bar")) + assert await store.exists("foo/zarr.json") + + async def test_delete(self, store: S) -> None: + if not store.supports_deletes: + pytest.skip("store does not support deletes") + await store.set("foo/zarr.json", self.buffer_cls.from_bytes(b"bar")) + assert await store.exists("foo/zarr.json") + await store.delete("foo/zarr.json") + assert not await store.exists("foo/zarr.json") + + async def test_delete_dir(self, store: S) -> None: + if not store.supports_deletes: + pytest.skip("store does not support deletes") + await store.set("zarr.json", self.buffer_cls.from_bytes(b"root")) + await store.set("foo-bar/zarr.json", self.buffer_cls.from_bytes(b"root")) + await store.set("foo/zarr.json", self.buffer_cls.from_bytes(b"bar")) + await store.set("foo/c/0", self.buffer_cls.from_bytes(b"chunk")) + await store.delete_dir("foo") + assert await store.exists("zarr.json") + assert await store.exists("foo-bar/zarr.json") + assert not await store.exists("foo/zarr.json") + assert not await store.exists("foo/c/0") + + async def test_delete_nonexistent_key_does_not_raise(self, store: S) -> None: + if not store.supports_deletes: + pytest.skip("store does not support deletes") + await store.delete("nonexistent_key") + + async def test_is_empty(self, store: S) -> None: + assert await store.is_empty("") + await self.set( + store, "foo/bar", self.buffer_cls.from_bytes(bytes("something", encoding="utf-8")) + ) + assert not await store.is_empty("") + assert await store.is_empty("fo") + assert not await store.is_empty("foo/") + assert not await store.is_empty("foo") + assert await store.is_empty("spam/") + + async def test_clear(self, store: S) -> None: + await self.set( + store, "key", self.buffer_cls.from_bytes(bytes("something", encoding="utf-8")) + ) + await store.clear() + assert await store.is_empty("") + + async def test_list(self, store: S) -> None: + assert await _collect_aiterator(store.list()) == () + prefix = "foo" + data = self.buffer_cls.from_bytes(b"") + store_dict = { + f"{prefix}/zarr.json": data, + **{f"{prefix}/c/{idx}": data for idx in range(10)}, + } + await store._set_many(store_dict.items()) + expected_sorted = sorted(store_dict.keys()) + observed = await _collect_aiterator(store.list()) + observed_sorted = sorted(observed) + assert observed_sorted == expected_sorted + + async def test_list_prefix(self, store: S) -> None: + """ + Test that the `list_prefix` method works as intended. Given a prefix, it should return + all the keys in storage that start with this prefix. + """ + prefixes = ("", "a/", "a/b/", "a/b/c/") + data = self.buffer_cls.from_bytes(b"") + fname = "zarr.json" + store_dict = {p + fname: data for p in prefixes} + + await store._set_many(store_dict.items()) + + for prefix in prefixes: + observed = tuple(sorted(await _collect_aiterator(store.list_prefix(prefix)))) + expected: tuple[str, ...] = () + for key in store_dict: + if key.startswith(prefix): + expected += (key,) + expected = tuple(sorted(expected)) + assert observed == expected + + async def test_list_empty_path(self, store: S) -> None: + """ + Verify that list and list_prefix work correctly when path is an empty string, + i.e. no unwanted replacement occurs. + """ + data = self.buffer_cls.from_bytes(b"") + store_dict = { + "foo/bar/zarr.json": data, + "foo/bar/c/1": data, + "foo/baz/c/0": data, + } + await store._set_many(store_dict.items()) + + # Test list() + observed_list = await _collect_aiterator(store.list()) + observed_list_sorted = sorted(observed_list) + expected_list_sorted = sorted(store_dict.keys()) + assert observed_list_sorted == expected_list_sorted + + # Test list_prefix() with an empty prefix + observed_prefix_empty = await _collect_aiterator(store.list_prefix("")) + observed_prefix_empty_sorted = sorted(observed_prefix_empty) + expected_prefix_empty_sorted = sorted(store_dict.keys()) + assert observed_prefix_empty_sorted == expected_prefix_empty_sorted + + # Test list_prefix() with a non-empty prefix + observed_prefix = await _collect_aiterator(store.list_prefix("foo/bar/")) + observed_prefix_sorted = sorted(observed_prefix) + expected_prefix_sorted = sorted(k for k in store_dict if k.startswith("foo/bar/")) + assert observed_prefix_sorted == expected_prefix_sorted + + async def test_list_dir(self, store: S) -> None: + roots_and_keys: list[tuple[str, dict[str, Buffer]]] = [ + ( + "foo", + { + "foo/zarr.json": self.buffer_cls.from_bytes(b"bar"), + "foo/c/1": self.buffer_cls.from_bytes(b"\x01"), + }, + ), + ( + "foo/bar", + { + "foo/bar/foobar_first_child": self.buffer_cls.from_bytes(b"1"), + "foo/bar/foobar_second_child/zarr.json": self.buffer_cls.from_bytes(b"2"), + }, + ), + ] + + assert await _collect_aiterator(store.list_dir("")) == () + + for root, store_dict in roots_and_keys: + assert await _collect_aiterator(store.list_dir(root)) == () + + await store._set_many(store_dict.items()) + + keys_observed = await _collect_aiterator(store.list_dir(root)) + keys_expected = {k.removeprefix(f"{root}/").split("/")[0] for k in store_dict} + assert sorted(keys_observed) == sorted(keys_expected) + + keys_observed = await _collect_aiterator(store.list_dir(f"{root}/")) + assert sorted(keys_expected) == sorted(keys_observed) + + async def test_set_if_not_exists(self, store: S) -> None: + key = "k" + data_buf = self.buffer_cls.from_bytes(b"0000") + await self.set(store, key, data_buf) + + new = self.buffer_cls.from_bytes(b"1111") + await store.set_if_not_exists("k", new) # no error + + result = await store.get(key, default_buffer_prototype()) + assert result == data_buf + + await store.set_if_not_exists("k2", new) # no error + + result = await store.get("k2", default_buffer_prototype()) + assert result == new + + # ------------------------------------------------------------------- + # Synchronous store methods (SupportsSyncStore protocol) + # ------------------------------------------------------------------- + + def test_get_sync(self, store: S) -> None: + getter = self._require_get_sync(store) + data_buf = self.buffer_cls.from_bytes(b"\x01\x02\x03\x04") + key = "sync_get" + sync(self.set(store, key, data_buf)) + result = getter.get_sync(key) + assert result is not None + assert_bytes_equal(result, data_buf) + + def test_get_sync_missing(self, store: S) -> None: + getter = self._require_get_sync(store) + result = getter.get_sync("nonexistent") + assert result is None + + def test_set_sync(self, store: S) -> None: + setter = self._require_set_sync(store) + data_buf = self.buffer_cls.from_bytes(b"\x01\x02\x03\x04") + key = "sync_set" + setter.set_sync(key, data_buf) + result = sync(self.get(store, key)) + assert_bytes_equal(result, data_buf) + + def test_delete_sync(self, store: S) -> None: + setter = self._require_set_sync(store) + deleter = self._require_delete_sync(store) + getter = self._require_get_sync(store) + if not store.supports_deletes: + pytest.skip("store does not support deletes") + data_buf = self.buffer_cls.from_bytes(b"\x01\x02\x03\x04") + key = "sync_delete" + setter.set_sync(key, data_buf) + deleter.delete_sync(key) + result = getter.get_sync(key) + assert result is None + + def test_delete_sync_missing(self, store: S) -> None: + deleter = self._require_delete_sync(store) + if not store.supports_deletes: + pytest.skip("store does not support deletes") + # should not raise + deleter.delete_sync("nonexistent_sync") + + # ------------------------------------------------------------------- + # Sync/async parity laws + # ------------------------------------------------------------------- + # A store's sync and async methods must observe the same key the same + # way. This is stronger than the individual test_get_sync/test_set_sync/ + # test_delete_sync tests above: those write and read back through the + # *same* API (sync-only or, via `self.set`/`self.get`, bypassing the + # store entirely), so a sync method that skips logic the async method + # applies (e.g. a path prefix) can still pass them. These laws write + # through one API and observe through the other. + + @pytest.mark.parametrize("direction", ["set_async_get_sync", "set_sync_get_async"]) + async def test_sync_async_set_get_parity(self, store: S, direction: str) -> None: + setter = self._require_set_sync(store) + getter = self._require_get_sync(store) + data_buf = self.buffer_cls.from_bytes(b"\x01\x02\x03\x04") + key = "parity_set_get" + if direction == "set_async_get_sync": + await store.set(key, data_buf) + result = getter.get_sync(key) + else: + setter.set_sync(key, data_buf) + result = await store.get(key, prototype=default_buffer_prototype()) + assert result is not None + assert_bytes_equal(result, data_buf) + + async def test_delete_sync_visible_to_async_get(self, store: S) -> None: + deleter = self._require_delete_sync(store) + if not store.supports_deletes: + pytest.skip("store does not support deletes") + data_buf = self.buffer_cls.from_bytes(b"\x01\x02\x03\x04") + key = "parity_delete" + await store.set(key, data_buf) + deleter.delete_sync(key) + result = await store.get(key, prototype=default_buffer_prototype()) + assert result is None + + @pytest.mark.parametrize( + "byte_range", + [ + None, + RangeByteRequest(1, 4), + OffsetByteRequest(1), + SuffixByteRequest(1), + RangeByteRequest(10, 20), + ], + ids=["none", "range", "offset", "suffix", "range-past-eof"], + ) + async def test_get_sync_byte_range_parity( + self, store: S, byte_range: ByteRequest | None + ) -> None: + getter = self._require_get_sync(store) + data_buf = self.buffer_cls.from_bytes(b"\x01\x02\x03\x04") + key = "parity_byte_range" + await store.set(key, data_buf) + sync_result = getter.get_sync(key, byte_range=byte_range) + async_result = await store.get( + key, prototype=default_buffer_prototype(), byte_range=byte_range + ) + if async_result is None: + assert sync_result is None + else: + assert sync_result is not None + assert_bytes_equal(sync_result, async_result) diff --git a/packages/zarr-storage/src/zarr_storage/testing/strategies.py b/packages/zarr-storage/src/zarr_storage/testing/strategies.py new file mode 100644 index 0000000000..ec2c708f0e --- /dev/null +++ b/packages/zarr-storage/src/zarr_storage/testing/strategies.py @@ -0,0 +1,826 @@ +import itertools +import math +import sys +from collections.abc import Callable, Mapping +from typing import Any, Literal + +import hypothesis.extra.numpy as npst +import hypothesis.strategies as st +import numpy as np +import numpy.typing as npt +import zarr +from hypothesis import event +from hypothesis.strategies import SearchStrategy +from zarr.codecs.bytes import BytesCodec +from zarr.codecs.crc32c_ import Crc32cCodec +from zarr.codecs.sharding import SUBCHUNK_WRITE_ORDER, ShardingCodec, SubchunkWriteOrder +from zarr.codecs.zstd import ZstdCodec +from zarr.core.array import Array, CompressorsLike, SerializerLike +from zarr.core.chunk_key_encodings import DefaultChunkKeyEncoding +from zarr.core.common import JSON, AccessModeLiteral, ZarrFormat +from zarr.core.dtype import get_data_type_from_native_dtype +from zarr.core.metadata import ArrayV2Metadata, ArrayV3Metadata +from zarr.core.metadata.v3 import RectilinearChunkGridMetadata, RegularChunkGridMetadata +from zarr.core.sync import sync +from zarr.types import AnyArray + +from zarr_storage.legacy import MemoryStore, StoreLike +from zarr_storage.legacy._abc import ( + ByteRequest, + OffsetByteRequest, + RangeByteRequest, + Store, + SuffixByteRequest, +) +from zarr_storage.legacy._utils import _join_paths, normalize_path + +TrueOrFalse = Literal[True, False] + +# Copied from Xarray +_attr_keys = st.text(st.characters(), min_size=1) +_attr_values = st.recursive( + st.none() | st.booleans() | st.text(st.characters(), max_size=5), + lambda children: st.lists(children) | st.dictionaries(_attr_keys, children), + max_leaves=3, +) + + +@st.composite +def keys(draw: st.DrawFn, *, max_num_nodes: int | None = None) -> str: + return draw(st.lists(node_names, min_size=1, max_size=max_num_nodes).map("/".join)) + + +@st.composite +def paths(draw: st.DrawFn, *, max_num_nodes: int | None = None) -> str: + return draw(st.just("/") | keys(max_num_nodes=max_num_nodes)) + + +def dtypes() -> st.SearchStrategy[np.dtype[Any]]: + return ( + npst.boolean_dtypes() + | npst.integer_dtypes(endianness="=") + | npst.unsigned_integer_dtypes(endianness="=") + | npst.floating_dtypes(endianness="=") + | npst.complex_number_dtypes(endianness="=") + | npst.byte_string_dtypes(endianness="=") + | npst.unicode_string_dtypes(endianness="=") + | npst.datetime64_dtypes(endianness="=") + | npst.timedelta64_dtypes(endianness="=") + ) + + +def v3_dtypes() -> st.SearchStrategy[np.dtype[Any]]: + return dtypes() + + +def v2_dtypes() -> st.SearchStrategy[np.dtype[Any]]: + return dtypes() + + +def safe_unicode_for_dtype(dtype: np.dtype[np.str_]) -> st.SearchStrategy[str]: + """Generate UTF-8-safe text constrained to max_len of dtype.""" + # account for utf-32 encoding (i.e. 4 bytes/character) + max_len = max(1, dtype.itemsize // 4) + + return st.text( + alphabet=st.characters( + exclude_categories=["Cs"], # Avoid *technically allowed* surrogates + min_codepoint=32, + ), + min_size=1, + max_size=max_len, + ) + + +def clear_store(x: Store) -> Store: + sync(x.clear()) + return x + + +# From https://zarr-specs.readthedocs.io/en/latest/v3/core/v3.0.html#node-names +# 1. must not be the empty string ("") +# 2. must not include the character "/" +# 3. must not be a string composed only of period characters, e.g. "." or ".." +# 4. must not start with the reserved prefix "__" +zarr_key_chars = st.sampled_from( + ".-0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ_abcdefghijklmnopqrstuvwxyz" +) +node_names = ( + st.text(zarr_key_chars, min_size=1) + .filter(lambda t: t not in (".", "..") and not t.startswith("__")) + .filter(lambda name: name.lower() != "zarr.json") +) +short_node_names = ( + st.text(zarr_key_chars, max_size=3, min_size=1) + .filter(lambda t: t not in (".", "..") and not t.startswith("__")) + .filter(lambda name: name.lower() != "zarr.json") +) +array_names = node_names +attrs: st.SearchStrategy[Mapping[str, JSON] | None] = st.none() | st.dictionaries( + _attr_keys, _attr_values +) +# st.builds will only call a new store constructor for different keyword arguments +# i.e. stores.examples() will always return the same object per Store class. +# So we map a clear to reset the store. +stores = st.builds(MemoryStore, st.just({})).map(clear_store) +compressors = st.sampled_from([None, "default"]) +zarr_formats: st.SearchStrategy[ZarrFormat] = st.sampled_from([3, 2]) +# We de-prioritize arrays having dim sizes 0, 1, 2 +array_shapes = npst.array_shapes(max_dims=4, min_side=3, max_side=5) | npst.array_shapes( + max_dims=4, min_side=0 +) + + +@st.composite +def dimension_names(draw: st.DrawFn, *, ndim: int | None = None) -> list[str | None] | None: + simple_text = st.text(zarr_key_chars, min_size=0) + return draw(st.none() | st.lists(st.none() | simple_text, min_size=ndim, max_size=ndim)) # type: ignore[arg-type] + + +subchunk_write_orders: st.SearchStrategy[SubchunkWriteOrder] = st.sampled_from(SUBCHUNK_WRITE_ORDER) + +# Inner codec chains for a ShardingCodec. We MUST sample the uncompressed, +# single-BytesCodec configuration (no Zstd) — that is the only configuration in +# which the FusedCodecPipeline's vectorized whole-shard "bulk decode" fast path +# engages, so it is the only one that can exercise (and regress-guard) that path +# against arbitrary indexing. Freezing the inner codecs to [BytesCodec, ZstdCodec] +# silently disables the fast path under every property test. +sharding_inner_codecs: st.SearchStrategy[list[BytesCodec | ZstdCodec]] = st.sampled_from( + [ + [BytesCodec()], + [BytesCodec(), ZstdCodec()], + ] +) + + +@st.composite +def array_metadata( + draw: st.DrawFn, + *, + array_shapes: Callable[..., st.SearchStrategy[tuple[int, ...]]] = npst.array_shapes, + zarr_formats: st.SearchStrategy[ZarrFormat] = zarr_formats, + attributes: SearchStrategy[Mapping[str, JSON] | None] = attrs, +) -> ArrayV2Metadata | ArrayV3Metadata: + zarr_format = draw(zarr_formats) + # separator = draw(st.sampled_from(['/', '\\'])) + shape = draw(array_shapes()) + ndim = len(shape) + np_dtype = draw(dtypes()) + dtype = get_data_type_from_native_dtype(np_dtype) + fill_value = draw(npst.from_dtype(np_dtype)) + if zarr_format == 2: + chunk_shape = draw(array_shapes(min_dims=ndim, max_dims=ndim, min_side=1)) + return ArrayV2Metadata( + shape=shape, + chunks=chunk_shape, + dtype=dtype, + fill_value=fill_value, + order=draw(st.sampled_from(["C", "F"])), + attributes=draw(attributes), # type: ignore[arg-type] + dimension_separator=draw(st.sampled_from([".", "/"])), + filters=None, + compressor=None, + ) + else: + chunk_grid = draw(chunk_grids(shape=shape)) + return ArrayV3Metadata( + shape=shape, + data_type=dtype, + chunk_grid=chunk_grid, + fill_value=fill_value, + attributes=draw(attributes), # type: ignore[arg-type] + dimension_names=draw(dimension_names(ndim=ndim)), + chunk_key_encoding=DefaultChunkKeyEncoding(separator="/"), # FIXME + codecs=[BytesCodec()], + storage_transformers=(), + ) + + +@st.composite +def numpy_arrays( + draw: st.DrawFn, + *, + shapes: st.SearchStrategy[tuple[int, ...]] = array_shapes, + dtype: np.dtype[Any] | None = None, +) -> npt.NDArray[Any]: + """ + Generate numpy arrays that can be saved in the provided Zarr format. + """ + if dtype is None: + dtype = draw(dtypes()) + if np.issubdtype(dtype, np.str_): + safe_unicode_strings = safe_unicode_for_dtype(dtype) + return draw(npst.arrays(dtype=dtype, shape=shapes, elements=safe_unicode_strings)) + + return draw(npst.arrays(dtype=dtype, shape=shapes)) + + +@st.composite +def chunk_shapes(draw: st.DrawFn, *, shape: tuple[int, ...]) -> tuple[int, ...]: + # We want this strategy to shrink towards arrays with smaller number of chunks + # 1. st.integers() shrinks towards smaller values. So we use that to generate number of chunks + numchunks = draw( + st.tuples( + *[ + st.integers(min_value=0 if size == 0 else 1, max_value=max(size, 1)) + for size in shape + ] + ) + ) + # 2. and now generate the chunks tuple + # Chunk sizes must be >= 1 per spec; for zero-extent dimensions use 1. + chunks = tuple( + max(1, size // nchunks) if nchunks > 0 else 1 + for size, nchunks in zip(shape, numchunks, strict=True) + ) + + for c in chunks: + event("chunk size", c) + + if any((c != 0 and s % c != 0) for s, c in zip(shape, chunks, strict=True)): + event("smaller last chunk") + + return chunks + + +@st.composite +def shard_shapes( + draw: st.DrawFn, *, shape: tuple[int, ...], chunk_shape: tuple[int, ...] +) -> tuple[int, ...]: + # We want this strategy to shrink towards arrays with smaller number of shards + # shards must be an integral number of chunks + assert all(c != 0 for c in chunk_shape) + numchunks = tuple(s // c for s, c in zip(shape, chunk_shape, strict=True)) + multiples = tuple(draw(st.integers(min_value=1, max_value=nc)) for nc in numchunks) + return tuple(m * c for m, c in zip(multiples, chunk_shape, strict=True)) + + +@st.composite +def np_array_and_chunks( + draw: st.DrawFn, + *, + arrays: st.SearchStrategy[npt.NDArray[Any]] = numpy_arrays(), # noqa: B008 +) -> tuple[np.ndarray[Any, Any], tuple[int, ...]]: + """A hypothesis strategy to generate small sized random arrays. + + Returns: a tuple of the array and a suitable random chunking for it. + """ + array = draw(arrays) + return (array, draw(chunk_shapes(shape=array.shape))) + + +@st.composite +def arrays( + draw: st.DrawFn, + *, + shapes: st.SearchStrategy[tuple[int, ...]] = array_shapes, + compressors: st.SearchStrategy = compressors, + stores: st.SearchStrategy[StoreLike] = stores, + paths: st.SearchStrategy[str] = paths(), # noqa: B008 + array_names: st.SearchStrategy = array_names, + arrays: st.SearchStrategy | None = None, + attrs: st.SearchStrategy = attrs, + zarr_formats: st.SearchStrategy = zarr_formats, + subchunk_write_orders: SearchStrategy[SubchunkWriteOrder] = subchunk_write_orders, + open_mode: AccessModeLiteral = "w", +) -> AnyArray: + store = draw(stores, label="store") + path = draw(paths, label="array parent") + name = draw(array_names, label="array name") + attributes = draw(attrs, label="attributes") + zarr_format = draw(zarr_formats, label="zarr format") + if arrays is None: + arrays = numpy_arrays(shapes=shapes) + nparray = draw(arrays, label="array data") + dim_names: list[str | None] | None = None + serializer: SerializerLike = "auto" + compressors_unsearched: CompressorsLike = "auto" + + # For v3 arrays, optionally use RectilinearChunkGridMetadata + chunk_grid_meta: RegularChunkGridMetadata | RectilinearChunkGridMetadata | None = None + + # test that None works too. + fill_value = draw(st.one_of([st.none(), npst.from_dtype(nparray.dtype)])) + # compressor = draw(compressors) + + expected_attrs = {} if attributes is None else attributes + + array_path = _join_paths([path, name]) + root = zarr.open_group(store, mode=open_mode, zarr_format=zarr_format) + + # Convert chunk grid metadata to a form create_array accepts: + # - RegularChunkGridMetadata -> flat tuple of ints + # - RectilinearChunkGridMetadata -> nested list of ints (triggers rectilinear path) + # - v2 -> flat tuple of ints + chunks_param: tuple[int, ...] | list[list[int]] + shard_shape = None + dim_names = None + if zarr_format == 3: + chunk_grid_meta = draw(st.none() | chunk_grids(shape=nparray.shape), label="chunk grid") + dim_names = draw(dimension_names(ndim=nparray.ndim), label="dimension names") + if isinstance(chunk_grid_meta, RectilinearChunkGridMetadata): + chunks_param = [ + list(dim) if isinstance(dim, tuple) else [dim] + for dim in chunk_grid_meta.chunk_shapes + ] + elif isinstance(chunk_grid_meta, RegularChunkGridMetadata): + chunks_param = chunk_grid_meta.chunk_shape + else: + chunks_param = draw(chunk_shapes(shape=nparray.shape), label="chunk shape") + + if all(s > c > 1 for s, c in zip(nparray.shape, chunks_param, strict=True)): + shard_shape = draw( + st.none() | shard_shapes(shape=nparray.shape, chunk_shape=chunks_param), + label="shard shape", + ) + if shard_shape is not None: + subchunk_write_order = draw(subchunk_write_orders) + inner_codecs = draw(sharding_inner_codecs, label="sharding inner codecs") + serializer = ShardingCodec( + subchunk_write_order=subchunk_write_order, + codecs=inner_codecs, + index_codecs=[BytesCodec(), Crc32cCodec()], + chunk_shape=chunks_param, + ) + compressors_unsearched = None + else: + chunks_param = draw(chunk_shapes(shape=nparray.shape), label="chunk shape") + a = root.create_array( + array_path, + shape=nparray.shape, + chunks=chunks_param, + shards=shard_shape, + dtype=nparray.dtype, + attributes=attributes, + compressors=compressors_unsearched, # FIXME + fill_value=fill_value, + dimension_names=dim_names, + serializer=serializer, + ) + + assert isinstance(a, Array) + if a.metadata.zarr_format == 3: + assert a.fill_value is not None + assert a.name is not None + assert a.path == normalize_path(array_path) + assert a.name == f"/{a.path}" + assert isinstance(root[array_path], Array) + assert nparray.shape == a.shape + + # Verify chunks — for rectilinear grids, .chunks raises + if zarr_format == 3: + assert shard_shape == a.shards + if isinstance(a.metadata.chunk_grid, RegularChunkGridMetadata): + assert a.metadata.chunk_grid.chunk_shape == ( + a.shards if shard_shape is not None else a.chunks + ) + assert shard_shape == a.shards + else: + assert isinstance(a.metadata.chunk_grid, RectilinearChunkGridMetadata) + assert shard_shape is None + + assert a.basename == name, (a.basename, name) + assert dict(a.attrs) == expected_attrs + + a[:] = nparray + + return a + + +@st.composite +def simple_arrays( + draw: st.DrawFn, + *, + shapes: st.SearchStrategy[tuple[int, ...]] = array_shapes, +) -> Any: + return draw( + arrays( + shapes=shapes, + paths=paths(max_num_nodes=2), + array_names=short_node_names, + attrs=st.none(), + compressors=st.sampled_from([None, "default"]), + ) + ) + + +@st.composite +def rectilinear_chunks(draw: st.DrawFn, *, shape: tuple[int, ...]) -> list[list[int]]: + """Generate valid rectilinear chunk shapes for a given array shape. + + Uses two modes per dimension: + - "expanded": random divider points create arbitrary chunk sizes + - "rle": uniform chunks with optional remainder, optionally shuffled + + Keeps max chunks per dimension <= 20 to avoid performance issues + in property tests. With higher dimensions, the total chunk count + grows multiplicatively. + """ + chunk_shapes: list[list[int]] = [] + for size in shape: + assert size > 0 + if size > 1: + mode = draw(st.sampled_from(["expanded", "rle"])) + if mode == "expanded": + event("rectilinear expanded") + max_chunks = min(size - 1, 20) + nchunks = draw(st.integers(min_value=1, max_value=max_chunks)) + dividers = sorted( + draw( + st.lists( + st.integers(min_value=1, max_value=size - 1), + min_size=nchunks - 1, + max_size=nchunks - 1, + unique=True, + ) + ) + ) + chunk_shapes.append( + [a - b for a, b in zip(dividers + [size], [0] + dividers, strict=False)] + ) + else: + # RLE mode: uniform chunks with optional remainder + max_chunk_size = min(size, 20) + chunk_size = draw(st.integers(min_value=1, max_value=max_chunk_size)) + n_full = size // chunk_size + remainder = size % chunk_size + chunks_list = [chunk_size] * n_full + if remainder > 0: + chunks_list.append(remainder) + # Optionally shuffle to create non-contiguous duplicate patterns + if draw(st.booleans()): + event("rectilinear rle shuffled") + chunks_list = draw(st.permutations(chunks_list)) + else: + event("rectilinear rle") + chunk_shapes.append(list(chunks_list)) + else: + chunk_shapes.append([1]) + return chunk_shapes + + +@st.composite +def chunk_grids( + draw: st.DrawFn, *, shape: tuple[int, ...] +) -> RegularChunkGridMetadata | RectilinearChunkGridMetadata: + """Generate either a RegularChunkGridMetadata or RectilinearChunkGridMetadata. + + This strategy depends on the global state of the config having rectilinear chunk grids enabled or not. + This means that it may be a possible source of a hypothesis FlakyStrategy error due dependence + on global state. However, in practice this seems unlikely to happen. + + This allows property tests to exercise both chunk grid types. + """ + # RectilinearChunkGridMetadata doesn't support zero-sized dimensions, + # so use RegularChunkGridMetadata if any dimension is 0 + if any(s == 0 for s in shape): + event("using RegularChunkGridMetadata (zero-sized dimensions)") + return RegularChunkGridMetadata(chunk_shape=draw(chunk_shapes(shape=shape))) + + if zarr.config.get("array.rectilinear_chunks") and draw(st.booleans()): + chunks = draw(rectilinear_chunks(shape=shape)) + event("using RectilinearChunkGridMetadata") + return RectilinearChunkGridMetadata(chunk_shapes=tuple(tuple(dim) for dim in chunks)) + else: + event("using RegularChunkGridMetadata") + return RegularChunkGridMetadata(chunk_shape=draw(chunk_shapes(shape=shape))) + + +# Rectilinear arrays need min_side >= 1 so every dimension has at least one element +_rectilinear_shapes = npst.array_shapes(max_dims=3, min_side=1, max_side=20) + + +@st.composite +def rectilinear_arrays( + draw: st.DrawFn, + *, + shapes: st.SearchStrategy[tuple[int, ...]] = _rectilinear_shapes, +) -> Any: + """Generate a zarr v3 array with rectilinear (variable) chunk grid.""" + shape = draw(shapes) + chunk_shapes = draw(rectilinear_chunks(shape=shape)) + + np_dtype = draw(dtypes()) + nparray = draw(numpy_arrays(shapes=st.just(shape), dtype=np_dtype)) + fill_value = draw(st.one_of([st.none(), npst.from_dtype(np_dtype)])) + dim_names = draw(dimension_names(ndim=len(shape))) + + store = MemoryStore() + with zarr.config.set({"array.rectilinear_chunks": True}): + a = zarr.create_array( + store=store, + shape=shape, + chunks=chunk_shapes, + dtype=np_dtype, + fill_value=fill_value, + dimension_names=dim_names, + ) + a[:] = nparray + + return a + + +def is_negative_slice(idx: Any) -> bool: + return isinstance(idx, slice) and idx.step is not None and idx.step < 0 + + +@st.composite +def end_slices(draw: st.DrawFn, *, shape: tuple[int, ...]) -> Any: + """ + A strategy that slices ranges that include the last chunk. + This is intended to stress-test handling of a possibly smaller last chunk. + """ + slicers = [] + for size in shape: + start = draw(st.integers(min_value=size // 2, max_value=size - 1)) + length = draw(st.integers(min_value=0, max_value=size - start)) + slicers.append(slice(start, start + length)) + event("drawing end slice") + return tuple(slicers) + + +@st.composite +def basic_indices( + draw: st.DrawFn, + *, + shape: tuple[int, ...], + min_dims: int = 0, + max_dims: int | None = None, + allow_newaxis: TrueOrFalse = False, + allow_ellipsis: TrueOrFalse = True, +) -> Any: + """Basic indices without unsupported negative slices.""" + strategy = npst.basic_indices( + shape=shape, + min_dims=min_dims, + max_dims=max_dims, + allow_newaxis=allow_newaxis, + allow_ellipsis=allow_ellipsis, + ).filter( + lambda idxr: ( + not ( + is_negative_slice(idxr) + or (isinstance(idxr, tuple) and any(is_negative_slice(idx) for idx in idxr)) + ) + ) + ) + if math.prod(shape) >= 3: + strategy = end_slices(shape=shape) | strategy + return draw(strategy) + + +@st.composite +def orthogonal_indices( + draw: st.DrawFn, *, shape: tuple[int, ...] +) -> tuple[tuple[np.ndarray[Any, Any], ...], tuple[np.ndarray[Any, Any], ...]]: + """ + Strategy that returns + (1) a tuple of integer arrays used for orthogonal indexing of Zarr arrays. + (2) a tuple of integer arrays that can be used for equivalent indexing of numpy arrays + """ + zindexer = [] + npindexer = [] + ndim = len(shape) + for axis, size in enumerate(shape): + if size != 0: + strategy = npst.integer_array_indices( + shape=(size,), result_shape=npst.array_shapes(min_side=1, max_side=size, max_dims=1) + ) | basic_indices(min_dims=1, shape=(size,), allow_ellipsis=False) + else: + strategy = basic_indices(min_dims=1, shape=(size,), allow_ellipsis=False) + + val = draw( + strategy + # bare ints, slices + .map(lambda x: (x,) if not isinstance(x, tuple) else x) + # skip empty tuple + .filter(bool) + ) + (idxr,) = val + if isinstance(idxr, int): + idxr = np.array([idxr]) + zindexer.append(idxr) + if isinstance(idxr, slice): + idxr = np.arange(*idxr.indices(size)) + elif isinstance(idxr, (tuple, int)): + idxr = np.array(idxr) + newshape = [1] * ndim + newshape[axis] = idxr.size + npindexer.append(idxr.reshape(newshape)) + + # casting the output of broadcast_arrays is needed for numpy < 2 + return tuple(zindexer), tuple(np.broadcast_arrays(*npindexer)) + + +@st.composite +def block_indices( + draw: st.DrawFn, *, chunk_sizes: tuple[tuple[int, ...], ...] +) -> tuple[tuple[int | slice, ...], tuple[slice, ...]]: + """ + Strategy for block-selection indexers over a chunk grid. + + Block indexing is basic indexing applied to the block grid (the grid of + chunks), so each axis is drawn with ``basic_indices`` over that axis's chunk + count, mirroring how ``orthogonal_indices`` reuses ``basic_indices`` per + axis. ``chunk_sizes`` gives the per-chunk data sizes of the array's *outer* + (block) grid for every axis — i.e. ``Array.write_chunk_sizes``, the grid that + ``Array.blocks`` addresses (the shard grid when sharding is used). For + example ``(3, 3, 3, 1)`` for a length-10 axis with a regular chunk size of 3, + or the explicit edges of a rectilinear axis; ``nchunks`` for an axis is + ``len(chunk_sizes[axis])``. + + The array-space translation uses the cumulative sum of those sizes, matching + ``BlockIndexer``'s use of ``dim_grid.chunk_offset``. Because the sizes are + clipped to the array extent, the final offset equals the extent and the + translation is exact for regular (uniform), rectilinear, and sharded grids + alike. + + Block indexing only supports integers and step-1 slices whose start + references an existing chunk, so strided slices and slices starting at the + grid edge are filtered out. + + Returns + ------- + block_indexer + A per-axis tuple of ints / step-1 slices addressing whole chunks, + suitable for ``Array.blocks`` / ``get_block_selection`` / ``set_block_selection``. + array_indexer + The equivalent array-space selection (a tuple of slices) for indexing + the corresponding numpy array, used as the comparison oracle. + """ + + def supported(nchunks: int) -> Callable[[tuple[Any, ...]], bool]: + # Block indexing only accepts step-1 slices whose start references an + # existing chunk (a slice starting at nchunks raises, unlike numpy). + def predicate(value: tuple[Any, ...]) -> bool: + dim_sel = value[0] + if isinstance(dim_sel, slice): + if dim_sel.step not in (None, 1): + return False + start = dim_sel.start or 0 + return 0 <= (start + nchunks if start < 0 else start) < nchunks + return True + + return predicate + + block_indexer: list[int | slice] = [] + array_indexer: list[slice] = [] + for sizes in chunk_sizes: + nchunks = len(sizes) + # offsets[i] is the array-space start of chunk i; length nchunks + 1. + offsets = list(itertools.accumulate(sizes, initial=0)) + dim_strategy = ( + basic_indices(min_dims=1, shape=(nchunks,), allow_ellipsis=False) + # normalize bare ints / slices to a 1-tuple, skip the empty tuple + .map(lambda x: (x,) if not isinstance(x, tuple) else x) + .filter(bool) + .filter(supported(nchunks)) + ) + # basic_indices draws slices far more often than bare integers, so the + # integer (single-block) branch below would only be hit on rare draws. + # Union in an explicit integer so it is reliably exercised — keeping + # coverage deterministic under the derandomized ``ci`` Hypothesis profile. + (dim_sel,) = draw( + dim_strategy | st.integers(min_value=0, max_value=nchunks - 1).map(lambda i: (i,)) + ) + block_indexer.append(dim_sel) + if isinstance(dim_sel, slice): + start, stop, _ = dim_sel.indices(nchunks) + array_indexer.append(slice(offsets[start], offsets[stop])) + else: + block = dim_sel % nchunks + array_indexer.append(slice(offsets[block], offsets[block + 1])) + return tuple(block_indexer), tuple(array_indexer) + + +@st.composite +def block_test_arrays( + draw: st.DrawFn, +) -> tuple[Array[Any], np.ndarray[Any, Any]]: + """Draw an array for block-indexing property tests, with its source contents. + + Two arms, selected with equal probability: + + - **regular**: a regular chunk grid, optionally wrapped in sharding. + - **rectilinear**: a variable (rectilinear) chunk grid, always unsharded. + + Returns ``(zarray, nparray)``. The per-axis block sizes the oracle needs are + ``zarray.write_chunk_sizes`` — the array's *outer* (block / shard) grid, which + is exactly the grid ``Array.blocks`` addresses; the caller reads it directly. + """ + chunks: tuple[int, ...] | list[list[int]] + if draw(st.booleans()): + # regular arm, optionally sharded + nparray, chunks = draw( + np_array_and_chunks( + arrays=numpy_arrays(shapes=npst.array_shapes(max_dims=4, min_side=1)) + ) + ) + # min_side=1 chunking guarantees shape // chunk >= 1 on every axis, which + # shard_shapes requires. + shards = draw(st.none() | shard_shapes(shape=nparray.shape, chunk_shape=chunks)) + event("block regular sharded" if shards is not None else "block regular unsharded") + rectilinear = False + else: + # rectilinear arm, always unsharded + event("block rectilinear") + shape = draw(_rectilinear_shapes) + chunks = draw(rectilinear_chunks(shape=shape)) + nparray = draw(numpy_arrays(shapes=st.just(shape), dtype=draw(dtypes()))) + shards, rectilinear = None, True + + store = draw(stores) + with zarr.config.set({"array.rectilinear_chunks": rectilinear}): + zarray = zarr.create_array( + store=store, + shape=nparray.shape, + chunks=chunks, + shards=shards, + dtype=nparray.dtype, + ) + zarray[...] = nparray + return zarray, nparray + + +def key_ranges( + keys: SearchStrategy[str] = node_names, max_size: int = sys.maxsize +) -> SearchStrategy[list[tuple[str, ByteRequest | None]]]: + """ + Function to generate key_ranges strategy for get_partial_values() + returns list strategy w/ form:: + + [(key, byte_request), + (key, byte_request),...] + + where ``byte_request`` is ``None`` or any of the concrete ``ByteRequest`` + subtypes. The bounds are drawn independently of each value's length, so the + offsets/suffixes routinely exceed the data and exercise the clamping logic + in ``_normalize_byte_range_index``. + """ + + def make_range(start: int, length: int) -> RangeByteRequest: + return RangeByteRequest(start, end=min(start + length, max_size)) + + bound = st.integers(min_value=0, max_value=max_size) + byte_ranges: SearchStrategy[ByteRequest | None] = st.one_of( + st.none(), + st.builds(make_range, start=bound, length=bound), + st.builds(OffsetByteRequest, offset=bound), + st.builds(SuffixByteRequest, suffix=bound), + ) + key_tuple = st.tuples(keys, byte_ranges) + return st.lists(key_tuple, min_size=1, max_size=10) + + +@st.composite +def complex_rectilinear_arrays( + draw: st.DrawFn, + *, + stores: st.SearchStrategy[StoreLike] = stores, + paths: st.SearchStrategy[str] = paths(), # noqa: B008 + array_names: st.SearchStrategy = array_names, + attrs: st.SearchStrategy = attrs, +) -> tuple[npt.NDArray[Any], AnyArray]: + """Generate a rectilinear array with many small chunks. + + The shape is derived from the chunk edges (5-10 chunks per dim, + sizes 1-5), exercising higher chunk counts than ``rectilinear_arrays``. + """ + ndim = draw(st.integers(min_value=1, max_value=3)) + nchunks = draw(st.integers(min_value=5, max_value=10)) + dim_chunks = st.lists(st.integers(min_value=1, max_value=5), min_size=nchunks, max_size=nchunks) + chunk_shapes = draw(st.lists(dim_chunks, min_size=ndim, max_size=ndim)) + + shape = tuple(sum(dim) for dim in chunk_shapes) + nparray = draw(numpy_arrays(shapes=st.just(shape))) + dim_names = draw(dimension_names(ndim=ndim)) + fill_value = draw(st.one_of([st.none(), npst.from_dtype(nparray.dtype)])) + attributes = draw(attrs) + + store = draw(stores, label="store") + path = draw(paths, label="array parent") + name = draw(array_names, label="array name") + array_path = _join_paths([path, name]) + + root = zarr.open_group(store, mode="w", zarr_format=3) + with zarr.config.set({"array.rectilinear_chunks": True}): + a = root.create_array( + array_path, + shape=shape, + chunks=chunk_shapes, + dtype=nparray.dtype, + fill_value=fill_value, + dimension_names=dim_names, + attributes=attributes, + ) + a[:] = nparray + return nparray, a + + +@st.composite +def chunk_paths(draw: st.DrawFn, ndim: int, numblocks: tuple[int, ...], subset: bool = True) -> str: + blockidx = draw( + st.tuples(*tuple(st.integers(min_value=0, max_value=max(0, b - 1)) for b in numblocks)) + ) + subset_slicer = slice(draw(st.integers(min_value=0, max_value=ndim))) if subset else slice(None) + return "/".join(map(str, blockidx[subset_slicer])) diff --git a/packages/zarr-storage/src/zarr_storage/testing/utils.py b/packages/zarr-storage/src/zarr_storage/testing/utils.py new file mode 100644 index 0000000000..b86801f90d --- /dev/null +++ b/packages/zarr-storage/src/zarr_storage/testing/utils.py @@ -0,0 +1,45 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, cast + +import pytest +from zarr.core.buffer import Buffer + +if TYPE_CHECKING: + from zarr.core.common import BytesLike + +__all__ = ["assert_bytes_equal"] + + +def assert_bytes_equal(b1: Buffer | BytesLike | None, b2: Buffer | BytesLike | None) -> None: + """Help function to assert if two bytes-like or Buffers are equal + + Warnings + -------- + Always copies data, only use for testing and debugging + """ + if isinstance(b1, Buffer): + b1 = b1.to_bytes() + if isinstance(b2, Buffer): + b2 = b2.to_bytes() + assert b1 == b2 + + +def has_cupy() -> bool: + try: + import cupy + + return cast("bool", cupy.cuda.runtime.getDeviceCount() > 0) + except ImportError: + return False + except cupy.cuda.runtime.CUDARuntimeError: + return False + + +gpu_mark = pytest.mark.gpu +skip_if_no_gpu = pytest.mark.skipif(not has_cupy(), reason="CuPy not installed or no GPU available") + + +# Decorator for GPU tests +def gpu_test[T](func: T) -> T: + return cast(T, gpu_mark(skip_if_no_gpu(func))) diff --git a/packages/zarr-storage/tests/test_legacy_contract.py b/packages/zarr-storage/tests/test_legacy_contract.py new file mode 100644 index 0000000000..6bb216cb1f --- /dev/null +++ b/packages/zarr-storage/tests/test_legacy_contract.py @@ -0,0 +1,52 @@ +"""Check the extracted contract against the Zarr checkout used for extraction.""" + +from __future__ import annotations + +import importlib +import importlib.util +import inspect + +import pytest +from zarr.abc import store as original + + +def test_storage_package_exists() -> None: + assert importlib.util.find_spec("zarr_storage") is not None + + +@pytest.mark.parametrize( + "name", + [ + "Store", + "ByteGetter", + "ByteSetter", + "SyncByteGetter", + "SyncByteSetter", + "SupportsGetSync", + "SupportsSetSync", + "SupportsDeleteSync", + "SupportsSyncStore", + "RangeByteRequest", + "OffsetByteRequest", + "SuffixByteRequest", + ], +) +def test_legacy_storage_signatures(name: str) -> None: + extracted = importlib.import_module("zarr_storage.legacy") + before = getattr(original, name) + after = getattr(extracted, name) + assert before is not after + assert str(inspect.signature(before)) == str(inspect.signature(after)) + assert getattr(before, "__abstractmethods__", None) == getattr( + after, "__abstractmethods__", None + ) + for member_name, member in vars(before).items(): + if isinstance(member, (classmethod, staticmethod)): + member = member.__func__ + replacement = vars(after)[member_name].__func__ + elif inspect.isfunction(member): + replacement = vars(after)[member_name] + else: + continue + assert str(inspect.signature(member)) == str(inspect.signature(replacement)) + assert inspect.iscoroutinefunction(member) == inspect.iscoroutinefunction(replacement) diff --git a/packages/zarr-storage/tests/test_storage_behavior.py b/packages/zarr-storage/tests/test_storage_behavior.py new file mode 100644 index 0000000000..1132b04c90 --- /dev/null +++ b/packages/zarr-storage/tests/test_storage_behavior.py @@ -0,0 +1,142 @@ +from __future__ import annotations + +import pickle +from typing import TYPE_CHECKING, Self + +import pytest + +if TYPE_CHECKING: + from collections.abc import AsyncIterator, Iterable +from zarr.core.buffer import Buffer, BufferPrototype, default_buffer_prototype +from zarr.core.buffer.cpu import Buffer as CpuBuffer + +from zarr_storage.legacy import ( + ByteRequest, + OffsetByteRequest, + RangeByteRequest, + Store, + SuffixByteRequest, +) + + +class DictStore(Store): + """A third-party-style store exercising inherited implementations.""" + + supports_writes = True + supports_deletes = True + supports_listing = True + + def __init__(self, *, read_only: bool = False) -> None: + super().__init__(read_only=read_only) + self.data: dict[str, bytes] = {} + self.reads = 0 + + def __eq__(self, other: object) -> bool: + return self is other + + def with_read_only(self, read_only: bool = False) -> Self: + result = type(self)(read_only=read_only) + result.data = self.data + return result + + async def get( + self, key: str, prototype: BufferPrototype, byte_range: ByteRequest | None = None + ) -> Buffer | None: + self.reads += 1 + value = self.data.get(key) + if value is None: + return None + match byte_range: + case RangeByteRequest(start, end): + value = value[start:end] + case OffsetByteRequest(offset): + value = value[offset:] + case SuffixByteRequest(suffix): + value = value[-suffix:] + return prototype.buffer.from_bytes(value) + + async def get_partial_values( + self, prototype: BufferPrototype, key_ranges: Iterable[tuple[str, ByteRequest | None]] + ) -> list[Buffer | None]: + return [await self.get(key, prototype, request) for key, request in key_ranges] + + async def exists(self, key: str) -> bool: + return key in self.data + + async def set(self, key: str, value: Buffer) -> None: + self._check_writable() + self.data[key] = value.to_bytes() + + async def delete(self, key: str) -> None: + self._check_writable() + self.data.pop(key, None) + + async def list(self) -> AsyncIterator[str]: + for key in list(self.data): + yield key + + async def list_prefix(self, prefix: str) -> AsyncIterator[str]: + for key in list(self.data): + if key.startswith(prefix): + yield key + + async def list_dir(self, prefix: str) -> AsyncIterator[str]: + prefix = prefix.rstrip("/") + "/" if prefix else "" + for child in sorted( + {key[len(prefix) :].split("/")[0] for key in self.data if key.startswith(prefix)} + ): + yield child + + +async def test_inherited_store_operations() -> None: + store = await DictStore.open() + assert store._is_open + await store.set_if_not_exists("a/one", CpuBuffer.from_bytes(b"0123456789")) + await store.set_if_not_exists("a/one", CpuBuffer.from_bytes(b"replacement")) + await store.set("a/two", CpuBuffer.from_bytes(b"abc")) + await store.set("b/three", CpuBuffer.from_bytes(b"xyz")) + assert await store.getsize("a/one") == 10 + assert await store.getsize_prefix("a/") == 13 + assert not await store.is_empty("a/") + await store.delete_dir("a") + assert store.data == {"b/three": b"xyz"} + await store.clear() + assert await store.is_empty("") + store.close() + assert not store._is_open + + +async def test_coalesced_ranges_use_extracted_request_types() -> None: + store = DictStore() + store.data["key"] = b"0123456789" + requests = [RangeByteRequest(5, 8), RangeByteRequest(0, 3), SuffixByteRequest(2)] + results = [ + item + async for batch in store.get_ranges("key", requests, prototype=default_buffer_prototype()) + for item in batch + ] + assert {index: value.to_bytes() for index, value in results if value is not None} == { + 0: b"567", + 1: b"012", + 2: b"89", + } + assert store.reads == 2 + + +async def test_getsize_missing_key() -> None: + with pytest.raises(FileNotFoundError, match="missing"): + await DictStore().getsize("missing") + + +async def test_read_only_write() -> None: + with pytest.raises(ValueError, match="read-only"): + await DictStore(read_only=True).set("key", CpuBuffer.from_bytes(b"data")) + + +@pytest.mark.parametrize( + "byte_request", [RangeByteRequest(1, 3), OffsetByteRequest(2), SuffixByteRequest(4)] +) +def test_request_pickle(byte_request: ByteRequest) -> None: + restored = pickle.loads(pickle.dumps(byte_request)) + assert type(restored) is type(byte_request) + assert restored == byte_request diff --git a/packages/zarr-storage/tests/test_store/__init__.py b/packages/zarr-storage/tests/test_store/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/packages/zarr-storage/tests/test_store/conftest.py b/packages/zarr-storage/tests/test_store/conftest.py new file mode 100644 index 0000000000..5abcb5bfcf --- /dev/null +++ b/packages/zarr-storage/tests/test_store/conftest.py @@ -0,0 +1,220 @@ +from __future__ import annotations + +import importlib +import os +import pathlib +import sys +from typing import TYPE_CHECKING + +import pytest +from hypothesis import HealthCheck, settings +from zarr import config +from zarr.core.sync import sync + +from zarr_storage.legacy import ( + FsspecStore, + LatencyStore, + LocalStore, + MemoryStore, + Store, + StorePath, + ZipStore, +) + +if TYPE_CHECKING: + from collections.abc import Generator + from typing import Any, Literal + + from zarr.core.common import ZarrFormat + + +@pytest.fixture(autouse=True) +def extracted_storage_integration(monkeypatch: pytest.MonkeyPatch) -> None: + """Exercise real Zarr arrays with the proposed storage re-exports. + + This is only a test harness for future runtime adoption. Rebind references + to the old storage definitions in already imported Zarr modules, including + aliases imported with ``from ... import ...``. No methods are mocked and + importing zarr_storage itself never changes Zarr's runtime bindings. + Pytest restores every binding at teardown; contract tests outside this + directory continue to compare against the unmodified original classes. + """ + module_pairs = { + "zarr.abc.store": "zarr_storage.legacy._abc", + "zarr.core._coalesce": "zarr_storage._coalesce", + "zarr.testing.store": "zarr_storage.testing.store", + "zarr.experimental.cache_store": "zarr_storage.legacy.experimental.cache_store", + } + for suffix in ( + "_common", + "_fsspec", + "_local", + "_logging", + "_memory", + "_obstore", + "_utils", + "_wrapper", + "_zip", + ): + module_pairs[f"zarr.storage.{suffix}"] = f"zarr_storage.legacy.{suffix}" + replacements: dict[int, object] = {} + for old_name, new_name in module_pairs.items(): + old = importlib.import_module(old_name) + new = importlib.import_module(new_name) + for name, value in vars(old).items(): + if ( + getattr(value, "__module__", None) == old_name + or name in {"StoreLike", "ByteRequest"} + ) and hasattr(new, name): + replacements[id(value)] = getattr(new, name) + for module_name, module in tuple(sys.modules.items()): + if module_name == "zarr" or module_name.startswith("zarr."): + for name, value in tuple(vars(module).items()): + if id(value) in replacements: + monkeypatch.setattr(module, name, replacements[id(value)]) + + +async def parse_store( + store: Literal["local", "memory", "fsspec", "zip", "memory_get_latency"], path: str +) -> LocalStore | MemoryStore | FsspecStore | ZipStore | LatencyStore: + if store == "local": + return await LocalStore.open(path) + if store == "memory": + return await MemoryStore.open() + if store == "fsspec": + return await FsspecStore.open(url=path) + if store == "zip": + return await ZipStore.open(f"{path}/zarr.zip", mode="w") + if store == "memory_get_latency": + return LatencyStore(MemoryStore(), get_latency=0.0001, set_latency=0.0) + raise AssertionError + + +@pytest.fixture(params=[str, pathlib.Path]) +def path_type(request: pytest.FixtureRequest) -> Any: + return request.param + + +@pytest.fixture +async def store_path(tmp_path: pathlib.Path) -> StorePath: + store = await LocalStore.open(str(tmp_path)) + return StorePath(store) + + +@pytest.fixture +async def local_store(tmp_path: pathlib.Path) -> LocalStore: + return await LocalStore.open(str(tmp_path)) + + +@pytest.fixture +async def memory_store() -> MemoryStore: + return await MemoryStore.open() + + +@pytest.fixture +async def zip_store(tmp_path: pathlib.Path) -> ZipStore: + return await ZipStore.open(str(tmp_path / "zarr.zip"), mode="w") + + +@pytest.fixture +async def store(request: pytest.FixtureRequest, tmp_path: pathlib.Path) -> Store: + param = request.param + return await parse_store(param, str(tmp_path)) + + +@pytest.fixture +async def store2(request: pytest.FixtureRequest, tmp_path: pathlib.Path) -> Store: + """Fixture to create a second store for testing copy operations between stores""" + param = request.param + store2_path = tmp_path / "store2" + store2_path.mkdir() + return await parse_store(param, str(store2_path)) + + +@pytest.fixture(params=["local", "memory", "zip"]) +def sync_store(request: pytest.FixtureRequest, tmp_path: pathlib.Path) -> Store: + result = sync(parse_store(request.param, str(tmp_path))) + if not isinstance(result, Store): + raise TypeError(f"Wrong store class returned by test fixture! got {result} instead") + return result + + +@pytest.fixture(params=["numpy", "cupy"]) +def xp(request: pytest.FixtureRequest) -> Any: + """Fixture to parametrize over numpy-like libraries""" + + if request.param == "cupy": + request.node.add_marker(pytest.mark.gpu) + + return pytest.importorskip(request.param) + + +@pytest.fixture(autouse=True) +def reset_config() -> Generator[None, None, None]: + config.reset() + yield + config.reset() + + +@pytest.fixture(params=(2, 3), ids=["zarr2", "zarr3"]) +def zarr_format(request: pytest.FixtureRequest) -> ZarrFormat: + if request.param == 2: + return 2 + elif request.param == 3: + return 3 + msg = f"Invalid zarr format requested. Got {request.param}, expected on of (2,3)." + raise ValueError(msg) + + +def pytest_addoption(parser: Any) -> None: + parser.addoption( + "--run-slow-hypothesis", + action="store_true", + default=False, + help="run slow hypothesis tests", + ) + + +def pytest_collection_modifyitems(config: Any, items: Any) -> None: + if config.getoption("--run-slow-hypothesis"): + return + skip_slow_hyp = pytest.mark.skip(reason="need --run-slow-hypothesis option to run") + for item in items: + if "slow_hypothesis" in item.keywords: + item.add_marker(skip_slow_hyp) + + +@pytest.fixture(scope="session") +def moto_server() -> Generator[str, None, None]: + """Start a session-scoped moto S3 server and yield its endpoint URL. + + The server binds an ephemeral port (port=0), so the endpoint is only known at + runtime; consumers must take it from this fixture rather than a constant. A fixed + port deadlocks under pytest-xdist: session-scoped fixtures run once per *worker*, so + concurrent workers race to bind the same port, and the losers block forever inside + ThreadedMotoServer.start(), whose server thread dies on "Address already in use" + before ever setting the ready event that start() waits on. + + importorskip lives inside the fixture so moto is only required when a test actually + requests an S3 backend, not for the whole test session.""" + moto_server_mod = pytest.importorskip("moto.moto_server.threaded_moto_server") + + server = moto_server_mod.ThreadedMotoServer(ip_address="127.0.0.1", port=0) + server.start() + host, port = server.get_host_and_port() + # moto needs *some* credentials present; use throwaway values if the environment has none. + os.environ.setdefault("AWS_SECRET_ACCESS_KEY", "foo") + os.environ.setdefault("AWS_ACCESS_KEY_ID", "foo") + try: + yield f"http://{host}:{port}/" + finally: + server.stop() + + +settings.register_profile( + "storage", + max_examples=50, + deadline=None, + suppress_health_check=[HealthCheck.filter_too_much, HealthCheck.too_slow], +) +settings.load_profile(os.environ.get("HYPOTHESIS_PROFILE", "storage")) diff --git a/packages/zarr-storage/tests/test_store/test_cache_store.py b/packages/zarr-storage/tests/test_store/test_cache_store.py new file mode 100644 index 0000000000..178cee288c --- /dev/null +++ b/packages/zarr-storage/tests/test_store/test_cache_store.py @@ -0,0 +1,1075 @@ +""" +Tests for the dual-store cache implementation. +""" + +import asyncio +import time + +import pytest +from zarr.core.buffer.core import default_buffer_prototype +from zarr.core.buffer.cpu import Buffer as CPUBuffer + +from zarr_storage.legacy import MemoryStore +from zarr_storage.legacy._abc import RangeByteRequest, Store, SuffixByteRequest +from zarr_storage.legacy.experimental.cache_store import CacheStore + + +class TestCacheStore: + """Test the dual-store cache implementation.""" + + @pytest.fixture + def source_store(self) -> MemoryStore: + """Create a source store with some test data.""" + return MemoryStore() + + @pytest.fixture + def cache_store(self) -> MemoryStore: + """Create an empty cache store.""" + return MemoryStore() + + @pytest.fixture + def cached_store(self, source_store: Store, cache_store: Store) -> CacheStore: + """Create a cached store instance.""" + return CacheStore(source_store, cache_store=cache_store) + + async def test_with_read_only_round_trip(self) -> None: + """ + Ensure that CacheStore.with_read_only returns another CacheStore with + the requested read_only state, shares cache state, and does not change + the original store's read_only flag. + """ + source = MemoryStore() + cache = MemoryStore() + + # Start from a read-only underlying store + source_ro = source.with_read_only(read_only=True) + cached_ro = CacheStore(store=source_ro, cache_store=cache) + assert cached_ro.read_only + + buf = CPUBuffer.from_bytes(b"0123") + + # Cannot write through the read-only cache store + with pytest.raises( + ValueError, match="store was opened in read-only mode and does not support writing" + ): + await cached_ro.set("foo", buf) + + # Create a writable cache store from the read-only one + writer = cached_ro.with_read_only(read_only=False) + assert isinstance(writer, CacheStore) + assert not writer.read_only + + # Cache configuration and state are shared + assert writer._cache is cached_ro._cache + assert writer._state is cached_ro._state + assert writer._state.key_insert_times is cached_ro._state.key_insert_times + + # Writes via the writable cache store succeed and are cached + await writer.set("foo", buf) + out = await writer.get("foo", default_buffer_prototype()) + assert out is not None + assert out.to_bytes() == buf.to_bytes() + + # The original cache store remains read-only + assert cached_ro.read_only + with pytest.raises( + ValueError, match="store was opened in read-only mode and does not support writing" + ): + await cached_ro.set("bar", buf) + + # Creating a read-only copy from the writable cache store works and is enforced + reader = writer.with_read_only(read_only=True) + assert isinstance(reader, CacheStore) + assert reader.read_only + with pytest.raises( + ValueError, match="store was opened in read-only mode and does not support writing" + ): + await reader.set("baz", buf) + + async def test_basic_caching(self, cached_store: CacheStore, source_store: Store) -> None: + """Test basic cache functionality.""" + # Store some data + test_data = CPUBuffer.from_bytes(b"test data") + await cached_store.set("test_key", test_data) + + # Verify it's in both stores + assert await source_store.exists("test_key") + assert await cached_store._cache.exists("test_key") + + # Retrieve and verify caching works + result = await cached_store.get("test_key", default_buffer_prototype()) + assert result is not None + assert result.to_bytes() == b"test data" + + async def test_cache_miss_and_population( + self, cached_store: CacheStore, source_store: Store + ) -> None: + """Test cache miss and subsequent population.""" + # Put data directly in source store (bypassing cache) + test_data = CPUBuffer.from_bytes(b"source data") + await source_store.set("source_key", test_data) + + # First access should miss cache but populate it + result = await cached_store.get("source_key", default_buffer_prototype()) + assert result is not None + assert result.to_bytes() == b"source data" + + # Verify data is now in cache + assert await cached_store._cache.exists("source_key") + + async def test_cache_expiration(self) -> None: + """Test cache expiration based on max_age_seconds.""" + source_store = MemoryStore() + cache_store = MemoryStore() + cached_store = CacheStore( + source_store, + cache_store=cache_store, + max_age_seconds=1, # 1 second expiration + ) + + # Store data + test_data = CPUBuffer.from_bytes(b"expiring data") + await cached_store.set("expire_key", test_data) + + # Should be fresh initially + assert cached_store._is_key_fresh("expire_key") + + # Wait for expiration + await asyncio.sleep(1.1) + + # Should now be stale + assert not cached_store._is_key_fresh("expire_key") + + async def test_cache_set_data_false(self, source_store: Store, cache_store: Store) -> None: + """Test behavior when cache_set_data=False.""" + cached_store = CacheStore(source_store, cache_store=cache_store, cache_set_data=False) + + test_data = CPUBuffer.from_bytes(b"no cache data") + await cached_store.set("no_cache_key", test_data) + + # Data should be in source but not cache + assert await source_store.exists("no_cache_key") + assert not await cache_store.exists("no_cache_key") + + async def test_delete_removes_from_both_stores(self, cached_store: CacheStore) -> None: + """Test that delete removes from both source and cache.""" + test_data = CPUBuffer.from_bytes(b"delete me") + await cached_store.set("delete_key", test_data) + + # Verify in both stores + assert await cached_store._store.exists("delete_key") + assert await cached_store._cache.exists("delete_key") + + # Delete + await cached_store.delete("delete_key") + + # Verify removed from both + assert not await cached_store._store.exists("delete_key") + assert not await cached_store._cache.exists("delete_key") + + async def test_exists_checks_source_store( + self, cached_store: CacheStore, source_store: Store + ) -> None: + """Test that exists() checks the source store (source of truth).""" + # Put data directly in source + test_data = CPUBuffer.from_bytes(b"exists test") + await source_store.set("exists_key", test_data) + + # Should exist even though not in cache + assert await cached_store.exists("exists_key") + + async def test_list_operations(self, cached_store: CacheStore, source_store: Store) -> None: + """Test listing operations delegate to source store.""" + # Add some test data + test_data = CPUBuffer.from_bytes(b"list test") + await cached_store.set("list/item1", test_data) + await cached_store.set("list/item2", test_data) + await cached_store.set("other/item3", test_data) + + # Test list_dir + list_items = [key async for key in cached_store.list_dir("list/")] + assert len(list_items) >= 2 # Should include our items + + # Test list_prefix + prefix_items = [key async for key in cached_store.list_prefix("list/")] + assert len(prefix_items) >= 2 + + async def test_stale_cache_refresh(self) -> None: + """Test that stale cache entries are refreshed from source.""" + source_store = MemoryStore() + cache_store = MemoryStore() + cached_store = CacheStore(source_store, cache_store=cache_store, max_age_seconds=1) + + # Store initial data + old_data = CPUBuffer.from_bytes(b"old data") + await cached_store.set("refresh_key", old_data) + + # Wait for expiration + await asyncio.sleep(1.1) + + # Update source store directly (simulating external update) + new_data = CPUBuffer.from_bytes(b"new data") + await source_store.set("refresh_key", new_data) + + # Access should refresh from source when cache is stale + result = await cached_store.get("refresh_key", default_buffer_prototype()) + assert result is not None + assert result.to_bytes() == b"new data" + + async def test_infinity_max_age(self, cached_store: CacheStore) -> None: + """Test that 'infinity' max_age means cache never expires.""" + test_data = CPUBuffer.from_bytes(b"eternal data") + await cached_store.set("eternal_key", test_data) + + # Should always be fresh + assert cached_store._is_key_fresh("eternal_key") + + # Even after time passes + await asyncio.sleep(0.1) + assert cached_store._is_key_fresh("eternal_key") + + async def test_cache_returns_cached_data_for_performance( + self, cached_store: CacheStore, source_store: Store + ) -> None: + """Test that cache returns cached data for performance, even if not in source.""" + # Put data in cache but not source (simulates orphaned cache entry) + test_data = CPUBuffer.from_bytes(b"orphaned data") + await cached_store._cache.set("orphan_key", test_data) + cached_store._state.key_insert_times["orphan_key"] = time.monotonic() + + # Cache should return data for performance (no source verification) + result = await cached_store.get("orphan_key", default_buffer_prototype()) + assert result is not None + assert result.to_bytes() == b"orphaned data" + + # Cache entry should remain (performance optimization) + assert await cached_store._cache.exists("orphan_key") + assert "orphan_key" in cached_store._state.key_insert_times + + async def test_cache_coherency_through_expiration(self) -> None: + """Test that cache coherency is managed through cache expiration, not source verification.""" + source_store = MemoryStore() + cache_store = MemoryStore() + cached_store = CacheStore( + source_store, + cache_store=cache_store, + max_age_seconds=1, # Short expiration for coherency + ) + + # Add data to both stores + test_data = CPUBuffer.from_bytes(b"original data") + await cached_store.set("coherency_key", test_data) + + # Remove from source (simulating external deletion) + await source_store.delete("coherency_key") + + # Cache should still return cached data (performance optimization) + result = await cached_store.get("coherency_key", default_buffer_prototype()) + assert result is not None + assert result.to_bytes() == b"original data" + + # Wait for cache expiration + await asyncio.sleep(1.1) + + # Now stale cache should be refreshed from source + result = await cached_store.get("coherency_key", default_buffer_prototype()) + assert result is None # Key no longer exists in source + + async def test_cache_info(self, cached_store: CacheStore) -> None: + """Test cache_info method returns correct information.""" + # Test initial state + info = cached_store.cache_info() + + # Check all expected keys are present + expected_keys = { + "cache_store_type", + "max_age_seconds", + "max_size", + "current_size", + "cache_set_data", + "tracked_keys", + "cached_keys", + } + assert set(info.keys()) == expected_keys + + # Check initial values + assert info["cache_store_type"] == "MemoryStore" + assert info["max_age_seconds"] == "infinity" + assert info["max_size"] is None # Default unlimited + assert info["current_size"] == 0 + assert info["cache_set_data"] is True + assert info["tracked_keys"] == 0 + assert info["cached_keys"] == 0 + + # Add some data and verify tracking + test_data = CPUBuffer.from_bytes(b"test data for cache info") + await cached_store.set("info_test_key", test_data) + + # Check updated info + updated_info = cached_store.cache_info() + assert updated_info["tracked_keys"] == 1 + assert updated_info["cached_keys"] == 1 + assert updated_info["current_size"] > 0 # Should have some size now + + async def test_cache_info_with_max_size(self) -> None: + """Test cache_info with max_size configuration.""" + source_store = MemoryStore() + cache_store = MemoryStore() + + # Create cache with specific max_size and max_age + cached_store = CacheStore( + source_store, + cache_store=cache_store, + max_size=1024, + max_age_seconds=300, + ) + + info = cached_store.cache_info() + assert info["max_size"] == 1024 + assert info["max_age_seconds"] == 300 + assert info["current_size"] == 0 + + async def test_clear_cache(self, cached_store: CacheStore) -> None: + """Test clear_cache method clears all cache data and tracking.""" + # Add some test data + test_data1 = CPUBuffer.from_bytes(b"test data 1") + test_data2 = CPUBuffer.from_bytes(b"test data 2") + + await cached_store.set("clear_test_1", test_data1) + await cached_store.set("clear_test_2", test_data2) + + # Verify data is cached + info_before = cached_store.cache_info() + assert info_before["tracked_keys"] == 2 + assert info_before["cached_keys"] == 2 + assert info_before["current_size"] > 0 + + # Verify data exists in cache + assert await cached_store._cache.exists("clear_test_1") + assert await cached_store._cache.exists("clear_test_2") + + # Clear the cache + await cached_store.clear_cache() + + # Verify cache is cleared + info_after = cached_store.cache_info() + assert info_after["tracked_keys"] == 0 + assert info_after["cached_keys"] == 0 + assert info_after["current_size"] == 0 + + # Verify data is removed from cache store (if it supports clear) + if hasattr(cached_store._cache, "clear"): + # If cache store supports clear, all data should be gone + assert not await cached_store._cache.exists("clear_test_1") + assert not await cached_store._cache.exists("clear_test_2") + + # Verify data still exists in source store + assert await cached_store._store.exists("clear_test_1") + assert await cached_store._store.exists("clear_test_2") + + async def test_max_age_infinity(self) -> None: + """Test cache with infinite max age.""" + source_store = MemoryStore() + cache_store = MemoryStore() + cached_store = CacheStore(source_store, cache_store=cache_store, max_age_seconds="infinity") + + # Add data and verify it never expires + test_data = CPUBuffer.from_bytes(b"test data") + await cached_store.set("test_key", test_data) + + # Even after time passes, key should be fresh + assert cached_store._is_key_fresh("test_key") + + async def test_max_age_numeric(self) -> None: + """Test cache with numeric max age.""" + source_store = MemoryStore() + cache_store = MemoryStore() + cached_store = CacheStore( + source_store, + cache_store=cache_store, + max_age_seconds=1, # 1 second + ) + + # Add data + test_data = CPUBuffer.from_bytes(b"test data") + await cached_store.set("test_key", test_data) + + # Key should be fresh initially + assert cached_store._is_key_fresh("test_key") + + # Manually set old timestamp to test expiration + cached_store._state.key_insert_times["test_key"] = time.monotonic() - 2 # 2 seconds ago + + # Key should now be stale + assert not cached_store._is_key_fresh("test_key") + + async def test_cache_set_data_disabled(self) -> None: + """Test cache behavior when cache_set_data is False.""" + source_store = MemoryStore() + cache_store = MemoryStore() + cached_store = CacheStore(source_store, cache_store=cache_store, cache_set_data=False) + + # Set data + test_data = CPUBuffer.from_bytes(b"test data") + await cached_store.set("test_key", test_data) + + # Data should be in source but not in cache + assert await source_store.exists("test_key") + assert not await cache_store.exists("test_key") + + # Cache info should show no cached data + info = cached_store.cache_info() + assert info["cache_set_data"] is False + assert info["cached_keys"] == 0 + + async def test_eviction_with_max_size(self) -> None: + """Test LRU eviction when max_size is exceeded.""" + source_store = MemoryStore() + cache_store = MemoryStore() + cached_store = CacheStore( + source_store, + cache_store=cache_store, + max_size=100, # Small cache size + ) + + # Add data that exceeds cache size + small_data = CPUBuffer.from_bytes(b"a" * 40) # 40 bytes + medium_data = CPUBuffer.from_bytes(b"b" * 40) # 40 bytes + large_data = CPUBuffer.from_bytes(b"c" * 40) # 40 bytes (would exceed 100 byte limit) + + # Set first two items + await cached_store.set("key1", small_data) + await cached_store.set("key2", medium_data) + + # Cache should have 2 items + info = cached_store.cache_info() + assert info["cached_keys"] == 2 + assert info["current_size"] == 80 + + # Add third item - should trigger eviction of first item + await cached_store.set("key3", large_data) + + # Cache should still have items but first one may be evicted + info = cached_store.cache_info() + assert info["current_size"] <= 100 + + async def test_value_exceeds_max_size(self) -> None: + """Test behavior when a single value exceeds max_size.""" + source_store = MemoryStore() + cache_store = MemoryStore() + cached_store = CacheStore( + source_store, + cache_store=cache_store, + max_size=50, # Small cache size + ) + + # Try to cache data larger than max_size + large_data = CPUBuffer.from_bytes(b"x" * 100) # 100 bytes > 50 byte limit + await cached_store.set("large_key", large_data) + + # Data should be in source but not cached + assert await source_store.exists("large_key") + info = cached_store.cache_info() + assert info["cached_keys"] == 0 + assert info["current_size"] == 0 + + async def test_get_nonexistent_key(self) -> None: + """Test getting a key that doesn't exist in either store.""" + source_store = MemoryStore() + cache_store = MemoryStore() + cached_store = CacheStore(source_store, cache_store=cache_store) + + # Try to get nonexistent key + result = await cached_store.get("nonexistent", default_buffer_prototype()) + assert result is None + + # Should not create any cache entries + info = cached_store.cache_info() + assert info["cached_keys"] == 0 + + async def test_delete_both_stores(self) -> None: + """Test that delete removes from both source and cache stores.""" + source_store = MemoryStore() + cache_store = MemoryStore() + cached_store = CacheStore(source_store, cache_store=cache_store) + + # Add data + test_data = CPUBuffer.from_bytes(b"test data") + await cached_store.set("test_key", test_data) + + # Verify it's in both stores + assert await source_store.exists("test_key") + assert await cache_store.exists("test_key") + + # Delete + await cached_store.delete("test_key") + + # Verify it's removed from both + assert not await source_store.exists("test_key") + assert not await cache_store.exists("test_key") + + # Verify tracking is updated + info = cached_store.cache_info() + assert info["cached_keys"] == 0 + + async def test_invalid_max_age_seconds(self) -> None: + """Test that invalid max_age_seconds values raise ValueError.""" + source_store = MemoryStore() + cache_store = MemoryStore() + + with pytest.raises(ValueError, match="max_age_seconds string value must be 'infinity'"): + CacheStore(source_store, cache_store=cache_store, max_age_seconds="invalid") + + async def test_unlimited_cache_size(self) -> None: + """Test behavior when max_size is None (unlimited).""" + source_store = MemoryStore() + cache_store = MemoryStore() + cached_store = CacheStore( + source_store, + cache_store=cache_store, + max_size=None, # Unlimited cache + ) + + # Add large amounts of data + for i in range(10): + large_data = CPUBuffer.from_bytes(b"x" * 1000) # 1KB each + await cached_store.set(f"large_key_{i}", large_data) + + # All should be cached since there's no size limit + info = cached_store.cache_info() + assert info["cached_keys"] == 10 + assert info["current_size"] == 10000 # 10 * 1000 bytes + + async def test_evict_key_exception_handling(self) -> None: + """Test exception handling in _evict_key method.""" + source_store = MemoryStore() + cache_store = MemoryStore() + cached_store = CacheStore(source_store, cache_store=cache_store, max_size=100) + + # Add some data + test_data = CPUBuffer.from_bytes(b"test data") + await cached_store.set("test_key", test_data) + + # Manually corrupt the tracking to trigger exception + # Remove from one structure but not others to create inconsistency + del cached_store._state.cache_order["test_key"] + + # Try to evict - should handle the KeyError gracefully + await cached_store._evict_key("test_key") + + # Should still work and not crash + info = cached_store.cache_info() + assert isinstance(info, dict) + + async def test_get_no_cache_delete_tracking(self) -> None: + """Test _get_no_cache when key doesn't exist and needs cleanup.""" + source_store = MemoryStore() + cache_store = MemoryStore() + cached_store = CacheStore(source_store, cache_store=cache_store) + + # First, add key to cache tracking but not to source + test_data = CPUBuffer.from_bytes(b"test data") + await cache_store.set("phantom_key", test_data) + await cached_store._track_entry("phantom_key", test_data) + + # Verify it's in tracking + assert "phantom_key" in cached_store._state.cache_order + assert "phantom_key" in cached_store._state.key_insert_times + + # Now try to get it - since it's not in source, should clean up tracking + result = await cached_store._get_no_cache("phantom_key", default_buffer_prototype()) + assert result is None + + # Should have cleaned up tracking + assert "phantom_key" not in cached_store._state.cache_order + assert "phantom_key" not in cached_store._state.key_insert_times + + async def test_accommodate_value_no_max_size(self) -> None: + """Test _accommodate_value early return when max_size is None.""" + source_store = MemoryStore() + cache_store = MemoryStore() + cached_store = CacheStore( + source_store, + cache_store=cache_store, + max_size=None, # No size limit + ) + + # This should return early without doing anything + await cached_store._accommodate_value(1000000) # Large value + + # Should not affect anything since max_size is None + info = cached_store.cache_info() + assert info["current_size"] == 0 + + async def test_concurrent_set_operations(self) -> None: + """Test that concurrent set operations don't corrupt cache size tracking.""" + source_store = MemoryStore() + cache_store = MemoryStore() + cached_store = CacheStore(source_store, cache_store=cache_store, max_size=1000) + + # Create 10 concurrent set operations + async def set_data(key: str) -> None: + data = CPUBuffer.from_bytes(b"x" * 50) + await cached_store.set(key, data) + + # Run concurrently + await asyncio.gather(*[set_data(f"key_{i}") for i in range(10)]) + + info = cached_store.cache_info() + # Expected: 10 keys * 50 bytes = 500 bytes + assert info["cached_keys"] == 10 + assert info["current_size"] == 500 # WOULD FAIL due to race condition + + async def test_concurrent_eviction_race(self) -> None: + """Test concurrent evictions don't corrupt size tracking.""" + source_store = MemoryStore() + cache_store = MemoryStore() + cached_store = CacheStore(source_store, cache_store=cache_store, max_size=200) + + # Fill cache to near capacity + data = CPUBuffer.from_bytes(b"x" * 80) + await cached_store.set("key1", data) + await cached_store.set("key2", data) + + # Now trigger two concurrent sets that both need to evict + async def set_large(key: str) -> None: + large_data = CPUBuffer.from_bytes(b"y" * 100) + await cached_store.set(key, large_data) + + await asyncio.gather(set_large("key3"), set_large("key4")) + + info = cached_store.cache_info() + # Size should be consistent with tracked keys + assert info["current_size"] <= 200 # Might pass + # But verify actual cache store size matches tracking + total_size = sum( + cached_store._state.key_sizes.get(k, 0) for k in cached_store._state.cache_order + ) + assert total_size == info["current_size"] # WOULD FAIL + + async def test_concurrent_get_and_evict(self) -> None: + """Test get operations during eviction don't cause corruption.""" + source_store = MemoryStore() + cache_store = MemoryStore() + cached_store = CacheStore(source_store, cache_store=cache_store, max_size=100) + + # Setup + data = CPUBuffer.from_bytes(b"x" * 40) + await cached_store.set("key1", data) + await cached_store.set("key2", data) + + # Concurrent: read key1 while adding key3 (triggers eviction) + async def read_key() -> None: + for _ in range(100): + await cached_store.get("key1", default_buffer_prototype()) + + async def write_key() -> None: + for i in range(10): + new_data = CPUBuffer.from_bytes(b"y" * 40) + await cached_store.set(f"new_{i}", new_data) + + await asyncio.gather(read_key(), write_key()) + + # Verify consistency + info = cached_store.cache_info() + assert info["current_size"] <= 100 + assert len(cached_store._state.cache_order) == len(cached_store._state.key_sizes) + + async def test_eviction_actually_deletes_from_cache_store(self) -> None: + """Test that eviction removes keys from cache_store, not just tracking.""" + source_store = MemoryStore() + cache_store = MemoryStore() + cached_store = CacheStore(source_store, cache_store=cache_store, max_size=100) + + # Add data that will be evicted + data1 = CPUBuffer.from_bytes(b"x" * 60) + data2 = CPUBuffer.from_bytes(b"y" * 60) + + await cached_store.set("key1", data1) + + # Verify key1 is in cache_store + assert await cache_store.exists("key1") + + # Add key2, which should evict key1 + await cached_store.set("key2", data2) + + # Check tracking - key1 should be removed + assert "key1" not in cached_store._state.cache_order + assert "key1" not in cached_store._state.key_sizes + + # CRITICAL: key1 should also be removed from cache_store + assert not await cache_store.exists("key1"), ( + "Evicted key still exists in cache_store! _evict_key doesn't actually delete." + ) + + # But key1 should still exist in source store + assert await source_store.exists("key1") + + async def test_eviction_no_orphaned_keys(self) -> None: + """Test that eviction doesn't leave orphaned keys in cache_store.""" + source_store = MemoryStore() + cache_store = MemoryStore() + cached_store = CacheStore(source_store, cache_store=cache_store, max_size=150) + + # Add multiple keys that will cause evictions + for i in range(10): + data = CPUBuffer.from_bytes(b"x" * 60) + await cached_store.set(f"key_{i}", data) + + # Check tracking + info = cached_store.cache_info() + tracked_keys = info["cached_keys"] + + # Count actual keys in cache_store + actual_keys = 0 + async for _ in cache_store.list(): + actual_keys += 1 + + # Cache store should have same number of keys as tracking + assert actual_keys == tracked_keys, ( + f"Cache store has {actual_keys} keys but tracking shows {tracked_keys}. " + f"Eviction doesn't delete from cache_store!" + ) + + async def test_size_accounting_with_key_updates(self) -> None: + """Test that updating the same key replaces size instead of accumulating.""" + source_store = MemoryStore() + cache_store = MemoryStore() + cached_store = CacheStore(source_store, cache_store=cache_store, max_size=500) + + # Set initial value + data1 = CPUBuffer.from_bytes(b"x" * 100) + await cached_store.set("same_key", data1) + + info1 = cached_store.cache_info() + assert info1["current_size"] == 100 + + # Update with different size + data2 = CPUBuffer.from_bytes(b"y" * 200) + await cached_store.set("same_key", data2) + + info2 = cached_store.cache_info() + + # Should be 200, not 300 (update replaces, doesn't accumulate) + assert info2["current_size"] == 200, ( + f"Expected size 200 but got {info2['current_size']}. " + "Updating same key should replace, not accumulate." + ) + + async def test_all_tracked_keys_exist_in_cache_store(self) -> None: + """Test invariant: all keys in tracking should exist in cache_store.""" + source_store = MemoryStore() + cache_store = MemoryStore() + cached_store = CacheStore(source_store, cache_store=cache_store, max_size=500) + + # Add some data + for i in range(5): + data = CPUBuffer.from_bytes(b"x" * 50) + await cached_store.set(f"key_{i}", data) + + # Every str key in tracking should exist in cache_store + # (tuple keys are byte-range entries stored in-memory, not in the Store) + for entry_key in cached_store._state.cache_order: + if isinstance(entry_key, str): + assert await cache_store.exists(entry_key), ( + f"Key '{entry_key}' is tracked but doesn't exist in cache_store" + ) + + # Every str key in _key_sizes should exist in cache_store + for entry_key in cached_store._state.key_sizes: + if isinstance(entry_key, str): + assert await cache_store.exists(entry_key), ( + f"Key '{entry_key}' has size tracked but doesn't exist in cache_store" + ) + + # Additional coverage tests for 100% coverage + + async def test_cache_store_requires_delete_support(self) -> None: + """Test that CacheStore validates cache_store supports deletes.""" + from unittest.mock import MagicMock + + # Create a mock store that doesn't support deletes + source_store = MemoryStore() + cache_store = MagicMock() + cache_store.supports_deletes = False + + with pytest.raises(ValueError, match="does not support deletes"): + CacheStore(store=source_store, cache_store=cache_store) + + async def test_evict_key_exception_handling_with_real_error( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Test _evict_key exception handling when deletion fails.""" + source_store = MemoryStore() + cache_store = MemoryStore() + cached_store = CacheStore(store=source_store, cache_store=cache_store, max_size=100) + + # Set up a key in tracking + buffer = CPUBuffer.from_bytes(b"test data") + await cached_store.set("test_key", buffer) + + # Mock the cache delete to raise an exception + async def failing_delete(key: str) -> None: + raise RuntimeError("Simulated cache deletion failure") + + monkeypatch.setattr(cache_store, "delete", failing_delete) + + # Attempt to evict should raise the exception + with pytest.raises(RuntimeError, match="Simulated cache deletion failure"): + async with cached_store._state.lock: + await cached_store._evict_key("test_key") + + async def test_cache_stats_method(self) -> None: + """Test cache_stats method returns correct statistics.""" + source_store = MemoryStore() + cache_store = MemoryStore() + cached_store = CacheStore(store=source_store, cache_store=cache_store, max_size=1000) + + # Initially, stats should be zero + stats = cached_store.cache_stats() + assert stats["hits"] == 0 + assert stats["misses"] == 0 + assert stats["evictions"] == 0 + assert stats["total_requests"] == 0 + assert stats["hit_rate"] == 0.0 + + # Perform some operations + buffer = CPUBuffer.from_bytes(b"x" * 100) + + # Write to source store directly to avoid affecting stats + await source_store.set("key1", buffer) + + # First get is a miss (not in cache yet) + result1 = await cached_store.get("key1", default_buffer_prototype()) + assert result1 is not None + + # Second get is a hit (now in cache) + result2 = await cached_store.get("key1", default_buffer_prototype()) + assert result2 is not None + + stats = cached_store.cache_stats() + assert stats["hits"] == 1 + assert stats["misses"] == 1 + assert stats["total_requests"] == 2 + assert stats["hit_rate"] == 0.5 + + async def test_cache_stats_with_evictions(self) -> None: + """Test cache_stats tracks evictions correctly.""" + source_store = MemoryStore() + cache_store = MemoryStore() + cached_store = CacheStore( + store=source_store, + cache_store=cache_store, + max_size=150, # Small size to force eviction + ) + + # Add items that will trigger eviction + buffer1 = CPUBuffer.from_bytes(b"x" * 100) + buffer2 = CPUBuffer.from_bytes(b"y" * 100) + + await cached_store.set("key1", buffer1) + await cached_store.set("key2", buffer2) # Should evict key1 + + stats = cached_store.cache_stats() + assert stats["evictions"] == 1 + + def test_repr_method(self) -> None: + """Test __repr__ returns useful string representation.""" + source_store = MemoryStore() + cache_store = MemoryStore() + cached_store = CacheStore( + store=source_store, cache_store=cache_store, max_age_seconds=60, max_size=1024 + ) + + repr_str = repr(cached_store) + + # Check that repr contains key information + assert "CacheStore" in repr_str + assert "max_age_seconds=60" in repr_str + assert "max_size=1024" in repr_str + assert "current_size=0" in repr_str + assert "cached_keys=0" in repr_str + + async def test_cache_stats_zero_division_protection(self) -> None: + """Test cache_stats handles zero requests correctly.""" + source_store = MemoryStore() + cache_store = MemoryStore() + cached_store = CacheStore(store=source_store, cache_store=cache_store) + + # With no requests, hit_rate should be 0.0 (not NaN or error) + stats = cached_store.cache_stats() + assert stats["hit_rate"] == 0.0 + assert stats["total_requests"] == 0 + + async def test_byte_range_does_not_corrupt_cache(self) -> None: + """Test that fetching a byte range does not store partial data under the full key. + + Reproduces https://github.com/zarr-developers/zarr-python/issues/3690: + when a byte-range read populates the cache, subsequent reads of different + ranges (or the full key) return wrong data. + """ + source_store = MemoryStore() + cache_store = MemoryStore() + cached_store = CacheStore(store=source_store, cache_store=cache_store) + + full_data = b"bar baz" + await source_store.set("foo", CPUBuffer.from_bytes(full_data)) + + proto = default_buffer_prototype() + + # First read: byte range [0, 3) -> b"bar" + bar = await cached_store.get("foo", proto, byte_range=RangeByteRequest(0, 3)) + assert bar is not None + assert bar.to_bytes() == b"bar" + + # Second read: different byte range [4, 7) -> b"baz" + baz = await cached_store.get("foo", proto, byte_range=RangeByteRequest(4, 7)) + assert baz is not None + assert baz.to_bytes() == b"baz" + + # Third read: full key -> full data + full = await cached_store.get("foo", proto) + assert full is not None + assert full.to_bytes() == full_data + + async def test_full_read_then_byte_range(self) -> None: + """Test that a cached full read correctly serves subsequent byte-range requests.""" + source_store = MemoryStore() + cache_store = MemoryStore() + cached_store = CacheStore(store=source_store, cache_store=cache_store) + + full_data = b"hello world" + await source_store.set("key", CPUBuffer.from_bytes(full_data)) + + proto = default_buffer_prototype() + + # Full read populates cache + full = await cached_store.get("key", proto) + assert full is not None + assert full.to_bytes() == full_data + + # Byte-range reads should return the correct slices + part = await cached_store.get("key", proto, byte_range=RangeByteRequest(0, 5)) + assert part is not None + assert part.to_bytes() == b"hello" + + part2 = await cached_store.get("key", proto, byte_range=RangeByteRequest(6, 11)) + assert part2 is not None + assert part2.to_bytes() == b"world" + + suffix = await cached_store.get("key", proto, byte_range=SuffixByteRequest(5)) + assert suffix is not None + assert suffix.to_bytes() == b"world" + + async def test_byte_range_set_then_read(self) -> None: + """Test that data written via set() can be read back with byte ranges.""" + source_store = MemoryStore() + cache_store = MemoryStore() + cached_store = CacheStore(store=source_store, cache_store=cache_store) + + full_data = b"abcdefghij" + await cached_store.set("key", CPUBuffer.from_bytes(full_data)) + + proto = default_buffer_prototype() + + # Byte-range reads from the cached data + mid = await cached_store.get("key", proto, byte_range=RangeByteRequest(3, 7)) + assert mid is not None + assert mid.to_bytes() == b"defg" + + # Full read should still work + full = await cached_store.get("key", proto) + assert full is not None + assert full.to_bytes() == full_data + + async def test_set_invalidates_cached_byte_ranges(self) -> None: + """Test that set() invalidates previously cached byte-range entries.""" + source_store = MemoryStore() + cache_store = MemoryStore() + cached_store = CacheStore(store=source_store, cache_store=cache_store) + + proto = default_buffer_prototype() + + # Populate source and cache some byte ranges + await source_store.set("key", CPUBuffer.from_bytes(b"old data!!")) + r1 = await cached_store.get("key", proto, byte_range=RangeByteRequest(0, 3)) + assert r1 is not None + assert r1.to_bytes() == b"old" + + # Byte-range entry should be in range_cache + assert ("key", RangeByteRequest(0, 3)) in cached_store._state.cache_order + + # Overwrite via set() — range entries must be invalidated + await cached_store.set("key", CPUBuffer.from_bytes(b"NEW DATA!!")) + + # The old range entry should be gone from tracking and range_cache + assert ("key", RangeByteRequest(0, 3)) not in cached_store._state.cache_order + assert "key" not in cached_store._state.range_cache + + # A fresh byte-range read should return the new data + r2 = await cached_store.get("key", proto, byte_range=RangeByteRequest(0, 3)) + assert r2 is not None + assert r2.to_bytes() == b"NEW" + + async def test_delete_invalidates_cached_byte_ranges(self) -> None: + """Test that delete() removes previously cached byte-range entries.""" + source_store = MemoryStore() + cache_store = MemoryStore() + cached_store = CacheStore(store=source_store, cache_store=cache_store) + + proto = default_buffer_prototype() + + # Populate and cache a byte range + await source_store.set("key", CPUBuffer.from_bytes(b"hello world")) + r = await cached_store.get("key", proto, byte_range=RangeByteRequest(0, 5)) + assert r is not None + assert r.to_bytes() == b"hello" + + assert ("key", RangeByteRequest(0, 5)) in cached_store._state.cache_order + + # Delete the key — range entries must be cleaned up + await cached_store.delete("key") + + assert ("key", RangeByteRequest(0, 5)) not in cached_store._state.cache_order + assert "key" not in cached_store._state.range_cache + + # Key is gone from source + result = await cached_store.get("key", proto) + assert result is None + + +def test_cache_store_opts_out_of_sync_io() -> None: + """`CacheStore` must not advertise sync IO capability. + + Its caching logic lives only in the async `get`/`set`/`delete` overrides, + while the inherited `WrapperStore` sync methods delegate straight to the + source store. If the fused codec pipeline took the sync fast path, writes + and deletes would bypass the cache and later async reads would serve stale + entries. The opt-out forces sync-capable consumers onto the async path, + which keeps the cache coherent. + """ + from zarr_storage.legacy import MemoryStore + from zarr_storage.legacy._abc import _store_supports_sync_io + + cached = CacheStore(MemoryStore(), cache_store=MemoryStore()) + assert _store_supports_sync_io(cached) is False + + +async def test_cache_coherent_after_fused_pipeline_write() -> None: + """Writing through the fused pipeline must not leave stale cache entries.""" + import numpy as np + import zarr + from zarr.core.config import config as zarr_config + + from zarr_storage.legacy import MemoryStore + + source = MemoryStore() + cached = CacheStore(source, cache_store=MemoryStore()) + with zarr_config.set({"codec_pipeline.path": "zarr.core.codec_pipeline.FusedCodecPipeline"}): + arr = zarr.create_array(cached, shape=(8,), chunks=(8,), dtype="int32", fill_value=0) + arr[:] = np.arange(8, dtype="int32") + np.testing.assert_array_equal(arr[:], np.arange(8)) + # Overwrite, then read back through the same cached handle: the read + # must observe the overwrite, not a cached copy of the first write. + arr[:] = np.arange(100, 108, dtype="int32") + np.testing.assert_array_equal(arr[:], np.arange(100, 108)) diff --git a/packages/zarr-storage/tests/test_store/test_core.py b/packages/zarr-storage/tests/test_store/test_core.py new file mode 100644 index 0000000000..fe07636b10 --- /dev/null +++ b/packages/zarr-storage/tests/test_store/test_core.py @@ -0,0 +1,417 @@ +import tempfile +from collections.abc import Awaitable, Callable, Generator +from pathlib import Path +from typing import Any, Literal + +import pytest +import zarr +from packaging.version import parse as parse_version +from zarr import Group +from zarr.core.buffer import cpu +from zarr.core.common import ZARR_JSON, AccessModeLiteral, ZarrFormat + +from zarr_storage.legacy import FsspecStore, LocalStore, MemoryStore, StoreLike, StorePath, ZipStore +from zarr_storage.legacy._abc import Store +from zarr_storage.legacy._common import ( + _contains_node_v3, + contains_array, + contains_group, + make_store, + make_store_path, +) +from zarr_storage.legacy._utils import ( + _join_paths, + _normalize_path_keys, + _normalize_paths, + _relativize_path, + normalize_path, +) + +# contains_array and contains_group share this signature. +_ContainsFunc = Callable[[StorePath, ZarrFormat], Awaitable[bool]] + + +@pytest.fixture( + params=["none", "temp_dir_str", "temp_dir_path", "store_path", "memory_store", "dict"] +) +def store_like( + request: pytest.FixtureRequest, +) -> Generator[str | Path | StorePath | MemoryStore | dict[Any, Any] | None, None, None]: + if request.param == "none": + yield None + elif request.param == "temp_dir_str": + with tempfile.TemporaryDirectory() as temp_dir: + yield temp_dir + elif request.param == "temp_dir_path": + with tempfile.TemporaryDirectory() as temp_dir: + yield Path(temp_dir) + elif request.param == "store_path": + yield StorePath(store=MemoryStore(store_dict={}), path="/") + elif request.param == "memory_store": + yield MemoryStore(store_dict={}) + elif request.param == "dict": + yield {} + + +@pytest.mark.parametrize("path", ["foo", "foo/bar"]) +@pytest.mark.parametrize("write_group", [True, False]) +@pytest.mark.parametrize("zarr_format", [2, 3]) +async def test_contains_group( + local_store: LocalStore, path: str, write_group: bool, zarr_format: ZarrFormat +) -> None: + """ + Test that the contains_group method correctly reports the existence of a group. + """ + root = Group.from_store(store=local_store, zarr_format=zarr_format) + if write_group: + root.create_group(path) + store_path = StorePath(local_store, path=path) + assert await contains_group(store_path, zarr_format=zarr_format) == write_group + + +@pytest.mark.parametrize("path", ["foo", "foo/bar"]) +@pytest.mark.parametrize("write_array", [True, False]) +@pytest.mark.parametrize("zarr_format", [2, 3]) +async def test_contains_array( + local_store: LocalStore, path: str, write_array: bool, zarr_format: ZarrFormat +) -> None: + """ + Test that the contains array method correctly reports the existence of an array. + """ + root = Group.from_store(store=local_store, zarr_format=zarr_format) + if write_array: + root.create_array(path, shape=(100,), chunks=(10,), dtype="i4") + store_path = StorePath(local_store, path=path) + assert await contains_array(store_path, zarr_format=zarr_format) == write_array + + +@pytest.mark.parametrize("func", [contains_array, contains_group]) +async def test_contains_invalid_format_raises(local_store: LocalStore, func: _ContainsFunc) -> None: + """ + Test contains_group and contains_array raise errors for invalid zarr_formats + """ + store_path = StorePath(local_store) + with pytest.raises(ValueError, match="Invalid zarr_format provided. Got 3.0, expected 2 or 3"): + assert await func(store_path, "3.0") # type: ignore[arg-type] + + +async def _write_zarr_json(store_path: StorePath, data: bytes) -> None: + """Write raw bytes to the v3 metadata key under `store_path`.""" + await (store_path / ZARR_JSON).set(cpu.Buffer.from_bytes(data)) + + +@pytest.mark.parametrize("func", [contains_array, contains_group]) +async def test_contains_malformed_json_returns_false( + local_store: LocalStore, func: _ContainsFunc +) -> None: + """A v3 metadata document that is not valid JSON reads as 'not present'.""" + store_path = StorePath(local_store, path="foo") + await _write_zarr_json(store_path, b"{not valid json") + assert await func(store_path, 3) is False + + +@pytest.mark.parametrize("func", [contains_array, contains_group]) +async def test_contains_non_object_json_returns_false( + local_store: LocalStore, func: _ContainsFunc +) -> None: + """A v3 metadata document that is valid JSON but not an object reads as 'not present'.""" + store_path = StorePath(local_store, path="foo") + await _write_zarr_json(store_path, b"[1, 2, 3]") + assert await func(store_path, 3) is False + + +@pytest.mark.parametrize("func", [contains_array, contains_group]) +async def test_contains_missing_node_type_returns_false( + local_store: LocalStore, func: _ContainsFunc +) -> None: + """A v3 metadata document with no 'node_type' key reads as 'not present'.""" + store_path = StorePath(local_store, path="foo") + await _write_zarr_json(store_path, b'{"zarr_format": 3}') + assert await func(store_path, 3) is False + + +@pytest.mark.parametrize("func", [contains_array, contains_group]) +async def test_contains_non_utf8_bytes_returns_false( + local_store: LocalStore, func: _ContainsFunc +) -> None: + """A v3 metadata document that is not valid UTF-8 reads as 'not present' (not an error).""" + store_path = StorePath(local_store, path="foo") + await _write_zarr_json(store_path, b"\x80\x81\x82\x83") + assert await func(store_path, 3) is False + + +async def test_contains_node_v3_malformed_json_returns_nothing(local_store: LocalStore) -> None: + """`_contains_node_v3` returns 'nothing' when the document is not valid JSON.""" + store_path = StorePath(local_store, path="foo") + await _write_zarr_json(store_path, b"{not valid json") + assert await _contains_node_v3(store_path) == "nothing" + + +async def test_contains_node_v3_non_object_json_returns_nothing(local_store: LocalStore) -> None: + """`_contains_node_v3` returns 'nothing' when the document is not a JSON object.""" + store_path = StorePath(local_store, path="foo") + await _write_zarr_json(store_path, b"[1, 2, 3]") + assert await _contains_node_v3(store_path) == "nothing" + + +async def test_contains_node_v3_missing_node_type_returns_nothing(local_store: LocalStore) -> None: + """`_contains_node_v3` returns 'nothing' when the document lacks a 'node_type' key.""" + store_path = StorePath(local_store, path="foo") + await _write_zarr_json(store_path, b'{"zarr_format": 3}') + assert await _contains_node_v3(store_path) == "nothing" + + +async def test_contains_node_v3_non_utf8_bytes_returns_nothing(local_store: LocalStore) -> None: + """`_contains_node_v3` returns 'nothing' when the document is not valid UTF-8.""" + store_path = StorePath(local_store, path="foo") + await _write_zarr_json(store_path, b"\x80\x81\x82\x83") + assert await _contains_node_v3(store_path) == "nothing" + + +@pytest.mark.parametrize("path", [None, "", "bar"]) +async def test_make_store_path_none(path: str) -> None: + """ + Test that creating a store_path with None creates a memorystore + """ + store_path = await make_store_path(None, path=path) + assert isinstance(store_path.store, MemoryStore) + assert store_path.path == normalize_path(path) + + +@pytest.mark.parametrize("path", [None, "", "bar"]) +@pytest.mark.parametrize("store_type", [str, Path]) +@pytest.mark.parametrize("mode", ["r", "w"]) +async def test_make_store_path_local( + tmp_path: Path, + store_type: type[str] | type[Path] | type[LocalStore], + path: str, + mode: AccessModeLiteral, +) -> None: + """ + Test the various ways of invoking make_store_path that create a LocalStore + """ + store_like = store_type(str(tmp_path)) + store_path = await make_store_path(store_like, path=path, mode=mode) + assert isinstance(store_path.store, LocalStore) + assert Path(store_path.store.root) == Path(tmp_path) + assert store_path.path == normalize_path(path) + assert store_path.read_only == (mode == "r") + + +@pytest.mark.parametrize("path", [None, "", "bar"]) +@pytest.mark.parametrize("mode", ["r", "w"]) +async def test_make_store_path_store_path( + tmp_path: Path, path: str, mode: AccessModeLiteral +) -> None: + """ + Test invoking make_store_path when the input is another store_path. In particular we want to ensure + that a new path is handled correctly. + """ + ro = mode == "r" + store_like = await StorePath.open( + LocalStore(str(tmp_path), read_only=ro), path="root", mode=mode + ) + store_path = await make_store_path(store_like, path=path, mode=mode) + assert isinstance(store_path.store, LocalStore) + assert Path(store_path.store.root) == tmp_path + path_normalized = normalize_path(path) + assert store_path.path == (store_like / path_normalized).path + assert store_path.read_only == ro + + +@pytest.mark.parametrize("modes", [(True, "w"), (False, "x")]) +async def test_store_path_invalid_mode_raises( + tmp_path: Path, modes: tuple[bool, Literal["w", "x"]] +) -> None: + """ + Test that ValueErrors are raise for invalid mode. + """ + with pytest.raises(ValueError): + await StorePath.open(LocalStore(str(tmp_path), read_only=modes[0]), path="", mode=modes[1]) # type: ignore[arg-type] + + +async def test_make_store_path_invalid() -> None: + """ + Test that invalid types raise TypeError + """ + with pytest.raises(TypeError, match="Unsupported type for store_like: 'int'"): + await make_store_path(1) + + +async def test_make_store_path_fsspec() -> None: + pytest.importorskip("fsspec") + pytest.importorskip("requests") + pytest.importorskip("aiohttp") + store_path = await make_store_path("http://foo.com/bar") + assert isinstance(store_path.store, FsspecStore) + + +async def test_make_store_path_storage_options_raises(store_like: StoreLike) -> None: + with pytest.raises(TypeError, match="storage_options"): + await make_store_path(store_like, storage_options={"foo": "bar"}) + + +# universal-pathlib 0.2.x emits this from its own subclass registry when a local UPath is built. +@pytest.mark.filterwarnings( + "ignore:Detected a customized `__new__` method in subclass:DeprecationWarning" +) +@pytest.mark.parametrize( + ("url", "expected"), + [ + ("memory://bucket/foo.zarr", FsspecStore), + ("s3://bucket/foo.zarr", FsspecStore), + ("file://{tmp}/foo.zarr", LocalStore), + ("{tmp}/foo.zarr", LocalStore), + ], +) +async def test_make_store_upath(url: str, expected: type[Store], tmp_path: Path) -> None: + """ + A remote UPath becomes an FsspecStore, and a local one becomes a LocalStore, so that + UPath("/data") and Path("/data") agree. See https://github.com/zarr-developers/zarr-python/issues/4244. + """ + upath = pytest.importorskip("upath") + fsspec = pytest.importorskip("fsspec") + if url.startswith("s3://"): + pytest.importorskip("s3fs") + if url.startswith("memory://") and parse_version(fsspec.__version__) < parse_version( + "2024.12.0" + ): + # MemoryFileSystem is synchronous, so it can only be used once fsspec is new enough to + # supply AsyncFileSystemWrapper. + pytest.skip("No AsyncFileSystemWrapper") + store = await make_store(upath.UPath(url.format(tmp=tmp_path))) + assert isinstance(store, expected) + if isinstance(store, LocalStore): + # The local branch rebuilds the root from the UPath, so a mangled path would still + # produce a LocalStore. Pin the root down too, since "file://{tmp}" has no leading + # slash on Windows. + assert store.root == tmp_path / "foo.zarr" + + +async def test_make_store_upath_storage_options_raises() -> None: + """A UPath carries its own storage options, so a separate mapping is ambiguous.""" + upath = pytest.importorskip("upath") + with pytest.raises(TypeError, match="storage_options"): + await make_store(upath.UPath("memory://bucket/foo.zarr"), storage_options={"foo": "bar"}) + + +async def test_unsupported() -> None: + with pytest.raises(TypeError, match="Unsupported type for store_like: 'int'"): + await make_store_path(1) + + +@pytest.mark.parametrize( + "path", + [ + "/foo/bar", + "//foo/bar", + "foo///bar", + "foo/bar///", + Path("foo/bar"), + b"foo/bar", + ], +) +def test_normalize_path_valid(path: str | bytes | Path) -> None: + assert normalize_path(path) == "foo/bar" + + +def test_normalize_path_upath() -> None: + upath = pytest.importorskip("upath") + assert normalize_path(upath.UPath("foo/bar", protocol="memory")) == "memory:/foo/bar" + + +def test_normalize_path_none() -> None: + assert normalize_path(None) == "" + + +@pytest.mark.parametrize("path", [".", ".."]) +def test_normalize_path_invalid(path: str) -> None: + with pytest.raises(ValueError, match="is invalid because its string representation contains"): + normalize_path(path) + + +@pytest.mark.parametrize("paths", [("", "foo"), ("foo", "bar")]) +def test_join_paths(paths: tuple[str, str]) -> None: + """ + Test that _join_paths joins paths in a way that is robust to an empty string + """ + observed = _join_paths(paths) + if paths[0] == "": + assert observed == paths[1] + else: + assert observed == "/".join(paths) + + +class TestNormalizePaths: + @staticmethod + def test_valid() -> None: + """ + Test that path normalization works as expected + """ + paths = ["a", "b", "c", "d", "", "//a///b//"] + assert _normalize_paths(paths) == tuple(normalize_path(p) for p in paths) + + @staticmethod + @pytest.mark.parametrize("paths", [("", "/"), ("///a", "a")]) + def test_invalid(paths: tuple[str, str]) -> None: + """ + Test that name collisions after normalization raise a ``ValueError`` + """ + msg = ( + f"After normalization, the value '{paths[1]}' collides with '{paths[0]}'. " + f"Both '{paths[1]}' and '{paths[0]}' normalize to the same value: '{normalize_path(paths[0])}'. " + f"You should use either '{paths[1]}' or '{paths[0]}', but not both." + ) + with pytest.raises(ValueError, match=msg): + _normalize_paths(paths) + + +def test_normalize_path_keys() -> None: + """ + Test that ``_normalize_path_keys`` just applies the normalize_path function to each key of its + input + """ + data = {"a": 10, "//b": 10} + assert _normalize_path_keys(data) == {normalize_path(k): v for k, v in data.items()} + + +@pytest.mark.parametrize( + ("path", "prefix", "expected"), + [ + ("a", "", "a"), + ("a/b/c", "a/b", "c"), + ("a/b/c", "a", "b/c"), + ], +) +def test_relativize_path_valid(path: str, prefix: str, expected: str) -> None: + """ + Test the normal behavior of the _relativize_path function. Prefixes should be removed from the + path argument. + """ + assert _relativize_path(path=path, prefix=prefix) == expected + + +def test_relativize_path_invalid() -> None: + path = "a/b/c" + prefix = "b" + msg = f"The first component of {path} does not start with {prefix}." + with pytest.raises(ValueError, match=msg): + _relativize_path(path="a/b/c", prefix="b") + + +def test_different_open_mode(tmp_path: Path) -> None: + # Test with a store that implements .with_read_only() + store = MemoryStore() + zarr.create((100,), store=store, zarr_format=2, path="a") + arr = zarr.open_array(store=store, path="a", zarr_format=2, mode="r") + assert arr.store.read_only + + # Test with a store that doesn't implement .with_read_only() + zarr_path = tmp_path / "foo.zarr" + zip_store = ZipStore(zarr_path, mode="w") + zarr.create((100,), store=zip_store, zarr_format=2, path="a") + with pytest.raises( + ValueError, + match="Store is not read-only but mode is 'r'. Unable to create a read-only copy of the store. Please use a read-only store or a storage class that implements .with_read_only().", + ): + zarr.open_array(store=zip_store, path="a", zarr_format=2, mode="r") diff --git a/packages/zarr-storage/tests/test_store/test_fsspec.py b/packages/zarr-storage/tests/test_store/test_fsspec.py new file mode 100644 index 0000000000..6252c3941b --- /dev/null +++ b/packages/zarr-storage/tests/test_store/test_fsspec.py @@ -0,0 +1,702 @@ +from __future__ import annotations + +import json +import re +import warnings +from typing import TYPE_CHECKING, Any + +import numpy as np +import pytest +import zarr.api.asynchronous +from packaging.version import parse as parse_version +from zarr import Array +from zarr.core.buffer import Buffer, cpu, default_buffer_prototype +from zarr.core.sync import _collect_aiterator, sync +from zarr.errors import ZarrUserWarning + +from zarr_storage.legacy import FsspecStore +from zarr_storage.legacy._abc import OffsetByteRequest +from zarr_storage.legacy._common import make_store +from zarr_storage.legacy._fsspec import _make_async +from zarr_storage.testing.store import StoreTests + +if TYPE_CHECKING: + import pathlib + from collections.abc import Generator + from pathlib import Path + + import botocore.client + import s3fs + from zarr.core.common import JSON + + +# Warning filter due to https://github.com/boto/boto3/issues/3889 +pytestmark = [ + pytest.mark.filterwarnings( + re.escape("ignore:datetime.datetime.utcnow() is deprecated:DeprecationWarning") + ), + # FsspecStore.from_url() and from_mapper() now close the aiohttp session on store.close(). + # This filter covers stores that are GC'd without an explicit close() call, and any + # residual sessions from aiobotocore's ClientCreatorContext (a separate upstream issue). + pytest.mark.filterwarnings("ignore:Unclosed client session:ResourceWarning"), + pytest.mark.filterwarnings( + "ignore:coroutine 'ClientCreatorContext.__aexit__' was never awaited:RuntimeWarning" + ), + # s3fs finalizers can fail when sessions are garbage collected without being entered + pytest.mark.filterwarnings( + "ignore:Exception ignored in.*finalize object.*:pytest.PytestUnraisableExceptionWarning" + ), +] + +fsspec = pytest.importorskip("fsspec") +s3fs = pytest.importorskip("s3fs") +requests = pytest.importorskip("requests") +# Skip this module entirely when moto is absent; the server itself comes from the shared +# `moto_server` fixture in tests/conftest.py. +pytest.importorskip("moto") +botocore = pytest.importorskip("botocore") + +# ### amended from s3fs ### # +test_bucket_name = "test" +secure_bucket_name = "test-secure" + + +@pytest.fixture +def endpoint_url(moto_server: str) -> str: + """Endpoint of the shared session-scoped moto server (see tests/conftest.py). + + A fixture rather than a module-level constant because the server binds an ephemeral + port, so the endpoint is only known once the server is running.""" + return moto_server + + +def get_boto3_client(endpoint_url: str) -> botocore.client.BaseClient: + # NB: we use the sync botocore client for setup + session = botocore.session.Session() + + # Prevent IllegalLocationConstraintException by explicitly setting region to + # "us-east-1", which does not require configuring LocationConstraint during + # bucket creation. (It is, in fact, forbidden for that region.) Necessary + # in the face of "ambient" AWS configuration in a development environment + # where the default region might be configured differently. + return session.create_client("s3", endpoint_url=endpoint_url, region_name="us-east-1") + + +@pytest.fixture(autouse=True) +def s3(endpoint_url: str) -> Generator[s3fs.S3FileSystem, None, None]: + """ + Quoting Martin Durant: + pytest-asyncio creates a new event loop for each async test. + When an async-mode s3fs instance is made from async, it will be assigned to the loop from + which it is made. That means that if you use s3fs again from a subsequent test, + you will have the same identical instance, but be running on a different loop - which fails. + + For the rest: it's very convenient to clean up the state of the store between tests, + make sure we start off blank each time. + + https://github.com/zarr-developers/zarr-python/pull/1785#discussion_r1634856207 + """ + client = get_boto3_client(endpoint_url) + client.create_bucket(Bucket=test_bucket_name, ACL="public-read") + s3fs.S3FileSystem.clear_instance_cache() + s3 = s3fs.S3FileSystem( + anon=False, + client_kwargs={"endpoint_url": endpoint_url}, + # Prevent "AssertionError: Session was never entered" from aiobotocore + # at end of test execution. Using clear_instance_cache is insufficient, + # although still necessary. + skip_instance_cache=True, + ) + session = sync(s3.set_session()) + s3.invalidate_cache() + yield s3 + requests.post(f"{endpoint_url}/moto-api/reset") + client.close() + sync(session.close()) + + +# ### end from s3fs ### # + + +async def test_basic(endpoint_url: str) -> None: + store = FsspecStore.from_url( + f"s3://{test_bucket_name}/foo/spam/", + storage_options={"endpoint_url": endpoint_url, "anon": False}, + ) + assert store.fs.asynchronous + assert store.path == f"{test_bucket_name}/foo/spam" + assert await _collect_aiterator(store.list()) == () + assert not await store.exists("foo") + data = b"hello" + await store.set("foo", cpu.Buffer.from_bytes(data)) + assert await store.exists("foo") + buf = await store.get("foo", prototype=default_buffer_prototype()) + assert buf is not None + assert buf.to_bytes() == data + out = await store.get_partial_values( + prototype=default_buffer_prototype(), key_ranges=[("foo", OffsetByteRequest(1))] + ) + assert out[0] is not None + assert out[0].to_bytes() == data[1:] + + +class TestFsspecStoreS3(StoreTests[FsspecStore, cpu.Buffer]): + store_cls = FsspecStore + buffer_cls = cpu.Buffer + + @pytest.fixture + def store_kwargs(self, endpoint_url: str) -> dict[str, str | bool]: + try: + from fsspec import url_to_fs + except ImportError: + # before fsspec==2024.3.1 + from fsspec.core import url_to_fs + fs, path = url_to_fs( + f"s3://{test_bucket_name}", endpoint_url=endpoint_url, anon=False, asynchronous=True + ) + return {"fs": fs, "path": path} + + @pytest.fixture + async def store(self, store_kwargs: dict[str, Any]) -> FsspecStore: + return self.store_cls(**store_kwargs) + + async def get(self, store: FsspecStore, key: str) -> Buffer: + # make a new, synchronous instance of the filesystem because this test is run in sync code + new_fs = fsspec.filesystem( + "s3", endpoint_url=store.fs.endpoint_url, anon=store.fs.anon, asynchronous=False + ) + return self.buffer_cls.from_bytes(new_fs.cat(f"{store.path}/{key}")) + + async def set(self, store: FsspecStore, key: str, value: Buffer) -> None: + # make a new, synchronous instance of the filesystem because this test is run in sync code + new_fs = fsspec.filesystem( + "s3", endpoint_url=store.fs.endpoint_url, anon=store.fs.anon, asynchronous=False + ) + new_fs.write_bytes(f"{store.path}/{key}", value.to_bytes()) + + def test_store_repr(self, store: FsspecStore) -> None: + assert str(store) == "" + + def test_store_supports_writes(self, store: FsspecStore) -> None: + assert store.supports_writes + + def test_store_supports_listing(self, store: FsspecStore) -> None: + assert store.supports_listing + + async def test_fsspec_store_from_uri(self, store: FsspecStore, endpoint_url: str) -> None: + storage_options = { + "endpoint_url": endpoint_url, + "anon": False, + } + + meta: dict[str, JSON] = { + "attributes": {"key": "value"}, + "zarr_format": 3, + "node_type": "group", + } + + await store.set( + "zarr.json", + self.buffer_cls.from_bytes(json.dumps(meta).encode()), + ) + group = await zarr.api.asynchronous.open_group( + store=f"s3://{test_bucket_name}", storage_options=storage_options + ) + assert dict(group.attrs) == {"key": "value"} + + meta = { + "attributes": {"key": "value-2"}, + "zarr_format": 3, + "node_type": "group", + } + await store.set( + "directory-2/zarr.json", + self.buffer_cls.from_bytes(json.dumps(meta).encode()), + ) + group = await zarr.api.asynchronous.open_group( + store=f"s3://{test_bucket_name}/directory-2", storage_options=storage_options + ) + assert dict(group.attrs) == {"key": "value-2"} + + meta = { + "attributes": {"key": "value-3"}, + "zarr_format": 3, + "node_type": "group", + } + await store.set( + "directory-3/zarr.json", + self.buffer_cls.from_bytes(json.dumps(meta).encode()), + ) + group = await zarr.api.asynchronous.open_group( + store=f"s3://{test_bucket_name}", path="directory-3", storage_options=storage_options + ) + assert dict(group.attrs) == {"key": "value-3"} + + @pytest.mark.skipif( + parse_version(fsspec.__version__) < parse_version("2024.03.01"), + reason="Prior bug in from_upath", + ) + def test_from_upath(self, endpoint_url: str) -> None: + upath = pytest.importorskip("upath") + path = upath.UPath( + f"s3://{test_bucket_name}/foo/bar/", + endpoint_url=endpoint_url, + anon=False, + asynchronous=True, + ) + result = FsspecStore.from_upath(path) + assert result.fs.endpoint_url == endpoint_url + assert result.fs.asynchronous + assert result.path == f"{test_bucket_name}/foo/bar" + + @pytest.mark.skipif( + parse_version(fsspec.__version__) < parse_version("2024.03.01"), + reason="Prior bug in from_upath", + ) + def test_from_upath_sync_filesystem(self, endpoint_url: str) -> None: + """ + A UPath built without ``asynchronous=True`` -- the common case -- yields an async-mode + filesystem that keeps the original storage options. + """ + upath = pytest.importorskip("upath") + path = upath.UPath( + f"s3://{test_bucket_name}/foo/bar/", + endpoint_url=endpoint_url, + anon=False, + ) + assert not path.fs.asynchronous + with warnings.catch_warnings(): + warnings.simplefilter("error", ZarrUserWarning) + result = FsspecStore.from_upath(path) + assert result.fs.asynchronous + assert result.fs.endpoint_url == endpoint_url + assert result.path == f"{test_bucket_name}/foo/bar" + + async def test_open_group_from_upath(self, endpoint_url: str) -> None: + """ + Passing a remote UPath to the top-level API works. + + Regression test for https://github.com/zarr-developers/zarr-python/issues/4244. + """ + upath = pytest.importorskip("upath") + path = upath.UPath( + f"s3://{test_bucket_name}/upath-group", + endpoint_url=endpoint_url, + anon=False, + ) + group = await zarr.api.asynchronous.open_group(path, mode="w", attributes={"key": "value"}) + assert isinstance(group.store_path.store, FsspecStore) + + reopened = await zarr.api.asynchronous.open_group(path, mode="r") + assert dict(reopened.attrs) == {"key": "value"} + + def test_init_warns_if_fs_asynchronous_is_false(self, endpoint_url: str) -> None: + try: + from fsspec import url_to_fs + except ImportError: + # before fsspec==2024.3.1 + from fsspec.core import url_to_fs + fs, path = url_to_fs( + f"s3://{test_bucket_name}", endpoint_url=endpoint_url, anon=False, asynchronous=False + ) + store_kwargs = {"fs": fs, "path": path} + with pytest.warns(ZarrUserWarning, match=r".* was not created with `asynchronous=True`.*"): + self.store_cls(**store_kwargs) + + async def test_empty_nonexistent_path(self, store_kwargs: dict[str, Any]) -> None: + # regression test for https://github.com/zarr-developers/zarr-python/pull/2343 + store_kwargs["path"] += "/abc" + store = await self.store_cls.open(**store_kwargs) + assert await store.is_empty("") + + async def test_delete_dir_unsupported_deletes(self, store: FsspecStore) -> None: + store.supports_deletes = False + with pytest.raises( + NotImplementedError, + match="This method is only available for stores that support deletes.", + ): + await store.delete_dir("test_prefix") + + # ── Filesystem lifecycle ────────────────────────────────────────────────── + + async def test_close_marks_store_closed(self, endpoint_url: str) -> None: + """close() must succeed and mark the store not-open.""" + store = FsspecStore.from_url( + f"s3://{test_bucket_name}/lifecycle/", + storage_options={"endpoint_url": endpoint_url, "anon": False}, + ) + await store.set("probe", cpu.Buffer.from_bytes(b"x")) + + store.close() + + assert not store._is_open + + +def array_roundtrip(store: FsspecStore) -> None: + """ + Round trip an array using a Zarr store + + Args: + store: FsspecStore + """ + data = np.ones((3, 3)) + arr = zarr.create_array(store=store, overwrite=True, data=data) + assert isinstance(arr, Array) + # Read set values + arr2 = zarr.open_array(store=store) + assert isinstance(arr2, Array) + np.testing.assert_array_equal(arr[:], data) + + +@pytest.mark.parametrize( + ("root", "key", "expected"), + [ + # `"/"` as root collapses so that bare-key backends (notably + # ReferenceFileSystem) get the right key. Regression test for + # https://github.com/zarr-developers/zarr-python/issues/3922 . + ("/", "zarr.json", "zarr.json"), + ("", "zarr.json", "zarr.json"), + # Trailing slashes on the root are stripped before joining. + ("foo/", "zarr.json", "foo/zarr.json"), + ("foo", "zarr.json", "foo/zarr.json"), + # Leading slashes on the root are preserved -- absolute filesystem + # paths must stay absolute. Regression test for the titiler-xarray + # breakage that #3924 introduced when `normalize_path` was applied to + # `FsspecStore.path`. + ("/home/runner/data.zarr", "zarr.json", "/home/runner/data.zarr/zarr.json"), + ("/home/runner/data.zarr/", "zarr.json", "/home/runner/data.zarr/zarr.json"), + # Multi-segment keys. + ("/home/foo", "a/b/zarr.json", "/home/foo/a/b/zarr.json"), + ("", "a/b/zarr.json", "a/b/zarr.json"), + # Trailing slash on the result is stripped (relevant when key is ""). + ("/home/foo", "", "/home/foo"), + ], +) +def test_dereference_path(root: str, key: str, expected: str) -> None: + """Verify the contract `_dereference_path` provides for `FsspecStore`. + + `FsspecStore.path` is stored verbatim; the join with a key must collapse a + sentinel `"/"` root, strip trailing slashes, and preserve leading + slashes on absolute paths. + """ + from zarr_storage.legacy._utils import _dereference_path + + assert _dereference_path(root, key) == expected + + +async def test_fsspec_store_open_group_via_reference_filesystem() -> None: + """End-to-end regression test for + https://github.com/zarr-developers/zarr-python/issues/3922 . + + ``ReferenceFileSystem`` keys its refs by bare strings like ``"zarr.json"``. + The bug was that ``FsspecStore(fs=ref_fs, path="/")`` produced + ``"//zarr.json"`` at the join site and failed to find the entry, raising + ``GroupNotFoundError``. This test pins ``path="/"`` explicitly to keep + coverage even if the default value changes later. + """ + import json + + from fsspec.implementations.reference import ReferenceFileSystem + + group_json = json.dumps({"zarr_format": 3, "node_type": "group", "attributes": {}}) + fs = ReferenceFileSystem( + fo={"version": 1, "refs": {"zarr.json": group_json}}, + asynchronous=True, + ) + store = FsspecStore(fs=fs, path="/", read_only=True) + group = await zarr.api.asynchronous.open_group(store, mode="r") + assert group.metadata.zarr_format == 3 + + +async def test_fsspec_store_read_array_chunk_via_reference_filesystem() -> None: + """End-to-end regression test that exercises the byte-range read path + against ``ReferenceFileSystem``. + + Beyond opening a group (covered by + ``test_fsspec_store_open_group_via_reference_filesystem``), this test + constructs a small zarr v3 array whose chunk lives in the refs dict and + reads it through the store. Path-handling bugs on the byte-range + fetch path (used by kerchunk-style virtualization) would surface here + rather than at metadata-open time. + """ + import json + + import numpy as np + from fsspec.implementations.reference import ReferenceFileSystem + + # Construct a minimal v3 zarr: a single 1-D uint8 array of length 4 with + # one chunk of size 4. The chunk bytes are little-endian uint8s 1..4. + array_meta = json.dumps( + { + "zarr_format": 3, + "node_type": "array", + "shape": [4], + "chunk_grid": {"name": "regular", "configuration": {"chunk_shape": [4]}}, + "data_type": "uint8", + "chunk_key_encoding": {"name": "default", "configuration": {"separator": "/"}}, + "fill_value": 0, + "codecs": [{"name": "bytes", "configuration": {"endian": "little"}}], + "attributes": {}, + } + ) + chunk_bytes = bytes([1, 2, 3, 4]) + + refs: dict[str, str] = { + "zarr.json": array_meta, + # ReferenceFileSystem accepts raw bytes via base64 encoding or + # latin-1-decoded strings; latin-1 round-trips bytes 1:1. + "c/0": chunk_bytes.decode("latin-1"), + } + + fs = ReferenceFileSystem( + fo={"version": 1, "refs": refs}, + asynchronous=True, + ) + store = FsspecStore(fs=fs, path="/", read_only=True) + array = await zarr.api.asynchronous.open_array(store=store, mode="r") + data = await array.getitem(slice(None)) + np.testing.assert_array_equal(data, np.array([1, 2, 3, 4], dtype="uint8")) + + +@pytest.mark.skipif( + parse_version(fsspec.__version__) < parse_version("2024.12.0"), + reason="No AsyncFileSystemWrapper", +) +def test_wrap_sync_filesystem(tmp_path: pathlib.Path) -> None: + """The local fs is not async so we should expect it to be wrapped automatically""" + from fsspec.implementations.asyn_wrapper import AsyncFileSystemWrapper + + store = FsspecStore.from_url(f"file://{tmp_path}", storage_options={"auto_mkdir": True}) + assert isinstance(store.fs, AsyncFileSystemWrapper) + assert store.fs.async_impl + array_roundtrip(store) + + +@pytest.mark.skipif( + parse_version(fsspec.__version__) >= parse_version("2024.12.0"), + reason="No AsyncFileSystemWrapper", +) +def test_wrap_sync_filesystem_raises(tmp_path: pathlib.Path) -> None: + """The local fs is not async so we should expect it to be wrapped automatically""" + with pytest.raises(ImportError, match="The filesystem .*"): + FsspecStore.from_url(f"file://{tmp_path}", storage_options={"auto_mkdir": True}) + + +@pytest.mark.skipif( + parse_version(fsspec.__version__) < parse_version("2024.12.0"), + reason="No AsyncFileSystemWrapper", +) +def test_no_wrap_async_filesystem(endpoint_url: str) -> None: + """An async fs should not be wrapped automatically; fsspec's s3 filesystem is such an fs""" + from fsspec.implementations.asyn_wrapper import AsyncFileSystemWrapper + + store = FsspecStore.from_url( + f"s3://{test_bucket_name}/foo/spam/", + storage_options={"endpoint_url": endpoint_url, "anon": False, "asynchronous": True}, + read_only=False, + ) + assert not isinstance(store.fs, AsyncFileSystemWrapper) + assert store.fs.async_impl + array_roundtrip(store) + + +@pytest.mark.skipif( + parse_version(fsspec.__version__) < parse_version("2024.12.0"), + reason="No AsyncFileSystemWrapper", +) +def test_open_fsmap_file(tmp_path: pathlib.Path) -> None: + min_fsspec_with_async_wrapper = parse_version("2024.12.0") + current_version = parse_version(fsspec.__version__) + + fs = fsspec.filesystem("file", auto_mkdir=True) + mapper = fs.get_mapper(tmp_path) + + if current_version < min_fsspec_with_async_wrapper: + # Expect ImportError for older versions + with pytest.raises( + ImportError, + match=r"The filesystem .* is synchronous, and the required AsyncFileSystemWrapper is not available.*", + ): + array_roundtrip(mapper) + else: + # Newer versions should work + array_roundtrip(mapper) + + +@pytest.mark.skipif( + parse_version(fsspec.__version__) < parse_version("2024.12.0"), + reason="No AsyncFileSystemWrapper", +) +def test_open_fsmap_file_raises(tmp_path: pathlib.Path) -> None: + fsspec = pytest.importorskip("fsspec.implementations.local") + fs = fsspec.LocalFileSystem(auto_mkdir=False) + mapper = fs.get_mapper(tmp_path) + with pytest.raises(FileNotFoundError, match="No such file or directory: .*"): + array_roundtrip(mapper) + + +@pytest.mark.parametrize("asynchronous", [True, False]) +def test_open_fsmap_s3(asynchronous: bool, endpoint_url: str) -> None: + s3_filesystem = s3fs.S3FileSystem( + asynchronous=asynchronous, endpoint_url=endpoint_url, anon=False + ) + mapper = s3_filesystem.get_mapper(f"s3://{test_bucket_name}/map/foo/") + array_roundtrip(mapper) + + +def test_open_s3map_raises(endpoint_url: str) -> None: + with pytest.raises(TypeError, match="Unsupported type for store_like:.*"): + zarr.open(store=0, mode="w", shape=(3, 3)) + s3_filesystem = s3fs.S3FileSystem(asynchronous=True, endpoint_url=endpoint_url, anon=False) + mapper = s3_filesystem.get_mapper(f"s3://{test_bucket_name}/map/foo/") + with pytest.raises( + ValueError, match="'path' was provided but is not used for FSMap store_like objects" + ): + zarr.open(store=mapper, path="bar", mode="w", shape=(3, 3)) + with pytest.raises( + TypeError, + match="'storage_options' is only used when the store is passed as an FSSpec URI string.", + ): + zarr.open(store=mapper, storage_options={"anon": True}, mode="w", shape=(3, 3)) + + +async def test_close_does_not_close_filesystem_session() -> None: + """close() must not touch the filesystem's session. + + fsspec caches and shares filesystem instances across callers, so the + session is not the store's to close. HTTP is used because its aiohttp + session is observably closed for good; s3fs transparently reconnects, which + would hide a regression. No request is issued — set_session() only + constructs the session. + """ + pytest.importorskip("aiohttp") + store = FsspecStore.from_url("http://example.com/a") + session = await store.fs.set_session() + + store.close() + + assert not session.closed + + +async def test_close_does_not_break_a_sibling_store() -> None: + """Closing one store must not close a session another store is using. + + Two stores from different URLs on one host are handed the same cached + filesystem; a store that closed it on close() would take the sibling's + session down too. This is the regression guard for that bug. + """ + pytest.importorskip("aiohttp") + s1 = FsspecStore.from_url("http://example.com/a") + s2 = FsspecStore.from_url("http://example.com/b") + session = await s2.fs.set_session() + + s1.close() + + assert not session.closed + + +@pytest.mark.skipif( + parse_version(fsspec.__version__) < parse_version("2024.12.0"), + reason="No AsyncFileSystemWrapper", +) +def test_from_mapper_wraps_sync_filesystem(tmp_path: pathlib.Path) -> None: + """from_mapper() with a sync fs wraps it in an AsyncFileSystemWrapper.""" + import fsspec as _fsspec + from fsspec.implementations.asyn_wrapper import AsyncFileSystemWrapper + + fs = _fsspec.filesystem("file", auto_mkdir=True) + mapper = fs.get_mapper(str(tmp_path)) + store = FsspecStore.from_mapper(mapper) + assert isinstance(store.fs, AsyncFileSystemWrapper) + + +@pytest.mark.skipif( + parse_version(fsspec.__version__) < parse_version("2024.12.0"), + reason="No AsyncFileSystemWrapper", +) +def test_with_read_only_shares_filesystem(tmp_path: pathlib.Path) -> None: + """with_read_only() returns a store sharing the source's filesystem.""" + source = FsspecStore.from_url(f"file://{tmp_path}", storage_options={"auto_mkdir": False}) + + derived = source.with_read_only(read_only=True) + + assert derived.fs is source.fs + assert derived.read_only + assert not source.read_only + + +def test_make_async_preserves_unserializable_storage_options() -> None: + """A sync instance of an async filesystem whose storage options hold objects that + cannot round-trip through JSON (e.g. an Azure credential) must still convert. + + See https://github.com/zarr-developers/zarr-python/issues/4220 + """ + pytest.importorskip("aiohttp") + credential = object() # stand-in for e.g. azure.identity.DefaultAzureCredential + sync_fs = fsspec.filesystem("http", client_kwargs={"auth": credential}) + assert sync_fs.async_impl + assert not sync_fs.asynchronous + + async_fs = _make_async(sync_fs) + + assert async_fs.asynchronous + assert async_fs.client_kwargs["auth"] is credential + + +@pytest.mark.parametrize("asynchronous", [True, False]) +def test_make_async(asynchronous: bool, endpoint_url: str) -> None: + s3_filesystem = s3fs.S3FileSystem( + asynchronous=asynchronous, endpoint_url=endpoint_url, anon=False + ) + fs = _make_async(s3_filesystem) + assert fs.asynchronous + + +@pytest.mark.skipif( + parse_version(fsspec.__version__) < parse_version("2024.12.0"), + reason="No AsyncFileSystemWrapper", +) +async def test_delete_dir_wrapped_filesystem(tmp_path: Path) -> None: + from fsspec.implementations.asyn_wrapper import AsyncFileSystemWrapper + from fsspec.implementations.local import LocalFileSystem + + wrapped_fs = AsyncFileSystemWrapper(LocalFileSystem(auto_mkdir=True)) + store = FsspecStore(wrapped_fs, read_only=False, path=f"{tmp_path}/test/path") + + assert isinstance(store.fs, AsyncFileSystemWrapper) + assert store.fs.asynchronous + + await store.set("zarr.json", cpu.Buffer.from_bytes(b"root")) + await store.set("foo-bar/zarr.json", cpu.Buffer.from_bytes(b"root")) + await store.set("foo/zarr.json", cpu.Buffer.from_bytes(b"bar")) + await store.set("foo/c/0", cpu.Buffer.from_bytes(b"chunk")) + await store.delete_dir("foo") + assert await store.exists("zarr.json") + assert await store.exists("foo-bar/zarr.json") + assert not await store.exists("foo/zarr.json") + assert not await store.exists("foo/c/0") + + +@pytest.mark.skipif( + parse_version(fsspec.__version__) < parse_version("2024.12.0"), + reason="No AsyncFileSystemWrapper", +) +async def test_with_read_only_auto_mkdir(tmp_path: Path) -> None: + """ + Test that creating a read-only copy of a store backed by the local file system does not error + if auto_mkdir is False. + """ + + store_w = FsspecStore.from_url(f"file://{tmp_path}", storage_options={"auto_mkdir": False}) + _ = store_w.with_read_only() + + +@pytest.mark.skipif( + parse_version(fsspec.__version__) < parse_version("2024.12.0"), + reason="No AsyncFileSystemWrapper", +) +async def test_memory_scheme() -> None: + """Test that the "memory" scheme creates a `MemoryFileSystem`-backed store""" + store = await make_store("memory://test") + assert isinstance(store, FsspecStore) + assert store.fs.protocol == "memory" diff --git a/packages/zarr-storage/tests/test_store/test_fsspec_get_ranges.py b/packages/zarr-storage/tests/test_store/test_fsspec_get_ranges.py new file mode 100644 index 0000000000..babfa50fe4 --- /dev/null +++ b/packages/zarr-storage/tests/test_store/test_fsspec_get_ranges.py @@ -0,0 +1,124 @@ +# tests/test_store/test_fsspec_get_ranges.py +"""Lightweight integration tests for FsspecStore.get_ranges using MemoryFileSystem. + +These don't need moto/s3 — they exercise the new method against an in-process +fsspec MemoryFileSystem wrapped in the async wrapper. +""" + +from __future__ import annotations + +import pytest +from packaging.version import parse as parse_version +from zarr.core.buffer import Buffer, default_buffer_prototype + +from zarr_storage.legacy import FsspecStore +from zarr_storage.legacy._abc import RangeByteRequest +from zarr_storage.legacy._fsspec import _make_async + +fsspec = pytest.importorskip("fsspec") + +# AsyncFileSystemWrapper (needed to wrap a sync MemoryFileSystem) landed in fsspec 2024.12.0. +# Older versions are pinned by the min-deps CI job, so skip the whole file there. +pytestmark = pytest.mark.skipif( + parse_version(fsspec.__version__) < parse_version("2024.12.0"), + reason="No AsyncFileSystemWrapper", +) + + +@pytest.fixture +def memory_store() -> FsspecStore: + """An FsspecStore backed by fsspec MemoryFileSystem (wrapped async).""" + from fsspec.implementations.memory import MemoryFileSystem + + # Each test gets a clean filesystem; MemoryFileSystem is a singleton per target_options, + # so clear state explicitly. + fs: MemoryFileSystem = MemoryFileSystem() + fs.store.clear() + fs.pseudo_dirs.clear() + async_fs = _make_async(fs) + return FsspecStore(fs=async_fs, path="/root") + + +async def _write(store: FsspecStore, key: str, data: bytes) -> None: + buf = default_buffer_prototype().buffer.from_bytes(data) + await store.set(key, buf) + + +async def test_get_ranges_happy_path(memory_store: FsspecStore) -> None: + blob = bytes(i % 256 for i in range(1024)) + await _write(memory_store, "blob", blob) + proto = default_buffer_prototype() + + ranges = [ + RangeByteRequest(0, 10), + RangeByteRequest(100, 110), + RangeByteRequest(500, 520), + ] + groups: list[list[tuple[int, Buffer | None]]] = [ + list(group) async for group in memory_store.get_ranges("blob", ranges, prototype=proto) + ] + + flat: dict[int, bytes] = {} + for group in groups: + for idx, buf in group: + assert buf is not None + flat[idx] = buf.to_bytes() + + assert flat[0] == blob[0:10] + assert flat[1] == blob[100:110] + assert flat[2] == blob[500:520] + + +async def test_get_ranges_missing_key_raises(memory_store: FsspecStore) -> None: + """A request against a missing key raises BaseExceptionGroup containing FileNotFoundError.""" + proto = default_buffer_prototype() + agen = memory_store.get_ranges("does-not-exist", [RangeByteRequest(0, 10)], prototype=proto) + with pytest.RaisesGroup(pytest.RaisesExc(FileNotFoundError)): + await anext(agen) + + +async def test_get_ranges_forwards_coalescing_kwargs(memory_store: FsspecStore) -> None: + """`max_gap_bytes=-1` forces no merging; we should see three groups for three ranges.""" + blob = bytes(i % 256 for i in range(1024)) + await _write(memory_store, "blob", blob) + proto = default_buffer_prototype() + + ranges = [ + RangeByteRequest(0, 10), + RangeByteRequest(11, 20), # adjacent: would merge under defaults + RangeByteRequest(21, 30), + ] + groups: list[list[tuple[int, Buffer | None]]] = [ + list(group) + async for group in memory_store.get_ranges( + "blob", ranges, prototype=proto, max_gap_bytes=-1 + ) + ] + # With merging disabled, every range becomes its own one-tuple group. + assert sorted(len(g) for g in groups) == [1, 1, 1] + + +async def test_get_ranges_mixed_range_types(memory_store: FsspecStore) -> None: + """Covers RangeByteRequest, OffsetByteRequest, SuffixByteRequest, and None in one call.""" + from zarr_storage.legacy._abc import ByteRequest, OffsetByteRequest, SuffixByteRequest + + blob = bytes(i % 256 for i in range(512)) + await _write(memory_store, "mixed", blob) + proto = default_buffer_prototype() + + ranges: list[ByteRequest | None] = [ + RangeByteRequest(0, 10), + OffsetByteRequest(500), + SuffixByteRequest(12), + None, + ] + flat: dict[int, bytes] = {} + async for group in memory_store.get_ranges("mixed", ranges, prototype=proto): + for idx, buf in group: + assert buf is not None + flat[idx] = buf.to_bytes() + + assert flat[0] == blob[0:10] + assert flat[1] == blob[500:] + assert flat[2] == blob[-12:] + assert flat[3] == blob diff --git a/packages/zarr-storage/tests/test_store/test_get_ranges.py b/packages/zarr-storage/tests/test_store/test_get_ranges.py new file mode 100644 index 0000000000..da339e01a3 --- /dev/null +++ b/packages/zarr-storage/tests/test_store/test_get_ranges.py @@ -0,0 +1,184 @@ +# tests/test_store/test_get_ranges.py +"""Tests for `Store.get_ranges` — the ABC default implementation and wrapper delegation. + +`Store.get_ranges` is defined on the ABC with a default implementation built +on `coalesced_get(self.get, ...)`, so every store inherits a working version. +These tests cover that inherited path and the explicit delegation in +`WrapperStore` (which ensures wrapped stores' optimized overrides are honored). +Store-specific overrides (e.g. `FsspecStore`) have their own test modules. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pytest +from zarr.core.buffer import default_buffer_prototype + +from zarr_storage.legacy import MemoryStore, ZipStore +from zarr_storage.legacy._abc import RangeByteRequest +from zarr_storage.legacy._wrapper import WrapperStore + +if TYPE_CHECKING: + from collections.abc import AsyncIterator, Sequence + from pathlib import Path + + from zarr.core.buffer import Buffer, BufferPrototype + + from zarr_storage.legacy._abc import ByteRequest + + +async def _write(store: MemoryStore, key: str, data: bytes) -> None: + buf = default_buffer_prototype().buffer.from_bytes(data) + await store.set(key, buf) + + +async def test_memory_store_inherits_get_ranges_from_abc() -> None: + """MemoryStore doesn't override `get_ranges`; the ABC default must work end-to-end.""" + store = MemoryStore() + blob = bytes(i % 256 for i in range(512)) + await _write(store, "blob", blob) + + ranges = [RangeByteRequest(0, 10), RangeByteRequest(100, 110)] + proto = default_buffer_prototype() + flat: dict[int, bytes] = {} + async for group in store.get_ranges("blob", ranges, prototype=proto): + for idx, buf in group: + assert buf is not None + flat[idx] = buf.to_bytes() + + assert flat[0] == blob[0:10] + assert flat[1] == blob[100:110] + + +async def test_memory_store_get_ranges_missing_key_raises() -> None: + """A missing key on a default-impl store raises BaseExceptionGroup containing FileNotFoundError.""" + store = MemoryStore() + proto = default_buffer_prototype() + agen = store.get_ranges("does-not-exist", [RangeByteRequest(0, 10)], prototype=proto) + with pytest.RaisesGroup(pytest.RaisesExc(FileNotFoundError)): + await anext(agen) + + +def test_get_ranges_sync_reads_multiple_ranges() -> None: + """The synchronous `get_ranges_sync` on a sync-capable store returns each + requested range, mirroring the async `get_ranges` happy path.""" + import asyncio + + store = MemoryStore() + blob = bytes(i % 256 for i in range(512)) + asyncio.run(_write(store, "blob", blob)) + + ranges = [RangeByteRequest(0, 10), RangeByteRequest(100, 110)] + proto = default_buffer_prototype() + flat: dict[int, bytes] = {} + for idx, buf in store.get_ranges_sync("blob", ranges, prototype=proto): + assert buf is not None + flat[idx] = buf.to_bytes() + + assert flat[0] == blob[0:10] + assert flat[1] == blob[100:110] + + +def test_get_ranges_sync_missing_key_raises() -> None: + """A missing key makes `get_ranges_sync` raise a BaseExceptionGroup + containing FileNotFoundError — the same contract as async `get_ranges`, so + callers handle a deleted shard uniformly across sync and async paths.""" + store = MemoryStore() + proto = default_buffer_prototype() + with pytest.RaisesGroup(pytest.RaisesExc(FileNotFoundError)): + store.get_ranges_sync("does-not-exist", [RangeByteRequest(0, 10)], prototype=proto) + + +def test_get_ranges_sync_on_non_sync_store_raises_type_error(tmp_path: Path) -> None: + """`get_ranges_sync` requires the store to support synchronous reads + (`SupportsGetSync`); a non-sync store raises TypeError rather than silently + falling back.""" + store = ZipStore(tmp_path / "store.zip", mode="w") + proto = default_buffer_prototype() + with pytest.raises(TypeError, match="does not support synchronous reads"): + store.get_ranges_sync("k", [RangeByteRequest(0, 10)], prototype=proto) + + +async def test_wrapper_store_delegates_get_ranges() -> None: + """WrapperStore.get_ranges must delegate to the wrapped store, not fall back to the default.""" + + class CountingMemoryStore(MemoryStore): + """Tallies get_ranges invocations so we can assert delegation.""" + + get_ranges_calls: int = 0 + + async def get_ranges( + self, + key: str, + byte_ranges: Sequence[ByteRequest | None], + *, + prototype: BufferPrototype, + max_concurrency: int = 10, + max_gap_bytes: int = 1 << 20, + max_coalesced_bytes: int = 16 << 20, + ) -> AsyncIterator[Sequence[tuple[int, Buffer | None]]]: + type(self).get_ranges_calls += 1 + async for group in super().get_ranges( + key, + byte_ranges, + prototype=prototype, + max_concurrency=max_concurrency, + max_gap_bytes=max_gap_bytes, + max_coalesced_bytes=max_coalesced_bytes, + ): + yield group + + inner = CountingMemoryStore() + blob = b"x" * 100 + await _write(inner, "k", blob) + wrapped = WrapperStore(inner) + + proto = default_buffer_prototype() + groups: list[list[tuple[int, Buffer | None]]] = [ + list(group) + async for group in wrapped.get_ranges("k", [RangeByteRequest(0, 5)], prototype=proto) + ] + + assert CountingMemoryStore.get_ranges_calls == 1 + assert len(groups) == 1 + assert groups[0][0][0] == 0 + + +async def test_wrapper_store_forwards_coalescing_kwargs() -> None: + """Coalescing kwargs flow through WrapperStore to the wrapped store's get_ranges.""" + + class SpyMemoryStore(MemoryStore): + last_max_gap_bytes: int | None = None + + async def get_ranges( + self, + key: str, + byte_ranges: Sequence[ByteRequest | None], + *, + prototype: BufferPrototype, + max_concurrency: int = 10, + max_gap_bytes: int = 1 << 20, + max_coalesced_bytes: int = 16 << 20, + ) -> AsyncIterator[Sequence[tuple[int, Buffer | None]]]: + type(self).last_max_gap_bytes = max_gap_bytes + async for group in super().get_ranges( + key, + byte_ranges, + prototype=prototype, + max_concurrency=max_concurrency, + max_gap_bytes=max_gap_bytes, + max_coalesced_bytes=max_coalesced_bytes, + ): + yield group + + inner = SpyMemoryStore() + await _write(inner, "k", b"y" * 100) + wrapped = WrapperStore(inner) + proto = default_buffer_prototype() + async for _ in wrapped.get_ranges( + "k", [RangeByteRequest(0, 5)], prototype=proto, max_gap_bytes=-1 + ): + pass + + assert SpyMemoryStore.last_max_gap_bytes == -1 diff --git a/packages/zarr-storage/tests/test_store/test_latency.py b/packages/zarr-storage/tests/test_store/test_latency.py new file mode 100644 index 0000000000..5017621e4b --- /dev/null +++ b/packages/zarr-storage/tests/test_store/test_latency.py @@ -0,0 +1,165 @@ +from __future__ import annotations + +import time +from unittest.mock import patch + +import numpy as np +import pytest +import zarr +from zarr.core.buffer import default_buffer_prototype +from zarr.core.codec_pipeline import FusedCodecPipeline +from zarr.core.config import config as zarr_config + +from zarr_storage.legacy import MemoryStore +from zarr_storage.legacy._abc import RangeByteRequest +from zarr_storage.testing.store import LatencyStore + + +async def test_latency_store_with_read_only_round_trip() -> None: + """ + Ensure that LatencyStore.with_read_only returns another LatencyStore with + the requested read_only state, preserves latency configuration, and does + not change the original wrapper. + """ + base = await MemoryStore.open() + # Start from a read-only underlying store + ro_base = base.with_read_only(read_only=True) + latency_ro = LatencyStore(ro_base, get_latency=0.01, set_latency=0.02) + + assert latency_ro.read_only + assert latency_ro.get_latency == pytest.approx(0.01) + assert latency_ro.set_latency == pytest.approx(0.02) + + buf = default_buffer_prototype().buffer.from_bytes(b"abcd") + + # Cannot write through the read-only wrapper + with pytest.raises( + ValueError, match="store was opened in read-only mode and does not support writing" + ): + await latency_ro.set("key", buf) + + # Create a writable wrapper from the read-only one + writer = latency_ro.with_read_only(read_only=False) + assert isinstance(writer, LatencyStore) + assert not writer.read_only + # Latency configuration is preserved + assert writer.get_latency == latency_ro.get_latency + assert writer.set_latency == latency_ro.set_latency + + # Writes via the writable wrapper succeed + await writer.set("key", buf) + out = await writer.get("key", prototype=default_buffer_prototype()) + assert out is not None + assert out.to_bytes() == buf.to_bytes() + + # Creating a read-only copy from the writable wrapper works and is enforced + reader = writer.with_read_only(read_only=True) + assert isinstance(reader, LatencyStore) + assert reader.read_only + with pytest.raises( + ValueError, match="store was opened in read-only mode and does not support writing" + ): + await reader.set("other", buf) + + # The original read-only wrapper remains read-only + assert latency_ro.read_only + + +@pytest.mark.parametrize( + ("get_latency", "set_latency"), + [ + (0.01, 0.02), + ((0.1, 0.05), (0.2, 0.01)), + ], + ids=["scalar", "distribution"], +) +def test_with_store_preserves_latency_config( + get_latency: float | tuple[float, float], set_latency: float | tuple[float, float] +) -> None: + """Derived stores (e.g. via `with_read_only`) keep the raw latency config — + a `(loc, scale)` distribution must not collapse to one sampled float.""" + store = LatencyStore(MemoryStore(), get_latency=get_latency, set_latency=set_latency) + derived = store.with_read_only(True) + assert derived._get_latency == store._get_latency + assert derived._set_latency == store._set_latency + + +def test_sync_methods_inject_latency(monkeypatch: pytest.MonkeyPatch) -> None: + """`get_sync`/`set_sync` sleep the configured latency on the calling thread + before delegating to the wrapped store.""" + sleeps: list[float] = [] + monkeypatch.setattr(time, "sleep", sleeps.append) + + store = LatencyStore(MemoryStore(), get_latency=0.123, set_latency=0.456) + buf = default_buffer_prototype().buffer.from_bytes(b"abcd") + store.set_sync("key", buf) + assert sleeps == [pytest.approx(0.456)] + out = store.get_sync("key", prototype=default_buffer_prototype()) + assert out is not None + assert out.to_bytes() == b"abcd" + assert sleeps == [pytest.approx(0.456), pytest.approx(0.123)] + + +async def test_get_ranges_pays_latency_per_fetch() -> None: + """`get_ranges` routes through the coalescing default built on `self.get`, + so each merged fetch pays the configured latency instead of bypassing it + via WrapperStore delegation. Two ranges further apart than `max_gap_bytes` + cannot coalesce -> exactly two `get` calls.""" + proto = default_buffer_prototype() + inner = MemoryStore() + await inner.set("blob", proto.buffer.from_bytes(bytes(4 << 20))) + store = LatencyStore(inner, get_latency=0.0) + + requests = [RangeByteRequest(0, 10), RangeByteRequest(2 << 20, (2 << 20) + 10)] + results: list[tuple[int, object]] = [] + with patch.object(store, "get", wraps=store.get) as get_spy: + async for group in store.get_ranges("blob", requests, prototype=proto): + results.extend(group) + assert get_spy.await_count == 2 + assert sorted(idx for idx, _ in results) == [0, 1] + for _, buf in results: + assert buf is not None + assert len(buf) == 10 # type: ignore[arg-type] + + +async def test_get_partial_values_routes_through_get() -> None: + """`get_partial_values` issues one `self.get` per key-range so each fetch + pays the configured latency instead of bypassing it via WrapperStore + delegation.""" + proto = default_buffer_prototype() + inner = MemoryStore() + await inner.set("blob", proto.buffer.from_bytes(b"0123456789")) + store = LatencyStore(inner, get_latency=0.0) + + with patch.object(store, "get", wraps=store.get) as get_spy: + results = await store.get_partial_values( + proto, [("blob", RangeByteRequest(0, 4)), ("blob", None)] + ) + assert get_spy.await_count == 2 + assert results[0] is not None + assert results[0].to_bytes() == b"0123" + assert results[1] is not None + assert results[1].to_bytes() == b"0123456789" + + +def test_latency_store_engages_fused_sync_path() -> None: + """A LatencyStore wrapping a sync-capable store must take the fused sync + fast path: reads go through the inner store's `get_sync`, not the async + fallback.""" + inner = MemoryStore() + store = LatencyStore(inner, get_latency=0.0, set_latency=0.0) + with zarr_config.set({"codec_pipeline.path": "zarr.core.codec_pipeline.FusedCodecPipeline"}): + arr = zarr.create_array( + store=store, + shape=(8,), + chunks=(4,), + dtype="uint8", + compressors=None, + fill_value=0, + ) + assert isinstance(arr._async_array.codec_pipeline, FusedCodecPipeline) + data = np.arange(8, dtype="uint8") + arr[:] = data + with patch.object(inner, "get_sync", wraps=inner.get_sync) as get_sync_spy: + np.testing.assert_array_equal(arr[:], data) + assert get_sync_spy.call_count == 2 # one per chunk diff --git a/packages/zarr-storage/tests/test_store/test_local.py b/packages/zarr-storage/tests/test_store/test_local.py new file mode 100644 index 0000000000..4a92fa244a --- /dev/null +++ b/packages/zarr-storage/tests/test_store/test_local.py @@ -0,0 +1,165 @@ +from __future__ import annotations + +import pathlib +import re + +import numpy as np +import pytest +import zarr +from zarr import create_array +from zarr.core.buffer import Buffer, cpu + +from zarr_storage.legacy import LocalStore +from zarr_storage.legacy._local import _atomic_write +from zarr_storage.testing.store import StoreTests +from zarr_storage.testing.utils import assert_bytes_equal + + +class TestLocalStore(StoreTests[LocalStore, cpu.Buffer]): + store_cls = LocalStore + buffer_cls = cpu.Buffer + + async def get(self, store: LocalStore, key: str) -> Buffer: + return self.buffer_cls.from_bytes((store.root / key).read_bytes()) + + async def set(self, store: LocalStore, key: str, value: Buffer) -> None: + parent = (store.root / key).parent + if not parent.exists(): + parent.mkdir(parents=True) + (store.root / key).write_bytes(value.to_bytes()) + + @pytest.fixture + def store_kwargs(self, tmp_path: pathlib.Path) -> dict[str, str]: + return {"root": str(tmp_path)} + + def test_store_repr(self, store: LocalStore) -> None: + assert str(store) == f"file://{store.root.as_posix()}" + + def test_store_supports_writes(self, store: LocalStore) -> None: + assert store.supports_writes + + def test_store_supports_listing(self, store: LocalStore) -> None: + assert store.supports_listing + + async def test_empty_with_empty_subdir(self, store: LocalStore) -> None: + assert await store.is_empty("") + (store.root / "foo/bar").mkdir(parents=True) + assert await store.is_empty("") + + def test_delete_sync_directory(self, store: LocalStore) -> None: + """`delete_sync` on a key that is a directory must remove the whole tree. + + Mirrors the async `delete_dir` behavior: deleting `"foo"` where + `"foo"` is a directory containing further nested paths should remove + everything under it, not just fail or delete a single file. + """ + (store.root / "foo" / "bar").mkdir(parents=True) + (store.root / "foo" / "bar" / "baz").write_bytes(b"data") + + store.delete_sync("foo") + + assert not (store.root / "foo").exists() + + def test_creates_new_directory(self, tmp_path: pathlib.Path) -> None: + target = tmp_path.joinpath("a", "b", "c") + assert not target.exists() + + store = self.store_cls(root=target) + zarr.group(store=store) + + def test_invalid_root_raises(self) -> None: + """ + Test that a TypeError is raised when a non-str/Path type is used for the `root` argument + """ + with pytest.raises( + TypeError, + match=r"'root' must be a string or Path instance. Got an instance of instead.", + ): + LocalStore(root=0) # type: ignore[arg-type] + + async def test_get_with_prototype_default(self, store: LocalStore) -> None: + """ + Ensure that data can be read via ``store.get`` if the prototype keyword argument is unspecified, i.e. set to ``None``. + """ + data_buf = self.buffer_cls.from_bytes(b"\x01\x02\x03\x04") + key = "c/0" + await self.set(store, key, data_buf) + observed = await store.get(key, prototype=None) + assert_bytes_equal(observed, data_buf) + + @pytest.mark.parametrize("ndim", [0, 1, 3]) + @pytest.mark.parametrize( + "destination", ["destination", "foo/bar/destintion", pathlib.Path("foo/bar/destintion")] + ) + async def test_move( + self, tmp_path: pathlib.Path, ndim: int, destination: pathlib.Path | str + ) -> None: + origin = tmp_path / "origin" + if isinstance(destination, str): + destination = str(tmp_path / destination) + else: + destination = tmp_path / destination + + print(type(destination)) + store = await LocalStore.open(root=origin) + shape = (4,) * ndim + chunks = (2,) * ndim + data = np.arange(4**ndim) + if ndim > 0: + data = data.reshape(*shape) + array = create_array(store, data=data, chunks=chunks or "auto") + + await store.move(destination) + + assert store.root == pathlib.Path(destination) + assert pathlib.Path(destination).exists() + assert not origin.exists() + assert np.array_equal(array[...], data) + + store2 = await LocalStore.open(root=origin) + with pytest.raises( + FileExistsError, match=re.escape(f"Destination root {destination} already exists") + ): + await store2.move(destination) + + +@pytest.mark.parametrize("exclusive", [True, False]) +def test_atomic_write_successful(tmp_path: pathlib.Path, exclusive: bool) -> None: + path = tmp_path / "data" + with _atomic_write(path, "wb", exclusive=exclusive) as f: + f.write(b"abc") + assert path.read_bytes() == b"abc" + assert list(path.parent.iterdir()) == [path] # no temp files + + +@pytest.mark.parametrize("exclusive", [True, False]) +def test_atomic_write_incomplete(tmp_path: pathlib.Path, exclusive: bool) -> None: + path = tmp_path / "data" + with pytest.raises(RuntimeError): # noqa: PT012 + with _atomic_write(path, "wb", exclusive=exclusive) as f: + f.write(b"a") + raise RuntimeError + assert not path.exists() + assert list(path.parent.iterdir()) == [] # no temp files + + +def test_atomic_write_non_exclusive_preexisting(tmp_path: pathlib.Path) -> None: + path = tmp_path / "data" + with path.open("wb") as f: + f.write(b"xyz") + assert path.read_bytes() == b"xyz" + with _atomic_write(path, "wb", exclusive=False) as f: + f.write(b"abc") + assert path.read_bytes() == b"abc" + assert list(path.parent.iterdir()) == [path] # no temp files + + +def test_atomic_write_exclusive_preexisting(tmp_path: pathlib.Path) -> None: + path = tmp_path / "data" + with path.open("wb") as f: + f.write(b"xyz") + assert path.read_bytes() == b"xyz" + with pytest.raises(FileExistsError), _atomic_write(path, "wb", exclusive=True) as f: + f.write(b"abc") + assert path.read_bytes() == b"xyz" + assert list(path.parent.iterdir()) == [path] # no temp files diff --git a/packages/zarr-storage/tests/test_store/test_logging.py b/packages/zarr-storage/tests/test_store/test_logging.py new file mode 100644 index 0000000000..ea1e660726 --- /dev/null +++ b/packages/zarr-storage/tests/test_store/test_logging.py @@ -0,0 +1,170 @@ +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING, TypedDict + +import pytest +import zarr +from zarr.core.buffer import Buffer, cpu, default_buffer_prototype + +from zarr_storage.legacy import LocalStore, LoggingStore +from zarr_storage.testing.store import StoreTests + +if TYPE_CHECKING: + from pathlib import Path + + from zarr_storage.legacy._abc import Store + + +class StoreKwargs(TypedDict): + store: LocalStore + log_level: str + + +class TestLoggingStore(StoreTests[LoggingStore[LocalStore], cpu.Buffer]): + # store_cls is needed to do an isinstance check, so can't be a subscripted generic + store_cls = LoggingStore # type: ignore[assignment] + buffer_cls = cpu.Buffer + + async def get(self, store: LoggingStore[LocalStore], key: str) -> Buffer: + return self.buffer_cls.from_bytes((store._store.root / key).read_bytes()) + + async def set(self, store: LoggingStore[LocalStore], key: str, value: Buffer) -> None: + parent = (store._store.root / key).parent + if not parent.exists(): + parent.mkdir(parents=True) + (store._store.root / key).write_bytes(value.to_bytes()) + + @pytest.fixture + def store_kwargs(self, tmp_path: Path) -> StoreKwargs: + return {"store": LocalStore(str(tmp_path)), "log_level": "DEBUG"} + + @pytest.fixture + def open_kwargs(self, tmp_path: Path) -> dict[str, type[LocalStore] | str]: + return {"store_cls": LocalStore, "root": str(tmp_path), "log_level": "DEBUG"} + + @pytest.fixture + def store(self, store_kwargs: StoreKwargs) -> LoggingStore[LocalStore]: + return self.store_cls(**store_kwargs) + + def test_store_supports_writes(self, store: LoggingStore[LocalStore]) -> None: + assert store.supports_writes + + def test_store_supports_listing(self, store: LoggingStore[LocalStore]) -> None: + assert store.supports_listing + + def test_store_repr(self, store: LoggingStore[LocalStore]) -> None: + assert f"{store!r}" == f"LoggingStore(LocalStore, 'file://{store._store.root.as_posix()}')" + + def test_store_str(self, store: LoggingStore[LocalStore]) -> None: + assert str(store) == f"logging-file://{store._store.root.as_posix()}" + + async def test_default_handler( + self, local_store: LocalStore, capsys: pytest.CaptureFixture[str] + ) -> None: + # Store and then remove existing handlers to enter default handler code path + handlers = logging.getLogger().handlers[:] + for h in handlers: + logging.getLogger().removeHandler(h) + # Test logs are sent to stdout + wrapped = LoggingStore(store=local_store) + buffer = default_buffer_prototype().buffer + res = await wrapped.set("foo/bar/c/0", buffer.from_bytes(b"\x01\x02\x03\x04")) # type: ignore[func-returns-value] + assert res is None + captured = capsys.readouterr() + assert len(captured) == 2 + assert "Calling LocalStore.set" in captured.out + assert "Finished LocalStore.set" in captured.out + # Restore handlers + for h in handlers: + logging.getLogger().addHandler(h) + + def test_is_open_setter_raises(self, store: LoggingStore[LocalStore]) -> None: + "Test that a user cannot change `_is_open` without opening the underlying store." + with pytest.raises( + NotImplementedError, match="LoggingStore must be opened via the `_open` method" + ): + store._is_open = True + + async def test_with_read_only_round_trip(self, local_store: LocalStore) -> None: + """ + Ensure that LoggingStore.with_read_only returns another LoggingStore with + the requested read_only state, preserves logging configuration, and does + not change the original store. + """ + # Start from a read-only underlying store + ro_store = local_store.with_read_only(read_only=True) + wrapped_ro = LoggingStore(store=ro_store, log_level="INFO") + assert wrapped_ro.read_only + + buf = default_buffer_prototype().buffer.from_bytes(b"0123") + + # Cannot write through the read-only wrapper + with pytest.raises( + ValueError, match="store was opened in read-only mode and does not support writing" + ): + await wrapped_ro.set("foo", buf) + + # Create a writable wrapper + writer = wrapped_ro.with_read_only(read_only=False) + assert isinstance(writer, LoggingStore) + assert not writer.read_only + # logging configuration is preserved + assert writer.log_level == wrapped_ro.log_level + assert writer.log_handler == wrapped_ro.log_handler + + # Writes via the writable wrapper succeed + await writer.set("foo", buf) + out = await writer.get("foo", prototype=default_buffer_prototype()) + assert out is not None + assert out.to_bytes() == buf.to_bytes() + + # The original wrapper remains read-only + assert wrapped_ro.read_only + with pytest.raises( + ValueError, match="store was opened in read-only mode and does not support writing" + ): + await wrapped_ro.set("bar", buf) + + +@pytest.mark.parametrize("store", ["local", "memory", "zip"], indirect=["store"]) +async def test_logging_store(store: Store, caplog: pytest.LogCaptureFixture) -> None: + wrapped = LoggingStore(store=store, log_level="DEBUG") + buffer = default_buffer_prototype().buffer + + caplog.clear() + res = await wrapped.set("foo/bar/c/0", buffer.from_bytes(b"\x01\x02\x03\x04")) # type: ignore[func-returns-value] + assert res is None + assert len(caplog.record_tuples) == 2 + for tup in caplog.record_tuples: + assert str(store) in tup[0] + assert f"Calling {type(store).__name__}.set" in caplog.record_tuples[0][2] + assert f"Finished {type(store).__name__}.set" in caplog.record_tuples[1][2] + + caplog.clear() + keys = [k async for k in wrapped.list()] + assert keys == ["foo/bar/c/0"] + assert len(caplog.record_tuples) == 2 + for tup in caplog.record_tuples: + assert str(store) in tup[0] + assert f"Calling {type(store).__name__}.list" in caplog.record_tuples[0][2] + assert f"Finished {type(store).__name__}.list" in caplog.record_tuples[1][2] + + +@pytest.mark.parametrize("store", ["local", "memory", "zip"], indirect=["store"]) +async def test_logging_store_counter(store: Store) -> None: + wrapped = LoggingStore(store=store, log_level="DEBUG") + + arr = zarr.create(shape=(10,), store=wrapped, overwrite=True) + arr[:] = 1 + + assert wrapped.counter["set"] == 2 + assert wrapped.counter["list"] == 0 + assert wrapped.counter["list_dir"] == 0 + assert wrapped.counter["list_prefix"] == 0 + if store.supports_deletes: + assert wrapped.counter["get"] == 0 # 1 if overwrite=False + assert wrapped.counter["delete_dir"] == 1 + else: + assert wrapped.counter["get"] == 1 + assert wrapped.counter["delete_dir"] == 0 diff --git a/packages/zarr-storage/tests/test_store/test_memory.py b/packages/zarr-storage/tests/test_store/test_memory.py new file mode 100644 index 0000000000..6e8c24b7f3 --- /dev/null +++ b/packages/zarr-storage/tests/test_store/test_memory.py @@ -0,0 +1,527 @@ +from __future__ import annotations + +import re +from typing import TYPE_CHECKING, Any + +import numpy as np +import numpy.typing as npt +import pytest +import zarr +from zarr.core.buffer import Buffer, cpu, default_buffer_prototype, gpu +from zarr.errors import ZarrUserWarning + +from zarr_storage.legacy import GpuMemoryStore, ManagedMemoryStore, MemoryStore +from zarr_storage.legacy._utils import _join_paths +from zarr_storage.testing.store import StoreTests +from zarr_storage.testing.utils import gpu_test + +if TYPE_CHECKING: + from zarr.core.common import ZarrFormat + + +# TODO: work out where this warning is coming from and fix it +@pytest.mark.filterwarnings( + re.escape("ignore:coroutine 'ClientCreatorContext.__aexit__' was never awaited") +) +class TestMemoryStore(StoreTests[MemoryStore, cpu.Buffer]): + store_cls = MemoryStore + buffer_cls = cpu.Buffer + + async def set(self, store: MemoryStore, key: str, value: Buffer) -> None: + store._store_dict[key] = value + + async def get(self, store: MemoryStore, key: str) -> Buffer: + return store._store_dict[key] + + @pytest.fixture(params=[None, True]) + def store_kwargs(self, request: pytest.FixtureRequest) -> dict[str, Any]: + kwargs: dict[str, Any] + if request.param is True: + kwargs = {"store_dict": {}} + else: + kwargs = {"store_dict": None} + return kwargs + + @pytest.fixture + async def store(self, store_kwargs: dict[str, Any]) -> MemoryStore: + return self.store_cls(**store_kwargs) + + def test_store_repr(self, store: MemoryStore) -> None: + assert str(store) == f"memory://{id(store._store_dict)}" + + def test_store_supports_writes(self, store: MemoryStore) -> None: + assert store.supports_writes + + def test_store_supports_listing(self, store: MemoryStore) -> None: + assert store.supports_listing + + async def test_list_prefix(self, store: MemoryStore) -> None: + assert True + + @pytest.mark.parametrize("dtype", ["uint8", "float32", "int64"]) + @pytest.mark.parametrize("zarr_format", [2, 3]) + async def test_deterministic_size( + self, store: MemoryStore, dtype: npt.DTypeLike, zarr_format: ZarrFormat + ) -> None: + a = zarr.empty( + store=store, + shape=(3,), + chunks=(1000,), + dtype=dtype, + zarr_format=zarr_format, + overwrite=True, + ) + a[...] = 1 + a.resize((1000,)) + + np.testing.assert_array_equal(a[:3], 1) + np.testing.assert_array_equal(a[3:], 0) + + @pytest.mark.parametrize("method", ["set", "set_sync", "set_if_not_exists"]) + async def test_set_does_not_retain_caller_buffer(self, store: MemoryStore, method: str) -> None: + """Writing a buffer must not alias the caller's memory. + + MemoryStore keeps whatever it is handed alive in a dict, so retaining + the caller's buffer lets a later mutation of that buffer rewrite data + already committed to the store. + """ + source = np.frombuffer(bytearray(b"\x01\x02\x03\x04"), dtype="B") + value = cpu.Buffer.from_array_like(source) + + if method == "set_sync": + store.set_sync("k", value) + else: + await getattr(store, method)("k", value) + + source[:] = 0xF # mutate the caller's memory after the write + stored = await store.get("k", prototype=default_buffer_prototype()) + assert stored is not None + assert stored.to_bytes() == b"\x01\x02\x03\x04" + + @pytest.mark.parametrize( + "pipeline", + [ + "zarr.core.codec_pipeline.BatchedCodecPipeline", + "zarr.core.codec_pipeline.FusedCodecPipeline", + ], + ) + @pytest.mark.parametrize(("shape", "chunks"), [((30,), (10,)), ((8,), (4,)), ((4,), (4,))]) + def test_write_does_not_alias_source_array( + self, pipeline: str, shape: tuple[int], chunks: tuple[int] + ) -> None: + """Mutating the source array after a write must not corrupt stored chunks. + + Without compression the encoded buffer is a zero-copy view of the + caller's array all the way down to the store, so this covers both the + single-chunk and multi-chunk write paths. + """ + with zarr.config.set({"codec_pipeline.path": pipeline}): + array = zarr.create_array( + store=MemoryStore(), shape=shape, chunks=chunks, dtype="i4", compressors=None + ) + source = np.arange(shape[0], dtype="i4") + expected = source.copy() + array[:] = source + source[:] = -1 + + np.testing.assert_array_equal(array[:], expected) + + +# TODO: fix this warning +@pytest.mark.filterwarnings("ignore:Unclosed client session:ResourceWarning") +@gpu_test +class TestGpuMemoryStore(StoreTests[GpuMemoryStore, gpu.Buffer]): + store_cls = GpuMemoryStore + buffer_cls = gpu.Buffer + + async def set(self, store: GpuMemoryStore, key: str, value: gpu.Buffer) -> None: # type: ignore[override] + store._store_dict[key] = value + + async def get(self, store: MemoryStore, key: str) -> Buffer: + return store._store_dict[key] + + @pytest.fixture(params=[None, True]) + def store_kwargs(self, request: pytest.FixtureRequest) -> dict[str, Any]: + kwargs: dict[str, Any] + if request.param is True: + kwargs = {"store_dict": {}} + else: + kwargs = {"store_dict": None} + return kwargs + + @pytest.fixture + async def store(self, store_kwargs: dict[str, Any]) -> GpuMemoryStore: + return self.store_cls(**store_kwargs) + + def test_store_repr(self, store: GpuMemoryStore) -> None: + assert str(store) == f"gpumemory://{id(store._store_dict)}" + + def test_store_supports_writes(self, store: GpuMemoryStore) -> None: + assert store.supports_writes + + def test_store_supports_listing(self, store: GpuMemoryStore) -> None: + assert store.supports_listing + + async def test_list_prefix(self, store: GpuMemoryStore) -> None: + assert True + + def test_dict_reference(self, store: GpuMemoryStore) -> None: + store_dict: dict[str, Any] = {} + result = GpuMemoryStore(store_dict=store_dict) + assert result._store_dict is store_dict + + def test_from_dict(self) -> None: + d = { + "a": gpu.Buffer.from_bytes(b"aaaa"), + "b": cpu.Buffer.from_bytes(b"bbbb"), + } + msg = "Creating a zarr.buffer.gpu.Buffer with an array that does not support the __cuda_array_interface__ for zero-copy transfers, falling back to slow copy based path" + with pytest.warns(ZarrUserWarning, match=msg): + result = GpuMemoryStore.from_dict(d) + for v in result._store_dict.values(): + assert type(v) is gpu.Buffer + + def test_set_sync_converts_to_gpu_buffer(self, store: GpuMemoryStore) -> None: + """`set_sync` must convert its value to a `gpu.Buffer`, mirroring `set`. + + `GpuMemoryStore`'s invariant is that every stored value is a + `gpu.Buffer`. Without this override, the inherited `MemoryStore.set_sync` + would store the CPU buffer it was given as-is, breaking that invariant + for whichever code path (e.g. the fused pipeline) uses the sync API. + """ + cpu_value = cpu.Buffer.from_bytes(b"aaaa") + msg = "Creating a zarr.buffer.gpu.Buffer with an array that does not support the __cuda_array_interface__ for zero-copy transfers, falling back to slow copy based path" + with pytest.warns(ZarrUserWarning, match=msg): + store.set_sync("k", cpu_value) + assert type(store._store_dict["k"]) is gpu.Buffer + + +class TestManagedMemoryStore(StoreTests[ManagedMemoryStore, cpu.Buffer]): + store_cls = ManagedMemoryStore + buffer_cls = cpu.Buffer + + async def set(self, store: ManagedMemoryStore, key: str, value: Buffer) -> None: + store._store_dict[_join_paths([store.path, key])] = value + + async def get(self, store: ManagedMemoryStore, key: str) -> Buffer: + return store._store_dict[_join_paths([store.path, key])] + + @pytest.fixture + def store_kwargs(self, request: pytest.FixtureRequest) -> dict[str, Any]: + # Use a unique name per test to avoid sharing state between tests + # but ensure the name is deterministic for equality tests + # Replace '/' with '-' since store names cannot contain '/' + # A non-empty path exercises prefix handling; a store with an + # unprefixed key in its backing dict would pass these tests + # vacuously with path="". + sanitized_name = request.node.name.replace("/", "-") + return {"name": f"test-{sanitized_name}", "path": "prefix"} + + @pytest.fixture + async def store(self, store_kwargs: dict[str, Any]) -> ManagedMemoryStore: + return self.store_cls(**store_kwargs) + + def test_store_repr(self, store: ManagedMemoryStore) -> None: + assert str(store) == _join_paths([f"memory://{store.name}", store.path]) + + async def test_serializable_store(self, store: ManagedMemoryStore) -> None: + """ + Test pickling semantics for ManagedMemoryStore. + + When pickled and unpickled within the same process (where the original + store still exists in the registry), the unpickled store reconnects to + the same backing dict. + """ + import pickle + + # Add some data to the store + await store.set("test-key", self.buffer_cls.from_bytes(b"test-value")) + + # Pickle and unpickle the store + pickled = pickle.dumps(store) + store2 = pickle.loads(pickled) + + # The unpickled store should reconnect to the same backing dict + assert store2._store_dict is store._store_dict + assert store2.name == store.name + assert store2.path == store.path + assert store2.read_only == store.read_only + + # The data should be accessible + result = await store2.get("test-key") + assert result is not None + assert result.to_bytes() == b"test-value" + + async def test_pickle_with_path(self) -> None: + """Test that path is preserved through pickle round-trip.""" + import pickle + + store = ManagedMemoryStore(name="pickle-path-test", path="some/path") + await store.set("key", self.buffer_cls.from_bytes(b"value")) + + pickled = pickle.dumps(store) + store2 = pickle.loads(pickled) + + assert store2.path == "some/path" + assert store2._store_dict is store._store_dict + + # Check that operations use the path correctly + result = await store2.get("key") + assert result is not None + assert result.to_bytes() == b"value" + + def test_pickle_after_gc(self) -> None: + """ + Test that unpickling after the original store is garbage collected + creates a new empty store with the same name (in the same process). + """ + import gc + import pickle + + # Create a store with a unique name and pickle it + store = ManagedMemoryStore(name="gc-pickle-test") + store._store_dict["key"] = self.buffer_cls.from_bytes(b"value") + pickled = pickle.dumps(store) + + # Delete the store and garbage collect + del store + gc.collect() + + # Unpickling should create a new store with an empty dict + store2 = pickle.loads(pickled) + assert store2.name == "gc-pickle-test" + # The dict is empty because the original was garbage collected + assert len(store2._store_dict) == 0 + + async def test_cross_process_detection(self) -> None: + """ + Test that unpickling a ManagedMemoryStore in a different process raises an error. + + This prevents silent data loss when a store is pickled and unpickled + in a different process (e.g., with multiprocessing). + """ + import os + + store = ManagedMemoryStore(name="cross-process-test") + await store.set("key", self.buffer_cls.from_bytes(b"value")) + + # Get the reduce tuple and modify the state to simulate a different process + cls, args, state = store.__reduce__() + state["created_pid"] = os.getpid() + 1 # Fake a different process ID + + # Manually reconstruct what pickle.loads would do + # This simulates unpickling data that was pickled in a different process + reconstructed = cls(*args) + with pytest.raises(RuntimeError, match="was created in process"): + reconstructed.__setstate__(state) + + def test_store_supports_writes(self, store: ManagedMemoryStore) -> None: + assert store.supports_writes + + def test_store_supports_listing(self, store: ManagedMemoryStore) -> None: + assert store.supports_listing + + @pytest.mark.parametrize("dtype", ["uint8", "float32", "int64"]) + @pytest.mark.parametrize("zarr_format", [2, 3]) + async def test_deterministic_size( + self, store: MemoryStore, dtype: npt.DTypeLike, zarr_format: ZarrFormat + ) -> None: + a = zarr.empty( + store=store, + shape=(3,), + chunks=(1000,), + dtype=dtype, + zarr_format=zarr_format, + overwrite=True, + ) + a[...] = 1 + a.resize((1000,)) + + np.testing.assert_array_equal(a[:3], 1) + np.testing.assert_array_equal(a[3:], 0) + + def test_from_url(self, store: ManagedMemoryStore) -> None: + """Test that from_url creates a store sharing the same dict.""" + url = str(store) + store2 = ManagedMemoryStore.from_url(url) + assert store2._store_dict is store._store_dict + + def test_from_url_with_path(self, store: ManagedMemoryStore) -> None: + """Test that from_url extracts path component from URL.""" + # Reconnect to the fixture's dict via its name, but with an empty + # path, so appending "/some/path" below yields exactly that path. + base = ManagedMemoryStore(name=store.name) + url = f"{base}/some/path" + store2 = ManagedMemoryStore.from_url(url) + assert store2._store_dict is store._store_dict + assert store2.path == "some/path" + assert str(store2) == url + + def test_from_url_invalid(self) -> None: + """Test that from_url raises ValueError for non-existent store.""" + with pytest.raises(ValueError, match="Memory store not found"): + ManagedMemoryStore.from_url("memory://nonexistent-store") + + def test_from_url_not_memory_scheme(self) -> None: + """Test that from_url raises ValueError for non-memory URLs.""" + with pytest.raises(ValueError, match="Expected a 'memory://' URL"): + ManagedMemoryStore.from_url("file:///tmp/test") + + def test_named_store(self) -> None: + """Test that stores can be created with explicit names.""" + store = ManagedMemoryStore(name="my-test-store") + assert store.name == "my-test-store" + assert str(store) == "memory://my-test-store" + + def test_named_store_shares_dict(self) -> None: + """Test that creating a store with the same name shares the dict.""" + store1 = ManagedMemoryStore(name="shared-store") + store2 = ManagedMemoryStore(name="shared-store") + assert store1._store_dict is store2._store_dict + assert store1.name == store2.name + + def test_auto_generated_name(self) -> None: + """Test that stores get auto-generated names when none provided.""" + store = ManagedMemoryStore() + assert store.name is not None + assert str(store) == f"memory://{store.name}" + + def test_with_read_only_shares_dict(self, store: ManagedMemoryStore) -> None: + """Test that with_read_only creates a store sharing the same dict.""" + store2 = store.with_read_only(True) + assert store2._store_dict is store._store_dict + assert store2.read_only is True + assert store.read_only is False + + def test_with_read_only_preserves_path(self) -> None: + """Test that with_read_only preserves the path.""" + store = ManagedMemoryStore(name="path-test", path="some/path") + store2 = store.with_read_only(True) + assert store2.path == "some/path" + assert store2._store_dict is store._store_dict + + async def test_path_prefix_operations(self) -> None: + """Test that store operations use the path prefix correctly.""" + store = ManagedMemoryStore(name="prefix-test") + store_with_path = ManagedMemoryStore.from_url("memory://prefix-test/subdir") + + # Write via store_with_path + await store_with_path.set("key", self.buffer_cls.from_bytes(b"value")) + + # The key should be stored with the prefix in the underlying dict + assert "subdir/key" in store._store_dict + assert "key" not in store._store_dict + + # Read via store_with_path should work + result = await store_with_path.get("key") + assert result is not None + assert result.to_bytes() == b"value" + + # Read via store without path should use full key + result2 = await store.get("subdir/key") + assert result2 is not None + assert result2.to_bytes() == b"value" + + async def test_path_list_operations(self) -> None: + """Test that list operations filter by path prefix.""" + store = ManagedMemoryStore(name="list-test") + + # Set up some keys at different paths + await store.set("a/key1", self.buffer_cls.from_bytes(b"v1")) + await store.set("a/key2", self.buffer_cls.from_bytes(b"v2")) + await store.set("b/key3", self.buffer_cls.from_bytes(b"v3")) + + # Create a store with path "a" + store_a = ManagedMemoryStore.from_url("memory://list-test/a") + + # list() should only return keys under "a", without the "a/" prefix + keys = [k async for k in store_a.list()] + assert sorted(keys) == ["key1", "key2"] + + async def test_path_exists(self) -> None: + """Test that exists() uses the path prefix.""" + store = ManagedMemoryStore(name="exists-test") + await store.set("prefix/key", self.buffer_cls.from_bytes(b"value")) + + store_with_path = ManagedMemoryStore.from_url("memory://exists-test/prefix") + assert await store_with_path.exists("key") + assert not await store_with_path.exists("prefix/key") + + def test_path_normalization(self) -> None: + """Test that paths are normalized.""" + store1 = ManagedMemoryStore(name="norm-test", path="a/b/") + store2 = ManagedMemoryStore(name="norm-test", path="/a/b") + store3 = ManagedMemoryStore(name="norm-test", path="a//b") + assert store1.path == "a/b" + assert store2.path == "a/b" + assert store3.path == "a/b" + + def test_name_cannot_contain_slash(self) -> None: + """Test that store names cannot contain '/'.""" + with pytest.raises(ValueError, match="cannot contain '/'"): + ManagedMemoryStore(name="foo/bar") + + def test_garbage_collection(self) -> None: + """Test that the dict is garbage collected when no stores reference it.""" + import gc + + store = ManagedMemoryStore() + url = str(store) + + # URL should resolve while store exists + store2 = ManagedMemoryStore.from_url(url) + assert store2._store_dict is store._store_dict + + # Delete both stores + del store + del store2 + gc.collect() + + # URL should no longer resolve + with pytest.raises(ValueError, match="garbage collected"): + ManagedMemoryStore.from_url(url) + + def test_sync_methods_respect_path_prefix(self) -> None: + """`get_sync`/`set_sync`/`delete_sync` must prefix keys with `self.path`, + exactly like the async `get`/`set`/`delete` methods. + + `ManagedMemoryStore` used to inherit these from `MemoryStore`, which + writes/reads the raw key. Two stores sharing a dict with different + `path` values would then cross-talk through the sync API. + """ + store = ManagedMemoryStore(name="sync-prefix-test", path="subdir") + data_buf = self.buffer_cls.from_bytes(b"value") + + store.set_sync("key", data_buf) + assert "subdir/key" in store._store_dict + assert "key" not in store._store_dict + + result = store.get_sync("key") + assert result is not None + assert result.to_bytes() == b"value" + + store.delete_sync("key") + assert "subdir/key" not in store._store_dict + + def test_fused_pipeline_respects_path_prefix(self) -> None: + """End-to-end regression: the fused pipeline's sync store fast path must + write chunks under the store's path prefix. + + `FusedCodecPipeline` uses `set_sync`/`get_sync` when a store implements + the sync protocols. If those methods skip the prefix that the async + methods apply, chunk data lands outside `self.path` and a fresh handle + re-reading through the prefix silently sees fill values instead. + """ + with zarr.config.set( + {"codec_pipeline.path": "zarr.core.codec_pipeline.FusedCodecPipeline"} + ): + store = ManagedMemoryStore(name="fused-prefix-test", path="subdir") + arr = zarr.create_array(store, shape=(4,), chunks=(4,), dtype="uint8", zarr_format=3) + arr[:] = np.arange(4, dtype="uint8") + + bad_keys = [k for k in store._store_dict if not k.startswith("subdir/")] + assert bad_keys == [], f"keys written outside the store's path prefix: {bad_keys}" + + store2 = ManagedMemoryStore.from_url("memory://fused-prefix-test/subdir") + arr2 = zarr.open_array(store2, mode="r") + np.testing.assert_array_equal(arr2[:], np.arange(4, dtype="uint8")) diff --git a/packages/zarr-storage/tests/test_store/test_object.py b/packages/zarr-storage/tests/test_store/test_object.py new file mode 100644 index 0000000000..b007a1473b --- /dev/null +++ b/packages/zarr-storage/tests/test_store/test_object.py @@ -0,0 +1,141 @@ +import re +from pathlib import Path +from typing import TypedDict + +import pytest + +obstore = pytest.importorskip("obstore") + +from hypothesis.stateful import ( + run_state_machine_as_test, +) +from obstore.store import LocalStore, MemoryStore, S3Store +from zarr.core.buffer import Buffer, cpu +from zarr.core.sync import _collect_aiterator + +from zarr_storage.legacy import ObjectStore +from zarr_storage.testing.stateful import ZarrHierarchyStateMachine +from zarr_storage.testing.store import StoreTests + + +class StoreKwargs(TypedDict): + store: LocalStore + read_only: bool + + +class TestObjectStore(StoreTests[ObjectStore[LocalStore], cpu.Buffer]): + # store_cls is needed to do an isinstance check, so can't be a subscripted generic + store_cls = ObjectStore # type: ignore[assignment] + buffer_cls = cpu.Buffer + + @pytest.fixture + def store_kwargs(self, tmp_path: Path) -> StoreKwargs: + store = LocalStore(prefix=tmp_path) + return {"store": store, "read_only": False} + + @pytest.fixture + def store(self, store_kwargs: StoreKwargs) -> ObjectStore[LocalStore]: + return self.store_cls(**store_kwargs) + + async def get(self, store: ObjectStore[LocalStore], key: str) -> Buffer: + assert isinstance(store.store, LocalStore) + new_local_store = LocalStore(prefix=store.store.prefix) + return self.buffer_cls.from_bytes(obstore.get(new_local_store, key).bytes()) + + async def set(self, store: ObjectStore[LocalStore], key: str, value: Buffer) -> None: + assert isinstance(store.store, LocalStore) + new_local_store = LocalStore(prefix=store.store.prefix) + obstore.put(new_local_store, key, value.to_bytes()) + + def test_store_repr(self, store: ObjectStore[LocalStore]) -> None: + from fnmatch import fnmatch + + pattern = "ObjectStore(object_store://LocalStore(*))" + assert fnmatch(f"{store!r}", pattern) + + def test_store_supports_writes(self, store: ObjectStore[LocalStore]) -> None: + assert store.supports_writes + + def test_store_supports_partial_writes(self, store: ObjectStore[LocalStore]) -> None: + assert not store.supports_partial_writes + + def test_store_supports_listing(self, store: ObjectStore[LocalStore]) -> None: + assert store.supports_listing + + def test_store_equal(self, store: ObjectStore[LocalStore]) -> None: + """Test store equality""" + # Test equality against a different instance type + assert store != 0 + # Test equality against a different store type + new_memory_store = ObjectStore(MemoryStore()) + assert store != new_memory_store + # Test equality against a read only store + assert isinstance(store.store, LocalStore) + new_local_store = ObjectStore(LocalStore(prefix=store.store.prefix), read_only=True) + assert store != new_local_store + # Test two memory stores cannot be equal + second_memory_store = ObjectStore(MemoryStore()) + assert new_memory_store != second_memory_store + + def test_store_init_raises(self) -> None: + """Test __init__ raises appropriate error for improper store type""" + with pytest.raises(TypeError): + ObjectStore("path/to/store") # type: ignore[type-var] + + async def test_store_getsize(self, store: ObjectStore[LocalStore]) -> None: + buf = cpu.Buffer.from_bytes(b"\x01\x02\x03\x04") + await self.set(store, "key", buf) + size = await store.getsize("key") + assert size == len(buf) + + async def test_store_getsize_prefix(self, store: ObjectStore[LocalStore]) -> None: + buf = cpu.Buffer.from_bytes(b"\x01\x02\x03\x04") + await self.set(store, "c/key1/0", buf) + await self.set(store, "c/key2/0", buf) + size = await store.getsize_prefix("c/key1") + assert size == len(buf) + total_size = await store.getsize_prefix("c") + assert total_size == len(buf) * 2 + + +@pytest.mark.filterwarnings( + re.escape("ignore:datetime.datetime.utcnow() is deprecated:DeprecationWarning") +) +async def test_list_dir_ignores_s3_prefix_marker(moto_server: str) -> None: + """Ensure obstore's exact-prefix S3 directory marker is not listed as a child.""" + boto3 = pytest.importorskip("boto3") + bucket = "object-store-prefix-marker" + client = boto3.client( + "s3", + endpoint_url=moto_server, + region_name="us-east-1", + aws_access_key_id="x", + aws_secret_access_key="x", + ) + client.create_bucket(Bucket=bucket) + client.put_object(Bucket=bucket, Key="g/", Body=b"") + + store = ObjectStore( + S3Store( + bucket=bucket, + endpoint=moto_server, + region="us-east-1", + access_key_id="x", + secret_access_key="x", + client_options={"allow_http": True}, + virtual_hosted_style_request=False, + ) + ) + + assert await _collect_aiterator(store.list_dir("g")) == () + assert await _collect_aiterator(store.list_dir("g/")) == () + + +@pytest.mark.slow_hypothesis +def test_zarr_hierarchy() -> None: + sync_store = ObjectStore(MemoryStore()) + + def mk_test_instance_sync() -> ZarrHierarchyStateMachine: + return ZarrHierarchyStateMachine(sync_store) + + run_state_machine_as_test(mk_test_instance_sync) # type: ignore[no-untyped-call] diff --git a/packages/zarr-storage/tests/test_store/test_stateful.py b/packages/zarr-storage/tests/test_store/test_stateful.py new file mode 100644 index 0000000000..98077813ab --- /dev/null +++ b/packages/zarr-storage/tests/test_store/test_stateful.py @@ -0,0 +1,51 @@ +# Stateful tests for arbitrary Zarr stores. +from collections.abc import Generator + +import pytest +import zarr +from hypothesis.stateful import ( + run_state_machine_as_test, +) + +from zarr_storage.legacy import LocalStore, ZipStore +from zarr_storage.legacy._abc import Store +from zarr_storage.testing.stateful import ZarrHierarchyStateMachine, ZarrStoreStateMachine + +pytestmark = [ + pytest.mark.slow_hypothesis, + # TODO: work out where this warning is coming from and fix + pytest.mark.filterwarnings("ignore:Unclosed client session:ResourceWarning"), +] + + +@pytest.fixture(autouse=True) +def _enable_rectilinear_chunks() -> Generator[None, None, None]: + """Enable rectilinear chunks since strategies may generate them.""" + with zarr.config.set({"array.rectilinear_chunks": True}): + yield + + +@pytest.mark.filterwarnings("ignore::zarr.core.dtype.common.UnstableSpecificationWarning") +def test_zarr_hierarchy(sync_store: Store) -> None: + def mk_test_instance_sync() -> ZarrHierarchyStateMachine: + return ZarrHierarchyStateMachine(sync_store) + + if isinstance(sync_store, ZipStore): + pytest.skip(reason="ZipStore does not support delete") + + run_state_machine_as_test(mk_test_instance_sync) # type: ignore[no-untyped-call] + + +def test_zarr_store(sync_store: Store) -> None: + def mk_test_instance_sync() -> ZarrStoreStateMachine: + return ZarrStoreStateMachine(sync_store) + + if isinstance(sync_store, ZipStore): + pytest.skip(reason="ZipStore does not support delete") + + if isinstance(sync_store, LocalStore): + # This test uses arbitrary keys, which are passed to `set` and `delete`. + # It assumes that `set` and `delete` are the only two operations that modify state. + # But LocalStore, directories can hang around even after a key is delete-d. + pytest.skip(reason="Test isn't suitable for LocalStore.") + run_state_machine_as_test(mk_test_instance_sync) # type: ignore[no-untyped-call] diff --git a/packages/zarr-storage/tests/test_store/test_utils.py b/packages/zarr-storage/tests/test_store/test_utils.py new file mode 100644 index 0000000000..cac7e366e2 --- /dev/null +++ b/packages/zarr-storage/tests/test_store/test_utils.py @@ -0,0 +1,128 @@ +from __future__ import annotations + +import sys +from unittest.mock import patch + +import pytest +from zarr.core.buffer.core import default_buffer_prototype + +from zarr_storage.legacy._abc import SuffixByteRequest +from zarr_storage.legacy._utils import ParsedStoreUrl, _normalize_byte_range_index, parse_store_url + + +class TestParseStoreUrl: + """Tests for parse_store_url.""" + + def test_memory_url(self) -> None: + result = parse_store_url("memory://mystore") + assert result == ParsedStoreUrl( + scheme="memory", name="mystore", path="", raw="memory://mystore" + ) + + def test_memory_url_with_path(self) -> None: + result = parse_store_url("memory://mystore/path/to/data") + assert result == ParsedStoreUrl( + scheme="memory", + name="mystore", + path="path/to/data", + raw="memory://mystore/path/to/data", + ) + + def test_memory_url_no_name(self) -> None: + result = parse_store_url("memory://") + assert result.scheme == "memory" + assert result.name is None + + def test_s3_url(self) -> None: + result = parse_store_url("s3://bucket/key") + assert result == ParsedStoreUrl( + scheme="s3", name="bucket", path="key", raw="s3://bucket/key" + ) + + def test_file_url(self) -> None: + result = parse_store_url("file:///tmp/test") + assert result.scheme == "file" + + def test_local_absolute_path(self) -> None: + result = parse_store_url("/local/path") + assert result == ParsedStoreUrl(scheme="", name=None, path="/local/path", raw="/local/path") + + def test_local_relative_path(self) -> None: + result = parse_store_url("relative/path") + assert result == ParsedStoreUrl( + scheme="", name=None, path="relative/path", raw="relative/path" + ) + + @pytest.mark.parametrize( + "url", + [ + "C:\\Users\\foo", + "C:/Users/foo", + "D:/data/zarr", + "c:/test", + ], + ) + def test_windows_drive_letter(self, url: str) -> None: + """On Windows, bare drive-letter paths must be treated as local paths.""" + with patch.object(sys, "platform", "win32"): + result = parse_store_url(url) + assert result.scheme == "" + assert result.name is None + assert result.path == url + assert result.raw == url + + @pytest.mark.parametrize( + "url", + [ + "file:///C:/Users/foo", + "file://C:/Users/foo", + ], + ) + def test_file_url_with_drive_letter_on_windows(self, url: str) -> None: + """file:// URLs with drive letters are not treated as bare paths.""" + with patch.object(sys, "platform", "win32"): + result = parse_store_url(url) + assert result.scheme == "file" + + @pytest.mark.parametrize( + "url", + [ + "C:\\Users\\foo", + "C:/Users/foo", + ], + ) + def test_drive_letter_not_special_on_non_windows(self, url: str) -> None: + """On non-Windows platforms, drive-letter paths go through urlparse.""" + with patch.object(sys, "platform", "linux"): + result = parse_store_url(url) + # urlparse interprets the drive letter as a scheme + assert result.scheme == "c" + + +class TestNormalizeByteRangeIndex: + """Tests for _normalize_byte_range_index.""" + + def test_suffix_larger_than_data_returns_all_bytes(self) -> None: + """Regression: SuffixByteRequest with suffix > len(data) must not produce a + negative start index that causes numpy to return fewer bytes than available.""" + prototype = default_buffer_prototype() + data = prototype.buffer.from_bytes(b"hello") # 5 bytes + byte_range = SuffixByteRequest(suffix=7) + start, stop = _normalize_byte_range_index(data, byte_range) + assert start == 0, f"start should be 0 (clamped), got {start}" + result = data[start:stop] + assert len(result) == 5, f"expected all 5 bytes, got {len(result)}" + + def test_suffix_exact_length(self) -> None: + """SuffixByteRequest with suffix == len(data) returns all bytes.""" + prototype = default_buffer_prototype() + data = prototype.buffer.from_bytes(b"hello") + start, _stop = _normalize_byte_range_index(data, SuffixByteRequest(suffix=5)) + assert start == 0 + + def test_suffix_shorter_than_data(self) -> None: + """SuffixByteRequest with suffix < len(data) returns the last n bytes.""" + prototype = default_buffer_prototype() + data = prototype.buffer.from_bytes(b"hello") + start, _stop = _normalize_byte_range_index(data, SuffixByteRequest(suffix=3)) + assert start == 2 diff --git a/packages/zarr-storage/tests/test_store/test_wrapper.py b/packages/zarr-storage/tests/test_store/test_wrapper.py new file mode 100644 index 0000000000..d80bb4f827 --- /dev/null +++ b/packages/zarr-storage/tests/test_store/test_wrapper.py @@ -0,0 +1,169 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, TypedDict + +import pytest +from zarr.core.buffer import Buffer +from zarr.core.buffer.cpu import Buffer as CPUBuffer +from zarr.core.buffer.cpu import buffer_prototype + +from zarr_storage.legacy import LocalStore, MemoryStore, WrapperStore, ZipStore +from zarr_storage.legacy._abc import ByteRequest, Store, _store_supports_sync_io +from zarr_storage.testing.store import LatencyStore, StoreTests + +if TYPE_CHECKING: + from pathlib import Path + + from zarr.core.buffer.core import BufferPrototype + + +class StoreKwargs(TypedDict): + store: LocalStore + + +class OpenKwargs(TypedDict): + store_cls: type[LocalStore] + root: str + + +# TODO: fix this warning +@pytest.mark.filterwarnings( + "ignore:coroutine 'ClientCreatorContext.__aexit__' was never awaited:RuntimeWarning" +) +class TestWrapperStore(StoreTests[WrapperStore[Any], Buffer]): + store_cls = WrapperStore + buffer_cls = CPUBuffer + + async def get(self, store: WrapperStore[LocalStore], key: str) -> Buffer: + return self.buffer_cls.from_bytes((store._store.root / key).read_bytes()) + + async def set(self, store: WrapperStore[LocalStore], key: str, value: Buffer) -> None: + parent = (store._store.root / key).parent + if not parent.exists(): + parent.mkdir(parents=True) + (store._store.root / key).write_bytes(value.to_bytes()) + + @pytest.fixture + def store_kwargs(self, tmp_path: Path) -> StoreKwargs: + return {"store": LocalStore(str(tmp_path))} + + @pytest.fixture + def open_kwargs(self, tmp_path: Path) -> OpenKwargs: + return {"store_cls": LocalStore, "root": str(tmp_path)} + + def test_store_supports_writes(self, store: WrapperStore[LocalStore]) -> None: + assert store.supports_writes + + def test_store_supports_listing(self, store: WrapperStore[LocalStore]) -> None: + assert store.supports_listing + + def test_store_repr(self, store: WrapperStore[LocalStore]) -> None: + assert f"{store!r}" == f"WrapperStore(LocalStore, 'file://{store._store.root.as_posix()}')" + + def test_store_str(self, store: WrapperStore[LocalStore]) -> None: + assert str(store) == f"wrapping-file://{store._store.root.as_posix()}" + + def test_check_writeable(self, store: WrapperStore[LocalStore]) -> None: + """ + Test _check_writeable() runs without errors. + """ + store._check_writable() + + def test_close(self, store: WrapperStore[LocalStore]) -> None: + "Test store can be closed" + store.close() + assert not store._is_open + + def test_is_open_setter_raises(self, store: WrapperStore[LocalStore]) -> None: + """ + Test that a user cannot change `_is_open` without opening the underlying store. + """ + with pytest.raises( + NotImplementedError, match="WrapperStore must be opened via the `_open` method" + ): + store._is_open = True + + +# TODO: work out where warning is coming from and fix +@pytest.mark.filterwarnings( + "ignore:coroutine 'ClientCreatorContext.__aexit__' was never awaited:RuntimeWarning" +) +@pytest.mark.parametrize("store", ["local", "memory", "zip"], indirect=True) +async def test_wrapped_set(store: Store, capsys: pytest.CaptureFixture[str]) -> None: + # define a class that prints when it sets + class NoisySetter(WrapperStore[Store]): + async def set(self, key: str, value: Buffer) -> None: + print(f"setting {key}") + await super().set(key, value) + + key = "foo" + value = CPUBuffer.from_bytes(b"bar") + store_wrapped = NoisySetter(store) + await store_wrapped.set(key, value) + captured = capsys.readouterr() + assert f"setting {key}" in captured.out + assert await store_wrapped.get(key, buffer_prototype) == value + + +@pytest.mark.filterwarnings("ignore:Unclosed client session:ResourceWarning") +@pytest.mark.parametrize("store", ["local", "memory", "zip"], indirect=True) +async def test_wrapped_get(store: Store, capsys: pytest.CaptureFixture[str]) -> None: + # define a class that prints when it sets + class NoisyGetter(WrapperStore[Any]): + async def get( + self, key: str, prototype: BufferPrototype, byte_range: ByteRequest | None = None + ) -> None: + print(f"getting {key}") + await super().get(key, prototype=prototype, byte_range=byte_range) + + key = "foo" + value = CPUBuffer.from_bytes(b"bar") + store_wrapped = NoisyGetter(store) + await store_wrapped.set(key, value) + await store_wrapped.get(key, buffer_prototype) + captured = capsys.readouterr() + assert f"getting {key}" in captured.out + + +@pytest.mark.parametrize( + ("store_factory", "expected"), + [ + (lambda tmp: MemoryStore(), True), + (lambda tmp: LocalStore(str(tmp)), True), + (lambda tmp: WrapperStore(MemoryStore()), True), + (lambda tmp: LatencyStore(MemoryStore()), True), + (lambda tmp: ZipStore(tmp / "store.zip", mode="w"), False), + (lambda tmp: WrapperStore(ZipStore(tmp / "store.zip", mode="w")), False), + ], + ids=[ + "memory", + "local", + "wrapper-of-memory", + "latency-wrapper-of-memory", + "zip", + "wrapper-of-zip", + ], +) +def test_supports_sync_io(store_factory: Any, expected: bool, tmp_path: Path | Any) -> None: + """`_store_supports_sync_io` is True only for stores implementing the full + sync surface (get_sync + set_sync + delete_sync); wrappers forward the + wrapped store's capability via `_supports_sync_io`.""" + assert _store_supports_sync_io(store_factory(tmp_path)) is expected + + +def test_wrapper_get_sync_without_inner_sync_raises(tmp_path: Any) -> None: + store = WrapperStore(ZipStore(tmp_path / "store.zip", mode="w")) + with pytest.raises(TypeError, match="does not support synchronous get"): + store.get_sync("key") + + +def test_wrapper_set_sync_without_inner_sync_raises(tmp_path: Any) -> None: + store = WrapperStore(ZipStore(tmp_path / "store.zip", mode="w")) + with pytest.raises(TypeError, match="does not support synchronous set"): + store.set_sync("key", CPUBuffer.from_bytes(b"data")) + + +def test_wrapper_delete_sync_without_inner_sync_raises(tmp_path: Any) -> None: + store = WrapperStore(ZipStore(tmp_path / "store.zip", mode="w")) + with pytest.raises(TypeError, match="does not support synchronous delete"): + store.delete_sync("key") diff --git a/packages/zarr-storage/tests/test_store/test_zip.py b/packages/zarr-storage/tests/test_store/test_zip.py new file mode 100644 index 0000000000..0be5798078 --- /dev/null +++ b/packages/zarr-storage/tests/test_store/test_zip.py @@ -0,0 +1,409 @@ +from __future__ import annotations + +import io +import os +import pickle +import shutil +import tempfile +import zipfile +from typing import TYPE_CHECKING + +import numpy as np +import pytest +import zarr +from hypothesis import settings +from hypothesis.stateful import ( + RuleBasedStateMachine, + initialize, + precondition, + rule, + run_state_machine_as_test, +) +from zarr import create_array +from zarr.core.buffer import Buffer, cpu, default_buffer_prototype +from zarr.core.sync import sync + +from zarr_storage.legacy import ZipStore +from zarr_storage.testing.store import StoreTests + +if TYPE_CHECKING: + from pathlib import Path + from typing import Any + + +# TODO: work out where this is coming from and fix +pytestmark = [ + pytest.mark.filterwarnings( + "ignore:coroutine method 'aclose' of 'ZipStore.list' was never awaited:RuntimeWarning" + ) +] + + +class TestZipStore(StoreTests[ZipStore, cpu.Buffer]): + store_cls = ZipStore + buffer_cls = cpu.Buffer + + @pytest.fixture + def store_kwargs(self) -> dict[str, str | bool]: + fd, temp_path = tempfile.mkstemp() + os.close(fd) + os.unlink(temp_path) + + return {"path": temp_path, "mode": "w", "read_only": False} + + async def get(self, store: ZipStore, key: str) -> Buffer: + buf = store._get(key, prototype=default_buffer_prototype()) + assert buf is not None + return buf + + async def set(self, store: ZipStore, key: str, value: Buffer) -> None: + return store._set(key, value) + + def test_store_read_only(self, store: ZipStore) -> None: + assert not store.read_only + + async def test_read_only_store_raises(self, store_kwargs: dict[str, Any]) -> None: + # we need to create the zipfile in write mode before switching to read mode + store = await self.store_cls.open(**store_kwargs) + store.close() + + kwargs = {**store_kwargs, "mode": "a", "read_only": True} + store = await self.store_cls.open(**kwargs) + assert store._zmode == "a" + assert store.read_only + + # set + with pytest.raises(ValueError): + await store.set("foo", cpu.Buffer.from_bytes(b"bar")) + + def test_store_repr(self, store: ZipStore) -> None: + assert str(store) == f"zip://{store.path}" + + def test_store_supports_writes(self, store: ZipStore) -> None: + assert store.supports_writes + + def test_store_supports_listing(self, store: ZipStore) -> None: + assert store.supports_listing + + # TODO: fix this warning + @pytest.mark.filterwarnings("ignore:Unclosed client session:ResourceWarning") + def test_api_integration(self, store: ZipStore) -> None: + root = zarr.open_group(store=store, mode="a") + + data = np.arange(10000, dtype=np.uint16).reshape(100, 100) + z = root.create_array( + shape=data.shape, chunks=(10, 10), name="foo", dtype=np.uint16, fill_value=99 + ) + z[:] = data + + assert np.array_equal(data, z[:]) + + # you can overwrite existing chunks but zipfile will issue a warning + with pytest.warns(UserWarning, match="Duplicate name: 'foo/c/0/0'"): + z[0, 0] = 100 + + # TODO: assigning an entire chunk to fill value ends up deleting the chunk which is not supported + # a work around will be needed here. + with pytest.raises(NotImplementedError): + z[0:10, 0:10] = 99 + + bar = root.create_group("bar", attributes={"hello": "world"}) + assert "hello" in dict(bar.attrs) + + # keys cannot be deleted + with pytest.raises(NotImplementedError): + del root["bar"] + + store.close() + + @pytest.mark.parametrize("read_only", [True, False]) + async def test_store_open_read_only( + self, store_kwargs: dict[str, Any], read_only: bool + ) -> None: + if read_only: + # create an empty zipfile + with zipfile.ZipFile(store_kwargs["path"], mode="w"): + pass + + await super().test_store_open_read_only(store_kwargs, read_only) + + @pytest.mark.parametrize(("zip_mode", "read_only"), [("w", False), ("a", False), ("x", False)]) + async def test_zip_open_mode_translation( + self, store_kwargs: dict[str, Any], zip_mode: str, read_only: bool + ) -> None: + kws = {**store_kwargs, "mode": zip_mode} + store = await self.store_cls.open(**kws) + assert store.read_only == read_only + + def test_externally_zipped_store(self, tmp_path: Path) -> None: + # See: https://github.com/zarr-developers/zarr-python/issues/2757 + zarr_path = tmp_path / "foo.zarr" + root = zarr.open_group(store=zarr_path, mode="w") + root.require_group("foo") + foo = root.get_group("foo") + foo["bar"] = np.array([1]) + shutil.make_archive(str(zarr_path), "zip", zarr_path) + zip_path = tmp_path / "foo.zarr.zip" + zipped = zarr.open_group(ZipStore(zip_path, mode="r"), mode="r") + assert list(zipped.keys()) == list(root.keys()) + group = zipped.get_group("foo") + assert list(group.keys()) == list(group.keys()) + + async def test_list_without_explicit_open(self, tmp_path: Path) -> None: + # ZipStore.list(), list_dir(), and exists() should auto-open + # the zip file just like _get() and _set() do. + zip_path = tmp_path / "data.zip" + zarr_path = tmp_path / "foo.zarr" + root = zarr.open_group(store=zarr_path, mode="w") + root["x"] = np.array([1, 2, 3]) + shutil.make_archive(str(zarr_path), "zip", zarr_path) + shutil.move(f"{zarr_path}.zip", zip_path) + + store = ZipStore(zip_path, mode="r") + assert not store._is_open + + keys = [k async for k in store.list()] + assert len(keys) > 0 + + store2 = ZipStore(zip_path, mode="r") + assert not store2._is_open + assert await store2.exists(keys[0]) + + store3 = ZipStore(zip_path, mode="r") + assert not store3._is_open + dir_keys = [k async for k in store3.list_dir("")] + assert len(dir_keys) > 0 + + async def test_move(self, tmp_path: Path) -> None: + origin = tmp_path / "origin.zip" + destination = tmp_path / "some_folder" / "destination.zip" + + store = await ZipStore.open(path=origin, mode="a") + array = create_array(store, data=np.arange(10)) + + await store.move(str(destination)) + + assert store.path == destination + assert destination.exists() + assert not origin.exists() + assert np.array_equal(array[...], np.arange(10)) + + +class TestZipStoreFileObj: + """ZipStore backed by an open binary file-like object instead of a path.""" + + @pytest.fixture + def zip_bytes(self, tmp_path: Path) -> bytes: + path = tmp_path / "data.zip" + store = ZipStore(path, mode="w") + zarr.create_array(store, data=np.arange(10), chunks=(5,)) + store.close() + return path.read_bytes() + + def test_read_from_fileobj(self, zip_bytes: bytes) -> None: + # an existing archive can be read through any seekable binary reader + store = ZipStore(io.BytesIO(zip_bytes), mode="r") + array = zarr.open_array(store, mode="r") + assert np.array_equal(array[...], np.arange(10)) + assert store.path is None + + def test_write_to_fileobj(self) -> None: + # a writable file object receives the archive; the bytes it holds + # after close() are a complete, reopenable zip + buffer = io.BytesIO() + store = ZipStore(buffer, mode="w", read_only=False) + zarr.create_array(store, data=np.arange(4)) + store.close() + + roundtrip = ZipStore(io.BytesIO(buffer.getvalue()), mode="r") + array = zarr.open_array(roundtrip, mode="r") + assert np.array_equal(array[...], np.arange(4)) + + async def test_clear_unsupported(self, zip_bytes: bytes) -> None: + # clear() requires a filesystem location, so it raises a clear error + # for file-object-backed stores + store = ZipStore(io.BytesIO(zip_bytes), mode="a", read_only=False) + store._sync_open() + with pytest.raises(NotImplementedError, match="clear.*file-like"): + await store.clear() + + async def test_move_unsupported(self, zip_bytes: bytes) -> None: + # move() requires a filesystem location, so it raises a clear error + # for file-object-backed stores + store = ZipStore(io.BytesIO(zip_bytes), mode="a", read_only=False) + store._sync_open() + with pytest.raises(NotImplementedError, match="move.*file-like"): + await store.move("elsewhere.zip") + + def test_invalid_file_object_rejected(self) -> None: + # objects without read/seek/tell are rejected at construction, not + # deep inside zipfile + with pytest.raises(TypeError, match="read/seek/tell"): + ZipStore(42, mode="r") # type: ignore[arg-type] + + @pytest.mark.parametrize("mode", ["w", "a", "x"]) + def test_non_iobase_reader_write_modes_rejected(self, zip_bytes: bytes, mode: str) -> None: + # readers that are not io.IOBase instances are adapted for reading + # only; write modes are rejected at construction with a clear error + class MinimalReader: + def __init__(self, data: bytes) -> None: + self._buffer = io.BytesIO(data) + + def read(self, size: int, /) -> bytes: + return self._buffer.read(size) + + def seek(self, pos: int, whence: int = 0, /) -> int: + return self._buffer.seek(pos, whence) + + def tell(self) -> int: + return self._buffer.tell() + + with pytest.raises(TypeError, match="opened for reading"): + ZipStore(MinimalReader(zip_bytes), mode=mode, read_only=False) # type: ignore[arg-type] + + def test_fsspec_file(self, tmp_path: Path, zip_bytes: bytes) -> None: + # a file opened through fsspec (already an io.IOBase) is used directly; + # fsspec's local filesystem stands in for a remote one + fsspec = pytest.importorskip("fsspec") + + path = tmp_path / "fsspec.zip" + path.write_bytes(zip_bytes) + with fsspec.open(f"local://{path}", "rb") as fileobj: + store = ZipStore(fileobj, mode="r") + array = zarr.open_array(store, mode="r") + assert np.array_equal(array[...], np.arange(10)) + assert store.path is None + + def test_obstore_reader(self, tmp_path: Path, zip_bytes: bytes) -> None: + # obstore's ReadableFile is not an io.IOBase and its read() returns a + # buffer-protocol object; ZipStore adapts it via _RawReaderAdapter + obstore = pytest.importorskip("obstore") + from obstore.store import LocalStore as ObstoreLocalStore + + (tmp_path / "obstore.zip").write_bytes(zip_bytes) + reader = obstore.open_reader(ObstoreLocalStore(str(tmp_path)), "obstore.zip") + store = ZipStore(reader, mode="r") + array = zarr.open_array(store, mode="r") + assert np.array_equal(array[...], np.arange(10)) + + def test_raw_reader_adapter_eof(self) -> None: + from zarr_storage.legacy._zip import _RawReaderAdapter + + class MinimalReader: + """Non-io.IOBase reader exposing only read/seek/tell, like obstore.""" + + def __init__(self, data: bytes) -> None: + self._buffer = io.BytesIO(data) + + def read(self, size: int, /) -> bytes: + return self._buffer.read(size) + + def seek(self, pos: int, whence: int = 0, /) -> int: + return self._buffer.seek(pos, whence) + + def tell(self) -> int: + return self._buffer.tell() + + # the adapter must clamp reads to EOF: some readers (obstore < 0.6) + # raise on short reads instead of returning fewer bytes + data = b"0123456789" + adapter = _RawReaderAdapter(MinimalReader(data)) # type: ignore[arg-type] + + # A read straddling EOF returns only the remaining bytes. + adapter.seek(len(data) - 3) + buf = bytearray(8) + assert adapter.readinto(buf) == 3 + assert bytes(buf[:3]) == data[-3:] + + # A read at EOF returns 0. + assert adapter.tell() == len(data) + assert adapter.readinto(bytearray(8)) == 0 + + def test_pickle_fileobj_raises(self, zip_bytes: bytes) -> None: + # an open file object cannot be reliably serialized, so pickling a + # file-object-backed store raises with a pointer at the alternative + store = ZipStore(io.BytesIO(zip_bytes), mode="r") + with pytest.raises(TypeError, match="cannot pickle a ZipStore backed by a file-like"): + pickle.dumps(store) + + def test_pickle_path_backed_roundtrip(self, tmp_path: Path, zip_bytes: bytes) -> None: + # path-backed stores remain picklable: the path is serialized and the + # archive is reopened on unpickling + path = tmp_path / "pickled.zip" + path.write_bytes(zip_bytes) + store = ZipStore(path, mode="r") + unpickled = pickle.loads(pickle.dumps(store)) + array = zarr.open_array(unpickled, mode="r") + assert np.array_equal(array[...], np.arange(10)) + + def test_str_and_eq(self, zip_bytes: bytes) -> None: + # file-object-backed stores stringify with the object repr and + # compare equal only when backed by the very same file object + fileobj = io.BytesIO(zip_bytes) + store = ZipStore(fileobj, mode="r") + assert str(store).startswith("zip://<") + assert store == ZipStore(fileobj, mode="r") + assert store != ZipStore(io.BytesIO(zip_bytes), mode="r") + + +class ZipStoreLifecycleMachine(RuleBasedStateMachine): + """Drive a ZipStore through construct / open / write / close transitions. + + 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``). + """ + + def __init__(self, tmp_path: Path) -> None: + super().__init__() + self._tmp_path = tmp_path + self._counter = 0 + self.store: ZipStore | None = None + self._opened = False + + @initialize() + def start(self) -> None: + self.store = None + self._opened = False + + @precondition(lambda self: self.store is None) + @rule() + def construct(self) -> None: + # Fresh path each time so mode="w" never clobbers a closed archive. + self._counter += 1 + self.store = ZipStore(self._tmp_path / f"s{self._counter}.zip", mode="w") + self._opened = False + + @precondition(lambda self: self.store is not None and not self._opened) + @rule() + def open(self) -> None: + assert self.store is not None + self.store._sync_open() + self._opened = True + + @precondition(lambda self: self.store is not None and not self._opened) + @rule() + def write(self) -> None: + assert self.store is not None + # store.set auto-opens the store. + sync(self.store.set("a", cpu.Buffer.from_bytes(b"hi"))) + self._opened = True + + @precondition(lambda self: self.store is not None) + @rule() + def close(self) -> None: + assert self.store is not None + # The property under test: close() must never raise, even with no + # prior open or I/O. + self.store.close() + self.store = None + self._opened = False + + +def test_zipstore_close_lifecycle(tmp_path: Path) -> None: + run_state_machine_as_test( # type: ignore[no-untyped-call] + lambda: ZipStoreLifecycleMachine(tmp_path), + settings=settings(max_examples=50, deadline=None), + ) diff --git a/packages/zarr-storage/tests/test_store_exports.py b/packages/zarr-storage/tests/test_store_exports.py new file mode 100644 index 0000000000..d72fc3f283 --- /dev/null +++ b/packages/zarr-storage/tests/test_store_exports.py @@ -0,0 +1,113 @@ +from __future__ import annotations + +import importlib +import inspect +import subprocess +import sys + +import pytest + +from zarr_storage import legacy + + +@pytest.mark.parametrize( + "name", + [ + "MemoryStore", + "ManagedMemoryStore", + "GpuMemoryStore", + "LocalStore", + "ZipStore", + "FsspecStore", + "ObjectStore", + "WrapperStore", + "LoggingStore", + "LatencyStore", + ], +) +def test_stores_belong_to_extracted_hierarchy(name: str) -> None: + store_class = getattr(legacy, name) + assert issubclass(store_class, legacy.Store) + assert store_class.__module__.startswith("zarr_storage.") + + +def test_cache_store_belongs_to_extracted_hierarchy() -> None: + module = importlib.import_module("zarr_storage.legacy.experimental.cache_store") + assert issubclass(module.CacheStore, legacy.WrapperStore) + + +def test_conformance_suite_is_distributed() -> None: + module = importlib.import_module("zarr_storage.testing") + assert module.StoreTests.__module__ == "zarr_storage.testing.store" + + +@pytest.mark.parametrize( + ("old_module", "new_module", "name"), + [ + ("zarr.storage", "zarr_storage.legacy", name) + for name in ( + "MemoryStore", + "ManagedMemoryStore", + "GpuMemoryStore", + "LocalStore", + "ZipStore", + "FsspecStore", + "ObjectStore", + "WrapperStore", + "LoggingStore", + "StorePath", + ) + ] + + [ + ("zarr.testing.store", "zarr_storage.legacy", "LatencyStore"), + ( + "zarr.experimental.cache_store", + "zarr_storage.legacy.experimental.cache_store", + "CacheStore", + ), + ], +) +def test_store_implementation_signatures(old_module: str, new_module: str, name: str) -> None: + old = getattr(importlib.import_module(old_module), name) + new = getattr(importlib.import_module(new_module), name) + assert old is not new + assert str(inspect.signature(old)) == str(inspect.signature(new)) + for method_name, method in vars(old).items(): + if isinstance(method, (classmethod, staticmethod)): + method = method.__func__ + replacement = vars(new)[method_name].__func__ + elif inspect.isfunction(method): + replacement = vars(new)[method_name] + elif isinstance(method, property): + method = method.fget + replacement = vars(new)[method_name].fget + else: + continue + assert str(inspect.signature(method)) == str(inspect.signature(replacement)) + assert inspect.iscoroutinefunction(method) == inspect.iscoroutinefunction(replacement) + + +def test_import_keeps_zarr_bindings_and_does_not_import_pytest() -> None: + subprocess.run( + [ + sys.executable, + "-c", + """ +import sys +import zarr.storage +from zarr.abc.store import Store + +original_memory_store = zarr.storage.MemoryStore +from zarr_storage.legacy import LatencyStore, MemoryStore + +store = LatencyStore(MemoryStore()) +assert type(store._store) is MemoryStore +assert zarr.storage.MemoryStore is original_memory_store +assert not isinstance(store._store, Store) +assert 'pytest' not in sys.modules +""", + ], + check=True, + capture_output=True, + text=True, + )