APIGatewayのWebSocket APIを使ってみる with Chalice

API GatewayのWebsocket APIとは

Amazon API GatewayのWebSocket APIは、双方向通信を可能にするリアルタイム通信プロトコルを提供するサービスです。これにより、クライアントとサーバーが持続的な接続を維持しながら、メッセージを効率的に送受信できます。

特徴

• 双方向通信
  • クライアントとサーバー間でリアルタイムにデータを送受信可能
• イベント駆動
  • 接続イベント (Connect)
  • メッセージイベント (Message)
  • 切断イベント (Disconnect)

各イベントに対してLambda関数などをトリガー可能。

• 持続的な接続

  • HTTPリクエストのように毎回接続を確立する必要がなく、接続を維持したまま通信を続けられる

• スケーラブル

  • AWSのマネージドサービスを活用し、大量の接続を処理可能

使いどころ

• チャットアプリケーション
• リアルタイム通知システム
• ゲームの状態同期
• IoTデバイスの制御

リアルタイムでの情報共有や双方向通信が必要なユースケースに非常に適しています。

AWS Chaliceとは

AWS Chaliceは、Pythonを使用してAWS Lambda関数やAPI Gatewayを簡単に構築・デプロイできるフレームワークです。サーバーレスアプリケーションの開発を効率化するために設計されており、API、イベント処理、バックグラウンドタスクなどを簡単に構築できます。

主な特徴

1. シンプルな開発

   • PythonコードでAPIエンドポイントやイベントハンドラーを簡単に記述
   • Flaskに似たシンタックスで直感的に利用可能

2. サーバーレスに最適化

   • AWS Lambda、API Gateway、DynamoDB、S3などと連携しやすい
   • インフラのスケーリングや管理をAWSに任せられる

3. イベント駆動アーキテクチャ

   • HTTPリクエスト（RESTまたはWebSocket）への応答
   • イベントトリガー（S3イベント、DynamoDB Streamsなど）
   • スケジュール実行（定期タスク）

4. デプロイが簡単

   • コマンド1つでAWSにデプロイ可能（chalice deploy）
   • 自動でIAMポリシーを生成し、必要な権限を設定

5. ローカル開発サポート

   • ローカルでテスト用サーバーを起動し、APIやイベント処理を試すことが可能

使いどころ

• REST API: シンプルなAPIバックエンドを素早く構築。
• WebSocket API: リアルタイム通信を実装。
• イベント駆動アプリケーション: S3、DynamoDB、SNS、SQSのイベントを処理。
• 定期タスク: AWS Lambdaのスケジュール実行。

実装

インストール＆プロジェクト作成

$ pip3 install chalice
$ chalice new-project chat --profile <your-profile>

以下がデフォルトのファイル構成

$ cd chat && ls
app.py                  requirements.txt

仮想環境の作成

$ python3 -m venv <your-venv-name>
$ ls

app.py                  requirements.txt        venv

仮想環境の起動

$ source venv/bin/activate

依存関係のインストール＆書き出し

$ pip3 install chalice boto3 boto3-stubs
$ pip3 freeze > requirements.txt

サンプルコード

• partition keyにconnection_idを指定します
• 本来はroomIdをクエリパラメーターで指定したができないので仕方なくmessageで指定
• event.bodyをそのままroomにいるユーザーへ送信

import json

import boto3.dynamodb
import boto3.dynamodb.types
import boto3.resources
from boto3.dynamodb.conditions import Key
from chalice import Chalice
from chalice.app import WebsocketEvent, Chalice
import boto3

app: Chalice = Chalice(app_name='chat_app')
app.experimental_feature_flags.update([
    'WEBSOCKETS'
])

dynamodb = boto3.resource('dynamodb')
TABLE_NAME = 'chat_app'
table = dynamodb.Table(TABLE_NAME)
session = boto3.Session()
app.websocket_api.session = session


@app.on_ws_connect()
def connect(event: WebsocketEvent):
    item = {
        'connectionId': event.connection_id
    }
    table.put_item(Item=item)


@app.on_ws_message()
def message(event: WebsocketEvent):
    body = json.loads(event.body)
    roomId = body.get('roomId', None)
    message = body.get('message', None)

    if message is None:
        table.put_item(
            Item={
                'connectionId': event.connection_id,
                'roomId': roomId
            }
        )
        return

    roomUsers = table.query(
        IndexName='roomId-index',
        KeyConditionExpression=Key('roomId').eq(roomId)
    )

    for user in roomUsers['Items']:
        connectionId = user['connectionId']
        if connectionId != event.connection_id:
            app.websocket_api.send(event.body)


@app.on_ws_disconnect()
def disconnect(event: WebsocketEvent):
    table.delete_item(
        Key={
            'connectionId': event.connection_id
        }
    )

デプロイ

$ chalice deploy --profile <your-profile>

Resources deployed:
  - Lambda ARN: xxx-chat-websocket_connect
  - Lambda ARN: xxx-chat-websocket_message
  - Lambda ARN: xxx-chat-websocket_disconnect
  - Websocket API URL: wss://xxxxxx.execute-api.ap-northeast-1.amazonaws.com/api/

使ってみた感想

• 超絶お手軽に実装できる
  • DynamoDBのテーブルを作成する必要はあるがそれ以外は特に面倒くさいことがないのでちょっとしたものを作るのであればこれでいい感はある
    　　
• サーバーレスなのでスケーラビリティを考える必要がない
  • ここが一番嬉しいかもしれない
    　　
• Policyを自動で解析してアタッチする機能があるが完全じゃないので結局自分で指定しなきゃいけない
  • それならもうTerraformで最初から...ってなる感もある
    　　
• クエリパラメーターだったりヘッダーだったり簡潔フレームワークが故の歯痒さがある
  • 認証とか諸々細かくやろうと思うとキツイ

最後に

自分はちょっとしたものを作るときは大体コイツを使っています。（Terraformから逃げているなんて言えない😇）
今度はちゃんとTerraformでも実装できるようにします。。。