enabling force parameter on screenshot endpoints

adding CacheScreenshotPayload state machine

superset/charts/api.py

@@ -30,7 +30,7 @@
 from werkzeug.wrappers import Response as WerkzeugResponse
 from werkzeug.wsgi import FileWrapper
 
-from superset import app, is_feature_enabled, thumbnail_cache
+from superset import app, is_feature_enabled
 from superset.charts.filters import (
     ChartAllTextFilter,
     ChartCertifiedFilter,
@@ -84,7 +84,12 @@
 from superset.tasks.thumbnails import cache_chart_thumbnail
 from superset.tasks.utils import get_current_user
 from superset.utils import json
-from superset.utils.screenshots import ChartScreenshot, DEFAULT_CHART_WINDOW_SIZE
+from superset.utils.screenshots import (
+    ChartScreenshot,
+    DEFAULT_CHART_WINDOW_SIZE,
+    ScreenshotCachePayload,
+    StatusValues,
+)
 from superset.utils.urls import get_url_path
 from superset.views.base_api import (
     BaseSupersetModelRestApi,
@@ -564,8 +569,14 @@ def cache_screenshot(self, pk: int, **kwargs: Any) -> WerkzeugResponse:
                 schema:
                   $ref: '#/components/schemas/screenshot_query_schema'
           responses:
+            200:
+                description: Chart async result
+                content:
+                  application/json:
+                    schema:
+                      $ref: "#/components/schemas/ChartCacheScreenshotResponseSchema"
             202:
-              description: Chart async result
+              description: Chart async created
               content:
                 application/json:
                   schema:
@@ -580,6 +591,7 @@ def cache_screenshot(self, pk: int, **kwargs: Any) -> WerkzeugResponse:
               $ref: '#/components/responses/500'
         """
         rison_dict = kwargs["rison"]
+        force = rison_dict.get("force")
         window_size = rison_dict.get("window_size") or DEFAULT_CHART_WINDOW_SIZE
 
         # Don't shrink the image if thumb_size is not specified
@@ -591,25 +603,38 @@ def cache_screenshot(self, pk: int, **kwargs: Any) -> WerkzeugResponse:
 
         chart_url = get_url_path("Superset.slice", slice_id=chart.id)
         screenshot_obj = ChartScreenshot(chart_url, chart.digest)
-        cache_key = screenshot_obj.cache_key(window_size, thumb_size)
+        cache_key = screenshot_obj.get_cache_key(window_size, thumb_size)
         image_url = get_url_path(
             "ChartRestApi.screenshot", pk=chart.id, digest=cache_key
         )
 
-        def trigger_celery() -> WerkzeugResponse:
-            logger.info("Triggering screenshot ASYNC")
-            cache_chart_thumbnail.delay(
-                current_user=get_current_user(),
-                chart_id=chart.id,
-                force=True,
-                window_size=window_size,
-                thumb_size=thumb_size,
-            )
-            return self.response(
-                202, cache_key=cache_key, chart_url=chart_url, image_url=image_url
-            )
+        if force or not screenshot_obj.get_from_cache_key(cache_key):
+            payload = ScreenshotCachePayload()
+            screenshot_obj.cache.set(cache_key, payload)
 
-        return trigger_celery()
+        if (cache_payload := screenshot_obj.get_from_cache_key(cache_key)) is not None:
+
+            def build_response(status_code: int) -> WerkzeugResponse:
+                return self.response(
+                    status_code,
+                    cache_key=cache_key,
+                    chart_url=chart_url,
+                    image_url=image_url,
+                    updated_at=cache_payload.get_timestamp(),
+                    update_status=cache_payload.get_status(),
+                )
+
+            if cache_payload.status != StatusValues.UPDATED:
+                logger.info("Triggering screenshot ASYNC")
+                cache_chart_thumbnail.delay(
+                    current_user=get_current_user(),
+                    chart_id=chart.id,
+                    window_size=window_size,
+                    thumb_size=thumb_size,
+                )
+                return build_response(202)
+            return build_response(200)
+        return self.response_500()
 
     @expose("/<pk>/screenshot/<digest>/", methods=("GET",))
     @protect()
@@ -635,7 +660,7 @@ def screenshot(self, pk: int, digest: str) -> WerkzeugResponse:
             name: digest
           responses:
             200:
-              description: Chart thumbnail image
+              description: Chart screenshot image
               content:
                image/*:
                  schema:
@@ -652,16 +677,16 @@ def screenshot(self, pk: int, digest: str) -> WerkzeugResponse:
         """
         chart = self.datamodel.get(pk, self._base_filters)
 
-        # Making sure the chart still exists
         if not chart:
             return self.response_404()
 
-        # fetch the chart screenshot using the current user and cache if set
-        if img := ChartScreenshot.get_from_cache_key(thumbnail_cache, digest):
-            return Response(
-                FileWrapper(img), mimetype="image/png", direct_passthrough=True
-            )
-        # TODO: return an empty image
+        if cache_payload := ChartScreenshot.get_from_cache_key(digest):
+            if cache_payload.status == StatusValues.UPDATED:
+                return Response(
+                    FileWrapper(cache_payload.get_image()),
+                    mimetype="image/png",
+                    direct_passthrough=True,
+                )
         return self.response_404()
 
     @expose("/<pk>/thumbnail/<digest>/", methods=("GET",))
@@ -713,33 +738,32 @@ def thumbnail(self, pk: int, digest: str, **kwargs: Any) -> WerkzeugResponse:
 
         current_user = get_current_user()
         url = get_url_path("Superset.slice", slice_id=chart.id)
-        if kwargs["rison"].get("force", False):
-            logger.info(
-                "Triggering thumbnail compute (chart id: %s) ASYNC", str(chart.id)
-            )
-            cache_chart_thumbnail.delay(
-                current_user=current_user,
-                chart_id=chart.id,
-                force=True,
-            )
-            return self.response(202, message="OK Async")
-        # fetch the chart screenshot using the current user and cache if set
-        screenshot = ChartScreenshot(url, chart.digest).get_from_cache(
-            cache=thumbnail_cache
-        )
-        # If not screenshot then send request to compute thumb to celery
-        if not screenshot:
+        screenshot_obj = ChartScreenshot(url, chart.digest)
+        cache_key = screenshot_obj.get_cache_key()
+        cache_payload = screenshot_obj.get_from_cache_key(cache_key)
+
+        if not cache_payload:
+            cache_payload = ScreenshotCachePayload()
+            screenshot_obj.cache.set(cache_key, cache_payload)
+
+        if (
+            kwargs["rison"].get("force", False)
+            or cache_payload.status != StatusValues.UPDATED
+        ):
             self.incr_stats("async", self.thumbnail.__name__)
             logger.info(
                 "Triggering thumbnail compute (chart id: %s) ASYNC", str(chart.id)
             )
             cache_chart_thumbnail.delay(
                 current_user=current_user,
                 chart_id=chart.id,
-                force=True,
             )
-            return self.response(202, message="OK Async")
-        # If digests
+            return self.response(
+                202,
+                updated_at=cache_payload.get_timestamp(),
+                update_status=cache_payload.get_status(),
+            )
+        # TODO remove digest from params... currently does nothing
         if chart.digest != digest:
             self.incr_stats("redirect", self.thumbnail.__name__)
             return redirect(
@@ -749,7 +773,9 @@ def thumbnail(self, pk: int, digest: str, **kwargs: Any) -> WerkzeugResponse:
             )
         self.incr_stats("from_cache", self.thumbnail.__name__)
         return Response(
-            FileWrapper(screenshot), mimetype="image/png", direct_passthrough=True
+            FileWrapper(cache_payload.get_image()),
+            mimetype="image/png",
+            direct_passthrough=True,
         )
 
     @expose("/export/", methods=("GET",))

superset/charts/schemas.py

@@ -304,6 +304,21 @@ class ChartCacheScreenshotResponseSchema(Schema):
     image_url = fields.String(
         metadata={"description": "The url to fetch the screenshot"}
     )
+    update_status = fields.String(
+        metadata={"description": "The status of the async screenshot"}
+    )
+    updated_at = fields.String(
+        metadata={"description": "The timestamp of the last change in status"}
+    )
+
+
+class ChartGetCachedScreenshotResponseSchema(Schema):
+    update_status = fields.String(
+        metadata={"description": "The status of the async screenshot"}
+    )
+    updated_at = fields.String(
+        metadata={"description": "The timestamp of the last change in status"}
+    )
 
 
 class ChartDataColumnSchema(Schema):

superset/config.py

@@ -716,6 +716,7 @@ class D3TimeFormat(TypedDict, total=False):
 
 THUMBNAIL_CACHE_CONFIG: CacheConfig = {
     "CACHE_TYPE": "NullCache",
+    "CACHE_DEFAULT_TIMEOUT": 3600,
     "CACHE_NO_NULL_WARNING": True,
 }
 

superset/dashboards/api.py

@@ -32,7 +32,7 @@
 from werkzeug.wrappers import Response as WerkzeugResponse
 from werkzeug.wsgi import FileWrapper
 
-from superset import db, is_feature_enabled, thumbnail_cache
+from superset import db, is_feature_enabled
 from superset.charts.schemas import ChartEntityResponseSchema
 from superset.commands.dashboard.copy import CopyDashboardCommand
 from superset.commands.dashboard.create import CreateDashboardCommand
@@ -116,6 +116,8 @@
 from superset.utils.screenshots import (
     DashboardScreenshot,
     DEFAULT_DASHBOARD_WINDOW_SIZE,
+    ScreenshotCachePayload,
+    StatusValues,
 )
 from superset.utils.urls import get_url_path
 from superset.views.base_api import (
@@ -1105,27 +1107,24 @@ def thumbnail(self, pk: int, digest: str, **kwargs: Any) -> WerkzeugResponse:
         )
         # If force, request a screenshot from the workers
         current_user = get_current_user()
-        if kwargs["rison"].get("force", False):
+        screenshot_obj = DashboardScreenshot(dashboard_url, digest)
+        cache_key = screenshot_obj.get_cache_key()
+        cache_payload = screenshot_obj.get_from_cache_key(cache_key)
+        if not cache_payload or kwargs["rison"].get("force", False):
+            cache_payload = ScreenshotCachePayload()
+            screenshot_obj.cache.set(cache_key, cache_payload)
+
+        if cache_payload.status != StatusValues.UPDATED:
             cache_dashboard_thumbnail.delay(
                 current_user=current_user,
                 dashboard_id=dashboard.id,
-                force=True,
             )
-            return self.response(202, message="OK Async")
-        # fetch the dashboard screenshot using the current user and cache if set
-        screenshot = DashboardScreenshot(
-            dashboard_url, dashboard.digest
-        ).get_from_cache(cache=thumbnail_cache)
-        # If the screenshot does not exist, request one from the workers
-        if not screenshot:
-            self.incr_stats("async", self.thumbnail.__name__)
-            cache_dashboard_thumbnail.delay(
-                current_user=current_user,
-                dashboard_id=dashboard.id,
-                force=True,
+            return self.response(
+                202,
+                updated_at=cache_payload.get_timestamp(),
+                update_status=cache_payload.get_status(),
             )
-            return self.response(202, message="OK Async")
-        # If digests
+
         if dashboard.digest != digest:
             self.incr_stats("redirect", self.thumbnail.__name__)
             return redirect(
@@ -1137,7 +1136,9 @@ def thumbnail(self, pk: int, digest: str, **kwargs: Any) -> WerkzeugResponse:
             )
         self.incr_stats("from_cache", self.thumbnail.__name__)
         return Response(
-            FileWrapper(screenshot), mimetype="image/png", direct_passthrough=True
+            FileWrapper(cache_payload.get_image()),
+            mimetype="image/png",
+            direct_passthrough=True,
         )
 
     @expose("/<pk>/cache_dashboard_screenshot/", methods=("POST",))
@@ -1150,7 +1151,9 @@ def thumbnail(self, pk: int, digest: str, **kwargs: Any) -> WerkzeugResponse:
         f".cache_dashboard_screenshot",
         log_to_statsd=False,
     )
-    def cache_dashboard_screenshot(self, pk: int, **kwargs: Any) -> WerkzeugResponse:
+    def cache_dashboard_screenshot(
+        self, pk: int, force: bool, **kwargs: Any
+    ) -> WerkzeugResponse:
         """Compute and cache a screenshot.
         ---
         post:
@@ -1185,7 +1188,6 @@ def cache_dashboard_screenshot(self, pk: int, **kwargs: Any) -> WerkzeugResponse
             payload = CacheScreenshotSchema().load(request.json)
         except ValidationError as error:
             return self.response_400(message=error.messages)
-
         dashboard = cast(Dashboard, self.datamodel.get(pk, self._base_filters))
         if not dashboard:
             return self.response_404()
@@ -1210,10 +1212,17 @@ 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, dashboard_state)
+        cache_key = screenshot_obj.get_cache_key(
+            window_size, thumb_size, dashboard_state
+        )
         image_url = get_url_path(
             "DashboardRestApi.screenshot", pk=dashboard.id, digest=cache_key
         )
+        force = kwargs["rison"].get("force", False)
+        cache_payload = screenshot_obj.get_from_cache_key(cache_key)
+        if force or cache_payload is None:
+            cache_payload = ScreenshotCachePayload()
+            screenshot_obj.cache.set(cache_key, cache_payload)
 
         def trigger_celery() -> WerkzeugResponse:
             logger.info("Triggering screenshot ASYNC")
@@ -1226,8 +1235,6 @@ def trigger_celery() -> WerkzeugResponse:
                 ),
                 dashboard_id=dashboard.id,
                 dashboard_url=dashboard_url,
-                cache_key=cache_key,
-                force=False,
                 thumb_size=thumb_size,
                 window_size=window_size,
             )
@@ -1236,6 +1243,8 @@ def trigger_celery() -> WerkzeugResponse:
                 cache_key=cache_key,
                 dashboard_url=dashboard_url,
                 image_url=image_url,
+                updated_at=cache_payload.get_timestamp(),
+                update_status=cache_payload.get_status(),
             )
 
         return trigger_celery()
@@ -1294,9 +1303,12 @@ def screenshot(self, pk: int, digest: str) -> WerkzeugResponse:
 
         # fetch the dashboard screenshot using the current user and cache if set
 
-        if img := DashboardScreenshot.get_from_cache_key(thumbnail_cache, digest):
+        if cache_payload := DashboardScreenshot.get_from_cache_key(digest):
+            image = cache_payload.get_image()
+            if not image:
+                return self.response_404()
             if download_format == "pdf":
-                pdf_img = img.getvalue()
+                pdf_img = image.getvalue()
                 # Convert the screenshot to PDF
                 pdf_data = build_pdf_from_screenshots([pdf_img])
 
@@ -1308,7 +1320,7 @@ def screenshot(self, pk: int, digest: str) -> WerkzeugResponse:
                 )
             if download_format == "png":
                 return Response(
-                    FileWrapper(img),
+                    FileWrapper(image),
                     mimetype="image/png",
                     direct_passthrough=True,
                 )

superset/dashboards/schemas.py

@@ -507,6 +507,12 @@ class DashboardCacheScreenshotResponseSchema(Schema):
     image_url = fields.String(
         metadata={"description": "The url to fetch the screenshot"}
     )
+    update_status = fields.String(
+        metadata={"description": "The status of the async screenshot"}
+    )
+    updated_at = fields.String(
+        metadata={"description": "The timestamp of the last change in status"}
+    )
 
 
 class CacheScreenshotSchema(Schema):

superset/models/dashboard.py

@@ -350,7 +350,6 @@ def update_thumbnail(self) -> None:
         cache_dashboard_thumbnail.delay(
             current_user=get_current_user(),
             dashboard_id=self.id,
-            force=True,
         )
 
     @classmethod