Add orjson options in JSONStore

src/maggma/stores/gridfs.py

@@ -2,7 +2,7 @@
 """
 Module containing various definitions of Stores.
 Stores are a default access pattern to data and provide
-various utillities
+various utilities
 """
 
 import copy
@@ -11,7 +11,7 @@
 from ruamel import yaml
 from datetime import datetime
 from pymongo.errors import ConfigurationError
-from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
+from typing import Dict, Iterator, List, Optional, Tuple, Union
 
 import gridfs
 from monty.json import jsanitize
@@ -162,7 +162,9 @@ def connect(self, force_reset: bool = False):
     def _collection(self):
         """Property referring to underlying pymongo collection"""
         if self._coll is None:
-            raise StoreError("Must connect Mongo-like store before attemping to use it")
+            raise StoreError(
+                "Must connect Mongo-like store before attempting to use it"
+            )
         return self._coll
 
     @property
@@ -244,7 +246,6 @@ def query(
             if properties is not None and prop_keys.issubset(set(doc.keys())):
                 yield {p: doc[p] for p in properties if p in doc}
             else:
-
                 metadata = doc.get("metadata", {})
 
                 data = self._collection.find_one(
@@ -354,7 +355,7 @@ def groupby(
 
     def ensure_index(self, key: str, unique: Optional[bool] = False) -> bool:
         """
-        Tries to create an index and return true if it suceeded
+        Tries to create an index and return true if it succeeded
         Currently operators on the GridFS files collection
         Args:
             key: single key to index
@@ -447,7 +448,8 @@ def remove_docs(self, criteria: Dict):
             self._collection.delete(_id)
 
     def close(self):
-        self._collection.database.client.close()
+        self._files_store.close()
+        self._coll = None
         if self.ssh_tunnel is not None:
             self.ssh_tunnel.stop()
 

src/maggma/stores/mongolike.py

@@ -10,7 +10,7 @@
 from itertools import chain, groupby
 from socket import socket
 import warnings
-from typing import Dict, Iterator, List, Optional, Tuple, Union
+from typing import Dict, Iterator, List, Optional, Tuple, Union, Any
 
 import mongomock
 import orjson
@@ -704,6 +704,8 @@ def __init__(
         self,
         paths: Union[str, List[str]],
         read_only: bool = True,
+        serialization_option: int | None = None,
+        serialization_default: Any | None = None,
         **kwargs,
     ):
         """
@@ -720,6 +722,10 @@ def __init__(
                        Note that when read_only=False, JSONStore only supports a single JSON
                        file. If the file does not exist, it will be automatically created
                        when the JSONStore is initialized.
+            serialization_option:
+                option that will be passed to the orjson.dump when saving to the json the file.
+            serialization_default:
+                default that will be passed to the orjson.dump when saving to the json the file.
         """
         paths = paths if isinstance(paths, (list, tuple)) else [paths]
         self.paths = paths
@@ -755,6 +761,8 @@ def __init__(
                 f.write(bytesdata.decode("utf-8"))
 
         self.default_sort = None
+        self.serialization_option = serialization_option
+        self.serialization_default = serialization_default
 
         super().__init__(**kwargs)
 
@@ -838,7 +846,11 @@ def update_json_file(self):
             data = [d for d in self.query()]
             for d in data:
                 d.pop("_id")
-            bytesdata = orjson.dumps(data)
+            bytesdata = orjson.dumps(
+                data,
+                option=self.serialization_option,
+                default=self.serialization_default,
+            )
             f.write(bytesdata.decode("utf-8"))
 
     def __hash__(self):

tests/stores/test_gridfs.py

@@ -6,6 +6,7 @@
 import numpy as np
 import numpy.testing.utils as nptu
 import pytest
+from maggma.core import StoreError
 
 from maggma.stores import GridFSStore, MongoStore
 from maggma.stores.gridfs import files_collection_fields, GridFSURIStore
@@ -218,7 +219,7 @@ def test_index(gridfsstore):
 
 def test_gfs_metadata(gridfsstore):
     """
-    Ensure metadata is put back in the docuemnt
+    Ensure metadata is put back in the document
     """
     tic = datetime(2018, 4, 12, 16)
 
@@ -244,7 +245,6 @@ def test_gridfsstore_from_launchpad_file(lp_file):
 
 
 def test_searchable_fields(gridfsstore):
-
     tic = datetime(2018, 4, 12, 16)
 
     data = [
@@ -259,7 +259,6 @@ def test_searchable_fields(gridfsstore):
 
 
 def test_additional_metadata(gridfsstore):
-
     tic = datetime(2018, 4, 12, 16)
 
     data = [
@@ -299,3 +298,12 @@ def test_gridfs_uri_dbname_parse():
     uri_with_db = "mongodb://uuu:xxxx@host:27017"
     with pytest.raises(ConfigurationError):
         GridFSURIStore(uri_with_db, "test")
+
+
+def test_close(gridfsstore):
+    assert gridfsstore.query_one() is None
+    gridfsstore.close()
+    with pytest.raises(StoreError):
+        gridfsstore.query_one()
+    # reconnect to allow the drop of the collection in the fixture
+    gridfsstore.connect()