Allow multiple MultiSourceDataStore instance in parallel

docs/examples/example1/utils.py

@@ -1,5 +1,3 @@
-import pyproj
-import pandas as pd
 import xarray as xr
 
 
@@ -10,4 +8,3 @@ def modify_cci_biomass(ds: xr.Dataset) -> xr.Dataset:
 def modify_gami(ds: xr.Dataset) -> xr.Dataset:
     ds = ds.mean(dim="members")
     return ds.transpose("time", "latitude", "longitude")
-

examples/example1/utils.py

@@ -1,5 +1,3 @@
-import pyproj
-import pandas as pd
 import xarray as xr
 
 
@@ -10,4 +8,3 @@ def modify_cci_biomass(ds: xr.Dataset) -> xr.Dataset:
 def modify_gami(ds: xr.Dataset) -> xr.Dataset:
     ds = ds.mean(dim="members")
     return ds.transpose("time", "latitude", "longitude")
-

xcube_multistore/accessor.py

@@ -36,10 +36,14 @@ class Accessor(ABC):
     """
 
     def __init__(
-        self, store: DataStore, storage: DataStore, identifier: str, notify: Callable
+        self,
+        store: DataStore,
+        storage_temp: DataStore,
+        identifier: str,
+        notify: Callable,
     ):
         self.store = store
-        self.storage = storage
+        self.storage_temp = storage_temp
         self.identifier = identifier
         self.notify = notify
 

xcube_multistore/accessors/stac.py

@@ -20,16 +20,14 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
-from collections.abc import Iterable
 import datetime
 import time
-import uuid
+from collections.abc import Iterable
 
+import xarray as xr
 from pystac_client.exceptions import APIError
 from requests.exceptions import ConnectionError, Timeout
 from urllib3.exceptions import ProtocolError
-import xarray as xr
-from xcube.core.store import DataStoreError
 
 from xcube_multistore.accessor import Accessor
 from xcube_multistore.visualization import GeneratorState
@@ -69,8 +67,7 @@ class StacAccessor(Accessor):
     def open_data(self, data_id: str, **open_params) -> xr.Dataset:
         time_ranges = self._split_time_range(data_id, open_params)
         nb_requests = len(time_ranges)
-        temp_id = uuid.uuid4().hex
-        temp_paths = [f"temp/{temp_id}/{i}.zarr" for i in range(nb_requests)]
+        temp_paths = [f"{self.identifier}/{i}.zarr" for i in range(nb_requests)]
 
         self.notify(
             GeneratorState(
@@ -95,7 +92,7 @@ def open_data(self, data_id: str, **open_params) -> xr.Dataset:
                 )
             )
 
-        dss = [self.storage.open_data(path) for path in valid_paths]
+        dss = [self.storage_temp.open_data(path) for path in valid_paths]
         ds = xr.concat(dss, dim="time", combine_attrs="drop_conflicts")
         return ds
 
@@ -124,7 +121,7 @@ def _open_and_store_with_retry(
                     )
                     return None
 
-                self.storage.write_data(ds, temp_path, replace=True)
+                self.storage_temp.write_data(ds, temp_path, replace=True)
                 return temp_path
 
             except _RETRYABLE_ERRORS as e:

xcube_multistore/constants.py

@@ -24,5 +24,6 @@
 
 LOG = logging.getLogger("xcube.multistore")
 NAME_WRITE_STORE = "storage"
+NAME_WRITE_STORE_TEMP = "storage_temp"
 COMPRESSED_FORMATS = ["zip", "tar", "tar.gz", "rar"]
 CRS_WGS84 = "epsg:4326"

xcube_multistore/multistore.py

@@ -44,7 +44,7 @@
 
 from .accessors import guess_accessor
 from .config import MultiSourceConfig
-from .constants import LOG, NAME_WRITE_STORE
+from .constants import LOG, NAME_WRITE_STORE, NAME_WRITE_STORE_TEMP
 from .gridmappings import GridMappings
 from .stores import DataStores
 from .utils import (
@@ -462,7 +462,10 @@ def _notify_error(self, identifier: str, exception: Any):
             )
         )
         if "Keyboard Interrupt caught! Exiting gracefully." == exception:
+            self._close()
             sys.exit(1)
+        else:
+            self._close(data_id=identifier)
 
     def _preload_datasets(self):
         for config_preload in self.config.preload_datasets:
@@ -601,9 +604,9 @@ def _open_single_dataset(self, config: dict) -> xr.Dataset | Exception:
         store = getattr(self.stores, config["store"])
         store_id = getattr(self.stores, f"{config['store']}_store_id")
         open_params = copy.deepcopy(config.get("open_params", {}))
-        storage_store = getattr(self.stores, NAME_WRITE_STORE)
+        storage_store_temp = getattr(self.stores, NAME_WRITE_STORE_TEMP)
         accessor = guess_accessor(
-            store_id, store, storage_store, config["identifier"], self._notify
+            store_id, store, storage_store_temp, config["identifier"], self._notify
         )
         ds = accessor.open_data(config["data_id"], **open_params)
         if isinstance(ds, MultiLevelDataset):
@@ -703,15 +706,15 @@ def _write_dataset(self, ds: xr.Dataset, config: dict) -> xr.Dataset | Exception
         store.write_data(ds, data_id, replace=True)
         return ds
 
-    def _close(self):
-        store = getattr(self.stores, NAME_WRITE_STORE)
+    def _close(self, data_id: str | None = None):
+        store = getattr(self.stores, NAME_WRITE_STORE_TEMP)
 
         try:
             fs = store.fs
-            temp_path = f"{store.root}/temp"
-
-            if fs.exists(temp_path):
-                fs.rm(temp_path, recursive=True)
-
+            root = store.root
+            if data_id is not None:
+                root = f"{root}/{data_id}"
+            if fs.exists(root):
+                fs.rm(root, recursive=True)
         except (AttributeError, OSError):
             pass

xcube_multistore/stores.py

@@ -20,7 +20,9 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
+import copy
 import json
+import uuid
 
 from xcube.core.store import new_data_store
 
@@ -36,11 +38,20 @@ def setup_data_stores(cls, config: MultiSourceConfig):
             if config_store["store_id"] == "clms":
                 with open(store_params["credentials"]) as f:
                     store_params["credentials"] = json.load(f)
-            if config_store["identifier"] == "storage" and (
-                config_store["store_id"] in ["file", "s3"]
-            ):
-                if not "max_depth" in store_params:
-                    store_params["max_depth"] = 10
+            if config_store["identifier"] == "storage":
+                # setup temp store
+                temp_store_params = copy.deepcopy(store_params)
+                temp_store_params["root"] += f"/temp_{uuid.uuid4().hex}"
+                setattr(
+                    cls,
+                    f"{identifier}_temp",
+                    new_data_store(config_store["store_id"], **temp_store_params),
+                )
+                setattr(cls, f"{identifier}_temp_store_id", config_store["store_id"])
+                if config_store["store_id"] in ["file", "s3"]:
+                    if "max_depth" not in store_params:
+                        store_params["max_depth"] = 10
+
             setattr(
                 cls,
                 identifier,