Merge pull request #1767 from getredash/patches

Google Spreadsheets: add timeout to request.

redash/query_runner/google_spreadsheets.py

@@ -3,6 +3,7 @@ import logging
 from base64 import b64decode
 
 from dateutil import parser
+from requests import Session
 
 from redash.query_runner import *
 from redash.utils import JSONEncoder
@@ -11,6 +12,7 @@ logger = logging.getLogger(__name__)
 
 try:
     import gspread
+    from gspread.httpsession import HTTPSession
     from oauth2client.service_account import ServiceAccountCredentials
     enabled = True
 except ImportError:
@@ -110,7 +112,7 @@ def parse_worksheet(worksheet):
         })
 
     if len(worksheet) > 1:
-        for j, value in enumerate(worksheet[HEADER_INDEX+1]):
+        for j, value in enumerate(worksheet[HEADER_INDEX + 1]):
             columns[j]['type'] = _guess_type(value)
 
     rows = [dict(zip(column_names, _value_eval_list(row))) for row in worksheet[HEADER_INDEX + 1:]]
@@ -130,6 +132,12 @@ def parse_spreadsheet(spreadsheet, worksheet_num):
     return parse_worksheet(worksheet)
 
 
+class TimeoutSession(Session):
+    def request(self, *args, **kwargs):
+        kwargs.setdefault('timeout', 300)
+        return super(TimeoutSession, self).request(*args, **kwargs)
+
+
 class GoogleSpreadsheet(BaseQueryRunner):
     @classmethod
     def annotate_query(cls):
@@ -167,7 +175,11 @@ class GoogleSpreadsheet(BaseQueryRunner):
 
         key = json.loads(b64decode(self.configuration['jsonKeyFile']))
         creds = ServiceAccountCredentials.from_json_keyfile_dict(key, scope)
-        spreadsheetservice = gspread.authorize(creds)
+
+        timeout_session = HTTPSession()
+        timeout_session.requests_session = TimeoutSession()
+        spreadsheetservice = gspread.Client(auth=creds, http_session=timeout_session)
+        spreadsheetservice.login()
         return spreadsheetservice
 
     def test_connection(self):
@@ -176,8 +188,9 @@ class GoogleSpreadsheet(BaseQueryRunner):
     def run_query(self, query, user):
         logger.debug("Spreadsheet is about to execute query: %s", query)
         values = query.split("|")
-        key = values[0] #key of the spreadsheet
-        worksheet_num = 0 if len(values) != 2 else int(values[1])# if spreadsheet contains more than one worksheet - this is the number of it
+        key = values[0]  # key of the spreadsheet
+        worksheet_num = 0 if len(values) != 2 else int(values[1])  # if spreadsheet contains more than one worksheet - this is the number of it
+
         try:
             spreadsheet_service = self._get_spreadsheet_service()
             spreadsheet = spreadsheet_service.open_by_key(key)

requirements_all_ds.txt

@@ -1,5 +1,5 @@
 google-api-python-client==1.5.1
-gspread==0.2.5
+gspread==0.6.2
 impyla==0.10.0
 influxdb==2.7.1
 MySQL-python==1.2.5