(superset) 09/22: add dashboard screenshot celery task

[di...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

diegopucci pushed a commit to branch geido/feat/dashboard-endpoint-screenshot
in repository https://gitbox.apache.org/repos/asf/superset.git

commit 03b09b0d3e09814c5696f507a20f4a5b3e30d8d7
Author: Edgar Ulloa <ul...@gmail.com>
AuthorDate: Thu Jun 6 23:23:38 2024 -0700

    add dashboard screenshot celery task
---
 superset/dashboards/api.py   | 15 ++++++++++++---
 superset/tasks/thumbnails.py | 36 ++++++++++++++++++++++++++++++++++++
 2 files changed, 48 insertions(+), 3 deletions(-)

diff --git a/superset/dashboards/api.py b/superset/dashboards/api.py
index de7b660cdf..bdcbcbf745 100644
--- a/superset/dashboards/api.py
+++ b/superset/dashboards/api.py
@@ -85,7 +85,10 @@ from superset.dashboards.schemas import (
 from superset.extensions import event_logger
 from superset.models.dashboard import Dashboard
 from superset.models.embedded_dashboard import EmbeddedDashboard
-from superset.tasks.thumbnails import cache_dashboard_thumbnail
+from superset.tasks.thumbnails import (
+    cache_dashboard_screenshot,
+    cache_dashboard_thumbnail,
+)
 from superset.tasks.utils import get_current_user
 from superset.utils import json
 from superset.utils.pdf import build_pdf_from_screenshots
@@ -888,7 +891,7 @@ class DashboardRestApi(BaseSupersetModelRestApi):
         )
 
     # TODO: cache_dashboard_screenshot endpoint
-    @expose("/<pk>/cache_screenshot/", methods=("POST",))
+    @expose("/<pk>/cache_dashboard_screenshot/", methods=("POST",))
     # @protect()
     @rison(screenshot_query_schema)
     @safe
@@ -948,10 +951,12 @@ class DashboardRestApi(BaseSupersetModelRestApi):
         """
         try:
             payload = CacheScreenshotSchema().load(request.json)
+            print("PAYLOAD: ", payload)
         except ValidationError as error:
             return self.response_400(message=error.messages)
 
         dashboard = cast(Dashboard, self.datamodel.get(pk, self._base_filters))
+        print("DASHBOARD: ", dashboard)
         if not dashboard:
             return self.response_404()
 
@@ -972,7 +977,9 @@ class DashboardRestApi(BaseSupersetModelRestApi):
             dashboard_id=str(dashboard.id),
             state=dashboard_state,
         ).run()
+        print("PERMALINK KEY:", permalink_key)
         dashboard_url = get_url_path("Superset.dashboard_permalink", key=permalink_key)
+        print("DASHBOARD_URL: ", dashboard_url)
         # dashboard_url = get_url_path(
         #     "Superset.dashboard", dashboard_id_or_slug=dashboard.id
         # )
@@ -984,9 +991,11 @@ class DashboardRestApi(BaseSupersetModelRestApi):
 
         def trigger_celery() -> WerkzeugResponse:
             logger.info("Triggering screenshot ASYNC")
-            cache_dashboard_thumbnail.delay(
+            cache_dashboard_screenshot.delay(
                 current_user=get_current_user(),
                 dashboard_id=dashboard.id,
+                dashboard_url=dashboard_url,
+                cache_key=cache_key,
                 force=True,
                 thumb_size=thumb_size,
                 window_size=window_size,
diff --git a/superset/tasks/thumbnails.py b/superset/tasks/thumbnails.py
index 145f4d1aad..6b70b92231 100644
--- a/superset/tasks/thumbnails.py
+++ b/superset/tasks/thumbnails.py
@@ -104,3 +104,39 @@ def cache_dashboard_thumbnail(
             window_size=window_size,
             thumb_size=thumb_size,
         )
+
+
+@celery_app.task(name="cache_dashboard_screenshot", soft_time_limit=300)
+def cache_dashboard_screenshot(
+    current_user: Optional[str],
+    dashboard_id: int,
+    dashboard_url: str,
+    cache_key: str,
+    force: bool = False,
+    thumb_size: Optional[WindowSize] = None,
+    window_size: Optional[WindowSize] = None,
+) -> None:
+    # pylint: disable=import-outside-toplevel
+    from superset.models.dashboard import Dashboard
+
+    if not thumbnail_cache:
+        logging.warning("No cache set, refusing to compute")
+        return
+    dashboard = Dashboard.get(dashboard_id)
+
+    logger.info("Caching dashboard: %s", dashboard_url)
+    _, username = get_executor(
+        executor_types=current_app.config["THUMBNAIL_EXECUTE_AS"],
+        model=dashboard,
+        current_user=current_user,
+    )
+    user = security_manager.find_user(username)
+    with override_user(user):
+        screenshot = DashboardScreenshot(dashboard_url, cache_key)
+        screenshot.compute_and_cache(
+            user=user,
+            cache=thumbnail_cache,
+            force=force,
+            window_size=window_size,
+            thumb_size=thumb_size,
+        )