[Enhance] enhance fetch worksheet by title parsing, improve code readability

.gitignore

@@ -18,6 +18,8 @@ _build
 .vscode
 .env
 
+.ruff_cache
+
 dump.rdb
 
 node_modules

redash/query_runner/google_spreadsheets.py

@@ -1,5 +1,4 @@
 import logging
-import re
 from base64 import b64decode
 
 from dateutil import parser
@@ -99,20 +98,35 @@ def __init__(self, worksheet_title):
 
 
 def parse_query(query):
-    values = query.split("|")
-    key = values[0]  # key of the spreadsheet
-    worksheet_num_or_title = 0  # 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_num_or_title = s[1:-1]
+    parts = query.split("|")
+    key = parts[0].strip()
+
+    if len(parts) < 2:
+        return query, 0, None
+
+    worksheet_num = None
+    worksheet_title = None
+    value = parts[1].strip()
+    if value:
+        # Check if the value is a quoted string
+        if value[0] == '"' and value[-1] == '"':
+            # Extract the worksheet title without quotes
+            worksheet_title = value[1:-1]
+        elif value.isdigit():
+            # Convert the value to an integer if it's a number
+            parsed_value = int(value)
+            # make is zero based indexing
+            if parsed_value > 0:
+                parsed_value -= 1
             else:
-                # if spreadsheet contains more than one worksheet - this is the number of it
-                worksheet_num_or_title = int(s)
-
-    return key, worksheet_num_or_title
+                parsed_value = 0
+            worksheet_num = parsed_value
+        else:
+            raise ValueError(f"Invalid value: {value}")
+    # If both worksheet_title and worksheet_num are not available, set worksheet_num default value to 0
+    if worksheet_title is None and worksheet_num is None:
+        worksheet_num = 0
+    return key, worksheet_num, worksheet_title
 
 
 def parse_worksheet(worksheet):
@@ -132,20 +146,18 @@ def parse_worksheet(worksheet):
     return data
 
 
-def parse_spreadsheet(spreadsheet, worksheet_num_or_title):
+def parse_spreadsheet(spreadsheet, worksheet_num, worksheet_title):
     worksheet = None
-    if isinstance(worksheet_num_or_title, int):
-        worksheet = spreadsheet.get_worksheet_by_index(worksheet_num_or_title)
+    if worksheet_title:
+        worksheet = spreadsheet.get_worksheet_by_title(worksheet_title)
         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 WorksheetNotFoundByTitleError(worksheet_title)
+    else:
+        worksheet = spreadsheet.get_worksheet_by_index(worksheet_num)
         if worksheet is None:
-            raise WorksheetNotFoundByTitleError(worksheet_num_or_title)
-
+            worksheet_count = len(spreadsheet.worksheets())
+            raise WorksheetNotFoundError(worksheet_num, worksheet_count)
     worksheet_values = worksheet.get_all_values()
-
     return parse_worksheet(worksheet_values)
 
 
@@ -245,7 +257,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_num, worksheet_title = parse_query(query)
 
         try:
             spreadsheet_service = self._get_spreadsheet_service()
@@ -255,7 +267,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_num, worksheet_title)
 
             return data, None
         except gspread.SpreadsheetNotFound:

tests/query_runner/test_google_spreadsheets.py

@@ -57,13 +57,13 @@ def test_returns_meaningful_error_for_missing_worksheet(self):
 
         spreadsheet.worksheets = MagicMock(return_value=[])
         spreadsheet.get_worksheet_by_index = MagicMock(return_value=None)
-        self.assertRaises(WorksheetNotFoundError, parse_spreadsheet, spreadsheet, 0)
+        self.assertRaises(WorksheetNotFoundError, parse_spreadsheet, spreadsheet, 0, None)
 
     def test_returns_meaningful_error_for_missing_worksheet_by_title(self):
         spreadsheet = MagicMock()
 
         spreadsheet.get_worksheet_by_title = MagicMock(return_value=None)
-        self.assertRaises(WorksheetNotFoundByTitleError, parse_spreadsheet, spreadsheet, "a")
+        self.assertRaises(WorksheetNotFoundByTitleError, parse_spreadsheet, spreadsheet, None, "a")
 
 
 empty_worksheet = []
@@ -111,36 +111,36 @@ def test_parse_worksheet_with_duplicate_column_names(self):
 class TestParseQuery(TestCase):
     def test_parse_query(self):
         parsed = parse_query("key|0")
-        self.assertEqual(("key", 0), parsed)
+        self.assertEqual(("key", 0, None), parsed)
 
     def test_parse_query_ignored(self):
         parsed = parse_query("key")
-        self.assertEqual(("key", 0), parsed)
+        self.assertEqual(("key", 0, None), parsed)
 
         parsed = parse_query("key|")
-        self.assertEqual(("key", 0), parsed)
+        self.assertEqual(("key", 0, None), parsed)
 
         parsed = parse_query("key|1|")
-        self.assertEqual(("key", 0), parsed)
+        self.assertEqual(("key", 0, None), parsed)
 
     def test_parse_query_title(self):
         parsed = parse_query('key|""')
-        self.assertEqual(("key", ""), parsed)
+        self.assertEqual(("key", None, ""), parsed)
 
         parsed = parse_query('key|"1"')
-        self.assertEqual(("key", "1"), parsed)
+        self.assertEqual(("key", None, "1"), parsed)
 
         parsed = parse_query('key|"abc"')
-        self.assertEqual(("key", "abc"), parsed)
+        self.assertEqual(("key", None, "abc"), parsed)
 
         parsed = parse_query('key|"あ"')
-        self.assertEqual(("key", "あ"), parsed)
+        self.assertEqual(("key", None, "あ"), parsed)
 
         parsed = parse_query('key|"1""')
-        self.assertEqual(("key", '1"'), parsed)
+        self.assertEqual(("key", None, '1"'), parsed)
 
         parsed = parse_query('key|""')
-        self.assertEqual(("key", ""), parsed)
+        self.assertEqual(("key", None, ""), parsed)
 
     def test_parse_query_failed(self):
         self.assertRaises(ValueError, parse_query, "key|0x01")