From 5fde3a29ca0a9f1a005eb5000d1e8585d4ca0bde Mon Sep 17 00:00:00 2001
From: Ian Carroll <carroll.ian@gmail.com>
Date: Fri, 5 Apr 2024 13:15:10 -0400
Subject: [PATCH] Make sure fs exceptions are raised if not MissingFs
 exceptions (clone) (#1604)

* Make sure fs exceptions are raised if not Missing

* lint

* add missing argument in tests, lint

* clear memory filesystem during test

* improve commenting

* add memory_store fixture, getitems performance

* Update release.rst

* improve FSStore.test_exception coverage

---------

Co-authored-by: Martin Durant <martin.durant@alumni.utoronto.ca>
Co-authored-by: Joe Hamman <joe@earthmover.io>
Co-authored-by: Josh Moore <josh@openmicroscopy.org>
Co-authored-by: Sanket Verma <svsanketverma5@gmail.com>
---
 docs/release.rst           |  3 +++
 zarr/storage.py            | 22 +++++++++++++++++-----
 zarr/tests/test_storage.py | 25 +++++++++++++++++++++++++
 3 files changed, 45 insertions(+), 5 deletions(-)

diff --git a/docs/release.rst b/docs/release.rst
index b011b0986b..da802651c2 100644
--- a/docs/release.rst
+++ b/docs/release.rst
@@ -41,6 +41,9 @@ Maintenance
 
 * Bump minimum supported NumPy version to 1.23 (per spec 0000)
   By :user:`Joe Hamman <jhamman>` :issue:`1719`.
+  
+* FSStore now raises rather than return bad data.
+  By :user:`Martin Durant <martindurant>` and :user:`Ian Carroll <itcarroll>` :issue:`1604`.
 
 .. _release_2.17.1:
 
diff --git a/zarr/storage.py b/zarr/storage.py
index f6903d29b2..10f55f0ba3 100644
--- a/zarr/storage.py
+++ b/zarr/storage.py
@@ -1417,11 +1417,23 @@ def _normalize_key(self, key):
     def getitems(
         self, keys: Sequence[str], *, contexts: Mapping[str, Context]
     ) -> Mapping[str, Any]:
-        keys_transformed = [self._normalize_key(key) for key in keys]
-        results = self.map.getitems(keys_transformed, on_error="omit")
-        # The function calling this method may not recognize the transformed keys
-        # So we send the values returned by self.map.getitems back into the original key space.
-        return {keys[keys_transformed.index(rk)]: rv for rk, rv in results.items()}
+        keys_transformed = {self._normalize_key(key): key for key in keys}
+        results_transformed = self.map.getitems(list(keys_transformed), on_error="return")
+        results = {}
+        for k, v in results_transformed.items():
+            if isinstance(v, self.exceptions):
+                # Cause recognized exceptions to prompt a KeyError in the
+                # function calling this method
+                continue
+            elif isinstance(v, Exception):
+                # Raise any other exception
+                raise v
+            else:
+                # The function calling this method may not recognize the transformed
+                # keys, so we send the values returned by self.map.getitems back into
+                # the original key space.
+                results[keys_transformed[k]] = v
+        return results
 
     def __getitem__(self, key):
         key = self._normalize_key(key)
diff --git a/zarr/tests/test_storage.py b/zarr/tests/test_storage.py
index 358d043ad6..ae8a56fa61 100644
--- a/zarr/tests/test_storage.py
+++ b/zarr/tests/test_storage.py
@@ -1098,6 +1098,12 @@ def mock_walker_no_slash(_path):
 
 @pytest.mark.skipif(have_fsspec is False, reason="needs fsspec")
 class TestFSStore(StoreTests):
+    @pytest.fixture
+    def memory_store(self):
+        store = FSStore("memory://")
+        yield store
+        store.fs.store.clear()
+
     def create_store(self, normalize_keys=False, dimension_separator=".", path=None, **kwargs):
         if path is None:
             path = tempfile.mkdtemp()
@@ -1337,6 +1343,25 @@ def test_s3_complex(self):
         )
         assert (a[:] == -np.ones((8, 8, 8))).all()
 
+    def test_exceptions(self, memory_store):
+        fs = memory_store.fs
+        group = zarr.open(memory_store, mode="w")
+        x = group.create_dataset("x", data=[1, 2, 3])
+        y = group.create_dataset("y", data=1)
+        fs.store["/x/0"] = None
+        fs.store["/y/0"] = None
+        # no exception from FSStore.getitems getting KeyError
+        assert group.store.getitems(["foo"], contexts={}) == {}
+        # exception from FSStore.getitems getting AttributeError
+        with pytest.raises(Exception):
+            group.store.getitems(["x/0"], contexts={})
+        # exception from FSStore.getitems getting AttributeError
+        with pytest.raises(Exception):
+            x[...]
+        # exception from FSStore.__getitem__ getting AttributeError
+        with pytest.raises(Exception):
+            y[...]
+
 
 @pytest.mark.skipif(have_fsspec is False, reason="needs fsspec")
 class TestFSStoreWithKeySeparator(StoreTests):