diff --git a/superset/charts/api.py b/superset/charts/api.py index a600a3ca7f93f..292575feac49e 100644 --- a/superset/charts/api.py +++ b/superset/charts/api.py @@ -623,7 +623,7 @@ def build_response(status_code: int) -> WerkzeugResponse: if cache_payload.should_trigger_task(force): logger.info("Triggering screenshot ASYNC") - screenshot_obj.cache.set(cache_key, ScreenshotCachePayload()) + screenshot_obj.cache.set(cache_key, ScreenshotCachePayload().to_dict()) cache_chart_thumbnail.delay( current_user=get_current_user(), chart_id=chart.id, @@ -755,7 +755,7 @@ def thumbnail(self, pk: int, digest: str, **kwargs: Any) -> WerkzeugResponse: logger.info( "Triggering thumbnail compute (chart id: %s) ASYNC", str(chart.id) ) - screenshot_obj.cache.set(cache_key, ScreenshotCachePayload()) + screenshot_obj.cache.set(cache_key, ScreenshotCachePayload().to_dict()) cache_chart_thumbnail.delay( current_user=current_user, chart_id=chart.id, diff --git a/superset/dashboards/api.py b/superset/dashboards/api.py index c8c744ec63e92..c15610010bab4 100644 --- a/superset/dashboards/api.py +++ b/superset/dashboards/api.py @@ -1115,7 +1115,7 @@ def build_response(status_code: int) -> WerkzeugResponse: if cache_payload.should_trigger_task(force): logger.info("Triggering screenshot ASYNC") - screenshot_obj.cache.set(cache_key, ScreenshotCachePayload()) + screenshot_obj.cache.set(cache_key, ScreenshotCachePayload().to_dict()) cache_dashboard_screenshot.delay( username=get_current_user(), guest_token=( @@ -1296,7 +1296,7 @@ def thumbnail(self, pk: int, digest: str, **kwargs: Any) -> WerkzeugResponse: "Triggering thumbnail compute (dashboard id: %s) ASYNC", str(dashboard.id), ) - screenshot_obj.cache.set(cache_key, ScreenshotCachePayload()) + screenshot_obj.cache.set(cache_key, ScreenshotCachePayload().to_dict()) cache_dashboard_thumbnail.delay( current_user=current_user, dashboard_id=dashboard.id, diff --git a/superset/utils/screenshots.py b/superset/utils/screenshots.py index 86f5a94ce7041..7f40c0a8205f6 100644 --- a/superset/utils/screenshots.py +++ b/superset/utils/screenshots.py @@ -20,7 +20,7 @@ from datetime import datetime from enum import Enum from io import BytesIO -from typing import TYPE_CHECKING +from typing import cast, TYPE_CHECKING, TypedDict from flask import current_app @@ -63,13 +63,37 @@ class StatusValues(Enum): ERROR = "Error" +class ScreenshotCachePayloadType(TypedDict): + image: bytes | None + timestamp: str + status: str + + class ScreenshotCachePayload: - def __init__(self, image: bytes | None = None): + def __init__( + self, + image: bytes | None = None, + status: StatusValues = StatusValues.PENDING, + timestamp: str = "", + ): self._image = image - self._timestamp = datetime.now().isoformat() - self.status = StatusValues.PENDING - if image: - self.status = StatusValues.UPDATED + self._timestamp = timestamp or datetime.now().isoformat() + self.status = StatusValues.UPDATED if image else status + + @classmethod + def from_dict(cls, payload: ScreenshotCachePayloadType) -> ScreenshotCachePayload: + return cls( + image=payload["image"], + status=StatusValues(payload["status"]), + timestamp=payload["timestamp"], + ) + + def to_dict(self) -> ScreenshotCachePayloadType: + return { + "image": self._image, + "timestamp": self._timestamp, + "status": self.status.value, + } def update_timestamp(self) -> None: self._timestamp = datetime.now().isoformat() @@ -177,9 +201,16 @@ def get_from_cache( def get_from_cache_key(cls, cache_key: str) -> ScreenshotCachePayload | None: logger.info("Attempting to get from cache: %s", cache_key) if payload := cls.cache.get(cache_key): - # for backwards compatability, byte objects should be converted - if not isinstance(payload, ScreenshotCachePayload): + # Initially, only bytes were stored. This was changed to store an instance + # of ScreenshotCachePayload, but since it can't be serialized in all + # backends it was further changed to a dict of attributes. + if isinstance(payload, bytes): payload = ScreenshotCachePayload(payload) + elif isinstance(payload, ScreenshotCachePayload): + pass + elif isinstance(payload, dict): + payload = cast(ScreenshotCachePayloadType, payload) + payload = ScreenshotCachePayload.from_dict(payload) return payload logger.info("Failed at getting from cache: %s", cache_key) return None @@ -217,7 +248,7 @@ def compute_and_cache( # pylint: disable=too-many-arguments thumb_size = thumb_size or self.thumb_size logger.info("Processing url for thumbnail: %s", cache_key) cache_payload.computing() - self.cache.set(cache_key, cache_payload) + self.cache.set(cache_key, cache_payload.to_dict()) image = None # Assuming all sorts of things can go wrong with Selenium try: @@ -239,7 +270,7 @@ def compute_and_cache( # pylint: disable=too-many-arguments logger.info("Caching thumbnail: %s", cache_key) with event_logger.log_context(f"screenshot.cache.{self.thumbnail_type}"): cache_payload.update(image) - self.cache.set(cache_key, cache_payload) + self.cache.set(cache_key, cache_payload.to_dict()) logger.info("Updated thumbnail cache; Status: %s", cache_payload.get_status()) return diff --git a/tests/unit_tests/utils/screenshot_test.py b/tests/unit_tests/utils/screenshot_test.py index 5d29d829a20f6..fa928fc73fb67 100644 --- a/tests/unit_tests/utils/screenshot_test.py +++ b/tests/unit_tests/utils/screenshot_test.py @@ -26,7 +26,7 @@ from superset.utils.screenshots import ( BaseScreenshot, ScreenshotCachePayload, - StatusValues, + ScreenshotCachePayloadType, ) BASE_SCREENSHOT_PATH = "superset.utils.screenshots.BaseScreenshot" @@ -121,24 +121,24 @@ def _setup_compute_and_cache(self, mocker: MockerFixture, screenshot_obj): def test_happy_path(self, mocker: MockerFixture, screenshot_obj): self._setup_compute_and_cache(mocker, screenshot_obj) screenshot_obj.compute_and_cache(force=False) - cache_payload: ScreenshotCachePayload = screenshot_obj.cache.get("key") - assert cache_payload.status == StatusValues.UPDATED + cache_payload: ScreenshotCachePayloadType = screenshot_obj.cache.get("key") + assert cache_payload["status"] == "Updated" def test_screenshot_error(self, mocker: MockerFixture, screenshot_obj): mocks = self._setup_compute_and_cache(mocker, screenshot_obj) get_screenshot: MagicMock = mocks.get("get_screenshot") get_screenshot.side_effect = Exception screenshot_obj.compute_and_cache(force=False) - cache_payload: ScreenshotCachePayload = screenshot_obj.cache.get("key") - assert cache_payload.status == StatusValues.ERROR + cache_payload: ScreenshotCachePayloadType = screenshot_obj.cache.get("key") + assert cache_payload["status"] == "Error" def test_resize_error(self, mocker: MockerFixture, screenshot_obj): mocks = self._setup_compute_and_cache(mocker, screenshot_obj) resize_image: MagicMock = mocks.get("resize_image") resize_image.side_effect = Exception screenshot_obj.compute_and_cache(force=False) - cache_payload: ScreenshotCachePayload = screenshot_obj.cache.get("key") - assert cache_payload.status == StatusValues.ERROR + cache_payload: ScreenshotCachePayloadType = screenshot_obj.cache.get("key") + assert cache_payload["status"] == "Error" def test_skips_if_computing(self, mocker: MockerFixture, screenshot_obj): mocks = self._setup_compute_and_cache(mocker, screenshot_obj) @@ -155,8 +155,8 @@ def test_skips_if_computing(self, mocker: MockerFixture, screenshot_obj): # Ensure that it processes when force = True screenshot_obj.compute_and_cache(force=True) get_screenshot.assert_called_once() - cache_payload: ScreenshotCachePayload = screenshot_obj.cache.get("key") - assert cache_payload.status == StatusValues.UPDATED + cache_payload: ScreenshotCachePayloadType = screenshot_obj.cache.get("key") + assert cache_payload["status"] == "Updated" def test_skips_if_updated(self, mocker: MockerFixture, screenshot_obj): mocks = self._setup_compute_and_cache(mocker, screenshot_obj) @@ -177,8 +177,8 @@ def test_skips_if_updated(self, mocker: MockerFixture, screenshot_obj): force=True, window_size=window_size, thumb_size=thumb_size ) get_screenshot.assert_called_once() - cache_payload: ScreenshotCachePayload = screenshot_obj.cache.get("key") - assert cache_payload._image != b"initial_value" + cache_payload: ScreenshotCachePayloadType = screenshot_obj.cache.get("key") + assert cache_payload["image"] != b"initial_value" def test_resize(self, mocker: MockerFixture, screenshot_obj): mocks = self._setup_compute_and_cache(mocker, screenshot_obj)