iop
/
superset
4
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: ScreenshotCachePayload serialization (#32156)

This commit is contained in:
Beto Dealmeida 2025-02-06 15:13:40 -05:00 committed by GitHub
parent acf91e1f60
commit e8990f4a36
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 56 additions and 25 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
superset/charts/api.py
View File

@ -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,

4
superset/dashboards/api.py
View File

@ -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,

51
superset/utils/screenshots.py
View File

@ -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

22
tests/unit_tests/utils/screenshot_test.py
View File

@ -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)