diff --git a/handwritten/bigquery/src/arrow.ts b/handwritten/bigquery/src/arrow.ts index 7ba0dec4218..3e7702fc018 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, Query} from './bigquery'; +import type {Job} from './job'; import type bigquery from './types'; type ArrowModule = typeof import('apache-arrow'); @@ -36,7 +37,29 @@ 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; +} + export interface BigQueryReadClient { + createReadSession( + request: CreateReadSessionRequest, + callback: (err: Error | null, session?: IReadSession | null) => void, + ): void; readRows(request: {readStream: string; offset?: number}): ReadRowsStream; } @@ -62,6 +85,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 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; +}; + /** * State shared by the helpers that stream rows for one * {@link BigQuery#createQueryArrowStream} call. @@ -255,10 +293,37 @@ export function decodeArrowPayload( : [createEmptyRecordBatch(arrow, reader.schema)]; } +/** + * Emits a 0-row `RecordBatch` (preserving `serializedSchema` when present) and + * ends `stream`. + * + * @private + */ +function emitZeroRowBatchAndEnd( + arrow: ArrowModule, + stream: ArrowRecordBatchStream, + serializedSchema?: Uint8Array | string, +): void { + try { + for (const batch of decodeArrowPayload(arrow, serializedSchema)) { + if (stream._ended || stream.destroyed) { + return; + } + stream.push(batch); + } + if (!stream._ended && !stream.destroyed) { + stream.end(); + } + } catch (parseErr) { + stream.destroy(parseErr as Error); + } +} + /** * 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 */ @@ -267,6 +332,7 @@ function streamDefaultJobStream( jobRef: bigquery.IJobReference | undefined, offset: number, initialSchema?: Uint8Array | string, + job?: Job, ): void { const {bq, stream} = ctx; const projectId = resolveProjectId(bq, jobRef?.projectId); @@ -281,7 +347,61 @@ function streamDefaultJobStream( } const readStream = `projects/${projectId}/locations/${location}/jobs/${jobId}/streams/_default`; - streamReadRows(ctx, {readStream, offset, initialSchema}); + streamReadRows(ctx, {readStream, offset, initialSchema, job}); +} + +/** + * 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 `callback` gets no schema. + * + * @private + */ +function readDestinationSchema( + ctx: ArrowStreamContext, + job: Job, + callback: (err: Error | null, schema?: Uint8Array | string) => void, +): void { + 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. + void job.getMetadata((err: Error | null, metadata?: bigquery.IJob) => { + if (err) { + callback(err); + return; + } + const destRef = metadata?.configuration?.query?.destinationTable; + if (!destRef) { + callback(null); + return; + } + let readClient: BigQueryReadClient; + try { + readClient = getStorageReadClient(bq, storage); + } catch (err) { + callback(err as Error); + return; + } + readClient.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, + }, + (err, session) => { + callback(err, session?.arrowSchema?.serializedSchema || undefined); + }, + ); + }); } /** @@ -296,10 +416,12 @@ function streamReadRows( readStream: string; offset: number; initialSchema?: Uint8Array | string; + /** The finished job, to read the schema of a result without rows. */ + job?: Job; }, ): void { const {bq, arrow, storage, stream} = ctx; - const {readStream, offset, initialSchema} = options; + const {readStream, offset, initialSchema, job} = options; let readClient: BigQueryReadClient; try { @@ -332,6 +454,7 @@ function streamReadRows( let serializedSchema: Uint8Array | undefined = initialSchema ? toUint8Array(initialSchema) : undefined; + let emittedBatch = false; grpcStream.on('data', (resp: IReadRowsResponse) => { if (stream._ended || stream.destroyed) { @@ -359,6 +482,7 @@ function streamReadRows( cancelGrpcStream(); return; } + emittedBatch = true; if (!stream.push(batch)) { grpcStream.pause(); } @@ -381,12 +505,79 @@ function streamReadRows( grpcStream.on('end', () => { stream.removeListener('close', cancelGrpcStream); finished = true; - if (!stream._ended && !stream.destroyed) { - stream.end(); + if (stream._ended || stream.destroyed) { + return; + } + // 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) { + if (serializedSchema || !job) { + emitZeroRowBatchAndEnd(arrow, stream, serializedSchema); + return; + } + readDestinationSchema(ctx, job, (err, schema) => { + if (stream._ended || stream.destroyed) { + return; + } + if (err) { + stream.destroy(err); + return; + } + emitZeroRowBatchAndEnd(arrow, stream, schema); + }); + return; } + stream.end(); }); } +/** + * Waits for an unfinished BigQuery job to complete using `job.getQueryResults` + * without fetching any `STRUCT_ENCODING` rows (`maxResults: 0`). If `timeoutMs` + * is set, fails like {@link BigQuery#query} when the job doesn't complete in + * time instead of polling again. + * + * @private + */ +function waitForJobCompletion( + job: Job, + stream: ArrowRecordBatchStream, + onComplete: () => void, + timeoutMs?: number, +): void { + const poll = () => { + if (stream._ended || stream.destroyed) { + return; + } + // Set maxResults: 0 so jobs.getQueryResults blocks up to the server's + // default wait interval and surfaces job errors without returning JSON rows. + 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}), + }, + (err, _rows, _nextQuery, res) => { + if (err) { + stream.destroy(err); + return; + } + if (stream._ended || stream.destroyed) { + return; + } + if (res?.jobComplete === false) { + poll(); + return; + } + onComplete(); + }, + ); + }; + poll(); +} + /** * Executes a query and returns a stream of `apache-arrow` `RecordBatch` objects. * @@ -394,7 +585,7 @@ function streamReadRows( */ export function createQueryArrowStream_( bq: BigQuery, - query?: Query | string, + query?: ArrowQuery | string, ): ArrowRecordBatchStream { const stream = new ResourceStream( {}, @@ -425,7 +616,8 @@ export function createQueryArrowStream_( return; } - const queryObj: Query = typeof query === 'string' ? {query} : {...query}; + const queryObj: ArrowQuery = + typeof query === 'string' ? {query} : {...query}; if (queryObj.writeDisposition === 'WRITE_APPEND') { stream.destroy( @@ -452,17 +644,59 @@ export function createQueryArrowStream_( let queryReq: bigquery.IQueryRequest | undefined; try { - queryReq = bq.buildQueryRequest_(queryObj, {}); + queryReq = bq.buildQueryRequest_(queryObj, { + timeoutMs: queryObj.timeoutMs, + }); } catch (err) { // For example, invalid query parameters. stream.destroy(err as Error); return; } if (!queryReq) { - // TODO: Route queries that cannot use jobs.query through jobs.insert (createQueryJob). - stream.destroy( - 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. 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; + try { + bq.createQueryJob(jobQuery, (err, job, resp) => { + if (err || !job) { + stream.destroy(err ?? new Error('jobs.insert returned no job.')); + return; + } + + 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) { + emitZeroRowBatchAndEnd(arrow, stream); + return; + } + + // Wait for the job to finish, then read all rows from streams/_default at offset = 0. + waitForJobCompletion( + job, + stream, + () => { + streamDefaultJobStream( + ctx, + resp?.jobReference, + 0, + undefined, + job, + ); + }, + timeoutMs, + ); + }); + } catch (err) { + // For example, a destination that isn't a Table, or invalid parameters. + stream.destroy(err as Error); + } return; } @@ -487,10 +721,27 @@ export function createQueryArrowStream_( } if (res?.jobComplete === false) { - // TODO: Wait for unfinished jobs before streaming Arrow results. - stream.destroy( - new Error('Query did not complete within the jobs.query timeout.'), - ); + // Match BigQuery#query: if the caller provided a non-zero timeoutMs, + // fail immediately instead of polling jobs.getQueryResults. + if (queryReq.timeoutMs) { + stream.destroy( + new Error( + `The query did not complete before ${queryReq.timeoutMs}ms`, + ), + ); + return; + } + if (!res.jobReference || !job) { + stream.destroy( + new Error( + 'Cannot wait for an incomplete query without a jobReference.', + ), + ); + return; + } + waitForJobCompletion(job, stream, () => { + streamDefaultJobStream(ctx, res.jobReference, 0, undefined, job); + }); return; } diff --git a/handwritten/bigquery/src/bigquery.ts b/handwritten/bigquery/src/bigquery.ts index 76cb1d0645e..2fc2e5bed7b 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 314ec9f5a97..bd47e961003 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 2614a19cbb9..e801ee4c039 100644 --- a/handwritten/bigquery/system-test/arrow.ts +++ b/handwritten/bigquery/system-test/arrow.ts @@ -398,4 +398,238 @@ 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/, + ); + }); + + 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/system-test/high_precision_timestamp.ts b/handwritten/bigquery/system-test/high_precision_timestamp.ts index 17c32320980..658be8d41ae 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 974956a8917..559787c5c22 100644 --- a/handwritten/bigquery/test/arrow.ts +++ b/handwritten/bigquery/test/arrow.ts @@ -17,10 +17,19 @@ 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, + JobCallback, + JobOptions, + Query, + QueryResultsOptions, + QueryRowsCallback, +} from '../src'; import { BigQueryReadClient, checkArrowDependencies, + CreateReadSessionRequest, getStorageReadClient, StorageModule, } from '../src/arrow'; @@ -107,11 +116,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, { @@ -124,6 +140,50 @@ function stubStorageReadClient( }); } +/** + * Stubs `job.getQueryResults` for callback-based polling tests. + */ +function stubJobGetQueryResults( + job: Job, + impl: (opts: QueryResultsOptions, cb: QueryRowsCallback) => void, +): void { + job.getQueryResults = impl as Job['getQueryResults']; +} + +/** + * Stubs `job.getMetadata` (`jobs.get`) to return the given job resource. + */ +function stubJobGetMetadata( + job: Job, + outcome: {err?: Error; metadata?: bigquery.IJob}, +): void { + job.getMetadata = (( + cb: (err: Error | null, metadata?: bigquery.IJob) => void, + ) => { + process.nextTick(() => cb(outcome.err || null, outcome.metadata)); + }) as Job['getMetadata']; +} + +/** + * 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` for callback-based `jobs.insert` tests. + */ +function stubCreateQueryJob( + bq: BigQuery, + impl: (opts: Query | string, cb: JobCallback) => void, +): void { + bq.createQueryJob = impl as BigQuery['createQueryJob']; +} + describe('Arrow Support', () => { let originalRequire: typeof Module.prototype.require; @@ -420,16 +480,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( @@ -439,16 +489,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)', () => { @@ -471,6 +511,9 @@ describe('Arrow Support', () => { readRows(): never { throw new Error('not called'); } + createReadSession(): never { + throw new Error('not called'); + } }, }; @@ -714,6 +757,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, @@ -867,4 +913,700 @@ 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); + 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; + }, + }); + + 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: () => { + 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; + }, + }); + + // 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: () => { + const grpcStream = new EventEmitter(); + process.nextTick(() => { + grpcStream.emit('end'); + }); + return grpcStream; + }, + }); + + 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: (req, cb) => { + sessionReqs.push(req); + process.nextTick(() => { + cb(null, { + arrowSchema: { + serializedSchema: Buffer.from( + SAMPLE_PAYLOAD.serializedSchema, + 'base64', + ), + }, + }); + }); + }, + readRows: () => { + const grpcStream = new EventEmitter(); + process.nextTick(() => { + grpcStream.emit('end'); + }); + return grpcStream; + }, + }); + + 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: () => { + const grpcStream = new EventEmitter(); + process.nextTick(() => { + grpcStream.emit('end'); + }); + return grpcStream; + }, + }); + + // 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/, + ); + }); + }); + + 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/, + ); + }); + }); });