diff --git a/handwritten/bigquery/src/arrow.ts b/handwritten/bigquery/src/arrow.ts index c108fe20ba5..7bad9894504 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; }; @@ -521,13 +521,16 @@ async 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 @@ -536,6 +539,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; @@ -543,6 +549,66 @@ 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); + // 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.'); + } + + 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. * @@ -611,8 +677,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. @@ -651,14 +720,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/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 980b1d86d9e..c7c6c9296e0 100644 --- a/handwritten/bigquery/system-test/arrow.ts +++ b/handwritten/bigquery/system-test/arrow.ts @@ -538,4 +538,135 @@ 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 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. + 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 bc754606105..be5b74e737f 100644 --- a/handwritten/bigquery/test/arrow.ts +++ b/handwritten/bigquery/test/arrow.ts @@ -17,7 +17,14 @@ 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} from '../src'; +import { + BigQuery, + Job, + JobCallback, + JobOptions, + Query, + QueryResultsOptions, +} from '../src'; import { BigQueryReadClient, checkArrowDependencies, @@ -202,6 +209,27 @@ 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`, which the code awaits, with an async `impl` + * that resolves to `[job, apiResponse]` or rejects. + */ +function stubCreateQueryJob( + bq: BigQuery, + impl: (opts: Query | string) => Promise<[Job, bigquery.IJob | undefined]>, +): void { + bq.createQueryJob = impl as unknown as BigQuery['createQueryJob']; +} + describe('Arrow Support', () => { let originalRequire: typeof Module.prototype.require; @@ -537,16 +565,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( @@ -1421,4 +1439,306 @@ 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, async opts => { + getQueryResultsCalled = true; + assert.strictEqual(capturedReadRowsReqs.length, 0); + assert.deepStrictEqual(opts, { + autoPaginate: false, + maxResults: 0, + }); + return [[], 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, 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. + 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, async () => { + getQueryResultsCalled = true; + return [[], null, {jobComplete: true}]; + }); + stubCreateQueryJob(bq, async () => [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, async opts => { + waitOpts = opts; + // Job.getQueryResults fails like this when timeoutMs passes before + // the job completes. + throw new Error('The query did not complete before 1000ms'); + }); + 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, async opts => { + waitOpts.push(opts); + throw new Error('stop after the first poll'); + }); + 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 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, async () => { + throw new Error('Permission denied creating job'); + }); + + // Verify errors from createQueryJob reject stream consumption. + await assert.rejects( + bq + .createQueryArrowStream({query: 'SELECT 1', jobId: 'custom-err-1'}) + .toArray(), + /Permission denied creating job/, + ); + }); + }); });