Skip to content

Commit

Permalink
MAPI and main app updates with versions
Browse files Browse the repository at this point in the history
  • Loading branch information
munrojm committed Jun 3, 2021
1 parent 3cf604a commit e2ad8b3
Show file tree
Hide file tree
Showing 2 changed files with 35 additions and 75 deletions.
22 changes: 16 additions & 6 deletions app.py
Expand Up @@ -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",
)

Expand All @@ -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"],
)

Expand Down Expand Up @@ -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

Expand Down
88 changes: 19 additions & 69 deletions 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": {
Expand All @@ -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)

0 comments on commit e2ad8b3

Please sign in to comment.