diff --git a/superset/charts/commands/export.py b/superset/charts/commands/export.py index bb594591e2bbb..39c3c7d46a77b 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", "query_context", "url_params"] +REMOVE_KEYS = ["datasource_type", "datasource_name", "url_params"] class ExportChartsCommand(ExportModelsCommand): diff --git a/superset/views/core.py b/superset/views/core.py index 94cabfa96c181..f1603837bb660 100755 --- a/superset/views/core.py +++ b/superset/views/core.py @@ -56,6 +56,7 @@ ) from superset.charts.commands.exceptions import ChartNotFoundError from superset.charts.dao import ChartDAO +from superset.charts.data.commands.get_data_command import ChartDataCommand from superset.common.chart_data import ChartDataResultFormat, ChartDataResultType from superset.common.db_query_status import QueryStatus from superset.connectors.base.models import BaseDatasource @@ -1740,28 +1741,35 @@ def warm_up_cache( # pylint: disable=too-many-locals,no-self-use for slc in slices: try: - form_data = get_form_data(slc.id, use_slice_data=True)[0] - if dashboard_id: - form_data["extra_filters"] = ( - json.loads(extra_filters) - if extra_filters - else get_dashboard_extra_filters(slc.id, dashboard_id) - ) + query_context = slc.get_query_context() + if query_context: + query_context.force = True + command = ChartDataCommand(query_context) + command.validate() + payload = command.run() + else: + form_data = get_form_data(slc.id, use_slice_data=True)[0] + if dashboard_id: + form_data["extra_filters"] = ( + json.loads(extra_filters) + if extra_filters + else get_dashboard_extra_filters(slc.id, dashboard_id) + ) - if not slc.datasource: - raise Exception("Slice's datasource does not exist") + if not slc.datasource: + raise Exception("Slice's datasource does not exist") - obj = get_viz( - datasource_type=slc.datasource.type, - datasource_id=slc.datasource.id, - form_data=form_data, - force=True, - ) + obj = get_viz( + datasource_type=slc.datasource.type, + datasource_id=slc.datasource.id, + form_data=form_data, + force=True, + ) + # pylint: disable=assigning-non-slot + g.form_data = form_data + payload = obj.get_payload() + delattr(g, "form_data") - # pylint: disable=assigning-non-slot - g.form_data = form_data - payload = obj.get_payload() - delattr(g, "form_data") error = payload["errors"] or None status = payload["status"] except Exception as ex: # pylint: disable=broad-except diff --git a/tests/integration_tests/charts/commands_tests.py b/tests/integration_tests/charts/commands_tests.py index 7e3991f375d37..da9a7550acb44 100644 --- a/tests/integration_tests/charts/commands_tests.py +++ b/tests/integration_tests/charts/commands_tests.py @@ -88,52 +88,7 @@ def test_export_chart_command(self, mock_g): "dataset_uuid": str(example_chart.table.uuid), "uuid": str(example_chart.uuid), "version": "1.0.0", - } - - @patch("superset.security.manager.g") - @pytest.mark.usefixtures("load_energy_table_with_slice") - def test_export_chart_with_query_context(self, mock_g): - """Test that charts that have a query_context are exported correctly""" - - mock_g.user = security_manager.find_user("alpha") - example_chart = db.session.query(Slice).filter_by(slice_name="Heatmap").one() - command = ExportChartsCommand([example_chart.id]) - - contents = dict(command.run()) - - expected = [ - "metadata.yaml", - f"charts/Heatmap_{example_chart.id}.yaml", - "datasets/examples/energy_usage.yaml", - "databases/examples.yaml", - ] - assert expected == list(contents.keys()) - - metadata = yaml.safe_load(contents[f"charts/Heatmap_{example_chart.id}.yaml"]) - - assert metadata == { - "slice_name": "Heatmap", - "description": None, - "certified_by": None, - "certification_details": None, - "viz_type": "heatmap", - "params": { - "all_columns_x": "source", - "all_columns_y": "target", - "canvas_image_rendering": "pixelated", - "collapsed_fieldsets": "", - "linear_color_scheme": "blue_white_yellow", - "metric": "sum__value", - "normalize_across": "heatmap", - "slice_name": "Heatmap", - "viz_type": "heatmap", - "xscale_interval": "1", - "yscale_interval": "1", - }, - "cache_timeout": None, - "dataset_uuid": str(example_chart.table.uuid), - "uuid": str(example_chart.uuid), - "version": "1.0.0", + "query_context": None, } @patch("superset.security.manager.g") @@ -179,6 +134,7 @@ def test_export_chart_command_key_order(self, mock_g): "certification_details", "viz_type", "params", + "query_context", "cache_timeout", "uuid", "version",