From e5144620d6615419f79ff6667e5c529a8d23dfeb Mon Sep 17 00:00:00 2001 From: ohmayr Date: Tue, 20 Oct 2026 19:55:03 +0000 Subject: [PATCH 1/2] feat(bigquery): support jobs.insert in createQueryArrowStream Routes queries that cannot use jobs.query (!queryReq, such as jobId, dryRun, jobTimeoutMs, or other jobs.insert-only options) through createQueryJob: - Keeps timeoutMs and jobCreationMode, which only apply to jobs.query, out of the job configuration. - Populates stream.jobReference and emits the job event once created. - Emits a single 0-row, 0-column RecordBatch immediately when dryRun is true without polling getQueryResults or calling ReadRows. - Otherwise waits for job completion via job.getQueryResults (maxResults: 0), bounded by timeoutMs like BigQuery#query, and streams Arrow batches from the anonymous default job stream (streams/_default) at offset = 0. --- handwritten/bigquery/src/arrow.ts | 86 +++++- handwritten/bigquery/system-test/arrow.ts | 94 +++++++ handwritten/bigquery/test/arrow.ts | 311 +++++++++++++++++++++- 3 files changed, 467 insertions(+), 24 deletions(-) diff --git a/handwritten/bigquery/src/arrow.ts b/handwritten/bigquery/src/arrow.ts index 0843d9ad7cd..9358e853e86 100644 --- a/handwritten/bigquery/src/arrow.ts +++ b/handwritten/bigquery/src/arrow.ts @@ -110,9 +110,9 @@ export type ArrowRecordBatchStream = ResourceStream & { */ 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. + * How long to wait for the query to complete, in milliseconds. If it doesn't + * complete in time, the stream fails. If unset or 0, the stream waits until + * the job completes. */ timeoutMs?: number; }; @@ -520,13 +520,16 @@ 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`. + * if the consumer destroys `stream`. If `timeoutMs` is set, fails like + * {@link BigQuery#query} when the job doesn't complete in time instead of + * polling again. * * @private */ async function waitForJobCompletion( job: Job, stream: ArrowRecordBatchStream, + timeoutMs?: number, ): Promise { while (!stream.destroyed) { // Set maxResults: 0 so jobs.getQueryResults blocks up to the server's @@ -535,6 +538,9 @@ async function waitForJobCompletion( const [, , res] = await job.getQueryResults({ autoPaginate: false, maxResults: 0, + // 0 means no limit, as in BigQuery#query. Sending 0 would make + // jobs.getQueryResults return at once and this loop spin. + ...(timeoutMs && {timeoutMs}), }); if (res?.jobComplete !== false) { return; @@ -542,6 +548,62 @@ async function waitForJobCompletion( } } +/** + * Waits for `job` to finish, then streams all of its rows from the default job + * stream. An unfinished jobs.query and jobs.insert return no rows or schema, + * so this reads every row from offset 0 and passes the job to look up the + * schema of an empty result. + * + * @private + */ +async function streamJobWhenComplete( + ctx: ArrowStreamContext, + job: Job, + jobRef: bigquery.IJobReference | undefined, + timeoutMs?: number, +): Promise { + await waitForJobCompletion(job, ctx.stream, timeoutMs); + if (ctx.stream.destroyed) { + return; + } + await streamDefaultJobStream(ctx, jobRef, 0, undefined, job); +} + +/** + * Runs `queryObj` with `jobs.insert` and streams its results. A dry run emits + * one 0-row, 0-column `RecordBatch` without waiting or streaming. + * + * @private + */ +async function streamQueryJob( + ctx: ArrowStreamContext, + queryObj: ArrowQuery, +): Promise { + const {bq, arrow, stream} = ctx; + // timeoutMs only bounds the wait below, and jobCreationMode only applies to + // jobs.query, so keep both out of the job configuration. + const {timeoutMs, ...jobQuery} = queryObj; + delete jobQuery.jobCreationMode; + // Throws for an invalid request, for example a destination that isn't a + // Table, or invalid parameters. + const [job, resp] = await bq.createQueryJob(jobQuery); + if (!job) { + throw new Error('jobs.insert returned no job.'); + } + + stream.jobReference = resp?.jobReference; + // For a dry run, job.metadata holds the query's statistics. + stream.emit('job', job); + + // A dry run validates the query without running it. + if (queryObj.dryRun) { + pushBatches(stream, decodeArrowPayload(arrow)); + return; + } + + await streamJobWhenComplete(ctx, job, resp?.jobReference, timeoutMs); +} + /** * Executes a query and returns a stream of `apache-arrow` `RecordBatch` objects. * @@ -594,8 +656,11 @@ export function createQueryArrowStream_( 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.'); + // buildQueryRequest_ returns undefined when the query sets an option + // that only jobs.insert accepts (e.g., jobId, dryRun, or jobTimeoutMs), + // or when FAST_QUERY_PATH=DISABLED. + await streamQueryJob(ctx, queryObj); + return; } // Default to JOB_CREATION_OPTIONAL unless overridden on the query or client. @@ -630,14 +695,7 @@ export function createQueryArrowStream_( '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); + await streamJobWhenComplete(ctx, job, res.jobReference); return; } diff --git a/handwritten/bigquery/system-test/arrow.ts b/handwritten/bigquery/system-test/arrow.ts index 980b1d86d9e..e801ee4c039 100644 --- a/handwritten/bigquery/system-test/arrow.ts +++ b/handwritten/bigquery/system-test/arrow.ts @@ -538,4 +538,98 @@ describe('Arrow Query Results', () => { /The query did not complete before 1ms/, ); }); + + it('should execute queries with jobs.insert-only options via jobs.insert and stream Arrow results from streams/_default at offset = 0', async () => { + const sql = + 'SELECT x, x * 2 AS doubled FROM UNNEST(GENERATE_ARRAY(1, 25)) AS x ORDER BY x'; + let emittedJobId: string | undefined; + + const stream = bigquery.createQueryArrowStream({ + query: sql, + jobTimeoutMs: 60000, + location: 'EU', + useQueryCache: false, + }); + stream.on('job', job => { + emittedJobId = job.id; + }); + + const batches: arrow.RecordBatch[] = await stream.toArray(); + const [expectedRows] = await bigquery.query({ + query: sql, + jobTimeoutMs: 60000, + location: 'EU', + useQueryCache: false, + }); + + // Verify the query ran via jobs.insert (no queryId, jobReference in EU) and matches STRUCT_ENCODING rows. + assert.strictEqual(stream.queryId, undefined); + assert.ok(stream.jobReference?.jobId); + assert.strictEqual(stream.jobReference?.location, 'EU'); + assert.strictEqual(emittedJobId, stream.jobReference?.jobId); + + const table = new arrow.Table(batches); + assert.strictEqual(table.numRows, expectedRows.length); + for (let i = 0; i < expectedRows.length; i++) { + assert.strictEqual( + Number(table.getChild('x')!.get(i)), + expectedRows[i].x, + ); + assert.strictEqual( + Number(table.getChild('doubled')!.get(i)), + expectedRows[i].doubled, + ); + } + }); + + it('should raise the same error as query() when a jobs.insert query fails during execution', async () => { + const failingQuery = { + query: 'SELECT 1 / 0', + jobTimeoutMs: 60000, + useQueryCache: false, + }; + + let expectedError: (Error & {code?: number}) | undefined; + try { + await bigquery.query(failingQuery); + } catch (err) { + expectedError = err as Error & {code?: number}; + } + assert.ok(expectedError, 'Expected bigquery.query() to reject'); + + let actualError: (Error & {code?: number}) | undefined; + try { + await bigquery.createQueryArrowStream(failingQuery).toArray(); + } catch (err) { + actualError = err as Error & {code?: number}; + } + + // Verify jobs.insert Arrow queries surface the same error code and message as query(). + assert.ok(actualError, 'Expected createQueryArrowStream() to reject'); + assert.strictEqual(actualError.code, expectedError.code); + assert.match(actualError.message, /division by zero/i); + }); + + it('should emit a 0-row RecordBatch preserving the schema for an empty result read via jobs.insert and streams/_default', async () => { + // jobs.insert returns no rows inline, so the schema of an empty result + // has to come from ReadRows on streams/_default. + const stream = bigquery.createQueryArrowStream({ + query: + 'SELECT x, CAST(x AS STRING) AS s FROM UNNEST(CAST([] AS ARRAY)) AS x', + jobTimeoutMs: 60000, + useQueryCache: false, + }); + const batches: arrow.RecordBatch[] = await stream.toArray(); + + // Verify the query ran via jobs.insert and emitted one 0-row batch with the column names and types. + assert.strictEqual(stream.queryId, undefined); + 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)); + }); }); diff --git a/handwritten/bigquery/test/arrow.ts b/handwritten/bigquery/test/arrow.ts index 365716af881..260d63cefe0 100644 --- a/handwritten/bigquery/test/arrow.ts +++ b/handwritten/bigquery/test/arrow.ts @@ -20,6 +20,8 @@ import * as arrow from 'apache-arrow'; import { BigQuery, Job, + JobCallback, + JobOptions, Query, QueryResultsOptions, QueryRowsCallback, @@ -211,6 +213,32 @@ const SAMPLE_READ_ROWS_RESPONSE = { }, }; +/** + * Stubs `bq.createJob` for callback-based `jobs.insert` tests. + */ +function stubCreateJob( + bq: BigQuery, + impl: (opts: JobOptions, cb: JobCallback) => void, +): void { + bq.createJob = impl as BigQuery['createJob']; +} + +/** + * Stubs `bq.createQueryJob` with a callback-style `impl`, returning a promise + * because the code calls the promisified method. + */ +function stubCreateQueryJob( + bq: BigQuery, + impl: (opts: Query | string, cb: JobCallback) => void, +): void { + bq.createQueryJob = ((opts: Query | string) => + new Promise((resolve, reject) => { + impl(opts, (err, job, resp) => + err ? reject(err) : resolve([job, resp]), + ); + })) as unknown as BigQuery['createQueryJob']; +} + describe('Arrow Support', () => { let originalRequire: typeof Module.prototype.require; @@ -507,16 +535,6 @@ describe('Arrow Support', () => { ); }); - it('should fail when the query configuration cannot be executed via jobs.query', async () => { - // TODO: Replace with jobs.insert (createQueryJob) fallback tests once added. - await assert.rejects( - bq - .createQueryArrowStream({query: 'SELECT 1', jobId: 'custom-job-1'}) - .toArray(), - /Query configuration cannot be executed via jobs\.query\./, - ); - }); - it('should forward errors thrown while building the jobs.query request', async () => { // A null parameter without an explicit type makes buildQueryRequest_ throw. await assert.rejects( @@ -1334,4 +1352,277 @@ describe('Arrow Support', () => { ); }); }); + + describe('BigQuery#createQueryArrowStream (jobs.insert)', () => { + let bq: BigQuery; + + beforeEach(() => { + bq = new BigQuery({projectId: 'test-project'}); + }); + + it('should route jobs.insert-only options through createQueryJob without queryResultsFormat, wait for completion, and stream from streams/_default at offset = 0', async () => { + const getCapturedQueryReq = stubJobsQuery(bq, {res: {}}); + + let capturedJobOpts: JobOptions | undefined; + let getQueryResultsCalled = false; + const capturedReadRowsReqs: Array<{ + readStream?: string; + offset?: number; + }> = []; + + const fakeJob = bq.job('custom-job-1', {location: 'EU'}); + fakeJob.metadata = { + jobReference: { + projectId: 'test-project', + jobId: 'custom-job-1', + location: 'EU', + }, + }; + stubJobGetQueryResults(fakeJob, (opts, cb) => { + getQueryResultsCalled = true; + assert.strictEqual(capturedReadRowsReqs.length, 0); + assert.deepStrictEqual(opts, { + autoPaginate: false, + maxResults: 0, + }); + process.nextTick(() => { + cb(null, [], null, {jobComplete: true, totalRows: '2'}); + }); + }); + + // Intercept createJob (called by createQueryJob) to inspect the outgoing jobs.insert payload. + stubCreateJob(bq, (reqOpts, cb) => { + capturedJobOpts = reqOpts; + process.nextTick(() => { + cb(null, fakeJob, fakeJob.metadata); + }); + }); + + stubStorageReadClient(bq, { + readRows: (req: {readStream?: string; offset?: number}) => { + capturedReadRowsReqs.push(req); + const grpcStream = new EventEmitter(); + process.nextTick(() => { + grpcStream.emit('data', { + arrowSchema: { + serializedSchema: Buffer.from( + SAMPLE_PAYLOAD.serializedSchema, + 'base64', + ), + }, + arrowRecordBatch: { + serializedRecordBatch: Buffer.from( + SAMPLE_PAYLOAD.serializedRecordBatch, + 'base64', + ), + }, + }); + grpcStream.emit('end'); + }); + return grpcStream; + }, + }); + + let emittedJob: Job | undefined; + const stream = bq.createQueryArrowStream({ + query: 'SELECT id, name FROM tbl', + jobId: 'custom-job-1', + location: 'EU', + jobCreationMode: 'JOB_CREATION_OPTIONAL', + }); + stream.on('job', j => { + emittedJob = j; + }); + + const batches: arrow.RecordBatch[] = await stream.toArray(); + + // Verify jobs.query was bypassed and jobs.insert did not receive queryResultsFormat or jobCreationMode. + assert.strictEqual(getCapturedQueryReq(), undefined); + const queryConfig = capturedJobOpts?.configuration?.query; + assert.ok(queryConfig); + assert.strictEqual(queryConfig.query, 'SELECT id, name FROM tbl'); + assert.strictEqual('queryResultsFormat' in queryConfig, false); + assert.strictEqual('jobCreationMode' in queryConfig, false); + + // Verify metadata, job wait, and ReadRows at offset = 0 in the job's EU location. + assert.strictEqual(stream.queryId, undefined); + assert.deepStrictEqual(stream.jobReference, { + projectId: 'test-project', + jobId: 'custom-job-1', + location: 'EU', + }); + assert.strictEqual(emittedJob, fakeJob); + assert.strictEqual(getQueryResultsCalled, true); + assert.deepStrictEqual(capturedReadRowsReqs, [ + { + readStream: + 'projects/test-project/locations/EU/jobs/custom-job-1/streams/_default', + offset: 0, + }, + ]); + assert.strictEqual(batches.length, 1); + assert.strictEqual(batches[0].numRows, 2); + }); + + it('should not retry a failed jobs.query call as jobs.insert', async () => { + stubJobsQuery(bq, {err: new Error('jobs.query failed')}); + let createQueryJobCalled = false; + stubCreateQueryJob(bq, () => { + createQueryJobCalled = true; + }); + + // Verify a failed jobs.query surfaces its error directly without falling back to jobs.insert. + await assert.rejects( + bq.createQueryArrowStream('SELECT 1').toArray(), + /jobs\.query failed/, + ); + assert.strictEqual(createQueryJobCalled, false); + }); + + it('should emit a 0-row, 0-column RecordBatch without waiting or calling ReadRows when dryRun is true', async () => { + const fakeJob = bq.job('dry-run-job-1', {location: 'US'}); + fakeJob.metadata = { + jobReference: { + projectId: 'test-project', + jobId: 'dry-run-job-1', + location: 'US', + }, + }; + let getQueryResultsCalled = false; + stubJobGetQueryResults(fakeJob, () => { + getQueryResultsCalled = true; + }); + stubCreateQueryJob(bq, (_opts, cb) => { + process.nextTick(() => { + cb(null, fakeJob, fakeJob.metadata); + }); + }); + + const stream = bq.createQueryArrowStream({ + query: 'SELECT 1', + dryRun: true, + }); + const batches: arrow.RecordBatch[] = await stream.toArray(); + + // Verify dryRun emits an empty batch immediately without polling getQueryResults. + assert.strictEqual(getQueryResultsCalled, false); + assert.strictEqual(stream.jobReference?.jobId, 'dry-run-job-1'); + assert.strictEqual(batches.length, 1); + assert.strictEqual(batches[0].numRows, 0); + assert.strictEqual(batches[0].numCols, 0); + }); + + it('should bound the jobs.insert wait with timeoutMs and keep timeoutMs out of configuration.query', async () => { + const fakeJob = bq.job('timeout-job-1', {location: 'US'}); + fakeJob.metadata = { + jobReference: { + projectId: 'test-project', + jobId: 'timeout-job-1', + location: 'US', + }, + }; + let waitOpts: QueryResultsOptions | undefined; + stubJobGetQueryResults(fakeJob, (opts, cb) => { + waitOpts = opts; + // Job.getQueryResults fails like this when timeoutMs passes before + // the job completes. + process.nextTick(() => { + cb(new Error('The query did not complete before 1000ms'), null); + }); + }); + let capturedJobOpts: JobOptions | undefined; + stubCreateJob(bq, (reqOpts, cb) => { + capturedJobOpts = reqOpts; + process.nextTick(() => { + cb(null, fakeJob, fakeJob.metadata); + }); + }); + + // Verify the wait fails like BigQuery#query when timeoutMs passes. + await assert.rejects( + bq + .createQueryArrowStream({ + query: 'SELECT 1', + jobTimeoutMs: 60000, + timeoutMs: 1000, + }) + .toArray(), + /The query did not complete before 1000ms/, + ); + // Verify timeoutMs bounded jobs.getQueryResults and wasn't sent as a + // jobs.insert configuration field. + assert.strictEqual(waitOpts?.timeoutMs, 1000); + assert.strictEqual( + (capturedJobOpts?.configuration?.query as Record) + ?.timeoutMs, + undefined, + ); + }); + + it('should wait without a timeout when timeoutMs is 0 on the jobs.insert path', async () => { + const fakeJob = bq.job('timeout-zero-job', {location: 'US'}); + fakeJob.metadata = { + jobReference: { + projectId: 'test-project', + jobId: 'timeout-zero-job', + location: 'US', + }, + }; + const waitOpts: QueryResultsOptions[] = []; + stubJobGetQueryResults(fakeJob, (opts, cb) => { + waitOpts.push(opts); + process.nextTick(() => + cb(new Error('stop after the first poll'), null), + ); + }); + stubCreateJob(bq, (_reqOpts, cb) => { + process.nextTick(() => cb(null, fakeJob, fakeJob.metadata)); + }); + + await assert.rejects( + bq + .createQueryArrowStream({ + query: 'SELECT 1', + jobTimeoutMs: 60000, + timeoutMs: 0, + }) + .toArray(), + /stop after the first poll/, + ); + // Verify 0 means "no limit", as in BigQuery#query: sending timeoutMs: 0 + // would make jobs.getQueryResults return at once and poll in a loop. + assert.deepStrictEqual(waitOpts, [{autoPaginate: false, maxResults: 0}]); + }); + + it('should forward errors thrown synchronously by createQueryJob', async () => { + // jobId forces jobs.insert; a null parameter without a type makes + // createQueryJob throw before sending the request. + await assert.rejects( + bq + .createQueryArrowStream({ + query: 'SELECT ?', + jobId: 'custom-job-throw', + params: [null], + }) + .toArray(), + /Parameter types must be provided for null values/, + ); + }); + + it('should forward createQueryJob errors to the stream', async () => { + stubCreateQueryJob(bq, (_opts, cb) => { + process.nextTick(() => { + cb(new Error('Permission denied creating job'), null); + }); + }); + + // Verify errors from createQueryJob reject stream consumption. + await assert.rejects( + bq + .createQueryArrowStream({query: 'SELECT 1', jobId: 'custom-err-1'}) + .toArray(), + /Permission denied creating job/, + ); + }); + }); }); From f41ec7d81049024cba0772de8c7451c4ca6ac69b Mon Sep 17 00:00:00 2001 From: ohmayr Date: Wed, 21 Oct 2026 00:39:39 +0000 Subject: [PATCH 2/2] fix(bigquery): stop Arrow jobs.insert path after destroy; async stubs Return after createQueryJob if the consumer destroyed the stream, so no 'job' event is emitted and the job isn't waited on. Stub createQueryJob with async functions in the unit tests, matching the awaited promise API. Also say in the createQueryArrowStream JSDoc that timeoutMs bounds the jobs.insert wait too and that jobs.insert queries read every row with the Storage Read API, and add system tests for timeoutMs and dryRun on the jobs.insert path. --- handwritten/bigquery/src/arrow.ts | 4 ++ handwritten/bigquery/src/bigquery.ts | 22 ++++---- handwritten/bigquery/system-test/arrow.ts | 37 +++++++++++++ handwritten/bigquery/test/arrow.ts | 67 ++++++++++++++++------- 4 files changed, 101 insertions(+), 29 deletions(-) diff --git a/handwritten/bigquery/src/arrow.ts b/handwritten/bigquery/src/arrow.ts index 372c31a1a02..7bad9894504 100644 --- a/handwritten/bigquery/src/arrow.ts +++ b/handwritten/bigquery/src/arrow.ts @@ -588,6 +588,10 @@ async function streamQueryJob( // Throws for an invalid request, for example a destination that isn't a // Table, or invalid parameters. const [job, resp] = await bq.createQueryJob(jobQuery); + // The consumer may have destroyed the stream while jobs.insert ran. + if (stream.destroyed) { + return; + } if (!job) { throw new Error('jobs.insert returned no job.'); } diff --git a/handwritten/bigquery/src/bigquery.ts b/handwritten/bigquery/src/bigquery.ts index 3470b7b4c32..da52ab6e0ea 100644 --- a/handwritten/bigquery/src/bigquery.ts +++ b/handwritten/bigquery/src/bigquery.ts @@ -391,19 +391,21 @@ export class BigQuery extends Service { * Requires the optional `@google-cloud/bigquery-storage` (>= 6.3.0) package * to be installed. * - * Rows that don't fit in the first `jobs.query` response are read with the - * BigQuery Storage Read API, which must be enabled in the project and is - * billed separately. The stream reads more rows only as the caller consumes - * batches, so iterate over batches to keep memory bounded; collecting every - * batch (for example with `toArray()`) holds the whole result in memory. - * If you stop reading before the stream ends, call `stream.destroy()` to - * cancel the read; breaking out of `for await` does this for you. + * Rows that don't fit in the first `jobs.query` response, and every row of a + * query that runs with `jobs.insert` (for example, one that sets `jobId` or + * `jobTimeoutMs`), are read with the BigQuery Storage Read API, which must be + * enabled in the project and is billed separately. The stream reads more + * rows only as the caller consumes batches, so iterate over batches to keep + * memory bounded; collecting every batch (for example with `toArray()`) + * holds the whole result in memory. If you stop reading before the stream + * ends, call `stream.destroy()` to cancel the read; breaking out of + * `for await` does this for you. * * @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. + * @param {number} [options.timeoutMs] How long to wait for the query to + * complete, in milliseconds. If it doesn't complete in time, the stream + * fails. If unset or 0, the stream waits until the job completes. * @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 diff --git a/handwritten/bigquery/system-test/arrow.ts b/handwritten/bigquery/system-test/arrow.ts index e801ee4c039..c7c6c9296e0 100644 --- a/handwritten/bigquery/system-test/arrow.ts +++ b/handwritten/bigquery/system-test/arrow.ts @@ -610,6 +610,43 @@ describe('Arrow Query Results', () => { assert.match(actualError.message, /division by zero/i); }); + it('should reject when a jobs.insert query does not complete within caller-specified timeoutMs', async () => { + // jobTimeoutMs forces jobs.insert; timeoutMs then bounds the wait on + // jobs.getQueryResults, like BigQuery#query. + await assert.rejects( + bigquery + .createQueryArrowStream({ + query: + 'DECLARE n INT64 DEFAULT 50; SELECT x FROM UNNEST(GENERATE_ARRAY(1, n)) AS x ORDER BY x;', + jobTimeoutMs: 60000, + timeoutMs: 1, + useQueryCache: false, + }) + .toArray(), + /The query did not complete before 1ms/, + ); + }); + + it('should emit one 0-row, 0-column RecordBatch and the job with dry-run statistics for dryRun', async () => { + let emittedJob: Job | undefined; + const stream = bigquery.createQueryArrowStream({ + query: 'SELECT x FROM UNNEST(GENERATE_ARRAY(1, 10)) AS x', + dryRun: true, + }); + stream.on('job', job => { + emittedJob = job; + }); + const batches: arrow.RecordBatch[] = await stream.toArray(); + + // Verify the dry run returned no rows or columns, and its statistics + // are available on the emitted job. + assert.strictEqual(batches.length, 1); + assert.strictEqual(batches[0].numRows, 0); + assert.strictEqual(batches[0].numCols, 0); + assert.ok(emittedJob); + assert.ok(emittedJob.metadata?.statistics?.totalBytesProcessed); + }); + it('should emit a 0-row RecordBatch preserving the schema for an empty result read via jobs.insert and streams/_default', async () => { // jobs.insert returns no rows inline, so the schema of an empty result // has to come from ReadRows on streams/_default. diff --git a/handwritten/bigquery/test/arrow.ts b/handwritten/bigquery/test/arrow.ts index d7fa33d696d..be5b74e737f 100644 --- a/handwritten/bigquery/test/arrow.ts +++ b/handwritten/bigquery/test/arrow.ts @@ -220,19 +220,14 @@ function stubCreateJob( } /** - * Stubs `bq.createQueryJob` with a callback-style `impl`, returning a promise - * because the code calls the promisified method. + * Stubs `bq.createQueryJob`, which the code awaits, with an async `impl` + * that resolves to `[job, apiResponse]` or rejects. */ function stubCreateQueryJob( bq: BigQuery, - impl: (opts: Query | string, cb: JobCallback) => void, + impl: (opts: Query | string) => Promise<[Job, bigquery.IJob | undefined]>, ): void { - bq.createQueryJob = ((opts: Query | string) => - new Promise((resolve, reject) => { - impl(opts, (err, job, resp) => - err ? reject(err) : resolve([job, resp]), - ); - })) as unknown as BigQuery['createQueryJob']; + bq.createQueryJob = impl as unknown as BigQuery['createQueryJob']; } describe('Arrow Support', () => { @@ -1557,8 +1552,9 @@ describe('Arrow Support', () => { it('should not retry a failed jobs.query call as jobs.insert', async () => { stubJobsQuery(bq, {err: new Error('jobs.query failed')}); let createQueryJobCalled = false; - stubCreateQueryJob(bq, () => { + stubCreateQueryJob(bq, async () => { createQueryJobCalled = true; + throw new Error('createQueryJob should not be called'); }); // Verify a failed jobs.query surfaces its error directly without falling back to jobs.insert. @@ -1583,11 +1579,7 @@ describe('Arrow Support', () => { getQueryResultsCalled = true; return [[], null, {jobComplete: true}]; }); - stubCreateQueryJob(bq, (_opts, cb) => { - process.nextTick(() => { - cb(null, fakeJob, fakeJob.metadata); - }); - }); + stubCreateQueryJob(bq, async () => [fakeJob, fakeJob.metadata]); const stream = bq.createQueryArrowStream({ query: 'SELECT 1', @@ -1696,11 +1688,48 @@ describe('Arrow Support', () => { ); }); + it('should not emit job or wait for the job when the stream is destroyed during jobs.insert', async () => { + const fakeJob = bq.job('destroyed-insert-job', {location: 'US'}); + fakeJob.metadata = { + jobReference: { + projectId: 'test-project', + jobId: 'destroyed-insert-job', + location: 'US', + }, + }; + let getQueryResultsCalled = false; + stubJobGetQueryResults(fakeJob, async () => { + getQueryResultsCalled = true; + return [[], null, {jobComplete: true}]; + }); + const stream = bq.createQueryArrowStream({ + query: 'SELECT 1', + jobId: 'destroyed-insert-job', + }); + stubCreateQueryJob(bq, async () => { + // The consumer gives up while jobs.insert is in flight. + stream.destroy(); + return [fakeJob, fakeJob.metadata]; + }); + let jobEmitted = false; + stream.on('job', () => { + jobEmitted = true; + }); + + stream.resume(); + await new Promise(resolve => stream.once('close', resolve)); + // Let the query finish its work after jobs.insert resolves. + await new Promise(resolve => setImmediate(resolve)); + + // Verify nothing happened on the destroyed stream. + assert.strictEqual(jobEmitted, false); + assert.strictEqual(stream.jobReference, undefined); + assert.strictEqual(getQueryResultsCalled, false); + }); + it('should forward createQueryJob errors to the stream', async () => { - stubCreateQueryJob(bq, (_opts, cb) => { - process.nextTick(() => { - cb(new Error('Permission denied creating job'), null); - }); + stubCreateQueryJob(bq, async () => { + throw new Error('Permission denied creating job'); }); // Verify errors from createQueryJob reject stream consumption.