From 15bb133c895f882bde944f4189e94e51b63a0087 Mon Sep 17 00:00:00 2001
From: Alvaro Viebrantz <aviebrantz@google.com>
Date: Tue, 10 Sep 2024 16:41:25 -0400
Subject: [PATCH 1/2] fix: throw timeout error when using jobs.query

---
 src/bigquery.ts         | 10 +++++++++-
 system-test/bigquery.ts | 37 +++++++++++++++++++++++++++++++++++++
 test/bigquery.ts        | 27 ++++++++++++++++++++++++++-
 3 files changed, 72 insertions(+), 2 deletions(-)

diff --git a/src/bigquery.ts b/src/bigquery.ts
index fb6bd5ea..238b34c3 100644
--- 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
index 971c0ed5..316a2a39 100644
--- a/system-test/bigquery.ts
+++ b/system-test/bigquery.ts
@@ -369,6 +369,43 @@ 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);
+    });
+
+    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);
+    });
+  });
+
   it('should allow querying in series', done => {
     bigquery.query(
       query,
diff --git a/test/bigquery.ts b/test/bigquery.ts
index 9e6ecca1..e1539897 100644
--- 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,

From b0c03f242bb997daa07b5ee04f27352d446efa2e Mon Sep 17 00:00:00 2001
From: Alvaro Viebrantz <aviebrantz@google.com>
Date: Thu, 19 Sep 2024 13:49:34 -0400
Subject: [PATCH 2/2] test: improve error assertion

---
 system-test/bigquery.ts | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/system-test/bigquery.ts b/system-test/bigquery.ts
index 316a2a39..1f6e57a4 100644
--- a/system-test/bigquery.ts
+++ b/system-test/bigquery.ts
@@ -388,6 +388,10 @@ describe('BigQuery', () => {
         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 () => {
@@ -403,6 +407,10 @@ describe('BigQuery', () => {
         foundError = error as Error;
       }
       assert.notEqual(foundError, null);
+      assert.equal(
+        foundError?.message,
+        'The query did not complete before 1000ms'
+      );
     });
   });