FastAPI: WebSocket Example

1. Websocket을 통해 서버로부터 일정 주기로 데이터 받아오기

 

 

* fastapi와 websocket 사용을 위한 uvicorn[standard] 설치

pip install fastapi "uvicorn[standard]"

 

 

* main.py

from fastapi import FastAPI, WebSocket
import asyncio

app = FastAPI()

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    counter = 0
    while True:
        counter += 1
        await websocket.send_json({"message": f"Update #{counter}"})
        await asyncio.sleep(1)

 

 

* index.html or .js

const ws = new WebSocket("ws://localhost:8000/ws");
ws.onmessage = function(event) {
    console.log("Message from server:", event.data);
};

 

 

2. 브라우저 console창을 활용한 간단한 채팅 기능 구현

 

*main.py

from fastapi import WebSocket, WebSocketDisconnect
from typing import List

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 broadcast(self, message: str):
        for connection in self.active_connections:
            await connection.send_text(message)

manager = ConnectionManager()

@app.websocket("/ws/chat")
async def websocket_endpoint(websocket: WebSocket):
    await manager.connect(websocket)
    try:
        while True:
            data = await websocket.receive_text()
            await manager.broadcast(f"Message: {data}")
    except WebSocketDisconnect:
        manager.disconnect(websocket)
        await manager.broadcast("A user has left the chat.")

 

* .html or javascript

const ws = new WebSocket("ws://127.0.0.1:8000/ws/chat");
ws.onmessage = function(event) {
    console.log("Received:", event.data);
};

 

브라우저를 여러개 키고, 콘솔 창에서

ws.send();

를 통해 메시지를 보내면 다른 브라우저 콘솔창에서도 뜨는게 확인 된다.

 

한쪽 브라우저를 종료하면 

FastAPI server에서는 'connection closed' 메시지 확인 가능

 

남아있는 유저의 브라우저 콘솔에서는 await manager.broadcast("A user has left the chat.")

을 통해 보내진 'A user has left the chat.' 메시지 확인 가능