aioredis fastapi websocket pubsub

ws_app.py

import uvicorn
from aioredis import create_pool
from fastapi import FastAPI
from starlette.websockets import WebSocket

app = FastAPI()


REDIS_URL = 'redis://redis:6379'
REDIS_DB = 0
REDIS_PASSWORD = None

POOL = None


async def get_pool():
    global POOL

    if not POOL:
        POOL = await create_pool(REDIS_URL, db=REDIS_DB, password=REDIS_PASSWORD, minsize=5, maxsize=1000)

    return POOL


async def handler(websocket: WebSocket):
    await websocket.accept()
    await websocket.send_json({"message": "Welcome to server!"})

    user_id = websocket.path_params["user_id"]
    pool = await get_pool()

    with await pool as connection:
        try:
            await connection.execute_pubsub('subscribe', f'user:{user_id}')
            print("Active users:", connection.pubsub_channels.keys())
            channel = connection.pubsub_channels[f'user:{user_id}']

            while await channel.wait_message():
                message = await channel.get()

                if message:
                    await websocket.send_bytes(message)

        except Exception as e:
            print("Closed in handler: ", e)
        finally:
            await connection.execute_pubsub('unsubscribe', f'user:{user_id}')
            await websocket.close()


app.add_websocket_route('/ws/{user_id}/', handler)

uvicorn.run(app, host='0.0.0.0', port=8080)