JupiterOne · austinkelleher · Dec 18, 2020 · Dec 13, 2020 · Dec 14, 2020 · Dec 16, 2020
diff --git a/packages/cli/package.json b/packages/cli/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@jupiterone/cli",
-  "version": "5.1.0",
+  "version": "5.2.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": "^5.1.0",
-    "@jupiterone/integration-sdk-runtime": "^5.1.0",
+    "@jupiterone/integration-sdk-core": "^5.2.0",
+    "@jupiterone/integration-sdk-runtime": "^5.2.0",
     "@lifeomic/attempt": "^3.0.0",
     "commander": "^5.0.0",
     "globby": "^11.0.1",

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": "5.1.0",
+  "version": "5.2.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": "^5.1.0",
+    "@jupiterone/integration-sdk-runtime": "^5.2.0",
     "commander": "^5.0.0",
     "globby": "^11.0.0",
     "lodash": "^4.17.19",
@@ -31,7 +31,7 @@
     "vis": "^4.21.0-EOL"
   },
   "devDependencies": {
-    "@jupiterone/integration-sdk-private-test-utils": "^5.1.0",
+    "@jupiterone/integration-sdk-private-test-utils": "^5.2.0",
     "@pollyjs/adapter-node-http": "^4.0.4",
     "@pollyjs/core": "^4.0.4",
     "@pollyjs/persister-fs": "^4.0.4",

diff --git a/packages/integration-sdk-cli/src/commands/collect.ts b/packages/integration-sdk-cli/src/commands/collect.ts
@@ -2,6 +2,7 @@ import { createCommand } from 'commander';
 
 import {
   executeIntegrationLocally,
+  FileSystemGraphObjectStore,
   prepareLocalStepCollection,
 } from '@jupiterone/integration-sdk-runtime';
 
@@ -30,6 +31,11 @@ export function collect() {
       const enableSchemaValidation = !options.disableSchemaValidation;
       const config = prepareLocalStepCollection(await loadConfig(), options);
       log.info('\nConfiguration loaded! Running integration...\n');
+
+      const graphObjectStore = new FileSystemGraphObjectStore({
+        prettifyFiles: true,
+      });
+
       const results = await executeIntegrationLocally(
         config,
         {
@@ -39,6 +45,7 @@ export function collect() {
         },
         {
           enableSchemaValidation,
+          graphObjectStore,
         },
       );
       log.displayExecutionResults(results);

diff --git a/packages/integration-sdk-cli/src/commands/run.ts b/packages/integration-sdk-cli/src/commands/run.ts
@@ -8,16 +8,19 @@ import {
   createIntegrationInstanceForLocalExecution,
   createIntegrationLogger,
   executeIntegrationInstance,
+  FileSystemGraphObjectStore,
   finalizeSynchronization,
   getAccountFromEnvironment,
   getApiBaseUrl,
   getApiKeyFromEnvironment,
   initiateSynchronization,
-  uploadCollectedData,
 } from '@jupiterone/integration-sdk-runtime';
 
 import { loadConfig } from '../config';
 import * as log from '../log';
+import { createPersisterApiStepGraphObjectDataUploader } from '@jupiterone/integration-sdk-runtime/dist/src/execution/uploader';
+
+const DEFAULT_UPLOAD_CONCURRENCY = 5;
 
 export function run() {
   return createCommand('run')
@@ -70,6 +73,10 @@ export function run() {
 
       const invocationConfig = await loadConfig();
 
+      const graphObjectStore = new FileSystemGraphObjectStore({
+        prettifyFiles: true,
+      });
+
       try {
         const executionResults = await executeIntegrationInstance(
           logger,
@@ -82,15 +89,21 @@ export function run() {
           },
           {
             enableSchemaValidation: true,
+            graphObjectStore,
+            createStepGraphObjectDataUploader(stepId) {
+              return createPersisterApiStepGraphObjectDataUploader({
+                stepId,
+                synchronizationJobContext: synchronizationContext,
+                uploadConcurrency: DEFAULT_UPLOAD_CONCURRENCY,
+              });
+            },
           },
         );
 
         await eventPublishingQueue.onIdle();
 
         log.displayExecutionResults(executionResults);
 
-        await uploadCollectedData(synchronizationContext);
-
         const synchronizationResult = await finalizeSynchronization({
           ...synchronizationContext,
           partialDatasets: executionResults.metadata.partialDatasets,

diff --git a/packages/integration-sdk-cli/src/visualization/generateVisualization.ts b/packages/integration-sdk-cli/src/visualization/generateVisualization.ts
@@ -27,9 +27,11 @@ export async function generateVisualization(
     log.warn(`Unable to find any files under path: ${resolvedIntegrationPath}`);
   }
 
-  const { entities, relationships, mappedRelationships } = await retrieveIntegrationData(
-    entitiesAndRelationshipPaths,
-  );
+  const {
+    entities,
+    relationships,
+    mappedRelationships,
+  } = await retrieveIntegrationData(entitiesAndRelationshipPaths);
 
   const nodeDataSets = entities.map((entity) => ({
     id: getNodeIdFromEntity(entity, []),
@@ -45,18 +47,21 @@ export async function generateVisualization(
   );
 
   const {
-    mappedRelationshipEdges, 
+    mappedRelationshipEdges,
     mappedRelationshipNodes,
   } = createMappedRelationshipNodesAndEdges({
-    mappedRelationships, 
+    mappedRelationships,
     explicitEntities: entities,
   });
 
   const htmlFileLocation = path.join(resolvedIntegrationPath, 'index.html');
 
   await writeFileToPath({
     path: htmlFileLocation,
-    content: generateVisHTML([...nodeDataSets, ...mappedRelationshipNodes], [...explicitEdgeDataSets, ...mappedRelationshipEdges]),
+    content: generateVisHTML(
+      [...nodeDataSets, ...mappedRelationshipNodes],
+      [...explicitEdgeDataSets, ...mappedRelationshipEdges],
+    ),
   });
 
   return htmlFileLocation;

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": "5.1.0",
+  "version": "5.2.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/jobState.ts b/packages/integration-sdk-core/src/types/jobState.ts
@@ -112,4 +112,10 @@ export interface JobState {
    * flushed to reduce memory consumption.
    */
   flush: () => Promise<void>;
+
+  /**
+   * A job state may be created with a graph object uploader. This function
+   * resolves when all uploads have been completed.
+   */
+  waitUntilUploadsComplete?: () => Promise<void>;
 }
diff --git a/packages/integration-sdk-core/src/types/storage.ts b/packages/integration-sdk-core/src/types/storage.ts
@@ -11,11 +11,16 @@ import { Relationship } from './relationship';
  * integration execution.
  */
 export interface GraphObjectStore {
-  addEntities(stepId: string, newEntities: Entity[]): Promise<void>;
+  addEntities(
+    stepId: string,
+    newEntities: Entity[],
+    onEntitiesFlushed?: (entities: Entity[]) => Promise<void>,
+  ): Promise<void>;
 
   addRelationships(
     stepId: string,
     newRelationships: Relationship[],
+    onRelationshipsFlushed?: (relationships: Relationship[]) => Promise<void>,
   ): Promise<void>;
 
   getEntity({ _key, _type }: GraphObjectLookupKey): Promise<Entity>;
@@ -30,5 +35,8 @@ export interface GraphObjectStore {
     iteratee: GraphObjectIteratee<T>,
   ): Promise<void>;
 
-  flush(): Promise<void>;
+  flush(
+    onEntitiesFlushed?: (entities: Entity[]) => Promise<void>,
+    onRelationshipsFlushed?: (relationships: Relationship[]) => Promise<void>,
+  ): Promise<void>;
 }
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": "5.1.0",
+  "version": "5.2.0",
   "description": "A collection of developer tools that will assist with building integrations.",
   "repository": "[email protected]:JupiterOne/sdk.git",
   "author": "JupiterOne <[email protected]>",
@@ -15,8 +15,8 @@
     "access": "public"
   },
   "dependencies": {
-    "@jupiterone/integration-sdk-cli": "^5.1.0",
-    "@jupiterone/integration-sdk-testing": "^5.1.0",
+    "@jupiterone/integration-sdk-cli": "^5.2.0",
+    "@jupiterone/integration-sdk-testing": "^5.2.0",
     "@types/jest": "^25.2.3",
     "@types/node": "^14.0.5",
     "@typescript-eslint/eslint-plugin": "^3.8.0",

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": "5.1.0",
+  "version": "5.2.0",
   "description": "The SDK for developing JupiterOne integrations",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",
@@ -12,7 +12,7 @@
     "node": "10.x || 12.x || 14.x"
   },
   "dependencies": {
-    "@jupiterone/integration-sdk-core": "^5.1.0",
+    "@jupiterone/integration-sdk-core": "^5.2.0",
     "lodash": "^4.17.15",
     "uuid": "^7.0.3"
   },

diff --git a/packages/integration-sdk-private-test-utils/src/graphObject.ts b/packages/integration-sdk-private-test-utils/src/graphObject.ts
@@ -0,0 +1,26 @@
+import { Entity, ExplicitRelationship } from '@jupiterone/integration-sdk-core';
+import { v4 as uuid } from 'uuid';
+
+export function createTestEntity(partial?: Partial<Entity>): Entity {
+  return {
+    _key: uuid(),
+    _class: uuid(),
+    _type: uuid(),
+    [uuid()]: uuid(),
+    ...partial,
+  };
+}
+
+export function createTestRelationship(
+  partial?: Partial<ExplicitRelationship>,
+): ExplicitRelationship {
+  return {
+    _key: uuid(),
+    _toEntityKey: uuid(),
+    _fromEntityKey: uuid(),
+    _class: uuid(),
+    _type: uuid(),
+    [uuid()]: uuid(),
+    ...partial,
+  };
+}
diff --git a/packages/integration-sdk-private-test-utils/src/index.ts b/packages/integration-sdk-private-test-utils/src/index.ts
@@ -1,3 +1,5 @@
 export * from './loadProjectStructure';
 export * from './toUnixPath';
 export * from './graphObjectStore';
+export * from './graphObject';
+export * from './util';
diff --git a/packages/integration-sdk-private-test-utils/src/util.ts b/packages/integration-sdk-private-test-utils/src/util.ts
@@ -0,0 +1,5 @@
+export function sleep(ms: number) {
+  return new Promise((resolve) => {
+    setTimeout(() => resolve(), ms);
+  });
+}
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": "5.1.0",
+  "version": "5.2.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": "^5.1.0",
+    "@jupiterone/integration-sdk-core": "^5.2.0",
     "@lifeomic/alpha": "^1.1.3",
     "@lifeomic/attempt": "^3.0.0",
     "async-sema": "^3.1.0",
@@ -43,7 +43,7 @@
     "uuid": "^7.0.3"
   },
   "devDependencies": {
-    "@jupiterone/integration-sdk-private-test-utils": "^5.1.0",
+    "@jupiterone/integration-sdk-private-test-utils": "^5.2.0",
     "@types/uuid": "^7.0.2",
     "get-port": "^5.1.1",
     "memfs": "^3.2.0",

diff --git a/packages/integration-sdk-runtime/src/__tests__/index.test.ts b/packages/integration-sdk-runtime/src/__tests__/index.test.ts
@@ -1,10 +1,6 @@
-import { BucketMap, FileSystemGraphObjectStore } from '../index';
+import { FileSystemGraphObjectStore } from '../index';
 
 describe('#storage', () => {
-  test('should expose BucketMap', () => {
-    expect(BucketMap).not.toEqual(undefined);
-  });
-
   test('should expose FileSystemGraphObjectStore', () => {
     expect(FileSystemGraphObjectStore).not.toEqual(undefined);
   });