feat: add pymongo connection

.gitignore

@@ -120,7 +120,7 @@ celerybeat.pid
 *.sage.py
 
 # Environments
-.env
+# .env
 .venv
 env/
 venv/

marble-sculp/.env

@@ -0,0 +1 @@
+MONGO_DB_URI="mongodb+srv://aicostsqo:iQ38WbjVEoYLCzI1@aicostsqo.woupdve.mongodb.net/?retryWrites=true&w=majority"

marble-sculp/api.py

@@ -15,8 +15,27 @@
 from models import *
 from utils import calculate_dip_and_dip_direction_from_unit_vec
 
+from dotenv import dotenv_values
+from contextlib import asynccontextmanager
+from pymongo import MongoClient
 
-app = FastAPI()
+from typing import Callable
+from pyinstrument import Profiler
+from pyinstrument.renderers.html import HTMLRenderer
+from pyinstrument.renderers.speedscope import SpeedscopeRenderer
+
+config = dotenv_values(".env")
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    app.client = MongoClient(config["MONGO_DB_URI"])
+    app.db = app.client["test"]
+    yield
+    app.client.close()
+
+
+app = FastAPI(lifespan=lifespan)
 app.mount("/static", StaticFiles(directory="static"), name="static")
 
 app.add_middleware(
@@ -28,6 +47,57 @@
     expose_headers=["*"],
 )
 
+PROFILING = True  # Set this from a settings model
+
+
+# @app.middleware("http")
+# async def profile_request(request: Request, call_next: Callable):
+#     """Profile the current request
+
+#     Taken from https://pyinstrument.readthedocs.io/en/latest/guide.html#profile-a-web-request-in-fastapi
+#     with small improvements.
+
+#     """
+#     # we map a profile type to a file extension, as well as a pyinstrument profile renderer
+#     profile_type_to_ext = {"html": "html", "speedscope": "speedscope.json"}
+#     profile_type_to_renderer = {
+#         "html": HTMLRenderer,
+#         "speedscope": SpeedscopeRenderer,
+#     }
+
+#     # if the `profile=true` HTTP query argument is passed, we profile the request
+#     if request.query_params.get("profile", False):
+
+#         # The default profile format is speedscope
+#         profile_type = request.query_params.get("profile_format", "speedscope")
+
+#         # we profile the request along with all additional middlewares, by interrupting
+#         # the program every 1ms1 and records the entire stack at that point
+#         with Profiler(interval=0.001, async_mode="enabled") as profiler:
+#             response = await call_next(request)
+
+#         # we dump the profiling into a file
+#         extension = profile_type_to_ext[profile_type]
+#         renderer = profile_type_to_renderer[profile_type]()
+#         with open(f"profile.{extension}", "w") as out:
+#             out.write(profiler.output(renderer=renderer))
+#         return response
+
+#     # Proceed without profiling
+#     return await call_next(request)
+
+
+@app.middleware("http")
+async def add_process_time_header(request: Request, call_next):
+    start_time = time.time()
+    response = await call_next(request)
+    print(
+        "Time took to process the request and return response is {} sec".format(
+            time.time() - start_time
+        )
+    )
+    return response
+
 
 @app.get("/test")
 async def test(request: Request):