From adebc0997b574f4b25a9ed9bd79543a021350956 Mon Sep 17 00:00:00 2001
From: Beto Dealmeida <roberto@dealmeida.net>
Date: Fri, 20 Aug 2021 10:03:31 -0700
Subject: [PATCH] fix: update table ID in query context on chart import
 (#16374)

* fix: update table ID in query context on chart import

* Fix test
---
 .../charts/commands/importers/v1/__init__.py  |  7 +++++
 .../charts/commands_tests.py                  | 28 +++++++++++++++++++
 .../fixtures/importexport.py                  |  1 +
 3 files changed, 36 insertions(+)

diff --git a/superset/charts/commands/importers/v1/__init__.py b/superset/charts/commands/importers/v1/__init__.py
index 0e2b5b3a8adf7..5c2b535a7b34a 100644
--- a/superset/charts/commands/importers/v1/__init__.py
+++ b/superset/charts/commands/importers/v1/__init__.py
@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
+import json
 from typing import Any, Dict, Set
 
 from marshmallow import Schema
@@ -95,4 +96,10 @@ 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)
+
                 import_chart(session, config, overwrite=overwrite)
diff --git a/tests/integration_tests/charts/commands_tests.py b/tests/integration_tests/charts/commands_tests.py
index cd6e01f3dc1d4..238a54e29c7f4 100644
--- a/tests/integration_tests/charts/commands_tests.py
+++ b/tests/integration_tests/charts/commands_tests.py
@@ -191,6 +191,34 @@ def test_import_v1_chart(self):
         )
         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()
diff --git a/tests/integration_tests/fixtures/importexport.py b/tests/integration_tests/fixtures/importexport.py
index 951ecf9bb4350..78f643c587af1 100644
--- a/tests/integration_tests/fixtures/importexport.py
+++ b/tests/integration_tests/fixtures/importexport.py
@@ -444,6 +444,7 @@
         },
         "viz_type": "deck_path",
     },
+    "query_context": '{"datasource":{"id":12,"type":"table"},"force":false,"queries":[{"time_range":" : ","filters":[],"extras":{"time_grain_sqla":null,"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"}',
     "cache_timeout": None,
     "uuid": "0c23747a-6528-4629-97bf-e4b78d3b9df1",
     "version": "1.0.0",