From 302e0cde4aa495ca9ba170ecf18f24c07e5433d9 Mon Sep 17 00:00:00 2001
From: Maxime Beauchemin <maximebeauchemin@gmail.com>
Date: Mon, 6 May 2024 16:47:33 -0700
Subject: [PATCH] fix: don't strip SQL comments

Following up in the footsteps of https://github.com/apache/superset/pull/27725 as comments can be used as a way to pass the database engine hints, or for logging purposes.
---
 superset/db_engine_specs/base.py | 5 ++---
 superset/models/helpers.py       | 2 +-
 superset/sqllab/query_render.py  | 1 -
 3 files changed, 3 insertions(+), 5 deletions(-)

diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py
index 4ee4e4dc0385b..5b8a1deaba085 100644
--- a/superset/db_engine_specs/base.py
+++ b/superset/db_engine_specs/base.py
@@ -1134,9 +1134,8 @@ def apply_top_to_sql(cls, sql: str, limit: int) -> str:
         cte = None
         sql_remainder = None
         sql = sql.strip(" \t\n;")
-        sql_statement = sqlparse.format(sql, strip_comments=True)
         query_limit: int | None = sql_parse.extract_top_from_query(
-            sql_statement, cls.top_keywords
+            sql, cls.top_keywords
         )
         if not limit:
             final_limit = query_limit
@@ -1145,7 +1144,7 @@ def apply_top_to_sql(cls, sql: str, limit: int) -> str:
         else:
             final_limit = limit
         if not cls.allows_cte_in_subquery:
-            cte, sql_remainder = sql_parse.get_cte_remainder_query(sql_statement)
+            cte, sql_remainder = sql_parse.get_cte_remainder_query(sql)
         if cte:
             str_statement = str(sql_remainder)
             cte = cte + "\n"
diff --git a/superset/models/helpers.py b/superset/models/helpers.py
index 100391086ccd4..5845a400c2939 100644
--- a/superset/models/helpers.py
+++ b/superset/models/helpers.py
@@ -1088,7 +1088,7 @@ def get_rendered_sql(
                 _("Virtual dataset query cannot consist of multiple statements")
             )
 
-        sql = script.statements[0].format(comments=False)
+        sql = script.statements[0].format()
         if not sql:
             raise QueryObjectValidationError(_("Virtual dataset query cannot be empty"))
         return sql
diff --git a/superset/sqllab/query_render.py b/superset/sqllab/query_render.py
index caf9a3cb2b206..67bea2b1fb4cd 100644
--- a/superset/sqllab/query_render.py
+++ b/superset/sqllab/query_render.py
@@ -60,7 +60,6 @@ def render(self, execution_context: SqlJsonExecutionContext) -> str:
 
             parsed_query = ParsedQuery(
                 query_model.sql,
-                strip_comments=True,
                 engine=query_model.database.db_engine_spec.engine,
             )
             rendered_query = sql_template_processor.process_template(