feat(stepfunctions): custom state as an escape hatch

Custom State which enables the capability to provide Amazon States Language (ASL) JSON as an escape hatch. Useful when there are capabilities that are offered through Step Functions such as service integrations, and state properties but there isn't support through the CDK yet.

It enables the usage of all service integrations we don't currently support.

Author: shivlaks

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

@@ -113,6 +113,7 @@ are supported:
 * [`Succeed`](#succeed)
 * [`Fail`](#fail)
 * [`Map`](#map)
+* [`Custom State`](#custom-state)
 
 An arbitrary JSON object (specified at execution start) is passed from state to
 state and transformed during the execution of the workflow. For more
@@ -256,6 +257,76 @@ const map = new stepfunctions.Map(this, 'Map State', {
 map.iterator(new stepfunctions.Pass(this, 'Pass State'));
 ```
 
+### Custom State
+
+It's possible that the high-level constructs for the states or `stepfunctions-tasks` do not have
+the states or service integrations you are looking for. The primary reasons for this lack of
+functionality are:
+
+* A [service integration](https://docs.aws.amazon.com/step-functions/latest/dg/concepts-service-integrations.html) is available through Amazon States Langauge, but not available as construct
+  classes in the CDK.
+* The state or state properties are available through Step Functions, but are not configurable
+  through constructs
+
+If a feature is not available, a `CustomState` can be used to supply any Amazon States Language
+JSON-based object as the state definition.
+
+[Code Snippets](https://docs.aws.amazon.com/step-functions/latest/dg/tutorial-code-snippet.html#tutorial-code-snippet-1) are available and can be plugged in as the state definition.
+
+Custom states can be chained together with any of the other states to create your state machine
+definition. You will also need to provide any permissions that are required to the `role` that
+the State Machine uses.
+
+The following example uses the `DynamoDB` service integration to insert data into a DynamoDB table.
+
+```ts
+import * as ddb from '@aws-cdk/aws-dynamodb';
+import * as cdk from '@aws-cdk/core';
+import * as sfn from '@aws-cdk/aws-stepfunctions';
+
+// create a table
+const table = new ddb.Table(this, 'montable', {
+  partitionKey: {
+    name: 'id',
+    type: ddb.AttributeType.STRING,
+  },
+});
+
+const finalStatus = new sfn.Pass(stack, 'final step');
+
+// States language JSON to put an item into DynamoDB
+// snippet generated from https://docs.aws.amazon.com/step-functions/latest/dg/tutorial-code-snippet.html#tutorial-code-snippet-1
+const stateJson = {
+  Type: 'Task',
+  Resource: 'arn:aws:states:::dynamodb:putItem',
+  Parameters: {
+    TableName: table.tableName,
+    Item: {
+      id: {
+        S: 'MyEntry',
+      },
+    },
+  },
+  ResultPath: null,
+};
+
+// custom state which represents a task to insert data into DynamoDB
+const custom = new sfn.CustomState(this, 'my custom task', {
+  stateJson,
+});
+
+const chain = sfn.Chain.start(custom)
+      .next(finalStatus);
+
+const sm = new sfn.StateMachine(this, 'StateMachine', {
+  definition: chain,
+  timeout: cdk.Duration.seconds(30),
+});
+
+// don't forget permissions. You need to assign them
+table.grantWriteData(sm.role);
+```
+
 ## Task Chaining
 
 To make defining work flows as convenient (and readable in a top-to-bottom way)

packages/@aws-cdk/aws-stepfunctions/lib/index.ts

@@ -19,6 +19,7 @@ export * from './states/succeed';
 export * from './states/task';
 export * from './states/wait';
 export * from './states/map';
+export * from './states/custom-state';
 
 // AWS::StepFunctions CloudFormation Resources:
 export * from './stepfunctions.generated';

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

@@ -0,0 +1,55 @@
+import * as cdk from '@aws-cdk/core';
+import { Chain } from '..';
+import { IChainable, INextable } from '../types';
+import { State } from './state';
+
+/**
+ * Properties for defining a custom state definition
+ */
+export interface CustomStateProps {
+  /**
+   * Amazon States Language (JSON-based) definition of the state
+   *
+   * @see https://docs.aws.amazon.com/step-functions/latest/dg/concepts-amazon-states-language.html
+   */
+  readonly stateJson: { [key: string]: any };
+}
+
+/**
+ * State defined by supplying Amazon States Language (ASL) in the state machine.
+ *
+ * @experimental
+ */
+export class CustomState extends State implements IChainable, INextable {
+  public readonly endStates: INextable[];
+
+  /**
+   * Amazon States Language (JSON-based) definition of the state
+   */
+  private readonly stateJson: { [key: string]: any};
+
+  constructor(scope: cdk.Construct, id: string, props: CustomStateProps) {
+    super(scope, id, {});
+
+    this.endStates = [this];
+    this.stateJson = props.stateJson;
+  }
+
+  /**
+   * Continue normal execution with the given state
+   */
+  public next(next: IChainable): Chain {
+    super.makeNext(next.startState);
+    return Chain.sequence(this, next);
+  }
+
+  /**
+   * Returns the Amazon States Language object for this state
+   */
+  public toStateJson(): object {
+    return {
+      ...this.renderNextEnd(),
+      ...this.stateJson,
+    };
+  }
+}

packages/@aws-cdk/aws-stepfunctions/test/custom-state.test.ts

@@ -0,0 +1,34 @@
+import '@aws-cdk/assert/jest';
+import * as cdk from '@aws-cdk/core';
+import * as stepfunctions from '../lib';
+
+describe('Custom State', () => {
+  test('maintains the state Json provided during construction', () => {
+    // GIVEN
+    const stack = new cdk.Stack();
+    const stateJson = {
+      Type: 'Task',
+      Resource: 'arn:aws:states:::dynamodb:putItem',
+      Parameters: {
+        TableName: 'MyTable',
+        Item: {
+          id: {
+            S: 'MyEntry',
+          },
+        },
+      },
+      ResultPath: null,
+    };
+
+    // WHEN
+    const customState = new stepfunctions.CustomState(stack, 'Custom', {
+      stateJson,
+    });
+
+    // THEN
+    expect(customState.toStateJson()).toStrictEqual({
+      ...stateJson,
+      End: true,
+    });
+  });
+});

packages/@aws-cdk/aws-stepfunctions/test/integ.custom-state.expected.json

@@ -0,0 +1,54 @@
+{
+  "Resources": {
+    "StateMachineRoleB840431D": {
+      "Type": "AWS::IAM::Role",
+      "Properties": {
+        "AssumeRolePolicyDocument": {
+          "Statement": [
+            {
+              "Action": "sts:AssumeRole",
+              "Effect": "Allow",
+              "Principal": {
+                "Service": {
+                  "Fn::Join": [
+                    "",
+                    [
+                      "states.",
+                      {
+                        "Ref": "AWS::Region"
+                      },
+                      ".amazonaws.com"
+                    ]
+                  ]
+                }
+              }
+            }
+          ],
+          "Version": "2012-10-17"
+        }
+      }
+    },
+    "StateMachine2E01A3A5": {
+      "Type": "AWS::StepFunctions::StateMachine",
+      "Properties": {
+        "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}",
+        "RoleArn": {
+          "Fn::GetAtt": [
+            "StateMachineRoleB840431D",
+            "Arn"
+          ]
+        }
+      },
+      "DependsOn": [
+        "StateMachineRoleB840431D"
+      ]
+    }
+  },
+  "Outputs": {
+    "StateMachineARN": {
+      "Value": {
+        "Ref": "StateMachine2E01A3A5"
+      }
+    }
+  }
+}

packages/@aws-cdk/aws-stepfunctions/test/integ.custom-state.ts

@@ -0,0 +1,43 @@
+import * as cdk from '@aws-cdk/core';
+import * as sfn from '../lib';
+
+/*
+ * Stack verification steps:
+ *
+ * -- aws stepfunctions describe-state-machine --state-machine-arn <stack-output> has a status of `ACTIVE`
+ */
+const app = new cdk.App();
+const stack = new cdk.Stack(app, 'aws-stepfunctions-custom-state-integ');
+
+const finalStatus = new sfn.Pass(stack, 'final step');
+
+const stateJson = {
+  Type: 'Task',
+  Resource: 'arn:aws:states:::dynamodb:putItem',
+  Parameters: {
+    TableName: 'my-cool-table',
+    Item: {
+      id: {
+        S: 'my-entry',
+      },
+    },
+  },
+  ResultPath: null,
+};
+
+const custom = new sfn.CustomState(stack, 'my custom task', {
+  stateJson,
+});
+
+const chain = sfn.Chain.start(custom).next(finalStatus);
+
+const sm = new sfn.StateMachine(stack, 'StateMachine', {
+  definition: chain,
+  timeout: cdk.Duration.seconds(30),
+});
+
+new cdk.CfnOutput(stack, 'StateMachineARN', {
+  value: sm.stateMachineArn,
+});
+
+app.synth();