Merge pull request #448 from munrojm/main

Maggma API additions

src/maggma/api/API.py

@@ -1,10 +1,11 @@
 from datetime import datetime
-from typing import Dict
+from typing import Dict, List, Optional
 
 import uvicorn
 from fastapi import FastAPI
 from monty.json import MSONable
 from starlette.responses import RedirectResponse
+from fastapi.middleware.cors import CORSMiddleware
 
 from maggma import __version__
 from maggma.api.resource import Resource
@@ -17,21 +18,24 @@ class API(MSONable):
 
     def __init__(
         self,
-        resources: Dict[str, Resource],
-        title="Generic API",
-        version="v0.0.0",
-        debug=False,
+        resources: Dict[str, List[Resource]],
+        title: str = "Generic API",
+        version: str = "v0.0.0",
+        debug: bool = False,
+        heartbeat_meta: Optional[Dict] = None,
     ):
         """
         Args:
             resources: dictionary of resource objects and http prefix they live in
             title: a string title for this API
             version: the version for this API
             debug: turns debug on in FastAPI
+            heartbeat_meta: dictionary of additional metadata to include in the heartbeat response
         """
         self.title = title
         self.version = version
         self.debug = debug
+        self.heartbeat_meta = heartbeat_meta
 
         if len(resources) == 0:
             raise RuntimeError("ERROR: There are no endpoints provided")
@@ -42,8 +46,9 @@ def on_startup(self):
         """
         Basic startup that runs the resource startup functions
         """
-        for resource in self.resources.values():
-            resource.on_startup()
+        for resource_list in self.resources.values():
+            for resource in resource_list:
+                resource.on_startup()
 
     @property
     def app(self):
@@ -56,14 +61,35 @@ def app(self):
             on_startup=[self.on_startup],
             debug=self.debug,
         )
-        for prefix, resource in self.resources.items():
-            app.include_router(resource.router, prefix=f"/{prefix}")
+
+        # Allow requests from other domains in debug mode. This allows
+        # testing with local deployments of other services. For production
+        # deployment, this will be taken care of by nginx.
+        if self.debug:
+            app.add_middleware(
+                CORSMiddleware,
+                allow_origins=["*"],
+                allow_methods=["GET"],
+                allow_headers=["*"],
+            )
+
+        for prefix, resource_list in self.resources.items():
+            main_resource = resource_list.pop(0)
+            for resource in resource_list:
+                main_resource.router.include_router(resource.router)
+
+            app.include_router(main_resource.router, prefix=f"/{prefix}")
 
         @app.get("/heartbeat", include_in_schema=False)
         def heartbeat():
             """ API Heartbeat for Load Balancing """
 
-            return {"status": "OK", "time": datetime.utcnow()}
+            return {
+                "status": "OK",
+                "time": datetime.utcnow(),
+                "version": self.version,
+                **self.heartbeat_meta,
+            }
 
         @app.get("/", include_in_schema=False)
         def redirect_docs():

src/maggma/api/query_operator/__init__.py

@@ -2,3 +2,5 @@
 from maggma.api.query_operator.dynamic import NumericQuery, StringQueryOperator
 from maggma.api.query_operator.pagination import PaginationQuery
 from maggma.api.query_operator.sparse_fields import SparseFieldsQuery
+from maggma.api.query_operator.sorting import SortQuery
+from maggma.api.query_operator.submission import SubmissionQuery

src/maggma/api/query_operator/core.py

@@ -1,5 +1,5 @@
 from abc import ABCMeta, abstractmethod
-from typing import Dict
+from typing import Dict, List
 
 from monty.json import MSONable
 
