sqme: chunk_store fix

odc/loader/_zarr.py

@@ -68,10 +68,12 @@ def _from_zarr_spec(
     drop_variables: Sequence[str] = (),
 ) -> xr.Dataset:
     fsspec_opts = fsspec_opts or {}
-    if target is not None and chunk_store is None:
+    if target is not None:
         if chunk_store is None:
             fs, target = url_to_fs(target, **fsspec_opts)
             chunk_store = fs.get_mapper(target)
+        elif isinstance(chunk_store, fsspec.AbstractFileSystem):
+            chunk_store = chunk_store.get_mapper(target)
 
     # TODO: deal with coordinates being loaded at open time.
     #
@@ -200,9 +202,14 @@ def __init__(
         self,
         src: RasterSource,
         chunks: Any | None = None,
-        chunk_store: fsspec.AbstractFileSystem | Mapping[str, Any] | None = None,
+        chunk_store: (
+            fsspec.AbstractFileSystem | fsspec.FSMap | Mapping[str, Any] | None
+        ) = None,
         drop_variables: Sequence[str] = (),
     ) -> None:
+        if isinstance(chunk_store, fsspec.AbstractFileSystem):
+            chunk_store = chunk_store.get_mapper(src.uri)
+
         driver_data: xr.DataArray | xr.Dataset | SomeDoc = src.driver_data
         self._spec: ZarrSpecDict | None = None
         self._ds: xr.Dataset | None = None
@@ -231,6 +238,9 @@ def __init__(
             assert subdataset is not None
             assert subdataset in ds.data_vars
 
+            if chunk_store is None:
+                chunk_store = fsspec.get_mapper(src.uri)
+
             # recreate xr.DataArray with all the dims/coords/attrs
             # but this time loadable from chunk_store
             xx = ds.data_vars[subdataset]