Day-25 WebSocket 接口

2023 iThome 鐵人賽

DAY 25

Modern Web

FastAPI 入門30天系列第 25 篇

15th鐵人賽 python fastapi 後端

Nesso

2023-09-30 11:29:27

396 瀏覽

分享至

今天我們要來講解如何在 FastAPI 中使用 WebSocket，首先我們先來了解甚麼是 WebSocket吧。

WebSocket

WebSocket 是一種網路通訊協定，與 http 不同的是，WebSocket 在單個 TCP 連接上進行全雙工的通訊，可以讓伺服器端向用戶端主動推送資料，也可以實現訊息流。

用戶端

在官方文件中有提供一個現成可供測試使用的簡易用戶端，我們可以直接拿來使用。

# client.py

from fastapi.responses import HTMLResponse

def get_websocket_client():
    html = """
        <!DOCTYPE html>
        <html>
            <head>
                <title>Chat</title>
            </head>
            <body>
                <h1>WebSocket Chat</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/websocket");
                    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>
    """
    return HTMLResponse(html)

我們將其放在獨立一個檔案，待會再從其他地方引用即可。

# main.py
from fastapi import FastAPI, WebSocket
from src.client import get_websocket_client

app = FastAPI()

@app.get("/client")
async def get():
    return get_websocket_client()

在定義接口的地方我們設置一個 /client 的接口讓我們可以存取用戶端。

WebSocket 接口

我們有用戶端後，就可以設置一個 WebSocket的接口來進行測試啦。

@app.websocket("/websocket")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    while True:
        data = await websocket.receive_text()
        await websocket.send_text(f"Message text was: {data}")

我們可以試著使用這個接口來發送資訊了，可以接收和發送二進位資料、文本和 JSON 資料。

可以使用用戶端發送訊息，伺服器端接收後會再發送回來。

你也可以在主控台上看到有人連接到你的伺服器上。

另外 WebSocket 的接口跟一般接口一樣，支援各種路徑參數跟依賴注入。

管理多個用戶端

首先我們先修改一下用戶端的html碼，以利我們等等進行測試

def get_websocket_client():
    html = """
				<!DOCTYPE html>
				<html>
				    <head>
				        <title>Chat</title>
				    </head>
				    <body>
				        <h1>WebSocket Chat</h1>
				        <h2>Your ID: <span id="ws-id"></span></h2>
				        <form action="" onsubmit="sendMessage(event)">
				            <input type="text" id="messageText" autocomplete="off"/>
				            <button>Send</button>
				        </form>
				        <ul id='messages'>
				        </ul>
				        <script>
				            var client_id = Date.now()
				            document.querySelector("#ws-id").textContent = client_id;
				            var ws = new WebSocket(`ws://localhost:8000/ws/${client_id}`);
				            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>
		"""
    return HTMLResponse(html)

當我們有多個用戶端時，其中一個如果關閉連線，await websocket.receive_text() 將拋出 WebSocketDisconnect 錯誤。我們可以捕捉該錯誤來確認哪個用戶端斷開了連線。

class ConnectionManager:
    def __init__(self):
        self.active_connections: list[WebSocket] = []

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)

    def disconnect(self, websocket: WebSocket):
        self.active_connections.remove(websocket)

    async def send_personal_message(self, message: str, websocket: WebSocket):
        await websocket.send_text(message)

    async def broadcast(self, message: str):
        for connection in self.active_connections:
            await connection.send_text(message)

透過 ConnectionManager 我們可以將連線存到內存中進行管理，透過這個類別也可以實現群發信息的功能。


@app.websocket("/websocket/{client_id}")
async def websocket_endpoint(websocket: WebSocket, client_id: int):
    await manager.connect(websocket)
    try:
        while True:
            data = await websocket.receive_text()
            await manager.send_personal_message(f"You wrote: {data}", websocket)
            await manager.broadcast(f"Client #{client_id} says: {data}")
    except WebSocketDisconnect:
        manager.disconnect(websocket)
        await manager.broadcast(f"Client #{client_id} left the chat")

我們將其套用到接口上，可以在連線時多紀錄一個client id ，這樣我們可以管理誰和我們發起了連線，我們要將訊息發送給誰。

這個接口在有人關閉連線時會群發訊息告訴現在有連線的人，我們可以透過開啟多個標籤頁來進行測試。

小結

今天跟大家介紹了 WebSocket 如何使用，當然實務上肯定沒有這麼簡單，如果想要在生產環境中使用 WebSocket 通常會使用其他套件與 FastAPI 集成，像是 https://github.com/encode/broadcaster，但今天就是簡單帶大家使用，沒有到這麼深入，有興趣的朋友們可以看看該套件的文檔，想必也會有一番收穫。