feat(bigquery): support jobs.insert in createQueryArrowStream

handwritten/bigquery/src/arrow.ts

@@ -110,9 +110,9 @@ export type ArrowRecordBatchStream = ResourceStream<RecordBatch> & {
  */
 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<void> {
   while (!stream.destroyed) {
     // Set maxResults: 0 so jobs.getQueryResults blocks up to the server's
@@ -536,13 +539,76 @@ 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;
     }
   }
 }
 
+/**
+ * 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<void> {
+  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<void> {
+  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;
     }
 

handwritten/bigquery/src/bigquery.ts

@@ -375,7 +375,7 @@
   private _universeDomain: string;
   private _defaultJobCreationMode: JobCreationMode;
 
   createQueryStream(options?: Query | string): ResourceStream<RowMetadata> {
     // placeholder body, overwritten in constructor
     return new ResourceStream<RowMetadata>({}, () => {});
   }
@@ -391,19 +391,21 @@
    * 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
@@ -439,12 +441,12 @@
     return createQueryArrowStream_(this, options);
   }
 
   getDatasetsStream(options?: GetDatasetsOptions): ResourceStream<Dataset> {
     // placeholder body, overwritten in constructor
     return new ResourceStream<Dataset>({}, () => {});
   }
 
   getJobsStream(options?: GetJobsOptions): ResourceStream<Job> {
     // placeholder body, overwritten in constructor
     return new ResourceStream<Job>({}, () => {});
   }
@@ -1663,7 +1665,7 @@
     const parameterMode = isArray(params) ? 'positional' : 'named';
     const queryParameters: bigquery.IQueryParameter[] = [];
     if (parameterMode === 'named') {
       const namedParams = params as {[param: string]: any};
       for (const namedParameter of Object.getOwnPropertyNames(namedParams)) {
         const value = namedParams[namedParameter];
         let queryParameter;
@@ -2310,7 +2312,7 @@
 
       options = extend({job}, queryOpts, options);
       if (res && res.jobComplete) {
         let rows: any = [];
         if (res.schema && res.rows) {
           try {
             /*

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<INT64>)) 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));
+  });
 });