fix: update table ID in query context on chart import (apache#16374)

* fix: update table ID in query context on chart import * Fix test
preset-io · Aug 20, 2021 · adebc09 · adebc09
1 parent 13a2ee3
commit adebc09
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 0 deletions.
diff --git 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
@@ -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
@@ -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",