Merge pull request #464 from munrojm/switch_to_orjson

Switch from monty to orjson for serialization

requirements.txt

@@ -16,3 +16,4 @@ uvicorn==0.13.4
 sshtunnel==0.4.0
 msgpack==1.0.2
 msgpack-python==0.5.6
+orjson==3.6.0

src/maggma/api/resource/read_resource.py

@@ -9,10 +9,10 @@
 from maggma.api.query_operator import PaginationQuery, QueryOperator, SparseFieldsQuery
 from maggma.api.resource import Resource
 from maggma.api.resource.utils import attach_query_ops
-from maggma.api.utils import STORE_PARAMS, merge_queries
+from maggma.api.utils import STORE_PARAMS, merge_queries, object_id_serilaization_helper
 from maggma.core import Store
 
-from monty.serialization import MontyEncoder
+import orjson
 
 
 class ReadOnlyResource(Resource):
@@ -32,7 +32,7 @@ def __init__(
         query: Optional[Dict] = None,
         enable_get_by_key: bool = True,
         enable_default_search: bool = True,
-        monty_encoded_response: bool = False,
+        disable_validation: bool = False,
         include_in_schema: Optional[bool] = True,
         sub_path: Optional[str] = "/",
     ):
@@ -46,7 +46,7 @@ def __init__(
                 to allow user to define these on-the-fly.
             enable_get_by_key: Enable default key route for endpoint.
             enable_default_search: Enable default endpoint search behavior.
-            monty_encoded_response: Whether to use MontyEncoder and provide a direct FastAPI response.
+            disable_validation: Whether to use ORJSON and provide a direct FastAPI response.
                 Note this will disable auto JSON serialization and response validation with the
                 provided model.
             include_in_schema: Whether the endpoint should be shown in the documented schema.
@@ -59,7 +59,7 @@ def __init__(
         self.versioned = False
         self.enable_get_by_key = enable_get_by_key
         self.enable_default_search = enable_default_search
-        self.monty_encoded_response = monty_encoded_response
+        self.disable_validation = disable_validation
         self.include_in_schema = include_in_schema
         self.sub_path = sub_path
         self.response_model = Response[model]  # type: ignore
@@ -138,9 +138,9 @@ async def get_by_key(
 
             response = {"data": item}
 
-            if self.monty_encoded_response:
+            if self.disable_validation:
                 response = BareResponse(  # type: ignore
-                    MontyEncoder().encode(response)
+                    orjson.dumps(response, default=object_id_serilaization_helper)
                 )
 
             return response
@@ -195,9 +195,9 @@ async def search(**queries: Dict[str, STORE_PARAMS]) -> Dict:
 
             response = {"data": data, "meta": {**meta.dict(), **operator_meta}}
 
-            if self.monty_encoded_response:
+            if self.disable_validation:
                 response = BareResponse(  # type: ignore
-                    MontyEncoder().encode(response)
+                    orjson.dumps(response, default=object_id_serilaization_helper)
                 )
 
             return response

src/maggma/api/utils.py

@@ -1,6 +1,7 @@
 import inspect
 import sys
 from typing import Any, Callable, Dict, List, Optional, Type
+from bson.objectid import ObjectId
 
 from monty.json import MSONable
 from pydantic import BaseModel
@@ -161,3 +162,9 @@ def validate_monty(cls, v):
     setattr(monty_cls, "validate_monty", classmethod(validate_monty))
 
     return monty_cls
+
+
+def object_id_serilaization_helper(obj):
+    if isinstance(obj, ObjectId):
+        return str(obj)
+    raise TypeError

tests/api/test_read_resource.py

@@ -69,7 +69,7 @@ def test_msonable(owner_store):
 
 
 def test_get_by_key(owner_store):
-    endpoint = ReadOnlyResource(owner_store, Owner, monty_encoded_response=True)
+    endpoint = ReadOnlyResource(owner_store, Owner, disable_validation=True)
     app = FastAPI()
     app.include_router(endpoint.router)
 
@@ -128,7 +128,7 @@ def search_helper(payload, base: str = "/?", debug=True) -> Response:
             NumericQuery(model=Owner),
             SparseFieldsQuery(model=Owner),
         ],
-        monty_encoded_response=True,
+        disable_validation=True,
     )
     app = FastAPI()
     app.include_router(endpoint.router)

tests/api/test_utils.py

@@ -5,9 +5,11 @@
 from monty.json import MSONable
 from pydantic import BaseModel, Field
 
-from maggma.api.utils import api_sanitize
+from maggma.api.utils import api_sanitize, object_id_serilaization_helper
 from typing import Union
 
+from bson import ObjectId
+
 
 class SomeEnum(Enum):
     A = 1
@@ -89,3 +91,14 @@ def test_api_sanitize():
     temp_pet_dict = AnotherPet(name="fido", age=3).as_dict()
 
     assert isinstance(AnotherPet.validate_monty(temp_pet_dict), dict)
+
+
+def test_object_id_serilaization_helper():
+    oid = ObjectId("60b7d47bb671aa7b01a2adf6")
+    assert object_id_serilaization_helper(oid) == "60b7d47bb671aa7b01a2adf6"
+
+
+@pytest.mark.xfail
+def test_object_id_serilaization_helper_xfail():
+    oid = "test"
+    object_id_serilaization_helper(oid)