diff --git a/superset/dashboards/api.py b/superset/dashboards/api.py index 3a3752830a285..f4644e68c2380 100644 --- a/superset/dashboards/api.py +++ b/superset/dashboards/api.py @@ -19,20 +19,19 @@ import logging from datetime import datetime from io import BytesIO -from typing import Any, Callable, cast, Optional +from typing import Any, Callable, cast from zipfile import is_zipfile, ZipFile from flask import g, redirect, request, Response, send_file, url_for from flask_appbuilder import permission_name from flask_appbuilder.api import expose, protect, rison, safe -from flask_appbuilder.hooks import before_request from flask_appbuilder.models.sqla.interface import SQLAInterface from flask_babel import gettext, ngettext from marshmallow import ValidationError from werkzeug.wrappers import Response as WerkzeugResponse from werkzeug.wsgi import FileWrapper -from superset import db, is_feature_enabled, thumbnail_cache +from superset import db, thumbnail_cache from superset.charts.schemas import ChartEntityResponseSchema from superset.commands.dashboard.copy import CopyDashboardCommand from superset.commands.dashboard.create import CreateDashboardCommand @@ -124,6 +123,7 @@ requires_form_data, requires_json, statsd_metrics, + validate_feature_flags, ) from superset.views.error_handling import handle_api_exception from superset.views.filters import ( @@ -160,18 +160,6 @@ def wraps(self: BaseSupersetModelRestApi, id_or_slug: str) -> Response: class DashboardRestApi(BaseSupersetModelRestApi): datamodel = SQLAInterface(Dashboard) - @before_request(only=["thumbnail", "cache_dashboard_screenshot", "screenshot"]) - def ensure_thumbnails_enabled(self) -> Optional[Response]: - if not is_feature_enabled("THUMBNAILS"): - return self.response_404() - return None - - @before_request(only=["cache_dashboard_screenshot", "screenshot"]) - def ensure_screenshots_enabled(self) -> Optional[Response]: - if not is_feature_enabled("ENABLE_DASHBOARD_SCREENSHOT_ENDPOINTS"): - return self.response_404() - return None - include_route_methods = RouteMethod.REST_MODEL_VIEW_CRUD_SET | { RouteMethod.EXPORT, RouteMethod.IMPORT, @@ -1038,6 +1026,7 @@ def export(self, **kwargs: Any) -> Response: return response @expose("//thumbnail//", methods=("GET",)) + @validate_feature_flags(["THUMBNAILS"]) @protect() @safe @rison(thumbnail_query_schema) @@ -1141,6 +1130,7 @@ def thumbnail(self, pk: int, digest: str, **kwargs: Any) -> WerkzeugResponse: ) @expose("//cache_dashboard_screenshot/", methods=("POST",)) + @validate_feature_flags(["THUMBNAILS", "ENABLE_DASHBOARD_SCREENSHOT_ENDPOINTS"]) @protect() @rison(screenshot_query_schema) @safe @@ -1241,6 +1231,7 @@ def trigger_celery() -> WerkzeugResponse: return trigger_celery() @expose("//screenshot//", methods=("GET",)) + @validate_feature_flags(["THUMBNAILS", "ENABLE_DASHBOARD_SCREENSHOT_ENDPOINTS"]) @protect() @safe @statsd_metrics diff --git a/superset/views/base_api.py b/superset/views/base_api.py index 8240481adaa61..682e8f8697130 100644 --- a/superset/views/base_api.py +++ b/superset/views/base_api.py @@ -31,6 +31,7 @@ from sqlalchemy import and_, distinct, func from sqlalchemy.orm.query import Query +from superset import is_feature_enabled from superset.exceptions import InvalidPayloadFormatError from superset.extensions import db, event_logger, security_manager, stats_logger_manager from superset.models.core import FavStar @@ -130,6 +131,29 @@ def wraps(self: BaseSupersetApiMixin, *args: Any, **kwargs: Any) -> Response: return functools.update_wrapper(wraps, f) +def validate_feature_flags( + feature_flags: list[str], +) -> Callable[[Callable[..., Response]], Callable[..., Response]]: + """ + A decorator to check if all given feature flags are enabled. + + :param feature_flags: List of feature flag names to be checked. + """ + + def decorate(f: Callable[..., Response]) -> Callable[..., Response]: + @functools.wraps(f) + def wrapper( + self: BaseSupersetModelRestApi, *args: Any, **kwargs: Any + ) -> Response: + if not all(is_feature_enabled(flag) for flag in feature_flags): + return self.response_404() + return f(self, *args, **kwargs) + + return wrapper + + return decorate + + class RelatedFieldFilter: # data class to specify what filter to use on a /related endpoint # pylint: disable=too-few-public-methods diff --git a/tests/integration_tests/dashboards/api_tests.py b/tests/integration_tests/dashboards/api_tests.py index be40b7dfb3ed4..2bab17d05cfe7 100644 --- a/tests/integration_tests/dashboards/api_tests.py +++ b/tests/integration_tests/dashboards/api_tests.py @@ -41,6 +41,7 @@ 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.constants import ( ADMIN_USERNAME, ALPHA_USERNAME, @@ -3027,10 +3028,9 @@ def _get_screenshot(self, dashboard_id, cache_key, download_format): uri = f"/api/v1/dashboard/{dashboard_id}/screenshot/{cache_key}/?download_format={download_format}" # noqa: E501 return self.client.get(uri) + @with_feature_flags(THUMBNAILS=True, ENABLE_DASHBOARD_SCREENSHOT_ENDPOINTS=True) @pytest.mark.usefixtures("create_dashboard_with_tag") - @patch("superset.dashboards.api.is_feature_enabled") - def test_cache_dashboard_screenshot_success(self, is_feature_enabled): - is_feature_enabled.return_value = True + def test_cache_dashboard_screenshot_success(self): self.login(ADMIN_USERNAME) dashboard = ( db.session.query(Dashboard) @@ -3040,10 +3040,9 @@ def test_cache_dashboard_screenshot_success(self, is_feature_enabled): response = self._cache_screenshot(dashboard.id) assert response.status_code == 202 + @with_feature_flags(THUMBNAILS=True, ENABLE_DASHBOARD_SCREENSHOT_ENDPOINTS=True) @pytest.mark.usefixtures("create_dashboard_with_tag") - @patch("superset.dashboards.api.is_feature_enabled") - def test_cache_dashboard_screenshot_dashboard_validation(self, is_feature_enabled): - is_feature_enabled.return_value = True + def test_cache_dashboard_screenshot_dashboard_validation(self): self.login(ADMIN_USERNAME) dashboard = ( db.session.query(Dashboard) @@ -3059,25 +3058,21 @@ def test_cache_dashboard_screenshot_dashboard_validation(self, is_feature_enable response = self._cache_screenshot(dashboard.id, invalid_payload) assert response.status_code == 400 - @patch("superset.dashboards.api.is_feature_enabled") - def test_cache_dashboard_screenshot_dashboard_not_found(self, is_feature_enabled): - is_feature_enabled.return_value = True + @with_feature_flags(THUMBNAILS=True, ENABLE_DASHBOARD_SCREENSHOT_ENDPOINTS=True) + def test_cache_dashboard_screenshot_dashboard_not_found(self): self.login(ADMIN_USERNAME) non_existent_id = 999 response = self._cache_screenshot(non_existent_id) assert response.status_code == 404 + @with_feature_flags(THUMBNAILS=True, ENABLE_DASHBOARD_SCREENSHOT_ENDPOINTS=True) @pytest.mark.usefixtures("create_dashboard_with_tag") @patch("superset.dashboards.api.cache_dashboard_screenshot") @patch("superset.dashboards.api.DashboardScreenshot.get_from_cache_key") - @patch("superset.dashboards.api.is_feature_enabled") - def test_screenshot_success_png( - self, is_feature_enabled, mock_get_cache, mock_cache_task - ): + def test_screenshot_success_png(self, mock_get_cache, mock_cache_task): """ Validate screenshot returns png """ - is_feature_enabled.return_value = True self.login(ADMIN_USERNAME) mock_cache_task.return_value = None mock_get_cache.return_value = BytesIO(b"fake image data") @@ -3096,18 +3091,17 @@ def test_screenshot_success_png( assert response.mimetype == "image/png" assert response.data == b"fake image data" + @with_feature_flags(THUMBNAILS=True, ENABLE_DASHBOARD_SCREENSHOT_ENDPOINTS=True) @pytest.mark.usefixtures("create_dashboard_with_tag") @patch("superset.dashboards.api.cache_dashboard_screenshot") @patch("superset.dashboards.api.build_pdf_from_screenshots") @patch("superset.dashboards.api.DashboardScreenshot.get_from_cache_key") - @patch("superset.dashboards.api.is_feature_enabled") def test_screenshot_success_pdf( - self, is_feature_enabled, mock_get_from_cache, mock_build_pdf, mock_cache_task + self, mock_get_from_cache, mock_build_pdf, mock_cache_task ): """ Validate screenshot can return pdf. """ - is_feature_enabled.return_value = True self.login(ADMIN_USERNAME) mock_cache_task.return_value = None mock_get_from_cache.return_value = BytesIO(b"fake image data") @@ -3127,14 +3121,11 @@ def test_screenshot_success_pdf( assert response.mimetype == "application/pdf" assert response.data == b"fake pdf data" + @with_feature_flags(THUMBNAILS=True, ENABLE_DASHBOARD_SCREENSHOT_ENDPOINTS=True) @pytest.mark.usefixtures("create_dashboard_with_tag") @patch("superset.dashboards.api.cache_dashboard_screenshot") @patch("superset.dashboards.api.DashboardScreenshot.get_from_cache_key") - @patch("superset.dashboards.api.is_feature_enabled") - def test_screenshot_not_in_cache( - self, is_feature_enabled, mock_get_cache, mock_cache_task - ): - is_feature_enabled.return_value = True + def test_screenshot_not_in_cache(self, mock_get_cache, mock_cache_task): self.login(ADMIN_USERNAME) mock_cache_task.return_value = None mock_get_cache.return_value = None @@ -3151,22 +3142,18 @@ def test_screenshot_not_in_cache( response = self._get_screenshot(dashboard.id, cache_key, "pdf") assert response.status_code == 404 - @patch("superset.dashboards.api.is_feature_enabled") - def test_screenshot_dashboard_not_found(self, is_feature_enabled): - is_feature_enabled.return_value = True + @with_feature_flags(THUMBNAILS=True, ENABLE_DASHBOARD_SCREENSHOT_ENDPOINTS=True) + def test_screenshot_dashboard_not_found(self): self.login(ADMIN_USERNAME) non_existent_id = 999 response = self._get_screenshot(non_existent_id, "some_cache_key", "png") assert response.status_code == 404 + @with_feature_flags(THUMBNAILS=True, ENABLE_DASHBOARD_SCREENSHOT_ENDPOINTS=True) @pytest.mark.usefixtures("create_dashboard_with_tag") @patch("superset.dashboards.api.cache_dashboard_screenshot") @patch("superset.dashboards.api.DashboardScreenshot.get_from_cache_key") - @patch("superset.dashboards.api.is_feature_enabled") - def test_screenshot_invalid_download_format( - self, is_feature_enabled, mock_get_cache, mock_cache_task - ): - is_feature_enabled.return_value = True + def test_screenshot_invalid_download_format(self, mock_get_cache, mock_cache_task): self.login(ADMIN_USERNAME) mock_cache_task.return_value = None mock_get_cache.return_value = BytesIO(b"fake png data") @@ -3184,10 +3171,41 @@ def test_screenshot_invalid_download_format( response = self._get_screenshot(dashboard.id, cache_key, "invalid") assert response.status_code == 404 + @with_feature_flags(THUMBNAILS=False, ENABLE_DASHBOARD_SCREENSHOT_ENDPOINTS=True) + @pytest.mark.usefixtures("create_dashboard_with_tag") + def test_cache_dashboard_screenshot_feature_thumbnails_ff_disabled(self): + self.login(ADMIN_USERNAME) + + dashboard = ( + db.session.query(Dashboard) + .filter(Dashboard.dashboard_title == "dash with tag") + .first() + ) + + assert dashboard is not None + + response = self._cache_screenshot(dashboard.id) + assert response.status_code == 404 + + @with_feature_flags(THUMBNAILS=True, ENABLE_DASHBOARD_SCREENSHOT_ENDPOINTS=False) + @pytest.mark.usefixtures("create_dashboard_with_tag") + def test_cache_dashboard_screenshot_feature_screenshot_ff_disabled(self): + self.login(ADMIN_USERNAME) + + dashboard = ( + db.session.query(Dashboard) + .filter(Dashboard.dashboard_title == "dash with tag") + .first() + ) + + assert dashboard is not None + + response = self._cache_screenshot(dashboard.id) + assert response.status_code == 404 + + @with_feature_flags(THUMBNAILS=False, ENABLE_DASHBOARD_SCREENSHOT_ENDPOINTS=False) @pytest.mark.usefixtures("create_dashboard_with_tag") - @patch("superset.dashboards.api.is_feature_enabled") - def test_cache_dashboard_screenshot_feature_disabled(self, is_feature_enabled): - is_feature_enabled.return_value = False + def test_cache_dashboard_screenshot_feature_both_ff_disabled(self): self.login(ADMIN_USERNAME) dashboard = (