From ee6dcab362f441113efbb9bb66eebf2a9b276df2 Mon Sep 17 00:00:00 2001 From: Omer Lachish Date: Wed, 4 Mar 2020 22:45:20 +0200 Subject: [PATCH] Cancel BigQuery Queries (#4701) * cancel BigQuery queries when user requests cancellation or when the job times out * create a new bigquery client to flush exising requests --- redash/query_runner/big_query.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/redash/query_runner/big_query.py b/redash/query_runner/big_query.py index eec4b75a4..1b29de7ff 100644 --- a/redash/query_runner/big_query.py +++ b/redash/query_runner/big_query.py @@ -192,12 +192,13 @@ class BigQuery(BaseQueryRunner): project_id = self._get_project_id() job_data = self._get_job_data(query) insert_response = jobs.insert(projectId=project_id, body=job_data).execute() + self.current_job_id = insert_response["jobReference"]["jobId"] current_row = 0 query_reply = _get_query_results( jobs, project_id=project_id, location=self._get_location(), - job_id=insert_response["jobReference"]["jobId"], + job_id=self.current_job_id, start_index=current_row, ) @@ -332,6 +333,15 @@ class BigQuery(BaseQueryRunner): error = json_loads(e.content)["error"]["message"] else: error = e.content + except (KeyboardInterrupt, InterruptException, JobTimeoutException): + if self.current_job_id: + self._get_bigquery_service().jobs().cancel( + projectId=self._get_project_id(), + jobId=self.current_job_id, + location=self._get_location(), + ).execute() + + raise return json_data, error