feat: query datasets from SQL Lab (#15241)

* feat: Jinja2 macro for querying datasets

* Add docs

* Address comments

tests/unit_tests/jinja_context_test.py

@@ -14,8 +14,15 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+# pylint: disable=invalid-name, unused-argument
 
-from superset.jinja_context import where_in
+import json
+
+import pytest
+from pytest_mock import MockFixture
+
+from superset.datasets.commands.exceptions import DatasetNotFoundError
+from superset.jinja_context import dataset_macro, where_in
 
 
 def test_where_in() -> None:
@@ -25,3 +32,95 @@ def test_where_in() -> None:
     assert where_in([1, "b", 3]) == "(1, 'b', 3)"
     assert where_in([1, "b", 3], '"') == '(1, "b", 3)'
     assert where_in(["O'Malley's"]) == "('O''Malley''s')"
+
+
+def test_dataset_macro(mocker: MockFixture, app_context: None) -> None:
+    """
+    Test the ``dataset_macro`` macro.
+    """
+    # pylint: disable=import-outside-toplevel
+    from superset.connectors.sqla.models import SqlaTable, SqlMetric, TableColumn
+    from superset.models.core import Database
+
+    columns = [
+        TableColumn(column_name="ds", is_dttm=1, type="TIMESTAMP"),
+        TableColumn(column_name="num_boys", type="INTEGER"),
+        TableColumn(column_name="revenue", type="INTEGER"),
+        TableColumn(column_name="expenses", type="INTEGER"),
+        TableColumn(
+            column_name="profit", type="INTEGER", expression="revenue-expenses"
+        ),
+    ]
+    metrics = [
+        SqlMetric(metric_name="cnt", expression="COUNT(*)"),
+    ]
+
+    dataset = SqlaTable(
+        table_name="old_dataset",
+        columns=columns,
+        metrics=metrics,
+        main_dttm_col="ds",
+        default_endpoint="https://www.youtube.com/watch?v=dQw4w9WgXcQ",  # not used
+        database=Database(database_name="my_database", sqlalchemy_uri="sqlite://"),
+        offset=-8,
+        description="This is the description",
+        is_featured=1,
+        cache_timeout=3600,
+        schema="my_schema",
+        sql=None,
+        params=json.dumps(
+            {
+                "remote_id": 64,
+                "database_name": "examples",
+                "import_time": 1606677834,
+            }
+        ),
+        perm=None,
+        filter_select_enabled=1,
+        fetch_values_predicate="foo IN (1, 2)",
+        is_sqllab_view=0,  # no longer used?
+        template_params=json.dumps({"answer": "42"}),
+        schema_perm=None,
+        extra=json.dumps({"warning_markdown": "*WARNING*"}),
+    )
+    DatasetDAO = mocker.patch("superset.datasets.dao.DatasetDAO")
+    DatasetDAO.find_by_id.return_value = dataset
+
+    assert (
+        dataset_macro(1)
+        == """(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"""
+    )
+
+    assert (
+        dataset_macro(1, include_metrics=True)
+        == """(SELECT ds AS ds,
+       num_boys AS num_boys,
+       revenue AS revenue,
+       expenses AS expenses,
+       revenue-expenses AS profit,
+       COUNT(*) AS cnt
+FROM my_schema.old_dataset
+GROUP BY ds,
+         num_boys,
+         revenue,
+         expenses,
+         revenue-expenses) AS dataset_1"""
+    )
+
+    assert (
+        dataset_macro(1, include_metrics=True, columns=["ds"])
+        == """(SELECT ds AS ds,
+       COUNT(*) AS cnt
+FROM my_schema.old_dataset
+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!"