AWS CDK StepFunction Tasks EvaluateExpression

README.md

EvaluateExpression

This is an alternative to the supplied EvaluateExpression construct. EvaluateExpression is intended to provide some sugar over creating a lambda when defining step functions, by allowing a text representation of javascript code to be passed (as expression).

This has a number of disadvantages, such as lack of IDE syntax checking, and access to state input is done via a limited regular expression matching. There have been issues raised and pull requests created, but none have made their way into the CDK (in part because they only solve parts of the overall problem).

The existing implementation in the CDK uses eval() to process the expression, and I think this futher hampered the design decisions. A safer way to implement is to use the ability of Function to create a callable function from a string, as this isolates any variables in scope to those explictly passed (whereas eval has access to all).

This allows us to inject variables for $ (state input) and $$ (step function context) and requires no manipulation of the provided expression.

Further, we can provide the desired input parameters to the EvaluateExpression instantiation, which allows for customisation of what variables are injected (still relative to $).

As a bonus, in this implementation we can specify expression as a function, which means we can preserve IDE syntax checking. We must only be careful to avoid referencing variables from the CDK scope, but even if we did that would provoke a runtime error.

Basic usage

A typical use case is getting a TTL value for Dynamo DB items.

const getExpiryTime = new EvaluateExpression<MyNodeFunctionProps>(scope, 'Get Expiry Time', {
    expression: 'Math.floor(new Date($$.Execution.StartTime).getTime() / 1000) + $.timeToLive',
    functionClass: MyNodeFunction,
    parameters: {
        timeToLive: TimeToLive.toSeconds(),
    },
    resultPath: '$.expiresAt',
})

Extended Usage

Here we use the ability to provide a function instead of a string expression. We keep syntax checking in our IDE, and we can add types, both to help document and enhance syntax checking. It also makes it easier for code to run over several lines for clarity.

const getExpiryTime = new EvaluateExpression<MyNodeFunctionProps>(scope, 'Get Expiry Time', {
    // NB: $ is inferred from parameters (Record<string, any> when not supplied)
    // $$ is Context (step function context, e.g. $$.Execution.Id etc)
    expression: ($, $$) => {
        const date = new Date($$.Execution.StartTime)
        date.setSeconds(date.getSeconds() + $.timeToLive)
        return Math.floor(date.getTime() / 1000)
    },
    functionClass: MyNodeFunction,
    parameters: {
        timeToLive: TimeToLive.toSeconds(),
    },
    resultPath: '$.expiresAt',
})

Caveats

EvaluateExpression is intended as a shortcut to implementating a lambda anytime you want to do some minor computation that is beyond that provided by Step Function Intrinsics.

We must remember that we cannot import modules (even though they might be available in the runtime, such as the AWS SDK), nor our own. If you find yourself needing to do this, that's when you could be taking the longer route to a full Lambda implementation (via LambdaInvoke etc).

Notes

Using EvaluateExpression will add a Lambda to your stack, but only one is added for each runtime (so only one if you stick to the default or only specify a particular runtime). It will have a log group created for it, as is normal for Lambda, and will be named after the stack name (and stage if used), and the runtime (e.g. "/aws/lambda/Dev-Replicated-SfnTasksEvaluateExpression-nodejs18x")

EvaluateExpression.handler.ts

import { Event } from './EvaluateExpression'

export async function handler(event: Event): Promise<any> {
    // console.log('event', JSON.stringify(event))

    const { expression, isFunction, expressionAttributes } = event

    const source = isFunction ? `return (${expression})($, $$)` : `return ${expression}`

    return Function(...Object.keys(expressionAttributes), source)(...Object.values(expressionAttributes))

}

EvaluateExpression.ts

import { FieldUtils, TaskMetricsConfig, TaskStateBase, TaskStateBaseProps } from "aws-cdk-lib/aws-stepfunctions"
import { Construct } from "constructs"
import { PolicyStatement } from "aws-cdk-lib/aws-iam"
import path from "path"
import { NodejsFunction, NodejsFunctionProps } from "aws-cdk-lib/aws-lambda-nodejs"
import { Runtime } from "aws-cdk-lib/aws-lambda"
import { Stack } from "aws-cdk-lib"

export interface Context {
    Execution: {
        Id: string
        Input: any
        StartTime: string
        Name: string
        RoleArn: string

    }
    StateMachine: {
        Id: string
        Name: string
    }
    State: {
        Name: string
        EnteredTime: string
        RetryCount: number
    }
}

type Params = Record<string, any>

type Handler<T extends Params> = ($: T, $$: Context) => any

export interface Event {
    expressionAttributes: Params,
    expression: string,
    isFunction: boolean,
}

type Newable = new (scope: Construct, id: string, ...args: any[]) => any

export interface EvaluateExpressionProps<Klass extends Newable = Newable, Parameters extends Params = Params> extends TaskStateBaseProps {
    expression: Handler<Parameters> | string
    runtime?: Runtime
    parameters?: Parameters
    functionProps?: ConstructorParameters<Klass>[2]
    functionClass?: Klass
}

export class EvaluateExpression<Klass extends Newable = Newable, Parameters extends Params = Params> extends TaskStateBase {

    protected readonly taskMetrics?: TaskMetricsConfig;
    protected readonly taskPolicies?: PolicyStatement[];

    private fn: NodejsFunction

    constructor(scope: Construct, id: string, private readonly props: EvaluateExpressionProps<Klass, Parameters>) {

        super(scope, id, props)

        this.fn = this.getSingletonFunction<Klass, Parameters>(props)

        this.taskPolicies = [
            new PolicyStatement({
                actions: ['lambda:InvokeFunction' ],
                resources: [ this.fn.functionArn ],
            })
        ]
    }

    protected _renderTask(): any {

        const isFunction = (value: unknown): value is Function => typeof this.props.expression === 'function'

        const parameters: Event = {
            expressionAttributes: {
                ...this.props.parameters
                    ? { '$': FieldUtils.renderObject(this.props.parameters) }
                    : { '$.$': '$' },
                '$$.$': '$$',
            },
            // NB: toString() is necessary to convert the function, and harmless if expression is a string
            expression: this.props.expression.toString(),
            isFunction: isFunction(this.props.expression),
        }

        return {
            Resource: this.fn.functionArn,
            Parameters: parameters,
        }
    }

    /**
     * Regardless how many times this construct is instantiated, only one Lambda Function will be created (per
     * stack & runtime).
     */
    private getSingletonFunction<Klass extends Newable = Newable, Properties extends Params = Params>(props: EvaluateExpressionProps<Klass, Properties>): NodejsFunction {

        const runtime = props.runtime ?? Runtime.NODEJS_18_X
        const constructId = `SfnTasksEvaluateExpression-${runtime.name}`
        const existing = Stack.of(this).node.tryFindChild(constructId) as NodejsFunction

        return existing ??
            new (props.functionClass ?? NodejsFunction)(this, constructId, {
                functionName: (Stack.of(this).stackName + '-' + constructId).replaceAll(/[^a-zA-Z0-9-_]/g, ''),
                entry: path.resolve(__dirname, './EvaluateExpression.handler.ts'),
                runtime: runtime,
                ...props.functionProps,
            })
    }
}