improve generic endpoint exception handler

Author: fynnos

backend/src/common/exception_handler.py

@@ -1,15 +1,45 @@
+import traceback
+from collections.abc import Callable
+
 from fastapi import Request
 from fastapi.responses import PlainTextResponse
 from loguru import logger
 
 exception_handlers = []
 
 
-def exception_handler(http_status_code: int):
+def exception_handler(
+    http_status_code: int | Callable[[Exception], int],
+    extract_message: Callable[[Exception], str] = lambda exc: str(exc),
+):
     def decorator(exception_class):
         def handle_exception(req: Request, exc: Exception):
-            logger.exception(exc)
-            return PlainTextResponse(str(exc), status_code=http_status_code)
+            *_, (frame, lineno) = traceback.walk_tb(exc.__traceback__)
+
+            # start traceback at our code by skipping any library code (starlette etc.)
+            tb = exc.__traceback__
+            while tb is not None and "site-packages" in tb.tb_frame.f_code.co_filename:
+                tb = tb.tb_next
+
+            # override line, function name, module etc. because loguru would
+            # otherwise use the current file location, function name etc.
+            logger.patch(
+                lambda record: record.update(
+                    line=lineno,  # type: ignore
+                    function=frame.f_code.co_name,
+                    module=frame.f_globals["__name__"].split(".")[-1],
+                    name=frame.f_globals["__name__"],
+                )
+            ).opt(exception=(type(exc), exc, tb)).error(extract_message(exc))
+
+            if isinstance(http_status_code, int):
+                status_code = http_status_code
+            elif isinstance(http_status_code, Callable):
+                status_code = http_status_code(exc)
+            else:
+                status_code = 500  # type: ignore
+
+            return PlainTextResponse(extract_message(exc), status_code)
 
         exception_handlers.append((exception_class, handle_exception))
         return exception_class

backend/src/main.py

@@ -8,15 +8,14 @@
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.middleware.gzip import GZipMiddleware
 from fastapi.openapi.utils import get_openapi
-from fastapi.responses import PlainTextResponse
 from fastapi.routing import APIRoute
 from loguru import logger
 from psycopg2.errors import UniqueViolation
 from sqlalchemy.exc import IntegrityError
 from starlette.middleware.sessions import SessionMiddleware
 from uvicorn.main import run
 
-from common.exception_handler import exception_handlers
+from common.exception_handler import exception_handler, exception_handlers
 from repos.elastic.elastic_repo import ElasticSearchRepo
 from repos.llm_repo import LLMRepo
 from utils.import_utils import import_by_suffix
@@ -118,21 +117,16 @@ def custom_openapi():
     app.include_router(em.router)
 
 
-@app.exception_handler(IntegrityError)
-async def integrity_error_handler(_, exc: IntegrityError):
-    logger.exception(exc)
-    if isinstance(exc.orig, UniqueViolation):
-        msg = str(exc.orig.pgerror).split("\n")[1]
-        return PlainTextResponse(msg, status_code=409)
-    else:
-        return PlainTextResponse(str(exc), status_code=500)
-
-
-@app.exception_handler(NoSuchJobError)
-async def no_such_job_error_handler(_, exc: NoSuchJobError):
-    logger.exception(exc)
-    return PlainTextResponse(str(exc), status_code=404)
+exception_handler(
+    http_status_code=lambda exc: 409
+    if isinstance(exc, IntegrityError) and isinstance(exc.orig, UniqueViolation)
+    else 500,
+    extract_message=lambda exc: str(exc.orig.pgerror).split("\n")[1]
+    if isinstance(exc, IntegrityError) and isinstance(exc.orig, UniqueViolation)
+    else str(exc),
+)(IntegrityError)
 
+exception_handler(404)(NoSuchJobError)
 
 # register all exception handlers in fastAPI
 for ex_class, handler_func in exception_handlers: