This is an automated email from the ASF dual-hosted git repository. diegopucci pushed a commit to branch geido/fix/dash-screenshot-celery-optional in repository https://gitbox.apache.org/repos/asf/superset.git
commit 6fa4792c87e2512f04c1b33d626b70ff09a2ba6c Author: Diego Pucci <[email protected]> AuthorDate: Wed Oct 30 16:35:05 2024 +0200 fix(Dashboard): Sync/Async Dashboard Screenshot Generation and Default Cache --- .../DownloadMenuItems/DownloadScreenshot.test.tsx | 5 +- .../menu/DownloadMenuItems/DownloadScreenshot.tsx | 85 ++++++++++++---------- superset/config.py | 18 ++++- superset/dashboards/api.py | 53 ++++++++------ tests/integration_tests/dashboards/api_tests.py | 19 ++++- tests/integration_tests/superset_test_config.py | 6 ++ 6 files changed, 124 insertions(+), 62 deletions(-) diff --git a/superset-frontend/src/dashboard/components/menu/DownloadMenuItems/DownloadScreenshot.test.tsx b/superset-frontend/src/dashboard/components/menu/DownloadMenuItems/DownloadScreenshot.test.tsx index e1851e6199..9c8922f621 100644 --- a/superset-frontend/src/dashboard/components/menu/DownloadMenuItems/DownloadScreenshot.test.tsx +++ b/superset-frontend/src/dashboard/components/menu/DownloadMenuItems/DownloadScreenshot.test.tsx @@ -130,6 +130,9 @@ describe('DownloadScreenshot component', () => { await waitFor(() => { expect(mockAddInfoToast).toHaveBeenCalledWith( 'The screenshot is being generated. Please, do not leave the page.', + { + noDuplicate: true, + }, ); }); }); @@ -202,7 +205,7 @@ describe('DownloadScreenshot component', () => { // Wait for the successful image retrieval message await waitFor(() => { expect(mockAddSuccessToast).toHaveBeenCalledWith( - 'The screenshot is now being downloaded.', + 'The screenshot has been downloaded.', ); }); }); diff --git a/superset-frontend/src/dashboard/components/menu/DownloadMenuItems/DownloadScreenshot.tsx b/superset-frontend/src/dashboard/components/menu/DownloadMenuItems/DownloadScreenshot.tsx index 85f3e1d2c4..c4721b7714 100644 --- a/superset-frontend/src/dashboard/components/menu/DownloadMenuItems/DownloadScreenshot.tsx +++ b/superset-frontend/src/dashboard/components/menu/DownloadMenuItems/DownloadScreenshot.tsx @@ -53,21 +53,54 @@ export default function DownloadScreenshot({ const activeTabs = useSelector( (state: RootState) => state.dashboardState.activeTabs || undefined, ); - const anchor = useSelector( (state: RootState) => last(state.dashboardState.directPathToChild) || undefined, ); - const dataMask = useSelector( (state: RootState) => state.dataMask || undefined, ); - const { addDangerToast, addSuccessToast, addInfoToast } = useToasts(); + const printLoadingToast = () => + addInfoToast( + t('The screenshot is being generated. Please, do not leave the page.'), + { + noDuplicate: true, + }, + ); + + const printFailureToast = () => + addDangerToast( + t('The screenshot could not be downloaded. Please, try again later.'), + ); + + const printSuccessToast = () => + addSuccessToast(t('The screenshot has been downloaded.')); + + const stopToastInterval = ( + toastIntervalId: NodeJS.Timeout, + withFailure = false, + ) => { + clearInterval(toastIntervalId); + + if (withFailure) { + printFailureToast(); + } else { + printSuccessToast(); + } + }; + const onDownloadScreenshot = () => { let retries = 0; + const toastIntervalId = setInterval( + () => printLoadingToast(), + RETRY_INTERVAL, + ); + + printLoadingToast(); + // this function checks if the image is ready const checkImageReady = (cacheKey: string) => SupersetClient.get({ @@ -85,6 +118,7 @@ export default function DownloadScreenshot({ a.click(); document.body.removeChild(a); window.URL.revokeObjectURL(url); + stopToastInterval(toastIntervalId); }) .catch(err => { if ((err as SupersetApiError).status === 404) { @@ -94,32 +128,16 @@ export default function DownloadScreenshot({ // this is the functions that handles the retries const fetchImageWithRetry = (cacheKey: string) => { - checkImageReady(cacheKey) - .then(() => { - addSuccessToast(t('The screenshot is now being downloaded.')); - }) - .catch(error => { - // we check how many retries have been made - if (retries < MAX_RETRIES) { - retries += 1; - addInfoToast( - t( - 'The screenshot is being generated. Please, do not leave the page.', - ), - { - noDuplicate: true, - }, - ); - setTimeout(() => fetchImageWithRetry(cacheKey), RETRY_INTERVAL); - } else { - addDangerToast( - t( - 'The screenshot could not be downloaded. Please, try again later.', - ), - ); - logging.error(error); - } - }); + checkImageReady(cacheKey).catch(error => { + // we check how many retries have been made + if (retries < MAX_RETRIES) { + retries += 1; + setTimeout(() => fetchImageWithRetry(cacheKey), RETRY_INTERVAL); + } else { + stopToastInterval(toastIntervalId, true); + logging.error(error); + } + }); }; SupersetClient.post({ @@ -136,18 +154,11 @@ export default function DownloadScreenshot({ if (!cacheKey) { throw new Error('No image URL in response'); } - addInfoToast( - t( - 'The screenshot is being generated. Please, do not leave the page.', - ), - ); fetchImageWithRetry(cacheKey); }) .catch(error => { logging.error(error); - addDangerToast( - t('The screenshot could not be downloaded. Please, try again later.'), - ); + stopToastInterval(toastIntervalId, true); }) .finally(() => { logEvent?.( diff --git a/superset/config.py b/superset/config.py index d19e30a5a5..9d01f2e957 100644 --- a/superset/config.py +++ b/superset/config.py @@ -54,7 +54,7 @@ from superset.advanced_data_type.plugins.internet_port import internet_port from superset.advanced_data_type.types import AdvancedDataType from superset.constants import CHANGE_ME_SECRET_KEY from superset.jinja_context import BaseTemplateProcessor -from superset.key_value.types import JsonKeyValueCodec +from superset.key_value.types import JsonKeyValueCodec, PickleKeyValueCodec from superset.stats_logger import DummyStatsLogger from superset.superset_typing import CacheConfig from superset.tasks.types import ExecutorType @@ -478,6 +478,8 @@ DEFAULT_FEATURE_FLAGS: dict[str, bool] = { "PRESTO_EXPAND_DATA": False, # Exposes API endpoint to compute thumbnails "THUMBNAILS": False, + # Generates screenshots async when True (requires Celery) + "DASHBOARD_ASYNC_SCREENSHOT_GENERATION": False, "SHARE_QUERIES_VIA_KV_STORE": False, "TAGGING_SYSTEM": False, "SQLLAB_BACKEND_PERSISTENCE": True, @@ -791,6 +793,20 @@ EXPLORE_FORM_DATA_CACHE_CONFIG: CacheConfig = { "CODEC": JsonKeyValueCodec(), } +# Cache for thumbnail images. `CACHE_TYPE` defaults to `SupersetMetastoreCache` +# that stores the values in the key-value table in the Superset metastore, as it's +# required for Superset to operate correctly, but can be replaced by any +# `Flask-Caching` backend. +THUMBNAIL_CACHE_CONFIG = { + "CACHE_TYPE": "SupersetMetastoreCache", + "CACHE_DEFAULT_TIMEOUT": 300, + # Should the timeout be reset when retrieving a cached value? + "REFRESH_TIMEOUT_ON_RETRIEVAL": True, + # The following parameter only applies to `MetastoreCache`: + # How should entries be serialized/deserialized? + "CODEC": PickleKeyValueCodec(), +} + # store cache keys by datasource UID (via CacheKey) for custom processing/invalidation STORE_CACHE_KEYS_IN_METADATA_DB = False diff --git a/superset/dashboards/api.py b/superset/dashboards/api.py index 733cc555a4..8039b65e32 100644 --- a/superset/dashboards/api.py +++ b/superset/dashboards/api.py @@ -17,6 +17,7 @@ # pylint: disable=too-many-lines import functools import logging +import threading from datetime import datetime from io import BytesIO from typing import Any, Callable, cast, Optional @@ -1121,30 +1122,38 @@ class DashboardRestApi(BaseSupersetModelRestApi): "DashboardRestApi.screenshot", pk=dashboard.id, digest=cache_key ) - def trigger_celery() -> WerkzeugResponse: + screenshot_kwargs = { + "username": get_current_user(), + "guest_token": ( + g.user.guest_token + if get_current_user() and isinstance(g.user, GuestUser) + else None + ), + "dashboard_id": dashboard.id, + "dashboard_url": dashboard_url, + "cache_key": cache_key, + "force": True, + "thumb_size": thumb_size, + "window_size": window_size, + } + + def generate_screenshot() -> None: + cache_dashboard_screenshot(**screenshot_kwargs) + + if is_feature_enabled("DASHBOARD_ASYNC_SCREENSHOT_GENERATION"): logger.info("Triggering screenshot ASYNC") - cache_dashboard_screenshot.delay( - username=get_current_user(), - guest_token=( - g.user.guest_token - if get_current_user() and isinstance(g.user, GuestUser) - else None - ), - dashboard_id=dashboard.id, - dashboard_url=dashboard_url, - cache_key=cache_key, - force=True, - thumb_size=thumb_size, - window_size=window_size, - ) - return self.response( - 202, - cache_key=cache_key, - dashboard_url=dashboard_url, - image_url=image_url, - ) + cache_dashboard_screenshot.delay(**screenshot_kwargs) + else: + logger.info("Triggering screenshot SYNC") + thread = threading.Thread(target=generate_screenshot) + thread.start() - return trigger_celery() + return self.response( + 202, + cache_key=cache_key, + dashboard_url=dashboard_url, + image_url=image_url, + ) @expose("/<pk>/screenshot/<digest>/", methods=("GET",)) @protect() diff --git a/tests/integration_tests/dashboards/api_tests.py b/tests/integration_tests/dashboards/api_tests.py index c3bbf97536..fd86bbb0fc 100644 --- a/tests/integration_tests/dashboards/api_tests.py +++ b/tests/integration_tests/dashboards/api_tests.py @@ -3025,7 +3025,24 @@ class TestDashboardApi(ApiOwnersTestCaseMixin, InsertChartMixin, SupersetTestCas return self.client.get(uri) @pytest.mark.usefixtures("create_dashboard_with_tag") - def test_cache_dashboard_screenshot_success(self): + @patch("superset.dashboards.api.is_feature_enabled") + def test_async_cache_dashboard_screenshot_success(self, is_feature_enabled): + # DASHBOARD_ASYNC_SCREENSHOT_GENERATION feature flag to True + is_feature_enabled.return_value = True + self.login(ADMIN_USERNAME) + dashboard = ( + db.session.query(Dashboard) + .filter(Dashboard.dashboard_title == "dash with tag") + .first() + ) + response = self._cache_screenshot(dashboard.id) + assert response.status_code == 202 + + @pytest.mark.usefixtures("create_dashboard_with_tag") + @patch("superset.dashboards.api.is_feature_enabled") + def test_sync_cache_dashboard_screenshot_success(self, is_feature_enabled): + # DASHBOARD_ASYNC_SCREENSHOT_GENERATION feature flag to True + is_feature_enabled.return_value = False self.login(ADMIN_USERNAME) dashboard = ( db.session.query(Dashboard) diff --git a/tests/integration_tests/superset_test_config.py b/tests/integration_tests/superset_test_config.py index e1b3ce016a..08ca643b9c 100644 --- a/tests/integration_tests/superset_test_config.py +++ b/tests/integration_tests/superset_test_config.py @@ -122,6 +122,12 @@ FILTER_STATE_CACHE_CONFIG = { "CACHE_DEFAULT_TIMEOUT": int(timedelta(minutes=10).total_seconds()), } +THUMBNAIL_CACHE_CONFIG = { + "CACHE_TYPE": "SimpleCache", + "CACHE_THRESHOLD": math.inf, + "CACHE_DEFAULT_TIMEOUT": int(timedelta(minutes=10).total_seconds()), +} + EXPLORE_FORM_DATA_CACHE_CONFIG = { "CACHE_TYPE": "SimpleCache", "CACHE_THRESHOLD": math.inf,
