How to raise custom exceptions in a FastAPI middleware?

Question:

I have a simple FastAPI setup with a custom middleware class inherited from BaseHTTPMiddleware. Inside this middleware class, I need to terminate the execution flow under certain conditions. So, I created a custom exception class named CustomError and raised the exception.

from fastapi import FastAPI, Request
from starlette.middleware.base import (
    BaseHTTPMiddleware,
    RequestResponseEndpoint
)
from starlette.responses import JSONResponse, Response

app = FastAPI()


class CustomError(Exception):
    def __init__(self, message):
        self.message = message

    def __str__(self):
        return self.message


class CustomMiddleware(BaseHTTPMiddleware):
    def execute_custom_logic(self, request: Request):
        raise CustomError("This is from `CustomMiddleware`")

    async def dispatch(
            self,
            request: Request,
            call_next: RequestResponseEndpoint,
    ) -> Response:
        self.execute_custom_logic(request=request)
        response = await call_next(request)
        return response


app.add_middleware(CustomMiddleware)


@app.exception_handler(CustomError)
async def custom_exception_handler(request: Request, exc: CustomError):
    return JSONResponse(
        status_code=418,
        content={"message": exc.message},
    )


@app.get(path="/")
def root_api():
    return {"message": "Hello World"}

Unfortunately, FastAPI couldn’t handle the CustomError even though I added custom_exception_handler(...) handler.

Questions

  1. What is the FastAPI way to handle such situations?
  2. Why is my code not working?

Versions

  • FastAPI – 0.95.2
  • Python – 3.8.13
Asked By: JPG

||
Source

Answers:

FastAPI’s custom exception handlers are not handling middleware level exceptions. Although this is not stated anywhere in the docs, there is part about HTTPException, which says that you can raise HTTPException if you are inside a utility function that you are calling inside of your path operation function. HTTPException has default exception handler that acts absolutely the same as custom exception handlers do.

You can either handle your error (with try/except) within the same middleware or have separate middleware e.g. ExceptionHandlerMiddleware (but you’ll have to keep the order of middleware chain correct).

Answered By: sudden_appearance

The obvious way would be to raise an HTTPException; however, in a FastAPI/Starlette middleware, this wouldn’t work, leading to Exception in ASGI application error on server side, and hence, an Internal Server Error would be returned to the client.

Option 1 – Using middleware and try/except block

You could use a try/except block to handle the custom exception raised in your custom function. Once the error is raised, you could return a JSONResponse (or custom Response, if you prefer), including the msg (and any other arguments) from CustomException, as well as the desired status_code (in the example given below, 500 status code is used, which could be replaced by the status code of your choice).

Working Example

from fastapi import FastAPI, Request, HTTPException
from fastapi.responses import JSONResponse

app = FastAPI()


class CustomException(Exception):
    def __init__(self, msg: str):
        self.msg = msg
 
 
def exec_custom_logic(request: Request):
    raise CustomException(msg='Something went wrong') 

    
@app.middleware("http")
async def custom_middleware(request: Request, call_next):
    try:    
        exec_custom_logic(request)
    except CustomException as e:
        return JSONResponse(status_code=500, content={'message': e.msg})
        
    return await call_next(request)
    
    
@app.get('/')
async def main(request: Request):
    return 'OK'

Option 2 – Using an APIRouter with a custom APIRoute class

You could use an APIRouter with a custom APIRoute class, as demonstrated in Option 4 of this answer, and either handle the custom exception inside a try/except block (as shown in the previous option above), or raise an HTTPException directly. The advantages of this approach are: (1) you could raise an HTTPException directly, and hence, there is no need for using try/except blocks, and (2) you could add to the APIRouter only those routes that you would like to handle that way, using, for instance, the @router.get() decorator, while the rest of the routes could be added to the app instance, using, for example, the @app.get() decorator.

Working Example

from fastapi import FastAPI, APIRouter, Response, Request, HTTPException
from fastapi.routing import APIRoute
from typing import Callable


def exec_custom_logic(request: Request):
    raise HTTPException(status_code=500, detail='Something went wrong')
    

class CustomAPIRoute(APIRoute):
    def get_route_handler(self) -> Callable:
        original_route_handler = super().get_route_handler()

        async def custom_route_handler(request: Request) -> Response:
            exec_custom_logic(request)
            return await original_route_handler(request)
       
        return custom_route_handler


app = FastAPI()
router = APIRouter(route_class=CustomAPIRoute)


@router.get('/')
async def main(request: Request):
    return 'OK'
    
app.include_router(router)
Answered By: Chris