Extract 401 error from service layer

Author: koldakov

futuramaapi/apps/fastapi.py

@@ -1,10 +1,11 @@
 import mimetypes
 from collections.abc import AsyncGenerator
 from contextlib import asynccontextmanager
-from typing import TYPE_CHECKING, Any, ClassVar, Literal, Self
+from typing import TYPE_CHECKING, Any, ClassVar, Literal, NamedTuple, Self
 
 import sentry_sdk
-from fastapi import FastAPI
+from fastapi import FastAPI, Request, Response, status
+from fastapi.responses import JSONResponse
 from fastapi.staticfiles import StaticFiles
 from fastapi_pagination import add_pagination
 from starlette.routing import Host, Mount, Route, WebSocketRoute
@@ -23,6 +24,11 @@
 mimetypes.add_type("image/webp", ".webp")
 
 
+class _ExceptionValue(NamedTuple):
+    status_code: int
+    default_message: str
+
+
 class FuturamaAPI(FastAPI):
     BOTS_FORBIDDEN_URLS: ClassVar[tuple[str, ...]] = (
         "/favicon.ico",
@@ -92,12 +98,36 @@ def _setup_static(self) -> None:
             name="static",
         )
 
+    def _exception_handler(self, _: Request, exc) -> Response:
+        from futuramaapi.routers.services import ServiceError, UnauthorizedError  # noqa: PLC0415
+
+        exception_to_value: dict[type[ServiceError], _ExceptionValue] = {
+            UnauthorizedError: _ExceptionValue(
+                status_code=status.HTTP_401_UNAUTHORIZED,
+                default_message="Unauthorized",
+            ),
+        }
+
+        exc_value = exception_to_value[type(exc)]
+        return JSONResponse(
+            status_code=exc_value.status_code,
+            content={
+                "detail": str(exc) or exc_value.default_message,
+            },
+        )
+
+    def _setup_exceptions(self) -> None:
+        from futuramaapi.routers.services import ServiceError  # noqa: PLC0415
+
+        self.add_exception_handler(ServiceError, self._exception_handler)
+
     def setup(self) -> None:
         super().setup()
 
         self._setup_middlewares()
         self._setup_routers()
         self._setup_static()
+        self._setup_exceptions()
 
         add_pagination(self)
 

futuramaapi/routers/services/__init__.py

@@ -2,6 +2,8 @@
     BaseService,
     BaseSessionService,
     BaseUserAuthenticatedService,
+    ServiceError,
+    UnauthorizedError,
 )
 from ._base_template import BaseTemplateService
 
@@ -10,4 +12,6 @@
     "BaseSessionService",
     "BaseTemplateService",
     "BaseUserAuthenticatedService",
+    "ServiceError",
+    "UnauthorizedError",
 ]

futuramaapi/routers/services/_base.py

@@ -3,7 +3,6 @@
 from typing import Any, TypeVar
 
 import jwt
-from fastapi import HTTPException, status
 from fastapi_pagination import Page
 from jwt import ExpiredSignatureError, InvalidSignatureError, InvalidTokenError
 from sqlalchemy import Select, select
@@ -21,6 +20,14 @@
 )
 
 
+class ServiceError(Exception):
+    """Service Error."""
+
+
+class UnauthorizedError(ServiceError):
+    """Unauthorized Error."""
+
+
 class BaseService[TResponse](BaseModel, ABC):
     context: dict[str, Any] | None = None
 
@@ -79,10 +86,10 @@ def __get_decoded_token(
                 algorithms=[algorithm],
             )
         except (ExpiredSignatureError, InvalidSignatureError, InvalidTokenError):
-            raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED) from None
+            raise UnauthorizedError() from None
 
         if decoded_token["type"] != "access":
-            raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED) from None
+            raise UnauthorizedError() from None
 
         return decoded_token
 
@@ -95,7 +102,7 @@ async def __set_user(self) -> None:
         try:
             self._user = (await self.session.execute(self.__get_user_statement)).scalars().one()
         except NoResultFound:
-            raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED) from None
+            raise UnauthorizedError() from None
 
     async def __call__(self, *args, **kwargs) -> TResponse:
         async with session_manager.session() as session:

futuramaapi/routers/services/tokens/get_auth_user_token.py

