FastAPI Learning Road (Forty-four) WebSockets

FastAPI Learning Road (Part 44) WebSockets

2024-07-12

Our previous analysis was based on http requests. So can websockets support it? The answer is yes. Let's see how it is implemented.


from fastapi import WebSocket, FastAPI
from fastapi.responses import HTMLResponse
 
app = FastAPI()
 
html = """
 
<!DOCTYPE html>
 
<html>
 
    <head>
 
        <title>Chat</title>
 
    </head>
 
    <body>
 
        <h1>WebSocket 聊天</h1>
 
        <form action="" onsubmit="sendMessage(event)">
 
            <input type="text" id="messageText" autocomplete="off"/>
 
            <button>Send</button>
 
        </form>
 
        <ul id='messages'>
 
        </ul>
 
        <script>
 
            var ws = new WebSocket("ws://localhost:8000/ws");
 
            ws.onmessage = function(event) {
 
                var messages = document.getElementById('messages')
 
                var message = document.createElement('li')
 
                var content = document.createTextNode(event.data)
 
                message.appendChild(content)
 
                messages.appendChild(message)
 
            };
 
            function sendMessage(event) {
 
                var input = document.getElementById("messageText")
 
                ws.send(input.value)
 
                input.value = ''
 
                event.preventDefault()
 
            }
</script>
 
    </body>
 
</html>
 
"""
 
 
@app.get("/")
async def get_page():
    return HTMLResponse(html)
 
 
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    while True:
        data = await websocket.receive_text()
        await websocket.send_text(f"接收到的消息是：{data}")

It's actually very simple. We write an interface to receive messages in the backend. When a message comes, we process it. Before, we processed it using http, but now we use websocket. We return the received message directly to the frontend.

Let's look at the effect:

So we have simply implemented a websocket function. So if we have a specific message, we close the chat, how do we deal with it?


# 接收特定的消息后，终止该聊天
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    while True:
        data = await websocket.receive_text()
        if data == "bye":
            await websocket.send_text(f"请注意：窗口即将关闭")
            await websocket.send_text("窗口已关闭")
            await websocket.close()
            break
        else:
            await websocket.send_text(f"接收到的消息是：{data}")

The front-end code has not been changed, only the back-end has been modified.

In fact, it is very simple to implement websocket communication here.

Of course, we can also use it in the request

Depends
Security
Cookie
Header
Path
Query

We can look at the following simple backend example


# 使用多种参数，比如获取cookie等
from typing import Optional
from fastapi import Cookie, Query, status, Depends
 
 
async def get_cookie_or_token(
    websocket: WebSocket,
    session: Optional[str] = Cookie(None),
    token: Optional[str] = Query(None),
):
    if session is None and token is None:
        await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
    return session or token
 
 
@app.websocket("item/ws")
async def websocket_endpoint(
    websocket: WebSocket,
    q: Optional[int] = None,
    cookie_or_token: str = Depends(get_cookie_or_token),
):
    await websocket.accept()
    while True:
        data = await websocket.receive_text()
        await websocket.send_text(
            f"Session cookie or query token value is: {cookie_or_token}"
        )
        if q is not None:
            await websocket.send_text(f"parameter q is: {q}")
        await websocket.send_text(f"Message text was: {data}, for : {q}")

Although we have only simply implemented the backend, it can be seen that fastapi has a good support for websocket

Technology Sharing

FastAPI Learning Road (Part 44) WebSockets

Personal profile

my contact information