【Rust】AWS LambdaでWebSocketクライアントにメッセージ送信

AWS LambdaをRustで実装する際のサンプルコードが少なすぎるのでメモ

内容

LambdaからAPI Gateway経由でWebSocketクライアントにメッセージを送信する。
送信先クライアントのコネクションIDは入力イベントやDBで動的に取得することが多いだろうが、ここでは簡易化のため固定とする。

コード

Cargo.toml

[package]
name = "sample"
version = "0.1.0"
edition = "2021"

[dependencies]
aws-config = "0.9.0"
aws-sdk-apigatewaymanagement = "0.9.0"
aws-types = "0.9.0"
http = "0.2.6"
lambda_runtime = "0.5.1"
serde_json = "1.0.79"
tokio = "1.17.0"

main.rs

use aws_config::meta::region::RegionProviderChain;
use aws_sdk_apigatewaymanagement::types::Blob;
use aws_sdk_apigatewaymanagement::{config, Client as ApiGwClient, Endpoint};
use aws_types::region::Region;
use lambda_runtime::{service_fn, Error, LambdaEvent};
use serde_json::{json, Value};

#[tokio::main]
async fn main() -> Result<(), Error> {
    let func = service_fn(func);
    lambda_runtime::run(func).await?;
    Ok(())
}

async fn func(_event: LambdaEvent<Value>) -> Result<Value, Error> {
    // リージョン設定の読み込み
    let region = "ap-northeast-1";
    let region_provider = RegionProviderChain::first_try(Region::new(region));
    let shared_config = aws_config::from_env().region(region_provider).load().await;

    // API Gateway クライアント用意
    // 対象　API Gateway　のエンドポイントURI(要変更)
    let api_uri = "https://<id>.execute-api.<region>.amazonaws.com/<stage>";
    let endpoint = Endpoint::immutable(http::Uri::from_static(api_uri));
    let api_management_config = config::Builder::from(&shared_config)
        .endpoint_resolver(endpoint)
        .build();
    let client = ApiGwClient::from_conf(api_management_config);

    // メッセージを送信
    let connection_id = "OFWR4eiLtjMCIpw="; // 送信先クライアントのコネクションID(要変更)
    let data = r#"{"message": "hello from lambda"}"#; // メッセージ内容
    client
        .post_to_connection()
        .connection_id(connection_id)
        .data(Blob::new(data))
        .send()
        .await?;

    Ok(json!({ "message": "success" }))
}