diff --git a/superset/commands/chart/importers/v1/__init__.py b/superset/commands/chart/importers/v1/__init__.py index 89fe5e7a7..dc5a70796 100644 --- a/superset/commands/chart/importers/v1/__init__.py +++ b/superset/commands/chart/importers/v1/__init__.py @@ -26,6 +26,7 @@ from superset.commands.chart.importers.v1.utils import import_chart from superset.commands.database.importers.v1.utils import import_database from superset.commands.dataset.importers.v1.utils import import_dataset from superset.commands.importers.v1 import ImportModelsCommand +from superset.commands.utils import update_chart_config_dataset from superset.connectors.sqla.models import SqlaTable from superset.daos.chart import ChartDAO from superset.databases.schemas import ImportV1DatabaseSchema @@ -86,16 +87,10 @@ class ImportChartsCommand(ImportModelsCommand): # update datasource id, type, and name dataset = datasets[config["dataset_uuid"]] - config.update( - { - "datasource_id": dataset.id, - "datasource_type": "table", - "datasource_name": dataset.table_name, - } - ) - config["params"].update({"datasource": dataset.uid}) - - if "query_context" in config: - config["query_context"] = None - + dataset_dict = { + "datasource_id": dataset.id, + "datasource_type": "table", + "datasource_name": dataset.table_name, + } + config = update_chart_config_dataset(config, dataset_dict) import_chart(config, overwrite=overwrite) diff --git a/superset/commands/dashboard/importers/v1/__init__.py b/superset/commands/dashboard/importers/v1/__init__.py index 48b4e93e8..18cbb7da8 100644 --- a/superset/commands/dashboard/importers/v1/__init__.py +++ b/superset/commands/dashboard/importers/v1/__init__.py @@ -34,6 +34,7 @@ from superset.commands.dashboard.importers.v1.utils import ( from superset.commands.database.importers.v1.utils import import_database from superset.commands.dataset.importers.v1.utils import import_dataset from superset.commands.importers.v1 import ImportModelsCommand +from superset.commands.utils import update_chart_config_dataset from superset.daos.dashboard import DashboardDAO from superset.dashboards.schemas import ImportV1DashboardSchema from superset.databases.schemas import ImportV1DatabaseSchema @@ -113,11 +114,7 @@ class ImportDashboardsCommand(ImportModelsCommand): ): # update datasource id, type, and name dataset_dict = dataset_info[config["dataset_uuid"]] - config.update(dataset_dict) - dataset_uid = f"{dataset_dict['datasource_id']}__{dataset_dict['datasource_type']}" - config["params"].update({"datasource": dataset_uid}) - if "query_context" in config: - config["query_context"] = None + config = update_chart_config_dataset(config, dataset_dict) chart = import_chart(config, overwrite=False) charts.append(chart) diff --git a/superset/commands/importers/v1/assets.py b/superset/commands/importers/v1/assets.py index 78a2251a2..c0be04a66 100644 --- a/superset/commands/importers/v1/assets.py +++ b/superset/commands/importers/v1/assets.py @@ -39,6 +39,7 @@ from superset.commands.importers.v1.utils import ( validate_metadata_type, ) from superset.commands.query.importers.v1.utils import import_saved_query +from superset.commands.utils import update_chart_config_dataset from superset.dashboards.schemas import ImportV1DashboardSchema from superset.databases.schemas import ImportV1DatabaseSchema from superset.datasets.schemas import ImportV1DatasetSchema @@ -113,11 +114,7 @@ class ImportAssetsCommand(BaseCommand): for file_name, config in configs.items(): if file_name.startswith("charts/"): dataset_dict = dataset_info[config["dataset_uuid"]] - config.update(dataset_dict) - dataset_uid = f"{dataset_dict['datasource_id']}__{dataset_dict['datasource_type']}" - config["params"].update({"datasource": dataset_uid}) - if "query_context" in config: - config["query_context"] = None + config = update_chart_config_dataset(config, dataset_dict) chart = import_chart(config, overwrite=True) charts.append(chart) chart_ids[str(chart.uuid)] = chart.id diff --git a/superset/commands/utils.py b/superset/commands/utils.py index 29a31aa2a..7486d657a 100644 --- a/superset/commands/utils.py +++ b/superset/commands/utils.py @@ -17,7 +17,7 @@ from __future__ import annotations from collections import Counter -from typing import Optional, TYPE_CHECKING +from typing import Any, Optional, TYPE_CHECKING from flask import g from flask_appbuilder.security.sqla.models import Role, User @@ -34,6 +34,7 @@ from superset.daos.datasource import DatasourceDAO from superset.daos.exceptions import DatasourceNotFound from superset.daos.tag import TagDAO from superset.tags.models import ObjectType, Tag, TagType +from superset.utils import json from superset.utils.core import DatasourceType, get_user_id if TYPE_CHECKING: @@ -185,3 +186,43 @@ def update_tags( TagDAO.create_custom_tagged_objects( object_type, object_id, [tag.name for tag in tags_to_add] ) + + +def update_chart_config_dataset( + config: dict[str, Any], dataset_info: dict[str, Any] +) -> dict[str, Any]: + """ + Update the chart configuration and query_context with new dataset information + + :param config: The original chart configuration + :param dataset_info: Dict with datasource_id, datasource_type, and datasource_name + :return: The updated chart configuration + """ + # Update datasource id, type, and name + config.update(dataset_info) + + dataset_uid = f"{dataset_info['datasource_id']}__{dataset_info['datasource_type']}" + config["params"].update({"datasource": dataset_uid}) + + if "query_context" in config and config["query_context"] is not None: + try: + query_context = json.loads(config["query_context"]) + + query_context["datasource"] = { + "id": dataset_info["datasource_id"], + "type": dataset_info["datasource_type"], + } + + if "form_data" in query_context: + query_context["form_data"]["datasource"] = dataset_uid + + if "queries" in query_context: + for query in query_context["queries"]: + if "datasource" in query: + query["datasource"] = query_context["datasource"] + + config["query_context"] = json.dumps(query_context) + except json.JSONDecodeError: + config["query_context"] = None + + return config diff --git a/tests/integration_tests/commands_test.py b/tests/integration_tests/commands_test.py index 77dd32b81..c3015520b 100644 --- a/tests/integration_tests/commands_test.py +++ b/tests/integration_tests/commands_test.py @@ -139,7 +139,29 @@ class TestImportAssetsCommand(SupersetTestCase): dataset = chart.table assert str(dataset.uuid) == dataset_config["uuid"] - assert chart.query_context is None + assert json.loads(chart.query_context) == { + "datasource": {"id": dataset.id, "type": "table"}, + "force": False, + "queries": [ + { + "annotation_layers": [], + "applied_time_extras": {}, + "columns": [], + "custom_form_data": {}, + "custom_params": {}, + "extras": {"having": "", "time_grain_sqla": None, "where": ""}, + "filters": [], + "metrics": [], + "order_desc": True, + "row_limit": 5000, + "time_range": " : ", + "timeseries_limit": 0, + "url_params": {}, + } + ], + "result_format": "json", + "result_type": "full", + } assert json.loads(chart.params)["datasource"] == dataset.uid database = dataset.database diff --git a/tests/integration_tests/dashboards/commands_tests.py b/tests/integration_tests/dashboards/commands_tests.py index 38606a130..8c2fbc05f 100644 --- a/tests/integration_tests/dashboards/commands_tests.py +++ b/tests/integration_tests/dashboards/commands_tests.py @@ -620,7 +620,29 @@ class TestImportDashboardsCommand(SupersetTestCase): dataset = chart.table assert str(dataset.uuid) == dataset_config["uuid"] - assert chart.query_context is None + assert json.loads(chart.query_context) == { + "datasource": {"id": dataset.id, "type": "table"}, + "force": False, + "queries": [ + { + "annotation_layers": [], + "applied_time_extras": {}, + "columns": [], + "custom_form_data": {}, + "custom_params": {}, + "extras": {"having": "", "time_grain_sqla": None, "where": ""}, + "filters": [], + "metrics": [], + "order_desc": True, + "row_limit": 5000, + "time_range": " : ", + "timeseries_limit": 0, + "url_params": {}, + } + ], + "result_format": "json", + "result_type": "full", + } assert json.loads(chart.params)["datasource"] == dataset.uid database = dataset.database