AWS サーバーレス（Python）でのWebSocket

AWSのサーバーレス（Python）でWebSocketを実現する方法を解説します。

システム構成

システム構成は下図になります。
AWS上ではAPI Gateway、Lambda（言語はPython）、DynamoDBを使用します。
Lambdaは3つ作成し、WebSocketの接続用・切断用・メッセージ送信用になります。

AWS環境の構築

DynamoDBの追加

接続情報の管理のため、DynamoDBを追加します。
AWSマネージメントコンソール ＞ サービス ＞ DynamoDB ＞ テーブルで、テーブルを作成します。

接続情報のキー（パーティションキー）は「id」とします。

IAMロールの用意

Lambdaに適用するIAMロールを用意します。
基本的なIAMロールの追加方法は、こちらの「IAMポリシーの追加」「IAMロールの追加」を参考にして下さい。

IAMポリシーは以下のように設定します。

Lambdaの追加

connect用

WebSocket接続用のLambdaを追加します。
基本的なLambdaの追加方法はこちらを参考にして下さい。ローカル環境で用意します。

LambdaのPythonファイル（lambda_websocket_connect.py）には、以下のように実装します。

lambda_websocket_connect.py

import sys
import json
import boto3

g_dynamodb = boto3.resource('dynamodb')

def lambda_handler(event, context):
    dbname = '作成したDynamoDBの名前'  # e.g. iso-dynamodb-websocket-test
    table = g_dynamodb.Table(dbname)
    connectionid = event.get('requestContext',{}).get('connectionId')
    ret = table.put_item(Item={ 'id': connectionid })
    return {
        'statusCode': 200,
        'headers': {
            "Access-Control-Allow-Origin": "*"
        },
        'body': json.dumps({"result": 0}, ensure_ascii=False)
    }

disconnect用

WebSocket切断用のLambdaを追加します。

LambdaのPythonファイル（lambda_websocket_disconnect.py）には、以下のように実装します。

lambda_websocket_disconnect.py

import sys
import json
import boto3

g_dynamodb = boto3.resource('dynamodb')

def lambda_handler(event, context):
    dbname = '作成したDynamoDBの名前'  # e.g. iso-dynamodb-websocket-test
    table = g_dynamodb.Table(dbname)
    connectionid = event.get('requestContext',{}).get('connectionId')
    ret = table.delete_item(Key={ 'id': connectionid })
    return {
        'statusCode': 200,
        'headers': {
            "Access-Control-Allow-Origin": "*"
        },
        'body': json.dumps({"result": 0}, ensure_ascii=False)
    }

sendmsg用

WebSocketでのメッセージ送信用のLambdaを追加します。

LambdaのPythonファイル（lambda_websocket_sendmsg.py）には、以下のように実装します。

lambda_websocket_sendmsg.py

import sys
import json
import boto3
import botocore

g_dynamodb = boto3.resource('dynamodb')
g_apigw_management = boto3.client('apigatewaymanagementapi', endpoint_url=F"この後出てくるAPI Gatewayの接続URL")
# 接続URL例 … https://hogehoge.execute-api.ap-northeast-1.amazonaws.com/production
# API Gatewayの画面で見えている@マーク以降は不要

def lambda_handler(event, context):
    dbname = '作成したDynamoDBの名前'  # e.g. iso-dynamodb-websocket-test
    table = g_dynamodb.Table(dbname)
    post_data = json.loads(event.get('body', '{}')).get('data')
    print(post_data)
    domain_name = event.get('requestContext',{}).get('domainName')
    stage       = event.get('requestContext',{}).get('stage')
    items = table.scan(ProjectionExpression='id').get('Items')
    if items is None:
        return {
            'statusCode': 500,
            'headers': {
                "Access-Control-Allow-Origin": "*"
            },
            'body': json.dumps({"result": 1}, ensure_ascii=False)
        }
    for item in items:
        try:
            print(item)
            _ = g_apigw_management.post_to_connection(ConnectionId=item['id'], Data=post_data)
        except botocore.exceptions.ClientError as e:
            print('Failed')
            print(e.response)
    return {
        'statusCode': 200,
        'headers': {
            "Access-Control-Allow-Origin": "*"
        },
        'body': json.dumps({"result": 0}, ensure_ascii=False)
    }

API Gatewayの追加

API Gatewayを追加します。
AWSマネージメントコンソール ＞ サービス ＞ API Gateway ＞ APIで、APIを作成します。
APIタイプを選択では「WebSocket API」を選択して構築します。
ルート選択式は「request.body.action」とします。

connectルートの設定

API Gatewayのconnectルートの設定を行います。
AWSマネージメントコンソール ＞ サービス ＞ API Gateway ＞ APIで、作成したAPIを選択します。さらに統合リクエストを選択します。

disconnectルートの設定

API Gatewayのdisconnectルートの設定を行います。

sendmsgルートの設定

Gatewayのsendmsgルートの設定を行います。

sendmsgルートを追加します。

統合リクエストを設定します。

APIのデプロイ

作成したAPIに外部からアクセスできるように、APIをデプロイします。

ステージエディターで「WebSocket URL」と「接続 URL」を控えておきます。

sendmsg用Lambdaの修正

先のlambda_websocket_sendmsg.pyについて、「この後出てくるAPI Gatewayの接続URL」の箇所を、上記の「接続 URL」で書き換えます。

Webサイトの用意

ローカルサイトにWebサイトを用意します。
基本的なWebサイトの構築方法はこちらを参考にして下さい。

ドキュメントルート以下に、index.htmlファイルを配置し、以下のように実装します。

index.html

<!DOCTYPE html>
<html>
<head>
    <meta charset="UTF-8">
    <title>AWS WebSocket チャット</title>
</head>
<body>
    <H3>AWS WebSocket チャット</H3>
    <input id="input" type="text" />
    <button onclick="send()">送信</button>
    <pre id="output"></pre>
    <script>
        var input = document.getElementById('input');
        var output = document.getElementById('output');
        var socket = new WebSocket("控えておいたAPI GatewayのWebSocket URL");

        // WebSocket URL例 … wss://hogehoge.execute-api.ap-northeast-1.amazonaws.com/production

        socket.onopen = function() {
           output.innerHTML += "接続に成功しました\n";
        };

        socket.onmessage = function(e) {
            output.innerHTML += "受信：" + e.data + "\n";
        };

        function send() {
            socket.send(JSON.stringify(
                {
                    "action":"sendmsg",  // API Gatewayのルート名（ここでは"sendmsg"）と合わせる必要がある
                    "data": input.value
                }
            ));
            input.value = "";
        };
    </script>
</body>
</html>

動作確認

ブラウザで動作確認します。

以上で、AWSのサーバーレス（Python）でWebSocket通信をすることができました。