This is an automated email from the ASF dual-hosted git repository. michaelsmolina pushed a commit to branch 5.0 in repository https://gitbox.apache.org/repos/asf/superset.git
commit 50705e72530e05f280e7df27f290d96e155d8048 Author: Beto Dealmeida <[email protected]> AuthorDate: Thu Feb 6 15:13:40 2025 -0500 fix: ScreenshotCachePayload serialization (#32156) (cherry picked from commit e8990f4a365e5b0cb5922d9bbdc21075a6685b5c) --- superset/charts/api.py | 4 +-- superset/dashboards/api.py | 4 +-- superset/utils/screenshots.py | 51 +++++++++++++++++++++++++------ tests/unit_tests/utils/screenshot_test.py | 22 ++++++------- 4 files changed, 56 insertions(+), 25 deletions(-) diff --git a/superset/charts/api.py b/superset/charts/api.py index a600a3ca7f..292575feac 100644 --- a/superset/charts/api.py +++ b/superset/charts/api.py @@ -623,7 +623,7 @@ class ChartRestApi(BaseSupersetModelRestApi): 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 @@ class ChartRestApi(BaseSupersetModelRestApi): 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 c8c744ec63..c15610010b 100644 --- a/superset/dashboards/api.py +++ b/superset/dashboards/api.py @@ -1115,7 +1115,7 @@ class DashboardRestApi(BaseSupersetModelRestApi): 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 @@ class DashboardRestApi(BaseSupersetModelRestApi): "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 86f5a94ce7..7f40c0a820 100644 --- a/superset/utils/screenshots.py +++ b/superset/utils/screenshots.py @@ -20,7 +20,7 @@ import logging 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 @@ class BaseScreenshot: 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 @@ class BaseScreenshot: 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 @@ class BaseScreenshot: 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 5d29d829a2..fa928fc73f 100644 --- a/tests/unit_tests/utils/screenshot_test.py +++ b/tests/unit_tests/utils/screenshot_test.py @@ -26,7 +26,7 @@ from superset.utils.hashing import md5_sha_from_dict from superset.utils.screenshots import ( BaseScreenshot, ScreenshotCachePayload, - StatusValues, + ScreenshotCachePayloadType, ) BASE_SCREENSHOT_PATH = "superset.utils.screenshots.BaseScreenshot" @@ -121,24 +121,24 @@ class TestComputeAndCache: 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 @@ class TestComputeAndCache: # 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 @@ class TestComputeAndCache: 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)
