Skip to content

Commit

Permalink
fix(alerts, thumbnails): celery tasks using flask-sqlalchemy sessions (
Browse files Browse the repository at this point in the history
…#11950)

* fix(alerts, thumbnails): celery tasks using flask-sqlalchemy sessions

* selenium user is case sensitive
  • Loading branch information
dpgaspar authored Dec 8, 2020
1 parent 05258eb commit 7e6f04f
Show file tree
Hide file tree
Showing 4 changed files with 30 additions and 4 deletions.
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,
)

0 comments on commit 7e6f04f

Please sign in to comment.