fix: return extra args when autoPaginate is on (#1365)

alvarowolfx · web-flow · commit 9d77cd85cf43 · 2024-05-31T15:22:17.000Z
Fixes https://togithub.com/googleapis/nodejs-bigquery/issues/1210 Fixes #1362
diff --git a/package.json b/package.json
@@ -49,7 +49,7 @@
   },
   "dependencies": {
     "@google-cloud/common": "^5.0.0",
-    "@google-cloud/paginator": "^5.0.0",
+    "@google-cloud/paginator": "^5.0.2",
     "@google-cloud/precise-date": "^4.0.0",
     "@google-cloud/promisify": "^4.0.0",
     "arrify": "^2.0.1",
diff --git a/src/bigquery.ts b/src/bigquery.ts
@@ -80,15 +80,19 @@ export type PagedRequest<P> = P & {
   maxApiCalls?: number;
 };
 
+export type QueryResultsResponse =
+  | bigquery.IGetQueryResultsResponse
+  | bigquery.IQueryResponse;
+
 export type QueryRowsResponse = PagedResponse<
   RowMetadata,
   Query,
-  bigquery.IGetQueryResultsResponse
+  QueryResultsResponse
 >;
 export type QueryRowsCallback = PagedCallback<
   RowMetadata,
   Query,
-  bigquery.IGetQueryResultsResponse
+  QueryResultsResponse
 >;
 
 export type SimpleQueryRowsResponse = [RowMetadata[], bigquery.IJob];
@@ -2185,6 +2189,7 @@ export class BigQuery extends Service {
         }
         this.trace_('[runJobsQuery] job complete');
         options._cachedRows = rows;
+        options._cachedResponse = res;
         if (res.pageToken) {
           this.trace_('[runJobsQuery] has more pages');
           options.pageToken = res.pageToken;
diff --git a/src/job.ts b/src/job.ts
@@ -56,6 +56,7 @@ export type QueryResultsOptions = {
      * internal properties
      */
     _cachedRows?: any[];
+    _cachedResponse?: bigquery.IQueryResponse;
   };
 
 /**
@@ -571,8 +572,9 @@ class Job extends Operation {
           pageToken: options.pageToken,
         });
         delete nextQuery._cachedRows;
+        delete nextQuery._cachedResponse;
       }
-      callback!(null, options._cachedRows, nextQuery);
+      callback!(null, options._cachedRows, nextQuery, options._cachedResponse);
       return;
     }
 
diff --git a/system-test/bigquery.ts b/system-test/bigquery.ts
@@ -34,7 +34,9 @@ import {
   Routine,
   Table,
   InsertRowsStreamResponse,
+  QueryOptions,
 } from '../src';
+import bq from '../src/types';
 
 const bigquery = new BigQuery();
 const storage = new Storage();
@@ -341,6 +343,32 @@ describe('BigQuery', () => {
     });
   });
 
+  it('should query with jobs.query and return all PagedResponse as positional parameters', async () => {
+    const [rows, q, response] = await bigquery.query(query);
+    const res: bq.IQueryResponse = response!;
+    assert.strictEqual(rows!.length, 100);
+    assert.notEqual(q?.job?.id, undefined);
+    assert.notEqual(res, undefined);
+    assert.strictEqual(res.kind, 'bigquery#queryResponse');
+    assert.notEqual(res.queryId, undefined);
+    assert.strictEqual(res.totalRows, '100');
+  });
+
+  it('should query without jobs.query and return all PagedResponse as positional parameters', async () => {
+    // force jobs.getQueryResult instead of fast query path
+    const jobId = generateName('job');
+    const qOpts: QueryOptions = {
+      job: bigquery.job(jobId),
+    };
+    const [rows, q, response] = await bigquery.query(query, qOpts);
+    const res: bq.IGetQueryResultsResponse = response!;
+    assert.strictEqual(rows!.length, 100);
+    assert.strictEqual(q?.job?.id, jobId);
+    assert.notEqual(res, undefined);
+    assert.strictEqual(res.kind, 'bigquery#getQueryResultsResponse');
+    assert.strictEqual(res.totalRows, '100');
+  });
+
   it('should allow querying in series', done => {
     bigquery.query(
       query,
diff --git a/test/bigquery.ts b/test/bigquery.ts
@@ -3152,24 +3152,26 @@ describe('BigQuery', () => {
       });
     });
 
-    it('should call job#getQueryResults with cached rows from jobs.query', done => {
+    it('should call job#getQueryResults with cached rows and response from jobs.query', done => {
       const fakeJob = {
         getQueryResults: (options: QueryResultsOptions, callback: Function) => {
-          callback(null, options._cachedRows, FAKE_RESPONSE);
+          callback(null, options._cachedRows, null, options._cachedResponse);
         },
       };
 
+      const fakeResponse = {
+        jobComplete: true,
+        schema: {
+          fields: [{name: 'value', type: 'INT64'}],
+        },
+        rows: [{f: [{v: 1}]}, {f: [{v: 2}]}, {f: [{v: 3}]}],
+      };
+
       bq.runJobsQuery = (query: {}, callback: Function) => {
-        callback(null, fakeJob, {
-          jobComplete: true,
-          schema: {
-            fields: [{name: 'value', type: 'INT64'}],
-          },
-          rows: [{f: [{v: 1}]}, {f: [{v: 2}]}, {f: [{v: 3}]}],
-        });
+        callback(null, fakeJob, fakeResponse);
       };
 
-      bq.query(QUERY_STRING, (err: Error, rows: {}, resp: {}) => {
+      bq.query(QUERY_STRING, (err: Error, rows: {}, query: {}, resp: {}) => {
         assert.ifError(err);
         assert.deepStrictEqual(rows, [
           {
@@ -3182,7 +3184,7 @@ describe('BigQuery', () => {
             value: 3,
           },
         ]);
-        assert.strictEqual(resp, FAKE_RESPONSE);
+        assert.strictEqual(resp, fakeResponse);
         done();
       });
     });
