Merge pull request #569 from materialsproject/header_processing

Add header processing abilities to certain `Resource` classes

src/maggma/api/resource/__init__.py

@@ -1,6 +1,7 @@
 # isort: off
 from maggma.api.resource.core import Resource
 from maggma.api.resource.core import HintScheme
+from maggma.api.resource.core import HeaderProcessor
 
 # isort: on
 

src/maggma/api/resource/aggregation.py

@@ -1,11 +1,12 @@
 from typing import Any, Dict, List, Optional, Type
 
-from fastapi import HTTPException
+from fastapi import HTTPException, Response, Request
 from pydantic import BaseModel
 
-from maggma.api.models import Meta, Response
+from maggma.api.models import Meta
+from maggma.api.models import Response as ResponseModel
 from maggma.api.query_operator import QueryOperator
-from maggma.api.resource import Resource
+from maggma.api.resource import Resource, HeaderProcessor
 from maggma.api.resource.utils import attach_query_ops
 from maggma.api.utils import STORE_PARAMS, merge_queries
 from maggma.core import Store
@@ -24,6 +25,7 @@ def __init__(
         tags: Optional[List[str]] = None,
         include_in_schema: Optional[bool] = True,
         sub_path: Optional[str] = "/",
+        header_processor: Optional[HeaderProcessor] = None,
     ):
         """
         Args:
@@ -39,9 +41,10 @@ def __init__(
 
         self.include_in_schema = include_in_schema
         self.sub_path = sub_path
-        self.response_model = Response[model]  # type: ignore
+        self.response_model = ResponseModel[model]  # type: ignore
 
         self.pipeline_query_operator = pipeline_query_operator
+        self.header_processor = header_processor
 
         super().__init__(model)
 
@@ -58,6 +61,8 @@ def build_dynamic_model_search(self):
         model_name = self.model.__name__
 
         async def search(**queries: Dict[str, STORE_PARAMS]) -> Dict:
+            request: Request = queries.pop("request")  # type: ignore
+            temp_response: Response = queries.pop("temp_response")  # type: ignore
 
             query: Dict[Any, Any] = merge_queries(list(queries.values()))  # type: ignore
 
@@ -78,6 +83,10 @@ async def search(**queries: Dict[str, STORE_PARAMS]) -> Dict:
 
             meta = Meta(total_doc=count)
             response = {"data": data, "meta": {**meta.dict(), **operator_meta}}
+
+            if self.header_processor is not None:
+                self.header_processor.process_header(temp_response, request)
+
             return response
 
         self.router.get(

src/maggma/api/resource/core.py

@@ -2,7 +2,7 @@
 from abc import ABCMeta, abstractmethod
 from typing import Dict, Type
 
-from fastapi import APIRouter, FastAPI
+from fastapi import APIRouter, FastAPI, Response, Request
 from monty.json import MontyDecoder, MSONable
 from pydantic import BaseModel
 from starlette.responses import RedirectResponse
@@ -97,3 +97,17 @@ def generate_hints(self, query: STORE_PARAMS) -> STORE_PARAMS:
         """
         This method takes in a MongoDB query and returns hints
         """
+
+
+class HeaderProcessor(MSONable, metaclass=ABCMeta):
+    """
+    Base class for generic header processing
+    """
+
+    @abstractmethod
+    def process_header(self, response: Response, request: Request):
+        """
+        This method takes in a FastAPI Response object and processes a new header for it in-place.
+        It can use data in the upstream request to generate the header.
+        (https://fastapi.tiangolo.com/advanced/response-headers/#use-a-response-parameter)
+        """

src/maggma/api/resource/post_resource.py

@@ -54,7 +54,10 @@ def __init__(
             if query_operators is not None
             else [
                 PaginationQuery(),
-                SparseFieldsQuery(model, default_fields=[self.store.key, self.store.last_updated_field],),
+                SparseFieldsQuery(
+                    model,
+                    default_fields=[self.store.key, self.store.last_updated_field],
+                ),
             ]
         )
 
@@ -74,16 +77,23 @@ def build_dynamic_model_search(self):
 
         async def search(**queries: Dict[str, STORE_PARAMS]) -> Dict:
             request: Request = queries.pop("request")  # type: ignore
+            queries.pop("temp_response")  # type: ignore
 
             query_params = [
-                entry for _, i in enumerate(self.query_operators) for entry in signature(i.query).parameters
+                entry
+                for _, i in enumerate(self.query_operators)
+                for entry in signature(i.query).parameters
             ]
 
-            overlap = [key for key in request.query_params.keys() if key not in query_params]
+            overlap = [
+                key for key in request.query_params.keys() if key not in query_params
+            ]
             if any(overlap):
                 raise HTTPException(
                     status_code=400,
-                    detail="Request contains query parameters which cannot be used: {}".format(", ".join(overlap)),
+                    detail="Request contains query parameters which cannot be used: {}".format(
+                        ", ".join(overlap)
+                    ),
                 )
 
             query: Dict[Any, Any] = merge_queries(list(queries.values()))  # type: ignore

src/maggma/api/resource/read_resource.py

@@ -1,13 +1,14 @@
 from inspect import signature
-from typing import Any, Dict, List, Optional, Type
+from typing import Any, Dict, List, Optional, Type, Union
 
 from fastapi import Depends, HTTPException, Path, Request
-from fastapi import Response as BareResponse
+from fastapi import Response
 from pydantic import BaseModel
 
-from maggma.api.models import Meta, Response
+from maggma.api.models import Meta
+from maggma.api.models import Response as ResponseModel
 from maggma.api.query_operator import PaginationQuery, QueryOperator, SparseFieldsQuery
-from maggma.api.resource import Resource, HintScheme
+from maggma.api.resource import Resource, HintScheme, HeaderProcessor
 from maggma.api.resource.utils import attach_query_ops
 from maggma.api.utils import STORE_PARAMS, merge_queries, object_id_serilaization_helper
 from maggma.core import Store
@@ -31,6 +32,7 @@ def __init__(
         query_operators: Optional[List[QueryOperator]] = None,
         key_fields: Optional[List[str]] = None,
         hint_scheme: Optional[HintScheme] = None,
+        header_processor: Optional[HeaderProcessor] = None,
         enable_get_by_key: bool = True,
         enable_default_search: bool = True,
         disable_validation: bool = False,
@@ -44,6 +46,7 @@ def __init__(
             tags: List of tags for the Endpoint
             query_operators: Operators for the query language
             hint_scheme: The hint scheme to use for this resource
+            header_processor: The header processor to use for this resource
             key_fields: List of fields to always project. Default uses SparseFieldsQuery
                 to allow user to define these on-the-fly.
             enable_get_by_key: Enable default key route for endpoint.
@@ -57,14 +60,16 @@ def __init__(
         self.store = store
         self.tags = tags or []
         self.hint_scheme = hint_scheme
+        self.header_processor = header_processor
         self.key_fields = key_fields
         self.versioned = False
         self.enable_get_by_key = enable_get_by_key
         self.enable_default_search = enable_default_search
         self.disable_validation = disable_validation
         self.include_in_schema = include_in_schema
         self.sub_path = sub_path
-        self.response_model = Response[model]  # type: ignore
+
+        self.response_model = ResponseModel[model]  # type: ignore
 
         if not isinstance(store, MongoStore) and self.hint_scheme is not None:
             raise ValueError("Hint scheme is only supported for MongoDB stores")
@@ -109,6 +114,8 @@ def field_input():
                 return {"properties": self.key_fields}
 
         async def get_by_key(
+            request: Request,
+            response: Response,
             key: str = Path(
                 ..., alias=key_name, title=f"The {key_name} of the {model_name} to get",
             ),
@@ -140,13 +147,16 @@ async def get_by_key(
             for operator in self.query_operators:
                 item = operator.post_process(item)
 
-            response = {"data": item}
+            response = {"data": item}  # type: ignore
 
             if self.disable_validation:
-                response = BareResponse(  # type: ignore
+                response = Response(  # type: ignore
                     orjson.dumps(response, default=object_id_serilaization_helper)
                 )
 
+            if self.header_processor is not None:
+                self.header_processor.process_header(response, request)
+
             return response
 
         self.router.get(
@@ -163,8 +173,9 @@ def build_dynamic_model_search(self):
 
         model_name = self.model.__name__
 
-        async def search(**queries: Dict[str, STORE_PARAMS]) -> Dict:
+        async def search(**queries: Dict[str, STORE_PARAMS]) -> Union[Dict, Response]:
             request: Request = queries.pop("request")  # type: ignore
+            response: Response = queries.pop("temp_response")  # type: ignore
 
             query_params = [
                 entry
@@ -208,13 +219,16 @@ async def search(**queries: Dict[str, STORE_PARAMS]) -> Dict:
 
             meta = Meta(total_doc=count)
 
-            response = {"data": data, "meta": {**meta.dict(), **operator_meta}}
+            response = {"data": data, "meta": {**meta.dict(), **operator_meta}}  # type: ignore
 
             if self.disable_validation:
-                response = BareResponse(  # type: ignore
+                response = Response(  # type: ignore
                     orjson.dumps(response, default=object_id_serilaization_helper)
                 )
 
+            if self.header_processor is not None:
+                self.header_processor.process_header(response, request)
+
             return response
 
         self.router.get(

src/maggma/api/resource/submission.py

@@ -175,6 +175,7 @@ def build_search_data(self):
         async def search(**queries: STORE_PARAMS):
 
             request: Request = queries.pop("request")  # type: ignore
+            queries.pop("temp_response")  # type: ignore
 
             query: STORE_PARAMS = merge_queries(list(queries.values()))
 
@@ -224,6 +225,7 @@ def build_post_data(self):
         async def post_data(**queries: STORE_PARAMS):
 
             request: Request = queries.pop("request")  # type: ignore
+            queries.pop("temp_response")  # type: ignore
 
             query: STORE_PARAMS = merge_queries(list(queries.values()))
 
@@ -274,8 +276,7 @@ async def post_data(**queries: STORE_PARAMS):
                 self.store.update(docs=query["criteria"])  # type: ignore
             except Exception:
                 raise HTTPException(
-                    status_code=400,
-                    detail="Problem when trying to post data.",
+                    status_code=400, detail="Problem when trying to post data.",
                 )
 
             response = {

src/maggma/api/resource/utils.py

@@ -1,6 +1,6 @@
 from typing import Callable, Dict, List
 
-from fastapi import Depends, Request
+from fastapi import Depends, Request, Response
 
 from maggma.api.query_operator import QueryOperator
 from maggma.api.utils import STORE_PARAMS, attach_signature
@@ -21,6 +21,7 @@ def attach_query_ops(
         annotations={
             **{f"dep{i}": STORE_PARAMS for i, _ in enumerate(query_ops)},
             "request": Request,
+            "temp_response": Response,
         },
         defaults={f"dep{i}": Depends(dep.query) for i, dep in enumerate(query_ops)},
     )