feat(stepfunctions): add comment ability for when condition (#27010)

Currently you can comment on a `Choice` state. This adds the ability to comment on a when condition. <img width="437" alt="Screenshot 2023-09-05 at 7 47 35 AM" src="https://github.com/aws/aws-cdk/assets/3310356/4cb92427-c16c-4994-9e86-0384c4c2b075"> <img width="473" alt="Screenshot 2023-09-05 at 7 54 29 AM" src="https://github.com/aws/aws-cdk/assets/3310356/4dc1857e-5515-4510-8fde-a5eb8bc93afe"> Closes #27005. ---- *By submitting this pull request, I confirm that my contribution is made under the terms of the Apache-2.0 license*
aws · Sep 19, 2023 · 3a4a7bb · 3a4a7bb
1 parent b58b9ce
commit 3a4a7bb
Show file tree

Hide file tree

Showing 11 changed files with 149 additions and 50 deletions.
diff --git a/...ws-stepfunctions/test/integ.state-machine.js.snapshot/aws-stepfunctions-integ.assets.json b/...ws-stepfunctions/test/integ.state-machine.js.snapshot/aws-stepfunctions-integ.assets.json
@@ -1,15 +1,15 @@
 {
-  "version": "32.0.0",
+  "version": "34.0.0",
   "files": {
-    "2b683032648fa932811d781e11acb71f29e8dcde21411e7cf232807cc2565770": {
+    "6fc80a94a8b95de238e82bf832c22a76af712d24dc10c332baa36571da0a4a41": {
       "source": {
         "path": "aws-stepfunctions-integ.template.json",
         "packaging": "file"
       },
       "destinations": {
         "current_account-current_region": {
           "bucketName": "cdk-hnb659fds-assets-${AWS::AccountId}-${AWS::Region}",
-          "objectKey": "2b683032648fa932811d781e11acb71f29e8dcde21411e7cf232807cc2565770.json",
+          "objectKey": "6fc80a94a8b95de238e82bf832c22a76af712d24dc10c332baa36571da0a4a41.json",
           "assumeRoleArn": "arn:${AWS::Partition}:iam::${AWS::AccountId}:role/cdk-hnb659fds-file-publishing-role-${AWS::AccountId}-${AWS::Region}"
         }
       }

diff --git a/...-stepfunctions/test/integ.state-machine.js.snapshot/aws-stepfunctions-integ.template.json b/...-stepfunctions/test/integ.state-machine.js.snapshot/aws-stepfunctions-integ.template.json
@@ -115,13 +115,13 @@
   "StateMachine2E01A3A5": {
    "Type": "AWS::StepFunctions::StateMachine",
    "Properties": {
+    "DefinitionString": "{\"StartAt\":\"wait time\",\"States\":{\"wait time\":{\"Type\":\"Wait\",\"SecondsPath\":\"$.waitSeconds\",\"Next\":\"choice\"},\"choice\":{\"Type\":\"Choice\",\"Comment\":\"this is a comment for the choice state\",\"Choices\":[{\"Variable\":\"$.success\",\"IsPresent\":true,\"Next\":\"success\",\"Comment\":\"this is a comment for the when condition\"},{\"Variable\":\"$.noComment\",\"IsPresent\":true,\"Next\":\"short wait time\"}],\"Default\":\"success\"},\"success\":{\"Type\":\"Succeed\"},\"short wait time\":{\"Type\":\"Wait\",\"Seconds\":1,\"Next\":\"success\"}},\"Comment\":\"a super cool state machine\"}",
     "RoleArn": {
      "Fn::GetAtt": [
       "StateMachineRoleB840431D",
       "Arn"
      ]
-    },
-    "DefinitionString": "{\"StartAt\":\"wait time\",\"States\":{\"wait time\":{\"Type\":\"Wait\",\"SecondsPath\":\"$.waitSeconds\",\"End\":true}},\"Comment\":\"a super cool state machine\"}"
+    }
    },
    "DependsOn": [
     "StateMachineRoleB840431D"

diff --git a/...sting/framework-integ/test/aws-stepfunctions/test/integ.state-machine.js.snapshot/cdk.out b/...sting/framework-integ/test/aws-stepfunctions/test/integ.state-machine.js.snapshot/cdk.out
@@ -1 +1 @@
-{"version":"32.0.0"}
+{"version":"34.0.0"}
diff --git a/...ng/framework-integ/test/aws-stepfunctions/test/integ.state-machine.js.snapshot/integ.json b/...ng/framework-integ/test/aws-stepfunctions/test/integ.state-machine.js.snapshot/integ.json
@@ -1,5 +1,5 @@
 {
-  "version": "32.0.0",
+  "version": "34.0.0",
   "testCases": {
     "integ.state-machine": {
       "stacks": [

diff --git a/...framework-integ/test/aws-stepfunctions/test/integ.state-machine.js.snapshot/manifest.json b/...framework-integ/test/aws-stepfunctions/test/integ.state-machine.js.snapshot/manifest.json
@@ -1,5 +1,5 @@
 {
-  "version": "32.0.0",
+  "version": "34.0.0",
   "artifacts": {
     "aws-stepfunctions-integ.assets": {
       "type": "cdk:asset-manifest",
@@ -17,7 +17,7 @@
         "validateOnSynth": false,
         "assumeRoleArn": "arn:${AWS::Partition}:iam::${AWS::AccountId}:role/cdk-hnb659fds-deploy-role-${AWS::AccountId}-${AWS::Region}",
         "cloudFormationExecutionRoleArn": "arn:${AWS::Partition}:iam::${AWS::AccountId}:role/cdk-hnb659fds-cfn-exec-role-${AWS::AccountId}-${AWS::Region}",
-        "stackTemplateAssetObjectUrl": "s3://cdk-hnb659fds-assets-${AWS::AccountId}-${AWS::Region}/2b683032648fa932811d781e11acb71f29e8dcde21411e7cf232807cc2565770.json",
+        "stackTemplateAssetObjectUrl": "s3://cdk-hnb659fds-assets-${AWS::AccountId}-${AWS::Region}/6fc80a94a8b95de238e82bf832c22a76af712d24dc10c332baa36571da0a4a41.json",
         "requiresBootstrapStackVersion": 6,
         "bootstrapStackVersionSsmParameter": "/cdk-bootstrap/hnb659fds/version",
         "additionalDependencies": [

diff --git a/...ing/framework-integ/test/aws-stepfunctions/test/integ.state-machine.js.snapshot/tree.json b/...ing/framework-integ/test/aws-stepfunctions/test/integ.state-machine.js.snapshot/tree.json
@@ -12,8 +12,32 @@
             "id": "wait time",
             "path": "aws-stepfunctions-integ/wait time",
             "constructInfo": {
-              "fqn": "aws-cdk-lib.aws_stepfunctions.Wait",
-              "version": "0.0.0"
+              "fqn": "constructs.Construct",
+              "version": "10.2.70"
+            }
+          },
+          "short wait time": {
+            "id": "short wait time",
+            "path": "aws-stepfunctions-integ/short wait time",
+            "constructInfo": {
+              "fqn": "constructs.Construct",
+              "version": "10.2.70"
+            }
+          },
+          "choice": {
+            "id": "choice",
+            "path": "aws-stepfunctions-integ/choice",
+            "constructInfo": {
+              "fqn": "constructs.Construct",
+              "version": "10.2.70"
+            }
+          },
+          "success": {
+            "id": "success",
+            "path": "aws-stepfunctions-integ/success",
+            "constructInfo": {
+              "fqn": "constructs.Construct",
+              "version": "10.2.70"
             }
           },
           "Role": {
@@ -24,8 +48,8 @@
                 "id": "ImportRole",
                 "path": "aws-stepfunctions-integ/Role/ImportRole",
                 "constructInfo": {
-                  "fqn": "aws-cdk-lib.Resource",
-                  "version": "0.0.0"
+                  "fqn": "constructs.Construct",
+                  "version": "10.2.70"
                 }
               },
               "Resource": {
@@ -49,8 +73,8 @@
                   }
                 },
                 "constructInfo": {
-                  "fqn": "aws-cdk-lib.aws_iam.CfnRole",
-                  "version": "0.0.0"
+                  "fqn": "constructs.Construct",
+                  "version": "10.2.70"
                 }
               },
               "DefaultPolicy": {
@@ -139,20 +163,20 @@
                       }
                     },
                     "constructInfo": {
-                      "fqn": "aws-cdk-lib.aws_iam.CfnPolicy",
-                      "version": "0.0.0"
+                      "fqn": "constructs.Construct",
+                      "version": "10.2.70"
                     }
                   }
                 },
                 "constructInfo": {
-                  "fqn": "aws-cdk-lib.aws_iam.Policy",
-                  "version": "0.0.0"
+                  "fqn": "constructs.Construct",
+                  "version": "10.2.70"
                 }
               }
             },
             "constructInfo": {
-              "fqn": "aws-cdk-lib.aws_iam.Role",
-              "version": "0.0.0"
+              "fqn": "constructs.Construct",
+              "version": "10.2.70"
             }
           },
           "StateMachine": {
@@ -167,8 +191,8 @@
                     "id": "ImportRole",
                     "path": "aws-stepfunctions-integ/StateMachine/Role/ImportRole",
                     "constructInfo": {
-                      "fqn": "aws-cdk-lib.Resource",
-                      "version": "0.0.0"
+                      "fqn": "constructs.Construct",
+                      "version": "10.2.70"
                     }
                   },
                   "Resource": {
@@ -192,14 +216,14 @@
                       }
                     },
                     "constructInfo": {
-                      "fqn": "aws-cdk-lib.aws_iam.CfnRole",
-                      "version": "0.0.0"
+                      "fqn": "constructs.Construct",
+                      "version": "10.2.70"
                     }
                   }
                 },
                 "constructInfo": {
-                  "fqn": "aws-cdk-lib.aws_iam.Role",
-                  "version": "0.0.0"
+                  "fqn": "constructs.Construct",
+                  "version": "10.2.70"
                 }
               },
               "Resource": {
@@ -208,60 +232,60 @@
                 "attributes": {
                   "aws:cdk:cloudformation:type": "AWS::StepFunctions::StateMachine",
                   "aws:cdk:cloudformation:props": {
+                    "definitionString": "{\"StartAt\":\"wait time\",\"States\":{\"wait time\":{\"Type\":\"Wait\",\"SecondsPath\":\"$.waitSeconds\",\"Next\":\"choice\"},\"choice\":{\"Type\":\"Choice\",\"Comment\":\"this is a comment for the choice state\",\"Choices\":[{\"Variable\":\"$.success\",\"IsPresent\":true,\"Next\":\"success\",\"Comment\":\"this is a comment for the when condition\"},{\"Variable\":\"$.noComment\",\"IsPresent\":true,\"Next\":\"short wait time\"}],\"Default\":\"success\"},\"success\":{\"Type\":\"Succeed\"},\"short wait time\":{\"Type\":\"Wait\",\"Seconds\":1,\"Next\":\"success\"}},\"Comment\":\"a super cool state machine\"}",
                     "roleArn": {
                       "Fn::GetAtt": [
                         "StateMachineRoleB840431D",
                         "Arn"
                       ]
-                    },
-                    "definitionString": "{\"StartAt\":\"wait time\",\"States\":{\"wait time\":{\"Type\":\"Wait\",\"SecondsPath\":\"$.waitSeconds\",\"End\":true}},\"Comment\":\"a super cool state machine\"}"
+                    }
                   }
                 },
                 "constructInfo": {
-                  "fqn": "aws-cdk-lib.aws_stepfunctions.CfnStateMachine",
-                  "version": "0.0.0"
+                  "fqn": "constructs.Construct",
+                  "version": "10.2.70"
                 }
               }
             },
             "constructInfo": {
-              "fqn": "aws-cdk-lib.aws_stepfunctions.StateMachine",
-              "version": "0.0.0"
+              "fqn": "constructs.Construct",
+              "version": "10.2.70"
             }
           },
           "BootstrapVersion": {
             "id": "BootstrapVersion",
             "path": "aws-stepfunctions-integ/BootstrapVersion",
             "constructInfo": {
-              "fqn": "aws-cdk-lib.CfnParameter",
-              "version": "0.0.0"
+              "fqn": "constructs.Construct",
+              "version": "10.2.70"
             }
           },
           "CheckBootstrapVersion": {
             "id": "CheckBootstrapVersion",
             "path": "aws-stepfunctions-integ/CheckBootstrapVersion",
             "constructInfo": {
-              "fqn": "aws-cdk-lib.CfnRule",
-              "version": "0.0.0"
+              "fqn": "constructs.Construct",
+              "version": "10.2.70"
             }
           }
         },
         "constructInfo": {
-          "fqn": "aws-cdk-lib.Stack",
-          "version": "0.0.0"
+          "fqn": "constructs.Construct",
+          "version": "10.2.70"
         }
       },
       "Tree": {
         "id": "Tree",
         "path": "Tree",
         "constructInfo": {
           "fqn": "constructs.Construct",
-          "version": "10.2.26"
+          "version": "10.2.70"
         }
       }
     },
     "constructInfo": {
-      "fqn": "aws-cdk-lib.App",
-      "version": "0.0.0"
+      "fqn": "constructs.Construct",
+      "version": "10.2.70"
     }
   }
 }
