From e2633a01e50e986aecc472f955528e5e0ebfeb8f Mon Sep 17 00:00:00 2001
From: profitroll <vozhd.kk@gmail.com>
Date: Thu, 16 Feb 2023 15:04:28 +0100
Subject: [PATCH] WIP: Better exception handling

---
 classes/models.py        |  5 ++++-
 extensions/exceptions.py | 20 +++++++++++---------
 extensions/users.py      | 15 +++++++++++++--
 3 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/classes/models.py b/classes/models.py
index 2b04e1b..8d90c2a 100644
--- a/classes/models.py
+++ b/classes/models.py
@@ -53,4 +53,7 @@ class SearchResultsPhoto(BaseModel):
 
 class SearchResultsVideo(BaseModel):
     results: List[VideoSearch]
-    next_page: Union[str, None]
\ No newline at end of file
+    next_page: Union[str, None]
+
+class EmailConfirmed(BaseModel):
+    detail: str
\ No newline at end of file
diff --git a/extensions/exceptions.py b/extensions/exceptions.py
index 9fa38ad..86dab08 100644
--- a/extensions/exceptions.py
+++ b/extensions/exceptions.py
@@ -1,66 +1,68 @@
+from fastapi import Request		
+from fastapi.responses import UJSONResponse
 from modules.app import app
 from classes.exceptions import *
 from starlette.status import HTTP_400_BAD_REQUEST, HTTP_401_UNAUTHORIZED, HTTP_404_NOT_FOUND, HTTP_406_NOT_ACCEPTABLE, HTTP_409_CONFLICT
 
 @app.exception_handler(AlbumNotFoundError)
 async def album_not_found_exception_handler(request: Request, exc: AlbumNotFoundError):
-    return JSONResponse(
+    return UJSONResponse(
         status_code=HTTP_404_NOT_FOUND,
         content={"detail": f"Could not find album with id '{exc.id}'."},
     )
 
 @app.exception_handler(AlbumAlreadyExistsError)
 async def album_already_exists_exception_handler(request: Request, exc: AlbumAlreadyExistsError):
-    return JSONResponse(
+    return UJSONResponse(
         status_code=HTTP_409_CONFLICT,
         content={"detail": f"Album with name '{exc.name}' already exists."},
     )
 
 @app.exception_handler(AlbumIncorrectError)
 async def album_incorrect_exception_handler(request: Request, exc: AlbumIncorrectError):
-    return JSONResponse(
+    return UJSONResponse(
         status_code=HTTP_406_NOT_ACCEPTABLE,
         content={"detail": f"Album {exc.place} invalid: {exc.error}"},
     )
 
 @app.exception_handler(PhotoNotFoundError)
 async def photo_not_found_exception_handler(request: Request, exc: PhotoNotFoundError):
-    return JSONResponse(
+    return UJSONResponse(
         status_code=HTTP_404_NOT_FOUND,
         content={"detail": f"Could not find photo with id '{exc.id}'."},
     )
 
 @app.exception_handler(SearchPageInvalidError)
 async def search_page_invalid_exception_handler(request: Request, exc: SearchPageInvalidError):
-    return JSONResponse(
+    return UJSONResponse(
         status_code=HTTP_400_BAD_REQUEST,
         content={"detail": "Parameters 'page' and 'page_size' must be greater or equal to 1."},
     )
 
 @app.exception_handler(SearchTokenInvalidError)
 async def search_token_invalid_exception_handler(request: Request, exc: SearchTokenInvalidError):
-    return JSONResponse(
+    return UJSONResponse(
         status_code=HTTP_401_UNAUTHORIZED,
         content={"detail": "Parameters 'page' and 'page_size' must be greater or equal to 1."},
     )
 
 @app.exception_handler(UserEmailCodeInvalid)
 async def user_email_code_invalid_exception_handler(request: Request, exc: UserEmailCodeInvalid):
-    return JSONResponse(
+    return UJSONResponse(
         status_code=HTTP_400_BAD_REQUEST,
         content={"detail": "Confirmation code is invalid."},
     )
 
 @app.exception_handler(UserAlreadyExists)
 async def user_already_exists_exception_handler(request: Request, exc: UserAlreadyExists):
-    return JSONResponse(
+    return UJSONResponse(
         status_code=HTTP_409_CONFLICT,
         content={"detail": "User with this username already exists."},
     )
 
 @app.exception_handler(UserCredentialsInvalid)
 async def user_credentials_invalid_exception_handler(request: Request, exc: UserCredentialsInvalid):
-    return JSONResponse(
+    return UJSONResponse(
         status_code=HTTP_401_UNAUTHORIZED,
         content={"detail": "Invalid credentials."},
     )
\ No newline at end of file
diff --git a/extensions/users.py b/extensions/users.py
index 53fdfab..eb82916 100644
--- a/extensions/users.py
+++ b/extensions/users.py
@@ -1,5 +1,6 @@
 from datetime import datetime, timedelta
 from classes.exceptions import UserAlreadyExists, UserCredentialsInvalid, UserEmailCodeInvalid
+from classes.models import EmailConfirmed
 from modules.database import col_users, col_albums, col_photos, col_emails, col_videos, col_emails
 from modules.app import app
 from modules.utils import configGet, logWrite
@@ -40,11 +41,21 @@ async def user_me(current_user: User = Depends(get_current_active_user)):
 
 
 user_confirm_responses = {
+    200: {
+        "description": "Successful Response",
+        "content": {
+            "application/json": {
+                "example": {
+                    "detail": configGet("email_confirmed", "messages")
+                }
+            }
+        }
+    },
     400: UserEmailCodeInvalid().openapi
 }
 if configGet("registration_requires_confirmation") is True:
-    @app.get("/users/{user}/confirm", responses=user_confirm_responses)
-    @app.patch("/users/{user}/confirm", responses=user_confirm_responses)
+    @app.get("/users/{user}/confirm", response_class=UJSONResponse, response_model=EmailConfirmed, responses=user_confirm_responses)
+    @app.patch("/users/{user}/confirm", response_class=UJSONResponse, response_model=EmailConfirmed, responses=user_confirm_responses)
     async def user_confirm(user: str, code: str):
         confirm_record = col_emails.find_one( {"user": user, "code": code, "used": False} )
         if confirm_record is None: