Rustで書いたAWS Lambda関数からS3 Selectしてみた

こんにちは。サービスグループの武田です。

先日RusotoがサポートしていないS3 Selectを、AWS CLIを外部コマンドとして呼び出すことでRustから実行するエントリを書きました。

今回は同様のことを、AWS Lambdaでもやってみました。

なお、ソースコードはGitHubに上がっています。

TAKEDA-Takashi/rust-lambda-call-aws-cli

検証環境

$ aws --version
aws-cli/2.1.10 Python/3.7.4 Darwin/19.6.0 exe/x86_64 prompt/off

$ rustc -V
rustc 1.50.0 (cb75ad5db 2021-02-10)

$ cargo -V
cargo 1.50.0 (f04e7fab7 2021-02-04)

$ docker --version
Docker version 20.10.0, build 7287ab3

$ sw_vers
ProductName:	Mac OS X
ProductVersion:	10.15.7
BuildVersion:	19H15

AWS LambdaでAWS CLIを実行するためには

まず大前提としてAWS CLIがインストールされている環境というものが必要です。RustをLambda関数として実行するためにはカスタムランタイムを利用します。詳しいことは平川のセッションが参考になりますので、ご覧ください。

さてカスタムランタイムのベースイメージはAmazon Linux 2ですので、もしかしてAWS CLI使えるのかな？という淡い期待もあったのですが、試してみると次のようなエラーが発生してダメでした。

{
  "errorType": "alloc::boxed::Box<dyn std::error::Error+core::marker::Sync+core::marker::Send>",
  "errorMessage": "Os { code: 2, kind: NotFound, message: \"No such file or directory\" }"
}

そんなわけでAWS CLIがインストールされた、Rustアプリケーションが実行できるLambda環境を用意する必要があります。皆さんは心当たりがありますね？そうです、Dockerです。

Rustアプリケーションのコンテナイメージ実行は、平川がエントリを書いていましたので、これを参考に進めましょう。

やってみた

まずはRustのプロジェクトを作成します。適当なディレクトリに移動してコマンドを実行しましょう。

$ cargo new rust-lambda-call-aws-cli

プロジェクトが作成できたら、必要な依存関係を書いておきます（dependencies以外は省略）。RustのLambda Runtimeはawslabsで公開されているRuntimeがあります。ただ最近は、Tokio 1.xへの対応遅れなど動きが遅いため、NetlifyがForkしたRuntimeを使用しました。

Cargo.toml

[dependencies]
lamedh_runtime = "0.3"
tokio = { version = "1", features = ["full"] }
serde = { version = "1", features = ["derive"] }
serde_json = "1"
log = "0.4"
env_logger = "0.8"

コードは先日のものをコピーし、Lambda用に少し書き換えます。またS3バケットやサンプルデータはそのときのものを再利用します。

main.rs

use lamedh_runtime::{Context, Error};
use log::{debug, info};
use serde::Deserialize;
use serde_json::Value;
use std::env;
use tokio::process::Command;

#[derive(Debug, Deserialize)]
struct TestData {
    name: String,
    code: u32,
    tags: Option<String>,
    lang: Option<String>,
}

#[tokio::main]
async fn main() -> Result<(), Error> {
    env::set_var("RUST_LOG", "rust_lambda_call_aws_cli=debug");
    env_logger::init();

    lamedh_runtime::run(lamedh_runtime::handler_fn(handler)).await?;

    Ok(())
}

async fn handler(_: Value, _: Context) -> Result<(), Error> {
    debug!("handler start");

    let output = Command::new("aws")
        .args(&[
            "s3api",
            "select-object-content",
            "--bucket=testdata-xxxx",
            "--key=test_data.json",
            "--input-serialization",
            r#"{"JSON":{"Type":"LINES"}}"#,
            "--output-serialization",
            r#"{"JSON":{"RecordDelimiter":"\n"}}"#,
            "--expression",
            "SELECT * FROM s3object s LIMIT 5",
            "--expression-type=SQL",
            "/tmp/output.json",
        ])
        .output()
        .await?;

    debug!("{:?}", output);

    if Some(0) != output.status.code() {
        panic!("{:?}", output);
    }

    let contents = tokio::fs::read("/tmp/output.json").await?;
    for line in String::from_utf8(contents)?.lines() {
        let d: TestData = serde_json::from_str(line)?;
        info!("{:?}", d);
    }

    Ok(())
}

