diff --git a/redash/query_runner/python.py b/redash/query_runner/python.py index 830d120c5..400d0f18d 100644 --- a/redash/query_runner/python.py +++ b/redash/query_runner/python.py @@ -298,6 +298,25 @@ class Python(BaseQueryRunner): def test_connection(self): pass + def validate_result(self, result): + """Validate the result after executing the query. + + Parameters: + :result dict: The result dict. + """ + if not result: + raise Exception("local variable `result` should not be empty.") + if not isinstance(result, dict): + raise Exception("local variable `result` should be of type `dict`.") + if "rows" not in result: + raise Exception("Missing `rows` field in `result` dict.") + if "columns" not in result: + raise Exception("Missing `columns` field in `result` dict.") + if not isinstance(result["rows"], list): + raise Exception("`rows` field should be of type `list`.") + if not isinstance(result["columns"], list): + raise Exception("`columns` field should be of type `list`.") + def run_query(self, query, user): self._current_user = user @@ -352,6 +371,7 @@ class Python(BaseQueryRunner): exec(code, restricted_globals, self._script_locals) result = self._script_locals["result"] + self.validate_result(result) result["log"] = self._custom_print.lines json_data = json_dumps(result) except Exception as e: diff --git a/tests/query_runner/test_python.py b/tests/query_runner/test_python.py index 283992451..66e4afc47 100644 --- a/tests/query_runner/test_python.py +++ b/tests/query_runner/test_python.py @@ -20,19 +20,44 @@ class TestPythonQueryRunner(TestCase): def test_empty_result(self): query_string = "result={}" result = self.python.run_query(query_string, "user") - self.assertEqual(result[0], '{"log": []}') + self.assertEqual(result[0], None) - def test_invalidate_result_type_string(self): + def test_none_result(self): + query_string = "result=None" + result = self.python.run_query(query_string, "user") + self.assertEqual(result[0], None) + + def test_invalid_result_type_string(self): query_string = "result='string'" result = self.python.run_query(query_string, "user") self.assertEqual(result[0], None) - def test_invalidate_result_type_int(self): + def test_invalid_result_type_int(self): query_string = "result=100" result = self.python.run_query(query_string, "user") self.assertEqual(result[0], None) - def test_validate_result_type(self): + def test_invalid_result_missing_rows(self): + query_string = "result={'columns': []}" + result = self.python.run_query(query_string, "user") + self.assertEqual(result[0], None) + + def test_invalid_result_not_list_rows(self): + query_string = "result={'rows': {}, 'columns': []}" + result = self.python.run_query(query_string, "user") + self.assertEqual(result[0], None) + + def test_invalid_result_missing_columns(self): + query_string = "result={'rows': []}" + result = self.python.run_query(query_string, "user") + self.assertEqual(result[0], None) + + def test_invalid_result_not_list_columns(self): + query_string = "result={'rows': [], 'columns': {}}" + result = self.python.run_query(query_string, "user") + self.assertEqual(result[0], None) + + def test_valid_result_type(self): query_string = ( "result=" '{"columns": [{"name": "col1", "type": TYPE_STRING},' @@ -51,7 +76,7 @@ class TestPythonQueryRunner(TestCase): ) @mock.patch("datetime.datetime") - def test_validate_result_type_with_print(self, mock_dt): + def test_valid_result_type_with_print(self, mock_dt): mock_dt.utcnow = mock.Mock(return_value=datetime(1901, 12, 21)) query_string = ( 'print("test")\n'