This is an automated email from the ASF dual-hosted git repository. hugh pushed a commit to branch fix-logs-sqllab in repository https://gitbox.apache.org/repos/asf/superset.git
commit be9874f68ff49fccfe625f35d87845c998fc99ca Author: hughhhh <[email protected]> AuthorDate: Wed Oct 5 15:30:27 2022 -0400 save --- superset/sql_lab.py | 18 ++++++++++ tests/integration_tests/sqllab_tests.py | 64 +++++++++++++++++++++++++++++++++ 2 files changed, 82 insertions(+) diff --git a/superset/sql_lab.py b/superset/sql_lab.py index ea91133270..ed78673c23 100644 --- a/superset/sql_lab.py +++ b/superset/sql_lab.py @@ -29,6 +29,7 @@ import simplejson as json from celery import Task from celery.exceptions import SoftTimeLimitExceeded from flask_babel import gettext as __ +from psycopg2.errors import SyntaxError from sqlalchemy.orm import Session from superset import ( @@ -84,6 +85,10 @@ class SqlLabQueryStoppedException(SqlLabException): pass +class SqlUserError(SyntaxError): + pass + + def handle_query_error( ex: Exception, query: Query, @@ -204,6 +209,7 @@ def execute_sql_statement( # pylint: disable=too-many-arguments,too-many-statem database: Database = query.database db_engine_spec = database.db_engine_spec + print('in 1') parsed_query = ParsedQuery(sql_statement) if is_feature_enabled("RLS_IN_SQLLAB"): # Insert any applicable RLS predicates @@ -306,6 +312,14 @@ def execute_sql_statement( # pylint: disable=too-many-arguments,too-many-statem level=ErrorLevel.ERROR, ) ) from ex + except SqlUserError as ex: + raise SupersetErrorException( + SupersetError( + message=ex, + error_type=SupersetErrorType.SYNTAX_ERROR, + level=ErrorLevel.ERROR, + ) + ) from ex except Exception as ex: # query is stopped in another thread/worker # stopping raises expected exceptions which we should skip @@ -494,6 +508,7 @@ def execute_sql_statements( # pylint: disable=too-many-arguments, too-many-loca query.set_extra_json_key("progress", msg) session.commit() try: + print('calling execute_sql_statement') result_set = execute_sql_statement( statement, query, @@ -506,6 +521,9 @@ def execute_sql_statements( # pylint: disable=too-many-arguments, too-many-loca payload.update({"status": QueryStatus.STOPPED}) return payload except Exception as ex: # pylint: disable=broad-except + print('in exc') + print(ex) + raise ex msg = str(ex) prefix_message = ( f"[Statement {i+1} out of {statement_count}]" diff --git a/tests/integration_tests/sqllab_tests.py b/tests/integration_tests/sqllab_tests.py index 0c4019e7d9..b78b33d04e 100644 --- a/tests/integration_tests/sqllab_tests.py +++ b/tests/integration_tests/sqllab_tests.py @@ -821,6 +821,70 @@ class TestSqlLab(SupersetTestCase): }, ) + @mock.patch("superset.sql_lab.get_query") + @mock.patch("superset.sql_lab.execute_sql_statement") + def test_execute_sql_statements_syntax_error( + self, mock_execute_sql_statement, mock_get_query + ): + sql = """ + -- comment + SET @value = 42; + SELECT @value AS foo; + -- comment + """ + mock_session = mock.MagicMock() + mock_query = mock.MagicMock() + mock_query.database.allow_run_async = False + mock_cursor = mock.MagicMock() + mock_query.database.get_sqla_engine.return_value.raw_connection.return_value.cursor.return_value = ( + mock_cursor + ) + mock_query.database.db_engine_spec.run_multiple_statements_as_one = False + mock_get_query.return_value = mock_query + + s2_error = SupersetError( + message="foo", + error_type=SupersetErrorType.SYNTAX_ERROR, + level=ErrorLevel.ERROR, + ) + mock_execute_sql_statement.side_effect = SupersetErrorException(s2_error) + with pytest.raises(SupersetErrorException) as excinfo: + execute_sql_statements( + query_id=1, + rendered_query=sql, + return_results=True, + store_results=False, + session=mock_session, + start_time=None, + expand_data=False, + log_params=None, + ) + + assert excinfo.value.error == s2_error + assert False is True + + + # mock_execute_sql_statement.assert_has_calls( + # [ + # mock.call( + # "SET @value = 42", + # mock_query, + # mock_session, + # mock_cursor, + # None, + # False, + # ), + # mock.call( + # "SELECT @value AS foo", + # mock_query, + # mock_session, + # mock_cursor, + # None, + # False, + # ), + # ] + # ) + @mock.patch("superset.sql_lab.get_query") @mock.patch("superset.sql_lab.execute_sql_statement") def test_execute_sql_statements_ctas(
