AWS Step Functionsでステート間の変数を共有する方法

AWS Step Functionsを使用する際、ステート間で変数を共有したい場合があります。Step FunctionsのAssignを使用して、あるステートの実行結果を別のステートで参照する方法について説明します。

サンプルとして、UUIDを生成するLambda関数を2回実行し、それぞれの結果を組み合わせて最終的な出力を生成するステートマシンを作成します。

ステートマシン

[!NOTE]
意図的にコメント入れていますが、ステートマシンはコメントを入れるとエラーになります。

{
    "Comment": "A description of my state machine",
    "StartAt": "SampleFirstLambda",
    "QueryLanguage": "JSONata",
    "States": {
        "SampleFirstLambda": {
            "Type": "Task",

            // 実行するLambda
            "Resource": "arn:aws:lambda:ap-northeast-1:999999999999:function:test-for-statemachine:$LATEST",

            // Lambdaに渡す引数
            // ステートマシンのパラメータは $state.input で取得可能
            "Arguments": {
                "sampleParam1": "{% $states.input.sampleParam1 %}"
            },

            // Lambdaの実行結果
            // $state.result で取得可能
            "Output": {
                "sampleParam1": "{% $states.input.sampleParam1 %}",
                "sampleParam2": "{% $states.input.sampleParam2 %}",
                "sampleResultId": "{% $states.result.hoge_id %}"
            },

            // ステートマシン内で扱う変数を定義できる
            "Assign": {
                "sampleFirstResult": {
                    "sampleParam1": "{% $states.input.sampleParam1 %}",
                    "sampleParam2": "{% $states.input.sampleParam2 %}",
                    "sampleResultId": "{% $states.result.hoge_id %}"
                }
            },
            "Next": "SampleSecondLambda"
        },
        "SampleSecondLambda": {
            "Type": "Task",
            "Resource": "arn:aws:lambda:ap-northeast-1:999999999999:function:test-for-statemachine:$LATEST",
            "Arguments": {
                "sampleParam1": "{% $states.input.sampleParam1 %}",
                "sampleParam2": "{% $states.input.sampleParam2 %}"
            },
            "Output": {
                "sampleParam1": "{% $states.input.sampleParam1 %}",
                
                // Assign で定義した変数を使っている
                "firstResultId": "{% $sampleFirstResult.sampleResultId %}",
                
                "secondResultId": "{% $states.result.hoge_id %}"
            },
            "End": true
        }
    }
}

Lambda

UUIDを生成して返すLambdaを作成しました。

import json
import uuid

def lambda_handler(event, context):
    event["hoge_id"] = str(uuid.uuid4())

    return event

ステートマシンのパラメータ

{
  "sampleParam1": "abc123",
  "sampleParam2": "xyz456"
}

結果

{
  "sampleParam1": "abc123",
  "firstResultId": "0961c121-7907-43a5-94e4-3d483b5aa61f",
  "secondResultId": "07e04490-7aec-43f1-b86f-da859c04f93b"
}