From d7e3a601b6b06b81e109c1b222ece009c7be5193 Mon Sep 17 00:00:00 2001
From: AAfghahi <48933336+AAfghahi@users.noreply.github.com>
Date: Tue, 30 Nov 2021 19:51:00 -0500
Subject: [PATCH] fix: Ch31968query context (#17600)

* a lot of console logs

* import and export of query context
---
 superset/charts/commands/export.py            |  8 +++--
 .../charts/commands/importers/v1/__init__.py  |  6 ++--
 superset/connectors/base/models.py            |  8 ++++-
 .../charts/commands_tests.py                  | 31 +------------------
 4 files changed, 15 insertions(+), 38 deletions(-)

diff --git a/superset/charts/commands/export.py b/superset/charts/commands/export.py
index b0dbc9f0d5036..960fd87762b6d 100644
--- a/superset/charts/commands/export.py
+++ b/superset/charts/commands/export.py
@@ -34,7 +34,7 @@
 
 
 # keys present in the standard export that are not needed
-REMOVE_KEYS = ["datasource_type", "datasource_name"]
+REMOVE_KEYS = ["datasource_type", "datasource_name", "query_context"]
 
 
 class ExportChartsCommand(ExportModelsCommand):
@@ -55,8 +55,10 @@ def _export(model: Slice) -> Iterator[Tuple[str, str]]:
         )
         # TODO (betodealmeida): move this logic to export_to_dict once this
         #  becomes the default export endpoint
-        for key in REMOVE_KEYS:
-            del payload[key]
+        payload = {
+            key: value for key, value in payload.items() if key not in REMOVE_KEYS
+        }
+
         if payload.get("params"):
             try:
                 payload["params"] = json.loads(payload["params"])
diff --git a/superset/charts/commands/importers/v1/__init__.py b/superset/charts/commands/importers/v1/__init__.py
index 5c2b535a7b34a..dc469a67690e8 100644
--- a/superset/charts/commands/importers/v1/__init__.py
+++ b/superset/charts/commands/importers/v1/__init__.py
@@ -96,10 +96,8 @@ def _import(
                     }
                 )
                 config["params"].update({"datasource": dataset.uid})
+
                 if config["query_context"]:
-                    # TODO (betodealmeida): export query_context as object, not string
-                    query_context = json.loads(config["query_context"])
-                    query_context["datasource"] = {"id": dataset.id, "type": "table"}
-                    config["query_context"] = json.dumps(query_context)
+                    del config["query_context"]
 
                 import_chart(session, config, overwrite=overwrite)
diff --git a/superset/connectors/base/models.py b/superset/connectors/base/models.py
index 0377401c99dd7..47d0a21657f19 100644
--- a/superset/connectors/base/models.py
+++ b/superset/connectors/base/models.py
@@ -26,6 +26,7 @@
 
 from superset import is_feature_enabled, security_manager
 from superset.constants import NULL_STRING
+from superset.datasets.commands.exceptions import DatasetNotFoundError
 from superset.models.helpers import AuditMixinNullable, ImportExportMixin, QueryResult
 from superset.models.slice import Slice
 from superset.typing import FilterValue, FilterValues, QueryObjectDict
@@ -319,8 +320,13 @@ def data_for_slices(  # pylint: disable=too-many-locals
                 if "column" in filter_config
             )
 
+            # for legacy dashboard imports which have the wrong query_context in them
+            try:
+                query_context = slc.get_query_context()
+            except DatasetNotFoundError:
+                query_context = None
+
             # legacy charts don't have query_context charts
-            query_context = slc.get_query_context()
             if query_context:
                 column_names.update(
                     [
diff --git a/tests/integration_tests/charts/commands_tests.py b/tests/integration_tests/charts/commands_tests.py
index bad8530488913..6d578b88aae43 100644
--- a/tests/integration_tests/charts/commands_tests.py
+++ b/tests/integration_tests/charts/commands_tests.py
@@ -79,7 +79,6 @@ def test_export_chart_command(self, mock_g):
                 "slice_name": "Energy Sankey",
                 "viz_type": "sankey",
             },
-            "query_context": None,
             "cache_timeout": None,
             "dataset_uuid": str(example_chart.table.uuid),
             "uuid": str(example_chart.uuid),
@@ -87,6 +86,7 @@ def test_export_chart_command(self, mock_g):
         }
 
     @patch("superset.security.manager.g")
+    @pytest.mark.usefixtures("load_energy_table_with_slice")
     def test_export_chart_command_no_access(self, mock_g):
         """Test that users can't export datasets they don't have access to"""
         mock_g.user = security_manager.find_user("gamma")
@@ -125,7 +125,6 @@ def test_export_chart_command_key_order(self, mock_g):
             "slice_name",
             "viz_type",
             "params",
-            "query_context",
             "cache_timeout",
             "uuid",
             "version",
@@ -191,34 +190,6 @@ def test_import_v1_chart(self, mock_g):
         )
         assert dataset.table_name == "imported_dataset"
         assert chart.table == dataset
-        assert json.loads(chart.query_context) == {
-            "datasource": {"id": dataset.id, "type": "table"},
-            "force": False,
-            "queries": [
-                {
-                    "time_range": " : ",
-                    "filters": [],
-                    "extras": {
-                        "time_grain_sqla": None,
-                        "having": "",
-                        "having_druid": [],
-                        "where": "",
-                    },
-                    "applied_time_extras": {},
-                    "columns": [],
-                    "metrics": [],
-                    "annotation_layers": [],
-                    "row_limit": 5000,
-                    "timeseries_limit": 0,
-                    "order_desc": True,
-                    "url_params": {},
-                    "custom_params": {},
-                    "custom_form_data": {},
-                }
-            ],
-            "result_format": "json",
-            "result_type": "full",
-        }
 
         database = (
             db.session.query(Database).filter_by(uuid=database_config["uuid"]).one()