villebro commented on code in PR #31757:
URL: https://github.com/apache/superset/pull/31757#discussion_r1937700247


##########
superset/utils/screenshots.py:
##########
@@ -54,6 +56,68 @@
     from flask_caching import Cache
 
 
+class StatusValues(Enum):
+    PENDING = "Pending"
+    COMPUTING = "Computing"
+    UPDATED = "Updated"
+    ERROR = "Error"
+
+
+class ScreenshotCachePayload:
+    def __init__(self, image: bytes | None = None):
+        self._image = image
+        self._timestamp = datetime.now().isoformat()
+        self.status = StatusValues.PENDING
+        if image:
+            self.status = StatusValues.UPDATED
+
+    def update_timestamp(self) -> None:
+        self._timestamp = datetime.now().isoformat()
+
+    def pending(self) -> None:
+        self.update_timestamp()
+        self._image = None
+        self.status = StatusValues.PENDING
+
+    def computing(self) -> None:
+        self.update_timestamp()
+        self._image = None
+        self.status = StatusValues.COMPUTING
+
+    def update(self, image: bytes) -> None:
+        self.update_timestamp()
+        self.status = StatusValues.UPDATED
+        self._image = image
+
+    def error(
+        self,
+    ) -> None:
+        self.update_timestamp()
+        self.status = StatusValues.ERROR
+
+    def get_image(self) -> BytesIO | None:
+        if not self._image:
+            return None
+        return BytesIO(self._image)
+
+    def get_timestamp(self) -> str:
+        return self._timestamp
+
+    def get_status(self) -> str:
+        return self.status.value
+
+    def should_trigger_task(self) -> bool:
+        def is_error_cache_ttl_expired() -> bool:
+            error_cache_ttl = app.config["THUMBNAIL_ERROR_CACHE_TTL"]
+            return (
+                datetime.now() - datetime.fromisoformat(self.get_timestamp())
+            ).total_seconds() > error_cache_ttl

Review Comment:
   Is there some reason we're nesting this? Couldn't it just be a regular 
method?



##########
superset/dashboards/api.py:
##########
@@ -1197,13 +1093,29 @@ 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
         )
+        cache_payload = (
+            screenshot_obj.get_from_cache_key(cache_key) or 
ScreenshotCachePayload()
+        )
+
+        def build_response(status_code: int) -> WerkzeugResponse:
+            return self.response(
+                status_code,
+                cache_key=cache_key,
+                dashboard_url=dashboard_url,
+                image_url=image_url,
+                task_updated_at=cache_payload.get_timestamp(),
+                task_status=cache_payload.get_status(),
+            )
 
-        def trigger_celery() -> WerkzeugResponse:
+        if cache_payload.should_trigger_task() or force:

Review Comment:
   same
   ```suggestion
           if force or cache_payload.should_trigger_task():
   ```



##########
superset/charts/api.py:
##########
@@ -591,25 +603,36 @@ 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)
+        cache_payload = (
+            screenshot_obj.get_from_cache_key(cache_key) or 
ScreenshotCachePayload()
+        )
         image_url = get_url_path(
             "ChartRestApi.screenshot", pk=chart.id, digest=cache_key
         )
 
-        def trigger_celery() -> WerkzeugResponse:
+        def build_response(status_code: int) -> WerkzeugResponse:
+            return self.response(
+                status_code,
+                cache_key=cache_key,
+                chart_url=chart_url,
+                image_url=image_url,
+                task_updated_at=cache_payload.get_timestamp(),
+                task_status=cache_payload.get_status(),
+            )
+
+        if cache_payload.should_trigger_task() or force:

Review Comment:
   very very nit, but flipping this around will likely save some CPU. To DRY it 
even more, we could maybe even consider making `force` an optional parameter in 
`should_trigger_task` that defaults to `False` (not highly opinionated here)
   ```suggestion
           if force or cache_payload.should_trigger_task():
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to