How to add background tasks when request fails and HTTPException is raised in FastAPI?

user_4117975 0
Viewed 231

I was trying to generate logs when an exception occurs in my FastAPI endpoint using a Background task as:

from fastapi import BackgroundTasks, FastAPI

app = FastAPI()

def write_notification(message=""):
    with open("log.txt", mode="w") as email_file:
        content = f"{message}"
        email_file.write(content)

@app.post("/send-notification/{email}")
async def send_notification(email: str, background_tasks: BackgroundTasks):
    if "hello" in email:
        background_tasks.add_task(write_notification, message="helloworld")
        raise HTTPException(status_code=500, detail="example error")

    background_tasks.add_task(write_notification, message="hello world.")
    return {"message": "Notification sent in the background"}

However, the logs are not generated because according to the documentation here and here, a background task runs "only" after the return statement is executed.

Is there any workaround to this? Thanks.

starlette fastapi logging python
1 Answers

The way to do this is to override the HTTPException error handler, and since there is no BackgroundTasks object in the exception_handler, you can add a background task to a response in the way it is described in Starlette documentation (FastAPI is actually Starlette underneath). Example below:

from fastapi import BackgroundTasks, FastAPI, HTTPException, Request
from fastapi.responses import PlainTextResponse
from starlette.exceptions import HTTPException as StarletteHTTPException
from starlette.background import BackgroundTask

app = FastAPI()

def write_notification(message):
    with open('log.txt', 'a') as f:
        f.write(f'{message}'+'\n')

@app.exception_handler(StarletteHTTPException)
async def http_exception_handler(request, exc):
    task = BackgroundTask(write_notification, message=exc.detail)
    return PlainTextResponse(str(exc.detail), status_code=exc.status_code, background=task)
 
@app.get("/{msg}")
def send_notification(msg: str, background_tasks: BackgroundTasks):
    if "hello" in msg:
        raise HTTPException(status_code=500, detail="Something went wrong")

    background_tasks.add_task(write_notification, message="Success")
    return {"message": "Request has been successfully submitted."}

If you need to add multiple background tasks to a response, then use:

@app.exception_handler(StarletteHTTPException)
async def http_exception_handler(request, exc):
    tasks = BackgroundTasks()
    tasks.add_task(write_notification, message=exc.detail)
    tasks.add_task(some_other_function, message="some other message")
    return PlainTextResponse(str(exc.detail), status_code=exc.status_code, background=tasks)

A variation of the above approach is the following (suggested here):

from starlette.background import BackgroundTask

@app.exception_handler(StarletteHTTPException)
async def http_exception_handler(request, exc):
    response = PlainTextResponse(str(exc.detail), status_code=exc.status_code)
    response.background = BackgroundTask(write_notification, message=exc.detail)
    return response

Some references that might prove useful for your task are: this answer that demonstrates how to add custom exception handlers, as well as this answer that shows a custom logging system for the incoming requests and outgoing responses.

user_178658040
Related