Permanently running background task that listens to Postgres notifications and sends data to websocket

First Check

• I added a very descriptive title to this issue.
• I used the GitHub search to find a similar issue and didn’t find it.
• I searched the FastAPI documentation, with the integrated search.
• I already searched in Google “How to X in FastAPI” and didn’t find any information.
• I already read and followed all the tutorial in the docs and didn’t find an answer.
• I already checked if it is not related to FastAPI but to Pydantic.
• I already checked if it is not related to FastAPI but to Swagger UI.
• I already checked if it is not related to FastAPI but to ReDoc.

Commit to Help

• I commit to help with one of those options 👆

Example Code

import asyncio
import aiopg
from fastapi import FastAPI, WebSocket


dsn = "dbname=aiopg user=aiopg password=passwd host=127.0.0.1"
app = FastAPI()


class ConnectionManager:
    self.count_connections = 0
    # other class functions and variables are taken from FastAPI docs
    ...


manager = ConnectionManager()


async def send_and_receive_data(websocket: WebSocket):
    data = await websocket.receive_json()
    await websocket.send_text('Thanks for the message')
    # then process received data


# taken from official aiopg documentation
# the function listens to PostgreSQL notifications
async def listen(conn):
    async with conn.cursor() as cur:
        await cur.execute("LISTEN channel")
        while True:
            msg = await conn.notifies.get()


async def postgres_listen():
    async with aiopg.connect(dsn) as listenConn:
        listener = listen(listenConn)
        await listener


@app.get("/")
def read_root():
    return {"Hello": "World"}


@app.websocket("/")
async def websocket_endpoint(websocket: WebSocket):
    await manager.connect(websocket)
    manager.count_connections += 1

    if manager.count_connections == 1:
        await asyncio.gather(
            send_and_receive_data(websocket),
            postgres_listen()
        )
    else:
        await send_and_receive_data(websocket)

Description

I am building an app with Vue.js, FastAPI and PostgreSQL. In this example I attempt to use listen/notify from Postgres and implement it in the websocket. I also use a lot of usual http endpoints along with the websocket endpoint.

I want to run a permanent background asynchronous function at the start of the FastAPI app that will then send messages to all websocket clients/connections. So, when I use uvicorn main:app it should not only run the FastAPI app but also my background function postgres_listen(), which notifies all websocket users, when a new row is added to the table in the database.

I know that I can use asyncio.create_task() and place it in the on_* event, or even place it after the manager = ConnectionManager() row, but it will not work in my case! Because after any http request (for instance, read_root() function), I will get the same error described below.

You see that I use a strange way to run my postgres_listen() function in my websocket_endpoint() function only when the first client connects to the websocket. Any subsequent client connection does not run/trigger this function again. And everything works fine… until the first client/user disconnects (for example, closes browser tab). When it happens, I immediately get the GeneratorExit error caused by psycopg2.OperationalError:

Future exception was never retrieved
future: <Future finished exception=OperationalError('Connection closed')>
psycopg2.OperationalError: Connection closed
Task was destroyed but it is pending!
task: <Task pending name='Task-18' coro=<Queue.get() done, defined at 
/home/user/anaconda3/lib/python3.8/asyncio/queues.py:154> wait_for=<Future cancelled>>

The error comes from the listen() function. After this error, I will not get any notification from the database as the asyncio’s Task is cancelled. There is nothing wrong with the psycopg2, aiopg or asyncio. The problem is that I don’t understand where to put the postgres_listen() function so it will not be cancelled after the first client disconnects. From my understanding, I can easily write a python script that will connect to the websocket (so I will be the first client of the websocket) and then run forever so I will not get the psycopg2.OperationalError exception again and the Task won’t be cancelled, but it does not seem right to do so.

My question is: where should I put postgres_listen() function, so the first connection to websocket may be disconnected with no consequences?

P.S. asyncio.shield() also does not work

Operating System

Linux

Operating System Details

No response

FastAPI Version

0.77.1

Python Version

3.8.5

Additional Context

No response

[JarroVGIT]

Okay this kept bugging me, so I’ve created a more relevant example for you. The code can be found here:

https://github.com/JarroVGIT/fastapi-github-issues/tree/master/5015

The app.py is basically the same as above, but with one change: it will listen to a websocket and publishes all incoming messages to all connections to the websockets in app.py.

The generator.py is a simple FastAPI app that has a websocket (that our example above listens to) that emits a message every 2 seconds to every connection it gets.

To try this out:

1. Start generator.py (e.g. python3 generator.py on your command line when in your working folder)
2. Start app.py (either debug mode in VScode or same as above)
3. Listen to http://localhost:8000/ws (= endpoint in app.py) with several clients, you will see that they will all join in the same message streak.

Let me know if this helps you out!

[artemonsh]

@JarroVGIT, huge respect for such well-documented code and reproducible example as well. It works as expected, it works like magic.