From aa961575ab07b66f5d079f9eee5b56e42e86188f Mon Sep 17 00:00:00 2001 From: Diego Pucci Date: Tue, 10 Sep 2024 22:49:19 +0200 Subject: [PATCH 1/3] 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 193901b48139f..d67da01436985 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, @@ -1037,10 +1038,13 @@ def cache_dashboard_screenshot(self, pk: int, **kwargs: Any) -> WerkzeugResponse 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 eb64ac49cb1e8..40a01b87209a6 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) From 4b1799f1fc97c718a911574d5697c22cc00659bd Mon Sep 17 00:00:00 2001 From: Diego Pucci Date: Tue, 10 Sep 2024 23:03:48 +0200 Subject: [PATCH 2/3] fix(Pylint): Lint --- superset/dashboards/api.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/superset/dashboards/api.py b/superset/dashboards/api.py index d67da01436985..43d4ecd8e8b35 100644 --- a/superset/dashboards/api.py +++ b/superset/dashboards/api.py @@ -1038,13 +1038,14 @@ def cache_dashboard_screenshot(self, pk: int, **kwargs: Any) -> WerkzeugResponse 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, + guest_token=g.user.guest_token + if isinstance(g.user, GuestUser) + else None, dashboard_id=dashboard.id, dashboard_url=dashboard_url, force=True, From bb1adae7f626af16c6fd14a4b7a677072ce5aae0 Mon Sep 17 00:00:00 2001 From: Diego Pucci Date: Wed, 11 Sep 2024 10:46:03 +0200 Subject: [PATCH 3/3] chore: Set gues_token to None default --- superset/tasks/thumbnails.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/superset/tasks/thumbnails.py b/superset/tasks/thumbnails.py index 40a01b87209a6..1c5db2a7f3d45 100644 --- a/superset/tasks/thumbnails.py +++ b/superset/tasks/thumbnails.py @@ -112,10 +112,10 @@ def cache_dashboard_thumbnail( @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, + guest_token: Optional[GuestToken] = None, thumb_size: Optional[WindowSize] = None, window_size: Optional[WindowSize] = None, ) -> None: