diff --git a/handwritten/bigquery/src/arrow.ts b/handwritten/bigquery/src/arrow.ts index 59ccf505fd3..c108fe20ba5 100644 --- a/handwritten/bigquery/src/arrow.ts +++ b/handwritten/bigquery/src/arrow.ts @@ -17,6 +17,7 @@ import {ResourceStream} from '@google-cloud/paginator'; import type {RecordBatch, Schema} from 'apache-arrow'; import type {BigQuery, JobsQueryResponse, Query} from './bigquery'; +import type {Job} from './job'; import type bigquery from './types'; type ArrowModule = typeof import('apache-arrow'); @@ -36,6 +37,30 @@ interface ReadRowsStream extends NodeJS.EventEmitter { resume(): void; } +interface IReadSession { + arrowSchema?: { + serializedSchema?: Uint8Array | string | null; + } | null; + streams?: Array<{ + name?: string | null; + }> | null; +} + +/** + * The `CreateReadSession` request that {@link BigQuery#createQueryArrowStream} + * sends. Exported for tests. + * + * @private + */ +export interface CreateReadSessionRequest { + parent: string; + readSession: { + table: string; + dataFormat: 'ARROW'; + }; + maxStreamCount: number; +} + /** * The part of `@google-cloud/bigquery-storage`'s `BigQueryReadClient` that * {@link BigQuery#createQueryArrowStream} uses. Exported for tests. @@ -43,6 +68,9 @@ interface ReadRowsStream extends NodeJS.EventEmitter { * @private */ export interface BigQueryReadClient { + createReadSession( + request: CreateReadSessionRequest, + ): Promise<[IReadSession, ...unknown[]]>; readRows(request: {readStream: string; offset?: number}): ReadRowsStream; } @@ -74,6 +102,21 @@ export type ArrowRecordBatchStream = ResourceStream & { jobReference?: bigquery.IJobReference; }; +/** + * [Beta] Options for {@link BigQuery#createQueryArrowStream}: the + * {@link Query} options plus options that only apply to Arrow queries. + * + * @beta + */ +export type ArrowQuery = Query & { + /** + * How long to wait for `jobs.query` to complete, in milliseconds. If the + * query doesn't complete in time, the stream fails instead of waiting for + * the job. + */ + timeoutMs?: number; +}; + /** * State shared by the helpers that stream rows for one * {@link BigQuery#createQueryArrowStream} call. @@ -264,10 +307,28 @@ export function decodeArrowPayload( : [createEmptyRecordBatch(arrow, reader.schema)]; } +/** + * Pushes `batches` to `stream`, stopping if the consumer destroys it. + * + * @private + */ +function pushBatches( + stream: ArrowRecordBatchStream, + batches: RecordBatch[], +): void { + for (const batch of batches) { + if (stream.destroyed) { + return; + } + stream.push(batch); + } +} + /** * Streams Arrow `RecordBatch`es from a job's sessionless default stream * (`projects/{projectId}/locations/{location}/jobs/{jobId}/streams/_default`) - * starting at `offset`. + * starting at `offset`. Pass the finished `job` when reading from offset 0 so + * a result without rows keeps its schema (see {@link readDestinationSchema}). * * @private */ @@ -276,8 +337,9 @@ async function streamDefaultJobStream( jobRef: bigquery.IJobReference | undefined, offset: number, initialSchema?: Uint8Array | string, + job?: Job, ): Promise { - const {bq} = ctx; + const {bq, arrow, stream} = ctx; const projectId = resolveProjectId(bq, jobRef?.projectId); const jobId = jobRef?.jobId; const location = jobRef?.location || bq.location; @@ -289,16 +351,67 @@ async function streamDefaultJobStream( } const readStream = `projects/${projectId}/locations/${location}/jobs/${jobId}/streams/_default`; - await streamReadRows(ctx, {readStream, offset, initialSchema}); + const {emittedBatch, serializedSchema} = await streamReadRows(ctx, { + readStream, + offset, + initialSchema, + }); + + // When reading from offset = 0 (e.g., after waiting for an unfinished job) + // and the stream finishes with 0 rows, still emit a 0-row RecordBatch. + if (offset !== 0 || emittedBatch || stream.destroyed) { + return; + } + const schema = + serializedSchema ?? (job && (await readDestinationSchema(ctx, job))); + pushBatches(stream, decodeArrowPayload(arrow, schema)); +} + +/** + * Reads the Arrow schema of a finished query job's results from its + * destination table. The default job stream sends no `ReadRowsResponse`, and + * so no schema, for a result without rows, but a read session on the + * destination table still returns it. Statements without a result table + * (e.g., DML) have no destination table, so this resolves to `undefined`. + * + * @private + */ +async function readDestinationSchema( + ctx: ArrowStreamContext, + job: Job, +): Promise { + const {bq, storage} = ctx; + // A script's destination table is only set once the job finishes, so read + // the job's current metadata instead of the response that created it. + const [metadata] = await job.getMetadata(); + const destRef = metadata?.configuration?.query?.destinationTable; + // Skip the read session if the consumer destroyed the stream meanwhile. + if (!destRef || ctx.stream.destroyed) { + return undefined; + } + const [session] = await getStorageReadClient(bq, storage).createReadSession({ + parent: `projects/${resolveProjectId( + bq, + metadata?.jobReference?.projectId, + destRef.projectId, + )}`, + readSession: { + table: `projects/${destRef.projectId}/datasets/${destRef.datasetId}/tables/${destRef.tableId}`, + dataFormat: 'ARROW', + }, + maxStreamCount: 1, + }); + return session?.arrowSchema?.serializedSchema || undefined; } /** * Streams Arrow `RecordBatch`es from a Storage Read API stream starting at * `offset`, pausing `ReadRows` while the consumer is behind. * - * Resolves when `ReadRows` ends or the consumer destroys the stream, and - * rejects on a `ReadRows` or decode error. In every case, the listeners added - * here are removed, and `ReadRows` is cancelled unless it already finished. + * Resolves when `ReadRows` ends or the consumer destroys the stream, with + * whether any batch was emitted and the last Arrow schema seen, and rejects + * on a `ReadRows` or decode error. In every case, the listeners added here + * are removed, and `ReadRows` is cancelled unless it already finished. * * @private */ @@ -309,7 +422,7 @@ async function streamReadRows( offset: number; initialSchema?: Uint8Array | string; }, -): Promise { +): Promise<{emittedBatch: boolean; serializedSchema?: Uint8Array}> { const {bq, arrow, storage, stream} = ctx; const {readStream, offset, initialSchema} = options; @@ -322,6 +435,7 @@ async function streamReadRows( let serializedSchema: Uint8Array | undefined = initialSchema ? toUint8Array(initialSchema) : undefined; + let emittedBatch = false; let settled = false; // Detaches every listener added here and settles the promise. Cancels @@ -348,7 +462,7 @@ async function streamReadRows( if (opts.err) { reject(opts.err); } else { - resolve(); + resolve({emittedBatch, serializedSchema}); } }; @@ -379,6 +493,7 @@ async function streamReadRows( settle({cancel: true}); return; } + emittedBatch = true; if (!stream.push(batch)) { grpcStream.pause(); } @@ -403,6 +518,31 @@ 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`. + * + * @private + */ +async function waitForJobCompletion( + job: Job, + stream: ArrowRecordBatchStream, +): Promise { + while (!stream.destroyed) { + // Set maxResults: 0 so jobs.getQueryResults blocks up to the server's + // default wait interval (about 10s) and surfaces job errors without + // returning JSON rows, so polling again right away doesn't spin. + const [, , res] = await job.getQueryResults({ + autoPaginate: false, + maxResults: 0, + }); + if (res?.jobComplete !== false) { + return; + } + } +} + /** * Executes a query and returns a stream of `apache-arrow` `RecordBatch` objects. * @@ -410,7 +550,7 @@ async function streamReadRows( */ export function createQueryArrowStream_( bq: BigQuery, - query?: Query | string, + query?: ArrowQuery | string, ): ArrowRecordBatchStream { let started = false; // Set once the query starts, so later reads can resume the active ReadRows. @@ -446,7 +586,8 @@ export function createQueryArrowStream_( throw new Error('A SQL query string is required.'); } - const queryObj: Query = typeof query === 'string' ? {query} : {...query}; + const queryObj: ArrowQuery = + typeof query === 'string' ? {query} : {...query}; if (queryObj.writeDisposition === 'WRITE_APPEND') { throw new Error( @@ -466,7 +607,9 @@ export function createQueryArrowStream_( readCtx = ctx; // Throws for an invalid request, for example invalid query parameters. - const queryReq = bq.buildQueryRequest_(queryObj, {}); + const queryReq = bq.buildQueryRequest_(queryObj, { + timeoutMs: queryObj.timeoutMs, + }); if (!queryReq) { // TODO: Route queries that cannot use jobs.query through jobs.insert (createQueryJob). throw new Error('Query configuration cannot be executed via jobs.query.'); @@ -493,8 +636,30 @@ export function createQueryArrowStream_( } if (res?.jobComplete === false) { - // TODO: Wait for unfinished jobs before streaming Arrow results. - throw new Error('Query did not complete within the jobs.query timeout.'); + // jobs.query already waited (about 10s by default) and the query is + // still running, so this depends on how long the query runs, not on + // how many rows it returns. + // Match BigQuery#query: if the caller provided a non-zero timeoutMs, + // fail immediately instead of polling jobs.getQueryResults. + if (queryReq.timeoutMs) { + throw new Error( + `The query did not complete before ${queryReq.timeoutMs}ms`, + ); + } + if (!res.jobReference || !job) { + throw new Error( + 'Cannot wait for an incomplete query without a jobReference.', + ); + } + await waitForJobCompletion(job, stream); + if (stream.destroyed) { + return; + } + // An unfinished jobs.query returns no rows or schema, so read every row + // from offset 0 and pass the job to look up the schema of an empty + // result. + await streamDefaultJobStream(ctx, res.jobReference, 0, undefined, job); + return; } // pageRowCount and totalRows are returned as uint64 strings in REST JSON. @@ -508,17 +673,14 @@ export function createQueryArrowStream_( // Emit the inline first page (or 0-row schema-preserving batch when totalRows is 0) // before resuming remaining rows from streams/_default at offset = pageRowCount. if (!hasRemainingRows || inlineRowsCount > 0n) { - const batches = decodeArrowPayload( - arrow, - res?.arrowSchema?.serializedSchema, - res?.arrowRecordBatch?.serializedRecordBatch, + pushBatches( + stream, + decodeArrowPayload( + arrow, + res?.arrowSchema?.serializedSchema, + res?.arrowRecordBatch?.serializedRecordBatch, + ), ); - for (const batch of batches) { - if (stream.destroyed) { - return; - } - stream.push(batch); - } } if (!hasRemainingRows || stream.destroyed) { diff --git a/handwritten/bigquery/src/bigquery.ts b/handwritten/bigquery/src/bigquery.ts index 5d6e201ec05..3470b7b4c32 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; @@ -395,8 +399,11 @@ export class BigQuery extends Service { * 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|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 @@ -426,7 +433,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..980b1d86d9e 100644 --- a/handwritten/bigquery/system-test/arrow.ts +++ b/handwritten/bigquery/system-test/arrow.ts @@ -398,4 +398,144 @@ describe('Arrow Query Results', () => { }, ); }); + + it('should wait for an unfinished jobs.query via jobs.getQueryResults (maxResults = 0) and stream all rows from streams/_default at offset = 0', async () => { + const origRequest = bigquery.request.bind(bigquery); + let getQueryResultsCalls = 0; + + // Force POST /queries to return jobComplete: false by setting a 1ms server wait on the wire + // while leaving the caller's Query object without timeoutMs so createQueryArrowStream waits. + bigquery.request = (reqOpts, callback) => { + if (reqOpts.uri === '/queries' && reqOpts.method === 'POST') { + reqOpts.json = {...reqOpts.json, timeoutMs: 1}; + } else if ( + typeof reqOpts.uri === 'string' && + reqOpts.uri.startsWith('/queries/') + ) { + getQueryResultsCalls++; + } + return origRequest(reqOpts, callback); + }; + + try { + const sql = + 'DECLARE n INT64 DEFAULT 50; SELECT x FROM UNNEST(GENERATE_ARRAY(1, n)) AS x ORDER BY x;'; + const batches: arrow.RecordBatch[] = await bigquery + .createQueryArrowStream({query: sql, useQueryCache: false}) + .toArray(); + + // Verify jobs.getQueryResults was polled to wait for completion and all 50 rows streamed from offset = 0. + assert.ok( + getQueryResultsCalls >= 1, + `Expected >= 1 jobs.getQueryResults call, got ${getQueryResultsCalls}`, + ); + const table = new arrow.Table(batches); + assert.strictEqual(table.numRows, 50); + assert.strictEqual(table.getChild('x')!.get(0), 1n); + assert.strictEqual(table.getChild('x')!.get(49), 50n); + } finally { + bigquery.request = origRequest; + } + }); + + it('should keep the schema of an empty result after waiting for an unfinished jobs.query', async () => { + // After the wait, rows come from streams/_default at offset = 0, which + // sends nothing for a result without rows, so the schema has to come from + // the job's destination table. + const origRequest = bigquery.request.bind(bigquery); + + // Force POST /queries to return jobComplete: false so the stream waits + // for the job before reading rows. + bigquery.request = (reqOpts, callback) => { + if (reqOpts.uri === '/queries' && reqOpts.method === 'POST') { + reqOpts.json = {...reqOpts.json, timeoutMs: 1}; + } + return origRequest(reqOpts, callback); + }; + + try { + const sql = `DECLARE n INT64 DEFAULT 0; + SELECT x, CAST(x AS STRING) AS s FROM UNNEST(GENERATE_ARRAY(1, n)) AS x;`; + const batches: arrow.RecordBatch[] = await bigquery + .createQueryArrowStream({query: sql, useQueryCache: false}) + .toArray(); + + // Verify one 0-row RecordBatch with the result's column names and types. + assert.strictEqual(batches.length, 1); + assert.strictEqual(batches[0].numRows, 0); + assert.deepStrictEqual( + batches[0].schema.fields.map(f => f.name), + ['x', 's'], + ); + assert.ok(arrow.DataType.isInt(batches[0].schema.fields[0].type)); + assert.ok(arrow.DataType.isUtf8(batches[0].schema.fields[1].type)); + } finally { + bigquery.request = origRequest; + } + }); + + it('should wait for an unfinished job with a picosecond TIMESTAMP(12) column using default jobs.getQueryResults format options', async function () { + // Guards the job wait against format option regressions: the wait calls + // jobs.getQueryResults with maxResults = 0 and the library's default + // formatOptions (useInt64Timestamp), which must still succeed when the + // result has a picosecond TIMESTAMP(12) column. Picosecond timestamps are + // only enabled in some projects, so this runs only when + // BIGQUERY_PICOSECOND_SUPPORT is set. + if (process.env.BIGQUERY_PICOSECOND_SUPPORT !== 'true') { + this.skip(); + } + const origRequest = bigquery.request.bind(bigquery); + let getQueryResultsCalls = 0; + + // Force POST /queries to return jobComplete: false so the stream waits + // for the job with jobs.getQueryResults before reading rows. + bigquery.request = (reqOpts, callback) => { + if (reqOpts.uri === '/queries' && reqOpts.method === 'POST') { + reqOpts.json = {...reqOpts.json, timeoutMs: 1}; + } else if ( + typeof reqOpts.uri === 'string' && + reqOpts.uri.startsWith('/queries/') + ) { + getQueryResultsCalls++; + } + return origRequest(reqOpts, callback); + }; + + try { + const sql = `DECLARE n INT64 DEFAULT 3; + SELECT x, CAST('2026-01-15 12:34:56.123456789123+00' AS TIMESTAMP(12)) AS ts + FROM UNNEST(GENERATE_ARRAY(1, n)) AS x ORDER BY x;`; + const batches: arrow.RecordBatch[] = await bigquery + .createQueryArrowStream({query: sql, useQueryCache: false}) + .toArray(); + + // Verify the wait ran and every row streamed with an Arrow timestamp. + assert.ok( + getQueryResultsCalls >= 1, + `Expected >= 1 jobs.getQueryResults call, got ${getQueryResultsCalls}`, + ); + const table = new arrow.Table(batches); + assert.strictEqual(table.numRows, 3); + assert.ok(arrow.DataType.isTimestamp(table.schema.fields[1].type)); + } finally { + bigquery.request = origRequest; + } + }); + + it('should reject when a query does not complete within caller-specified timeoutMs', async () => { + const sql = + 'DECLARE n INT64 DEFAULT 50; SELECT x FROM UNNEST(GENERATE_ARRAY(1, n)) AS x ORDER BY x;'; + + // Verify specifying timeoutMs: 1 rejects when jobs.query returns jobComplete = false. + await assert.rejects( + bigquery + .createQueryArrowStream({ + query: sql, + timeoutMs: 1, + useQueryCache: false, + }) + .toArray(), + /The query did not complete before 1ms/, + ); + }); }); diff --git a/handwritten/bigquery/system-test/high_precision_timestamp.ts b/handwritten/bigquery/system-test/high_precision_timestamp.ts index 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 ce5c4293969..bc754606105 100644 --- a/handwritten/bigquery/test/arrow.ts +++ b/handwritten/bigquery/test/arrow.ts @@ -17,10 +17,11 @@ import {EventEmitter} from 'events'; import {describe, it, before, after, beforeEach} from 'mocha'; import * as Module from 'module'; import * as arrow from 'apache-arrow'; -import {BigQuery, Job, Query} from '../src'; +import {BigQuery, Job, Query, QueryResultsOptions} from '../src'; import { BigQueryReadClient, checkArrowDependencies, + CreateReadSessionRequest, getStorageReadClient, StorageModule, } from '../src/arrow'; @@ -107,11 +108,18 @@ type FakeReadRowsStream = EventEmitter & */ function stubStorageReadClient( bq: BigQuery, - client: {readRows: (req: ReadRowsRequest) => FakeReadRowsStream}, + client: { + readRows: (req: ReadRowsRequest) => FakeReadRowsStream; + createReadSession?: BigQueryReadClient['createReadSession']; + }, ): BigQueryReadClient { const noop = () => {}; + const notStubbed = (): never => { + throw new Error('createReadSession is not stubbed'); + }; return getStorageReadClient(bq, { BigQueryReadClient: class { + createReadSession = client.createReadSession ?? notStubbed; readRows(req: ReadRowsRequest) { const grpcStream = client.readRows(req); return Object.assign(grpcStream, { @@ -137,6 +145,63 @@ function assertReadRowsListenersRemoved(grpcStream: EventEmitter): void { assert.doesNotThrow(() => grpcStream.emit('error', cancelled)); } +/** + * Stubs `job.getQueryResults`, which the code awaits, with an async `impl` + * that resolves to `[rows, nextQuery, response]` or rejects. + */ +function stubJobGetQueryResults( + job: Job, + impl: ( + opts: QueryResultsOptions, + ) => Promise<[unknown[], null, bigquery.IGetQueryResultsResponse]>, +): void { + job.getQueryResults = impl as unknown as Job['getQueryResults']; +} + +/** + * Stubs `job.getMetadata` (`jobs.get`) to resolve with the given job resource + * or reject with the given error. + */ +function stubJobGetMetadata( + job: Job, + outcome: {err?: Error; metadata?: bigquery.IJob}, +): void { + job.getMetadata = (async () => { + if (outcome.err) { + throw outcome.err; + } + return [outcome.metadata]; + }) as unknown as Job['getMetadata']; +} + +/** + * Returns a fake ReadRows stream that emits `responses`, then ends, on the + * next tick. + */ +function fakeReadRowsStream(responses: object[] = []): EventEmitter { + const grpcStream = new EventEmitter(); + process.nextTick(() => { + for (const resp of responses) { + grpcStream.emit('data', resp); + } + grpcStream.emit('end'); + }); + return grpcStream; +} + +/** A ReadRowsResponse carrying `SAMPLE_PAYLOAD` as bytes. */ +const SAMPLE_READ_ROWS_RESPONSE = { + arrowSchema: { + serializedSchema: Buffer.from(SAMPLE_PAYLOAD.serializedSchema, 'base64'), + }, + arrowRecordBatch: { + serializedRecordBatch: Buffer.from( + SAMPLE_PAYLOAD.serializedRecordBatch, + 'base64', + ), + }, +}; + describe('Arrow Support', () => { let originalRequire: typeof Module.prototype.require; @@ -491,16 +556,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)', () => { @@ -523,6 +578,9 @@ describe('Arrow Support', () => { readRows(): never { throw new Error('not called'); } + createReadSession(): never { + throw new Error('not called'); + } }, }; @@ -772,6 +830,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, @@ -939,4 +1000,425 @@ 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, async opts => { + getQueryResultsCalls.push(opts); + // Verify ReadRows has not been called while waiting for the job to complete. + assert.strictEqual(capturedReadRowsReqs.length, 0); + const isDone = getQueryResultsCalls.length >= 2; + return [ + [], + null, + {jobComplete: isDone, totalRows: isDone ? '2' : undefined}, + ]; + }); + + stubJobsQuery(bq, { + job: fakeJob, + res: { + jobComplete: false, + jobReference: { + projectId: 'test-project', + jobId: 'job-wait-1', + location: 'US', + }, + }, + }); + + stubStorageReadClient(bq, { + readRows: (req: {readStream?: string; offset?: number}) => { + capturedReadRowsReqs.push(req); + return fakeReadRowsStream([SAMPLE_READ_ROWS_RESPONSE]); + }, + }); + + const batches: arrow.RecordBatch[] = await bq + .createQueryArrowStream('SELECT id, name FROM tbl') + .toArray(); + + // Verify job.getQueryResults was polled twice with maxResults: 0 and autoPaginate: false. + assert.strictEqual(getQueryResultsCalls.length, 2); + for (const opts of getQueryResultsCalls) { + assert.deepStrictEqual(opts, { + autoPaginate: false, + maxResults: 0, + }); + } + + // Verify ReadRows was called at offset = 0 only after the job completed. + assert.deepStrictEqual(capturedReadRowsReqs, [ + { + readStream: + 'projects/test-project/locations/US/jobs/job-wait-1/streams/_default', + offset: 0, + }, + ]); + assert.strictEqual(batches.length, 1); + assert.strictEqual(batches[0].numRows, 2); + }); + + it('should fail with a timeout error when jobs.query returns jobComplete = false and timeoutMs was specified', async () => { + const fakeJob = bq.job('job-timeout-1', {location: 'US'}); + let getQueryResultsCalled = false; + stubJobGetQueryResults(fakeJob, async () => { + getQueryResultsCalled = true; + return [[], null, {jobComplete: 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, async () => { + polls++; + // The consumer gives up while this poll is in flight. + stream.destroy(); + return [[], 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(); + // Wait for the in-flight poll to resolve. + await new Promise(resolve => setImmediate(resolve)); + + // Verify the in-flight poll didn't start another one or call ReadRows. + 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, async opts => { + getQueryResultsCalls.push(opts); + return [[], null, {jobComplete: true}]; + }); + + stubJobsQuery(bq, { + job: fakeJob, + res: { + jobComplete: false, + jobReference: { + projectId: 'test-project', + jobId: 'job-timeout-0', + location: 'US', + }, + }, + }); + + stubStorageReadClient(bq, { + readRows: () => fakeReadRowsStream([SAMPLE_READ_ROWS_RESPONSE]), + }); + + // BigQuery#query treats a falsy timeoutMs as "no override", so 0 must + // not reject with "did not complete before 0ms"; poll and stream instead. + const batches: arrow.RecordBatch[] = await bq + .createQueryArrowStream({query: 'SELECT 1', timeoutMs: 0}) + .toArray(); + assert.deepStrictEqual(getQueryResultsCalls, [ + {autoPaginate: false, maxResults: 0}, + ]); + assert.strictEqual(batches.length, 1); + assert.strictEqual(batches[0].numRows, 2); + }); + + it('should forward job.getQueryResults errors to the stream without calling ReadRows', async () => { + const fakeJob = bq.job('job-err-1', {location: 'US'}); + stubJobGetQueryResults(fakeJob, async () => { + throw 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, async () => [ + [], + null, + {jobComplete: true, totalRows: '0'}, + ]); + // Statements without a result table, like DECLARE, have no destination. + stubJobGetMetadata(fakeJob, {metadata: {configuration: {query: {}}}}); + + stubJobsQuery(bq, { + job: fakeJob, + res: { + jobComplete: false, + jobReference: { + projectId: 'test-project', + jobId: 'job-empty-1', + location: 'US', + }, + }, + }); + + // createReadSession isn't stubbed, so calling it fails the test. + stubStorageReadClient(bq, { + readRows: () => fakeReadRowsStream(), + }); + + const batches: arrow.RecordBatch[] = await bq + .createQueryArrowStream('DECLARE x INT64 DEFAULT 1') + .toArray(); + + // Verify a 0-row RecordBatch is still emitted when streams/_default returns 0 messages at offset = 0. + assert.strictEqual(batches.length, 1); + assert.strictEqual(batches[0].numRows, 0); + assert.strictEqual(batches[0].numCols, 0); + }); + + it('should read the schema from the destination table when ReadRows at offset = 0 ends without rows or a schema', async () => { + // streams/_default sends no ReadRowsResponse for a result without rows, + // so the schema has to come from a read session on the destination table. + const fakeJob = bq.job('job-empty-2', {location: 'US'}); + stubJobGetQueryResults(fakeJob, async () => [ + [], + null, + {jobComplete: true, totalRows: '0'}, + ]); + stubJobGetMetadata(fakeJob, { + metadata: { + jobReference: {projectId: 'test-project', jobId: 'job-empty-2'}, + configuration: { + query: { + destinationTable: { + projectId: 'test-project', + datasetId: '_anon', + tableId: 'anon_table', + }, + }, + }, + }, + }); + + stubJobsQuery(bq, { + job: fakeJob, + res: { + jobComplete: false, + jobReference: { + projectId: 'test-project', + jobId: 'job-empty-2', + location: 'US', + }, + }, + }); + + const sessionReqs: CreateReadSessionRequest[] = []; + stubStorageReadClient(bq, { + createReadSession: async req => { + sessionReqs.push(req); + return [ + { + arrowSchema: { + serializedSchema: Buffer.from( + SAMPLE_PAYLOAD.serializedSchema, + 'base64', + ), + }, + }, + ]; + }, + readRows: () => fakeReadRowsStream(), + }); + + const batches: arrow.RecordBatch[] = await bq + .createQueryArrowStream('SELECT id, name FROM tbl WHERE FALSE') + .toArray(); + + // Verify the 0-row RecordBatch carries the destination table's schema. + assert.strictEqual(batches.length, 1); + assert.strictEqual(batches[0].numRows, 0); + assert.deepStrictEqual( + batches[0].schema.fields.map(f => f.name), + ['id', 'name'], + ); + assert.strictEqual(sessionReqs.length, 1); + assert.strictEqual(sessionReqs[0].parent, 'projects/test-project'); + assert.strictEqual( + sessionReqs[0].readSession.table, + 'projects/test-project/datasets/_anon/tables/anon_table', + ); + assert.strictEqual(sessionReqs[0].maxStreamCount, 1); + }); + + it('should not create a read session if the stream is destroyed while reading the job metadata', async () => { + const fakeJob = bq.job('job-empty-destroy', {location: 'US'}); + const stream = bq.createQueryArrowStream( + 'SELECT id FROM tbl WHERE FALSE', + ); + stubJobGetQueryResults(fakeJob, async () => [ + [], + null, + {jobComplete: true, totalRows: '0'}, + ]); + fakeJob.getMetadata = (async () => { + // The consumer gives up while jobs.get is in flight. + stream.destroy(); + return [ + { + configuration: { + query: { + destinationTable: { + projectId: 'test-project', + datasetId: '_anon', + tableId: 'anon_table', + }, + }, + }, + }, + ]; + }) as unknown as Job['getMetadata']; + stubJobsQuery(bq, { + job: fakeJob, + res: { + jobComplete: false, + jobReference: { + projectId: 'test-project', + jobId: 'job-empty-destroy', + location: 'US', + }, + }, + }); + let readSessions = 0; + stubStorageReadClient(bq, { + createReadSession: async () => { + readSessions++; + return [{}]; + }, + readRows: () => fakeReadRowsStream(), + }); + + stream.resume(); + await new Promise(resolve => stream.once('close', resolve)); + // Let the query finish its work after jobs.get resolves. + await new Promise(resolve => setImmediate(resolve)); + + // Verify no read session was created for the destroyed stream. + assert.strictEqual(readSessions, 0); + }); + + it('should forward jobs.get errors when reading the schema of a result without rows', async () => { + const fakeJob = bq.job('job-empty-3', {location: 'US'}); + stubJobGetQueryResults(fakeJob, async () => [ + [], + null, + {jobComplete: true, totalRows: '0'}, + ]); + stubJobGetMetadata(fakeJob, {err: new Error('jobs.get failed')}); + + stubJobsQuery(bq, { + job: fakeJob, + res: { + jobComplete: false, + jobReference: { + projectId: 'test-project', + jobId: 'job-empty-3', + location: 'US', + }, + }, + }); + + stubStorageReadClient(bq, { + readRows: () => fakeReadRowsStream(), + }); + + // Verify the jobs.get error reaches the stream instead of a schema-less batch. + await assert.rejects( + bq.createQueryArrowStream('SELECT id FROM tbl WHERE FALSE').toArray(), + /jobs.get failed/, + ); + }); + }); });