@@ -3,15 +3,14 @@
 from typing import Any, ClassVar, Self
 
 import jwt
-from fastapi import HTTPException, status
 from pydantic import Field, SecretStr
 from sqlalchemy import Result, Select, select
 from sqlalchemy.exc import NoResultFound
 
 from futuramaapi.core import settings
 from futuramaapi.helpers.pydantic import BaseModel
 from futuramaapi.repositories.models import UserModel
-from futuramaapi.routers.services import BaseSessionService
+from futuramaapi.routers.services import BaseSessionService, UnauthorizedError
 
 
 class GetAuthUserTokenResponse(BaseModel):
@@ -76,13 +75,11 @@ async def _get_user(self) -> UserModel:
         try:
             return result.scalars().one()
         except NoResultFound:
-            raise HTTPException(
-                status_code=status.HTTP_401_UNAUTHORIZED,
-            ) from None
+            raise UnauthorizedError() from None
 
     async def process(self, *args, **kwargs) -> GetAuthUserTokenResponse:
         user: UserModel = await self._get_user()
         if not self.hasher.verify(self.password.get_secret_value(), user.password):
-            raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED)
+            raise UnauthorizedError()
 
         return GetAuthUserTokenResponse.from_user_model(user)

futuramaapi/routers/services/tokens/get_refreshed_auth_user_token.py

@@ -1,7 +1,6 @@
 from typing import Any, ClassVar
 
 import jwt
-from fastapi import HTTPException, status
 from jwt import ExpiredSignatureError, InvalidSignatureError, InvalidTokenError
 from pydantic import Field
 from sqlalchemy import Result, Select, select
@@ -10,7 +9,7 @@
 from futuramaapi.core import settings
 from futuramaapi.helpers.pydantic import BaseModel
 from futuramaapi.repositories.models import UserModel
-from futuramaapi.routers.services import BaseSessionService
+from futuramaapi.routers.services import BaseSessionService, UnauthorizedError
 
 from .get_auth_user_token import GetAuthUserTokenResponse
 
@@ -38,7 +37,7 @@ async def _get_user(self, decoded_token: dict[str, Any], /) -> UserModel:
         try:
             user: UserModel = result.scalars().one()
         except NoResultFound:
-            raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED) from None
+            raise UnauthorizedError() from None
 
         return user
 
@@ -50,10 +49,10 @@ async def process(self, *args, **kwargs) -> GetRefreshedAuthUserTokenResponse:
                 algorithms=[self.algorithm],
             )
         except (ExpiredSignatureError, InvalidSignatureError, InvalidTokenError):
-            raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED) from None
+            raise UnauthorizedError() from None
 
         if decoded_token["type"] != "refresh":
-            raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED)
+            raise UnauthorizedError()
 
         user: UserModel = await self._get_user(decoded_token)
 

futuramaapi/routers/services/users/activate_signature_user.py

@@ -1,14 +1,14 @@
 from typing import Any, ClassVar
 
 import jwt
-from fastapi import HTTPException, Request, status
+from fastapi import Request, status
 from jwt import ExpiredSignatureError, InvalidSignatureError, InvalidTokenError
 from sqlalchemy import Result, Select, Update, select, update
 from starlette.responses import RedirectResponse
 
 from futuramaapi.core import settings
 from futuramaapi.repositories.models import UserModel
-from futuramaapi.routers.services import BaseSessionService
+from futuramaapi.routers.services import BaseSessionService, UnauthorizedError
 
 
 class TokenDecodeError(Exception):
@@ -59,18 +59,12 @@ async def process(self, *args, **kwargs) -> RedirectResponse:
         try:
             token_: dict[str, Any] = self._get_decoded_token()
         except TokenDecodeError:
-            raise HTTPException(
-                status_code=status.HTTP_401_UNAUTHORIZED,
-                detail="Token expired or invalid.",
-            ) from None
+            raise UnauthorizedError("Token expired or invalid.") from None
 
         user: UserModel = await self._get_current_user(token_["user"]["id"])
 
         if user.is_confirmed:
-            raise HTTPException(
-                status_code=status.HTTP_401_UNAUTHORIZED,
-                detail="User already activated.",
-            )
+            raise UnauthorizedError("User already activated.")
 
         await self.session.execute(self.__get_update_user_statement(user.id))
         await self.session.commit()