From 0679454b4871bf3d594d3707e9e8e51f589b92fd Mon Sep 17 00:00:00 2001
From: Geido <60598000+geido@users.noreply.github.com>
Date: Mon, 16 Sep 2024 16:09:01 +0200
Subject: [PATCH] fix(Screenshot): Dashboard screenshot cache key to include
 state (#30265)

---
 superset/dashboards/api.py    |  3 ++-
 superset/tasks/thumbnails.py  |  2 ++
 superset/utils/screenshots.py | 22 +++++++++++++++++++++-
 3 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/superset/dashboards/api.py b/superset/dashboards/api.py
index 9076303b2..5ff5b17b2 100644
--- a/superset/dashboards/api.py
+++ b/superset/dashboards/api.py
@@ -1036,7 +1036,7 @@ class DashboardRestApi(BaseSupersetModelRestApi):
 
         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 @@ class DashboardRestApi(BaseSupersetModelRestApi):
                 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
index 35781608b..34c4fc737 100644
--- 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
index a7411ef78..b31d9c0e4 100644
--- a/superset/utils/screenshots.py
+++ b/superset/utils/screenshots.py
@@ -23,6 +23,7 @@ from typing import TYPE_CHECKING
 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 @@ class BaseScreenshot:
         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 @@ class BaseScreenshot:
         :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 @@ class DashboardScreenshot(BaseScreenshot):
         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)