From ed43b79932ba8db605c29a41f5e67d664a668466 Mon Sep 17 00:00:00 2001
From: Beto Dealmeida <roberto@dealmeida.net>
Date: Wed, 1 Jun 2022 13:37:35 -0700
Subject: [PATCH] Address comments

---
 docs/docs/installation/sql-templating.mdx | 4 ++--
 superset/jinja_context.py                 | 5 ++---
 tests/unit_tests/jinja_context_test.py    | 2 +-
 3 files changed, 5 insertions(+), 6 deletions(-)

diff --git a/docs/docs/installation/sql-templating.mdx b/docs/docs/installation/sql-templating.mdx
index 4135403f80770..905f99cb9f5de 100644
--- a/docs/docs/installation/sql-templating.mdx
+++ b/docs/docs/installation/sql-templating.mdx
@@ -286,7 +286,7 @@ Once you have the ID you can query it as if it were a table:
 SELECT * FROM {{ dataset(42) }} LIMIT 10
 ```
 
-IF you want to select the metric definitions as well, in addition to the columns, you need to pass an additional keyword argument:
+If you want to select the metric definitions as well, in addition to the columns, you need to pass an additional keyword argument:
 
 ```
 SELECT * FROM {{ dataset(42, include_metrics=True) }} LIMIT 10
@@ -295,5 +295,5 @@ SELECT * FROM {{ dataset(42, include_metrics=True) }} LIMIT 10
 Since metrics are aggregations, the resulting SQL expression will be grouped by all non-metric columns. You can specify a subset of columns to group by instead:
 
 ```
-SELECT * FROM {{ dataset(42, include_metrics=True, groupby=["ds", "category"]) }} LIMIT 10
+SELECT * FROM {{ dataset(42, include_metrics=True, columns=["ds", "category"]) }} LIMIT 10
 ```
diff --git a/superset/jinja_context.py b/superset/jinja_context.py
index f2ad5378f00fb..42f6809c7402f 100644
--- a/superset/jinja_context.py
+++ b/superset/jinja_context.py
@@ -609,7 +609,7 @@ def get_template_processor(
 def dataset_macro(
     dataset_id: int,
     include_metrics: bool = False,
-    groupby: Optional[List[str]] = None,
+    columns: Optional[List[str]] = None,
 ) -> str:
     """
     Given a dataset ID, return the SQL that represents it.
@@ -624,14 +624,13 @@ def dataset_macro(
     if not dataset:
         raise DatasetNotFoundError(f"Dataset {dataset_id} not found!")
 
-    columns = [column.column_name for column in dataset.columns]
+    columns = columns or [column.column_name for column in dataset.columns]
     metrics = [metric.metric_name for metric in dataset.metrics]
     query_obj = {
         "is_timeseries": False,
         "filter": [],
         "metrics": metrics if include_metrics else None,
         "columns": columns,
-        "groupby": groupby,
     }
     sqla_query = dataset.get_query_str_extended(query_obj)
     sql = sqla_query.sql
diff --git a/tests/unit_tests/jinja_context_test.py b/tests/unit_tests/jinja_context_test.py
index 58791d0d6ccfd..75c49f0977bf6 100644
--- a/tests/unit_tests/jinja_context_test.py
+++ b/tests/unit_tests/jinja_context_test.py
@@ -113,7 +113,7 @@ def test_dataset_macro(mocker: MockFixture, app_context: None) -> None:
     )
 
     assert (
-        dataset_macro(1, include_metrics=True, groupby=["ds"])
+        dataset_macro(1, include_metrics=True, columns=["ds"])
         == """(SELECT ds AS ds,
        COUNT(*) AS cnt
 FROM my_schema.old_dataset