feat: custom error SQL Lab timeout (#15342)

* feat: custom error SQL Lab timeout

* Update test

docs/src/pages/docs/Miscellaneous/issue_codes.mdx

@@ -247,3 +247,19 @@ CVAS (create view as select) query is not a SELECT statement.
 ```
 
 When running a CVAS (create view as select) the query should be a SELECT statement. Please make sure the query has a single statement and it's a SELECT statement.
+
+## Issue 1026
+
+```
+Query is too complex and takes too long to run.
+```
+
+The submitted query might be too complex to run under the time limit defined by your Superset administrator. Please double check your query and verify if it can be optimized. Alternatively, contact your administrator to increase the timeout period.
+
+## Issue 1027
+
+```
+The database is currently running too many queries.
+```
+
+The database might be under heavy load, running too many queries. Please try again later, or contact an administrator for further assistance.

superset-frontend/src/components/ErrorMessage/types.ts

@@ -61,6 +61,7 @@ export const ErrorTypeEnum = {
   DML_NOT_ALLOWED_ERROR: 'DML_NOT_ALLOWED_ERROR',
   INVALID_CTAS_QUERY_ERROR: 'INVALID_CTAS_QUERY_ERROR',
   INVALID_CVAS_QUERY_ERROR: 'INVALID_CVAS_QUERY_ERROR',
+  SQLLAB_TIMEOUT_ERROR: 'SQLLAB_TIMEOUT_ERROR',
 
   // Generic errors
   GENERIC_COMMAND_ERROR: 'GENERIC_COMMAND_ERROR',

superset-frontend/src/setup/setupErrorMessages.ts

@@ -71,6 +71,10 @@ export default function setupErrorMessages() {
     ErrorTypeEnum.CONNECTION_INVALID_HOSTNAME_ERROR,
     DatabaseErrorMessage,
   );
+  errorMessageComponentRegistry.registerValue(
+    ErrorTypeEnum.SQLLAB_TIMEOUT_ERROR,
+    DatabaseErrorMessage,
+  );
   errorMessageComponentRegistry.registerValue(
     ErrorTypeEnum.CONNECTION_PORT_CLOSED_ERROR,
     DatabaseErrorMessage,

superset/errors.py

@@ -70,6 +70,7 @@ class SupersetErrorType(str, Enum):
     DML_NOT_ALLOWED_ERROR = "DML_NOT_ALLOWED_ERROR"
     INVALID_CTAS_QUERY_ERROR = "INVALID_CTAS_QUERY_ERROR"
     INVALID_CVAS_QUERY_ERROR = "INVALID_CVAS_QUERY_ERROR"
+    SQLLAB_TIMEOUT_ERROR = "SQLLAB_TIMEOUT_ERROR"
 
     # Generic errors
     GENERIC_COMMAND_ERROR = "GENERIC_COMMAND_ERROR"
@@ -295,6 +296,20 @@ ERROR_TYPES_TO_ISSUE_CODES_MAPPING = {
             ),
         },
     ],
+    SupersetErrorType.SQLLAB_TIMEOUT_ERROR: [
+        {
+            "code": 1026,
+            "message": _(
+                "Issue 1026 - Query is too complex and takes too long to run."
+            ),
+        },
+        {
+            "code": 1027,
+            "message": _(
+                "Issue 1027 - The database is currently running too many queries."
+            ),
+        },
+    ],
 }
 
 

superset/sql_lab.py

@@ -28,7 +28,7 @@ import pyarrow as pa
 import simplejson as json
 from celery import Task
 from celery.exceptions import SoftTimeLimitExceeded
-from flask_babel import gettext as __, lazy_gettext as _
+from flask_babel import gettext as __
 from sqlalchemy.orm import Session
 from werkzeug.local import LocalProxy
 
@@ -83,10 +83,6 @@ class SqlLabSecurityException(SqlLabException):
     pass
 
 
-class SqlLabTimeoutException(SqlLabException):
-    pass
-
-
 def handle_query_error(
     ex: Exception,
     query: Query,
@@ -184,15 +180,6 @@ def get_sql_results(  # pylint: disable=too-many-arguments
                 expand_data=expand_data,
                 log_params=log_params,
             )
-        except SoftTimeLimitExceeded as ex:
-            logger.warning("Query %d: Time limit exceeded", query_id)
-            logger.debug("Query %d: %s", query_id, ex)
-            raise SqlLabTimeoutException(
-                _(
-                    "SQL Lab timeout. This environment's policy is to kill queries "
-                    "after {} seconds.".format(SQLLAB_TIMEOUT)
-                )
-            )
         except Exception as ex:  # pylint: disable=broad-except
             logger.debug("Query %d: %s", query_id, ex)
             stats_logger.incr("error_sqllab_unhandled")
@@ -287,6 +274,19 @@ def execute_sql_statement(
             else:
                 # return 1 row less than increased_query
                 data = data[:-1]
+    except SoftTimeLimitExceeded as ex:
+        logger.warning("Query %d: Time limit exceeded", query.id)
+        logger.debug("Query %d: %s", query.id, ex)
+        raise SupersetErrorException(
+            SupersetError(
+                message=__(
+                    f"The query was killed after {SQLLAB_TIMEOUT} seconds. It might "
+                    "be too complex, or the database might be under heavy load."
+                ),
+                error_type=SupersetErrorType.SQLLAB_TIMEOUT_ERROR,
+                level=ErrorLevel.ERROR,
+            )
+        )
     except Exception as ex:
         logger.error("Query %d: %s", query.id, type(ex), exc_info=True)
         logger.debug("Query %d: %s", query.id, ex)

tests/sqllab_tests.py

@@ -20,6 +20,7 @@ import json
 from datetime import datetime, timedelta
 
 import pytest
+from celery.exceptions import SoftTimeLimitExceeded
 from parameterized import parameterized
 from random import random
 from unittest import mock
@@ -39,7 +40,6 @@ from superset.sql_lab import (
     execute_sql_statement,
     get_sql_results,
     SqlLabException,
-    SqlLabTimeoutException,
 )
 from superset.sql_parse import CtasMethod
 from superset.utils.core import (
@@ -945,25 +945,41 @@ class TestSqlLab(SupersetTestCase):
             },
         )
 
-    @mock.patch("superset.sql_lab.get_query")
-    @mock.patch("superset.sql_lab.execute_sql_statement")
-    def test_get_sql_results_soft_time_limit(
-        self, mock_execute_sql_statement, mock_get_query
-    ):
-        from celery.exceptions import SoftTimeLimitExceeded
+    @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")
+    def test_sql_json_soft_timeout(self):
+        examples_db = get_example_database()
+        if examples_db.backend == "sqlite":
+            return
 
-        sql = """
-            -- comment
-            SET @value = 42;
-            SELECT @value AS foo;
-            -- comment
-        """
-        mock_get_query.side_effect = SoftTimeLimitExceeded()
-        with pytest.raises(SqlLabTimeoutException) as excinfo:
-            get_sql_results(
-                1, sql, return_results=True, store_results=False,
-            )
-        assert (
-            str(excinfo.value)
-            == "SQL Lab timeout. This environment's policy is to kill queries after 21600 seconds."
-        )
+        self.login("admin")
+
+        with mock.patch.object(
+            examples_db.db_engine_spec, "handle_cursor"
+        ) as handle_cursor:
+            handle_cursor.side_effect = SoftTimeLimitExceeded()
+            data = self.run_sql("SELECT * FROM birth_names LIMIT 1", "1")
+
+        assert data == {
+            "errors": [
+                {
+                    "message": (
+                        "The query was killed after 21600 seconds. It might be too complex, "
+                        "or the database might be under heavy load."
+                    ),
+                    "error_type": SupersetErrorType.SQLLAB_TIMEOUT_ERROR,
+                    "level": ErrorLevel.ERROR,
+                    "extra": {
+                        "issue_codes": [
+                            {
+                                "code": 1026,
+                                "message": "Issue 1026 - Query is too complex and takes too long to run.",
+                            },
+                            {
+                                "code": 1027,
+                                "message": "Issue 1027 - The database is currently running too many queries.",
+                            },
+                        ]
+                    },
+                }
+            ]
+        }