Support for continuous integration data uploads

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);

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,

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;

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>;
 }

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>;
 }

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,
+  };
+}

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';

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);
+  });
+}

packages/integration-sdk-runtime/src/execution/__tests__/dependencyGraph.test.ts

@@ -31,6 +31,12 @@ import {
 import { LOCAL_INTEGRATION_INSTANCE } from '../instance';
 import { DuplicateKeyTracker } from '../jobState';
 import { getDefaultStepStartStates } from '../step';
+import {
+  CreateStepGraphObjectDataUploaderFunction,
+  StepGraphObjectDataUploader,
+} from '../uploader';
+import { FlushedGraphObjectData } from '../../storage/types';
+import { createTestEntity } from '@jupiterone/integration-sdk-private-test-utils';
 
 jest.mock('fs');
 
@@ -134,13 +140,15 @@ describe('executeStepDependencyGraph', () => {
     steps: IntegrationStep[],
     stepStartStates: StepStartStates = getDefaultStepStartStates(steps),
     graphObjectStore: GraphObjectStore = new FileSystemGraphObjectStore(),
+    createStepGraphObjectDataUploader?: CreateStepGraphObjectDataUploaderFunction,
   ) {
     return executeStepDependencyGraph({
       executionContext,
       inputGraph: buildStepDependencyGraph(steps),
       stepStartStates,
       duplicateKeyTracker: new DuplicateKeyTracker(),
       graphObjectStore,
+      createStepGraphObjectDataUploader,
     });
   }
 
@@ -421,7 +429,7 @@ describe('executeStepDependencyGraph', () => {
 
       // each step should have just generated one file
       const writtenData = await fs.readFile(`${directory}/${files[0]}`, 'utf8');
-      expect(writtenData).toEqual(JSON.stringify({ [type]: data }, null, 2));
+      expect(writtenData).toEqual(JSON.stringify({ [type]: data }));
     }
   });
 
@@ -622,6 +630,156 @@ describe('executeStepDependencyGraph', () => {
     expect(spyB).toHaveBeenCalledBefore(spyC);
   });
 
+  test('should mark steps with failed executionHandlers with status FAILURE and dependent steps with status PARTIAL_SUCCESS_DUE_TO_DEPENDENCY_FAILURE when step upload fails', async () => {
+    const spyA = jest.fn();
+    const spyB = jest.fn();
+    const spyC = jest.fn();
+
+    const eA = createTestEntity();
+    const eB = createTestEntity();
+    const eC = createTestEntity();
+
+    const steps: IntegrationStep[] = [
+      {
+        id: 'a',
+        name: 'a',
+        entities: [],
+        relationships: [],
+        async executionHandler({ jobState }) {
+          await jobState.addEntity(eA);
+          spyA();
+        },
+      },
+      {
+        id: 'b',
+        name: 'b',
+        entities: [],
+        relationships: [],
+        dependsOn: ['a'],
+        async executionHandler({ jobState }) {
+          await jobState.addEntity(eB);
+          spyB();
+        },
+      },
+      {
+        id: 'c',
+        name: 'c',
+        entities: [],
+        relationships: [],
+        dependsOn: ['b'],
+        async executionHandler({ jobState }) {
+          await jobState.addEntity(eC);
+          spyC();
+        },
+      },
+    ];
+
+    const stepStartStates = getDefaultStepStartStates(steps);
+    const graphObjectStore = new FileSystemGraphObjectStore();
+
+    function createPassingUploader(
+      stepId: string,
+      collector: FlushedGraphObjectData[],
+    ): StepGraphObjectDataUploader {
+      return {
+        stepId,
+        async enqueue(graphObjectData) {
+          collector.push(graphObjectData);
+          return Promise.resolve();
+        },
+        waitUntilUploadsComplete() {
+          return Promise.resolve();
+        },
+      };
+    }
+
+    function createFailingUploader(
+      stepId: string,
+    ): StepGraphObjectDataUploader {
+      return {
+        stepId,
+        async enqueue() {
+          return Promise.resolve();
+        },
+        waitUntilUploadsComplete() {
+          return Promise.reject(new Error('expected upload wait failure'));
+        },
+      };
+    }
+
+    const passingUploaderCollector: FlushedGraphObjectData[] = [];
+
+    /**
+     * Graph:
+     * a - b - c
+     *
+     * In this situation, 'a' is the leaf node
+     * 'b' depends on 'a',
+     * 'c' depends on 'b'
+     */
+    const result = await executeSteps(
+      steps,
+      stepStartStates,
+      graphObjectStore,
+      (stepId) => {
+        if (stepId === 'b') {
+          return createFailingUploader(stepId);
+        } else {
+          return createPassingUploader(stepId, passingUploaderCollector);
+        }
+      },
+    );
+
+    const expectedCollected: FlushedGraphObjectData[] = [
+      {
+        entities: [eA],
+        relationships: [],
+      },
+      {
+        entities: [eC],
+        relationships: [],
+      },
+    ];
+
+    expect(passingUploaderCollector).toEqual(expectedCollected);
+
+    expect(result).toEqual([
+      {
+        id: 'a',
+        name: 'a',
+        declaredTypes: [],
+        partialTypes: [],
+        encounteredTypes: [eA._type],
+        status: StepResultStatus.SUCCESS,
+      },
+      {
+        id: 'b',
+        name: 'b',
+        declaredTypes: [],
+        partialTypes: [],
+        encounteredTypes: [eB._type],
+        dependsOn: ['a'],
+        status: StepResultStatus.FAILURE,
+      },
+      {
+        id: 'c',
+        name: 'c',
+        declaredTypes: [],
+        partialTypes: [],
+        encounteredTypes: [eC._type],
+        dependsOn: ['b'],
+        status: StepResultStatus.PARTIAL_SUCCESS_DUE_TO_DEPENDENCY_FAILURE,
+      },
+    ]);
+
+    expect(spyA).toHaveBeenCalledTimes(1);
+    expect(spyB).toHaveBeenCalledTimes(1);
+    expect(spyC).toHaveBeenCalledTimes(1);
+
+    expect(spyA).toHaveBeenCalledBefore(spyB);
+    expect(spyB).toHaveBeenCalledBefore(spyC);
+  });
+
   test('logs error after step fails', async () => {
     const error = new IntegrationError({
       code: 'ABC-123',