次にRustコードのビルド用イメージと、実行用イメージをそれぞれ用意するためにDockerfileを作成します。まずはビルド用。

Dockerfile.build

FROM public.ecr.aws/lambda/provided:al2

# リンカーとしてgccを利用する
RUN yum install -y gcc

# rustupでRustツールチェーンをインストールする
RUN curl https://sh.rustup.rs -sSf | sh -s -- -y --default-toolchain stable
ENV PATH $PATH:/root/.cargo/bin
RUN rustup install stable

# ビルド対象のソースツリーをマウントする
VOLUME /code

# ローカル環境にRustを導入している場合は以下をコメントアウトするとビルドが早くなります
VOLUME /root/.cargo/registry
VOLUME /root/.cargo/git

WORKDIR /code
# provided:al2 はランタイム用の設定になっているので、ENTRYPOINTをビルド用に書き換える
ENTRYPOINT ["cargo", "build", "--release"]

次に実行用。この中でAWS CLIをインストールしているのがポイントですね。なお、2系がどうしても使いたかったわけではなく、pipでインストールしようとしたらコマンドがないと怒られたため、しかたなく手動でインストールしました。

Dockerfile

FROM public.ecr.aws/lambda/provided:al2

RUN yum install -y unzip

RUN curl "https://awscli.amazonaws.com/awscli-exe-linux-x86_64.zip" -o "awscliv2.zip" \
    && unzip awscliv2.zip \
    && ./aws/install

# 実行ファイルを起動するようにするため、ファイル名を "bootstrap" に変更する
COPY ./target/release/rust-lambda-call-aws-cli ${LAMBDA_RUNTIME_DIR}/bootstrap

# カスタムランタイム同様ハンドラ名は利用しないため、適当な文字列を指定する。
CMD ["lambda-handler"]

これで諸々準備が完了しました。次のコマンドを順番に実行していきます。

# ビルド用イメージをビルドします
$ docker image build -t rust-lambda-call-aws-cli-build -f Dockerfile.build .

# ビルド用コンテナを使ってRustアプリケーションをビルドします
$ docker run --rm \
    -v $PWD:/code \
    -v $HOME/.cargo/registry:/root/.cargo/registry \
    -v $HOME/.cargo/git:/root/.cargo/git \
    rust-lambda-call-aws-cli-build

# 実行用イメージをビルドします
$ docker image build -t "rust-lambda-call-aws-cli" .

# ECRリポジトリを作成します（この作業は一度のみ）
$ aws ecr create-repository --repository-name rust-lambda-call-aws-cli

# リポジトリにログインします
$ aws ecr get-login-password | docker login --username AWS --password-stdin 123456789012.dkr.ecr.ap-northeast-1.amazonaws.com/rust-lambda-call-aws-cli

# リポジトリに実行用イメージをプッシュします
$ docker image tag rust-lambda-call-aws-cli:latest 123456789012.dkr.ecr.ap-northeast-1.amazonaws.com/rust-lambda-call-aws-cli:latest

イメージの準備ができましたので、実際に実行してみましょう。マネジメントコンソールでLambda関数を作成し、用意したイメージを指定します（細かい手順は省略します。s3:GetObjectの権限が必要なので、そこだけ注意）。

初期値で実行するとタイムアウトしてしまいましたので、メモリ割り当てとタイムアウトを変更します。

最後に、適当なテストデータを指定して実行してみましょう。

無事に実行できました！

まとめ

なんとかRusotoを使わずにLambdaでS3 Selectできないものかなと、試行錯誤してみました。結果としては無事に取得できたのですが、予想より実行時間が長い結果となりました。とはいえ比較できるデータも多くないため、RusotoがS3 Selectをサポートしたら、実行時間の差分も検証してみます。