Building Custom Resources with Lambda in AWS Amplify (Gen 1) via CloudFormation

technologies

Important Tip

Before making any changes to your CloudFormation template or parameters file, run amplify env checkout to ensure that changes are detected when executing amplify push. If you skip this, the CLI might miss your updates during amplify status.

Why this tutorial?

The AWS Amplify CLI, while powerful, doesn’t natively support every AWS service. For instance, while it’s easy to add a Lambda function using amplify add function, there’s no corresponding command to add Step Functions (amplify add sfn). This limitation applies to other AWS services not supported by the CLI, requiring you to create custom resources.

To address these gaps, Amplify provides two primary methods for integrating custom resources:

  1. Using the Cloud Development Kit (CDK): Write custom resources in JavaScript, which are then converted into CloudFormation templates.
  2. Using CloudFormation directly: Provide a custom CloudFormation template (in JSON format) to be deployed via amplify push.

While both methods extend the functionality of Amplify apps, I’ve discovered a third, less-documented approach that utilizes custom Cloudformation templates in YAML format. I personally find YAML easier to read and perfer to use it. This tutorial focuses on that approach.

Architecture

Declaring the Step Functions State Machine

To integrate a Step Functions (SFN) state machine into an Amplify-based application, begin by declaring the SFN in Amplify’s toolchain. Amplify categories are defined in the amplify/backend/backend-config.json file. You can declare the SFN similarly to other categories like api, auth, and function.

Here’s an example configuration:

{
    "api": {},
    "auth": {},
    "function": {
        "exampleHandler": {
            "providerPlugin": "awscloudformation",
            "service": "Lambda"
        }
    },
    "sfn": {
        "stateMachine": {
            "dependsOn": [ // list of resources that this resource access or depends on
                {
                    "attributes": [ // list of attributes (variables output) that this resource access or depends on
                        "Arn",
                        "Name"
                    ],
                    "category": "function",
                    "resourceName": "exampleHandler"
                }
            ],
            "providerPlugin": "awscloudformation",
            "service": "StepFunctions"
        }
    }
}

In this setup, we define a custom category called sfn and a resource named stateMachine. Amplify will follow the same naming conventions for your custom resources as it does for other categories. To reference outputs from other resources, declare them as parameters in your template and add them to the dependsOn array in backend-config.json.

Adding the SFN State Machine

Next, mimic Amplify’s directory structure for other categories and backend resources:

/amplify
  |- backend
      |- sfn
          |- stateMachine
              |- template.yml
              |- parameters.json

Both template.yml and parameters.json are required. The CloudFormation template (template.yml) defines the resources needed for the SFN state machine. Here’s an example:

AWSTemplateFormatVersion: '2010-09-09'
Parameters:
  env:
    Type: String
# the format <category><resourceName><output_variables>
  functionexampleHandlerArn: 
    Type: String
  functionexampleHandlerName: 
    Type: String
# <some others parameters>
Resources:
  StateMachineRole:
    Type: AWS::IAM::Role
    Properties:
      RoleName:
        Fn::Join:
          - ''
          - - sfn-
            - Ref: env
            - -role
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
          - Effect: Allow
            Principal:
              Service:
                - states.amazonaws.com
            Action:
              - sts:AssumeRole
      Policies:
        - PolicyName: LambdaInvokeScopedAccessPolicy
          PolicyDocument:
            Statement:
              - Effect: Allow
                Action:
                  - lambda:InvokeFunction
                Resource:
                  - Fn::Sub: ${functionExampleHandlerArn}:*
              - Effect: Allow
                Action:
                  - lambda:InvokeFunction
                Resource:
                  - Ref: functionExampleHandlerArn
        - PolicyName: XRayAccessPolicy
  	 PolicyDocument:
    	 Version: '2012-10-17'
    	 Statement:
      	   - Effect: Allow
              Action:
                - xray:PutTraceSegments
                - xray:PutTelemetryRecords
                - xray:GetSamplingRules
                - xray:GetSamplingTargets
              Resource: '*'
  StateMachine:
   Type: AWS::StepFunctions::StateMachine
   Properties:
    StateMachineName:
      Fn::Join:
        - ''
        - - state-machine-
          - Ref: env
     RoleArn:
       Fn::GetAtt:
         - StateMachineRole
         - Arn
