From e2ad8b304aea39a553ec15e7a935a980b07c5c10 Mon Sep 17 00:00:00 2001
From: Jason Munro <jmunro@lbl.gov>
Date: Thu, 3 Jun 2021 15:44:43 -0700
Subject: [PATCH] MAPI and main app updates with versions

---
 app.py                 | 22 ++++++++---
 src/mp_api/core/api.py | 88 +++++++++---------------------------------
 2 files changed, 35 insertions(+), 75 deletions(-)

diff --git a/app.py b/app.py
index 1e875cf9..82efff20 100644
--- a/app.py
+++ b/app.py
@@ -243,7 +243,7 @@
     s3_bs_index = MongoURIStore(
         uri=f"mongodb+srv://{db_uri}",
         database="mp_core",
-        key="fs_id",
+        key="task_id",
         collection_name="s3_bandstructure_index",
     )
 
@@ -258,7 +258,7 @@
         index=s3_bs_index,
         bucket="mp-bandstructures",
         compress=True,
-        key="fs_id",
+        key="task_id",
         searchable_fields=["task_id", "fs_id"],
     )
 
@@ -493,15 +493,25 @@
 
 # Electronic Structure
 from mp_api.routes.electronic_structure.resources import (
+    dos_obj_resource,
     es_resource,
     bs_resource,
+    bs_obj_resource,
     dos_resource,
+    dos_obj_resource,
 )
 
-resources.update({"electronic_structure": [es_resource(es_store)]})
-resources.update({"bandstructure": [bs_resource(es_store, s3_bs)]})
-resources.update({"dos": [dos_resource(es_store, s3_dos)]})
-
+resources.update(
+    {
+        "electronic_structure": [
+            bs_resource(es_store),
+            dos_resource(es_store),
+            es_resource(es_store),
+            bs_obj_resource(s3_bs),
+            dos_obj_resource(s3_dos),
+        ]
+    }
+)
 # MPComplete
 from mp_api.routes.mpcomplete.resources import mpcomplete_resource
 
diff --git a/src/mp_api/core/api.py b/src/mp_api/core/api.py
index 502dc4f1..6221d601 100644
--- a/src/mp_api/core/api.py
+++ b/src/mp_api/core/api.py
@@ -1,77 +1,43 @@
-import os
-import uvicorn
-from starlette.responses import RedirectResponse
-from fastapi import FastAPI
-from typing import Dict, Union
-from datetime import datetime
-from monty.json import MSONable
-from mp_api.core.resource import ConsumerPostResource, GetResource
-from mp_api.core.settings import MAPISettings
+from typing import Dict, List
+from maggma.api.resource.core import Resource
 from pymatgen.core import __version__ as pmg_version  # type: ignore
+from mp_api import __version__ as api_version
 from fastapi.openapi.utils import get_openapi
-from fastapi.middleware.cors import CORSMiddleware
+from maggma.api.API import API
 
 
-class MAPI(MSONable):
+class MAPI(API):
     """
     Core Materials API that orchestrates resources together
-
-    TODO:
-        Build cross-resource relationships
-        Global Query Operators?
     """
 
     def __init__(
         self,
-        resources: Dict[str, Union[GetResource, ConsumerPostResource]],
+        resources: Dict[str, List[Resource]],
         title="Materials Project API",
-        version="3.0.0-dev",
+        version="v3.0",
         debug=False,
+        additional_meta={"pymatgen": pmg_version},
     ):
-        self.resources = resources
-        self.title = title
-        self.version = version
-        self.debug = debug
+        super().__init__(
+            resources=resources,
+            title=title,
+            version=version,
+            debug=debug,
+            additional_meta=additional_meta,
+        )
 
     @property
     def app(self):
         """
         App server for the cluster manager
         """
-
-        # TODO this should run on `not self.debug`!
-        on_startup = [resource.setup_indexes for resource in self.resources.values()] if self.debug else []
-        app = FastAPI(title=self.title, version=self.version, on_startup=on_startup)
-        if self.debug:
-            app.add_middleware(
-                CORSMiddleware, allow_origins=["*"], allow_methods=["GET"], allow_headers=["*"],
-            )
-
-        if len(self.resources) == 0:
-            raise RuntimeError("ERROR: There are no resources provided")
-
-        for prefix, endpoint in self.resources.items():
-            app.include_router(endpoint.router, prefix=f"/{prefix}")
-
-        @app.get("/heartbeat", include_in_schema=False)
-        def heartbeat():
-            """ API Heartbeat for Load Balancing """
-
-            return {
-                "status": "OK",
-                "time": datetime.utcnow(),
-                "api": self.version,
-                "database": os.environ.get("DB_VERSION", MAPISettings().db_version).replace("_", "."),
-                "pymatgen": pmg_version,
-            }
-
-        @app.get("/", include_in_schema=False)
-        def redirect_docs():
-            """ Redirects the root end point to the docs """
-            return RedirectResponse(url=app.docs_url, status_code=301)
+        app = super().app
 
         def custom_openapi():
-            openapi_schema = get_openapi(title=self.title, version=self.version, routes=app.routes)
+            openapi_schema = get_openapi(
+                title=self.title, version=self.version, routes=app.routes
+            )
 
             openapi_schema["components"]["securitySchemes"] = {
                 "ApiKeyAuth": {
@@ -94,19 +60,3 @@ def custom_openapi():
         app.openapi = custom_openapi
 
         return app
-
-    def run(self, ip: str = "127.0.0.1", port: int = 8000, log_level: str = "info"):
-        """
-        Runs the Cluster Manager locally
-        Meant for debugging purposes only
-
-        Args:
-            ip: Local IP to listen on
-            port: Local port to listen on
-            log_level: Logging level for the webserver
-
-        Returns:
-            None
-        """
-
-        uvicorn.run(self.app, host=ip, port=port, log_level=log_level, reload=False)