Merge pull request #762 from JupiterOne/INT-4666-beforeAddRelationship-async

austinkelleher · web-flow · commit 2139b03dc669 · 2022-07-20T14:20:24.000-04:00
INT-4666 - Support async `beforeAddRelationship` hook
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -9,6 +9,11 @@ and this project adheres to
 
 ## Unreleased
 
+## 8.22.0 - 2021-07-20
+
+- Support async `beforeAddRelationship` hook in `IntegrationInvocationConfig`.
+  See the development documentation for more information on its usage.
+
 ## 8.21.0 - 2021-07-20
 
 ### Added
diff --git a/docs/integrations/development.md b/docs/integrations/development.md
@@ -286,7 +286,7 @@ data retrieved in the step as a partial dataset. See the
 [Failure handling](#failure-handling) section below for more information on
 partial datasets.
 
-#### `beforeAddEntity`
+#### `beforeAddEntity(context: IntegrationExecutionContext<IntegrationConfig>, e: Entity): Entity`
 
 `beforeAddEntity` is an optional hook function that can be provided. The
 function is called before an entity is added to the job state internally and the
@@ -323,7 +323,7 @@ export const invocationConfig: IntegrationInvocationConfig<IntegrationConfig> =
   };
 ```
 
-#### `beforeAddRelationship`
+#### `beforeAddRelationship(context: IntegrationExecutionContext<IntegrationConfig>, r: Relationship): Promise<Relationship> | Relationship`
 
 `beforeAddRelationship` is an optional hook function that can be provided. The
 function is called before a relationship is added to the job state internally
diff --git a/packages/cli/package.json b/packages/cli/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@jupiterone/cli",
-  "version": "8.21.0",
+  "version": "8.22.0",
   "description": "The JupiterOne cli",
   "main": "dist/src/index.js",
   "types": "dist/src/index.d.ts",
@@ -24,8 +24,8 @@
     "test": "jest"
   },
   "dependencies": {
-    "@jupiterone/integration-sdk-core": "^8.21.0",
-    "@jupiterone/integration-sdk-runtime": "^8.21.0",
+    "@jupiterone/integration-sdk-core": "^8.22.0",
+    "@jupiterone/integration-sdk-runtime": "^8.22.0",
     "@lifeomic/attempt": "^3.0.3",
     "commander": "^5.0.0",
     "globby": "^11.0.1",
diff --git a/packages/integration-sdk-benchmark/package.json b/packages/integration-sdk-benchmark/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@jupiterone/integration-sdk-benchmark",
-  "version": "8.21.0",
+  "version": "8.22.0",
   "private": true,
   "description": "SDK benchmarking scripts",
   "main": "./src/index.js",
@@ -15,8 +15,8 @@
     "benchmark": "for file in ./src/benchmarks/*; do yarn prebenchmark && node $file; done"
   },
   "dependencies": {
-    "@jupiterone/integration-sdk-core": "^8.21.0",
-    "@jupiterone/integration-sdk-runtime": "^8.21.0",
+    "@jupiterone/integration-sdk-core": "^8.22.0",
+    "@jupiterone/integration-sdk-runtime": "^8.22.0",
     "benchmark": "^2.1.4"
   }
 }
diff --git a/packages/integration-sdk-cli/package.json b/packages/integration-sdk-cli/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@jupiterone/integration-sdk-cli",
-  "version": "8.21.0",
+  "version": "8.22.0",
   "description": "The SDK for developing JupiterOne integrations",
   "main": "dist/src/index.js",
   "types": "dist/src/index.d.ts",
@@ -22,7 +22,7 @@
     "prepack": "yarn build:dist"
   },
   "dependencies": {
-    "@jupiterone/integration-sdk-runtime": "^8.21.0",
+    "@jupiterone/integration-sdk-runtime": "^8.22.0",
     "chalk": "^4",
     "commander": "^5.0.0",
     "fs-extra": "^10.1.0",
@@ -37,7 +37,7 @@
     "vis": "^4.21.0-EOL"
   },
   "devDependencies": {
-    "@jupiterone/integration-sdk-private-test-utils": "^8.21.0",
+    "@jupiterone/integration-sdk-private-test-utils": "^8.22.0",
     "@pollyjs/adapter-node-http": "^6.0.5",
     "@pollyjs/core": "^6.0.5",
     "@pollyjs/persister-fs": "^6.0.5",
diff --git a/packages/integration-sdk-core/package.json b/packages/integration-sdk-core/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@jupiterone/integration-sdk-core",
-  "version": "8.21.0",
+  "version": "8.22.0",
   "description": "The SDK for developing JupiterOne integrations",
   "main": "dist/src/index.js",
   "types": "dist/src/index.d.ts",
diff --git a/packages/integration-sdk-core/src/types/config.ts b/packages/integration-sdk-core/src/types/config.ts
@@ -28,7 +28,10 @@ export type BeforeAddEntityHookFunction<
 
 export type BeforeAddRelationshipHookFunction<
   TExecutionContext extends ExecutionContext,
-> = (context: TExecutionContext, relationship: Relationship) => Relationship;
+> = (
+  context: TExecutionContext,
+  relationship: Relationship,
+) => Promise<Relationship> | Relationship;
 
 export type LoadExecutionConfigFunction<
   TInstanceConfig extends IntegrationInstanceConfig = IntegrationInstanceConfig,
diff --git a/packages/integration-sdk-dev-tools/package.json b/packages/integration-sdk-dev-tools/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@jupiterone/integration-sdk-dev-tools",
-  "version": "8.21.0",
+  "version": "8.22.0",
   "description": "A collection of developer tools that will assist with building integrations.",
   "repository": "git@github.com:JupiterOne/sdk.git",
   "author": "JupiterOne <dev@jupiterone.io>",
@@ -15,8 +15,8 @@
     "access": "public"
   },
   "dependencies": {
-    "@jupiterone/integration-sdk-cli": "^8.21.0",
-    "@jupiterone/integration-sdk-testing": "^8.21.0",
+    "@jupiterone/integration-sdk-cli": "^8.22.0",
+    "@jupiterone/integration-sdk-testing": "^8.22.0",
     "@types/jest": "^27.1.0",
     "@types/node": "^14.0.5",
     "@typescript-eslint/eslint-plugin": "^4.22.0",
diff --git a/packages/integration-sdk-http-client/package.json b/packages/integration-sdk-http-client/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@jupiterone/integration-sdk-http-client",
-  "version": "8.21.0",
+  "version": "8.22.0",
   "description": "The HTTP client for use in JupiterOne integrations",
   "main": "dist/src/index.js",
   "types": "dist/src/index.d.ts",
@@ -27,8 +27,8 @@
     "node-fetch": "^2.6.0"
   },
   "devDependencies": {
-    "@jupiterone/integration-sdk-dev-tools": "^8.21.0",
-    "@jupiterone/integration-sdk-private-test-utils": "^8.21.0",
+    "@jupiterone/integration-sdk-dev-tools": "^8.22.0",
+    "@jupiterone/integration-sdk-private-test-utils": "^8.22.0",
     "fetch-mock-jest": "^1.5.1"
   },
   "bugs": {
diff --git a/packages/integration-sdk-private-test-utils/package.json b/packages/integration-sdk-private-test-utils/package.json
@@ -1,7 +1,7 @@
 {
   "name": "@jupiterone/integration-sdk-private-test-utils",
   "private": true,
-  "version": "8.21.0",
+  "version": "8.22.0",
   "description": "The SDK for developing JupiterOne integrations",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",
@@ -15,7 +15,7 @@
     "build:dist": "tsc -p tsconfig.json --declaration"
   },
   "dependencies": {
-    "@jupiterone/integration-sdk-core": "^8.21.0",
+    "@jupiterone/integration-sdk-core": "^8.22.0",
     "lodash": "^4.17.15",
     "uuid": "^7.0.3"
   },
diff --git a/packages/integration-sdk-runtime/package.json b/packages/integration-sdk-runtime/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@jupiterone/integration-sdk-runtime",
-  "version": "8.21.0",
+  "version": "8.22.0",
   "description": "The SDK for developing JupiterOne integrations",
   "main": "dist/src/index.js",
   "types": "dist/src/index.d.ts",
@@ -24,7 +24,7 @@
     "prepack": "yarn build:dist"
   },
   "dependencies": {
-    "@jupiterone/integration-sdk-core": "^8.21.0",
+    "@jupiterone/integration-sdk-core": "^8.22.0",
     "@lifeomic/alpha": "^1.4.0",
     "@lifeomic/attempt": "^3.0.3",
     "async-sema": "^3.1.0",
@@ -43,7 +43,7 @@
     "uuid": "^7.0.3"
   },
   "devDependencies": {
-    "@jupiterone/integration-sdk-private-test-utils": "^8.21.0",
+    "@jupiterone/integration-sdk-private-test-utils": "^8.22.0",
     "@types/uuid": "^7.0.2",
     "get-port": "^5.1.1",
     "memfs": "^3.2.0",
diff --git a/packages/integration-sdk-runtime/src/execution/__tests__/executeIntegration.test.ts b/packages/integration-sdk-runtime/src/execution/__tests__/executeIntegration.test.ts
@@ -706,6 +706,123 @@ describe('executeIntegrationInstance', () => {
     ]);
   });
 
+  test('should call async "beforeAddRelationship" hook if provided to config', async () => {
+    const config = createInstanceConfiguration({
+      invocationConfig: {
+        async beforeAddRelationship(_, relationship) {
+          await sleep(50);
+
+          return {
+            ...relationship,
+            customProp:
+              typeof relationship.customProp === 'undefined'
+                ? true
+                : relationship.customProp,
+          };
+        },
+        integrationSteps: [
+          {
+            id: 'my-step',
+            name: 'My awesome step',
+            entities: [
+              {
+                resourceName: 'The Test',
+                _type: 'test',
+                _class: 'Test',
+              },
+            ],
+            relationships: [],
+            async executionHandler({ jobState }) {
+              const e1 = await jobState.addEntity({
+                _key: 'test',
+                _type: 'test',
+                _class: 'Test',
+              });
+
+              const e2 = await jobState.addEntity({
+                _key: 'test1',
+                _type: 'test',
+                _class: 'Test',
+              });
+
+              const e3 = await jobState.addEntity({
+                _key: 'test2',
+                _type: 'test',
+                _class: 'Test',
+              });
+
+              await jobState.addRelationship(
+                createDirectRelationship({
+                  _class: RelationshipClass.HAS,
+                  from: e2,
+                  to: e1,
+                }),
+              );
+
+              await jobState.addRelationship(
+                createDirectRelationship({
+                  _class: RelationshipClass.HAS,
+                  from: e3,
+                  to: e2,
+                  properties: {
+                    customProp: false,
+                  },
+                }),
+              );
+            },
+          },
+        ],
+      },
+    });
+
+    await executeIntegrationInstanceWithConfig(config);
+    const sortedLocalGraphData = await getSortedLocalGraphData();
+
+    expect(sortedLocalGraphData).toEqual([
+      {
+        entities: [
+          {
+            _key: 'test',
+            _type: 'test',
+            _class: 'Test',
+          },
+          {
+            _key: 'test1',
+            _type: 'test',
+            _class: 'Test',
+          },
+          {
+            _key: 'test2',
+            _type: 'test',
+            _class: 'Test',
+          },
+        ],
+      },
+      {
+        relationships: [
+          {
+            _key: 'test1|has|test',
+            _type: 'test_has_',
+            _class: 'HAS',
+            _fromEntityKey: 'test1',
+            _toEntityKey: 'test',
+            displayName: 'HAS',
+            customProp: true,
+          },
+          {
+            _key: 'test2|has|test1',
+            _type: 'test_has_',
+            _class: 'HAS',
+            _fromEntityKey: 'test2',
+            _toEntityKey: 'test1',
+            displayName: 'HAS',
+            customProp: false,
+          },
+        ],
+      },
+    ]);
+  });
+
   test('publishes disk usage metric', async () => {
     const config = createInstanceConfiguration({
       invocationConfig: {
diff --git a/packages/integration-sdk-runtime/src/execution/dependencyGraph.ts b/packages/integration-sdk-runtime/src/execution/dependencyGraph.ts
@@ -470,7 +470,8 @@ function buildStepContext<
 
   const jobStateBeforeAddRelationship =
     typeof beforeAddRelationship !== 'undefined'
-      ? (r: Relationship): Relationship => beforeAddRelationship(context, r)
+      ? (r: Relationship): Promise<Relationship> | Relationship =>
+          beforeAddRelationship(context, r)
       : undefined;
 
   const jobState = createStepJobState({
diff --git a/packages/integration-sdk-runtime/src/execution/jobState.ts b/packages/integration-sdk-runtime/src/execution/jobState.ts
@@ -163,7 +163,9 @@ export interface CreateStepJobStateParams {
   dataStore: MemoryDataStore;
   uploader?: StepGraphObjectDataUploader;
   beforeAddEntity?: (entity: Entity) => Entity;
-  beforeAddRelationship?: (relationship: Relationship) => Relationship;
+  beforeAddRelationship?: (
+    relationship: Relationship,
+  ) => Promise<Relationship> | Relationship;
 }
 
 export function createStepJobState({
@@ -203,23 +205,37 @@ export function createStepJobState({
     return entities;
   };
 
-  const addRelationships = (relationships: Relationship[]) => {
+  function registerRelationshipInTrackers(r: Relationship) {
+    duplicateKeyTracker.registerKey(r._key, {
+      _type: r._type,
+      _key: r._key,
+    });
+
+    typeTracker.addStepGraphObjectType({
+      stepId,
+      _type: r._type,
+      count: 1,
+    });
+  }
+
+  const addRelationships = async (relationshipsToAdd: Relationship[]) => {
+    let relationships: Relationship[];
+
     if (beforeAddRelationship) {
-      relationships = relationships.map(beforeAddRelationship);
-    }
+      relationships = [];
 
-    relationships.forEach((r) => {
-      duplicateKeyTracker.registerKey(r._key, {
-        _type: r._type,
-        _key: r._key,
-      });
+      for (const relationship of relationshipsToAdd) {
+        const newRelationship = await beforeAddRelationship(relationship);
+        relationships.push(newRelationship);
 
-      typeTracker.addStepGraphObjectType({
-        stepId,
-        _type: r._type,
-        count: 1,
-      });
-    });
+        // Avoid iterating the entire set of relationships again later by
+        // registering now.
+        registerRelationshipInTrackers(newRelationship);
+      }
+    } else {
+      relationships = relationshipsToAdd;
+      relationships.forEach(registerRelationshipInTrackers);
+    }
 
     return graphObjectStore.addRelationships(
       stepId,
diff --git a/packages/integration-sdk-testing/package.json b/packages/integration-sdk-testing/package.json

Original file line number	Diff line number	Diff line change
`@@ -1,6 +1,6 @@`
`1`	`1`	`{`
`2`	`2`	`"name": "@jupiterone/integration-sdk-benchmark",`
`3`		`- "version": "8.21.0",`
	`3`	`+ "version": "8.22.0",`
`4`	`4`	`"private": true,`
`5`	`5`	`"description": "SDK benchmarking scripts",`
`6`	`6`	`"main": "./src/index.js",`
`@@ -15,8 +15,8 @@`
`15`	`15`	`"benchmark": "for file in ./src/benchmarks/*; do yarn prebenchmark && node $file; done"`
`16`	`16`	`},`
`17`	`17`	`"dependencies": {`
`18`		`- "@jupiterone/integration-sdk-core": "^8.21.0",`
`19`		`- "@jupiterone/integration-sdk-runtime": "^8.21.0",`
	`18`	`+ "@jupiterone/integration-sdk-core": "^8.22.0",`
	`19`	`+ "@jupiterone/integration-sdk-runtime": "^8.22.0",`
`20`	`20`	`"benchmark": "^2.1.4"`
`21`	`21`	`}`
`22`	`22`	`}`
Original file line number	Diff line number	Diff line change
`@@ -1,6 +1,6 @@`
`1`	`1`	`{`
`2`	`2`	`"name": "@jupiterone/integration-sdk-core",`
`3`		`- "version": "8.21.0",`
	`3`	`+ "version": "8.22.0",`
`4`	`4`	`"description": "The SDK for developing JupiterOne integrations",`
`5`	`5`	`"main": "dist/src/index.js",`
`6`	`6`	`"types": "dist/src/index.d.ts",`