More than 3 years have passed since last update.

FastAPI シンプルチャット - Websocket

Posted at 2021-11-30

以下の過去記事でも述べていますが、FastAPIはStarletteを基盤として構築されており、WebSocketのサポートも継承しています。
Starletteで作る Simple Web Server - Qiita
WebSocket - FastAPI公式サイト

今回は公式サイトにあるチャットプログラムを少し修正し、特にクライアントは簡単なVueプログラムに置き替えています。

1. クライアント - Vue.js

過去記事のFastAPI OAuth2 クライアント - Qiitaと同じディレクトリ構成で、FastAPIのディレクトリのサブディレクトリstaticにindex.htmlとindex.jsを置きます。

index.htmlは2つのフィールドから構成されています。ログイン画面とメッセージ送信画面です。メッセージ送信ボタンはログインするまでは無効化されています。

static/index.html

<html>
    <head>
        <link rel="stylesheet" href="index.css">
        <script src="https://cdn.jsdelivr.net/npm/vue/dist/vue.js"></script>
    </head>
    <body>
        
    <div id="app">
        <div id="login">
            <h1>Login</h1>
            <p><label for="username">ユーザ名</label>
               <input id="username" v-model="username"></p>
            <p><label for="password">パスワード</label>
               <input id="password" v-model="password"></p>
            <button v-on:click="doLogin">ログイン</button>
            <hr>
            <h1>Chat message</h1>
            <p><input v-model="message"></p>
            <button v-bind:disabled="isButtonDisabled" v-on:click="sendMessage">メッセージ送信</button>
        </div>
        <ul id='messages'>
            <li v-for="(message,index) in logs">
                {{index}} {{message}}
            </li>
        </ul>
    </div>        
    <script src="index.js"></script>
    </body>
</html>

Vueのメッソドは2つあります。ログインを行うdoLoginと、メッセージを送信するsendMessageです。ログイン時の処理として、WebSocketのエンドポイントにはクエリパラメータでユーザ名とパスワードを渡します。これはWebSocketのインスタンスを作成するときに行われます。

index.js

var app = new Vue({
    el: '#app',
    data: {
      ws: null,
      username: 'yamada',
      password: 'pass123',
      message: '',
      logs: [],
      isButtonDisabled: true
    },
    methods: {
      doLogin: function(event) {
          this.ws = new WebSocket("ws://localhost:8000/ws?username=" + this.username + "&password=" + this.password);
          this.isButtonDisabled = false
          that = this
          this.ws.onmessage = function(event) {
            that.logs.push(event.data)
          };
          event.preventDefault()
        },
      sendMessage: function(event) {
        this.ws.send(this.message)
        this.message = ''
        event.preventDefault()
      }
    }
  })

2. サーバプログラム - FastAPI

サーバ側で注意すべき点は、path operation function ( websocket_endpoint ) の第一引数にWebSocketインスタンスが渡されることです。その後にクエリパラメータのusernameとpasswordが続きます。ここではナンチャッテ認証としてusernameとpasswordを検証していますが、もう少しそれらしくしたいのなら「FastAPI OAuth2 クライアント」でのコードをそのまま適用できるでしょう。

またブロードキャストを実現したい目的で、クラスConnectionManagerを定義し、接続クライアント( connection )を管理するようにしています。

main.py

from typing import List

from fastapi import FastAPI, Query, WebSocket, WebSocketDisconnect
from fastapi.staticfiles import StaticFiles

app = FastAPI()
app.mount("/static", StaticFiles(directory="static"), name="static")

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)

manager = ConnectionManager()


@app.websocket("/ws")
async def websocket_endpoint(
    websocket: WebSocket,
    username: str = Query(..., max_length=50),
    password: str = Query(..., max_length=50)
):
    print(f'username={username} password={password}')
    if username not in ['yamada','tanaka'] or password != 'pass123':
        return None
    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 #{username} says: {data}")
    except WebSocketDisconnect:
        manager.disconnect(websocket)
        await manager.broadcast(f"Client #{username} left the chat")