(superset) branch geido/fix/guest-user-celery created (now aa961575ab)

[di...@apache.org]

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

diegopucci pushed a change to branch geido/fix/guest-user-celery
in repository https://gitbox.apache.org/repos/asf/superset.git


      at aa961575ab fix(Celery): Pass guest_token as user context is not available in Celery

This branch includes the following new commits:

     new aa961575ab fix(Celery): Pass guest_token as user context is not available in Celery

The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails.  The revisions
listed as "add" were already present in the repository and have only
been added to this reference.

(superset) 01/01: fix(Celery): Pass guest_token as user context is not available in Celery

[di...@apache.org]

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

diegopucci pushed a commit to branch geido/fix/guest-user-celery
in repository https://gitbox.apache.org/repos/asf/superset.git

commit aa961575ab07b66f5d079f9eee5b56e42e86188f
Author: Diego Pucci <di...@gmail.com>
AuthorDate: Tue Sep 10 22:49:19 2024 +0200

    fix(Celery): Pass guest_token as user context is not available in Celery
---
 superset/dashboards/api.py   |  6 +++++-
 superset/tasks/thumbnails.py | 18 +++++++++++-------
 2 files changed, 16 insertions(+), 8 deletions(-)

diff --git a/superset/dashboards/api.py b/superset/dashboards/api.py
index 193901b481..d67da01436 100644
--- a/superset/dashboards/api.py
+++ b/superset/dashboards/api.py
@@ -22,7 +22,7 @@ from io import BytesIO
 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.dashboards.schemas import (
 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,
@@ -1037,10 +1038,13 @@ class DashboardRestApi(BaseSupersetModelRestApi):
         image_url = get_url_path(
             "DashboardRestApi.screenshot", pk=dashboard.id, digest=cache_key
         )
+        is_guest_user = isinstance(g.user, GuestUser)
 
         def trigger_celery() -> WerkzeugResponse:
             logger.info("Triggering screenshot ASYNC")
             cache_dashboard_screenshot.delay(
+                username=get_current_user(),
+                guest_token=g.user.guest_token if is_guest_user 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 eb64ac49cb..40a01b8720 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,6 +111,8 @@ 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,
+    guest_token: Optional[GuestToken],
     dashboard_id: int,
     dashboard_url: str,
     force: bool = True,
@@ -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)