Skip to content

Commit

Permalink
fix(Celery): Pass guest_token as user context is not available in Cel…
Browse files Browse the repository at this point in the history
…ery (#30224)

(cherry picked from commit 1b34ad6)
  • Loading branch information
geido authored and sadpandajoe committed Sep 11, 2024
1 parent 40568fd commit 789ca73
Show file tree
Hide file tree
Showing 2 changed files with 17 additions and 8 deletions.
7 changes: 6 additions & 1 deletion superset/dashboards/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down
18 changes: 11 additions & 7 deletions superset/tasks/thumbnails.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)

Expand All @@ -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:
Expand All @@ -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)
Expand Down

0 comments on commit 789ca73

Please sign in to comment.