JSONStore: write file on init; add tests; verbose KeyError

src/maggma/stores/mongolike.py

@@ -7,7 +7,7 @@
 
 import json
 from pathlib import Path
-
+from datetime import datetime
 import yaml
 from itertools import chain, groupby
 from socket import socket
@@ -697,7 +697,13 @@ def __init__(self, paths: Union[str, List[str]], file_writable=False, **kwargs):
             )
         self.file_writable = file_writable
         self.kwargs = kwargs
-        super().__init__(collection_name="collection", **kwargs)
+
+        # create the .json file if it does not exist
+        if self.file_writable and not Path(self.paths[0]).exists():
+            with zopen(self.paths[0], "w") as f:
+                data: List[dict] = []
+                json.dump(data, f, default=json_serial)
+        super().__init__(**kwargs)
 
     def connect(self, force_reset=False):
         """
@@ -710,7 +716,17 @@ def connect(self, force_reset=False):
                 data = data.decode() if isinstance(data, bytes) else data
                 objects = json.loads(data)
                 objects = [objects] if not isinstance(objects, list) else objects
-                self.update(objects)
+                try:
+                    self.update(objects)
+                except KeyError:
+                    raise KeyError(
+                        f"""
+                        Key field '{self.key}' not found in {f.name}. This
+                        could mean that this JSONStore was initially created with a different key field.
+                        The keys found in the .json file are {list(objects[0].keys())}. Try
+                        re-initializing your JSONStore using one of these as the key arguments.
+                        """
+                    )
 
     def update(self, docs: Union[List[Dict], Dict], key: Union[List, str, None] = None):
         """
@@ -750,7 +766,7 @@ def update_json_file(self):
             data = [d for d in self.query()]
             for d in data:
                 d.pop("_id")
-            json.dump(data, f)
+            json.dump(data, f, default=json_serial)
 
     def __hash__(self):
         return hash((*self.paths, self.last_updated_field))
@@ -893,3 +909,13 @@ def _find_free_port(address="0.0.0.0"):
     s = socket()
     s.bind((address, 0))  # Bind to a free port provided by the host.
     return s.getsockname()[1]  # Return the port number assigned.
+
+
+# Included for now to make it possible to serialize datetime objects. Probably
+# maggma already has a solution to this somewhere.
+def json_serial(obj):
+    """JSON serializer for objects not serializable by default json code"""
+
+    if isinstance(obj, (datetime)):
+        return obj.isoformat()
+    raise TypeError("Type %s not serializable" % type(obj))

tests/stores/test_mongolike.py

@@ -2,6 +2,7 @@
 import shutil
 from datetime import datetime
 from unittest import mock
+from pathlib import Path
 
 import mongomock.collection
 from monty.tempfile import ScratchDir
@@ -414,9 +415,20 @@ def test_json_store_load(jsonstore, test_dir):
     jsonstore.connect()
     assert len(list(jsonstore.query())) == 20
 
+    # confirm descriptive error raised if you get a KeyError
+    with pytest.raises(KeyError, match="Key field 'random key' not found"):
+        jsonstore = JSONStore(test_dir / "test_set" / "c.json.gz", key="random_key")
+
 
 def test_json_store_writeable(test_dir):
     with ScratchDir("."):
+        # if the .json does not exist, it should be created
+        jsonstore = JSONStore("a.json", file_writable=True)
+        assert Path("a.json").exists()
+        jsonstore.connect()
+        # confirm RunTimeError with multiple paths
+        with pytest.raises(RuntimeError, match="multiple JSON"):
+            jsonstore = JSONStore(["a.json", "d.json"], file_writable=True)
         shutil.copy(test_dir / "test_set" / "d.json", ".")
         jsonstore = JSONStore("d.json", file_writable=True)
         jsonstore.connect()