-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
77 lines (61 loc) · 2.06 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
import uvicorn
import argparse
parser = argparse.ArgumentParser(description="Run the database-store")
parser.add_argument('--env', default="dev", choices=["dev", "docker"])
args = parser.parse_args()
env = args.env
from app.internal.config import PROJECT_DBNAME
from fastapi import FastAPI, status
from fastapi.middleware.cors import CORSMiddleware
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
from app.MessageQueue import main
import asyncio
import argparse
from app.routers import router
from fastapi.middleware.gzip import GZipMiddleware
import traceback
class DatasetStore(FastAPI):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
# TODO: adapt to specific origins
self.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"]
)
self.add_middleware(GZipMiddleware, minimum_size=1000)
self.include_router(
router,
prefix="/ds"
)
app = DatasetStore()
rabbitMQTask = None
@app.exception_handler(ValueError)
async def value_error_exception_handler(request: Request, exc: ValueError):
return JSONResponse(
status_code=status.HTTP_400_BAD_REQUEST,
content={"message": str(exc)},
)
@app.exception_handler(TypeError)
async def type_error_exception_handler(request: Request, exc: TypeError):
print(exc)
print(traceback.format_exc())
return JSONResponse(
status_code=status.HTTP_400_BAD_REQUEST, content={"message": "Invalid input"}
)
@app.on_event('startup')
async def startup():
global rabbitMQTask
loop = asyncio.get_event_loop()
rabbitMQTask = asyncio.create_task(main(loop))
@app.on_event('shutdown')
async def shutdown():
rabbitMQTask.cancel()
if __name__ == "__main__":
if env == "dev":
uvicorn.run("main:app", host="0.0.0.0", port=3004, reload=True)
if env == "docker":
uvicorn.run("main:app", host="0.0.0.0", port=3004, workers=20)