Update auth.py

auth.py

@@ -7,13 +7,14 @@ from typing import List, Optional, Any
 
 from dotenv import load_dotenv
 from fastapi import APIRouter, HTTPException, Depends, Request, UploadFile, File, Form
-from fastapi.responses import JSONResponse
+from fastapi.responses import JSONResponse, StreamingResponse
 from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
 from jose import JWTError, jwt
 from passlib.context import CryptContext
 from pydantic import BaseModel, EmailStr, Field, validator
 from pymongo import MongoClient
 import os.path
+import gridfs  # For storing binary files in MongoDB
 
 load_dotenv()
 
@@ -27,6 +28,8 @@ MONGO_URL = os.getenv("CONNECTION_STRING").replace("${PASSWORD}", password)
 client = MongoClient(MONGO_URL)
 db = client.users_database
 users_collection = db.users
+# Create a GridFS instance for storing avatars in the "avatars" collection
+fs = gridfs.GridFS(db, collection="avatars")
 
 # OAuth2 setup
 oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
@@ -132,12 +135,8 @@ def get_current_user(token: str = Depends(oauth2_scheme)) -> dict:
     except JWTError:
         raise HTTPException(status_code=401, detail="Invalid token")
 
-# Setup for avatar file saving
-AVATAR_DIR = "avatars"
-if not os.path.exists(AVATAR_DIR):
-    os.makedirs(AVATAR_DIR)
-
-def save_avatar_file(file: UploadFile) -> str:
+# Instead of saving avatar to disk, store it in GridFS.
+async def save_avatar_file_to_gridfs(file: UploadFile) -> str:
     allowed_types = ["image/jpeg", "image/png", "image/gif"]
     if file.content_type not in allowed_types:
         logger.error(f"Unsupported file type: {file.content_type}")
@@ -145,20 +144,14 @@ def save_avatar_file(file: UploadFile) -> str:
             status_code=400,
             detail="Invalid image format. Only JPEG, PNG, and GIF are accepted."
         )
-    file_extension = os.path.splitext(file.filename)[1]
-    unique_filename = f"{uuid.uuid4()}{file_extension}"
-    file_path = os.path.join(AVATAR_DIR, unique_filename)
     try:
-        contents = file.file.read()
-        with open(file_path, "wb") as f:
-            f.write(contents)
-        logger.info(f"Avatar saved as {file_path}")
+        contents = await file.read()
+        file_id = fs.put(contents, filename=file.filename, contentType=file.content_type)
+        logger.info(f"Avatar stored in GridFS with file_id: {file_id}")
+        return str(file_id)
     except Exception as e:
-        logger.exception("Failed to save avatar file")
-        raise HTTPException(status_code=500, detail="Could not save avatar file.")
-    finally:
-        file.file.close()
-    return file_path
+        logger.exception("Failed to store avatar in GridFS")
+        raise HTTPException(status_code=500, detail="Could not store avatar file in MongoDB.")
 
 # ----- Auth Endpoints -----
 
@@ -186,8 +179,8 @@ async def signup(
         "chat_histories": []
     }
     if avatar:
-        avatar_path = save_avatar_file(avatar)
-        user_data["avatar"] = f"/avatars/{os.path.basename(avatar_path)}"
+        file_id = await save_avatar_file_to_gridfs(avatar)
+        user_data["avatar"] = file_id  # store the GridFS file id
     users_collection.insert_one(user_data)
     logger.info(f"New user registered: {email}")
     return {
@@ -203,20 +196,26 @@ async def login(request: Request, form_data: OAuth2PasswordRequestForm = Depends
         logger.warning(f"Failed login attempt for: {form_data.username}")
         raise HTTPException(status_code=401, detail="Incorrect username or password")
     logger.info(f"User logged in: {user['email']}")
+    avatar_url = None
+    if "avatar" in user and user["avatar"]:
+        avatar_url = f"/auth/avatar/{user['avatar']}"
     return {
         "access_token": create_access_token(user["email"]),
         "refresh_token": create_refresh_token(user["email"]),
         "token_type": "bearer",
         "name": user["name"],
-        "avatar": user.get("avatar")
+        "avatar": avatar_url
     }
 
 @router.get("/user/data")
 async def get_user_data(request: Request, current_user: dict = Depends(get_current_user)):
+    avatar_url = None
+    if "avatar" in current_user and current_user["avatar"]:
+        avatar_url = f"/auth/avatar/{current_user['avatar']}"
     return {
         "name": current_user["name"],
         "email": current_user["email"],
-        "avatar": current_user.get("avatar"),
+        "avatar": avatar_url,
         "chat_histories": current_user.get("chat_histories", [])
     }
 
@@ -242,8 +241,8 @@ async def update_user(
             raise HTTPException(status_code=400, detail=str(e))
         update_data["hashed_password"] = get_password_hash(password)
     if avatar:
-        avatar_path = save_avatar_file(avatar)
-        update_data["avatar"] = f"/avatars/{os.path.basename(avatar_path)}"
+        file_id = await save_avatar_file_to_gridfs(avatar)
+        update_data["avatar"] = file_id
     if not update_data:
         logger.info("No update parameters provided")
         raise HTTPException(status_code=400, detail="No update parameters provided")
@@ -255,3 +254,12 @@ async def update_user(
 async def logout(request: Request, current_user: dict = Depends(get_current_user)):
     logger.info(f"User logged out: {current_user['email']}")
     return {"message": "User logged out successfully"}
+
+@router.get("/avatar/{file_id}")
+async def get_avatar(file_id: str):
+    try:
+        file = fs.get(file_id)
+        return StreamingResponse(file, media_type=file.content_type)
+    except Exception as e:
+        logger.error(f"Avatar not found for file_id {file_id}: {e}")
+        raise HTTPException(status_code=404, detail="Avatar not found")