fix: dataset_macro (#23376)

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,

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}"

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"""
+    )