feat(stepfunctions): add support JSONata and variables

packages/aws-cdk-lib/aws-stepfunctions-tasks/lib/bedrock/invoke-model.ts

@@ -64,7 +64,7 @@ export interface BedrockInvokeModelOutputProps {
 /**
  * Properties for invoking a Bedrock Model
  */
-export interface BedrockInvokeModelProps extends sfn.TaskStateBaseProps {
+export interface BedrockInvokeModelProps extends sfn.TaskStateBaseJsonPathProps {
 
   /**
    * The Bedrock model that the task will invoke.

packages/aws-cdk-lib/aws-stepfunctions/README.md

@@ -23,8 +23,8 @@ const submitJob = new tasks.LambdaInvoke(this, 'Submit Job', {
   outputPath: '$.guid',
 });
 
-const waitX = new sfn.Wait(this, 'Wait X Seconds', {
-  time: sfn.WaitTime.secondsPath('$.waitSeconds'),
+const waitX = sfn.Wait.jsonata(this, 'Wait X Seconds', {
+  time: sfn.WaitTime.seconds('{% $states.input.waitSeconds %}'),
 });
 
 const getStatus = new tasks.LambdaInvoke(this, 'Get Job Status', {
@@ -50,10 +50,10 @@ const finalStatus = new tasks.LambdaInvoke(this, 'Get Final Job Status', {
 const definition = submitJob
   .next(waitX)
   .next(getStatus)
-  .next(new sfn.Choice(this, 'Job Complete?')
+  .next(sfn.Choice.jsonata(this, 'Job Complete?')
     // Look at the "status" field
-    .when(sfn.Condition.stringEquals('$.status', 'FAILED'), jobFailed)
-    .when(sfn.Condition.stringEquals('$.status', 'SUCCEEDED'), finalStatus)
+    .when(sfn.Condition.jsonata('{% $states.input.status = FAILED %}'), jobFailed)
+    .when(sfn.Condition.jsonata('{% $states.input.status = SUCCEEDED %}'), finalStatus)
     .otherwise(waitX));
 
 new sfn.StateMachine(this, 'StateMachine', {
@@ -73,7 +73,7 @@ definition. The definition is specified by its start state, and encompasses
 all states reachable from the start state:
 
 ```ts
-const startState = new sfn.Pass(this, 'StartState');
+const startState = sfn.Pass.jsonata(this, 'StartState');
 
 new sfn.StateMachine(this, 'StateMachine', {
   definitionBody: sfn.DefinitionBody.fromChainable(startState),
@@ -100,7 +100,7 @@ Alternatively you can specify an existing step functions definition by providing
 
 ```ts
 new sfn.StateMachine(this, 'StateMachineFromString', {
-  definitionBody: sfn.DefinitionBody.fromString('{"StartAt":"Pass","States":{"Pass":{"Type":"Pass","End":true}}}'),
+  definitionBody: sfn.DefinitionBody.fromString('{"StartAt":"Pass","States":{"Pass":{"Type":"Pass","QueryLanguage":"JSONata","End":true}}}'),
 });
 
 new sfn.StateMachine(this, 'StateMachineFromFile', {
@@ -117,10 +117,41 @@ gets modified by individual steps as the state machine progresses, and finally
 is produced as output.
 
 By default, the entire Data object is passed into every state, and the return data of the step
-becomes new the new Data object. This behavior can be modified by supplying values for `inputPath`,
-`resultSelector`, `resultPath` and `outputPath`.
+becomes new the new Data object. Step Functions supports JSONPath and JSONata for these data operations. Data can be changed using the operation methods provided for each query language.
 
-### Manipulating state machine data using inputPath, resultSelector, resultPath and outputPath
+### (JSONata) Transforming data with JSONata
+With JSONata, you gain a powerful open source query and expression language to select and transform data in your workflows. For a brief introduction and complete JSONata reference, see [JSONata.org documentation](https://docs.jsonata.org/overview.html). After selecting JSONata, your workflow fields will be reduced from five JSONPath fields (`InputPath`, `Parameters`, `ResultSelector`, `ResultPath`, and `OutputPath`) to only two fields: `Arguments` and `Output`. Also, you will not use `.$` on JSON object key names.
+
+If you are new to Step Functions, you only need to know that JSONata expressions use the following syntax:
+
+JSONata syntax: `"{% <JSONata expression> %}"`
+
+The following code samples show a conversion from JSONPath to JSONata:
+
+```json
+ // Original sample using JSONPath
+{
+  "static": "Hello",
+  "title.$": "$.title",
+  "name.$": "$customerName",  // With $customerName declared as a variable
+  "not-evaluated": "$customerName"
+}
+```
+
+```json
+// Sample after conversion to JSONata
+{
+  "static": "Hello",
+  "title": "{% $states.input.title %}", 
+  "name": "{% $customerName %}",   // With $customerName declared as a variable
+  "not-evaluated": "$customerName"
+}
+```
+
+See the official documentation on [Transforming data with JSONata in Step Functions](https://docs.aws.amazon.com/step-functions/latest/dg/transforming-data.html).
+
+
+### (JSONPath) Manipulating state machine data using inputPath, resultSelector, resultPath and outputPath
 
 These properties impact how each individual step interacts with the state machine data:
 
@@ -160,6 +191,8 @@ See the official documentation on [input and output processing in Step Functions
 Tasks take parameters, whose values can be taken from the State Machine Data object. For example, your
 workflow may want to start a CodeBuild with an environment variable that is taken from the State Machine data, or pass part of the State Machine Data into an AWS Lambda Function.
 
+#### (JSONPath) Passing Parameters to Tasks
+
 In the original JSON-based states language used by AWS Step Functions, you would
 add `.$` to the end of a key to indicate that a value needs to be interpreted as
 a JSON path. In the CDK API you do not change the names of any keys. Instead, you

packages/aws-cdk-lib/aws-stepfunctions/lib/condition.ts

@@ -338,6 +338,13 @@ export abstract class Condition {
     return new NotCondition(condition);
   }
 
+  /**
+   * JSONata expression condition
+   */
+  public static jsonata(conditon: string): Condition {
+    return new JsonataCondition(conditon);
+  }
+
   /**
    * Render Amazon States Language JSON for the condition
    */
@@ -452,3 +459,21 @@ class NotCondition extends Condition {
     };
   }
 }
+
+/**
+ * JSONata for Condition
+ */
+class JsonataCondition extends Condition {
+  constructor(private readonly condition: string) {
+    super();
+    if (!/^{%(.*)%}$/.test(condition)) {
+      throw new Error(`Variable reference must be '$', start with '$.', or start with '$[', got '${condition}'`);
+    }
+  }
+
+  public renderCondition(): any {
+    return {
+      Condition: this.condition,
+    };
+  }
+}

packages/aws-cdk-lib/aws-stepfunctions/lib/state-graph.ts

@@ -1,6 +1,7 @@
 import { StateMachine } from './state-machine';
 import { DistributedMap } from './states/distributed-map';
 import { State } from './states/state';
+import { QueryLanguage } from './types';
 import * as iam from '../../aws-iam';
 import { Duration } from '../../core';
 
@@ -105,10 +106,10 @@ export class StateGraph {
   /**
    * Return the Amazon States Language JSON for this graph
    */
-  public toGraphJson(): object {
+  public toGraphJson(queryLanguage?: QueryLanguage): object {
     const states: any = {};
     for (const state of this.allStates) {
-      states[state.stateId] = state.toStateJson();
+      states[state.stateId] = state.toStateJson(queryLanguage);
     }
 
     return {

packages/aws-cdk-lib/aws-stepfunctions/lib/state-machine.ts

@@ -5,7 +5,7 @@ import { buildEncryptionConfiguration } from './private/util';
 import { StateGraph } from './state-graph';
 import { StatesMetrics } from './stepfunctions-canned-metrics.generated';
 import { CfnStateMachine } from './stepfunctions.generated';
-import { IChainable } from './types';
+import { IChainable, QueryLanguage } from './types';
 import * as cloudwatch from '../../aws-cloudwatch';
 import * as iam from '../../aws-iam';
 import * as logs from '../../aws-logs';
@@ -129,6 +129,15 @@ export interface StateMachineProps {
    */
   readonly comment?: string;
 
+  /**
+   * The name of the query language used by the state machine.
+   * If the state does not contain a `queryLanguage` field,
+   * then it will use the query language specified in the top-level `queryLanguage` field.
+   *
+   * @default - JSONPATH
+   */
+  readonly queryLanguage?: QueryLanguage;
+
   /**
    * Type of the state machine
    *
@@ -786,9 +795,13 @@ export class ChainDefinitionBody extends DefinitionBody {
   }
 
   public bind(scope: Construct, _sfnPrincipal: iam.IPrincipal, sfnProps: StateMachineProps, graph?: StateGraph): DefinitionConfig {
-    const graphJson = graph!.toGraphJson();
+    const graphJson = graph!.toGraphJson(sfnProps.queryLanguage);
     return {
-      definitionString: Stack.of(scope).toJsonString({ ...graphJson, Comment: sfnProps.comment }),
+      definitionString: Stack.of(scope).toJsonString({
+        ...graphJson,
+        Comment: sfnProps.comment,
+        QueryLanguage: sfnProps.queryLanguage,
+      }),
     };
   }
 }

packages/aws-cdk-lib/aws-stepfunctions/lib/states/choice.ts

@@ -1,48 +1,28 @@
 import { Construct } from 'constructs';
 import { StateType } from './private/state-type';
-import { ChoiceTransitionOptions, State } from './state';
+import { ChoiceTransitionOptions, JsonataCommonOptions, JsonPathCommonOptions, State, StateBaseProps } from './state';
 import { Chain } from '../chain';
 import { Condition } from '../condition';
-import { IChainable, INextable } from '../types';
+import { IChainable, INextable, QueryLanguage } from '../types';
+
+interface ChoiceBaseProps extends StateBaseProps {}
+interface ChoiceJsonPathOptions extends JsonPathCommonOptions {}
+interface ChoiceJsonataOptions extends JsonataCommonOptions {}
 
 /**
- * Properties for defining a Choice state
+ * Properties for defining a Choice state that using JSONPath
  */
-export interface ChoiceProps {
-  /**
-   * Optional name for this state
-   *
-   * @default - The construct ID will be used as state name
-   */
-  readonly stateName?: string;
-
-  /**
-   * An optional description for this state
-   *
-   * @default No comment
-   */
-  readonly comment?: string;
+export interface ChoiceJsonPathProps extends ChoiceBaseProps, ChoiceJsonPathOptions { }
 
-  /**
-   * JSONPath expression to select part of the state to be the input to this state.
-   *
-   * May also be the special value JsonPath.DISCARD, which will cause the effective
-   * input to be the empty object {}.
-   *
-   * @default $
-   */
-  readonly inputPath?: string;
+/**
+ * Properties for defining a Choice state that using JSONata
+ */
+export interface ChoiceJsonataProps extends ChoiceBaseProps, ChoiceJsonataOptions { }
 
-  /**
-   * JSONPath expression to select part of the state to be the output to this state.
-   *
-   * May also be the special value JsonPath.DISCARD, which will cause the effective
-   * output to be the empty object {}.
-   *
-   * @default $
-   */
-  readonly outputPath?: string;
-}
+/**
+ * Properties for defining a Choice state
+ */
+export interface ChoiceProps extends ChoiceBaseProps, ChoiceJsonPathOptions, ChoiceJsonataOptions { }
 
 /**
  * Define a Choice in the state machine
@@ -51,6 +31,27 @@ export interface ChoiceProps {
  * state.
  */
 export class Choice extends State {
+  /**
+   * Define a Choice using JSONPath in the state machine
+   *
+   * A choice state can be used to make decisions based on the execution
+   * state.
+   */
+  public static jsonPath(scope: Construct, id: string, props: ChoiceJsonPathProps = {}) {
+    return new Choice(scope, id, props);
+  }
+  /**
+   * Define a Choice using JSONata in the state machine
+   *
+   * A choice state can be used to make decisions based on the execution
+   * state.
+   */
+  public static jsonata(scope: Construct, id: string, props: ChoiceJsonataProps = {}) {
+    return new Choice(scope, id, {
+      ...props,
+      queryLanguage: QueryLanguage.JSONATA,
+    });
+  }
   public readonly endStates: INextable[] = [];
 
   constructor(scope: Construct, id: string, props: ChoiceProps = {}) {
@@ -95,9 +96,10 @@ export class Choice extends State {
   /**
    * Return the Amazon States Language object for this state
    */
-  public toStateJson(): object {
+  public toStateJson(queryLanguage?: QueryLanguage): object {
     return {
       Type: StateType.CHOICE,
+      ...this.renderQueryLanguage(queryLanguage),
       Comment: this.comment,
       ...this.renderInputOutput(),
       ...this.renderChoices(),

packages/aws-cdk-lib/aws-stepfunctions/lib/states/custom-state.ts

@@ -2,7 +2,7 @@ import { Construct } from 'constructs';
 import { Chain } from '..';
 import { State } from './state';
 import { Annotations } from '../../../core/';
-import { CatchProps, IChainable, INextable, RetryProps } from '../types';
+import { CatchProps, IChainable, INextable, QueryLanguage, RetryProps } from '../types';
 
 /**
  * Properties for defining a custom state definition
@@ -68,8 +68,9 @@ export class CustomState extends State implements IChainable, INextable {
   /**
    * Returns the Amazon States Language object for this state
    */
-  public toStateJson(): object {
+  public toStateJson(queryLanguage?: QueryLanguage): object {
     const state = {
+      ...this.renderQueryLanguage(queryLanguage),
       ...this.renderNextEnd(),
       ...this.stateJson,
       ...this.renderRetryCatch(),