feat: use correct mongodb client lifecycle

server/app.py

@@ -7,7 +7,8 @@
 
 from server.config import settings
 from server.database.laf import laf_db_setup
-from server.database.valkey import valkey_setup
+from server.database.mongo import mongo_setup, mongo_shutdown
+from server.database.valkey import valkey_setup, valkey_shutdown
 from server.routes.auth import router as AuthRouter
 from server.routes.backtest import router as BacktestRouter
 from server.routes.laf import router as LAFRouter
@@ -44,11 +45,13 @@ def profiles_sampler(sampling_context) -> float:
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     # This runs during the startup phase
-    await laf_db_setup()
+    await mongo_setup(app)
+    await laf_db_setup(app)
     await valkey_setup(app)
     yield  # Application runs here
     # This runs during the shutdown phase
-    # Any cleanup logic can go here
+    await valkey_shutdown(app)
+    await mongo_shutdown(app)
 
 
 app = FastAPI(lifespan=lifespan, root_path=settings.ROOT_PATH)

server/database/__init__.py

@@ -1,7 +0,0 @@
-from pymongo import AsyncMongoClient
-
-from server.config import settings
-
-client = AsyncMongoClient(settings.MONGO_DETAILS)
-
-database = client.apo_main

server/database/backtest.py

@@ -1,16 +1,23 @@
 from aiocache import cached
 from bson import ObjectId
-from fastapi import HTTPException
+from fastapi import FastAPI, HTTPException, Request
 
-from server.database import database
+from server.helpers.cache import cache_key_exclude_request
 from server.helpers.sanitize import is_valid_object_id
 from server.models.backtest import Backtests, Course
 
-backtest_course_code_collection = database.get_collection(
-    "backtest_course_code_collection"
-)
-backtest_courses_collection = database.get_collection("backtest_courses_collection")
-backtest_collection = database.get_collection("backtest_collection")
+
+async def backtest_db_setup(app: FastAPI) -> None:
+    database = app.state.mongo_database
+    backtest_course_code_collection = database.get_collection(
+        "backtest_course_code_collection"
+    )
+    backtest_courses_collection = database.get_collection("backtest_courses_collection")
+    backtest_collection = database.get_collection("backtest_collection")
+
+    await backtest_course_code_collection.create_index("course_code")
+    await backtest_courses_collection.create_index("course_code")
+    await backtest_collection.create_index("course_ids")
 
 
 async def course_helper(course) -> Course:
@@ -21,25 +28,34 @@ async def course_helper(course) -> Course:
 
 
 # Retrieve all course codes present in the database
-@cached(ttl=86400)
-async def retrieve_coursecodes() -> list[str]:
+@cached(ttl=86400, key_builder=cache_key_exclude_request)
+async def retrieve_coursecodes(request: Request) -> list[str]:
+    backtest_course_code_collection = request.app.state.mongo_database.get_collection(
+        "backtest_course_code_collection"
+    )
     course_codes = await backtest_course_code_collection.distinct("course_code")
     return sorted(course_codes)
 
 
-@cached(ttl=86400)
-async def retrieve_courses(course_code: str) -> list[Course]:
+@cached(ttl=86400, key_builder=cache_key_exclude_request)
+async def retrieve_courses(request: Request, course_code: str) -> list[Course]:
+    backtest_courses_collection = request.app.state.mongo_database.get_collection(
+        "backtest_courses_collection"
+    )
     cursor = backtest_courses_collection.find({"course_code": course_code})
     courses = [await course_helper(course) async for course in cursor]
 
     return sorted(courses, key=lambda x: x["name"])
 
 
-@cached(ttl=86400)
-async def retrieve_backtest(backtest_id: str) -> list[Backtests]:
+@cached(ttl=86400, key_builder=cache_key_exclude_request)
+async def retrieve_backtest(request: Request, backtest_id: str) -> list[Backtests]:
     if not is_valid_object_id(backtest_id):
         raise HTTPException(status_code=404, detail="Backtest not found")
 
+    backtest_collection = request.app.state.mongo_database.get_collection(
+        "backtest_collection"
+    )
     backtest = await backtest_collection.find_one(
         {"course_ids": {"$in": [ObjectId(backtest_id)]}}
     )