From 241ee32f5625ef16051bb582969f74ef93b1703b Mon Sep 17 00:00:00 2001 From: Beto Dealmeida Date: Thu, 24 Jun 2021 08:02:49 -0700 Subject: [PATCH] feat: custom error SQL Lab timeout (#15342) * feat: custom error SQL Lab timeout * Update test --- .../pages/docs/Miscellaneous/issue_codes.mdx | 16 +++++ .../src/components/ErrorMessage/types.ts | 1 + .../src/setup/setupErrorMessages.ts | 4 ++ superset/errors.py | 15 +++++ superset/sql_lab.py | 28 ++++----- tests/sqllab_tests.py | 60 ++++++++++++------- 6 files changed, 88 insertions(+), 36 deletions(-) diff --git a/docs/src/pages/docs/Miscellaneous/issue_codes.mdx b/docs/src/pages/docs/Miscellaneous/issue_codes.mdx index cfe44493f..b51ce9b24 100644 --- a/docs/src/pages/docs/Miscellaneous/issue_codes.mdx +++ b/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. diff --git a/superset-frontend/src/components/ErrorMessage/types.ts b/superset-frontend/src/components/ErrorMessage/types.ts index c277c59c5..e2d2f2923 100644 --- a/superset-frontend/src/components/ErrorMessage/types.ts +++ b/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', diff --git a/superset-frontend/src/setup/setupErrorMessages.ts b/superset-frontend/src/setup/setupErrorMessages.ts index 6c005e97a..ab2fb2d83 100644 --- a/superset-frontend/src/setup/setupErrorMessages.ts +++ b/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, diff --git a/superset/errors.py b/superset/errors.py index 92ba980f4..02f420794 100644 --- a/superset/errors.py +++ b/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." + ), + }, + ], } diff --git a/superset/sql_lab.py b/superset/sql_lab.py index 7fa5de203..a40e132e8 100644 --- a/superset/sql_lab.py +++ b/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) diff --git a/tests/sqllab_tests.py b/tests/sqllab_tests.py index 35d6f2ded..a69c1c8e8 100644 --- a/tests/sqllab_tests.py +++ b/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.", + }, + ] + }, + } + ] + }