chore(sqllab): Do not strip comments when executing SQL statements (#…

…27725)
apache · Apr 4, 2024 · 51d8863 · 51d8863
1 parent 863461d
commit 51d8863
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 24 deletions.
diff --git a/superset/sql_lab.py b/superset/sql_lab.py
@@ -422,7 +422,6 @@ def execute_sql_statements(
     # Breaking down into multiple statements
     parsed_query = ParsedQuery(
         rendered_query,
-        strip_comments=True,
         engine=db_engine_spec.engine,
     )
     if not db_engine_spec.run_multiple_statements_as_one:

diff --git a/tests/integration_tests/databases/api_tests.py b/tests/integration_tests/databases/api_tests.py
@@ -1280,7 +1280,7 @@ def test_get_invalid_table_table_metadata(self):
                     "indexes": [],
                     "name": "wrong_table",
                     "primaryKey": {"constrained_columns": None, "name": None},
-                    "selectStar": "SELECT\n  *\nFROM wrong_table\nLIMIT 100\nOFFSET 0",
+                    "selectStar": "SELECT *\nFROM wrong_table\nLIMIT 100\nOFFSET 0",
                 },
             )
         elif example_db.backend == "mysql":
@@ -2293,9 +2293,9 @@ def test_import_database_masked_password(self):
         uri = "api/v1/database/import/"
 
         masked_database_config = database_config.copy()
-        masked_database_config[
-            "sqlalchemy_uri"
-        ] = "postgresql://username:XXXXXXXXXX@host:12345/db"
+        masked_database_config["sqlalchemy_uri"] = (
+            "postgresql://username:XXXXXXXXXX@host:12345/db"
+        )
 
         buf = BytesIO()
         with ZipFile(buf, "w") as bundle:
@@ -2350,9 +2350,9 @@ def test_import_database_masked_password_provided(self):
         uri = "api/v1/database/import/"
 
         masked_database_config = database_config.copy()
-        masked_database_config[
-            "sqlalchemy_uri"
-        ] = "vertica+vertica_python://hackathon:XXXXXXXXXX@host:5433/dbname?ssl=1"
+        masked_database_config["sqlalchemy_uri"] = (
+            "vertica+vertica_python://hackathon:XXXXXXXXXX@host:5433/dbname?ssl=1"
+        )
 
         buf = BytesIO()
         with ZipFile(buf, "w") as bundle:

diff --git a/tests/integration_tests/sqllab_tests.py b/tests/integration_tests/sqllab_tests.py
@@ -18,6 +18,7 @@
 """Unit tests for Sql Lab"""
 import json
 from datetime import datetime
+from textwrap import dedent
 
 import pytest
 from celery.exceptions import SoftTimeLimitExceeded
@@ -578,12 +579,13 @@ def test_sql_json_parameter_forbidden(self):
     @mock.patch("superset.sql_lab.get_query")
     @mock.patch("superset.sql_lab.execute_sql_statement")
     def test_execute_sql_statements(self, mock_execute_sql_statement, mock_get_query):
-        sql = """
+        sql = dedent(
+            """
             -- comment
             SET @value = 42;
-            SELECT @value AS foo;
-            -- comment
+            SELECT /*+ hint */ @value AS foo;
         """
+        )
         mock_session = mock.MagicMock()
         mock_query = mock.MagicMock()
         mock_query.database.allow_run_async = False
@@ -607,15 +609,15 @@ def test_execute_sql_statements(self, mock_execute_sql_statement, mock_get_query
         mock_execute_sql_statement.assert_has_calls(
             [
                 mock.call(
-                    "SET @value = 42",
+                    "-- comment\nSET @value = 42",
                     mock_query,
                     mock_session,
                     mock_cursor,
                     None,
                     False,
                 ),
                 mock.call(
-                    "SELECT @value AS foo",
+                    "SELECT /*+ hint */ @value AS foo",
                     mock_query,
                     mock_session,
                     mock_cursor,
@@ -631,12 +633,13 @@ def test_execute_sql_statements(self, mock_execute_sql_statement, mock_get_query
     def test_execute_sql_statements_no_results_backend(
         self, mock_execute_sql_statement, mock_get_query
     ):
-        sql = """
+        sql = dedent(
+            """
             -- comment
             SET @value = 42;
-            SELECT @value AS foo;
-            -- comment
+            SELECT /*+ hint */ @value AS foo;
         """
+        )
         mock_session = mock.MagicMock()
         mock_query = mock.MagicMock()
         mock_query.database.allow_run_async = True
@@ -681,12 +684,13 @@ def test_execute_sql_statements_no_results_backend(
     def test_execute_sql_statements_ctas(
         self, mock_execute_sql_statement, mock_get_query
     ):
-        sql = """
+        sql = dedent(
+            """
             -- comment
             SET @value = 42;
-            SELECT @value AS foo;
-            -- comment
+            SELECT /*+ hint */ @value AS foo;
         """
+        )
         mock_session = mock.MagicMock()
         mock_query = mock.MagicMock()
         mock_query.database.allow_run_async = False
@@ -714,15 +718,15 @@ def test_execute_sql_statements_ctas(
         mock_execute_sql_statement.assert_has_calls(
             [
                 mock.call(
-                    "SET @value = 42",
+                    "-- comment\nSET @value = 42",
                     mock_query,
                     mock_session,
                     mock_cursor,
                     None,
                     False,
                 ),
                 mock.call(
-                    "SELECT @value AS foo",
+                    "SELECT /*+ hint */ @value AS foo",
                     mock_query,
                     mock_session,
                     mock_cursor,
@@ -761,12 +765,13 @@ def test_execute_sql_statements_ctas(
 
         # try invalid CVAS
         mock_query.ctas_method = CtasMethod.VIEW
-        sql = """
+        sql = dedent(
+            """
             -- comment
             SET @value = 42;
-            SELECT @value AS foo;
-            -- comment
+            SELECT /*+ hint */ @value AS foo;
         """
+        )
         with pytest.raises(SupersetErrorException) as excinfo:
             execute_sql_statements(
                 query_id=1,