Update google_spreadsheets.py

redash/query_runner/google_spreadsheets.py

@@ -100,28 +100,20 @@ def __init__(self, worksheet_title):
 def parse_query(query):
     values = query.split("|")
     key = values[0]  # key of the spreadsheet
-    worksheet_num = 0  # A default value for when a number of inputs is invalid
-    worksheet_title = '' # No title initially
+    worksheet_pointer = ' 1'  # A default value for when a number of inputs is invalid
     if len(values) == 2:
         s = values[1].strip()
         if len(s) > 0:
             if re.match(r"^\"(.*?)\"$", s):
                 # A string quoted by " means a title of worksheet
-                worksheet_title = s[1:-1]
-            elif re.match(r"^\d+$", s):
-                # if spreadsheet contains more than one worksheet - this is the number of it
-                worksheet_num = int(s)
+                worksheet_pointer = s[1:-1]
             else:
-                # To match for non-quoted one-word file name
-                if len(s.split()) == 1 :
-                    worksheet_title = s
-            if worksheet_title != '':
-                return key, worksheet_title
-            else:
-                return key, worksheet_num
-
-        else:
-            return key, worksheet_num
+                if re.match(r"^\d+$", s.split()[0]):
+                    # index is initialized with a space to differentiate from the title of worksheet
+                    worksheet_pointer = ' '+s.split()[0]
+                else:
+                    worksheet_pointer = s.split()[0]
+    return key, worksheet_pointer
 
 
 def parse_worksheet(worksheet):
@@ -141,17 +133,17 @@ def parse_worksheet(worksheet):
     return data
 
 
-def parse_spreadsheet(spreadsheet, worksheet_num_or_title):
+def parse_spreadsheet(spreadsheet, worksheet_pointer):
     worksheet = None
-    if isinstance(worksheet_num_or_title, int):
-        worksheet = spreadsheet.get_worksheet_by_index(worksheet_num_or_title)
+    if worksheet_pointer[0] == ' ':  # An index of worksheet
+        worksheet = spreadsheet.get_worksheet_by_index(int(worksheet_pointer))
         if worksheet is None:
             worksheet_count = len(spreadsheet.worksheets())
-            raise WorksheetNotFoundError(worksheet_num_or_title, worksheet_count)
-    elif isinstance(worksheet_num_or_title, str):
-        worksheet = spreadsheet.get_worksheet_by_title(worksheet_num_or_title)
+            raise WorksheetNotFoundError(int(worksheet_pointer), worksheet_count)
+    elif worksheet_pointer[0] == ' ':  # A title of worksheet
+        worksheet = spreadsheet.get_worksheet_by_title(worksheet_pointer)
         if worksheet is None:
-            raise WorksheetNotFoundByTitleError(worksheet_num_or_title)
+            raise WorksheetNotFoundByTitleError(worksheet_pointer)
 
     worksheet_values = worksheet.get_all_values()
 
@@ -250,7 +242,7 @@ def test_connection(self):
 
     def run_query(self, query, user):
         logger.debug("Spreadsheet is about to execute query: %s", query)
-        key, worksheet_num_or_title = parse_query(query)
+        key, worksheet_pointer = parse_query(query)
 
         try:
             spreadsheet_service = self._get_spreadsheet_service()
@@ -260,7 +252,7 @@ def run_query(self, query, user):
             else:
                 spreadsheet = spreadsheet_service.open_by_key(key)
 
-            data = parse_spreadsheet(SpreadsheetWrapper(spreadsheet), worksheet_num_or_title)
+            data = parse_spreadsheet(SpreadsheetWrapper(spreadsheet), worksheet_pointer)
 
             return json_dumps(data), None
         except gspread.SpreadsheetNotFound: