fix: throw timeout error when using jobs.query (#1402)

alvarowolfx · web-flow · commit cf962a5a94ec · 2024-09-23T12:38:17.000-04:00
diff --git a/src/bigquery.ts b/src/bigquery.ts
@@ -2206,7 +2206,7 @@ export class BigQuery extends Service {
     this.runJobsQuery(queryReq, (err, job, res) => {
       this.trace_('[runJobsQuery callback]: ', query, err, job, res);
       if (err) {
-        (callback as SimpleQueryRowsCallback)(err, null, res);
+        (callback as SimpleQueryRowsCallback)(err, null, job);
         return;
       }
 
@@ -2232,6 +2232,14 @@ export class BigQuery extends Service {
         job!.getQueryResults(options, callback as QueryRowsCallback);
         return;
       }
+      // If timeout override was provided, return error.
+      if (queryReq.timeoutMs) {
+        const err = new Error(
+          `The query did not complete before ${queryReq.timeoutMs}ms`
+        );
+        (callback as SimpleQueryRowsCallback)(err, null, job);
+        return;
+      }
       delete options.timeoutMs;
       this.trace_('[runJobsQuery] job not complete');
       job!.getQueryResults(options, callback as QueryRowsCallback);
diff --git a/system-test/bigquery.ts b/system-test/bigquery.ts
@@ -369,6 +369,51 @@ describe('BigQuery', () => {
     assert.strictEqual(res.totalRows, '100');
   });
 
+  describe('Query timeout', () => {
+    let longRunningQuery = '';
+
+    beforeEach(() => {
+      const tableId = generateName('table');
+      longRunningQuery = `CREATE TABLE ${dataset.projectId}.${dataset.id}.${tableId} AS SELECT num FROM UNNEST(GENERATE_ARRAY(1,1000000)) as num`;
+    });
+
+    it('should throw a timeout error with jobs.query', async () => {
+      const qOpts: QueryOptions = {
+        timeoutMs: 1000,
+      };
+      let foundError: Error | null = null;
+      try {
+        await bigquery.query(longRunningQuery, qOpts);
+      } catch (error: unknown) {
+        foundError = error as Error;
+      }
+      assert.notEqual(foundError, null);
+      assert.equal(
+        foundError?.message,
+        'The query did not complete before 1000ms'
+      );
+    });
+
+    it('should throw a timeout error without jobs.query', async () => {
+      const jobId = generateName('job');
+      const qOpts: QueryOptions = {
+        job: bigquery.job(jobId),
+        timeoutMs: 1000,
+      };
+      let foundError: Error | null = null;
+      try {
+        await bigquery.query(longRunningQuery, qOpts);
+      } catch (error: unknown) {
+        foundError = error as Error;
+      }
+      assert.notEqual(foundError, null);
+      assert.equal(
+        foundError?.message,
+        'The query did not complete before 1000ms'
+      );
+    });
+  });
+
   it('should allow querying in series', done => {
     bigquery.query(
       query,
diff --git a/test/bigquery.ts b/test/bigquery.ts
@@ -3108,7 +3108,7 @@ describe('BigQuery', () => {
       const error = new Error('err');
 
       bq.runJobsQuery = (query: {}, callback: Function) => {
-        callback(error, null, FAKE_RESPONSE);
+        callback(error, FAKE_RESPONSE, {});
       };
 
       bq.query(QUERY_STRING, (err: Error, rows: {}, resp: {}) => {
@@ -3119,6 +3119,31 @@ describe('BigQuery', () => {
       });
     });
 
+    it('should return throw error when jobs.query times out', done => {
+      const fakeJob = {};
+
+      bq.runJobsQuery = (query: {}, callback: Function) => {
+        callback(null, fakeJob, {
+          queryId: uuid.v1(),
+          jobComplete: false,
+        });
+      };
+
+      bq.query(
+        QUERY_STRING,
+        {timeoutMs: 1000},
+        (err: Error, rows: {}, resp: {}) => {
+          assert.strictEqual(
+            err.message,
+            'The query did not complete before 1000ms'
+          );
+          assert.strictEqual(rows, null);
+          assert.strictEqual(resp, fakeJob);
+          done();
+        }
+      );
+    });
+
     it('should exit early if dryRun is set', done => {
       const options = {
         query: QUERY_STRING,
