Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix(Celery): Pass guest_token as user context is not available in Celery #30224

Merged
merged 3 commits into from
Sep 11, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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
Loading