chore: Cleanup database sessions (#10427)

Co-authored-by: John Bodley <john.bodley@airbnb.com>

superset/tasks/cache.py

@@ -134,8 +134,7 @@ class DummyStrategy(Strategy):
     name = "dummy"
 
     def get_urls(self) -> List[str]:
-        session = db.create_scoped_session()
-        charts = session.query(Slice).all()
+        charts = db.session.query(Slice).all()
 
         return [get_url(chart) for chart in charts]
 
@@ -167,10 +166,9 @@ class TopNDashboardsStrategy(Strategy):
 
     def get_urls(self) -> List[str]:
         urls = []
-        session = db.create_scoped_session()
 
         records = (
-            session.query(Log.dashboard_id, func.count(Log.dashboard_id))
+            db.session.query(Log.dashboard_id, func.count(Log.dashboard_id))
             .filter(and_(Log.dashboard_id.isnot(None), Log.dttm >= self.since))
             .group_by(Log.dashboard_id)
             .order_by(func.count(Log.dashboard_id).desc())
@@ -178,7 +176,9 @@ class TopNDashboardsStrategy(Strategy):
             .all()
         )
         dash_ids = [record.dashboard_id for record in records]
-        dashboards = session.query(Dashboard).filter(Dashboard.id.in_(dash_ids)).all()
+        dashboards = (
+            db.session.query(Dashboard).filter(Dashboard.id.in_(dash_ids)).all()
+        )
         for dashboard in dashboards:
             for chart in dashboard.slices:
                 form_data_with_filters = get_form_data(chart.id, dashboard)
@@ -211,14 +211,13 @@ class DashboardTagsStrategy(Strategy):
 
     def get_urls(self) -> List[str]:
         urls = []
-        session = db.create_scoped_session()
 
-        tags = session.query(Tag).filter(Tag.name.in_(self.tags)).all()
+        tags = db.session.query(Tag).filter(Tag.name.in_(self.tags)).all()
         tag_ids = [tag.id for tag in tags]
 
         # add dashboards that are tagged
         tagged_objects = (
-            session.query(TaggedObject)
+            db.session.query(TaggedObject)
             .filter(
                 and_(
                     TaggedObject.object_type == "dashboard",
@@ -228,14 +227,16 @@ class DashboardTagsStrategy(Strategy):
             .all()
         )
         dash_ids = [tagged_object.object_id for tagged_object in tagged_objects]
-        tagged_dashboards = session.query(Dashboard).filter(Dashboard.id.in_(dash_ids))
+        tagged_dashboards = db.session.query(Dashboard).filter(
+            Dashboard.id.in_(dash_ids)
+        )
         for dashboard in tagged_dashboards:
             for chart in dashboard.slices:
                 urls.append(get_url(chart))
 
         # add charts that are tagged
         tagged_objects = (
-            session.query(TaggedObject)
+            db.session.query(TaggedObject)
             .filter(
                 and_(
                     TaggedObject.object_type == "chart",
@@ -245,7 +246,7 @@ class DashboardTagsStrategy(Strategy):
             .all()
         )
         chart_ids = [tagged_object.object_id for tagged_object in tagged_objects]
-        tagged_charts = session.query(Slice).filter(Slice.id.in_(chart_ids))
+        tagged_charts = db.session.query(Slice).filter(Slice.id.in_(chart_ids))
         for chart in tagged_charts:
             urls.append(get_url(chart))
 

superset/tasks/schedules.py

@@ -48,7 +48,6 @@ from flask_login import login_user
 from retry.api import retry_call
 from selenium.common.exceptions import WebDriverException
 from selenium.webdriver import chrome, firefox
-from sqlalchemy.orm import Session
 from werkzeug.http import parse_cookie
 
 from superset import app, db, security_manager, thumbnail_cache
@@ -543,8 +542,7 @@ def schedule_alert_query(  # pylint: disable=unused-argument
     is_test_alert: Optional[bool] = False,
 ) -> None:
     model_cls = get_scheduler_model(report_type)
-    dbsession = db.create_scoped_session()
-    schedule = dbsession.query(model_cls).get(schedule_id)
+    schedule = db.session.query(model_cls).get(schedule_id)
 
     # The user may have disabled the schedule. If so, ignore this
     if not schedule or not schedule.active:
@@ -556,7 +554,7 @@ def schedule_alert_query(  # pylint: disable=unused-argument
             deliver_alert(schedule, recipients)
             return
 
-        if run_alert_query(schedule, dbsession):
+        if run_alert_query(schedule):
             # deliver_dashboard OR deliver_slice
             return
     else:
@@ -614,7 +612,7 @@ def deliver_alert(alert: Alert, recipients: Optional[str] = None) -> None:
     _deliver_email(recipients, deliver_as_group, subject, body, data, images)
 
 
-def run_alert_query(alert: Alert, dbsession: Session) -> Optional[bool]:
+def run_alert_query(alert: Alert) -> Optional[bool]:
     """
     Execute alert.sql and return value if any rows are returned
     """
@@ -666,7 +664,7 @@ def run_alert_query(alert: Alert, dbsession: Session) -> Optional[bool]:
             state=state,
         )
     )
-    dbsession.commit()
+    db.session.commit()
 
     return None
 
@@ -706,8 +704,7 @@ def schedule_window(
     if not model_cls:
         return None
 
-    dbsession = db.create_scoped_session()
-    schedules = dbsession.query(model_cls).filter(model_cls.active.is_(True))
+    schedules = db.session.query(model_cls).filter(model_cls.active.is_(True))
 
     for schedule in schedules:
         logging.info("Processing schedule %s", schedule)