Merge pull request #255 from JupiterOne/254-async-getStepStartStates

Fixes #254 - Add support for async 'getStepStartStates' in 'invocatio…
JupiterOne · Jul 28, 2020 · 7614529 · 7614529
2 parents 18979b3 + 401b463
commit 7614529
Show file tree

Hide file tree

Showing 5 changed files with 88 additions and 6 deletions.
diff --git a/packages/integration-sdk-core/src/types/step.ts b/packages/integration-sdk-core/src/types/step.ts
@@ -17,7 +17,7 @@ export type StepStartStates = Record<string, StepStartState>;
 
 export type GetStepStartStatesFunction<T extends ExecutionContext> = (
   context: T,
-) => StepStartStates;
+) => StepStartStates | Promise<StepStartStates>;
 
 export type ExecutionHandlerFunction<T extends StepExecutionContext> = (
   context: T,

diff --git a/packages/integration-sdk-runtime/src/execution/__tests__/executeIntegration.test.ts b/packages/integration-sdk-runtime/src/execution/__tests__/executeIntegration.test.ts
@@ -25,6 +25,12 @@ import {
 
 jest.mock('fs');
 
+function sleep(ms: number) {
+  return new Promise((resolve) => {
+    setTimeout(() => resolve(), ms);
+  });
+}
+
 afterEach(() => {
   vol.reset();
   delete process.env.ENABLE_GRAPH_OBJECT_SCHEMA_VALIDATION;
@@ -284,6 +290,60 @@ describe('executeIntegrationInstance', () => {
     });
   });
 
+  test('does not include partial data sets for disabled steps in async "getStepStartStates"', async () => {
+    invocationConfig = {
+      ...invocationConfig,
+      getStepStartStates: async () => {
+        await sleep(5);
+
+        return {
+          'my-step-b': { disabled: true },
+          'my-step-a': { disabled: false },
+        };
+      },
+      integrationSteps: [
+        {
+          id: 'my-step-a',
+          name: 'My awesome step',
+          types: ['test_a'],
+          executionHandler: jest
+            .fn()
+            .mockRejectedValue(new Error('something broke')),
+        },
+        {
+          id: 'my-step-b',
+          name: 'My awesome step',
+          types: ['test_b'],
+          executionHandler: jest.fn(),
+        },
+      ],
+    };
+
+    await expect(execute()).resolves.toEqual({
+      integrationStepResults: [
+        {
+          id: 'my-step-a',
+          name: 'My awesome step',
+          declaredTypes: ['test_a'],
+          encounteredTypes: [],
+          status: StepResultStatus.FAILURE,
+        },
+        {
+          id: 'my-step-b',
+          name: 'My awesome step',
+          declaredTypes: ['test_b'],
+          encounteredTypes: [],
+          status: StepResultStatus.DISABLED,
+        },
+      ],
+      metadata: {
+        partialDatasets: {
+          types: ['test_a'],
+        },
+      },
+    });
+  });
+
   test('clears out the storage directory prior to performing collection', async () => {
     const previousContentFilePath = path.resolve(
       getRootStorageDirectory(),

diff --git a/packages/integration-sdk-runtime/src/execution/executeIntegration.ts b/packages/integration-sdk-runtime/src/execution/executeIntegration.ts
@@ -7,6 +7,7 @@ import {
   IntegrationLogger,
   ExecutionContext,
   StepExecutionContext,
+  StepStartStates,
 } from '@jupiterone/integration-sdk-core';
 
 import {
@@ -101,8 +102,8 @@ export async function executeWithContext<
     throw err;
   }
 
-  const stepStartStates =
-    config.getStepStartStates?.(context) ??
+  const stepStartStates: StepStartStates =
+    (await config.getStepStartStates?.(context)) ??
     getDefaultStepStartStates(config.integrationSteps);
 
   validateStepStartStates(config.integrationSteps, stepStartStates);

diff --git a/packages/integration-sdk-runtime/src/execution/step.ts b/packages/integration-sdk-runtime/src/execution/step.ts
@@ -93,8 +93,10 @@ export function prepareLocalStepCollection<
   const originalGetStepStartStates = config.getStepStartStates;
 
   config.getStepStartStates = stepsToRun.length
-    ? (ctx) => {
-        const originalEnabledRecord = originalGetStepStartStates?.(ctx) ?? {};
+    ? async (ctx) => {
+        const originalEnabledRecord = await (originalGetStepStartStates?.(
+          ctx,
+        ) ?? {});
         const enabledRecord: StepStartStates = {};
         for (const stepId of allStepIds) {
           const originalValue = originalEnabledRecord[stepId] ?? {};
@@ -112,7 +114,8 @@ export function prepareLocalStepCollection<
         }
         return enabledRecord;
       }
-    : originalGetStepStartStates;
+    : originalGetStepStartStates &&
+      (async (ctx) => originalGetStepStartStates(ctx));
 
   return config;
 }
diff --git a/packages/integration-sdk/CHANGELOG.md b/packages/integration-sdk/CHANGELOG.md
@@ -9,6 +9,24 @@ and this project adheres to
 
 ## Unreleased
 
+### Changed
+
+- Updated `IntegrationConfig` to support asynchronous `getStepStartStates`. See
+  [#254](https://github.com/JupiterOne/sdk/issues/254) for more information.
+
+Example:
+
+```typescript
+export const invocationConfig: IntegrationInvocationConfig<IntegrationConfig> = {
+  async getStepStartStates(ctx) {
+    return {
+      'fetch-users': { disabled: await checkFetchUsersStepDisabled(ctx) }
+    };
+  },
+  ...
+};
+```
+
 ## 2.4.0 - 2020-07-22
 
 ### Added