@@ -28,8 +28,8 @@ def meta(self) -> Dict:
         """
         return {}
 
-    def post_process(self, doc: Dict) -> Dict:
+    def post_process(self, docs: List[Dict]) -> List[Dict]:
         """
         An optional post-processing function for the data
         """
-        return doc
+        return docs

src/maggma/api/query_operator/dynamic.py

@@ -1,4 +1,5 @@
 import inspect
+from typing import Type
 from abc import abstractmethod
 from typing import Any, Callable, Dict, List, Optional, Tuple
 
@@ -17,7 +18,7 @@ class DynamicQueryOperator(QueryOperator):
 
     def __init__(
         self,
-        model: BaseModel,
+        model: Type[BaseModel],
         fields: Optional[List[str]] = None,
         excluded_fields: Optional[List[str]] = None,
     ):
@@ -54,7 +55,15 @@ def query(**kwargs) -> STORE_PARAMS:
                             f"Cannot find key {k} in current query to database mapping"
                         )
 
-            return {"criteria": {k: v for crit in criteria for k, v in crit.items()}}
+            final_crit = {}
+            for entry in criteria:
+                for key, value in entry.items():
+                    if key not in final_crit:
+                        final_crit[key] = value
+                    else:
+                        final_crit[key].update(value)
+
+            return {"criteria": final_crit}
 
         # building the signatures for FastAPI Swagger UI
         signatures: List = [
@@ -128,22 +137,20 @@ def field_to_operator(
 
             ops = [
                 (
-                    f"{field.name}_lt",
+                    f"{field.name}_max",
                     field_type,
                     Query(
-                        default=None,
-                        description=f"Querying for {title} is less than a value",
+                        default=None, description=f"Query for maximum value of {title}",
                     ),
-                    lambda val: {f"{field.name}": {"$lt": val}},
+                    lambda val: {f"{field.name}": {"$lte": val}},
                 ),
                 (
-                    f"{field.name}_gt",
+                    f"{field.name}_min",
                     field_type,
                     Query(
-                        default=None,
-                        description=f"Querying for {title} is greater than a value",
+                        default=None, description=f"Query for minimum value of {title}",
                     ),
-                    lambda val: {f"{field.name}": {"$gt": val}},
+                    lambda val: {f"{field.name}": {"$gte": val}},
                 ),
             ]
 
@@ -155,7 +162,7 @@ def field_to_operator(
                         field_type,
                         Query(
                             default=None,
-                            description=f"Querying for {title} is equal to a value",
+                            description=f"Query for {title} being equal to an exact value",
                         ),
                         lambda val: {f"{field.name}": val},
                     ),
@@ -164,27 +171,36 @@ def field_to_operator(
                         field_type,
                         Query(
                             default=None,
-                            description=f"Querying for {title} is not equal to a value",
+                            description=f"Query for {title} being not equal to an exact value",
                         ),
                         lambda val: {f"{field.name}": {"$ne": val}},
                     ),
                     (
                         f"{field.name}_eq_any",
-                        List[field_type],  # type: ignore
+                        str,  # type: ignore
                         Query(
                             default=None,
-                            description=f"Querying for {title} is any of these values",
+                            description=f"Query for {title} being any of these values. Provide a comma separated list.",
                         ),
-                        lambda val: {f"{field.name}": {"$in": val}},
+                        lambda val: {
+                            f"{field.name}": {
+                                "$in": [int(entry.strip()) for entry in val.split(",")]
+                            }
+                        },
                     ),
                     (
                         f"{field.name}_neq_any",
-                        List[field_type],  # type: ignore
+                        str,  # type: ignore
                         Query(
                             default=None,
-                            description=f"Querying for {title} is not any of these values",
+                            description=f"Query for {title} being not any of these values. \
+                            Provide a comma separated list.",
                         ),
-                        lambda val: {f"{field.name}": {"$nin": val}},
+                        lambda val: {
+                            f"{field.name}": {
+                                "$nin": [int(entry.strip()) for entry in val.split(",")]
+                            }
+                        },
                     ),
                 ]
             )
@@ -219,7 +235,7 @@ def field_to_operator(
                     field_type,
                     Query(
                         default=None,
-                        description=f"Query for {title} is equal to a value",
+                        description=f"Query for {title} being equal to a value",
                     ),
                     lambda val: {f"{field.name}": val},
                 ),
@@ -228,27 +244,35 @@ def field_to_operator(
                     field_type,
                     Query(
                         default=None,
-                        description=f"Querying for {title} is not equal to a value",
+                        description=f"Query for {title} being not equal to a value",
                     ),
                     lambda val: {f"{field.name}": {"$ne": val}},
                 ),
                 (
                     f"{field.name}_eq_any",
-                    List[field_type],  # type: ignore
+                    str,  # type: ignore
                     Query(
                         default=None,
-                        description=f"Querying for {title} is any of these values",
+                        description=f"Query for {title} being any of these values. Provide a comma separated list.",
                     ),
-                    lambda val: {f"{field.name}": {"$in": val}},
+                    lambda val: {
+                        f"{field.name}": {
+                            "$in": [entry.strip() for entry in val.split(",")]
+                        }
+                    },
                 ),
                 (
                     f"{field.name}_neq_any",
-                    List[field_type],  # type: ignore
+                    str,  # type: ignore
                     Query(
                         default=None,
-                        description=f"Querying for {title} is not any of these values",
+                        description=f"Query for {title} being not any of these values. Provide a comma separated list",
                     ),
-                    lambda val: {f"{field.name}": {"$nin": val}},
+                    lambda val: {
+                        f"{field.name}": {
+                            "$nin": [entry.strip() for entry in val.split(",")]
+                        }
+                    },
                 ),
             ]
 

src/maggma/api/query_operator/pagination.py

@@ -9,9 +9,7 @@
 class PaginationQuery(QueryOperator):
     """Query opertators to provides Pagination"""
 
-    def __init__(
-        self, default_skip: int = 0, default_limit: int = 10, max_limit: int = 100
-    ):
+    def __init__(self, default_skip: int = 0, default_limit: int = 100, max_limit: int = 1000):
         """
         Args:
             default_skip: the default number of documents to skip
