diff --git a/.gitignore b/.gitignore index b324689c96..1f5b58930e 100644 --- a/.gitignore +++ b/.gitignore @@ -18,6 +18,8 @@ _build .vscode .env +.ruff_cache + dump.rdb node_modules diff --git a/redash/query_runner/google_spreadsheets.py b/redash/query_runner/google_spreadsheets.py index 6ea9757c4e..0ae75fcbc1 100644 --- a/redash/query_runner/google_spreadsheets.py +++ b/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: diff --git a/tests/query_runner/test_google_spreadsheets.py b/tests/query_runner/test_google_spreadsheets.py index db16d8fc3d..bc9ef0f256 100644 --- a/tests/query_runner/test_google_spreadsheets.py +++ b/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")