start migrating to Beanie ORM

Author: peterrauscher

pyproject.toml

@@ -13,11 +13,11 @@ requires-python = ">=3.12"
 dependencies = [
   "fastapi",
   "motor",
+  "beanie",
   "python-dotenv",
   "uvicorn",
   "pydantic",
   "mangum",
-  "Authlib",
   "requests",
   "Starlette",
   "itsdangerous",

src/api/data/instance.py

@@ -1,10 +1,23 @@
+import logging
 import os
 
-import motor.motor_asyncio
+from beanie import init_beanie
+from motor.motor_asyncio import AsyncIOMotorClient
 
-client = motor.motor_asyncio.AsyncIOMotorClient(
-    os.environ.get("SOLESEARCH_DB_CONNECTION_STRING")
-)
+logger = logging.getLogger(__name__)
+
+CONNECTION_STRING = os.environ.get("SOLESEARCH_DB_CONNECTION_STRING")
+if not CONNECTION_STRING:
+    raise EnvironmentError(
+        "Please set the SOLESEARCH_DB_CONNECTION_STRING environment variable."
+    )
+client = AsyncIOMotorClient(CONNECTION_STRING)
+DATABASE_NAME = os.environ.get("SOLESEARCH_DB_NAME")
+if not DATABASE_NAME:
+    DATABASE_NAME = "Sneakers"
+    logger.warning(
+        f"SOLESEARCH_DB_NAME environment variable not set, defaulting to {DATABASE_NAME}."
+    )
 db = client[os.environ.get("SOLESEARCH_DB_NAME")]
 sneakers = db[os.environ.get("SOLESEARCH_DB_PRIMARY_COLLECTION")]
 DEFAULT_LIMIT = int(os.environ.get("SOLESEARCH_DEFAULT_LIMIT", 10))

src/api/main.py

@@ -1,14 +1,18 @@
 import os
 
 from dotenv import load_dotenv
-from fastapi import FastAPI
-from mangum import Mangum
-from starlette.middleware.sessions import SessionMiddleware
 
 if not os.environ.get("AWS_EXECUTION_ENV"):
     load_dotenv(os.path.join(os.getcwd(), ".env"))
 
+from beanie import init_beanie
+from fastapi import FastAPI
+from mangum import Mangum
+from starlette.middleware.sessions import SessionMiddleware
+
+from api.data.instance import DATABASE_NAME, client
 from api.routes import auth, sneakers
+from core.models.shoes import Sneaker
 
 app = FastAPI(
     redoc_url=None,
@@ -17,8 +21,18 @@
 
 app.add_middleware(SessionMiddleware, secret_key="some secret key here")
 
-app.include_router(sneakers.router)
-app.include_router(auth.router)
+
+@app.on_event("startup")
+async def startup_event():
+    await init_beanie(
+        database=client.get_database(DATABASE_NAME),
+        document_models=[Sneaker],
+    )
+
+    app.include_router(sneakers.router)
+    app.include_router(auth.router)
+
+
 handler = Mangum(app)
 
 if __name__ == "__main__":

src/api/routes/sneakers.py

@@ -1,11 +1,12 @@
 from typing import Annotated
 
-from fastapi import APIRouter, Query
+from fastapi import APIRouter, HTTPException, Query
 
 from api.data.instance import DEFAULT_LIMIT, DEFAULT_OFFSET
 from api.data.models import SortKey, SortOrder
-from api.data.queries import find_sneaker_by_id, find_sneaker_by_sku, find_sneakers
+from api.data.queries import find_sneakers
 from core.models.details import Audience
+from core.models.shoes import Sneaker
 
 router = APIRouter(
     prefix="/sneakers",
@@ -43,9 +44,16 @@ async def get_sneakers(
 
 @router.get("/{product_id}")
 async def get_sneaker_by_id(product_id: str):
-    return await find_sneaker_by_id(product_id)
+    if not product_id:
+        raise HTTPException(status_code=400, detail="Invalid product_id")
+    return await Sneaker.get(product_id)
 
 
-@router.get("/sku/{product_id}")
+@router.get("/sku/{sku}")
 async def get_sneaker_by_sku(sku: str, brand: str | None = None):
-    return await find_sneaker_by_sku(sku, brand)
+    if not sku:
+        raise HTTPException(status_code=400, detail="Invalid sku")
+    if brand:
+        return await Sneaker.find_one(Sneaker.sku == sku, Sneaker.brand == brand)
+    else:
+        return await Sneaker.find_one(Sneaker.sku == sku)