AWS Step Functionsで5分後にDynamoDBのデータを削除する

はじめに

API Gateway + Lambda + DynamoDBでwebAPIを開発しています。
Lambda関数でDynamoDBへデータを書き込み、そのデータを5分後に削除する必要がありました。
そこで、以下の方法を考えました。
・DyanamoDBのTTL機能を使って5分後にデータを削除する
・データを書き込むLambdaから5分後にデータを削除するLambdaを呼び出す

DynamoDBのTTL機能

まずDynamoDBのTTL機能を調べたのですが、TTLでは有効期限から48時間以内に期限切れのアイテムを削除しますとの記載があり、5分後に必ず削除する設定は不可能でした。
https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/howitworks-ttl.html

AWS Step Functionsを使う

次に、データを書き込むLambdaからデータを削除するLambdaを呼び出すを検討します。
調べた結果、AWS Step Functionsを使うことで実現できそう。実装してみます。

データを書き込むLambda関数の作成

ランタイムはPython3.8、DynamoDBとStepFunctionsへのアクセス権を付与します。
以下のLambda関数でDynamoDBへデータを書き込み、StateMachineを呼び出します。
userIdをパラメータとしてStateMachineに渡します。

lambda_function.py

import  boto3

client = boto3.client('stepfunctions')
dynamoDB = boto3.resource('dynamodb')

def lambda_handler(event, context):
    # データを書き込むテーブル
    table = dynamoDB.Table('test')
    # 呼び出すStateMachine
    stateMachineArn = 'StateMachineのARN'
    # DynamoDBにデータを書き込む
    table.put_item(Item={'data': data})
    # StateMachineを呼び出す
    input = {"parameters":{"userId":userId}}
    client.start_execution(
        stateMachineArn = stateMachineArn,
        input = json.dumps(input)
        )
    return {'statusCode': 200, 'body': "writeData"}

StateMachineの作成

以下の定義でStateMachineを作成します。
StateMachineにLambdaへのアクセス権を付与してください。

{
  "StartAt": "wait_5_miutes",
  "States": {
    "wait_5_miutes": {
      "Type": "Wait",
      "Seconds": 300,
      "Next": "deleteData"
    },
    "deleteData": {
      "Type": "Task",
      "Resource": "データを削除するLambda関数のARN",
      "Parameters": {
          "userId.$": "$.parameters.userId"
      },
      "End": true
    }
  }
}

データを削除するLambda関数の作成

StateMachineから、以下のLambda関数を実行することでデータを削除します。

lambda_function.py

import boto3

dynamoDB = boto3.resource('dynamodb')
table = dynamoDB.Table('test')

# StateMachineで呼び出される
def lambda_handler(event, context):
    # StateMachineからパラメータを受け取る
    userId = event['userId']
    # 削除対象のデータを探す
    searchResult = table.get_item(Key={'userId': userId})
    if "Item" in searchResult:
        table.delete_item(Key={'userId': userId})
    return {'statusCode': 200, 'body': "deleteData"}

おわりに

AWS Step Functionsを使うことで、DynamoDBへデータを書き込み、そのデータを5分後に削除するが実現できました。
もっと複雑な処理もできそうなので試してみたいです。

参考

https://qiita.com/ketancho/items/147a141c9f8a6de86c97
https://aws.amazon.com/jp/getting-started/hands-on/create-a-serverless-workflow-step-functions-lambda/