CloudFormationでLambdaの非同期呼び出しを含むステートマシンを定義する

Parameters:
  LambdaSourceBucket:
    Type: String
    Default: lambda-sources-test1
    Description: Lambda関数を格納するS3バケット名

SampleDataTable:
    Type: AWS::DynamoDB::Table
    Properties:
      TableName: SampleDataTable
      AttributeDefinitions:
        - AttributeName: id
          AttributeType: S
      KeySchema:
        - AttributeName: id
          KeyType: HASH # パーティションキー
      BillingMode: PAY_PER_REQUEST

RegisterDataRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: RegisterDataRole
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: sts:AssumeRole
      Policies:
        - PolicyName: RegisterDataPolicy
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: Allow
                Action:
                  - dynamodb:PutItem
                Resource: !Sub arn:aws:dynamodb:${AWS::Region}:${AWS::AccountId}:table/${SampleDataTable}
              - Effect: Allow
                Action:
                  - logs:CreateLogGroup
                  - logs:CreateLogStream
                  - logs:PutLogEvents
                Resource: !Sub arn:aws:logs:${AWS::Region}:${AWS::AccountId}:log-group:/aws/lambda/RegisterData:*

InvokeRegisterDataRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: InvokeRegisterDataRole
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal:
              Service: lambda.amazonaws.com
            Action: sts:AssumeRole
      Policies:
        - PolicyName: InvokeRegisterDataPolicy
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: Allow
                Action:
                  - lambda:InvokeFunction
                Resource: !GetAtt RegisterData.Arn
              - Effect: Allow
                Action:
                  - logs:CreateLogGroup
                  - logs:CreateLogStream
                  - logs:PutLogEvents
                Resource: !Sub arn:aws:logs:${AWS::Region}:${AWS::AccountId}:log-group:/aws/lambda/InvokeRegisterData:*

RegisterData:
    Type: AWS::Lambda::Function
    Properties:
      FunctionName: RegisterData
      Handler: register_data.lambda_handler
      Runtime: python3.10
      Role: !GetAtt RegisterDataRole.Arn
      Code:
        S3Bucket: !Ref LambdaSourceBucket
        S3Key: register_data.zip
      Timeout: 300

InvokeRegisterData:
    Type: AWS::Lambda::Function
    Properties:
      FunctionName: InvokeRegisterData
      Handler: invoke_register_data.lambda_handler
      Runtime: python3.10
      Role: !GetAtt InvokeRegisterDataRole.Arn
      Code:
        S3Bucket: !Ref LambdaSourceBucket
        S3Key: invoke_register_data.zip
      Timeout: 30
      Environment:
        Variables:
          TARGET_FUNCTION_NAME: !Ref RegisterData
          TABLE_NAME: !Ref SampleDataTable

StepFunctionExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName: StepFunctionExecutionRole
      AssumeRolePolicyDocument:
        Version: "2012-10-17"
        Statement:
          - Effect: Allow
            Principal:
              Service: states.amazonaws.com
            Action: sts:AssumeRole
      Policies:
        - PolicyName: StepFunctionPolicy
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: Allow
                Action: lambda:InvokeFunction
                Resource: !GetAtt InvokeRegisterData.Arn

StateMachine:
    Type: AWS::StepFunctions::StateMachine
    Properties:
      StateMachineName: LambdaChainStarter
      RoleArn: !GetAtt StepFunctionExecutionRole.Arn
      DefinitionString:
        Fn::Sub: |
          {
            "StartAt": "ExecuteLambda",
            "States": {
              "ExecuteLambda": {
                "Type": "Task",
                "Resource": "arn:aws:states:::lambda:invoke",
                "Parameters": {
                  "FunctionName": "${InvokeRegisterData.Arn}",
                  "Payload": {
                    "Input.$": "$"
                  }
                },
                "End": true
              }
            }
          }

# invoke_register_data.py
import boto3
import os
import json


def lambda_handler(event, context):
    try:
        target_function_name = os.environ.get('TARGET_FUNCTION_NAME')
        table_name = os.environ.get('TABLE_NAME')
        payload = {
            'records': event['Input'],  # 登録データのリスト
            'table_name': table_name  # 登録先のテーブル名
        }
        client = boto3.client('lambda')
        res = client.invoke(
            FunctionName=target_function_name,
            InvocationType='Event',  # 非同期呼び出し
            Payload=json.dumps(payload)
        )
        return {'status': 'Target lambda invoked'}
    except Exception as e:
        print(f'Error occurred: {e}')
        raise

# register_data.py
import boto3
import time

dynamodb = boto3.resource('dynamodb')


def lambda_handler(event, context):
    try:
        data_list = event.get('records', [])
        table_name = event.get('table_name')
        if not data_list or not table_name:
            raise ValueError(
                'Invalid input: "records" and "table_name" are required.')
        table = dynamodb.Table(table_name)
        for item in data_list:
            table.put_item(Item=item)
        print(f'Data registration completed. - Total: {len(data_list)}')
        return {'status': 'done', 'count': len(data_list)}
    except Exception as e:
        print(f'Error occurred: {e}')
        raise