diff --git a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions/test/integ.state-machine.ts b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions/test/integ.state-machine.ts
@@ -14,12 +14,30 @@ const wait = new sfn.Wait(stack, 'wait time', {
   time: sfn.WaitTime.secondsPath('$.waitSeconds'),
 });
 
+const shortWait = new sfn.Wait(stack, 'short wait time', {
+  time: sfn.WaitTime.duration(cdk.Duration.seconds(1)),
+});
+
+const choice = new sfn.Choice(stack, 'choice', {
+  comment: 'this is a comment for the choice state',
+});
+
+const success = new sfn.Succeed(stack, 'success');
+
+choice.when(sfn.Condition.isPresent('$.success'), success, {
+  comment: 'this is a comment for the when condition',
+});
+choice.when(sfn.Condition.isPresent('$.noComment'), shortWait);
+choice.otherwise(success);
+wait.next(choice);
+shortWait.next(success);
+
 const role = new iam.Role(stack, 'Role', {
   assumedBy: new iam.ServicePrincipal('lambda.amazonaws.com'),
 });
 
 const stateMachine = new sfn.StateMachine(stack, 'StateMachine', {
-  definition: wait,
+  definitionBody: sfn.DefinitionBody.fromChainable(wait),
   comment: 'a super cool state machine',
 });
 

