Skip to content

Commit

Permalink
fix: return extra args when autoPaginate is on (#1365)
Browse files Browse the repository at this point in the history
  • Loading branch information
alvarowolfx authored May 31, 2024
1 parent 247fc15 commit 9d77cd8
Show file tree
Hide file tree
Showing 5 changed files with 52 additions and 15 deletions.
2 changes: 1 addition & 1 deletion package.json
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down
9 changes: 7 additions & 2 deletions src/bigquery.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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];
Expand Down Expand Up @@ -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;
Expand Down
4 changes: 3 additions & 1 deletion src/job.ts
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,7 @@ export type QueryResultsOptions = {
* internal properties
*/
_cachedRows?: any[];

Check warning on line 58 in src/job.ts

View workflow job for this annotation

GitHub Actions / lint

Unexpected any. Specify a different type
_cachedResponse?: bigquery.IQueryResponse;
};

/**
Expand Down Expand Up @@ -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;
}

Expand Down
28 changes: 28 additions & 0 deletions system-test/bigquery.ts
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,9 @@ import {
Routine,
Table,
InsertRowsStreamResponse,
QueryOptions,
} from '../src';
import bq from '../src/types';

const bigquery = new BigQuery();
const storage = new Storage();
Expand Down Expand Up @@ -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,
Expand Down
24 changes: 13 additions & 11 deletions test/bigquery.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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, [
{
Expand All @@ -3182,7 +3184,7 @@ describe('BigQuery', () => {
value: 3,
},
]);
assert.strictEqual(resp, FAKE_RESPONSE);
assert.strictEqual(resp, fakeResponse);
done();
});
});
Expand Down

0 comments on commit 9d77cd8

Please sign in to comment.