From 3452b8100b66c430167b494bb47fb71cef90539b Mon Sep 17 00:00:00 2001 From: ohmayr Date: Tue, 20 Oct 2026 19:15:41 +0000 Subject: [PATCH 1/3] feat(bigquery): wait for unfinished jobs before streaming Arrow results Implements unfinished job polling (jobComplete === false) in createQueryArrowStream: - Adds ArrowQuery (Query plus Arrow-only options, beta like createQueryArrowStream) with timeoutMs, and forwards it into buildQueryRequest_ so jobs.query receives the caller's server-side timeout. - When jobs.query returns jobComplete === false: - Fails immediately with a timeout error if a non-zero timeoutMs was specified, matching BigQuery#query. - Otherwise polls job.getQueryResults with maxResults: 0 (so BigQuery blocks until completion without serializing any STRUCT_ENCODING rows) and then streams all Arrow batches from streams/_default at offset = 0. - streams/_default sends nothing for a result without rows, so the 0-row RecordBatch reads its schema from the job's destination table via CreateReadSession (none for statements without a result table). - Adds a system test that waits on a job with a picosecond TIMESTAMP(12) column, run when BIGQUERY_PICOSECOND_SUPPORT is set, and adds the Arrow system tests to the high-precision re-run in high_precision_timestamp.ts so CI runs it. --- handwritten/bigquery/src/arrow.ts | 201 +++++++- handwritten/bigquery/src/bigquery.ts | 17 +- handwritten/bigquery/src/index.ts | 1 + handwritten/bigquery/system-test/arrow.ts | 140 ++++++ .../system-test/high_precision_timestamp.ts | 1 + handwritten/bigquery/test/arrow.ts | 461 +++++++++++++++++- 6 files changed, 782 insertions(+), 39 deletions(-) diff --git a/handwritten/bigquery/src/arrow.ts b/handwritten/bigquery/src/arrow.ts index dd99d2853542..9d849b4b9f22 100644 --- a/handwritten/bigquery/src/arrow.ts +++ b/handwritten/bigquery/src/arrow.ts @@ -17,6 +17,7 @@ import {ResourceStream} from '@google-cloud/paginator'; import type {RecordBatch, Schema} from 'apache-arrow'; import type {BigQuery, JobsQueryResponse, Query} from './bigquery'; +import type {Job} from './job'; import type bigquery from './types'; type ArrowModule = typeof import('apache-arrow'); @@ -36,6 +37,24 @@ interface ReadRowsStream extends NodeJS.EventEmitter { resume(): void; } +export interface IReadSession { + arrowSchema?: { + serializedSchema?: Uint8Array | string | null; + } | null; + streams?: Array<{ + name?: string | null; + }> | null; +} + +export interface CreateReadSessionRequest { + parent: string; + readSession: { + table: string; + dataFormat: 'ARROW'; + }; + maxStreamCount: number; +} + /** * The part of `@google-cloud/bigquery-storage`'s `BigQueryReadClient` that * {@link BigQuery#createQueryArrowStream} uses. Exported for tests. @@ -43,6 +62,9 @@ interface ReadRowsStream extends NodeJS.EventEmitter { * @private */ export interface BigQueryReadClient { + createReadSession( + request: CreateReadSessionRequest, + ): Promise<[IReadSession, ...unknown[]]>; readRows(request: {readStream: string; offset?: number}): ReadRowsStream; } @@ -74,6 +96,21 @@ export type ArrowRecordBatchStream = ResourceStream & { jobReference?: bigquery.IJobReference; }; +/** + * [Beta] Options for {@link BigQuery#createQueryArrowStream}: the + * {@link Query} options plus options that only apply to Arrow queries. + * + * @beta + */ +export type ArrowQuery = Query & { + /** + * How long to wait for `jobs.query` to complete, in milliseconds. If the + * query doesn't complete in time, the stream fails instead of waiting for + * the job. + */ + timeoutMs?: number; +}; + /** * State shared by the helpers that stream rows for one * {@link BigQuery#createQueryArrowStream} call. @@ -267,10 +304,28 @@ export function decodeArrowPayload( : [createEmptyRecordBatch(arrow, reader.schema)]; } +/** + * Pushes `batches` to `stream`, stopping if the consumer destroys it. + * + * @private + */ +function pushBatches( + stream: ArrowRecordBatchStream, + batches: RecordBatch[], +): void { + for (const batch of batches) { + if (stream.destroyed) { + return; + } + stream.push(batch); + } +} + /** * Streams Arrow `RecordBatch`es from a job's sessionless default stream * (`projects/{projectId}/locations/{location}/jobs/{jobId}/streams/_default`) - * starting at `offset`. + * starting at `offset`. Pass the finished `job` when reading from offset 0 so + * a result without rows keeps its schema (see {@link readDestinationSchema}). * * @private */ @@ -279,8 +334,9 @@ async function streamDefaultJobStream( jobRef: bigquery.IJobReference | undefined, offset: number, initialSchema?: Uint8Array | string, + job?: Job, ): Promise { - const {bq} = ctx; + const {bq, arrow, stream} = ctx; const projectId = resolveProjectId(bq, jobRef?.projectId); const jobId = jobRef?.jobId; const location = jobRef?.location || bq.location; @@ -292,16 +348,66 @@ async function streamDefaultJobStream( } const readStream = `projects/${projectId}/locations/${location}/jobs/${jobId}/streams/_default`; - await streamReadRows(ctx, {readStream, offset, initialSchema}); + const {emittedBatch, serializedSchema} = await streamReadRows(ctx, { + readStream, + offset, + initialSchema, + }); + + // When reading from offset = 0 (e.g., after waiting for an unfinished job) + // and the stream finishes with 0 rows, still emit a 0-row RecordBatch. + if (offset !== 0 || emittedBatch || stream.destroyed) { + return; + } + const schema = + serializedSchema ?? (job && (await readDestinationSchema(ctx, job))); + pushBatches(stream, decodeArrowPayload(arrow, schema)); +} + +/** + * Reads the Arrow schema of a finished query job's results from its + * destination table. The default job stream sends no `ReadRowsResponse`, and + * so no schema, for a result without rows, but a read session on the + * destination table still returns it. Statements without a result table + * (e.g., DML) have no destination table, so this resolves to `undefined`. + * + * @private + */ +async function readDestinationSchema( + ctx: ArrowStreamContext, + job: Job, +): Promise { + const {bq, storage} = ctx; + // A script's destination table is only set once the job finishes, so read + // the job's current metadata instead of the response that created it. + const [metadata] = await job.getMetadata(); + const destRef = metadata?.configuration?.query?.destinationTable; + if (!destRef) { + return undefined; + } + const [session] = await getStorageReadClient(bq, storage).createReadSession({ + parent: `projects/${resolveProjectId( + bq, + metadata?.jobReference?.projectId, + destRef.projectId, + )}`, + readSession: { + table: `projects/${destRef.projectId}/datasets/${destRef.datasetId}/tables/${destRef.tableId}`, + dataFormat: 'ARROW', + }, + maxStreamCount: 1, + }); + return session?.arrowSchema?.serializedSchema || undefined; } /** * Streams Arrow `RecordBatch`es from a Storage Read API stream starting at * `offset`, pausing `ReadRows` while the consumer is behind. * - * Resolves when `ReadRows` ends or the consumer destroys the stream, and - * rejects on a `ReadRows` or decode error. In every case, the listeners added - * here are removed, and `ReadRows` is cancelled unless it already finished. + * Resolves when `ReadRows` ends or the consumer destroys the stream, with + * whether any batch was emitted and the last Arrow schema seen, and rejects + * on a `ReadRows` or decode error. In every case, the listeners added here + * are removed, and `ReadRows` is cancelled unless it already finished. * * @private */ @@ -312,7 +418,7 @@ function streamReadRows( offset: number; initialSchema?: Uint8Array | string; }, -): Promise { +): Promise<{emittedBatch: boolean; serializedSchema?: Uint8Array}> { const {bq, arrow, storage, stream} = ctx; const {readStream, offset, initialSchema} = options; @@ -325,6 +431,7 @@ function streamReadRows( let serializedSchema: Uint8Array | undefined = initialSchema ? toUint8Array(initialSchema) : undefined; + let emittedBatch = false; let settled = false; // Detaches every listener added here and settles the promise. Cancels @@ -348,7 +455,7 @@ function streamReadRows( if (opts.err) { reject(opts.err); } else { - resolve(); + resolve({emittedBatch, serializedSchema}); } }; @@ -379,6 +486,7 @@ function streamReadRows( settle({cancel: true}); return; } + emittedBatch = true; if (!stream.push(batch)) { grpcStream.pause(); } @@ -403,6 +511,31 @@ function streamReadRows( }); } +/** + * Waits for an unfinished BigQuery job to complete using `job.getQueryResults` + * without fetching any `STRUCT_ENCODING` rows (`maxResults: 0`). Returns early + * if the consumer destroys `stream`. + * + * @private + */ +async function waitForJobCompletion( + job: Job, + stream: ArrowRecordBatchStream, +): Promise { + while (!stream.destroyed) { + // Set maxResults: 0 so jobs.getQueryResults blocks up to the server's + // default wait interval (about 10s) and surfaces job errors without + // returning JSON rows, so polling again right away doesn't spin. + const [, , res] = await job.getQueryResults({ + autoPaginate: false, + maxResults: 0, + }); + if (res?.jobComplete !== false) { + return; + } + } +} + /** * Executes a query and returns a stream of `apache-arrow` `RecordBatch` objects. * @@ -410,7 +543,7 @@ function streamReadRows( */ export function createQueryArrowStream_( bq: BigQuery, - query?: Query | string, + query?: ArrowQuery | string, ): ArrowRecordBatchStream { const stream = new ResourceStream( {}, @@ -430,7 +563,8 @@ export function createQueryArrowStream_( throw new Error('A SQL query string is required.'); } - const queryObj: Query = typeof query === 'string' ? {query} : {...query}; + const queryObj: ArrowQuery = + typeof query === 'string' ? {query} : {...query}; if (queryObj.writeDisposition === 'WRITE_APPEND') { throw new Error( @@ -450,7 +584,9 @@ export function createQueryArrowStream_( readCtx = ctx; // Throws for an invalid request, for example invalid query parameters. - const queryReq = bq.buildQueryRequest_(queryObj, {}); + const queryReq = bq.buildQueryRequest_(queryObj, { + timeoutMs: queryObj.timeoutMs, + }); if (!queryReq) { // TODO: Route queries that cannot use jobs.query through jobs.insert (createQueryJob). throw new Error('Query configuration cannot be executed via jobs.query.'); @@ -473,8 +609,30 @@ export function createQueryArrowStream_( } if (res?.jobComplete === false) { - // TODO: Wait for unfinished jobs before streaming Arrow results. - throw new Error('Query did not complete within the jobs.query timeout.'); + // jobs.query already waited (about 10s by default) and the query is + // still running, so this depends on how long the query runs, not on + // how many rows it returns. + // Match BigQuery#query: if the caller provided a non-zero timeoutMs, + // fail immediately instead of polling jobs.getQueryResults. + if (queryReq.timeoutMs) { + throw new Error( + `The query did not complete before ${queryReq.timeoutMs}ms`, + ); + } + if (!res.jobReference || !job) { + throw new Error( + 'Cannot wait for an incomplete query without a jobReference.', + ); + } + await waitForJobCompletion(job, stream); + if (stream.destroyed) { + return; + } + // An unfinished jobs.query returns no rows or schema, so read every row + // from offset 0 and pass the job to look up the schema of an empty + // result. + await streamDefaultJobStream(ctx, res.jobReference, 0, undefined, job); + return; } // pageRowCount and totalRows are returned as uint64 strings in REST JSON. @@ -488,17 +646,14 @@ export function createQueryArrowStream_( // Emit the inline first page (or 0-row schema-preserving batch when totalRows is 0) // before resuming remaining rows from streams/_default at offset = pageRowCount. if (!hasRemainingRows || inlineRowsCount > 0n) { - const batches = decodeArrowPayload( - arrow, - res?.arrowSchema?.serializedSchema, - res?.arrowRecordBatch?.serializedRecordBatch, + pushBatches( + stream, + decodeArrowPayload( + arrow, + res?.arrowSchema?.serializedSchema, + res?.arrowRecordBatch?.serializedRecordBatch, + ), ); - for (const batch of batches) { - if (stream.destroyed) { - return; - } - stream.push(batch); - } } if (!hasRemainingRows || stream.destroyed) { diff --git a/handwritten/bigquery/src/bigquery.ts b/handwritten/bigquery/src/bigquery.ts index 76cb1d0645ee..2fc2e5bed7b1 100644 --- a/handwritten/bigquery/src/bigquery.ts +++ b/handwritten/bigquery/src/bigquery.ts @@ -56,10 +56,14 @@ import { } from '@google-cloud/common/build/src/util'; import bigquery from './types'; import {logger, setLogFunction} from './logger'; -import {ArrowRecordBatchStream, createQueryArrowStream_} from './arrow'; +import { + ArrowQuery, + ArrowRecordBatchStream, + createQueryArrowStream_, +} from './arrow'; // Third-Party Re-exports -export {common, ArrowRecordBatchStream}; +export {common, ArrowQuery, ArrowRecordBatchStream}; export interface RequestCallback { (err: Error | null, response?: T | null): void; @@ -393,8 +397,11 @@ export class BigQuery extends Service { * batches, so iterate over batches to keep memory bounded; collecting every * batch (for example with `toArray()`) holds the whole result in memory. * - * @param {string|object} [options] A string SQL query or configuration + * @param {string|ArrowQuery} [options] A string SQL query or configuration * object. See {@link BigQuery#query} for a complete list of options. + * @param {number} [options.timeoutMs] How long to wait for `jobs.query` to + * complete, in milliseconds. If the query doesn't complete in time, the + * stream fails instead of waiting for the job. * @returns {ArrowRecordBatchStream} A stream of `RecordBatch` objects. The * stream emits a `job` event with the {@link Job} when BigQuery creates * a job, and sets `stream.queryId` and `stream.jobReference` (when @@ -424,7 +431,9 @@ export class BigQuery extends Service { * * @beta */ - createQueryArrowStream(options?: Query | string): ArrowRecordBatchStream { + createQueryArrowStream( + options?: ArrowQuery | string, + ): ArrowRecordBatchStream { return createQueryArrowStream_(this, options); } diff --git a/handwritten/bigquery/src/index.ts b/handwritten/bigquery/src/index.ts index 314ec9f5a97a..bd47e9610036 100644 --- a/handwritten/bigquery/src/index.ts +++ b/handwritten/bigquery/src/index.ts @@ -15,6 +15,7 @@ */ export { + ArrowQuery, ArrowRecordBatchStream, BigQuery, BigQueryDate, diff --git a/handwritten/bigquery/system-test/arrow.ts b/handwritten/bigquery/system-test/arrow.ts index 2614a19cbb92..980b1d86d9ef 100644 --- a/handwritten/bigquery/system-test/arrow.ts +++ b/handwritten/bigquery/system-test/arrow.ts @@ -398,4 +398,144 @@ describe('Arrow Query Results', () => { }, ); }); + + it('should wait for an unfinished jobs.query via jobs.getQueryResults (maxResults = 0) and stream all rows from streams/_default at offset = 0', async () => { + const origRequest = bigquery.request.bind(bigquery); + let getQueryResultsCalls = 0; + + // Force POST /queries to return jobComplete: false by setting a 1ms server wait on the wire + // while leaving the caller's Query object without timeoutMs so createQueryArrowStream waits. + bigquery.request = (reqOpts, callback) => { + if (reqOpts.uri === '/queries' && reqOpts.method === 'POST') { + reqOpts.json = {...reqOpts.json, timeoutMs: 1}; + } else if ( + typeof reqOpts.uri === 'string' && + reqOpts.uri.startsWith('/queries/') + ) { + getQueryResultsCalls++; + } + return origRequest(reqOpts, callback); + }; + + try { + const sql = + 'DECLARE n INT64 DEFAULT 50; SELECT x FROM UNNEST(GENERATE_ARRAY(1, n)) AS x ORDER BY x;'; + const batches: arrow.RecordBatch[] = await bigquery + .createQueryArrowStream({query: sql, useQueryCache: false}) + .toArray(); + + // Verify jobs.getQueryResults was polled to wait for completion and all 50 rows streamed from offset = 0. + assert.ok( + getQueryResultsCalls >= 1, + `Expected >= 1 jobs.getQueryResults call, got ${getQueryResultsCalls}`, + ); + const table = new arrow.Table(batches); + assert.strictEqual(table.numRows, 50); + assert.strictEqual(table.getChild('x')!.get(0), 1n); + assert.strictEqual(table.getChild('x')!.get(49), 50n); + } finally { + bigquery.request = origRequest; + } + }); + + it('should keep the schema of an empty result after waiting for an unfinished jobs.query', async () => { + // After the wait, rows come from streams/_default at offset = 0, which + // sends nothing for a result without rows, so the schema has to come from + // the job's destination table. + const origRequest = bigquery.request.bind(bigquery); + + // Force POST /queries to return jobComplete: false so the stream waits + // for the job before reading rows. + bigquery.request = (reqOpts, callback) => { + if (reqOpts.uri === '/queries' && reqOpts.method === 'POST') { + reqOpts.json = {...reqOpts.json, timeoutMs: 1}; + } + return origRequest(reqOpts, callback); + }; + + try { + const sql = `DECLARE n INT64 DEFAULT 0; + SELECT x, CAST(x AS STRING) AS s FROM UNNEST(GENERATE_ARRAY(1, n)) AS x;`; + const batches: arrow.RecordBatch[] = await bigquery + .createQueryArrowStream({query: sql, useQueryCache: false}) + .toArray(); + + // Verify one 0-row RecordBatch with the result's column names and types. + assert.strictEqual(batches.length, 1); + assert.strictEqual(batches[0].numRows, 0); + assert.deepStrictEqual( + batches[0].schema.fields.map(f => f.name), + ['x', 's'], + ); + assert.ok(arrow.DataType.isInt(batches[0].schema.fields[0].type)); + assert.ok(arrow.DataType.isUtf8(batches[0].schema.fields[1].type)); + } finally { + bigquery.request = origRequest; + } + }); + + it('should wait for an unfinished job with a picosecond TIMESTAMP(12) column using default jobs.getQueryResults format options', async function () { + // Guards the job wait against format option regressions: the wait calls + // jobs.getQueryResults with maxResults = 0 and the library's default + // formatOptions (useInt64Timestamp), which must still succeed when the + // result has a picosecond TIMESTAMP(12) column. Picosecond timestamps are + // only enabled in some projects, so this runs only when + // BIGQUERY_PICOSECOND_SUPPORT is set. + if (process.env.BIGQUERY_PICOSECOND_SUPPORT !== 'true') { + this.skip(); + } + const origRequest = bigquery.request.bind(bigquery); + let getQueryResultsCalls = 0; + + // Force POST /queries to return jobComplete: false so the stream waits + // for the job with jobs.getQueryResults before reading rows. + bigquery.request = (reqOpts, callback) => { + if (reqOpts.uri === '/queries' && reqOpts.method === 'POST') { + reqOpts.json = {...reqOpts.json, timeoutMs: 1}; + } else if ( + typeof reqOpts.uri === 'string' && + reqOpts.uri.startsWith('/queries/') + ) { + getQueryResultsCalls++; + } + return origRequest(reqOpts, callback); + }; + + try { + const sql = `DECLARE n INT64 DEFAULT 3; + SELECT x, CAST('2026-01-15 12:34:56.123456789123+00' AS TIMESTAMP(12)) AS ts + FROM UNNEST(GENERATE_ARRAY(1, n)) AS x ORDER BY x;`; + const batches: arrow.RecordBatch[] = await bigquery + .createQueryArrowStream({query: sql, useQueryCache: false}) + .toArray(); + + // Verify the wait ran and every row streamed with an Arrow timestamp. + assert.ok( + getQueryResultsCalls >= 1, + `Expected >= 1 jobs.getQueryResults call, got ${getQueryResultsCalls}`, + ); + const table = new arrow.Table(batches); + assert.strictEqual(table.numRows, 3); + assert.ok(arrow.DataType.isTimestamp(table.schema.fields[1].type)); + } finally { + bigquery.request = origRequest; + } + }); + + it('should reject when a query does not complete within caller-specified timeoutMs', async () => { + const sql = + 'DECLARE n INT64 DEFAULT 50; SELECT x FROM UNNEST(GENERATE_ARRAY(1, n)) AS x ORDER BY x;'; + + // Verify specifying timeoutMs: 1 rejects when jobs.query returns jobComplete = false. + await assert.rejects( + bigquery + .createQueryArrowStream({ + query: sql, + timeoutMs: 1, + useQueryCache: false, + }) + .toArray(), + /The query did not complete before 1ms/, + ); + }); }); diff --git a/handwritten/bigquery/system-test/high_precision_timestamp.ts b/handwritten/bigquery/system-test/high_precision_timestamp.ts index 17c323209807..658be8d41aea 100644 --- a/handwritten/bigquery/system-test/high_precision_timestamp.ts +++ b/handwritten/bigquery/system-test/high_precision_timestamp.ts @@ -66,6 +66,7 @@ describe('BigQuery System Tests (High Precision)', () => { // Programmatically require the tests. Mocha will discover and register them inside this 'describe' block. require('./bigquery'); require('./timestamp_output_format'); + require('./arrow'); }); }); diff --git a/handwritten/bigquery/test/arrow.ts b/handwritten/bigquery/test/arrow.ts index 49491697b529..365716af8812 100644 --- a/handwritten/bigquery/test/arrow.ts +++ b/handwritten/bigquery/test/arrow.ts @@ -17,10 +17,17 @@ import {EventEmitter} from 'events'; import {describe, it, before, after, beforeEach} from 'mocha'; import * as Module from 'module'; import * as arrow from 'apache-arrow'; -import {BigQuery, Job, Query} from '../src'; +import { + BigQuery, + Job, + Query, + QueryResultsOptions, + QueryRowsCallback, +} from '../src'; import { BigQueryReadClient, checkArrowDependencies, + CreateReadSessionRequest, getStorageReadClient, StorageModule, } from '../src/arrow'; @@ -107,11 +114,18 @@ type FakeReadRowsStream = EventEmitter & */ function stubStorageReadClient( bq: BigQuery, - client: {readRows: (req: ReadRowsRequest) => FakeReadRowsStream}, + client: { + readRows: (req: ReadRowsRequest) => FakeReadRowsStream; + createReadSession?: BigQueryReadClient['createReadSession']; + }, ): BigQueryReadClient { const noop = () => {}; + const notStubbed = (): never => { + throw new Error('createReadSession is not stubbed'); + }; return getStorageReadClient(bq, { BigQueryReadClient: class { + createReadSession = client.createReadSession ?? notStubbed; readRows(req: ReadRowsRequest) { const grpcStream = client.readRows(req); return Object.assign(grpcStream, { @@ -137,6 +151,66 @@ function assertReadRowsListenersRemoved(grpcStream: EventEmitter): void { assert.doesNotThrow(() => grpcStream.emit('error', cancelled)); } +/** + * Stubs `job.getQueryResults` with a callback-style `impl`, returning a + * promise because the code calls the promisified method. + */ +function stubJobGetQueryResults( + job: Job, + impl: (opts: QueryResultsOptions, cb: QueryRowsCallback) => void, +): void { + job.getQueryResults = ((opts: QueryResultsOptions) => + new Promise((resolve, reject) => { + impl(opts, (err, rows, nextQuery, res) => + err ? reject(err) : resolve([rows, nextQuery, res]), + ); + })) as unknown as Job['getQueryResults']; +} + +/** + * Stubs `job.getMetadata` (`jobs.get`) to resolve with the given job resource + * or reject with the given error. + */ +function stubJobGetMetadata( + job: Job, + outcome: {err?: Error; metadata?: bigquery.IJob}, +): void { + job.getMetadata = (async () => { + if (outcome.err) { + throw outcome.err; + } + return [outcome.metadata]; + }) as unknown as Job['getMetadata']; +} + +/** + * Returns a fake ReadRows stream that emits `responses`, then ends, on the + * next tick. + */ +function fakeReadRowsStream(responses: object[] = []): EventEmitter { + const grpcStream = new EventEmitter(); + process.nextTick(() => { + for (const resp of responses) { + grpcStream.emit('data', resp); + } + grpcStream.emit('end'); + }); + return grpcStream; +} + +/** A ReadRowsResponse carrying `SAMPLE_PAYLOAD` as bytes. */ +const SAMPLE_READ_ROWS_RESPONSE = { + arrowSchema: { + serializedSchema: Buffer.from(SAMPLE_PAYLOAD.serializedSchema, 'base64'), + }, + arrowRecordBatch: { + serializedRecordBatch: Buffer.from( + SAMPLE_PAYLOAD.serializedRecordBatch, + 'base64', + ), + }, +}; + describe('Arrow Support', () => { let originalRequire: typeof Module.prototype.require; @@ -452,16 +526,6 @@ describe('Arrow Support', () => { /Parameter types must be provided for null values/, ); }); - - it('should fail when jobs.query returns jobComplete = false', async () => { - stubJobsQuery(bq, {res: {jobComplete: false}}); - - // TODO: Replace with job wait (jobComplete === false) tests once unfinished job polling is added. - await assert.rejects( - bq.createQueryArrowStream('SELECT 1').toArray(), - /Query did not complete within the jobs\.query timeout\./, - ); - }); }); describe('BigQuery#createQueryArrowStream (default job stream)', () => { @@ -484,6 +548,9 @@ describe('Arrow Support', () => { readRows(): never { throw new Error('not called'); } + createReadSession(): never { + throw new Error('not called'); + } }, }; @@ -733,6 +800,9 @@ describe('Arrow Support', () => { readRows(): never { throw new Error('should not instantiate a new client'); } + createReadSession(): never { + throw new Error('should not instantiate a new client'); + } }, }), sharedClient, @@ -897,4 +967,371 @@ describe('Arrow Support', () => { assert.strictEqual(new arrow.Table(batches).numRows, 2); }); }); + + describe('BigQuery#createQueryArrowStream (job wait)', () => { + let bq: BigQuery; + + beforeEach(() => { + bq = new BigQuery({projectId: 'test-project'}); + }); + + it('should poll job.getQueryResults with maxResults = 0 and call ReadRows at offset = 0 only after the job completes', async () => { + const fakeJob = bq.job('job-wait-1', {location: 'US'}); + const getQueryResultsCalls: QueryResultsOptions[] = []; + const capturedReadRowsReqs: Array<{ + readStream?: string; + offset?: number; + }> = []; + + // First poll returns jobComplete: false; second poll returns jobComplete: true. + stubJobGetQueryResults(fakeJob, (opts, cb) => { + getQueryResultsCalls.push(opts); + // Verify ReadRows has not been called while waiting for the job to complete. + assert.strictEqual(capturedReadRowsReqs.length, 0); + const isDone = getQueryResultsCalls.length >= 2; + process.nextTick(() => { + cb(null, [], null, { + jobComplete: isDone, + totalRows: isDone ? '2' : undefined, + }); + }); + }); + + stubJobsQuery(bq, { + job: fakeJob, + res: { + jobComplete: false, + jobReference: { + projectId: 'test-project', + jobId: 'job-wait-1', + location: 'US', + }, + }, + }); + + stubStorageReadClient(bq, { + readRows: (req: {readStream?: string; offset?: number}) => { + capturedReadRowsReqs.push(req); + return fakeReadRowsStream([SAMPLE_READ_ROWS_RESPONSE]); + }, + }); + + const batches: arrow.RecordBatch[] = await bq + .createQueryArrowStream('SELECT id, name FROM tbl') + .toArray(); + + // Verify job.getQueryResults was polled twice with maxResults: 0 and autoPaginate: false. + assert.strictEqual(getQueryResultsCalls.length, 2); + for (const opts of getQueryResultsCalls) { + assert.deepStrictEqual(opts, { + autoPaginate: false, + maxResults: 0, + }); + } + + // Verify ReadRows was called at offset = 0 only after the job completed. + assert.deepStrictEqual(capturedReadRowsReqs, [ + { + readStream: + 'projects/test-project/locations/US/jobs/job-wait-1/streams/_default', + offset: 0, + }, + ]); + assert.strictEqual(batches.length, 1); + assert.strictEqual(batches[0].numRows, 2); + }); + + it('should fail with a timeout error when jobs.query returns jobComplete = false and timeoutMs was specified', async () => { + const fakeJob = bq.job('job-timeout-1', {location: 'US'}); + let getQueryResultsCalled = false; + stubJobGetQueryResults(fakeJob, () => { + getQueryResultsCalled = true; + }); + + const getCapturedReq = stubJobsQuery(bq, { + job: fakeJob, + res: { + jobComplete: false, + jobReference: { + projectId: 'test-project', + jobId: 'job-timeout-1', + location: 'US', + }, + }, + }); + + // Verify caller-provided timeoutMs fails immediately when jobComplete is false. + await assert.rejects( + bq + .createQueryArrowStream({query: 'SELECT 1', timeoutMs: 250}) + .toArray(), + /The query did not complete before 250ms/, + ); + assert.strictEqual(getCapturedReq()?.timeoutMs, 250); + assert.strictEqual(getQueryResultsCalled, false); + }); + + it('should fail when jobs.query returns jobComplete = false without a jobReference', async () => { + stubJobsQuery(bq, {res: {jobComplete: false, queryId: 'stateless-1'}}); + + // Verify an actionable error instead of waiting on a job that doesn't exist. + await assert.rejects( + bq.createQueryArrowStream('SELECT 1').toArray(), + /Cannot wait for an incomplete query without a jobReference\./, + ); + }); + + it('should stop polling and never call ReadRows when the stream is destroyed during the wait', async () => { + const fakeJob = bq.job('job-wait-destroy', {location: 'US'}); + const stream = bq.createQueryArrowStream('SELECT id, name FROM tbl'); + let polls = 0; + stubJobGetQueryResults(fakeJob, (_opts, cb) => { + polls++; + // The consumer gives up while this poll is in flight. + stream.destroy(); + process.nextTick(() => cb(null, [], null, {jobComplete: false})); + }); + stubJobsQuery(bq, { + job: fakeJob, + res: { + jobComplete: false, + jobReference: { + projectId: 'test-project', + jobId: 'job-wait-destroy', + location: 'US', + }, + }, + }); + let readRowsCalls = 0; + stubStorageReadClient(bq, { + readRows: () => { + readRowsCalls++; + return new EventEmitter(); + }, + }); + + stream.resume(); + // The poll's callback runs on the next tick; wait past it. + await new Promise(resolve => setImmediate(resolve)); + + // Verify the in-flight poll didn't start another one or call ReadRows. + assert.strictEqual(polls, 1); + assert.strictEqual(readRowsCalls, 0); + }); + + it('should wait for the job instead of failing when timeoutMs is 0, matching BigQuery#query', async () => { + const fakeJob = bq.job('job-timeout-0', {location: 'US'}); + const getQueryResultsCalls: QueryResultsOptions[] = []; + stubJobGetQueryResults(fakeJob, (opts, cb) => { + getQueryResultsCalls.push(opts); + process.nextTick(() => cb(null, [], null, {jobComplete: true})); + }); + + stubJobsQuery(bq, { + job: fakeJob, + res: { + jobComplete: false, + jobReference: { + projectId: 'test-project', + jobId: 'job-timeout-0', + location: 'US', + }, + }, + }); + + stubStorageReadClient(bq, { + readRows: () => fakeReadRowsStream([SAMPLE_READ_ROWS_RESPONSE]), + }); + + // BigQuery#query treats a falsy timeoutMs as "no override", so 0 must + // not reject with "did not complete before 0ms"; poll and stream instead. + const batches: arrow.RecordBatch[] = await bq + .createQueryArrowStream({query: 'SELECT 1', timeoutMs: 0}) + .toArray(); + assert.deepStrictEqual(getQueryResultsCalls, [ + {autoPaginate: false, maxResults: 0}, + ]); + assert.strictEqual(batches.length, 1); + assert.strictEqual(batches[0].numRows, 2); + }); + + it('should forward job.getQueryResults errors to the stream without calling ReadRows', async () => { + const fakeJob = bq.job('job-err-1', {location: 'US'}); + stubJobGetQueryResults(fakeJob, (_opts, cb) => { + process.nextTick(() => { + cb(new Error('Resources exceeded during query execution')); + }); + }); + + let readRowsCalled = false; + stubStorageReadClient(bq, { + readRows: () => { + readRowsCalled = true; + return new EventEmitter(); + }, + }); + + stubJobsQuery(bq, { + job: fakeJob, + res: { + jobComplete: false, + jobReference: { + projectId: 'test-project', + jobId: 'job-err-1', + location: 'US', + }, + }, + }); + + // Verify errors from job.getQueryResults reject stream consumption and skip ReadRows. + await assert.rejects( + bq.createQueryArrowStream('SELECT 1').toArray(), + /Resources exceeded during query execution/, + ); + assert.strictEqual(readRowsCalled, false); + }); + + it('should emit a 0-row, 0-column RecordBatch when ReadRows at offset = 0 ends without rows and the job has no destination table', async () => { + const fakeJob = bq.job('job-empty-1', {location: 'US'}); + stubJobGetQueryResults(fakeJob, (_opts, cb) => { + process.nextTick(() => { + cb(null, [], null, {jobComplete: true, totalRows: '0'}); + }); + }); + // Statements without a result table, like DECLARE, have no destination. + stubJobGetMetadata(fakeJob, {metadata: {configuration: {query: {}}}}); + + stubJobsQuery(bq, { + job: fakeJob, + res: { + jobComplete: false, + jobReference: { + projectId: 'test-project', + jobId: 'job-empty-1', + location: 'US', + }, + }, + }); + + // createReadSession isn't stubbed, so calling it fails the test. + stubStorageReadClient(bq, { + readRows: () => fakeReadRowsStream(), + }); + + const batches: arrow.RecordBatch[] = await bq + .createQueryArrowStream('DECLARE x INT64 DEFAULT 1') + .toArray(); + + // Verify a 0-row RecordBatch is still emitted when streams/_default returns 0 messages at offset = 0. + assert.strictEqual(batches.length, 1); + assert.strictEqual(batches[0].numRows, 0); + assert.strictEqual(batches[0].numCols, 0); + }); + + it('should read the schema from the destination table when ReadRows at offset = 0 ends without rows or a schema', async () => { + // streams/_default sends no ReadRowsResponse for a result without rows, + // so the schema has to come from a read session on the destination table. + const fakeJob = bq.job('job-empty-2', {location: 'US'}); + stubJobGetQueryResults(fakeJob, (_opts, cb) => { + process.nextTick(() => { + cb(null, [], null, {jobComplete: true, totalRows: '0'}); + }); + }); + stubJobGetMetadata(fakeJob, { + metadata: { + jobReference: {projectId: 'test-project', jobId: 'job-empty-2'}, + configuration: { + query: { + destinationTable: { + projectId: 'test-project', + datasetId: '_anon', + tableId: 'anon_table', + }, + }, + }, + }, + }); + + stubJobsQuery(bq, { + job: fakeJob, + res: { + jobComplete: false, + jobReference: { + projectId: 'test-project', + jobId: 'job-empty-2', + location: 'US', + }, + }, + }); + + const sessionReqs: CreateReadSessionRequest[] = []; + stubStorageReadClient(bq, { + createReadSession: async req => { + sessionReqs.push(req); + return [ + { + arrowSchema: { + serializedSchema: Buffer.from( + SAMPLE_PAYLOAD.serializedSchema, + 'base64', + ), + }, + }, + ]; + }, + readRows: () => fakeReadRowsStream(), + }); + + const batches: arrow.RecordBatch[] = await bq + .createQueryArrowStream('SELECT id, name FROM tbl WHERE FALSE') + .toArray(); + + // Verify the 0-row RecordBatch carries the destination table's schema. + assert.strictEqual(batches.length, 1); + assert.strictEqual(batches[0].numRows, 0); + assert.deepStrictEqual( + batches[0].schema.fields.map(f => f.name), + ['id', 'name'], + ); + assert.strictEqual(sessionReqs.length, 1); + assert.strictEqual(sessionReqs[0].parent, 'projects/test-project'); + assert.strictEqual( + sessionReqs[0].readSession.table, + 'projects/test-project/datasets/_anon/tables/anon_table', + ); + assert.strictEqual(sessionReqs[0].maxStreamCount, 1); + }); + + it('should forward jobs.get errors when reading the schema of a result without rows', async () => { + const fakeJob = bq.job('job-empty-3', {location: 'US'}); + stubJobGetQueryResults(fakeJob, (_opts, cb) => { + process.nextTick(() => { + cb(null, [], null, {jobComplete: true, totalRows: '0'}); + }); + }); + stubJobGetMetadata(fakeJob, {err: new Error('jobs.get failed')}); + + stubJobsQuery(bq, { + job: fakeJob, + res: { + jobComplete: false, + jobReference: { + projectId: 'test-project', + jobId: 'job-empty-3', + location: 'US', + }, + }, + }); + + stubStorageReadClient(bq, { + readRows: () => fakeReadRowsStream(), + }); + + // Verify the jobs.get error reaches the stream instead of a schema-less batch. + await assert.rejects( + bq.createQueryArrowStream('SELECT id FROM tbl WHERE FALSE').toArray(), + /jobs.get failed/, + ); + }); + }); }); From 6fce0fddfebae30f7be52c6a1192bb7f45d32d04 Mon Sep 17 00:00:00 2001 From: ohmayr Date: Tue, 20 Oct 2026 20:43:31 +0000 Subject: [PATCH 2/3] refactor(bigquery): keep Arrow read-session types internal IReadSession is only used inside arrow.ts, so stop exporting it. CreateReadSessionRequest stays exported for the unit tests and is now documented as @private, like the other test-only exports in arrow.ts. --- handwritten/bigquery/src/arrow.ts | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/handwritten/bigquery/src/arrow.ts b/handwritten/bigquery/src/arrow.ts index 9d849b4b9f22..0843d9ad7cd9 100644 --- a/handwritten/bigquery/src/arrow.ts +++ b/handwritten/bigquery/src/arrow.ts @@ -37,7 +37,7 @@ interface ReadRowsStream extends NodeJS.EventEmitter { resume(): void; } -export interface IReadSession { +interface IReadSession { arrowSchema?: { serializedSchema?: Uint8Array | string | null; } | null; @@ -46,6 +46,12 @@ export interface IReadSession { }> | null; } +/** + * The `CreateReadSession` request that {@link BigQuery#createQueryArrowStream} + * sends. Exported for tests. + * + * @private + */ export interface CreateReadSessionRequest { parent: string; readSession: { From b31bbd7c914d0193842215af1aa702019be333cb Mon Sep 17 00:00:00 2001 From: ohmayr Date: Wed, 21 Oct 2026 00:34:46 +0000 Subject: [PATCH 3/3] fix(bigquery): skip Arrow read session after destroy; async test stubs Return from readDestinationSchema without a CreateReadSession call if the consumer destroyed the stream during jobs.get. Stub job.getQueryResults with async functions in the unit tests, matching the awaited promise API instead of adapting callbacks. --- handwritten/bigquery/src/arrow.ts | 3 +- handwritten/bigquery/test/arrow.ts | 141 +++++++++++++++++++---------- 2 files changed, 95 insertions(+), 49 deletions(-) diff --git a/handwritten/bigquery/src/arrow.ts b/handwritten/bigquery/src/arrow.ts index 8cd4ba59fbf7..c108fe20ba59 100644 --- a/handwritten/bigquery/src/arrow.ts +++ b/handwritten/bigquery/src/arrow.ts @@ -385,7 +385,8 @@ async function readDestinationSchema( // the job's current metadata instead of the response that created it. const [metadata] = await job.getMetadata(); const destRef = metadata?.configuration?.query?.destinationTable; - if (!destRef) { + // Skip the read session if the consumer destroyed the stream meanwhile. + if (!destRef || ctx.stream.destroyed) { return undefined; } const [session] = await getStorageReadClient(bq, storage).createReadSession({ diff --git a/handwritten/bigquery/test/arrow.ts b/handwritten/bigquery/test/arrow.ts index 58e249aefd3d..bc754606105e 100644 --- a/handwritten/bigquery/test/arrow.ts +++ b/handwritten/bigquery/test/arrow.ts @@ -17,13 +17,7 @@ import {EventEmitter} from 'events'; import {describe, it, before, after, beforeEach} from 'mocha'; import * as Module from 'module'; import * as arrow from 'apache-arrow'; -import { - BigQuery, - Job, - Query, - QueryResultsOptions, - QueryRowsCallback, -} from '../src'; +import {BigQuery, Job, Query, QueryResultsOptions} from '../src'; import { BigQueryReadClient, checkArrowDependencies, @@ -152,19 +146,16 @@ function assertReadRowsListenersRemoved(grpcStream: EventEmitter): void { } /** - * Stubs `job.getQueryResults` with a callback-style `impl`, returning a - * promise because the code calls the promisified method. + * Stubs `job.getQueryResults`, which the code awaits, with an async `impl` + * that resolves to `[rows, nextQuery, response]` or rejects. */ function stubJobGetQueryResults( job: Job, - impl: (opts: QueryResultsOptions, cb: QueryRowsCallback) => void, + impl: ( + opts: QueryResultsOptions, + ) => Promise<[unknown[], null, bigquery.IGetQueryResultsResponse]>, ): void { - job.getQueryResults = ((opts: QueryResultsOptions) => - new Promise((resolve, reject) => { - impl(opts, (err, rows, nextQuery, res) => - err ? reject(err) : resolve([rows, nextQuery, res]), - ); - })) as unknown as Job['getQueryResults']; + job.getQueryResults = impl as unknown as Job['getQueryResults']; } /** @@ -1026,17 +1017,16 @@ describe('Arrow Support', () => { }> = []; // First poll returns jobComplete: false; second poll returns jobComplete: true. - stubJobGetQueryResults(fakeJob, (opts, cb) => { + stubJobGetQueryResults(fakeJob, async opts => { getQueryResultsCalls.push(opts); // Verify ReadRows has not been called while waiting for the job to complete. assert.strictEqual(capturedReadRowsReqs.length, 0); const isDone = getQueryResultsCalls.length >= 2; - process.nextTick(() => { - cb(null, [], null, { - jobComplete: isDone, - totalRows: isDone ? '2' : undefined, - }); - }); + return [ + [], + null, + {jobComplete: isDone, totalRows: isDone ? '2' : undefined}, + ]; }); stubJobsQuery(bq, { @@ -1086,8 +1076,9 @@ describe('Arrow Support', () => { it('should fail with a timeout error when jobs.query returns jobComplete = false and timeoutMs was specified', async () => { const fakeJob = bq.job('job-timeout-1', {location: 'US'}); let getQueryResultsCalled = false; - stubJobGetQueryResults(fakeJob, () => { + stubJobGetQueryResults(fakeJob, async () => { getQueryResultsCalled = true; + return [[], null, {jobComplete: true}]; }); const getCapturedReq = stubJobsQuery(bq, { @@ -1127,11 +1118,11 @@ describe('Arrow Support', () => { const fakeJob = bq.job('job-wait-destroy', {location: 'US'}); const stream = bq.createQueryArrowStream('SELECT id, name FROM tbl'); let polls = 0; - stubJobGetQueryResults(fakeJob, (_opts, cb) => { + stubJobGetQueryResults(fakeJob, async () => { polls++; // The consumer gives up while this poll is in flight. stream.destroy(); - process.nextTick(() => cb(null, [], null, {jobComplete: false})); + return [[], null, {jobComplete: false}]; }); stubJobsQuery(bq, { job: fakeJob, @@ -1153,7 +1144,7 @@ describe('Arrow Support', () => { }); stream.resume(); - // The poll's callback runs on the next tick; wait past it. + // Wait for the in-flight poll to resolve. await new Promise(resolve => setImmediate(resolve)); // Verify the in-flight poll didn't start another one or call ReadRows. @@ -1164,9 +1155,9 @@ describe('Arrow Support', () => { it('should wait for the job instead of failing when timeoutMs is 0, matching BigQuery#query', async () => { const fakeJob = bq.job('job-timeout-0', {location: 'US'}); const getQueryResultsCalls: QueryResultsOptions[] = []; - stubJobGetQueryResults(fakeJob, (opts, cb) => { + stubJobGetQueryResults(fakeJob, async opts => { getQueryResultsCalls.push(opts); - process.nextTick(() => cb(null, [], null, {jobComplete: true})); + return [[], null, {jobComplete: true}]; }); stubJobsQuery(bq, { @@ -1199,10 +1190,8 @@ describe('Arrow Support', () => { it('should forward job.getQueryResults errors to the stream without calling ReadRows', async () => { const fakeJob = bq.job('job-err-1', {location: 'US'}); - stubJobGetQueryResults(fakeJob, (_opts, cb) => { - process.nextTick(() => { - cb(new Error('Resources exceeded during query execution')); - }); + stubJobGetQueryResults(fakeJob, async () => { + throw new Error('Resources exceeded during query execution'); }); let readRowsCalled = false; @@ -1235,11 +1224,11 @@ describe('Arrow Support', () => { it('should emit a 0-row, 0-column RecordBatch when ReadRows at offset = 0 ends without rows and the job has no destination table', async () => { const fakeJob = bq.job('job-empty-1', {location: 'US'}); - stubJobGetQueryResults(fakeJob, (_opts, cb) => { - process.nextTick(() => { - cb(null, [], null, {jobComplete: true, totalRows: '0'}); - }); - }); + stubJobGetQueryResults(fakeJob, async () => [ + [], + null, + {jobComplete: true, totalRows: '0'}, + ]); // Statements without a result table, like DECLARE, have no destination. stubJobGetMetadata(fakeJob, {metadata: {configuration: {query: {}}}}); @@ -1274,11 +1263,11 @@ describe('Arrow Support', () => { // streams/_default sends no ReadRowsResponse for a result without rows, // so the schema has to come from a read session on the destination table. const fakeJob = bq.job('job-empty-2', {location: 'US'}); - stubJobGetQueryResults(fakeJob, (_opts, cb) => { - process.nextTick(() => { - cb(null, [], null, {jobComplete: true, totalRows: '0'}); - }); - }); + stubJobGetQueryResults(fakeJob, async () => [ + [], + null, + {jobComplete: true, totalRows: '0'}, + ]); stubJobGetMetadata(fakeJob, { metadata: { jobReference: {projectId: 'test-project', jobId: 'job-empty-2'}, @@ -1344,13 +1333,69 @@ describe('Arrow Support', () => { assert.strictEqual(sessionReqs[0].maxStreamCount, 1); }); + it('should not create a read session if the stream is destroyed while reading the job metadata', async () => { + const fakeJob = bq.job('job-empty-destroy', {location: 'US'}); + const stream = bq.createQueryArrowStream( + 'SELECT id FROM tbl WHERE FALSE', + ); + stubJobGetQueryResults(fakeJob, async () => [ + [], + null, + {jobComplete: true, totalRows: '0'}, + ]); + fakeJob.getMetadata = (async () => { + // The consumer gives up while jobs.get is in flight. + stream.destroy(); + return [ + { + configuration: { + query: { + destinationTable: { + projectId: 'test-project', + datasetId: '_anon', + tableId: 'anon_table', + }, + }, + }, + }, + ]; + }) as unknown as Job['getMetadata']; + stubJobsQuery(bq, { + job: fakeJob, + res: { + jobComplete: false, + jobReference: { + projectId: 'test-project', + jobId: 'job-empty-destroy', + location: 'US', + }, + }, + }); + let readSessions = 0; + stubStorageReadClient(bq, { + createReadSession: async () => { + readSessions++; + return [{}]; + }, + readRows: () => fakeReadRowsStream(), + }); + + stream.resume(); + await new Promise(resolve => stream.once('close', resolve)); + // Let the query finish its work after jobs.get resolves. + await new Promise(resolve => setImmediate(resolve)); + + // Verify no read session was created for the destroyed stream. + assert.strictEqual(readSessions, 0); + }); + it('should forward jobs.get errors when reading the schema of a result without rows', async () => { const fakeJob = bq.job('job-empty-3', {location: 'US'}); - stubJobGetQueryResults(fakeJob, (_opts, cb) => { - process.nextTick(() => { - cb(null, [], null, {jobComplete: true, totalRows: '0'}); - }); - }); + stubJobGetQueryResults(fakeJob, async () => [ + [], + null, + {jobComplete: true, totalRows: '0'}, + ]); stubJobGetMetadata(fakeJob, {err: new Error('jobs.get failed')}); stubJobsQuery(bq, {