Integrating AWS Step Functions ASL state machine definitions into CDK stacks in Python

AWS Step Functions state machines are defined in the JSON-based Amazon States Language (ASL). However, writing and maintaining ASL by hand is cumbersome. Therefore, most developers will prefer using the AWS Step Functions Workflow Studio in the AWS console, which provides a low-code editor and interactive testing facilities, such as the ability to inspect the state at different steps or to execute single steps in isolation.

While CloudFormation merely allows including an ASL snippet directly into a template, AWS CDK provides an abstraction to construct state machines. This is much more manageable than editing ASL, but prevents editing and testing of the state machines in the Step Functions Workflow Studio. In particular, it's not possible to convert an ASL state machine as exported from the low-code editor back to AWS CDK code.

Against this backdrop, in a recent project we were looking for a setup to...

• ... maintain the state machine's ASL definition so that it can be edited and tested in the Step Functions Workflow Studio's low-code editor. It should be possible to simply copy and paste the definition between the editor and the code repository.

• ... define the state machine in an AWS CDK stack, so that all roles and permissions can be handled through the CDK.

• ... adjust the state machine's definition based on dynamic values in the AWS CDK stack, such as Lambda function ARNs.

After several iterations, we arrived at a pattern where we keep the ASL JSON as-is and overwrite the relevant parts with a helper function:

import json
import pathlib

PROJECT_ROOT: pathlib.Path = ...


def generate_my_state_machine(
    my_lambda_fn_arn: str,
    my_bucket_name: str,
) -> str:
    # Prepare parameters
    output_path = f"s3://{my_bucket_name}/output"

    # Load ASL template
    with open(
        PROJECT_ROOT / "sfn" / "my_state_machine.json",
        "rt",
    ) as f:
        my_state_machine = json.load(f)

    # Modify template
    call_fn = my_state_machine["States"]["Call"]["Parameters"]
    call_fn["FunctionName"] = my_lambda_fn_arn
    call_fn["Payload"]["output_path"] = output_path

    return json.dumps(my_state_machine)

This approach works well for moderately nested state machines with few parameters where the names of steps and the overall structure (think Maps and Parallel) do not change often.

Using the helper function, we can generate the state machine definition and use it within an AWS CDK stack as follows:

import aws_cdk as cdk
from aws_cdk import aws_iam as iam
from aws_cdk import aws_lambda as lambda_
from aws_cdk import aws_s3 as s3
from aws_cdk import aws_stepfunctions as stepfunctions


class MyStack(cdk.Stack):
    def __init__(self, scope: cdk.App, stack_name: str, **kwargs):
        super().__init__(scope, stack_name, **kwargs)

        # Create resources used by the state machine
        my_lambda_fn = lambda_.Function(self, ...)
        my_bucket = s3.Bucket(self, ...)

        # Create a role and assign required permissions
        my_step_function_role = iam.Role(
            self,
            "MyStepFunctionRole",
            assumed_by=iam.ServicePrincipal("states.amazonaws.com"),
            role_name="MyStepFunctionRole",
        )
        my_lambda_fn.grant_invoke(my_step_function_role)
        my_bucket.grant_write(my_step_function_role)
        my_step_function_role.add_to_policy(iam.PolicyStatement(...))

        # Define the state machine
        my_state_machine_definition = generate_my_state_machine(
            my_lambda_fn_arn=my_lambda_fn.function_arn,
            my_bucket_name=my_bucket.bucket_name,
        )

        my_state_machine = stepfunctions.CfnStateMachine(
            self,
            "MyStateMachine",
            definition_string=my_state_machine_definition,
            state_machine_name="MyStateMachine",
            role_arn=my_step_function_role.arn,
        )

        my_state_machine.grant_execution(...)