From fced54ca13531e181c1667080f5b97c3452d568c Mon Sep 17 00:00:00 2001
From: Ryan Kingsbury <rkingsbury@users.noreply.github.com>
Date: Sun, 27 Aug 2023 21:42:22 -0400
Subject: [PATCH] JSONStore: enabled reading extended JSON

---
 src/maggma/stores/mongolike.py               | 3 ++-
 tests/stores/test_mongolike.py               | 7 +++++++
 tests/test_files/test_set/extended_json.json | 7 +++++++
 3 files changed, 16 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_files/test_set/extended_json.json

diff --git a/src/maggma/stores/mongolike.py b/src/maggma/stores/mongolike.py
index 6cd734b7b..a1e7520a3 100644
--- a/src/maggma/stores/mongolike.py
+++ b/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.
diff --git a/tests/stores/test_mongolike.py b/tests/stores/test_mongolike.py
index 68784a639..cd5e4e9c2 100644
--- a/tests/stores/test_mongolike.py
+++ b/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("."):
diff --git a/tests/test_files/test_set/extended_json.json b/tests/test_files/test_set/extended_json.json
new file mode 100644
index 000000000..971dc33fd
--- /dev/null
+++ b/tests/test_files/test_set/extended_json.json
@@ -0,0 +1,7 @@
+[{
+  "_id": {
+    "$oid": "64ebee18bd0b1265fe418be2"
+  },
+  "hello": "world",
+  "task_id": "1"
+}]