From 928471512c3330d186e088cae4f871fc5f19d66b Mon Sep 17 00:00:00 2001
From: John Bodley <4567245+john-bodley@users.noreply.github.com>
Date: Wed, 3 Apr 2024 19:24:39 -0700
Subject: [PATCH] chore(sqllab): Do not strip comments when executing SQL
 statements (#27725)

---
 superset/sql_lab.py                           |  1 -
 .../integration_tests/databases/api_tests.py  |  2 +-
 tests/integration_tests/sqllab_tests.py       | 37 +++++++++++--------
 3 files changed, 22 insertions(+), 18 deletions(-)

diff --git a/superset/sql_lab.py b/superset/sql_lab.py
index e9b4d406f8a8a..059e436962539 100644
--- 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
index 9d28bd941bff8..ed96e208e51bd 100644
--- 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":
diff --git a/tests/integration_tests/sqllab_tests.py b/tests/integration_tests/sqllab_tests.py
index 03c3bdfc235e9..b54d018894151 100644
--- 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,7 +609,7 @@ 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,
@@ -615,7 +617,7 @@ def test_execute_sql_statements(self, mock_execute_sql_statement, mock_get_query
                     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,7 +718,7 @@ 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,
@@ -722,7 +726,7 @@ def test_execute_sql_statements_ctas(
                     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,