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(alerts, thumbnails): celery tasks using flask-sqlalchemy sessions #11950

Merged
merged 2 commits into from
Dec 8, 2020
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
2 changes: 1 addition & 1 deletion superset/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -369,7 +369,7 @@ def _try_json_readsha( # pylint: disable=unused-argument
# ---------------------------------------------------
# Thumbnail config (behind feature flag)
# ---------------------------------------------------
THUMBNAIL_SELENIUM_USER = "Admin"
THUMBNAIL_SELENIUM_USER = "admin"
THUMBNAIL_CACHE_CONFIG: CacheConfig = {
"CACHE_TYPE": "null",
"CACHE_NO_NULL_WARNING": True,
Expand Down
17 changes: 17 additions & 0 deletions superset/security/manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
assoc_permissionview_role,
assoc_user_role,
PermissionView,
User,
)
from flask_appbuilder.security.views import (
PermissionModelView,
Expand All @@ -38,6 +39,7 @@
from flask_appbuilder.widgets import ListWidget
from sqlalchemy import and_, or_
from sqlalchemy.engine.base import Connection
from sqlalchemy.orm import Session
from sqlalchemy.orm.mapper import Mapper
from sqlalchemy.orm.query import Query as SqlaQuery

Expand Down Expand Up @@ -986,6 +988,21 @@ def raise_for_access( # pylint: disable=too-many-arguments,too-many-branches
self.get_datasource_access_error_object(datasource)
)

def get_user_by_username(
self, username: str, session: Session = None
) -> Optional[User]:
"""
Retrieves a user by it's username case sensitive. Optional session parameter
utility method normally useful for celery tasks where the session
need to be scoped
"""
session = session or self.get_session
return (
session.query(self.user_model)
.filter(self.user_model.username == username)
.one_or_none()
)

def get_rls_filters(self, table: "BaseDatasource") -> List[SqlaQuery]:
"""
Retrieves the appropriate row level security filters for the current user and
Expand Down
4 changes: 3 additions & 1 deletion superset/tasks/schedules.py
Original file line number Diff line number Diff line change
Expand Up @@ -385,7 +385,9 @@ def _get_slice_screenshot(slice_id: int, session: Session) -> ScreenshotData:
"Superset.slice", user_friendly=True, slice_id=slice_obj.id,
)

user = security_manager.find_user(current_app.config["THUMBNAIL_SELENIUM_USER"])
user = security_manager.get_user_by_username(
current_app.config["THUMBNAIL_SELENIUM_USER"], session=session
)
image_data = screenshot.compute_and_cache(
user=user, cache=thumbnail_cache, force=True,
)
Expand Down
11 changes: 9 additions & 2 deletions superset/tasks/thumbnails.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@

from superset import app, security_manager, thumbnail_cache
from superset.extensions import celery_app
from superset.utils.celery import session_scope
from superset.utils.screenshots import ChartScreenshot, DashboardScreenshot
from superset.utils.webdriver import WindowSize

Expand All @@ -44,7 +45,10 @@ def cache_chart_thumbnail(
return None
logger.info("Caching chart: %s", url)
screenshot = ChartScreenshot(url, digest)
user = security_manager.find_user(current_app.config["THUMBNAIL_SELENIUM_USER"])
with session_scope(nullpool=True) as session:
user = security_manager.get_user_by_username(
current_app.config["THUMBNAIL_SELENIUM_USER"], session=session
)
screenshot.compute_and_cache(
user=user,
cache=thumbnail_cache,
Expand All @@ -65,7 +69,10 @@ def cache_dashboard_thumbnail(
return
logger.info("Caching dashboard: %s", url)
screenshot = DashboardScreenshot(url, digest)
user = security_manager.find_user(current_app.config["THUMBNAIL_SELENIUM_USER"])
with session_scope(nullpool=True) as session:
user = security_manager.get_user_by_username(
current_app.config["THUMBNAIL_SELENIUM_USER"], session=session
)
screenshot.compute_and_cache(
user=user, cache=thumbnail_cache, force=force, thumb_size=thumb_size,
)