使用AWS Chalice构建WebSocket聊天服务器教程

使用AWS Chalice构建WebSocket聊天服务器教程

chalice Python Serverless Microframework for AWS 项目地址: https://gitcode.com/gh_mirrors/ch/chalice

概述

本文将详细介绍如何使用AWS Chalice框架构建一个基于WebSocket的聊天服务器应用。这个应用将实现多聊天室功能，用户可以使用昵称加入不同房间进行交流。我们将利用DynamoDB存储连接状态信息，并通过API Gateway WebSocket API实现实时通信。

环境准备

首先需要创建一个干净的Python虚拟环境并安装必要的依赖：

pip install -U chalice
chalice new-project chalice-chat-example
cd chalice-chat-example

由于我们的应用需要使用boto3与DynamoDB交互并通过API Gateway发送WebSocket消息，需要将其添加到依赖中：

echo "boto3>=1.9.91" > requirements.txt
pip install -r requirements.txt

项目结构

完成上述步骤后，项目目录结构如下：

.
├── .chalice
│   └── config.json
├── .gitignore
├── app.py
├── resources.json
└── requirements.txt

自定义CloudFormation模板

Chalice 1.10及以上版本支持通过--merge-template参数合并自定义JSON到生成的CloudFormation模板中。我们需要创建一个resources.json文件来定义DynamoDB表和相关配置：

{
  "Resources": {
    "ChaliceChatTable": {
      "Type": "AWS::DynamoDB::Table",
      "Properties": {
        "AttributeDefinitions": [
          {"AttributeName": "PK", "AttributeType": "S"},
          {"AttributeName": "SK", "AttributeType": "S"}
        ],
        "KeySchema": [
          {"AttributeName": "PK", "KeyType": "HASH"},
          {"AttributeName": "SK", "KeyType": "RANGE"}
        ],
        "GlobalSecondaryIndexes": [
          {
            "IndexName": "ReverseLookup",
            "KeySchema": [
              {"AttributeName": "SK", "KeyType": "HASH"},
              {"AttributeName": "PK", "KeyType": "RANGE"}
            ],
            "Projection": {"ProjectionType": "ALL"},
            "ProvisionedThroughput": {
              "ReadCapacityUnits": 1,
              "WriteCapacityUnits": 1
            }
          }
        ],
        "ProvisionedThroughput": {
          "ReadCapacityUnits": 1,
          "WriteCapacityUnits": 1
        },
        "TableName": "ChaliceChat"
      }
    },
    "WebsocketConnect": {
      "Properties": {
        "Environment": {"Variables": {"TABLE": {"Ref": "ChaliceChatTable"}}}
      }
    },
    "WebsocketMessage": {
      "Properties": {
        "Environment": {"Variables": {"TABLE": {"Ref": "ChaliceChatTable"}}}
      }
    },
    "WebsocketDisconnect": {
      "Properties": {
        "Environment": {"Variables": {"TABLE": {"Ref": "ChaliceChatTable"}}}
      }
    },
    "DefaultRole": {
      "Type": "AWS::IAM::Role",
      "Properties": {
        "AssumeRolePolicyDocument": {
          "Version": "2012-10-17",
          "Statement": [
            {
              "Sid": "",
              "Effect": "Allow",
              "Principal": {"Service": "lambda.amazonaws.com"},
              "Action": "sts:AssumeRole"
            }
          ]
        },
        "Policies": [
          {
            "PolicyName": "DefaultRolePolicy",
            "PolicyDocument": {
              "Version": "2012-10-17",
              "Statement": [
                {
                  "Effect": "Allow",
                  "Action": [
                    "logs:CreateLogGroup",
                    "logs:CreateLogStream",
                    "logs:PutLogEvents"
                  ],
                  "Resource": "arn:aws:logs:*:*:*"
                },
                {
                  "Effect": "Allow",
                  "Action": ["execute-api:ManageConnections"],
                  "Resource": "arn:aws:execute-api:*:*:*/@connections/*"
                },
                {
                  "Effect": "Allow",
                  "Action": [
                    "dynamodb:DeleteItem",
                    "dynamodb:PutItem",
                    "dynamodb:GetItem",
                    "dynamodb:UpdateItem",
                    "dynamodb:Query",
                    "dynamodb:Scan"
                  ],
                  "Resource": [
                    {"Fn::Sub": "arn:aws:dynamodb:${AWS::Region}:${AWS::AccountId}:table/${ChaliceChatTable}"},
                    {"Fn::Sub": "arn:aws:dynamodb:${AWS::Region}:${AWS::AccountId}:table/${ChaliceChatTable}/index/ReverseLookup"}
                  ]
                }
              ]
            }
          }
        ]
      }
    }
  }
}

主应用代码

