How to implement with encoder/broadcaster for Websocket pub/sub?

[Buuntu]

I see https://github.com/encode/broadcaster mentioned a lot on here (as well as in the documentation) as a project to use for more complicated Websocket integration. However I haven't been able to find a single example of doing this - I think it would be super useful if there was a quick example in the documentation of how to plug it in.

Encoder is made for Starlette but I imagine it can easily be modified to work with FastAPI, though for someone who is not an expert in both it doesn't seem trivial. Especially once you want to add testing and more complex backends such as Redis.

[Kludex]

On their repository (https://github.com/encode/broadcaster) you can see this example:

# Requires: `starlette`, `uvicorn`, `jinja2`
# Run with `uvicorn example:app`
from broadcaster import Broadcast
from starlette.applications import Starlette
from starlette.concurrency import run_until_first_complete
from starlette.routing import Route, WebSocketRoute
from starlette.templating import Jinja2Templates


broadcast = Broadcast("redis://localhost:6379")
templates = Jinja2Templates("templates")


async def homepage(request):
    template = "index.html"
    context = {"request": request}
    return templates.TemplateResponse(template, context)


async def chatroom_ws(websocket):
    await websocket.accept()
    await run_until_first_complete(
        (chatroom_ws_receiver, {"websocket": websocket}),
        (chatroom_ws_sender, {"websocket": websocket}),
    )


async def chatroom_ws_receiver(websocket):
    async for message in websocket.iter_text():
        await broadcast.publish(channel="chatroom", message=message)


async def chatroom_ws_sender(websocket):
    async with broadcast.subscribe(channel="chatroom") as subscriber:
        async for event in subscriber:
            await websocket.send_text(event.message)


routes = [
    Route("/", homepage),
    WebSocketRoute("/", chatroom_ws, name='chatroom_ws'),
]


app = Starlette(
    routes=routes, on_startup=[broadcast.connect], on_shutdown=[broadcast.disconnect],
)

If you want to adapt for FastAPI, just change the last three lines:

app = FastAPI(
    routes=routes, on_startup=[broadcast.connect], on_shutdown=[broadcast.disconnect],
)

I guess you want a "pure" FastAPI implementation? 🤔

# Requires: `starlette`, `uvicorn`, `jinja2`
# Run with `uvicorn example:app`
from broadcaster import Broadcast
from starlette.applications import Starlette
from starlette.concurrency import run_until_first_complete
from starlette.routing import Route, WebSocketRoute
from starlette.templating import Jinja2Templates


broadcast = Broadcast("redis://localhost:6379")
templates = Jinja2Templates("templates")


async def chatroom_ws_receiver(websocket):
    async for message in websocket.iter_text():
        await broadcast.publish(channel="chatroom", message=message)


async def chatroom_ws_sender(websocket):
    async with broadcast.subscribe(channel="chatroom") as subscriber:
        async for event in subscriber:
            await websocket.send_text(event.message)


app = FastAPI(on_startup=[broadcast.connect], on_shutdown=[broadcast.disconnect])

@app.get("/")
async def homepage(request: Request):
    template = "index.html"
    context = {"request": request}
    return templates.TemplateResponse(template, context)
    
@app.websocket("/")
async def chatroom_ws(websocket: WebSocket):
    await websocket.accept()
    await run_until_first_complete(
        (chatroom_ws_receiver, {"websocket": websocket}),
        (chatroom_ws_sender, {"websocket": websocket}),
    )

EDIT: I'm on my cellphone now, but please take in considerations that I didn't add the right imports. 😅

[Buuntu]

Thank you! I had tried this almost verbatim but I kept having it the Websocket disconnect almost immediately. Though now it seems to work when using memory, so I must have been missing an import or something trivial :). When I use this with redis it just gives me 502 errors when trying to connect to the Websocket.

Another issue I run into when using the above code is hanging tests. For example, I have a test like this that never finishes:

def test_websocket():
    client = TestClient(app)
    with client.websocket_connect("/ws") as websocket:
        data = websocket.receive_json()
        assert data == {"msg": "Hello WebSocket"}

I still think it might be good to have some examples of having multiple channels for example and handling real data, using Redis for example.

[lefnire]

Per another convo, dropping https://github.com/authorizon/fastapi_websocket_pubsub here for reference. It's a robust, user-friendly tool for pub/sub, and uses Broadcaster under the hood.

[Buuntu]

Woops, redis wasn't working because of docker-compose (calling the redis container localhost instead of the name of the Docker container). I still can't get tests to work with Redis though, maybe I need to mock Redis somehow. Currently have this as a test:

def test_websocket():
    client = TestClient(app)
    with client.websocket_connect("/api/ws/lobby") as websocket:
        data = websocket.receive_json()
        assert data == {"msg": "Hello WebSocket"}

and get the error message:

    async def subscribe(self, channel: str) -> None:
>       await self._subscriber.subscribe([channel])
E       AttributeError: 'RedisBackend' object has no attribute '_subscriber'

In case someone stumbles on to this, these are the exact files I needed to make it work (tests aside):

main.py:

from fastapi import FastAPI, Request, WebSocket

from broadcaster import Broadcast
from starlette.applications import Starlette
from starlette.concurrency import run_until_first_complete
from starlette.routing import Route, WebSocketRoute
from starlette.templating import Jinja2Templates


broadcast = Broadcast("redis://localhost:6379")
templates = Jinja2Templates("templates")


async def chatroom_ws_receiver(websocket):
    async for message in websocket.iter_text():
        await broadcast.publish(channel="chatroom", message=message)


async def chatroom_ws_sender(websocket):
    async with broadcast.subscribe(channel="chatroom") as subscriber:
        async for event in subscriber:
            await websocket.send_text(event.message)


app = FastAPI(on_startup=[broadcast.connect], on_shutdown=[broadcast.disconnect])

@app.get("/")
async def homepage(request: Request):
    template = "index.html"
    context = {"request": request}
    return templates.TemplateResponse(template, context)
    
@app.websocket("/")
async def chatroom_ws(websocket: WebSocket):
    await websocket.accept()
    await run_until_first_complete(
        (chatroom_ws_receiver, {"websocket": websocket}),
        (chatroom_ws_sender, {"websocket": websocket}),
    )

I needed to add websockets to the requirements file as well:

requirements.txt:

fastapi
uvicorn
websockets
jinja2
broadcaster[redis,postgres,kafka]

and the templates directory from https://github.com/encode/broadcaster/tree/master/example.

Once I get the tests working I'd like to add a PR to the FastAPI docs to show how to use this correctly, since broadcaster is mentioned in the docs already.

[tiangolo]

Thanks for the help here everyone! 👏 🙇

I think this would be good material for a blog post. 🤓 📝

Sorry for the long delay! 🙈 I wanted to personally address each issue/PR and they piled up through time, but now I'm checking each one in order.

[github-actions]

Assuming the original need was handled, this will be automatically closed now. But feel free to add more comments or create new issues or PRs.