Fix python query runner crashing the GUI due to missing 'rows' and/or… (

#5749)

* Fix python query runner crashing the GUI due to missing 'rows' and/or 'columns' in the JSON returned data.

* Fix typo of previous commit.

* Throw exception when python query runner has invalid result.

* Update test_python.py
---------

Co-authored-by: YuhengChen <[email protected]>

redash/query_runner/python.py

@@ -298,6 +298,25 @@ def get_current_user(self):
     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 @@ def run_query(self, query, user):
             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:

tests/query_runner/test_python.py

@@ -20,19 +20,44 @@ def test_print_in_query_string_success(self, mock_dt):
     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_none_result(self):
+        query_string = "result=None"
+        result = self.python.run_query(query_string, "user")
+        self.assertEqual(result[0], None)
 
-    def test_invalidate_result_type_string(self):
+    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 @@ def test_validate_result_type(self):
         )
 
     @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'