app.py是应用的主入口文件，主要处理WebSocket的连接、断开和消息事件：

from boto3.session import Session
from chalice import Chalice
from chalicelib import Storage, Sender, Handler

app = Chalice(app_name="chalice-chat-example")
app.websocket_api.session = Session()
app.experimental_feature_flags.update(['WEBSOCKETS'])

STORAGE = Storage.from_env()
SENDER = Sender(app, STORAGE)
HANDLER = Handler(STORAGE, SENDER)

@app.on_ws_connect()
def connect(event):
    STORAGE.create_connection(event.connection_id)

@app.on_ws_disconnect()
def disconnect(event):
    STORAGE.delete_connection(event.connection_id)

@app.on_ws_message()
def message(event):
    HANDLER.handle(event.connection_id, event.body)

核心功能实现

在chalicelib目录中，我们实现了三个核心类：

1. Storage类

负责与DynamoDB交互，管理连接状态：

class Storage(object):
    def __init__(self, table):
        self._table = table

    @classmethod
    def from_env(cls):
        table_name = os.environ.get('TABLE', '')
        table = boto3.resource('dynamodb').Table(table_name)
        return cls(table)

    def create_connection(self, connection_id):
        self._table.put_item(
            Item={'PK': connection_id, 'SK': 'username_'}
        )

    def set_username(self, connection_id, old_name, username):
        self._table.delete_item(
            Key={'PK': connection_id, 'SK': 'username_%s' % old_name}
        )
        self._table.put_item(
            Item={'PK': connection_id, 'SK': 'username_%s' % username}
        )

    def list_rooms(self):
        r = self._table.scan()
        rooms = set([item['SK'].split('_', 1)[1] for item in r['Items']
                     if item['SK'].startswith('room_')])
        return rooms

    def set_room(self, connection_id, room):
        self._table.put_item(
            Item={'PK': connection_id, 'SK': 'room_%s' % room}
        )

    def remove_room(self, connection_id, room):
        self._table.delete_item(
            Key={'PK': connection_id, 'SK': 'room_%s' % room}
        )

    def get_connection_ids_by_room(self, room):
        r = self._table.query(
            IndexName='ReverseLookup',
            KeyConditionExpression=(Key('SK').eq('room_%s' % room)),
            Select='ALL_ATTRIBUTES',
        )
        return [item['PK'] for item in r['Items']]

    def delete_connection(self, connection_id):
        try:
            r = self._table.query(
                KeyConditionExpression=(Key('PK').eq(connection_id)),
                Select='ALL_ATTRIBUTES',
            )
            for item in r['Items']:
                self._table.delete_item(
                    Key={'PK': connection_id, 'SK': item['SK']}
                )
        except Exception as e:
            print(e)

    def get_record_by_connection(self, connection_id):
        r = self._table.query(
            KeyConditionExpression=(Key('PK').eq(connection_id)),
            Select='ALL_ATTRIBUTES',
        )
        return {
            entry['SK'].split('_', 1)[0]: entry['SK'].split('_', 1)[1]
            for entry in r['Items']
        }

2. Sender类

负责通过WebSocket发送消息：

class Sender(object):
    def __init__(self, app, storage):
        self._app = app
        self._storage = storage

    def send(self, connection_id, message):
        try:
            self._app.websocket_api.send(connection_id, message)
        except WebsocketDisconnectedError as e:
            self._storage.delete_connection(e.connection_id)

    def broadcast(self, connection_ids, message):
        for cid in connection_ids:
            self.send(cid, message)

3. Handler类

处理接收到的WebSocket消息并执行相应操作：

class Handler(object):
    def __init__(self, storage, sender):
        self._storage = storage
        self._sender = sender
        self._command_table = {
            'help': self._help,
            'nick': self._nick,
            'join': self._join,
            'room': self._room,
            'quit': self._quit,
            'ls': self._list,
        }

    def handle(self, connection_id, message):
        record = self._storage.get_record_by_connection(connection_id)
        # 处理消息逻辑...

部署与测试

完成代码编写后，可以使用以下命令打包并部署应用：

chalice package --merge-template resources.json --pkg-format cloudformation out/

然后使用AWS CLI部署生成的CloudFormation模板。

总结

通过本教程，我们学习了如何使用AWS Chalice框架构建一个功能完整的WebSocket聊天应用。这个应用展示了：

1. WebSocket API的基本使用方法
2. DynamoDB作为状态存储的实现
3. 多聊天室和用户昵称功能
4. 消息广播机制

虽然这个示例为了演示目的简化了一些实现，但它提供了构建实时应用的良好起点。开发者可以在此基础上扩展更多功能，如消息持久化、用户认证等。

chalice Python Serverless Microframework for AWS 项目地址: https://gitcode.com/gh_mirrors/ch/chalice