Merge pull request #162 from JupiterOne/NO-TICKET-fix-sdk-to-use-proper-pagination

Replace skip limit pagination with cursor

Author: chrichts

CHANGELOG.md

@@ -8,6 +8,11 @@ and this project adheres to
 
 ## [Unreleased]
 
+## [2.1.0] - 2024-07-23
+
+- Update the `queryv1` function to use cursors instead of the legacy `LIMIT SKIP` pagination approach
+- Add support for a progress bar
+
 ## [2.0.1] - 2024-07-02
 
 - Converted `Content-Type` request header to lower case. This conforms to https://www.rfc-editor.org/rfc/rfc7540#section-8.1.2. And also addresses an issue where Content-Type and content-type headers were both being added.

package.json

@@ -1,6 +1,6 @@
 {
   "name": "@jupiterone/jupiterone-client-nodejs",
-  "version": "2.0.1",
+  "version": "2.1.0",
   "description": "A node.js client wrapper for JupiterOne public API",
   "repository": {
     "type": "git",
@@ -48,6 +48,7 @@
     "apollo-link-retry": "^2.2.13",
     "bunyan-category": "^0.4.0",
     "chalk": "^4.1.2",
+    "cli-progress": "^3.12.0",
     "commander": "^5.0.0",
     "graphql": "^14.6.0",
     "graphql-tag": "^2.10.1",

src/example-testing-data/example-data.json

@@ -3,8 +3,7 @@
     "queryV1": {
       "type": "deferred",
       "data": null,
-      "url": "https://an-s3-bucket/temp/j1dev/deferred/abc/state.json",
-      "__typename": "QueryV1Response"
+      "url": "https://an-s3-bucket/temp/j1dev/deferred/abc/state.json"
     }
   },
   "loading": false,

src/example-testing-data/example-deferred-result.json

@@ -1,5 +1,32 @@
 {
   "status": "COMPLETED",
   "url": "https://an-s3-bucket/temp/j1dev/deferred/abc/state.json",
-  "correlationId": "abc"
+  "correlationId": "abc",
+  "data": {
+    "id": "abc",
+    "entity": {
+      "_class": ["Class"],
+      "_type": ["an_example_type"],
+      "_key": "abc",
+      "displayName": "display_name",
+      "_integrationType": "github",
+      "_integrationClass": ["ITS", "SCM", "VCS", "VersionControl"],
+      "_integrationDefinitionId": "def",
+      "_integrationName": "JupiterOne",
+      "_beginOn": "2021-12-03T01:10:33.604Z",
+      "_id": "ghi",
+      "_integrationInstanceId": "nmi",
+      "_version": 11,
+      "_accountId": "an_account",
+      "_deleted": false,
+      "_source": "source",
+      "_createdOn": "2021-08-20T20:15:09.583Z"
+    },
+    "properties": {
+      "disabled": false,
+      "empty": false,
+      "fork": false,
+      "forkingAllowed": false
+    }
+  }
 }

src/index.ts

@@ -6,6 +6,7 @@ import { BatchHttpLink } from 'apollo-link-batch-http';
 import fetch, { RequestInit, Response as FetchResponse } from 'node-fetch';
 import { retry } from '@lifeomic/attempt';
 import gql from 'graphql-tag';
+import cliProgress from 'cli-progress';
 
 import Logger, { createLogger } from 'bunyan-category';
 
@@ -44,8 +45,6 @@ import {
 import { query, QueryTypes } from './util/query';
 
 const QUERY_RESULTS_TIMEOUT = 1000 * 60 * 5; // Poll s3 location for 5 minutes before timeout.
-const J1QL_SKIP_COUNT = 250;
-const J1QL_LIMIT_COUNT = 250;
 
 const JobStatus = {
   IN_PROGRESS: 'IN_PROGRESS',
@@ -75,10 +74,8 @@ export class FetchError extends Error {
     nameForLogging?: string;
   }) {
     super(
-      `JupiterOne API error. Response not OK (requestName=${
-        options.nameForLogging || '(none)'
-      }, status=${options.response.status}, url=${options.url}, method=${
-        options.method
+      `JupiterOne API error. Response not OK (requestName=${options.nameForLogging || '(none)'
+      }, status=${options.response.status}, url=${options.url}, method=${options.method
       }). Response: ${options.responseBody}`,
     );
     this.httpStatusCode = options.response.status;
@@ -327,7 +324,7 @@ export class JupiterOneClient {
     const token = this.accessToken;
     this.headers = {
       Authorization: `Bearer ${token}`,
-      'LifeOmic-Account': this.account,
+      'JupiterOne-Account': this.account,
       'content-type': 'application/json',
     };
 
@@ -351,84 +348,88 @@ export class JupiterOneClient {
   async queryV1(
     j1ql: string,
     options: QueryOptions | Record<string, unknown> = {},
+    /**
+     * include a progress bar to show progress of getting data.
+     */
+    showProgress = false,
     /** because this method queries repeatedly with its own LIMIT,
-     * this limits the looping to stop after at least {stopAfter} results are found */
+     * this limits the looping to stop after at least {stopAfter} results are found
+     * @deprecated This property is no longer supported.
+     */
     stopAfter = Number.MAX_SAFE_INTEGER,
     /** same as above, this gives more fine-grained control over the starting point of the query,
      * since this method controls the `SKIP` clause in the query
+     * @deprecated This property is no longer supported.
      */
     startPage = 0,
   ) {
+
+    let cursor: string;
     let complete = false;
-    let page = startPage;
     let results: any[] = [];
 
-    while (!complete && results.length < stopAfter) {
-      const j1qlForPage = `${j1ql} SKIP ${
-        page * J1QL_SKIP_COUNT
-      } LIMIT ${J1QL_LIMIT_COUNT}`;
+    const limitCheck = j1ql.match(/limit (\d+)/i);
+
+    let progress: any;
 
+    do {
       const res = await this.graphClient.query({
         query: QUERY_V1,
         variables: {
-          query: j1qlForPage,
+          query: j1ql,
           deferredResponse: 'FORCE',
+          flags: {
+            variableResultSize: true
+          },
+          cursor
         },
         ...options,
       });
-      page++;
       if (res.errors) {
         throw new Error(`JupiterOne returned error(s) for query: '${j1ql}'`);
       }
 
       const deferredUrl = res.data.queryV1.url;
       let status = JobStatus.IN_PROGRESS;
-      let statusFile;
+      let statusFile: any;
       const startTimeInMs = Date.now();
       do {
         if (Date.now() - startTimeInMs > QUERY_RESULTS_TIMEOUT) {
           throw new Error(
-            `Exceeded request timeout of ${
-              QUERY_RESULTS_TIMEOUT / 1000
+            `Exceeded request timeout of ${QUERY_RESULTS_TIMEOUT / 1000
             } seconds.`,
           );
         }
         this.logger.trace('Sleeping to wait for JobCompletion');
-        await sleep(200);
+        await sleep(100);
         statusFile = await networkRequest(deferredUrl);
         status = statusFile.status;
+        cursor = statusFile.cursor;
       } while (status === JobStatus.IN_PROGRESS);
 
-      let result;
-      if (status === JobStatus.COMPLETED) {
-        result = await networkRequest(statusFile.url);
-      } else {
-        // JobStatus.FAILED
-        throw new Error(
-          statusFile.error || 'Job failed without an error message.',
-        );
+      if (status === JobStatus.FAILED) {
+        throw new Error(`JupiterOne returned error(s) for query: '${statusFile.error}'`);
       }
 
-      const { data } = result;
+      const result = statusFile.data;
 
-      // data will assume tree shape if you specify 'return tree' in J1QL
-      const isTree = data.vertices && data.edges;
+      if (showProgress && !limitCheck) {
+        if (results.length === 0) {
+          progress = new cliProgress.SingleBar({}, cliProgress.Presets.shades_classic);
+          progress.start(Number(statusFile.totalCount), 0);
+        }
+        progress.update(results.length);
+      }
+
+      if (result) {
+        results = results.concat(result)
+      }
 
-      if (isTree) {
+      if (status === JobStatus.COMPLETED && (cursor == null || limitCheck)) {
         complete = true;
-        results = data;
-      } else {
-        // data is array-shaped, possibly paginated
-        if (data.length < J1QL_SKIP_COUNT) {
-          complete = true;
-        }
-        results = results.concat(data);
       }
-      this.logger.debug(
-        { resultsCount: results.length, pageCount: data.length },
-        'Query received page of results',
-      );
-    }
+
+    } while (complete === false);
     return results;
   }
 
@@ -887,7 +888,7 @@ export class JupiterOneClient {
     const headers = this.headers;
     const response = await makeFetchRequest(
       this.apiUrl +
-        `/persister/synchronization/jobs/${options.syncJobId}/upload`,
+      `/persister/synchronization/jobs/${options.syncJobId}/upload`,
       {
         method: 'POST',
         headers,

src/queries.ts

@@ -111,13 +111,17 @@ export const QUERY_V1 = gql`
     $includeDeleted: Boolean
     $deferredResponse: DeferredResponseOption
     $deferredFormat: DeferredResponseFormat
+    $cursor: String
+    $flags: QueryV1Flags
   ) {
     queryV1(
       query: $query
       variables: $variables
       includeDeleted: $includeDeleted
       deferredResponse: $deferredResponse
       deferredFormat: $deferredFormat
+      cursor: $cursor
+      flags: $flags
     ) {
       type
       data

yarn.lock

@@ -1344,6 +1344,13 @@ cli-cursor@^3.1.0:
   dependencies:
     restore-cursor "^3.1.0"
 
+cli-progress@^3.12.0:
+  version "3.12.0"
+  resolved "https://registry.npmjs.org/cli-progress/-/cli-progress-3.12.0.tgz#807ee14b66bcc086258e444ad0f19e7d42577942"
+  integrity sha512-tRkV3HJ1ASwm19THiiLIXLO7Im7wlTuKnvkYaTkyoAPefqjNg7W7DHKUlGRxy9vxDvbyCYQkQozvptuMkGCg8A==
+  dependencies:
+    string-width "^4.2.3"
+
 cli-spinners@^2.5.0:
   version "2.6.1"
   resolved "https://registry.npmjs.org/cli-spinners/-/cli-spinners-2.6.1.tgz"
@@ -3589,7 +3596,7 @@ string-length@^4.0.1:
     char-regex "^1.0.2"
     strip-ansi "^6.0.0"
 
-string-width@^4.1.0, string-width@^4.2.0:
+string-width@^4.1.0, string-width@^4.2.0, string-width@^4.2.3:
   version "4.2.3"
   resolved "https://registry.npmjs.org/string-width/-/string-width-4.2.3.tgz"
   integrity sha512-wKyQRQpjJ0sIp62ErSZdGsjMJWsap5oRNihHhu6G7JVO/9jIB6UyevL+tXuOqrng8j/cxKTWyWUwvSTriiZz/g==