Lambdaをローカルで実行する時だけ引数を追加したい(Go言語) - 株式会社CINC 開発本部(旧：開発部) エンジニアブログ

Lambdaをローカルで実行する時だけ引数を追加したい(Go言語) こんにちは、株式会社CINC開発本部です。

今回は、AWS Lambdaに関するTipsです。

はじめに

Go言語でAWS Lambdaの関数に手を入れる際に、「事故を起こしそうな項目があるので、eventファイルや環境変数には入れたくない。対象のローカルでしか実行できないようにしたい。」という要望がありました。

そこで、コマンドラインから実行するときに特定のフラグを立てた場合だけ動く処理を追加しようと考えました。

前提・背景

AWS Lambdaをローカルでテストする場合、「本番環境で受け取るイベントデータをあらかじめファイル（eventファイル）として用意しておき、handlerに渡す」という方法をよく使います。今回は、この eventファイルを用いて簡易的にローカルテストを行う仕組みが既に存在していました。

しかし、この仕組みだと「フラグを指定して動的に処理を切り替える」といった柔軟な対応が難しかったのです。大規模なリファクタリングは避けたい、という事情もあり、context を使って最小限の変更で目的を達成できないか検討しました。

やったこと

1. 関数の呼び出し部分の調整

下記のように main 関数で、環境変数 ENVIRONMENT をチェックし、

local だった場合は handlerWrap それ以外なら lambda.Start(handler) を呼ぶようにしています。

func main() {
    env := os.Getenv("ENVIRONMENT")
    switch env {
    case "local":
        res, err := handlerWrap(context.Background())
        if err != nil {
            os.Exit(1)
        }
    default:
        lambda.Start(handler)
    }
}

2. コマンドライン引数の受け取り

ローカル実行時に呼ばれる handlerWrap 関数は以下のようになっています。（簡単にするためエラー処理は省略）

type ctxTestKey struct{}

func handlerWrap(ctx context.Context) (*string, error) {
    // ① コマンドライン引数の取得
    eventPath := flag.String("event", "", "path to event file")
    testFlag := flag.Bool("test", false, "test flag")
    flag.Parse()

    // ② contextへの値の設定
    if *testFlag {
        ctx = context.WithValue(ctx, ctxTestKey{}, true)
    }

    // ③ eventファイルを読み込んで handlerを呼ぶ
    f, err := os.ReadFile(*eventPath)
    var ev event
    if err := json.Unmarshal(f, &ev); err != nil {
        return nil, err
    }
    res, err := handler(ctx, &ev)
    // 略
}

flag パッケージを使って、コマンドラインから渡される --event と --test の値を取得しています。 --test が true のときだけ context.WithValue で context に値を設定しています。 ctxTestKey という独自の型をkeyにすることで、パッケージ間衝突を防いでいます。

読み込んだ eventファイルを json.Unmarshal で構造体に変換し、 handler 関数を呼び出しています。

これにより、Lambdaで受け取る引数とほぼ同じ形でローカルテストが可能です。

3. contextからコマンドライン引数の取り出し

handler 関数側では、以下のように ctx.Value で値を取り出しています。返ってくるのは interface{} 型なので、キャストが必要ですが、今回は nil かどうかだけを見れば十分です。

func handler(ctx context.Context, ev *event) (*Response, error) {
    if ctx.Value(ctxTestKey{}) != nil {
        // 引数が設定された時だけ実行したいコード
    }
    // 略
}

まとめ

今回は context を利用して、「コマンドラインのフラグが立ったときだけ実行する処理」をローカル限定で表現しました。ただし、公式ドキュメントには「contextはリクエストスコープのデータにのみ使うことが推奨されている」と書かれています。（たとえばリクエストヘッダから渡されたIDやセッショントークンなど。） context package - context - Go Packages

今回のように、「関数のオプションのデータを渡す」目的でのcontext利用は、あまりお行儀が良くない方法です。（ローカル限定データとはいえ、厳密にはリクエストスコープとは言いづらい面もあります。）

ベストプラクティスとしては、リファクタリングを行って、handler内のロジックを別の関数に切り出し、コマンドラインからはその関数を直接呼ぶようにするのが望ましいでしょう。とはいえ、今回のケースでは大掛かりな変更を避けたい事情もあったため、contextを使う実装を選びました。