Merge pull request #847 from rkingsbury/ext_json

JSONStore: enabled reading of MongoDB extended JSON files

src/maggma/stores/mongolike.py

@@ -18,6 +18,7 @@
 
     from typing_extensions import Literal
 
+import bson
 import mongomock
 import orjson
 from monty.dev import requires
@@ -767,7 +768,7 @@ def read_json_file(self, path) -> List:
         with zopen(path) as f:
             data = f.read()
             data = data.decode() if isinstance(data, bytes) else data
-            objects = orjson.loads(data)
+            objects = bson.json_util.loads(data) if "$oid" in data else orjson.loads(data)
             objects = [objects] if not isinstance(objects, list) else objects
             # datetime objects deserialize to str. Try to convert the last_updated
             # field back to datetime.

tests/stores/test_mongolike.py

@@ -8,6 +8,7 @@
 import orjson
 import pymongo.collection
 import pytest
+from bson.objectid import ObjectId
 from maggma.core import StoreError
 from maggma.stores import JSONStore, MemoryStore, MongoStore, MongoURIStore, MontyStore
 from maggma.validators import JSONSchemaValidator
@@ -424,6 +425,12 @@ def test_json_store_load(jsonstore, test_dir):
         jsonstore = JSONStore("a.json", file_writable=False)
         assert jsonstore.read_only is True
 
+    # test loading an extended JSON file exported from MongoDB
+    js2 = JSONStore(test_dir / "test_set" / "extended_json.json")
+    js2.connect()
+    assert js2.count() == 1
+    assert js2.query_one()["_id"] == ObjectId("64ebee18bd0b1265fe418be2")
+
 
 def test_json_store_writeable(test_dir):
     with ScratchDir("."):

tests/test_files/test_set/extended_json.json

@@ -0,0 +1,7 @@
+[{
+  "_id": {
+    "$oid": "64ebee18bd0b1265fe418be2"
+  },
+  "hello": "world",
+  "task_id": "1"
+}]