From 789ca738dc4331400cca042f088c1af823f11526 Mon Sep 17 00:00:00 2001
From: Geido <60598000+geido@users.noreply.github.com>
Date: Wed, 11 Sep 2024 17:35:32 +0200
Subject: [PATCH] fix(Celery): Pass guest_token as user context is not
 available in Celery (#30224)

(cherry picked from commit 1b34ad65fa9890eddf07a93738a3bc968d70a117)
---
 superset/dashboards/api.py   |  7 ++++++-
 superset/tasks/thumbnails.py | 18 +++++++++++-------
 2 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/superset/dashboards/api.py b/superset/dashboards/api.py
index 193901b48139f..43d4ecd8e8b35 100644
--- a/superset/dashboards/api.py
+++ b/superset/dashboards/api.py
@@ -22,7 +22,7 @@
 from typing import Any, Callable, cast, Optional
 from zipfile import is_zipfile, ZipFile
 
-from flask import redirect, request, Response, send_file, url_for
+from flask import g, redirect, request, Response, send_file, url_for
 from flask_appbuilder import permission_name
 from flask_appbuilder.api import expose, protect, rison, safe
 from flask_appbuilder.hooks import before_request
@@ -93,6 +93,7 @@
 from superset.extensions import event_logger
 from superset.models.dashboard import Dashboard
 from superset.models.embedded_dashboard import EmbeddedDashboard
+from superset.security.guest_token import GuestUser
 from superset.tasks.thumbnails import (
     cache_dashboard_screenshot,
     cache_dashboard_thumbnail,
@@ -1041,6 +1042,10 @@ def cache_dashboard_screenshot(self, pk: int, **kwargs: Any) -> WerkzeugResponse
         def trigger_celery() -> WerkzeugResponse:
             logger.info("Triggering screenshot ASYNC")
             cache_dashboard_screenshot.delay(
+                username=get_current_user(),
+                guest_token=g.user.guest_token
+                if isinstance(g.user, GuestUser)
+                else None,
                 dashboard_id=dashboard.id,
                 dashboard_url=dashboard_url,
                 force=True,
diff --git a/superset/tasks/thumbnails.py b/superset/tasks/thumbnails.py
index eb64ac49cb1e8..1c5db2a7f3d45 100644
--- a/superset/tasks/thumbnails.py
+++ b/superset/tasks/thumbnails.py
@@ -20,11 +20,11 @@
 import logging
 from typing import cast, Optional
 
-from flask import current_app, g
+from flask import current_app
 
 from superset import security_manager, thumbnail_cache
 from superset.extensions import celery_app
-from superset.security.guest_token import GuestUser
+from superset.security.guest_token import GuestToken
 from superset.tasks.utils import get_executor
 from superset.utils.core import override_user
 from superset.utils.screenshots import ChartScreenshot, DashboardScreenshot
@@ -86,6 +86,7 @@ def cache_dashboard_thumbnail(
     if not thumbnail_cache:
         logging.warning("No cache set, refusing to compute")
         return
+
     dashboard = Dashboard.get(dashboard_id)
     url = get_url_path("Superset.dashboard", dashboard_id_or_slug=dashboard.id)
 
@@ -110,9 +111,11 @@ def cache_dashboard_thumbnail(
 # pylint: disable=too-many-arguments
 @celery_app.task(name="cache_dashboard_screenshot", soft_time_limit=300)
 def cache_dashboard_screenshot(
+    username: str,
     dashboard_id: int,
     dashboard_url: str,
     force: bool = True,
+    guest_token: Optional[GuestToken] = None,
     thumb_size: Optional[WindowSize] = None,
     window_size: Optional[WindowSize] = None,
 ) -> None:
@@ -124,18 +127,19 @@ def cache_dashboard_screenshot(
         return
 
     dashboard = Dashboard.get(dashboard_id)
-    current_user = g.user
 
     logger.info("Caching dashboard: %s", dashboard_url)
 
     # Requests from Embedded should always use the Guest user
-    if not isinstance(current_user, GuestUser):
-        _, username = get_executor(
+    if guest_token:
+        current_user = security_manager.get_guest_user_from_token(guest_token)
+    else:
+        _, exec_username = get_executor(
             executor_types=current_app.config["THUMBNAIL_EXECUTE_AS"],
             model=dashboard,
-            current_user=current_user.username,
+            current_user=username,
         )
-        current_user = security_manager.find_user(username)
+        current_user = security_manager.find_user(exec_username)
 
     with override_user(current_user):
         screenshot = DashboardScreenshot(dashboard_url, dashboard.digest)