Merge pull request #546 from akshaydk/distributed_map_support

feat: add IAM role to distributed_map to start the state machine
serverless-operations · Jan 16, 2023 · 1c0b347 · 1c0b347
2 parents b1cb3b6 + c7d22ba
commit 1c0b347
Show file tree

Hide file tree

Showing 2 changed files with 79 additions and 7 deletions.
diff --git a/lib/deploy/stepFunctions/compileIamRole.js b/lib/deploy/stepFunctions/compileIamRole.js
@@ -8,19 +8,27 @@ const { getArnPartition } = require('../../utils/arn');
 
 const logger = require('../../utils/logger');
 
-function getTaskStates(states) {
+function getTaskStates(states, stateMachineName) {
   return _.flatMap(states, (state) => {
     switch (state.Type) {
       case 'Task': {
         return [state];
       }
       case 'Parallel': {
         const parallelStates = _.flatMap(state.Branches, branch => _.values(branch.States));
-        return getTaskStates(parallelStates);
+        return getTaskStates(parallelStates, stateMachineName);
       }
       case 'Map': {
         const mapStates = state.ItemProcessor ? state.ItemProcessor.States : state.Iterator.States;
-        const taskStates = getTaskStates(mapStates);
+        const taskStates = getTaskStates(mapStates, stateMachineName);
+        if (state.ItemProcessor && state.ItemProcessor.ProcessorConfig
+          && state.ItemProcessor.ProcessorConfig.Mode === 'DISTRIBUTED') {
+          taskStates.push({
+            Resource: 'arn:aws:states:::states:startExecution',
+            Mode: 'DISTRIBUTED',
+            StateMachineName: stateMachineName,
+          });
+        }
         if (state.ItemReader) {
           taskStates.push(state.ItemReader);
         }
@@ -303,9 +311,16 @@ function getLambdaPermissions(state) {
 }
 
 function getStepFunctionsPermissions(state) {
-  const stateMachineArn = state.Parameters['StateMachineArn.$']
-    ? '*'
-    : state.Parameters.StateMachineArn;
+  let stateMachineArn = state.Mode === 'DISTRIBUTED' ? {
+    'Fn::Sub': [
+      `arn:aws:states:\${AWS::Region}:\${AWS::AccountId}:stateMachine:${state.StateMachineName}`,
+    ],
+  } : null;
+
+  if (!stateMachineArn) {
+    stateMachineArn = state.Parameters['StateMachineArn.$'] ? '*'
+      : state.Parameters.StateMachineArn;
+  }
 
   return [{
     action: 'states:StartExecution',
@@ -591,7 +606,7 @@ module.exports = {
         throw new Error(`Missing "definition" for state machine ${stateMachineName}`);
       }
 
-      const taskStates = getTaskStates(stateMachineObj.definition.States);
+      const taskStates = getTaskStates(stateMachineObj.definition.States, stateMachineName);
       let iamPermissions = getIamPermissions.bind(this)(taskStates);
 
       if (stateMachineObj.loggingConfig) {

diff --git a/lib/deploy/stepFunctions/compileIamRole.test.js b/lib/deploy/stepFunctions/compileIamRole.test.js
@@ -2286,6 +2286,63 @@ describe('#compileIamRole', () => {
     expect(lambdaPermissions[0].Resource).to.deep.equal(lambdaArns);
   });
 
+  it('should support Distributed Map state type', () => {
+    const getStateMachine = (id, lambdaArn) => ({
+      id,
+      definition: {
+        StartAt: 'A',
+        States: {
+          A: {
+            Type: 'Map',
+            ItemProcessor: {
+              ProcessorConfig: {
+                Mode: 'DISTRIBUTED',
+              },
+              StartAt: 'B',
+              States: {
+                B: {
+                  Type: 'Task',
+                  Resource: lambdaArn,
+                  End: true,
+                },
+              },
+            },
+            End: true,
+          },
+        },
+      },
+    });
+
+    serverless.service.stepFunctions = {
+      stateMachines: {
+        myStateMachine: getStateMachine('StateMachine1', 'arn:aws:lambda:us-west-2:1234567890:function:foo'),
+      },
+    };
+
+    serverlessStepFunctions.compileIamRole();
+    const statements = serverlessStepFunctions.serverless.service
+      .provider.compiledCloudFormationTemplate.Resources.StateMachine1Role
+      .Properties.Policies[0].PolicyDocument.Statement;
+
+    const lambdaPermissions = statements.filter(s => _.isEqual(s.Action, ['lambda:InvokeFunction']));
+    expect(lambdaPermissions).to.have.lengthOf(1);
+
+    const lambdaArns = [
+      'arn:aws:lambda:us-west-2:1234567890:function:foo',
+      getAlias('arn:aws:lambda:us-west-2:1234567890:function:foo'),
+    ];
+    expect(lambdaPermissions[0].Resource).to.deep.equal(lambdaArns);
+
+    const stepFunctionPermission = statements.filter(s => _.isEqual(s.Action, ['states:StartExecution']));
+    expect(stepFunctionPermission).to.have.lengthOf(1);
+    expect(stepFunctionPermission[0].Resource).to.deep.eq([{
+      'Fn::Sub': [
+        'arn:aws:states:${AWS::Region}:${AWS::AccountId}:stateMachine:myStateMachine',
+      ],
+    },
+    ]);
+  });
+
   it('should support nested Map state type', () => {
     const getStateMachine = (id, lambdaArn1, lambdaArn2) => ({
       id,