You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by hu...@apache.org on 2022/10/06 18:27:26 UTC

[superset] 01/03: save

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 <hu...@gmail.com>
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(