feat(stepfunctions): CustomState addCatch (#28422)

Adds a public `addCatch` method to the stepfunctions `CustomState` state. Closes #25798. ---- *By submitting this pull request, I confirm that my contribution is made under the terms of the Apache-2.0 license*
aws · Dec 19, 2023 · cf923bc · cf923bc
1 parent c687778
commit cf923bc
Show file tree

Hide file tree

Showing 5 changed files with 74 additions and 3 deletions.
diff --git a/...ns/test/integ.custom-state.js.snapshot/aws-stepfunctions-custom-state-integ.template.json b/...ns/test/integ.custom-state.js.snapshot/aws-stepfunctions-custom-state-integ.template.json
@@ -26,7 +26,7 @@
       "Arn"
      ]
     },
-    "DefinitionString": "{\"StartAt\":\"my custom task\",\"States\":{\"my custom task\":{\"Next\":\"final step\",\"Type\":\"Task\",\"Resource\":\"arn:aws:states:::dynamodb:putItem\",\"Parameters\":{\"TableName\":\"my-cool-table\",\"Item\":{\"id\":{\"S\":\"my-entry\"}}},\"ResultPath\":null},\"final step\":{\"Type\":\"Pass\",\"End\":true}},\"TimeoutSeconds\":30}"
+    "DefinitionString": "{\"StartAt\":\"my custom task\",\"States\":{\"my custom task\":{\"Next\":\"final step\",\"Type\":\"Task\",\"Resource\":\"arn:aws:states:::dynamodb:putItem\",\"Parameters\":{\"TableName\":\"my-cool-table\",\"Item\":{\"id\":{\"S\":\"my-entry\"}}},\"ResultPath\":null,\"Catch\":[{\"ErrorEquals\":[\"States.ALL\"],\"Next\":\"failed\"}]},\"final step\":{\"Type\":\"Pass\",\"End\":true},\"failed\":{\"Type\":\"Fail\",\"Error\":\"DidNotWork\",\"Cause\":\"We got stuck\"}},\"TimeoutSeconds\":30}"
    },
    "DependsOn": [
     "StateMachineRoleB840431D"

diff --git a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions/test/integ.custom-state.ts b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions/test/integ.custom-state.ts
@@ -25,10 +25,17 @@ const stateJson = {
   ResultPath: null,
 };
 
+const failure = new sfn.Fail(stack, 'failed', {
+  error: 'DidNotWork',
+  cause: 'We got stuck',
+});
+
 const custom = new sfn.CustomState(stack, 'my custom task', {
   stateJson,
 });
 
+custom.addCatch(failure);
+
 const chain = sfn.Chain.start(custom).next(finalStatus);
 
 const sm = new sfn.StateMachine(stack, 'StateMachine', {

diff --git a/packages/aws-cdk-lib/aws-stepfunctions/README.md b/packages/aws-cdk-lib/aws-stepfunctions/README.md
@@ -598,6 +598,10 @@ const custom = new sfn.CustomState(this, 'my custom task', {
   stateJson,
 });
 
+// catch errors with addCatch
+const errorHandler = new sfn.Pass(this, 'handle failure');
+custom.addCatch(errorHandler);
+
 const chain = sfn.Chain.start(custom)
   .next(finalStatus);
 

diff --git a/packages/aws-cdk-lib/aws-stepfunctions/lib/states/custom-state.ts b/packages/aws-cdk-lib/aws-stepfunctions/lib/states/custom-state.ts
@@ -1,7 +1,7 @@
 import { Construct } from 'constructs';
 import { State } from './state';
 import { Chain } from '..';
-import { IChainable, INextable } from '../types';
+import { CatchProps, IChainable, INextable } from '../types';
 
 /**
  * Properties for defining a custom state definition
@@ -25,7 +25,7 @@ export class CustomState extends State implements IChainable, INextable {
   /**
    * Amazon States Language (JSON-based) definition of the state
    */
-  private readonly stateJson: { [key: string]: any};
+  private readonly stateJson: { [key: string]: any };
 
   constructor(scope: Construct, id: string, props: CustomStateProps) {
     super(scope, id, {});
@@ -34,6 +34,17 @@ export class CustomState extends State implements IChainable, INextable {
     this.stateJson = props.stateJson;
   }
 
+  /**
+   * Add a recovery handler for this state
+   *
+   * When a particular error occurs, execution will continue at the error
+   * handler instead of failing the state machine execution.
+   */
+  public addCatch(handler: IChainable, props: CatchProps = {}): CustomState {
+    super._addCatch(handler.startState, props);
+    return this;
+  }
+
   /**
    * Continue normal execution with the given state
    */
@@ -49,6 +60,7 @@ export class CustomState extends State implements IChainable, INextable {
     return {
       ...this.renderNextEnd(),
       ...this.stateJson,
+      ...this.renderRetryCatch(),
     };
   }
 }
diff --git a/packages/aws-cdk-lib/aws-stepfunctions/test/custom-state.test.ts b/packages/aws-cdk-lib/aws-stepfunctions/test/custom-state.test.ts
@@ -33,6 +33,7 @@ describe('Custom State', () => {
     // THEN
     expect(customState.toStateJson()).toStrictEqual({
       ...stateJson,
+      ...{ Catch: undefined, Retry: undefined },
       End: true,
     });
   });
@@ -72,4 +73,51 @@ describe('Custom State', () => {
       },
     );
   });
+
+  test('can add a catch state', () => {
+    // GIVEN
+    const failure = new sfn.Fail(stack, 'failed', {
+      error: 'DidNotWork',
+      cause: 'We got stuck',
+    });
+    const custom = new sfn.CustomState(stack, 'Custom', {
+      stateJson,
+    });
+    const chain = sfn.Chain.start(custom);
+
+    // WHEN
+    custom.addCatch(failure);
+
+    // THEN
+    expect(render(stack, chain)).toStrictEqual(
+      {
+        StartAt: 'Custom',
+        States: {
+          Custom: {
+            Type: 'Task',
+            Resource: 'arn:aws:states:::dynamodb:putItem',
+            Parameters: {
+              TableName: 'MyTable',
+              Item: {
+                id: {
+                  S: 'MyEntry',
+                },
+              },
+            },
+            ResultPath: null,
+            Catch: [{
+              ErrorEquals: ['States.ALL'],
+              Next: 'failed',
+            }],
+            End: true,
+          },
+          failed: {
+            Type: 'Fail',
+            Error: 'DidNotWork',
+            Cause: 'We got stuck',
+          },
+        },
+      },
+    );
+  });
 });