fix: ScreenshotCachePayload serialization (#32156)

(cherry picked from commit e8990f4)

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,

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,

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
 

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)