refactor: Ensure Flask framework leverages the Flask-SQLAlchemy session (Phase II) (#26909)

tests/integration_tests/explore/api_tests.py

@@ -21,6 +21,7 @@ import pytest
 from flask_appbuilder.security.sqla.models import User
 from sqlalchemy.orm import Session
 
+from superset import db
 from superset.commands.explore.form_data.state import TemporaryExploreState
 from superset.connectors.sqla.models import SqlaTable
 from superset.explore.exceptions import DatasetAccessDeniedError
@@ -39,25 +40,22 @@ FORM_DATA = {"test": "test value"}
 @pytest.fixture
 def chart_id(load_world_bank_dashboard_with_slices) -> int:
     with app.app_context() as ctx:
-        session: Session = ctx.app.appbuilder.get_session
-        chart = session.query(Slice).filter_by(slice_name="World's Population").one()
+        chart = db.session.query(Slice).filter_by(slice_name="World's Population").one()
         return chart.id
 
 
 @pytest.fixture
 def admin_id() -> int:
     with app.app_context() as ctx:
-        session: Session = ctx.app.appbuilder.get_session
-        admin = session.query(User).filter_by(username="admin").one()
+        admin = db.session.query(User).filter_by(username="admin").one()
         return admin.id
 
 
 @pytest.fixture
 def dataset() -> int:
     with app.app_context() as ctx:
-        session: Session = ctx.app.appbuilder.get_session
         dataset = (
-            session.query(SqlaTable)
+            db.session.query(SqlaTable)
             .filter_by(table_name="wb_health_population")
             .first()
         )

tests/integration_tests/explore/form_data/api_tests.py

@@ -21,6 +21,7 @@ import pytest
 from flask_appbuilder.security.sqla.models import User
 from sqlalchemy.orm import Session
 
+from superset import db
 from superset.commands.dataset.exceptions import DatasetAccessDeniedError
 from superset.commands.explore.form_data.state import TemporaryExploreState
 from superset.connectors.sqla.models import SqlaTable
@@ -41,25 +42,22 @@ UPDATED_FORM_DATA = json.dumps({"test": "updated value"})
 @pytest.fixture
 def chart_id(load_world_bank_dashboard_with_slices) -> int:
     with app.app_context() as ctx:
-        session: Session = ctx.app.appbuilder.get_session
-        chart = session.query(Slice).filter_by(slice_name="World's Population").one()
+        chart = db.session.query(Slice).filter_by(slice_name="World's Population").one()
         return chart.id
 
 
 @pytest.fixture
 def admin_id() -> int:
     with app.app_context() as ctx:
-        session: Session = ctx.app.appbuilder.get_session
-        admin = session.query(User).filter_by(username="admin").one()
+        admin = db.session.query(User).filter_by(username="admin").one()
         return admin.id
 
 
 @pytest.fixture
 def datasource() -> int:
     with app.app_context() as ctx:
-        session: Session = ctx.app.appbuilder.get_session
         dataset = (
-            session.query(SqlaTable)
+            db.session.query(SqlaTable)
             .filter_by(table_name="wb_health_population")
             .first()
         )

tests/integration_tests/explore/form_data/commands_tests.py

@@ -45,22 +45,22 @@ class TestCreateFormDataCommand(SupersetTestCase):
                 schema=get_example_default_schema(),
                 sql="select 123 as intcol, 'abc' as strcol",
             )
-            session = db.session
-            session.add(dataset)
-            session.commit()
+            db.session.add(dataset)
+            db.session.commit()
 
             yield dataset
 
             # rollback
-            session.delete(dataset)
-            session.commit()
+            db.session.delete(dataset)
+            db.session.commit()
 
     @pytest.fixture()
     def create_slice(self):
         with self.create_app().app_context():
-            session = db.session
             dataset = (
-                session.query(SqlaTable).filter_by(table_name="dummy_sql_table").first()
+                db.session.query(SqlaTable)
+                .filter_by(table_name="dummy_sql_table")
+                .first()
             )
             slice = Slice(
                 datasource_id=dataset.id,
@@ -69,34 +69,32 @@ class TestCreateFormDataCommand(SupersetTestCase):
                 slice_name="slice_name",
             )
 
-            session.add(slice)
-            session.commit()
+            db.session.add(slice)
+            db.session.commit()
 
             yield slice
 
             # rollback
-            session.delete(slice)
-            session.commit()
+            db.session.delete(slice)
+            db.session.commit()
 
     @pytest.fixture()
     def create_query(self):
         with self.create_app().app_context():
-            session = db.session
-
             query = Query(
                 sql="select 1 as foo;",
                 client_id="sldkfjlk",
                 database=get_example_database(),
             )
 
-            session.add(query)
-            session.commit()
+            db.session.add(query)
+            db.session.commit()
 
             yield query
 
             # rollback
-            session.delete(query)
-            session.commit()
+            db.session.delete(query)
+            db.session.commit()
 
     @patch("superset.security.manager.g")
     @pytest.mark.usefixtures("create_dataset", "create_slice")

tests/integration_tests/explore/permalink/api_tests.py

@@ -38,8 +38,7 @@ from tests.integration_tests.test_app import app
 
 @pytest.fixture
 def chart(app_context, load_world_bank_dashboard_with_slices) -> Slice:
-    session: Session = app_context.app.appbuilder.get_session
-    chart = session.query(Slice).filter_by(slice_name="World's Population").one()
+    chart = db.session.query(Slice).filter_by(slice_name="World's Population").one()
     return chart
 
 

tests/integration_tests/explore/permalink/commands_tests.py

@@ -43,22 +43,22 @@ class TestCreatePermalinkDataCommand(SupersetTestCase):
                 schema=get_example_default_schema(),
                 sql="select 123 as intcol, 'abc' as strcol",
             )
-            session = db.session
-            session.add(dataset)
-            session.commit()
+            db.session.add(dataset)
+            db.session.commit()
 
             yield dataset
 
             # rollback
-            session.delete(dataset)
-            session.commit()
+            db.session.delete(dataset)
+            db.session.commit()
 
     @pytest.fixture()
     def create_slice(self):
         with self.create_app().app_context():
-            session = db.session
             dataset = (
-                session.query(SqlaTable).filter_by(table_name="dummy_sql_table").first()
+                db.session.query(SqlaTable)
+                .filter_by(table_name="dummy_sql_table")
+                .first()
             )
             slice = Slice(
                 datasource_id=dataset.id,
@@ -67,34 +67,32 @@ class TestCreatePermalinkDataCommand(SupersetTestCase):
                 slice_name="slice_name",
             )
 
-            session.add(slice)
-            session.commit()
+            db.session.add(slice)
+            db.session.commit()
 
             yield slice
 
             # rollback
-            session.delete(slice)
-            session.commit()
+            db.session.delete(slice)
+            db.session.commit()
 
     @pytest.fixture()
     def create_query(self):
         with self.create_app().app_context():
-            session = db.session
-
             query = Query(
                 sql="select 1 as foo;",
                 client_id="sldkfjlk",
                 database=get_example_database(),
             )
 
-            session.add(query)
-            session.commit()
+            db.session.add(query)
+            db.session.commit()
 
             yield query
 
             # rollback
-            session.delete(query)
-            session.commit()
+            db.session.delete(query)
+            db.session.commit()
 
     @patch("superset.security.manager.g")
     @pytest.mark.usefixtures("create_dataset", "create_slice")