Merge pull request #158 from pastas/allow_loading_pastastore_from_file

Allow loading pastastore from file

pastastore/connectors.py

@@ -763,6 +763,7 @@ def __init__(self, name: str, uri: str, verbose: bool = True):
         """
         try:
             import arcticdb
+
         except ModuleNotFoundError as e:
             print("Please install arcticdb with `pip install arcticdb`!")
             raise e
@@ -776,6 +777,9 @@ def __init__(self, name: str, uri: str, verbose: bool = True):
         # for older versions of PastaStore, if oseries_models library is empty
         # populate oseries - models database
         self._update_all_oseries_model_links()
+        # write pstore file to store database info that can be used to load pstore
+        if "lmdb" in self.uri:
+            self.write_pstore_config_file()
 
     def _initialize(self, verbose: bool = True) -> None:
         """Initialize the libraries (internal method)."""
@@ -791,6 +795,29 @@ def _initialize(self, verbose: bool = True) -> None:
                     )
             self.libs[libname] = self._get_library(libname)
 
+    def write_pstore_config_file(self, path: str = None) -> None:
+        """Write pstore configuration file to store database info."""
+        # NOTE: method is not private as theoretically an ArcticDB
+        # database could also be hosted in the cloud, in which case,
+        # writing this config in the folder holding the database
+        # is no longer possible. For those situations, the user can
+        # write this config file and specify the path it should be
+        # written to.
+        config = {
+            "connector_type": self.conn_type,
+            "name": self.name,
+            "uri": self.uri,
+        }
+        if path is None and "lmdb" in self.uri:
+            path = self.uri.split("://")[1]
+        elif path is None and "lmdb" not in self.uri:
+            raise ValueError("Please provide a path to write the pastastore file!")
+
+        with open(
+            os.path.join(path, f"{self.name}.pastastore"), "w", encoding="utf-8"
+        ) as f:
+            json.dump(config, f)
+
     def _library_name(self, libname: str) -> str:
         """Get full library name according to ArcticDB (internal method)."""
         return ".".join([self.name, libname])
@@ -1164,13 +1191,16 @@ def __init__(self, name: str, path: str, verbose: bool = True):
             whether to print message when database is initialized, by default True
         """
         self.name = name
+        self.parentdir = path
         self.path = os.path.abspath(os.path.join(path, self.name))
         self.relpath = os.path.relpath(self.path)
         self._initialize(verbose=verbose)
         self.models = ModelAccessor(self)
         # for older versions of PastaStore, if oseries_models library is empty
         # populate oseries_models library
         self._update_all_oseries_model_links()
+        # write pstore file to store database info that can be used to load pstore
+        self._write_pstore_config_file()
 
     def _initialize(self, verbose: bool = True) -> None:
         """Initialize the libraries (internal method)."""
@@ -1188,6 +1218,18 @@ def _initialize(self, verbose: bool = True) -> None:
                     )
             setattr(self, f"lib_{val}", os.path.join(self.path, val))
 
+    def _write_pstore_config_file(self):
+        """Write pstore configuration file to store database info."""
+        config = {
+            "connector_type": self.conn_type,
+            "name": self.name,
+            "path": self.parentdir,
+        }
+        with open(
+            os.path.join(self.path, f"{self.name}.pastastore"), "w", encoding="utf-8"
+        ) as f:
+            json.dump(config, f)
+
     def _get_library(self, libname: str):
         """Get path to directory holding data.
 

pastastore/store.py

@@ -15,7 +15,7 @@
 from tqdm.auto import tqdm
 
 from pastastore.base import BaseConnector
-from pastastore.connectors import DictConnector
+from pastastore.connectors import ArcticDBConnector, DictConnector, PasConnector
 from pastastore.plotting import Maps, Plots
 from pastastore.util import _custom_warning
 from pastastore.version import PASTAS_GEQ_150, PASTAS_LEQ_022
@@ -79,6 +79,24 @@ def __init__(
         self.plots = Plots(self)
         self.yaml = PastastoreYAML(self)
 
+    @classmethod
+    def from_pastastore_config_file(cls, fname):
+        """Create a PastaStore from a pastastore config file."""
+        with open(fname, "r") as f:
+            cfg = json.load(f)
+
+        conn_type = cfg.pop("connector_type")
+        if conn_type == "pas":
+            conn = PasConnector(**cfg)
+        elif conn_type == "arcticdb":
+            conn = ArcticDBConnector(**cfg)
+        else:
+            raise ValueError(
+                f"Cannot load connector type: '{conn_type}'. "
+                "This is only supported for PasConnector and ArcticDBConnector."
+            )
+        return cls(conn)
+
     @property
     def empty(self) -> bool:
         """Check if the PastaStore is empty."""

pastastore/util.py

@@ -1,6 +1,7 @@
 """Useful utilities for pastastore."""
 
 import os
+import shutil
 from typing import Dict, List, Optional, Union
 
 import numpy as np
@@ -43,8 +44,6 @@ def delete_arcticdb_connector(
         list of library names to delete, by default None which deletes
         all libraries
     """
-    import shutil
-
     import arcticdb
 
     if conn is not None:
@@ -75,9 +74,17 @@ def delete_arcticdb_connector(
             print()
             print(f" - deleted: {lib}")
 
-    remaining = [ilib for ilib in arc.list_libraries() if ilib.split(".") == name]
+    # delete .pastastore file if entire pastastore is deleted
+    remaining_libs = [
+        ilib for ilib in arc.list_libraries() if ilib.split(".")[0] == name
+    ]
+    if remaining_libs == 0:
+        os.unlink(os.path.join(uri.split("//")[-1], f"{name}.pastastore"))
+
+    # check if any remaining libraries in lmdb dir, if none, delete entire folder
+    remaining = arc.list_libraries()
     if len(remaining) == 0:
-        shutil.rmtree(os.path.join(conn.uri.split("//")[-1], name))
+        shutil.rmtree(os.path.join(conn.uri.split("//")[-1]))
 
     print("Done!")
 
@@ -98,8 +105,6 @@ def delete_dict_connector(conn, libraries: Optional[List[str]] = None) -> None:
 
 def delete_pas_connector(conn, libraries: Optional[List[str]] = None) -> None:
     """Delete PasConnector object."""
-    import shutil
-
     print(f"Deleting PasConnector database: '{conn.name}' ... ", end="")
     if libraries is None:
         shutil.rmtree(conn.path)

tests/test_003_pastastore.py

@@ -296,6 +296,16 @@ def test_to_from_zip(pstore):
         os.remove(zipname)
 
 
+def test_load_pastastore_from_config_file(pstore):
+    if pstore.type == "pas" or pstore.type == "arcticdb":
+        path = (
+            pstore.conn.path if pstore.type == "pas" else pstore.conn.uri.split("//")[1]
+        )
+        fname = os.path.join(path, f"{pstore.conn.name}.pastastore")
+        pstore2 = pst.PastaStore.from_pastastore_config_file(fname)
+        assert not pstore2.empty
+
+
 def test_example_pastastore():
     from pastastore.datasets import example_pastastore
 

tests/test_006_benchmark.py

@@ -176,6 +176,3 @@ def test_benchmark_read_model_arcticdb(benchmark):
     conn = pst.ArcticDBConnector("test", uri)
     _ = benchmark(read_model, conn=conn)
     pst.util.delete_arcticdb_connector(conn=conn)
-    import shutil
-
-    shutil.rmtree("./arctic_db/")