From cf923bc4eacd3bb0efdc469708dc389a67d38bcc Mon Sep 17 00:00:00 2001 From: Grayden <38144548+graydenshand@users.noreply.github.com> Date: Tue, 19 Dec 2023 16:25:06 -0500 Subject: [PATCH] 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* --- ...functions-custom-state-integ.template.json | 2 +- .../test/integ.custom-state.ts | 7 +++ .../aws-cdk-lib/aws-stepfunctions/README.md | 4 ++ .../lib/states/custom-state.ts | 16 ++++++- .../test/custom-state.test.ts | 48 +++++++++++++++++++ 5 files changed, 74 insertions(+), 3 deletions(-) diff --git a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions/test/integ.custom-state.js.snapshot/aws-stepfunctions-custom-state-integ.template.json b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions/test/integ.custom-state.js.snapshot/aws-stepfunctions-custom-state-integ.template.json index fde50d40dc7b7..b090502f0294b 100644 --- a/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions/test/integ.custom-state.js.snapshot/aws-stepfunctions-custom-state-integ.template.json +++ b/packages/@aws-cdk-testing/framework-integ/test/aws-stepfunctions/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 index d65b0e373945b..5d28b12436975 100644 --- 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 index ab142c0f1ddbd..aef1e2ffaa631 100644 --- 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 index 3bf14fc249b80..db66a5a8de17f 100644 --- 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 index 98ccf387bad49..766b8ac659abe 100644 --- 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', + }, + }, + }, + ); + }); }); \ No newline at end of file