[AWS] Create a Step Function with CDK Typescript

Sample Step Function

We need a sample step function to work with and it has to have multiple steps. So we have added three lambda functions and a choice in-between them to ensure that Step Function has multiple steps and is complex enough and diverse enough.

Create Common Lambda Function

We will create 3 lambda functions and uniting the creation of lambda functions and lambda role into one function, which will help us create step function in the next steps.

  private createLambdaFunctionAndRole(
    functionName: string,
    emailNotificationTopic: any,
    sampleCodeBucket: any,
    envVariable: any = {},
    memorySize = 1024
  ): any {
    // Create an IAM role for the Lambda function
    const lambdaRole = new iam.Role(this, `${functionName}Role`, {
      assumedBy: new iam.ServicePrincipal('lambda.amazonaws.com'),
      description: `${functionName}Lambda Function Role`,
      roleName: `${functionName}LambdaFunctionRole`,
      managedPolicies: [
        iam.ManagedPolicy.fromAwsManagedPolicyName(
          'service-role/AWSLambdaBasicExecutionRole'
        )
      ]
    });
    emailNotificationTopic.grantPublish(lambdaRole);

    const lambdaFunction = new lambda.Function(this, functionName, {
      functionName,
      description: functionName,
      runtime: lambda.Runtime.PROVIDED_AL2,
      memorySize,
      timeout: cdk.Duration.minutes(15),
      retryAttempts: 0,
      handler: 'bootstrap',
      code: lambda.Code.fromBucket(
        sampleCodeBucket,
        'sample-golang-project/golang-sample.zip'
      ),
      environment: envVariable,
      role: lambdaRole,
      deadLetterTopic: emailNotificationTopic
    });

    // Add tags to DynamoDB Table
    cdk.Tags.of(lambdaFunction).add('site', 'phase-two-core');
    cdk.Tags.of(lambdaFunction).add('author', 'Nagarjun Nagesh');

    // Add CloudWatch logs permission to the Lambda function's role
    this.createLogGroups(lambdaFunction, functionName);

    return { lambdaFunction, lambdaRole };
  }

1. Add the Lambda Role for the lambda function.
2. Fetch the SNS for the Dead letter notification.
3. Create a Lambda Function.
4. Create Log Groups and add it to the Lambda Role.

This code can be used to create a lambda function.

1. “runtime:” — Choose the run time — in this case Golang / the programming language of your choice.
2. “handler:” — Choose the handler for the function.
3. “code:” — The code in this case is fetched from the bucket where sample code is stored.
4. “memorySize:” — Choose the memory size.
5. “timeout:” — The timeout of the lambda function.
6. “role:” — Add the role that has to be attached to the lambda function.
7. “deadLetterTopic:” — Add the SNS topic for your dead letter topic.