# <other config>
Outputs:
  Arn:
    Description: ARN of the state machine
    Value:
      Fn::GetAtt:
        - StateMachine
        - Arn

The parameters.json file should contain an empty object {}.

Letting the CLI know about Custom Resource

After setting up your custom category and resources, notify the CLI by running

amplify env checkout <current-env-name> and then amplify push.

This will deploy the SFN state machine to your Amplify project.

Adding a Lambda Function with SFN Dependency

Now, let’s add a Lambda function that will trigger the Step Function. Use the Amplify CLI to add the Lambda function:

amplify add function

In your backend-config.json, make sure to indicate the dependency between the Lambda function and the Step Function. This ensures that the Lambda function is created after the SFN machine state:

{
    "api": {},
    "auth": {},
    "function": {
        "exampleHandler": {
            "providerPlugin": "awscloudformation",
            "service": "Lambda"
        },
	"sfnHandler": {
		"dependsOn": [ // list of resources that this resource access or depends on
                {
                    "attributes": [ // list of attributes (variables output) that this resource access or depends on
                        "Arn"
                    ],
                    "category": "sfn",
                    "resourceName": "stateMachine"
                }
            ],
            "providerPlugin": "awscloudformation",
            "service": "Lambda"
        },

    },
    "sfn": {
        "stateMachine": {
            "dependsOn": [ // list of resources that this resource access or depends on
                {
                    "attributes": [ // list of attributes (variables output) that this resource access or depends on
                        "Arn",
                        "Name"
                    ],
                    "category": "function",
                    "resourceName": "exampleHandler"
                }
            ],
            "providerPlugin": "awscloudformation",
            "service": "StepFunctions"
        }
    }
}

Here, the sfnHandler Lambda function depends on the SFN machine state, ensuring the resources are created in the proper order.

Linking the SFN Machine State with Lambda

Next, integrate the SFN state machine with the Lambda function. In the CloudFormation template for the Lambda function (amplify/backend/function/sfnHandler/sfnHandler-cloudformation-template.json), add the following environment variable configuration:

{
  "AWSTemplateFormatVersion": "2010-09-09",
  "Parameters": {
    "env": {
      "Type": "String"
    },
# the format <category><resourceName><output_variables>
    "sfnstateMachineArn": {
      "Type": "String",
      "Default": "sfnstateMachineArn"
    },
# <some others parameters>
},
  "Conditions": {
    "ShouldNotCreateEnvResources": {
      "Fn::Equals": [
        {
          "Ref": "env"
        },
        "NONE"
      ]
    }
  },
  "Resources": {
    "LambdaFunction": {
      "Type": "AWS::Lambda::Function",
      "Metadata": {
        "aws:asset:path": "./src",
        "aws:asset:property": "Code"
      },
      "Properties": {
# <other config>
        "Environment": {
          "Variables": {
            "ENV": {
              "Ref": "env"
            },
            "REGION": {
              "Ref": "AWS::Region"
            },
            "STATE_MACHINE_ARN": {
              "Ref": "sfnstateMachineArn"
            }
          }
        }
# <other config>
      }
    },
    "LambdaExecutionRole": {},
    "lambdaexecutionpolicy": {},
    "AmplifyResourcesPolicy": {},
    "CustomResourcePolicy": {
      "DependsOn": [
        "LambdaExecutionRole"
      ],
      "Type": "AWS::IAM::Policy",
      "Properties": {
        "PolicyName": "custom-resource-lambda-execution-policy",
        "Roles": [
          {
            "Ref": "LambdaExecutionRole"
          }
        ],
        "PolicyDocument": {
          "Version": "2012-10-17",
          "Statement": [
            {
              "Effect": "Allow",
              "Action": [
                "states:StartExecution"
              ],
              "Resource": {
                "Ref": "sfnstateMachineArn"
              }
            }
          ]
        }
      }
    }
  },
  "Outputs": {}
}

Use the STATE_MACHINE_ARN in your Lambda function’s handler (index.js):

const StateMachineArn = process.env.STATE_MACHINE_ARN

Finally, push the changes using the Amplify CLI:

amplify env checkout <current-env-name>

amplify push

After a successful push, your Lambda function and SFN state machine will be ready to work together in the cloud.

Related posts