From 5f49e0fdd06b558a9837d6fe07739d3989de9f61 Mon Sep 17 00:00:00 2001 From: John Bodley <4567245+john-bodley@users.noreply.github.com> Date: Wed, 19 Jul 2023 11:12:36 -0700 Subject: [PATCH] fix(cache): Add cache warmup for non-legacy charts (#24671) --- superset/charts/commands/export.py | 2 +- superset/charts/commands/warm_up_cache.py | 73 ++++++++---- superset/views/core.py | 59 ++++------ tests/integration_tests/charts/api_tests.py | 105 ++++++++++++++++-- .../charts/commands_tests.py | 48 +------- tests/integration_tests/core_tests.py | 89 +++++++++------ 6 files changed, 221 insertions(+), 155 deletions(-) diff --git a/superset/charts/commands/export.py b/superset/charts/commands/export.py index d1183a999c813..c942aa96c9a69 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/charts/commands/warm_up_cache.py b/superset/charts/commands/warm_up_cache.py index 97d5c4fe1e62a..3ef05d50a5be8 100644 --- a/superset/charts/commands/warm_up_cache.py +++ b/superset/charts/commands/warm_up_cache.py @@ -21,12 +21,17 @@ import simplejson as json from flask import g -from superset.charts.commands.exceptions import WarmUpCacheChartNotFoundError +from superset.charts.commands.exceptions import ( + ChartInvalidError, + WarmUpCacheChartNotFoundError, +) +from superset.charts.data.commands.get_data_command import ChartDataCommand from superset.commands.base import BaseCommand from superset.extensions import db from superset.models.slice import Slice from superset.utils.core import error_msg_from_exception from superset.views.utils import get_dashboard_extra_filters, get_form_data, get_viz +from superset.viz import viz_types class ChartWarmUpCacheCommand(BaseCommand): @@ -43,31 +48,51 @@ def __init__( def run(self) -> dict[str, Any]: self.validate() chart: Slice = self._chart_or_id # type: ignore + try: form_data = get_form_data(chart.id, use_slice_data=True)[0] - if self._dashboard_id: - form_data["extra_filters"] = ( - json.loads(self._extra_filters) - if self._extra_filters - else get_dashboard_extra_filters(chart.id, self._dashboard_id) - ) - - if not chart.datasource: - raise Exception("Chart's datasource does not exist") - - obj = get_viz( - datasource_type=chart.datasource.type, - datasource_id=chart.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") - error = payload["errors"] or None - status = payload["status"] + + if form_data.get("viz_type") in viz_types: + # Legacy visualizations. + if not chart.datasource: + raise ChartInvalidError("Chart's datasource does not exist") + + if self._dashboard_id: + form_data["extra_filters"] = ( + json.loads(self._extra_filters) + if self._extra_filters + else get_dashboard_extra_filters(chart.id, self._dashboard_id) + ) + + g.form_data = form_data # pylint: disable=assigning-non-slot + payload = get_viz( + datasource_type=chart.datasource.type, + datasource_id=chart.datasource.id, + form_data=form_data, + force=True, + ).get_payload() + delattr(g, "form_data") + error = payload["errors"] or None + status = payload["status"] + else: + # Non-legacy visualizations. + query_context = chart.get_query_context() + + if not query_context: + raise ChartInvalidError("Chart's query context does not exist") + + query_context.force = True + command = ChartDataCommand(query_context) + command.validate() + payload = command.run() + + # Report the first error. + for query in payload["queries"]: + error = query["error"] + status = query["status"] + + if error is not None: + break except Exception as ex: # pylint: disable=broad-except error = error_msg_from_exception(ex) status = None diff --git a/superset/views/core.py b/superset/views/core.py index 33b794a0bb2b1..f9bf51d54c905 100755 --- a/superset/views/core.py +++ b/superset/views/core.py @@ -43,6 +43,7 @@ security_manager, ) from superset.charts.commands.exceptions import ChartNotFoundError +from superset.charts.commands.warm_up_cache import ChartWarmUpCacheCommand from superset.common.chart_data import ChartDataResultFormat, ChartDataResultType from superset.connectors.base.models import BaseDatasource from superset.connectors.sqla.models import SqlaTable @@ -72,6 +73,7 @@ from superset.utils.async_query_manager import AsyncQueryTokenException from superset.utils.cache import etag_cache from superset.utils.core import ( + base_json_conv, DatasourceType, get_user_id, get_username, @@ -95,7 +97,6 @@ check_datasource_perms, check_explore_cache_perms, check_resource_permissions, - get_dashboard_extra_filters, get_datasource_info, get_form_data, get_viz, @@ -769,7 +770,8 @@ def save_or_overwrite_slice( @api @has_access_api @expose("/warm_up_cache/", methods=("GET",)) - def warm_up_cache( # pylint: disable=too-many-locals,no-self-use + @deprecated(new_target="api/v1/chart/warm_up_cache/") + def warm_up_cache( # pylint: disable=no-self-use self, ) -> FlaskResponse: """Warms up the cache for the slice or table. @@ -825,43 +827,22 @@ def warm_up_cache( # pylint: disable=too-many-locals,no-self-use .all() ) - result = [] - - 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) - ) - - 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, - ) - - # 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 - error = utils.error_msg_from_exception(ex) - status = None - - result.append( - {"slice_id": slc.id, "viz_error": error, "viz_status": status} - ) - - return json_success(json.dumps(result)) + return json_success( + json.dumps( + [ + { + "slice_id" if key == "chart_id" else key: value + for key, value in ChartWarmUpCacheCommand( + slc, dashboard_id, extra_filters + ) + .run() + .items() + } + for slc in slices + ], + default=base_json_conv, + ), + ) @has_access @expose("/dashboard//") diff --git a/tests/integration_tests/charts/api_tests.py b/tests/integration_tests/charts/api_tests.py index 734990a1eecc3..2db0cc7de301a 100644 --- a/tests/integration_tests/charts/api_tests.py +++ b/tests/integration_tests/charts/api_tests.py @@ -14,37 +14,41 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -# isort:skip_file """Unit tests for Superset""" import json from io import BytesIO +from unittest import mock from zipfile import is_zipfile, ZipFile import prison import pytest import yaml +from flask_babel import lazy_gettext as _ +from parameterized import parameterized from sqlalchemy import and_ from sqlalchemy.sql import func +from superset.charts.commands.exceptions import ChartDataQueryFailedError +from superset.charts.data.commands.get_data_command import ChartDataCommand from superset.connectors.sqla.models import SqlaTable from superset.extensions import cache_manager, db, security_manager from superset.models.core import Database, FavStar, FavStarClassName from superset.models.dashboard import Dashboard -from superset.reports.models import ReportSchedule, ReportScheduleType from superset.models.slice import Slice +from superset.reports.models import ReportSchedule, ReportScheduleType from superset.utils.core import get_example_default_schema from superset.utils.database import get_example_database - -from tests.integration_tests.conftest import with_feature_flags +from superset.viz import viz_types from tests.integration_tests.base_api_tests import ApiOwnersTestCaseMixin from tests.integration_tests.base_tests import SupersetTestCase +from tests.integration_tests.conftest import with_feature_flags from tests.integration_tests.fixtures.birth_names_dashboard import ( load_birth_names_dashboard_with_slices, load_birth_names_data, ) from tests.integration_tests.fixtures.energy_dashboard import ( - load_energy_table_with_slice, load_energy_table_data, + load_energy_table_with_slice, ) from tests.integration_tests.fixtures.importexport import ( chart_config, @@ -1710,12 +1714,16 @@ def test_gets_owned_created_favorited_by_me_filter(self): assert data["result"][0]["slice_name"] == "name0" assert data["result"][0]["datasource_id"] == 1 - @pytest.mark.usefixtures( - "load_energy_table_with_slice", "load_birth_names_dashboard_with_slices" + @parameterized.expand( + [ + "Top 10 Girl Name Share", # Legacy chart + "Pivot Table v2", # Non-legacy chart + ], ) - def test_warm_up_cache(self): + @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices") + def test_warm_up_cache(self, slice_name): self.login() - slc = self.get_slice("Top 10 Girl Name Share", db.session) + slc = self.get_slice(slice_name, db.session) rv = self.client.put("/api/v1/chart/warm_up_cache", json={"chart_id": slc.id}) self.assertEqual(rv.status_code, 200) data = json.loads(rv.data.decode("utf-8")) @@ -1780,7 +1788,6 @@ def test_warm_up_cache_payload_validation(self): ) self.assertEqual(rv.status_code, 400) data = json.loads(rv.data.decode("utf-8")) - print(data) self.assertEqual( data, { @@ -1791,3 +1798,81 @@ def test_warm_up_cache_payload_validation(self): } }, ) + + @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices") + def test_warm_up_cache_error(self) -> None: + self.login() + slc = self.get_slice("Pivot Table v2", db.session) + + with mock.patch.object(ChartDataCommand, "run") as mock_run: + mock_run.side_effect = ChartDataQueryFailedError( + _( + "Error: %(error)s", + error=_("Empty query?"), + ) + ) + + assert json.loads( + self.client.put( + "/api/v1/chart/warm_up_cache", + json={"chart_id": slc.id}, + ).data + ) == { + "result": [ + { + "chart_id": slc.id, + "viz_error": "Error: Empty query?", + "viz_status": None, + }, + ], + } + + @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices") + def test_warm_up_cache_no_query_context(self) -> None: + self.login() + slc = self.get_slice("Pivot Table v2", db.session) + + with mock.patch.object(Slice, "get_query_context") as mock_get_query_context: + mock_get_query_context.return_value = None + + assert json.loads( + self.client.put( + f"/api/v1/chart/warm_up_cache", + json={"chart_id": slc.id}, + ).data + ) == { + "result": [ + { + "chart_id": slc.id, + "viz_error": "Chart's query context does not exist", + "viz_status": None, + }, + ], + } + + @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices") + def test_warm_up_cache_no_datasource(self) -> None: + self.login() + slc = self.get_slice("Top 10 Girl Name Share", db.session) + + with mock.patch.object( + Slice, + "datasource", + new_callable=mock.PropertyMock, + ) as mock_datasource: + mock_datasource.return_value = None + + assert json.loads( + self.client.put( + f"/api/v1/chart/warm_up_cache", + json={"chart_id": slc.id}, + ).data + ) == { + "result": [ + { + "chart_id": slc.id, + "viz_error": "Chart's datasource does not exist", + "viz_status": None, + }, + ], + } diff --git a/tests/integration_tests/charts/commands_tests.py b/tests/integration_tests/charts/commands_tests.py index 9580c2bf33e6b..f9785a4dd6c20 100644 --- a/tests/integration_tests/charts/commands_tests.py +++ b/tests/integration_tests/charts/commands_tests.py @@ -96,52 +96,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") @@ -187,6 +142,7 @@ def test_export_chart_command_key_order(self, mock_g): "certification_details", "viz_type", "params", + "query_context", "cache_timeout", "uuid", "version", diff --git a/tests/integration_tests/core_tests.py b/tests/integration_tests/core_tests.py index 0acb19969a23d..e036602d0f973 100644 --- a/tests/integration_tests/core_tests.py +++ b/tests/integration_tests/core_tests.py @@ -14,72 +14,67 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -# isort:skip_file """Unit tests for Superset""" import datetime import doctest import html import json import logging -from urllib.parse import quote - -import prison -import superset.utils.database -from superset.utils.core import backend -from tests.integration_tests.fixtures.public_role import public_role_like_gamma -from tests.integration_tests.fixtures.birth_names_dashboard import ( - load_birth_names_dashboard_with_slices, - load_birth_names_data, -) -from sqlalchemy import Table - -import pytest -import pytz import random import unittest from unittest import mock +from urllib.parse import quote import pandas as pd +import prison +import pytest +import pytz import sqlalchemy as sqla +from flask_babel import lazy_gettext as _ +from sqlalchemy import Table from sqlalchemy.exc import SQLAlchemyError -from superset.models.cache import CacheKey -from superset.utils.database import get_example_database -from tests.integration_tests.conftest import with_feature_flags -from tests.integration_tests.fixtures.energy_dashboard import ( - load_energy_table_with_slice, - load_energy_table_data, -) -from tests.integration_tests.insert_chart_mixin import InsertChartMixin -from tests.integration_tests.test_app import app + +import superset.utils.database import superset.views.utils -from superset import ( - dataframe, - db, - security_manager, - sql_lab, -) +from superset import dataframe, db, security_manager, sql_lab +from superset.charts.commands.exceptions import ChartDataQueryFailedError +from superset.charts.data.commands.get_data_command import ChartDataCommand from superset.common.db_query_status import QueryStatus from superset.connectors.sqla.models import SqlaTable from superset.db_engine_specs.base import BaseEngineSpec from superset.db_engine_specs.mssql import MssqlEngineSpec -from superset.exceptions import SupersetException +from superset.exceptions import QueryObjectValidationError, SupersetException from superset.extensions import async_query_manager, cache_manager from superset.models import core as models from superset.models.annotations import Annotation, AnnotationLayer +from superset.models.cache import CacheKey from superset.models.dashboard import Dashboard from superset.models.slice import Slice from superset.models.sql_lab import Query from superset.result_set import SupersetResultSet from superset.utils import core as utils +from superset.utils.core import backend +from superset.utils.database import get_example_database from superset.views import core as views from superset.views.database.views import DatabaseView - -from .base_tests import SupersetTestCase +from tests.integration_tests.conftest import CTAS_SCHEMA_NAME, with_feature_flags +from tests.integration_tests.fixtures.birth_names_dashboard import ( + load_birth_names_dashboard_with_slices, + load_birth_names_data, +) +from tests.integration_tests.fixtures.energy_dashboard import ( + load_energy_table_data, + load_energy_table_with_slice, +) +from tests.integration_tests.fixtures.public_role import public_role_like_gamma from tests.integration_tests.fixtures.world_bank_dashboard import ( load_world_bank_dashboard_with_slices, load_world_bank_data, ) -from tests.integration_tests.conftest import CTAS_SCHEMA_NAME +from tests.integration_tests.insert_chart_mixin import InsertChartMixin +from tests.integration_tests.test_app import app + +from .base_tests import SupersetTestCase logger = logging.getLogger(__name__) @@ -389,7 +384,8 @@ def test_databaseview_edit(self, username="admin"): db.session.commit() @pytest.mark.usefixtures( - "load_energy_table_with_slice", "load_birth_names_dashboard_with_slices" + "load_birth_names_dashboard_with_slices", + "load_energy_table_with_slice", ) def test_warm_up_cache(self): self.login() @@ -415,6 +411,29 @@ def test_warm_up_cache(self): + quote(json.dumps([{"col": "name", "op": "in", "val": ["Jennifer"]}])) ) == [{"slice_id": slc.id, "viz_error": None, "viz_status": "success"}] + @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices") + def test_warm_up_cache_error(self) -> None: + self.login() + slc = self.get_slice("Pivot Table v2", db.session) + + with mock.patch.object( + ChartDataCommand, + "run", + side_effect=ChartDataQueryFailedError( + _( + "Error: %(error)s", + error=_("Empty query?"), + ) + ), + ): + assert self.get_json_resp(f"/superset/warm_up_cache?slice_id={slc.id}") == [ + { + "slice_id": slc.id, + "viz_error": "Error: Empty query?", + "viz_status": None, + } + ] + @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices") def test_cache_logging(self): self.login("admin")