fix(Screenshot): Dashboard screenshot cache key to include state (#30265

)
apache · Sep 16, 2024 · 0679454 · 0679454
1 parent 3ad0fd4
commit 0679454
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 2 deletions.
diff --git a/superset/dashboards/api.py b/superset/dashboards/api.py
@@ -1036,7 +1036,7 @@ def cache_dashboard_screenshot(self, pk: int, **kwargs: Any) -> WerkzeugResponse
 
         dashboard_url = get_url_path("Superset.dashboard_permalink", key=permalink_key)
         screenshot_obj = DashboardScreenshot(dashboard_url, dashboard.digest)
-        cache_key = screenshot_obj.cache_key(window_size, thumb_size)
+        cache_key = screenshot_obj.cache_key(window_size, thumb_size, dashboard_state)
         image_url = get_url_path(
             "DashboardRestApi.screenshot", pk=dashboard.id, digest=cache_key
         )
@@ -1053,6 +1053,7 @@ def trigger_celery() -> WerkzeugResponse:
                 force=True,
                 thumb_size=thumb_size,
                 window_size=window_size,
+                cache_key=cache_key,
             )
             return self.response(
                 202,

diff --git a/superset/tasks/thumbnails.py b/superset/tasks/thumbnails.py
@@ -117,6 +117,7 @@ def cache_dashboard_screenshot(  # pylint: disable=too-many-arguments
     guest_token: Optional[GuestToken] = None,
     thumb_size: Optional[WindowSize] = None,
     window_size: Optional[WindowSize] = None,
+    cache_key: Optional[str] = None,
 ) -> None:
     # pylint: disable=import-outside-toplevel
     from superset.models.dashboard import Dashboard
@@ -148,4 +149,5 @@ def cache_dashboard_screenshot(  # pylint: disable=too-many-arguments
             force=force,
             window_size=window_size,
             thumb_size=thumb_size,
+            cache_key=cache_key,
         )
diff --git a/superset/utils/screenshots.py b/superset/utils/screenshots.py
@@ -23,6 +23,7 @@
 from flask import current_app
 
 from superset import feature_flag_manager
+from superset.dashboards.permalink.types import DashboardPermalinkState
 from superset.utils.hashing import md5_sha_from_dict
 from superset.utils.urls import modify_url_query
 from superset.utils.webdriver import (
@@ -144,6 +145,7 @@ def compute_and_cache(  # pylint: disable=too-many-arguments
         thumb_size: WindowSize | None = None,
         cache: Cache = None,
         force: bool = True,
+        cache_key: str | None = None,
     ) -> bytes | None:
         """
         Fetches the screenshot, computes the thumbnail and caches the result
@@ -155,7 +157,7 @@ def compute_and_cache(  # pylint: disable=too-many-arguments
         :param force: Will force the computation even if it's already cached
         :return: Image payload
         """
-        cache_key = self.cache_key(window_size, thumb_size)
+        cache_key = cache_key or self.cache_key(window_size, thumb_size)
         window_size = window_size or self.window_size
         thumb_size = thumb_size or self.thumb_size
         if not force and cache and cache.get(cache_key):
@@ -251,3 +253,21 @@ def __init__(
         super().__init__(url, digest)
         self.window_size = window_size or DEFAULT_DASHBOARD_WINDOW_SIZE
         self.thumb_size = thumb_size or DEFAULT_DASHBOARD_THUMBNAIL_SIZE
+
+    def cache_key(
+        self,
+        window_size: bool | WindowSize | None = None,
+        thumb_size: bool | WindowSize | None = None,
+        dashboard_state: DashboardPermalinkState | None = None,
+    ) -> str:
+        window_size = window_size or self.window_size
+        thumb_size = thumb_size or self.thumb_size
+        args = {
+            "thumbnail_type": self.thumbnail_type,
+            "digest": self.digest,
+            "type": "thumb",
+            "window_size": window_size,
+            "thumb_size": thumb_size,
+            "dashboard_state": dashboard_state,
+        }
+        return md5_sha_from_dict(args)