Merge pull request #870 from materialsproject/bugfix/header_processor

Ensure header processor alters the correct object
materialsproject · Oct 12, 2023 · db7e614 · db7e614
2 parents ca5afd8 + 2e769a9
commit db7e614
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 4 deletions.
diff --git a/src/maggma/api/resource/aggregation.py b/src/maggma/api/resource/aggregation.py
@@ -1,4 +1,5 @@
 from typing import Any, Dict, List, Optional, Type
+import orjson
 
 from fastapi import HTTPException, Request, Response
 from pydantic import BaseModel
@@ -11,6 +12,7 @@
 from maggma.api.resource import HeaderProcessor, Resource
 from maggma.api.resource.utils import attach_query_ops
 from maggma.api.utils import STORE_PARAMS, merge_queries
+from maggma.api.utils import serialization_helper
 from maggma.core import Store
 
 
@@ -94,9 +96,10 @@ def search(**queries: Dict[str, STORE_PARAMS]) -> Dict:
 
             meta = Meta(total_doc=count)
             response = {"data": data, "meta": {**meta.dict(), **operator_meta}}
+            response = Response(orjson.dumps(response, default=serialization_helper))  # type: ignore
 
             if self.header_processor is not None:
-                self.header_processor.process_header(temp_response, request)
+                self.header_processor.process_header(response, request)
 
             return response
 

diff --git a/src/maggma/api/resource/read_resource.py b/src/maggma/api/resource/read_resource.py
@@ -173,7 +173,7 @@ def get_by_key(
                 response = Response(orjson.dumps(response, default=serialization_helper))  # type: ignore
 
             if self.header_processor is not None:
-                self.header_processor.process_header(temp_response, request)
+                self.header_processor.process_header(response, request)
 
             return response
 
@@ -269,7 +269,7 @@ def search(**queries: Dict[str, STORE_PARAMS]) -> Union[Dict, Response]:
                 response = Response(orjson.dumps(response, default=serialization_helper))  # type: ignore
 
             if self.header_processor is not None:
-                self.header_processor.process_header(temp_response, request)
+                self.header_processor.process_header(response, request)
 
             return response
 

diff --git a/src/maggma/api/resource/s3_url.py b/src/maggma/api/resource/s3_url.py
@@ -123,7 +123,7 @@ def get_by_key(
                 response = Response(orjson.dumps(response, default=serialization_helper))  # type: ignore
 
             if self.header_processor is not None:
-                self.header_processor.process_header(temp_response, request)
+                self.header_processor.process_header(response, request)
 
             return response