[superset] 01/01: fix: dataset_macro

[be...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

beto pushed a commit to branch fix_dataset_macro
in repository https://gitbox.apache.org/repos/asf/superset.git

commit dcdbe4885f5ae72866a21f75276da3a85d06bd40
Author: Beto Dealmeida <ro...@dealmeida.net>
AuthorDate: Tue Mar 14 15:02:54 2023 -0700

    fix: dataset_macro
---
 superset/jinja_context.py              |  2 +-
 tests/unit_tests/jinja_context_test.py | 41 +++++++++++++++++++++++++++++-----
 2 files changed, 36 insertions(+), 7 deletions(-)

diff --git a/superset/jinja_context.py b/superset/jinja_context.py
index d9409e297b..49b31da9d3 100644
--- a/superset/jinja_context.py
+++ b/superset/jinja_context.py
@@ -656,4 +656,4 @@ def dataset_macro(
     }
     sqla_query = dataset.get_query_str_extended(query_obj)
     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 13b3ae9e9c..3478a9e3f0 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 @@ GROUP BY ds,
          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"""
+    )