Create Step Function

  private createStepFunction(
    props: PropsPhaseTwoCore,
    emailNotificationTopic: any,
    sampleCodeBucket: any
  ): any {
    // Step 1: Lambda Function:
    const step1EnvVariable = {
      TABLE_NAME: "TableName"
    };
    const step1LambdaFunction = this.createLambdaFunctionAndRole(
      props.step1LambdaFunction,
      emailNotificationTopic,
      sampleCodeBucket,
      step1EnvVariable
    );

    // Step 2: Lambda Function
    const step2EnvVariable = {
      OTHER_ENV_VARIABLE: "VALUE"
    };
    const step2LambdaFunction = this.createLambdaFunctionAndRole(
      props.step2LambdaFunction,
      sampleCodeBucket,
      emailNotificationTopic,
      props,
      step2EnvVariable
    );
    

    // Step 3: Lambda Function
    const step3EnvVariable = {
      SOME_ENV_VARIABLE: "env Variable value"
    };
    const memorySize = 6144;
    const step3LambdaFunction = this.createLambdaFunctionAndRole(
      props.step3LambdaFunction,
      emailNotificationTopic,
      sampleCodeBucket,
      step3EnvVariable,
      memorySize
    );
    
    // Step Function Definition
    const definition = new tasks.LambdaInvoke(
      this,
      'step1LambdaFunction',
      {
        lambdaFunction: step1LambdaFunction.lambdaFunction,
        stateName: 'step1LambdaFunction',
        comment: 'Step 1 Lambda Function',
        invocationType: tasks.LambdaInvocationType.REQUEST_RESPONSE,
        outputPath: '$.Payload'
      }
    ).next(
      new sfn.Choice(this, 'Step1LambdaFunctionChoice', {
        stateName: 'Step1LambdaFunctionChoice',
        comment: 'Step 1 Lambda Function Choice',
        inputPath: '$',
        outputPath: '$'
      })
        .when(
          sfn.Condition.stringEquals('$.state', 'Yes'),
          new sfn.Succeed(this, 'ProcessEnded')
        )
        .otherwise(
          new tasks.LambdaInvoke(this, 'Step2LambdaFunction', {
            lambdaFunction: step2LambdaFunction.lambdaFunction,
            stateName: 'step2LambdaFunction',
            comment: 'Step 2 Lambda Function',
            invocationType: tasks.LambdaInvocationType.REQUEST_RESPONSE,
            outputPath: '$.Payload',
            inputPath: '$'
          }).next(
            new tasks.LambdaInvoke(
              this,
              'step3LambdaFunction',
              {
                lambdaFunction:
                  step3LambdaFunction.lambdaFunction,
                stateName: 'step3LambdaFunction',
                comment: 'Step 3 Lambda Function',
                invocationType: tasks.LambdaInvocationType.EVENT,
                outputPath: '$',
                inputPath: '$'
              }
            ).next(new sfn.Succeed(this, 'WorkflowCompleted'))
          )
        )
    );

    // Step Function State Machine
    // Create an IAM role for the Step Functions State Machine
    const stateMachineRole = new iam.Role(
      this,
      'StepFunctionStateMachineRole',
      {
        assumedBy: new iam.ServicePrincipal('states.amazonaws.com'),
        // Add any additional policies or permissions as needed
        description: 'Step Function State Machine Role',
        roleName: 'step-function-state-machine-role',
        inlinePolicies: {
          createCloudWatchLogsPolicy: new iam.PolicyDocument({
            statements: [
              new iam.PolicyStatement({
                actions: [
                  'logs:CreateLogDelivery',
                  'logs:GetLogDelivery',
                  'logs:UpdateLogDelivery',
                  'logs:DeleteLogDelivery',
                  'logs:ListLogDeliveries',
                  'logs:PutResourcePolicy',
                  'logs:DescribeResourcePolicies',
                  'logs:DescribeLogGroups'
                ],
                resources: ['*']
              })
            ]
          }),
          invokeLambdaFunction: new iam.PolicyDocument({
            statements: [
              new iam.PolicyStatement({
                actions: ['lambda:InvokeFunction'],
                resources: [
                  'arn:aws:lambda:eu-west-1:065741335689:function:step1LambdaFunction:*',
                  'arn:aws:lambda:eu-west-1:065741335689:function:step1LambdaFunction'
                ]
              })
            ]
          }),
          snsInvoke: new iam.PolicyDocument({
            statements: [
              new iam.PolicyStatement({
                actions: ['sns:*'],
                resources: ['*']
              })
            ]
          }),
          xraySegments: new iam.PolicyDocument({
            statements: [
              new iam.PolicyStatement({
                actions: [
                  'xray:PutTraceSegments',
                  'xray:PutTelemetryRecords',
                  'xray:GetSamplingRules',
                  'xray:GetSamplingTargets'
                ],
                resources: ['*']
              })
            ]
          })
        }
      }
    );

    const stepFunctionsLogGroup = new logs.LogGroup(
      this,
      'StepFunctionStateMachineLogGroup'
    );

    const stateMachine = new sfn.StateMachine(
      this,
      'StepFunctionStateMachine',
      {
        comment: 'Step Function State machine',
        stateMachineName: 'StepFunctionStateMachine',
        stateMachineType: sfn.StateMachineType.STANDARD,
        definitionBody: sfn.DefinitionBody.fromChainable(definition),
        timeout: cdk.Duration.minutes(30),
        role: stateMachineRole,
        logs: {
          destination: stepFunctionsLogGroup,
          level: sfn.LogLevel.ALL,
          includeExecutionData: true
        }
      }
    );

This is a large function. Let us break it down to a function by function level, in order to understand the code better.

Step 1: Lambda Function

// Step 1: Lambda Function:
const step1EnvVariable = {
  TABLE_NAME: "TableName"
};
const step1LambdaFunction = this.createLambdaFunctionAndRole(
  props.step1LambdaFunction,
  emailNotificationTopic,
  sampleCodeBucket,
  step1EnvVariable
);

This code creates the first lambda function towards the step function out of the three that we have to create, along with the environment variable.

Step 2: Lambda Function

// Step 2: Lambda Function
const step2EnvVariable = {
  OTHER_ENV_VARIABLE: "VALUE"
};
const step2LambdaFunction = this.createLambdaFunctionAndRole(
  props.step2LambdaFunction,
  sampleCodeBucket,
  emailNotificationTopic,
  props,
  step2EnvVariable
);

This code creates the second lambda function towards the step function out of the three that we have to create, along with the environment variable.

Step 3: Lambda Function

// Step 3: Lambda Function
const step3EnvVariable = {
  SOME_ENV_VARIABLE: "env Variable value"
};
const memorySize = 6144;
const step3LambdaFunction = this.createLambdaFunctionAndRole(
  props.step3LambdaFunction,
  emailNotificationTopic,
  sampleCodeBucket,
  step3EnvVariable,
  memorySize
);

This code creates the third lambda function towards the step function out of the three that we have to create, along with the environment variable.

Create State Machine Definition

// Step Function Definition
const definition = new tasks.LambdaInvoke(
  this,
  'step1LambdaFunction',
  {
    lambdaFunction: step1LambdaFunction.lambdaFunction,
    stateName: 'step1LambdaFunction',
    comment: 'Step 1 Lambda Function',
    invocationType: tasks.LambdaInvocationType.REQUEST_RESPONSE,
    outputPath: '$.Payload'
  }
).next(
  new sfn.Choice(this, 'Step1LambdaFunctionChoice', {
    stateName: 'Step1LambdaFunctionChoice',
    comment: 'Step 1 Lambda Function Choice',
    inputPath: '$',
    outputPath: '$'
  })
    .when(
      sfn.Condition.stringEquals('$.state', 'Yes'),
      new sfn.Succeed(this, 'ProcessEnded')
    )
    .otherwise(
      new tasks.LambdaInvoke(this, 'Step2LambdaFunction', {
        lambdaFunction: step2LambdaFunction.lambdaFunction,
        stateName: 'step2LambdaFunction',
        comment: 'Step 2 Lambda Function',
        invocationType: tasks.LambdaInvocationType.REQUEST_RESPONSE,
        outputPath: '$.Payload',
        inputPath: '$'
      }).next(
        new tasks.LambdaInvoke(
          this,
          'step3LambdaFunction',
          {
            lambdaFunction:
              step3LambdaFunction.lambdaFunction,
            stateName: 'step3LambdaFunction',
            comment: 'Step 3 Lambda Function',
            invocationType: tasks.LambdaInvocationType.EVENT,
            outputPath: '$',
            inputPath: '$'
          }
        ).next(new sfn.Succeed(this, 'WorkflowCompleted'))
      )
    )
);

1. Start the Step function with triggering the ‘step1LambdaFunction’ Lambda Function.
2. ‘Step1LambdaFunctionChoice’ is uses the output data from the previous lambda to make a choice — Stop the process or to Execute ‘Step2LambdaFunction’.
3. If the Choice made based on the output is to execute the ‘Step2LambdaFunction’ then the lambda function is executed.
4. Then the ‘step3LambdaFunction’ is also triggered from based on the output of the previous lambda function.

In some of the scenarios, we can see that the ‘$.Payload’ is used as an inputPath towards the Step Function lambda, this would ensure that in a JSON as follows.

{
  "Payload": {
    "someData": "someValue"
  },
  "MetaData": {
    "someMetaData": "someMetaValue"
  }
}

The input to the lambda function will be only the payload. i.e

{
  "someData": "someValue"
}

Invocation Type

The Invocation type of the lambda function can be either of the below values.

invocationType: tasks.LambdaInvocationType.EVENT,
invocationType: tasks.LambdaInvocationType.REQUEST_RESPONSE,

If the LambdaInvocationType is LambdaInvocationType.EVENT. It is triggered in a asynchronous manner.

If the LambdaInvocationType is LambdaInvocationType.REQUEST_RESPONSE. The Lambda is executed in a synchronous manner.

Create Step Machine Role

// Step Function State Machine
// Create an IAM role for the Step Functions State Machine
const stateMachineRole = new iam.Role(
  this,
  'StepFunctionStateMachineRole',
  {
    assumedBy: new iam.ServicePrincipal('states.amazonaws.com'),
    // Add any additional policies or permissions as needed
    description: 'Step Function State Machine Role',
    roleName: 'step-function-state-machine-role',
    inlinePolicies: {
      createCloudWatchLogsPolicy: new iam.PolicyDocument({
        statements: [
          new iam.PolicyStatement({
            actions: [
              'logs:CreateLogDelivery',
              'logs:GetLogDelivery',
              'logs:UpdateLogDelivery',
              'logs:DeleteLogDelivery',
              'logs:ListLogDeliveries',
              'logs:PutResourcePolicy',
              'logs:DescribeResourcePolicies',
              'logs:DescribeLogGroups'
            ],
            resources: ['*']
          })
        ]
      }),
      invokeLambdaFunction: new iam.PolicyDocument({
        statements: [
          new iam.PolicyStatement({
            actions: ['lambda:InvokeFunction'],
            resources: [
              'arn:aws:lambda:eu-west-1:065741335689:function:step1LambdaFunction:*',
              'arn:aws:lambda:eu-west-1:065741335689:function:step1LambdaFunction'
            ]
          })
        ]
      }),
      snsInvoke: new iam.PolicyDocument({
        statements: [
          new iam.PolicyStatement({
            actions: ['sns:*'],
            resources: ['*']
          })
        ]
      }),
      xraySegments: new iam.PolicyDocument({
        statements: [
          new iam.PolicyStatement({
            actions: [
              'xray:PutTraceSegments',
              'xray:PutTelemetryRecords',
              'xray:GetSamplingRules',
              'xray:GetSamplingTargets'
            ],
            resources: ['*']
          })
        ]
      })
    }
  }
);

We have created a State Machine role with an inline policy, where cloudwatch log roles, SNS, Xray Segments, Invoke the first lambda function are added.

assumedBy: new iam.ServicePrincipal('states.amazonaws.com'),

The above is the statement that adds principal to ensure that the role is executed for the service principal.

Create: State Machine

const stepFunctionsLogGroup = new logs.LogGroup(
  this,
  'StepFunctionStateMachineLogGroup'
);

const stateMachine = new sfn.StateMachine(
  this,
  'StepFunctionStateMachine',
  {
    comment: 'Step Function State machine',
    stateMachineName: 'StepFunctionStateMachine',
    stateMachineType: sfn.StateMachineType.STANDARD,
    definitionBody: sfn.DefinitionBody.fromChainable(definition),
    timeout: cdk.Duration.minutes(30),
    role: stateMachineRole,
    logs: {
      destination: stepFunctionsLogGroup,
      level: sfn.LogLevel.ALL,
      includeExecutionData: true
    }
  }
);

This piece of code is used to create a state machine and link the state machine role that we had already created to link it to the state machine.

Github Code

https://github.com/SkillSageAcademy/CDK-templates/blob/main/state-functions.ts

Conclusion

This article can be successfully used to create a Step Function using the CDK.

We have bloated the step function to ensure that some of the choices and triggers and invocation types can be seen and uploading this CDK to the AWS would ensure that the infrastructure with

1. 3 Lambda Function
2. 3 Relevant Lambda Roles for these Lambda Functions
3. 3 Log Groups in the Cloud Watch Log groups
4. A State Machine in the Step Function
5. A State Machine Role for the State Machine
6. A State Machine Log Group

Are created with the help of the above CDK typescript.