diff --git a/package.json b/package.json index e9fe937..32a5df0 100644 --- a/package.json +++ b/package.json @@ -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", diff --git a/src/example-testing-data/example-data.json b/src/example-testing-data/example-data.json index ddad0be..784e7f8 100644 --- a/src/example-testing-data/example-data.json +++ b/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, diff --git a/src/example-testing-data/example-deferred-result.json b/src/example-testing-data/example-deferred-result.json index 39babc4..9b8a48b 100644 --- a/src/example-testing-data/example-deferred-result.json +++ b/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 + } + } } diff --git a/src/index.ts b/src/index.ts index 11a1c73..1f62e2c 100644 --- a/src/index.ts +++ b/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 = {}, + /** + * 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, diff --git a/src/queries.ts b/src/queries.ts index e076e93..52a0e18 100644 --- a/src/queries.ts +++ b/src/queries.ts @@ -111,6 +111,8 @@ export const QUERY_V1 = gql` $includeDeleted: Boolean $deferredResponse: DeferredResponseOption $deferredFormat: DeferredResponseFormat + $cursor: String + $flags: QueryV1Flags ) { queryV1( query: $query @@ -118,6 +120,8 @@ export const QUERY_V1 = gql` includeDeleted: $includeDeleted deferredResponse: $deferredResponse deferredFormat: $deferredFormat + cursor: $cursor + flags: $flags ) { type data diff --git a/yarn.lock b/yarn.lock index b9399ef..5974d02 100644 --- a/yarn.lock +++ b/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==