feat(bigquery): support jobs.insert in createQueryArrowStream

handwritten/bigquery/src/arrow.ts

@@ -93,9 +93,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;
 };
@@ -533,14 +533,17 @@ function streamReadRows(
 
 /**
  * Waits for an unfinished BigQuery job to complete using `job.getQueryResults`
- * without fetching any `STRUCT_ENCODING` rows (`maxResults: 0`).
+ * 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) {
@@ -552,6 +555,9 @@ function waitForJobCompletion(
       {
         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) {
@@ -647,10 +653,50 @@ export function createQueryArrowStream_(
       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;
     }
 

handwritten/bigquery/system-test/arrow.ts

@@ -538,4 +538,98 @@ describe('Arrow Query Results', () => {
       /The query did not complete before 1ms/,
     );
   });
+
+  it('should execute queries with jobs.insert-only options via jobs.insert and stream Arrow results from streams/_default at offset = 0', async () => {
+    const sql =
+      'SELECT x, x * 2 AS doubled FROM UNNEST(GENERATE_ARRAY(1, 25)) AS x ORDER BY x';
+    let emittedJobId: string | undefined;
+
+    const stream = bigquery.createQueryArrowStream({
+      query: sql,
+      jobTimeoutMs: 60000,
+      location: 'EU',
+      useQueryCache: false,
+    });
+    stream.on('job', job => {
+      emittedJobId = job.id;
+    });
+
+    const batches: arrow.RecordBatch[] = await stream.toArray();
+    const [expectedRows] = await bigquery.query({
+      query: sql,
+      jobTimeoutMs: 60000,
+      location: 'EU',
+      useQueryCache: false,
+    });
+
+    // Verify the query ran via jobs.insert (no queryId, jobReference in EU) and matches STRUCT_ENCODING rows.
+    assert.strictEqual(stream.queryId, undefined);
+    assert.ok(stream.jobReference?.jobId);
+    assert.strictEqual(stream.jobReference?.location, 'EU');
+    assert.strictEqual(emittedJobId, stream.jobReference?.jobId);
+
+    const table = new arrow.Table(batches);
+    assert.strictEqual(table.numRows, expectedRows.length);
+    for (let i = 0; i < expectedRows.length; i++) {
+      assert.strictEqual(
+        Number(table.getChild('x')!.get(i)),
+        expectedRows[i].x,
+      );
+      assert.strictEqual(
+        Number(table.getChild('doubled')!.get(i)),
+        expectedRows[i].doubled,
+      );
+    }
+  });
+
+  it('should raise the same error as query() when a jobs.insert query fails during execution', async () => {
+    const failingQuery = {
+      query: 'SELECT 1 / 0',
+      jobTimeoutMs: 60000,
+      useQueryCache: false,
+    };
+
+    let expectedError: (Error & {code?: number}) | undefined;
+    try {
+      await bigquery.query(failingQuery);
+    } catch (err) {
+      expectedError = err as Error & {code?: number};
+    }
+    assert.ok(expectedError, 'Expected bigquery.query() to reject');
+
+    let actualError: (Error & {code?: number}) | undefined;
+    try {
+      await bigquery.createQueryArrowStream(failingQuery).toArray();
+    } catch (err) {
+      actualError = err as Error & {code?: number};
+    }
+
+    // Verify jobs.insert Arrow queries surface the same error code and message as query().
+    assert.ok(actualError, 'Expected createQueryArrowStream() to reject');
+    assert.strictEqual(actualError.code, expectedError.code);
+    assert.match(actualError.message, /division by zero/i);
+  });
+
+  it('should emit a 0-row RecordBatch preserving the schema for an empty result read via jobs.insert and streams/_default', async () => {
+    // jobs.insert returns no rows inline, so the schema of an empty result
+    // has to come from ReadRows on streams/_default.
+    const stream = bigquery.createQueryArrowStream({
+      query:
+        'SELECT x, CAST(x AS STRING) AS s FROM UNNEST(CAST([] AS ARRAY<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));
+  });
 });