From 1874f9a3b6049a5328fa579e7f89d4ad8cfc64f1 Mon Sep 17 00:00:00 2001
From: Beto Dealmeida <roberto@dealmeida.net>
Date: Tue, 14 Mar 2023 16:29:12 -0700
Subject: [PATCH] fix: dataset_macro (#23376)

---
 superset/connectors/sqla/models.py     |  9 ++++--
 superset/jinja_context.py              |  4 +--
 tests/unit_tests/jinja_context_test.py | 41 ++++++++++++++++++++++----
 3 files changed, 44 insertions(+), 10 deletions(-)

diff --git a/superset/connectors/sqla/models.py b/superset/connectors/sqla/models.py
index aeaee612cdfeb..b76e423caf441 100644
--- a/superset/connectors/sqla/models.py
+++ b/superset/connectors/sqla/models.py
@@ -876,12 +876,17 @@ def mutate_query_from_config(self, sql: str) -> str:
     def get_template_processor(self, **kwargs: Any) -> BaseTemplateProcessor:
         return get_template_processor(table=self, database=self.database, **kwargs)
 
-    def get_query_str_extended(self, query_obj: QueryObjectDict) -> QueryStringExtended:
+    def get_query_str_extended(
+        self,
+        query_obj: QueryObjectDict,
+        mutate: bool = True,
+    ) -> QueryStringExtended:
         sqlaq = self.get_sqla_query(**query_obj)
         sql = self.database.compile_sqla_query(sqlaq.sqla_query)
         sql = self._apply_cte(sql, sqlaq.cte)
         sql = sqlparse.format(sql, reindent=True)
-        sql = self.mutate_query_from_config(sql)
+        if mutate:
+            sql = self.mutate_query_from_config(sql)
         return QueryStringExtended(
             applied_template_filters=sqlaq.applied_template_filters,
             applied_filter_columns=sqlaq.applied_filter_columns,
diff --git a/superset/jinja_context.py b/superset/jinja_context.py
index d9409e297b5fb..45786856c6a66 100644
--- a/superset/jinja_context.py
+++ b/superset/jinja_context.py
@@ -654,6 +654,6 @@ def dataset_macro(
         "metrics": metrics if include_metrics else None,
         "columns": columns,
     }
-    sqla_query = dataset.get_query_str_extended(query_obj)
+    sqla_query = dataset.get_query_str_extended(query_obj, mutate=False)
     sql = sqla_query.sql
-    return f"({sql}) AS dataset_{dataset_id}"
+    return f"(\n{sql}\n) AS dataset_{dataset_id}"
diff --git a/tests/unit_tests/jinja_context_test.py b/tests/unit_tests/jinja_context_test.py
index 13b3ae9e9c948..3478a9e3f0932 100644
--- a/tests/unit_tests/jinja_context_test.py
+++ b/tests/unit_tests/jinja_context_test.py
@@ -88,17 +88,20 @@ def test_dataset_macro(mocker: MockFixture) -> None:
 
     assert (
         dataset_macro(1)
-        == """(SELECT ds AS ds,
+        == """(
+SELECT ds AS ds,
        num_boys AS num_boys,
        revenue AS revenue,
        expenses AS expenses,
        revenue-expenses AS profit
-FROM my_schema.old_dataset) AS dataset_1"""
+FROM my_schema.old_dataset
+) AS dataset_1"""
     )
 
     assert (
         dataset_macro(1, include_metrics=True)
-        == """(SELECT ds AS ds,
+        == """(
+SELECT ds AS ds,
        num_boys AS num_boys,
        revenue AS revenue,
        expenses AS expenses,
@@ -109,18 +112,44 @@ def test_dataset_macro(mocker: MockFixture) -> None:
          num_boys,
          revenue,
          expenses,
-         revenue-expenses) AS dataset_1"""
+         revenue-expenses
+) AS dataset_1"""
     )
 
     assert (
         dataset_macro(1, include_metrics=True, columns=["ds"])
-        == """(SELECT ds AS ds,
+        == """(
+SELECT ds AS ds,
        COUNT(*) AS cnt
 FROM my_schema.old_dataset
-GROUP BY ds) AS dataset_1"""
+GROUP BY ds
+) AS dataset_1"""
     )
 
     DatasetDAO.find_by_id.return_value = None
     with pytest.raises(DatasetNotFoundError) as excinfo:
         dataset_macro(1)
     assert str(excinfo.value) == "Dataset 1 not found!"
+
+
+def test_dataset_macro_mutator_with_comments(mocker: MockFixture) -> None:
+    """
+    Test ``dataset_macro`` when the mutator adds comment.
+    """
+
+    def mutator(sql: str) -> str:
+        """
+        A simple mutator that wraps the query in comments.
+        """
+        return f"-- begin\n{sql}\n-- end"
+
+    DatasetDAO = mocker.patch("superset.datasets.dao.DatasetDAO")
+    DatasetDAO.find_by_id().get_query_str_extended().sql = mutator("SELECT 1")
+    assert (
+        dataset_macro(1)
+        == """(
+-- begin
+SELECT 1
+-- end
+) AS dataset_1"""
+    )