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

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

* selenium user is case sensitive

superset/config.py

@@ -369,7 +369,7 @@ GET_FEATURE_FLAGS_FUNC: Optional[Callable[[Dict[str, bool]], Dict[str, bool]]] =
 # ---------------------------------------------------
 # 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,

superset/security/manager.py

@@ -27,6 +27,7 @@ from flask_appbuilder.security.sqla.models import (
     assoc_permissionview_role,
     assoc_user_role,
     PermissionView,
+    User,
 )
 from flask_appbuilder.security.views import (
     PermissionModelView,
@@ -38,6 +39,7 @@ from flask_appbuilder.security.views import (
 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
 
@@ -986,6 +988,21 @@ class SupersetSecurityManager(  # pylint: disable=too-many-public-methods
                     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

superset/tasks/schedules.py

@@ -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,
     )

superset/tasks/thumbnails.py

@@ -24,6 +24,7 @@ from flask import current_app
 
 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
 
@@ -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,
@@ -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,
         )