diff --git a/packages/aws-cdk-lib/aws-stepfunctions/README.md b/packages/aws-cdk-lib/aws-stepfunctions/README.md
@@ -353,6 +353,20 @@ const shipTheItem = new sfn.Pass(this, 'ShipTheItem');
 choice.afterwards().next(shipTheItem);
 ```
 
+You can add comments to `Choice` states as well as conditions that use `choice.when`.
+
+```ts
+const choice = new sfn.Choice(this, 'What color is it?', {
+  comment: 'color comment',
+});
+const handleBlueItem = new sfn.Pass(this, 'HandleBlueItem');
+const handleOtherItemColor = new sfn.Pass(this, 'HanldeOtherItemColor');
+choice.when(sfn.Condition.stringEquals('$.color', 'BLUE'), handleBlueItem, {
+  comment: 'blue item comment',
+});
+choice.otherwise(handleOtherItemColor);
+```
+
 If your `Choice` doesn't have an `otherwise()` and none of the conditions match
 the JSON state, a `NoChoiceMatched` error will be thrown. Wrap the state machine
 in a `Parallel` state if you want to catch and recover from this.

diff --git a/packages/aws-cdk-lib/aws-stepfunctions/lib/states/choice.ts b/packages/aws-cdk-lib/aws-stepfunctions/lib/states/choice.ts
@@ -1,6 +1,6 @@
 import { Construct } from 'constructs';
 import { StateType } from './private/state-type';
-import { State } from './state';
+import { ChoiceTransitionOptions, State } from './state';
 import { Chain } from '../chain';
 import { Condition } from '../condition';
 import { IChainable, INextable } from '../types';
@@ -53,8 +53,8 @@ export class Choice extends State {
   /**
    * If the given condition matches, continue execution with the given state
    */
-  public when(condition: Condition, next: IChainable): Choice {
-    super.addChoice(condition, next.startState);
+  public when(condition: Condition, next: IChainable, options?: ChoiceTransitionOptions): Choice {
+    super.addChoice(condition, next.startState, options);
     return this;
   }
 

diff --git a/packages/aws-cdk-lib/aws-stepfunctions/lib/states/state.ts b/packages/aws-cdk-lib/aws-stepfunctions/lib/states/state.ts
@@ -318,16 +318,16 @@ export abstract class State extends Construct implements IChainable {
   /**
    * Add a choice branch to this state
    */
-  protected addChoice(condition: Condition, next: State) {
-    this.choices.push({ condition, next });
+  protected addChoice(condition: Condition, next: State, options?: ChoiceTransitionOptions) {
+    this.choices.push({ condition, next, ...options });
     next.startState.addIncoming(this);
     if (this.containingGraph) {
       next.startState.bindToGraph(this.containingGraph);
     }
   }
 
   /**
-   * Add a paralle branch to this state
+   * Add a parallel branch to this state
    */
   protected addBranch(branch: StateGraph) {
     this.branches.push(branch);
@@ -479,7 +479,7 @@ export interface FindStateOptions {
 /**
  * A Choice Transition
  */
-interface ChoiceTransition {
+interface ChoiceTransition extends ChoiceTransitionOptions {
   /**
    * State to transition to
    */
@@ -491,13 +491,26 @@ interface ChoiceTransition {
   condition: Condition;
 }
 
+/**
+ * Options for Choice Transition
+ */
+export interface ChoiceTransitionOptions {
+  /**
+   * An optional description for the choice transition
+   *
+   * @default No comment
+   */
+  readonly comment?: string;
+}
+
 /**
  * Render a choice transition
  */
 function renderChoice(c: ChoiceTransition) {
   return {
     ...c.condition.renderCondition(),
     Next: c.next.stateId,
+    Comment: c.comment,
   };
 }