upgrade exceptions

app/auth/auth_api.py

@@ -1,7 +1,7 @@
 from datetime import timedelta
 from typing import Annotated
 
-from fastapi import APIRouter, Depends, HTTPException, status
+from fastapi import APIRouter, Depends, status
 from fastapi.security import OAuth2PasswordRequestForm
 
 from app.auth.models import USER_SERVICE
@@ -12,7 +12,8 @@
     create_access_token,
     get_current_active_user,
 )
-from app.database.unit_of_work import unit
+from app.database.unit_of_work import unit_api
+from app.exceptions import CannotCreateStillExistsException, UnauthorizedException
 from app.settings import ACCESS_TOKEN_EXPIRE_MINUTES
 
 auth_router = APIRouter(
@@ -25,18 +26,13 @@
 def login_for_access_token(
     form_data: Annotated[OAuth2PasswordRequestForm, Depends()],
 ) -> Token:
-    with unit() as session:
+    with unit_api("Trying to authenticate user") as session:
         user = authenticate_user(session, form_data.username, form_data.password)
         if user is None:
-            raise HTTPException(
-                status_code=status.HTTP_401_UNAUTHORIZED,
-                detail="Incorrect username or password",
-                headers={"WWW-Authenticate": "Bearer"},
-            )
+            raise UnauthorizedException("Incorrect username or password")
+
         access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
-        access_token = create_access_token(
-            data={"sub": user.username}, expires_delta=access_token_expires
-        )
+        access_token = create_access_token({"sub": user.username}, access_token_expires)
 
         return Token(access_token=access_token, token_type="bearer")
 
@@ -47,25 +43,22 @@ def login_for_access_token(
 def register_user(
     form_data: Annotated[OAuth2PasswordRequestForm, Depends()],
 ) -> Token:
-    with unit() as session:
+    with unit_api("Trying to register user") as session:
         existing_user = USER_SERVICE.get_or_none(session, username=form_data.username)
 
-    if existing_user is not None:
-        raise HTTPException(
-            status_code=status.HTTP_400_BAD_REQUEST,
-            detail="Username already registered",
-        )
+        if existing_user is not None:
+            raise CannotCreateStillExistsException("Username already registered")
 
-    # Create new user
-    with unit() as session:
         new_user = User(username=form_data.username, password=form_data.password)
         new_user = USER_SERVICE.create(session, new_user)
 
         # Create and return token
         access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
         access_token = create_access_token(
-            data={"sub": new_user.username}, expires_delta=access_token_expires
+            data={"sub": new_user.username},
+            expires_delta=access_token_expires,
         )
+
     return Token(access_token=access_token, token_type="bearer")
 
 

app/database/unit_of_work.py

@@ -1,11 +1,19 @@
 import contextlib
+import logging
 
+from fastapi import HTTPException, status
 from sqlalchemy import create_engine
 from sqlmodel import Session, SQLModel
 
-from app.exceptions import ParentsListMakerException
+from app.exceptions import (
+    CannotCreateStillExistsException,
+    ParentsListMakerException,
+    UnauthorizedException,
+)
 from app.settings import DB_URL
 
+logger = logging.getLogger(__name__)
+
 # Create the database
 
 engine = create_engine(DB_URL)
@@ -20,9 +28,43 @@ def unit():
         session.commit()
     except ParentsListMakerException as e:
         session.rollback()
+        logger.exception(e)
         raise e
     except Exception as e:
         session.rollback()
+        logger.exception(e)
         raise ValueError(f"Rolling back, cause : {str(e)}") from e
     finally:
         session.close()
+
+
+@contextlib.contextmanager
+def unit_api(attempt_message: str):
+    session = Session(engine)
+    try:
+        yield session
+        session.commit()
+
+    except CannotCreateStillExistsException as e:
+        session.rollback()
+        logger.exception(e)
+        raise HTTPException(
+            status_code=status.HTTP_403_FORBIDDEN,
+            detail=f"{attempt_message} FAILED",
+        )
+    except UnauthorizedException as e:
+        session.rollback()
+        logger.exception(e)
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail=f"{attempt_message} FAILED",
+        )
+    except Exception as e:
+        session.rollback()
+        logger.exception(e)
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail=f"{attempt_message} FAILED",
+        )
+    finally:
+        session.close()

app/exceptions.py

@@ -18,3 +18,18 @@ class NotUniqueException(DatabaseException):
 
 class NotFoundException(DatabaseException):
     pass
+
+
+# API Exception
+
+
+class APIException(ParentsListMakerException):
+    pass
+
+
+class UnauthorizedException(APIException):
+    pass
+
+
+class CannotCreateStillExistsException(APIException):
+    pass