fix(Jinja metric macro): Support Drill By and Excel/CSV download with…

…out a dataset ID (#30443)
apache · Oct 11, 2024 · 9c12b1c · 9c12b1c
1 parent 0db59b4
commit 9c12b1c
Show file tree

Hide file tree

Showing 3 changed files with 284 additions and 167 deletions.
diff --git a/superset/jinja_context.py b/superset/jinja_context.py
@@ -25,7 +25,7 @@
 from typing import Any, Callable, cast, Optional, TYPE_CHECKING, TypedDict, Union
 
 import dateutil
-from flask import current_app, has_request_context, request
+from flask import current_app, g, has_request_context, request
 from flask_babel import gettext as _
 from jinja2 import DebugUndefined, Environment
 from jinja2.sandbox import SandboxedEnvironment
@@ -847,35 +847,45 @@ def dataset_macro(
 
 def get_dataset_id_from_context(metric_key: str) -> int:
     """
-    Retrives the Dataset ID from the request context.
+    Retrieves the Dataset ID from the request context.
 
     :param metric_key: the metric key.
     :returns: the dataset ID.
     """
     # pylint: disable=import-outside-toplevel
     from superset.daos.chart import ChartDAO
-    from superset.views.utils import get_form_data
+    from superset.views.utils import loads_request_json
 
+    form_data: dict[str, Any] = {}
     exc_message = _(
         "Please specify the Dataset ID for the ``%(name)s`` metric in the Jinja macro.",
         name=metric_key,
     )
 
-    form_data, chart = get_form_data()
-    if not (form_data or chart):
-        raise SupersetTemplateException(exc_message)
+    if has_request_context():
+        if payload := request.get_json(cache=True) if request.is_json else None:
+            if dataset_id := payload.get("datasource", {}).get("id"):
+                return dataset_id
+            form_data.update(payload.get("form_data", {}))
+        request_form = loads_request_json(request.form.get("form_data"))
+        form_data.update(request_form)
+        request_args = loads_request_json(request.args.get("form_data"))
+        form_data.update(request_args)
+
+    if form_data := (form_data or getattr(g, "form_data", {})):
+        if datasource_info := form_data.get("datasource"):
+            if isinstance(datasource_info, dict):
+                return datasource_info["id"]
+            return datasource_info.split("__")[0]
+        url_params = form_data.get("queries", [{}])[0].get("url_params", {})
+        if dataset_id := url_params.get("datasource_id"):
+            return dataset_id
+        if chart_id := (form_data.get("slice_id") or url_params.get("slice_id")):
+            chart_data = ChartDAO.find_by_id(chart_id)
+            if not chart_data:
+                raise SupersetTemplateException(exc_message)
+            return chart_data.datasource_id
 
-    if chart and chart.datasource_id:
-        return chart.datasource_id
-    if dataset_id := form_data.get("url_params", {}).get("datasource_id"):
-        return dataset_id
-    if chart_id := (
-        form_data.get("slice_id") or form_data.get("url_params", {}).get("slice_id")
-    ):
-        chart_data = ChartDAO.find_by_id(chart_id)
-        if not chart_data:
-            raise SupersetTemplateException(exc_message)
-        return chart_data.datasource_id
     raise SupersetTemplateException(exc_message)
 
 

diff --git a/tests/integration_tests/sqla_models_tests.py b/tests/integration_tests/sqla_models_tests.py
@@ -200,8 +200,8 @@ def test_jinja_metrics_and_calc_columns(self, mock_username):
         db.session.delete(table)
         db.session.commit()
 
-    @patch("superset.views.utils.get_form_data")
-    def test_jinja_metric_macro(self, mock_form_data_context):
+    @patch("superset.jinja_context.get_dataset_id_from_context")
+    def test_jinja_metric_macro(self, mock_dataset_id_from_context):
         self.login(username="admin")
         table = self.get_table(name="birth_names")
         metric = SqlMetric(
@@ -234,14 +234,8 @@ def test_jinja_metric_macro(self, mock_form_data_context):
             "filter": [],
             "extras": {"time_grain_sqla": "P1D"},
         }
-        mock_form_data_context.return_value = [
-            {
-                "url_params": {
-                    "datasource_id": table.id,
-                }
-            },
-            None,
-        ]
+        mock_dataset_id_from_context.return_value = table.id
+
         sqla_query = table.get_sqla_query(**base_query_obj)
         query = table.database.compile_sqla_query(sqla_query.sqla_query)