stepfunctions: Step function definition does not include query language when using JSONata

[joerino96]

Describe the bug

When using JSONata, CDK does not always write the query language parameter to the step function definition, leading to undefined behavior.

This becomes a problem when using the DescribeStateMachine API and TestState API, where TestState using the definition from DescribeStateMachine from will assume the state is using JSONPath instead of JSONata.

This behavior is not consistent, as I have noticed that states within Maps do write the query language parameter to the state definition, while non nested states do not.

The work around is to manually add the QueryLanguage parameter to the definition before invoking TestState API.

Regression Issue

• Select this option if this issue appears to be a regression.

Last Known Working CDK Version

No response

Expected Behavior

The following CDK code:

    return LambdaInvoke.jsonata(scope, "MyStep", {
        lambdaFunction: lambda,
        payload: TaskInput.fromObject({}),
    });

consistently produces the step function step definition ASL:

"MyStep": {
      "QueryLanguage": "JSONata",
    {...}
    },

Current Behavior

The following CDK code:

    return LambdaInvoke.jsonata(scope, "MyStep", {
        lambdaFunction: lambda,
        payload: TaskInput.fromObject({}),
    });

produces the step function step definition ASL:

"MyStep": {
    {...}
    },

Reproduction Steps

   LambdaInvoke.jsonata(scope, "MyStep", {
        lambdaFunction: lambda,
        payload: TaskInput.fromObject({}),
    });

Possible Solution

jsonata constructors are mostly not writing the query language to ASL, with this behavior being undefined.

Additional Information/Context

No response

CDK CLI Version

2.178.1

Framework Version

No response

Node.js Version

Noe 18

OS

macOS 15.3

Language

TypeScript

Language Version

No response

Other information

No response

[pahud]

After analyzing the codebase, the issue stems from how query language handling is implemented in the AWS CDK Step Functions library. Here's what's happening:

1. The LambdaInvoke.jsonata() static method is correctly setting the query language.
   2.However, in the _renderTask method of LambdaInvoke, the query language is not being consistently included in the ASL output:

aws-cdk/packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/lambda/invoke.ts

Lines 176 to 192 in 47a3f90

	protected _renderTask(topLevelQueryLanguage?: sfn.QueryLanguage): any {
	const queryLanguage = sfn._getActualQueryLanguage(topLevelQueryLanguage, this.props.queryLanguage);
	const [resource, paramOrArg] = this.props.payloadResponseOnly ?
	[this.props.lambdaFunction.functionArn, this.props.payload?.value] :
	[integrationResourceArn('lambda', 'invoke', this.integrationPattern), {
	FunctionName: this.props.lambdaFunction.functionArn,
	Payload: this.props.payload?.value ??
	(queryLanguage === sfn.QueryLanguage.JSONATA ? '{% $states.input %}'
	: sfn.TaskInput.fromJsonPathAt('$').value),
	InvocationType: this.props.invocationType,
	ClientContext: this.props.clientContext,
	Qualifier: this.props.qualifier,
	}];
	return {
	Resource: resource,
	...this._renderParametersOrArguments(paramOrArg, queryLanguage),
	};

3. The issue is that while the query language is used to determine the payload format, it's not being explicitly included in the ASL output as a QueryLanguage field.

To fix this, we need to ensure the QueryLanguage field is consistently included in the ASL output.

[joerino96]

Thank you @pahud for the quick deep dive. Your attentiveness to the bugs I've reported is much appreciated!

[GavinZZ]

@joerino96 Hello, I am unable to reproduce this error. Consider this following stack I'm using which defines LambdaInvoke task using jsonata.

export class CdkAppStack extends cdk.Stack {
  constructor(scope: Construct, id: string, props?: cdk.StackProps) {
    super(scope, id, props);

    const submitJobLambda = new Function(this, 'submitJobLambda', {
      code: Code.fromInline(`exports.handler = async (event, context) => {
            return {
              statusCode: '200',
              body: 'hello, world!',
              ...event,
            };
          };`),
      runtime: Runtime.NODEJS_18_X,
      handler: 'index.handler',
    });
    
    const submitJob = tasks.LambdaInvoke.jsonata(this, 'Invoke Handler', {
      lambdaFunction: submitJobLambda,
      payload: sfn.TaskInput.fromObject({
        execId: '{% $states.context.Execution.Id %}',
      }),
      outputs: '{% $states.result.Payload %}',
    });
    const chain = sfn.Chain.start(submitJob)

    const stateMachine = new sfn.StateMachine(this, 'StateMachine', {
      definition: chain,
    });
  }
}

The snippet of generated CFN template is as follows:

"StateMachine2E01A3A5": {
   "Type": "AWS::StepFunctions::StateMachine",
   "Properties": {
    "DefinitionString": {
     "Fn::Join": [
      "",
      [
       "{\"StartAt\":\"Invoke Handler\",\"States\":{\"Invoke Handler\":{\"QueryLanguage\":\"JSONata\",\"End\":true,\"Retry\":[{\"ErrorEquals\":[\"Lambda.ClientExecutionTimeoutException\",\"Lambda.ServiceException\",\"Lambda.AWSLambdaException\",\"Lambda.SdkClientException\"],\"IntervalSeconds\":2,\"MaxAttempts\":6,\"BackoffRate\":2}],\"Type\":\"Task\",\"Arguments\":{\"FunctionName\":\"",
       {
        "Fn::GetAtt": [
         "submitJobLambdaEFB00F3C",
         "Arn"
        ]
       },
       "\",\"Payload\":{\"execId\":\"{% $states.context.Execution.Id %}\"}},\"Output\":\"{% $states.result.Payload %}\",\"Resource\":\"arn:",
       {
        "Ref": "AWS::Partition"
       },
       ":states:::lambda:invoke\"}}}"
      ]
     ]
    },

As you can see above, the DefinitionString with Invoke Handler state does contain a query language JSONATA.

Furthermore, looking at the code of TaskStateBase class which is the parent class of LambdaInvoke class will always render the state json with query language.

Would you provide some reproducible code for me so that I can further look into the issue?

[joerino96]

Hi @GavinZZ, I believe we can close this issue.

When looking into my code I found I was using queryLanguage in the step function definition at the top level. This will set the default language for the entire step function. The queryLanguage inside the step can override the top-level default according to the step function ASL, which means that queryLanguage does not necessarily need to be set within the step.

Unfortunately, for users to determine the query language being used, a user will need to invoke DescribeStateMachine API and then:

1. Check the step function top level queryLanguage.
2. Check the step's queryLanguage.

This is un-optimal for the testState API, which takes in a single step's definition and is unaware of the top level queryLanguage being set by the step function, but that's not a CDK problem :).

[GavinZZ]

Thanks for getting back to me! This exact issue happened to me when I was fixing other issues in JSONata and confuses me as I was testing at state-level! I can understand the frustration and I'll convey the finding with stepfunctions team and see if there's anything that they can do.

That being said, will mark this issue as closed.

[github-actions]

Comments on closed issues and PRs are hard for our team to see.
If you need help, please open a new issue that references this one.