fix(embedded): sankey charts (apache#30491)

preset-io · Oct 2, 2024 · e0172a2 · e0172a2
1 parent bd42e33
commit e0172a2
Show file tree

Hide file tree

Showing 2 changed files with 155 additions and 4 deletions.
diff --git a/superset/security/manager.py b/superset/security/manager.py
@@ -177,7 +177,11 @@ def query_context_modified(query_context: "QueryContext") -> bool:
     )
 
     # compare columns and metrics in form_data with stored values
-    for key in ["metrics", "columns", "groupby"]:
+    for key, equivalent in [
+        ("metrics", ["metrics"]),
+        ("columns", ["columns", "groupby"]),
+        ("groupby", ["columns", "groupby"]),
+    ]:
         requested_values = {freeze_value(value) for value in form_data.get(key) or []}
         stored_values = {
             freeze_value(value) for value in stored_chart.params_dict.get(key) or []
@@ -193,9 +197,10 @@ def query_context_modified(query_context: "QueryContext") -> bool:
         }
         if stored_query_context:
             for query in stored_query_context.get("queries") or []:
-                stored_values.update(
-                    {freeze_value(value) for value in query.get(key) or []}
-                )
+                for key in equivalent:
+                    stored_values.update(
+                        {freeze_value(value) for value in query.get(key) or []}
+                    )
 
         if not queries_values.issubset(stored_values):
             return True

diff --git a/tests/unit_tests/security/manager_test.py b/tests/unit_tests/security/manager_test.py
@@ -17,6 +17,8 @@
 
 # pylint: disable=invalid-name, unused-argument, redefined-outer-name
 
+import json
+
 import pytest
 from flask_appbuilder.security.sqla.models import Role, User
 from pytest_mock import MockerFixture
@@ -691,6 +693,150 @@ def test_query_context_modified_mixed_chart(mocker: MockerFixture) -> None:
     assert not query_context_modified(query_context)
 
 
+def test_query_context_modified_sankey_tampered(mocker: MockerFixture) -> None:
+    """
+    Test the `query_context_modified` function for a sankey chart request.
+    """
+    query_context = mocker.MagicMock()
+    query_context.queries = [
+        {
+            "apply_fetch_values_predicate": False,
+            "columns": ["bot_id", "channel_id"],
+            "extras": {"having": "", "where": ""},
+            "filter": [
+                {
+                    "col": "bot_profile__updated",
+                    "op": "TEMPORAL_RANGE",
+                    "val": "No filter",
+                }
+            ],
+            "from_dttm": None,
+            "granularity": None,
+            "inner_from_dttm": None,
+            "inner_to_dttm": None,
+            "is_rowcount": False,
+            "is_timeseries": False,
+            "metrics": ["count"],
+            "order_desc": True,
+            "orderby": [],
+            "row_limit": 10000,
+            "row_offset": 0,
+            "series_columns": [],
+            "series_limit": 0,
+            "series_limit_metric": None,
+            "time_shift": None,
+            "to_dttm": None,
+        }
+    ]
+    query_context.form_data = {
+        "datasource": "12__table",
+        "viz_type": "sankey_v2",
+        "slice_id": 97,
+        "url_params": {},
+        "source": "bot_id",
+        "target": "channel_id",
+        "metric": "count",
+        "adhoc_filters": [
+            {
+                "clause": "WHERE",
+                "comparator": "No filter",
+                "expressionType": "SIMPLE",
+                "operator": "TEMPORAL_RANGE",
+                "subject": "bot_profile__updated",
+            }
+        ],
+        "row_limit": 10000,
+        "color_scheme": "supersetColors",
+        "dashboards": [11],
+        "extra_form_data": {},
+        "label_colors": {},
+        "shared_label_colors": {},
+        "extra_filters": [],
+        "dashboardId": 11,
+        "force": False,
+        "result_format": "json",
+        "result_type": "full",
+    }
+    query_context.slice_.id = 97
+    query_context.slice_.params_dict = {
+        "datasource": "12__table",
+        "viz_type": "sankey_v2",
+        "slice_id": 97,
+        "source": "bot_id",
+        "target": "channel_id",
+        "metric": "count",
+        "adhoc_filters": [
+            {
+                "clause": "WHERE",
+                "comparator": "No filter",
+                "expressionType": "SIMPLE",
+                "operator": "TEMPORAL_RANGE",
+                "subject": "bot_profile__updated",
+            }
+        ],
+        "row_limit": 10000,
+        "color_scheme": "supersetColors",
+        "extra_form_data": {},
+        "dashboards": [11],
+    }
+    query_context.slice_.query_context = json.dumps(
+        {
+            "datasource": {"id": 12, "type": "table"},
+            "force": False,
+            "queries": [
+                {
+                    "filters": [
+                        {
+                            "col": "bot_profile__updated",
+                            "op": "TEMPORAL_RANGE",
+                            "val": "No filter",
+                        }
+                    ],
+                    "extras": {"having": "", "where": ""},
+                    "applied_time_extras": {},
+                    "columns": [],
+                    "metrics": ["count"],
+                    "annotation_layers": [],
+                    "row_limit": 10000,
+                    "series_limit": 0,
+                    "order_desc": True,
+                    "url_params": {},
+                    "custom_params": {},
+                    "custom_form_data": {},
+                    "groupby": ["bot_id", "channel_id"],
+                }
+            ],
+            "form_data": {
+                "datasource": "12__table",
+                "viz_type": "sankey_v2",
+                "slice_id": 97,
+                "source": "bot_id",
+                "target": "channel_id",
+                "metric": "count",
+                "adhoc_filters": [
+                    {
+                        "clause": "WHERE",
+                        "comparator": "No filter",
+                        "expressionType": "SIMPLE",
+                        "operator": "TEMPORAL_RANGE",
+                        "subject": "bot_profile__updated",
+                    }
+                ],
+                "row_limit": 10000,
+                "color_scheme": "supersetColors",
+                "extra_form_data": {},
+                "dashboards": [11],
+                "force": False,
+                "result_format": "json",
+                "result_type": "full",
+            },
+            "result_format": "json",
+            "result_type": "full",
+        }
+    )
+    assert not query_context_modified(query_context)
+
+
 def test_get_catalog_perm() -> None:
     """
     Test the `get_catalog_perm` method.