@@ -23,13 +21,10 @@ def __init__(
         self.max_limit = max_limit
 
         def query(
-            skip: int = Query(
-                default_skip, description="Number of entries to skip in the search"
-            ),
+            skip: int = Query(default_skip, description="Number of entries to skip in the search"),
             limit: int = Query(
                 default_limit,
-                description="Max number of entries to return in a single query."
-                f" Limited to {max_limit}",
+                description="Max number of entries to return in a single query." f" Limited to {max_limit}",
             ),
         ) -> STORE_PARAMS:
             """
@@ -39,7 +34,7 @@ def query(
                 raise HTTPException(
                     status_code=400,
                     detail="Requested more data per query than allowed by this endpoint."
-                    f"The max limit is {max_limit} entries",
+                    f" The max limit is {max_limit} entries",
                 )
             return {"skip": skip, "limit": limit}
 

src/maggma/api/query_operator/sorting.py

@@ -0,0 +1,28 @@
+from typing import Optional
+from maggma.api.query_operator import QueryOperator
+from maggma.api.utils import STORE_PARAMS
+from fastapi import HTTPException, Query
+
+
+class SortQuery(QueryOperator):
+    """
+    Method to generate the sorting portion of a query
+    """
+
+    def query(
+        self,
+        field: Optional[str] = Query(None, description="Field to sort with"),
+        ascending: Optional[bool] = Query(None, description="Whether the sorting should be ascending",),
+    ) -> STORE_PARAMS:
+
+        sort = {}
+
+        if field and ascending is not None:
+            sort.update({field: 1 if ascending else -1})
+
+        elif field or ascending is not None:
+            raise HTTPException(
+                status_code=400, detail="Must specify both a field and order for sorting.",
+            )
+
+        return {"sort": sort}

src/maggma/api/query_operator/submission.py

@@ -0,0 +1,49 @@
+from typing import Optional
+from maggma.api.query_operator import QueryOperator
+from maggma.api.utils import STORE_PARAMS
+from fastapi import Query
+from datetime import datetime
+
+
+class SubmissionQuery(QueryOperator):
+    """
+    Method to generate a query for submission data using status and datetime
+    """
+
+    def __init__(self, status_enum):
+
+        self.status_enum = status_enum
+
+        def query(
+            state: Optional[status_enum] = Query(
+                None, description="Latest status of the submission"
+            ),
+            last_updated: Optional[datetime] = Query(
+                None, description="Minimum datetime of status update for submission",
+            ),
+        ) -> STORE_PARAMS:
+
+            crit = {}  # type: dict
+
+            if state:
+                s_dict = {"$expr": {"$eq": [{"$arrayElemAt": ["$state", -1]}, state.value]}}  # type: ignore
+                crit.update(s_dict)
+
+            if last_updated:
+                l_dict = {
+                    "$expr": {
+                        "$gt": [{"$arrayElemAt": ["$last_updated", -1]}, last_updated]
+                    }
+                }
+                crit.update(l_dict)
+
+            if state and last_updated:
+                crit = {"$and": [s_dict, l_dict]}
+
+            return {"criteria": crit}
+
+        self.query = query
+
+    def query(self):
+        " Stub query function for abstract class "
+        pass