chore: Cleanup database sessions (#10427)

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

tests/dict_import_export_tests.py

@@ -47,14 +47,13 @@ class TestDictImportExport(SupersetTestCase):
     def delete_imports(cls):
         with app.app_context():
             # Imported data clean up
-            session = db.session
-            for table in session.query(SqlaTable):
+            for table in db.session.query(SqlaTable):
                 if DBREF in table.params_dict:
-                    session.delete(table)
-            for datasource in session.query(DruidDatasource):
+                    db.session.delete(table)
+            for datasource in db.session.query(DruidDatasource):
                 if DBREF in datasource.params_dict:
-                    session.delete(datasource)
-            session.commit()
+                    db.session.delete(datasource)
+            db.session.commit()
 
     @classmethod
     def setUpClass(cls):
@@ -90,9 +89,7 @@ class TestDictImportExport(SupersetTestCase):
 
     def create_druid_datasource(self, name, id=0, cols_names=[], metric_names=[]):
         cluster_name = "druid_test"
-        cluster = self.get_or_create(
-            DruidCluster, {"cluster_name": cluster_name}, db.session
-        )
+        cluster = self.get_or_create(DruidCluster, {"cluster_name": cluster_name})
 
         name = "{0}{1}".format(NAME_PREFIX, name)
         params = {DBREF: id, "database_name": cluster_name}
@@ -159,7 +156,7 @@ class TestDictImportExport(SupersetTestCase):
 
     def test_import_table_no_metadata(self):
         table, dict_table = self.create_table("pure_table", id=ID_PREFIX + 1)
-        new_table = SqlaTable.import_from_dict(db.session, dict_table)
+        new_table = SqlaTable.import_from_dict(dict_table)
         db.session.commit()
         imported_id = new_table.id
         imported = self.get_table_by_id(imported_id)
@@ -173,7 +170,7 @@ class TestDictImportExport(SupersetTestCase):
             cols_names=["col1"],
             metric_names=["metric1"],
         )
-        imported_table = SqlaTable.import_from_dict(db.session, dict_table)
+        imported_table = SqlaTable.import_from_dict(dict_table)
         db.session.commit()
         imported = self.get_table_by_id(imported_table.id)
         self.assert_table_equals(table, imported)
@@ -189,7 +186,7 @@ class TestDictImportExport(SupersetTestCase):
             cols_names=["c1", "c2"],
             metric_names=["m1", "m2"],
         )
-        imported_table = SqlaTable.import_from_dict(db.session, dict_table)
+        imported_table = SqlaTable.import_from_dict(dict_table)
         db.session.commit()
         imported = self.get_table_by_id(imported_table.id)
         self.assert_table_equals(table, imported)
@@ -199,7 +196,7 @@ class TestDictImportExport(SupersetTestCase):
         table, dict_table = self.create_table(
             "table_override", id=ID_PREFIX + 3, cols_names=["col1"], metric_names=["m1"]
         )
-        imported_table = SqlaTable.import_from_dict(db.session, dict_table)
+        imported_table = SqlaTable.import_from_dict(dict_table)
         db.session.commit()
         table_over, dict_table_over = self.create_table(
             "table_override",
@@ -207,7 +204,7 @@ class TestDictImportExport(SupersetTestCase):
             cols_names=["new_col1", "col2", "col3"],
             metric_names=["new_metric1"],
         )
-        imported_over_table = SqlaTable.import_from_dict(db.session, dict_table_over)
+        imported_over_table = SqlaTable.import_from_dict(dict_table_over)
         db.session.commit()
 
         imported_over = self.get_table_by_id(imported_over_table.id)
@@ -227,7 +224,7 @@ class TestDictImportExport(SupersetTestCase):
         table, dict_table = self.create_table(
             "table_override", id=ID_PREFIX + 3, cols_names=["col1"], metric_names=["m1"]
         )
-        imported_table = SqlaTable.import_from_dict(db.session, dict_table)
+        imported_table = SqlaTable.import_from_dict(dict_table)
         db.session.commit()
         table_over, dict_table_over = self.create_table(
             "table_override",
@@ -236,7 +233,7 @@ class TestDictImportExport(SupersetTestCase):
             metric_names=["new_metric1"],
         )
         imported_over_table = SqlaTable.import_from_dict(
-            session=db.session, dict_rep=dict_table_over, sync=["metrics", "columns"]
+            dict_rep=dict_table_over, sync=["metrics", "columns"]
         )
         db.session.commit()
 
@@ -260,7 +257,7 @@ class TestDictImportExport(SupersetTestCase):
             cols_names=["new_col1", "col2", "col3"],
             metric_names=["new_metric1"],
         )
-        imported_table = SqlaTable.import_from_dict(db.session, dict_table)
+        imported_table = SqlaTable.import_from_dict(dict_table)
         db.session.commit()
         copy_table, dict_copy_table = self.create_table(
             "copy_cat",
@@ -268,7 +265,7 @@ class TestDictImportExport(SupersetTestCase):
             cols_names=["new_col1", "col2", "col3"],
             metric_names=["new_metric1"],
         )
-        imported_copy_table = SqlaTable.import_from_dict(db.session, dict_copy_table)
+        imported_copy_table = SqlaTable.import_from_dict(dict_copy_table)
         db.session.commit()
         self.assertEqual(imported_table.id, imported_copy_table.id)
         self.assert_table_equals(copy_table, self.get_table_by_id(imported_table.id))
@@ -281,10 +278,7 @@ class TestDictImportExport(SupersetTestCase):
         self.delete_fake_db()
 
         cli_export = export_to_dict(
-            session=db.session,
-            recursive=True,
-            back_references=False,
-            include_defaults=False,
+            recursive=True, back_references=False, include_defaults=False,
         )
         self.get_resp("/login/", data=dict(username="admin", password="general"))
         resp = self.get_resp(
@@ -303,7 +297,7 @@ class TestDictImportExport(SupersetTestCase):
         datasource, dict_datasource = self.create_druid_datasource(
             "pure_druid", id=ID_PREFIX + 1
         )
-        imported_cluster = DruidDatasource.import_from_dict(db.session, dict_datasource)
+        imported_cluster = DruidDatasource.import_from_dict(dict_datasource)
         db.session.commit()
         imported = self.get_datasource(imported_cluster.id)
         self.assert_datasource_equals(datasource, imported)
@@ -315,7 +309,7 @@ class TestDictImportExport(SupersetTestCase):
             cols_names=["col1"],
             metric_names=["metric1"],
         )
-        imported_cluster = DruidDatasource.import_from_dict(db.session, dict_datasource)
+        imported_cluster = DruidDatasource.import_from_dict(dict_datasource)
         db.session.commit()
         imported = self.get_datasource(imported_cluster.id)
         self.assert_datasource_equals(datasource, imported)
@@ -331,7 +325,7 @@ class TestDictImportExport(SupersetTestCase):
             cols_names=["c1", "c2"],
             metric_names=["m1", "m2"],
         )
-        imported_cluster = DruidDatasource.import_from_dict(db.session, dict_datasource)
+        imported_cluster = DruidDatasource.import_from_dict(dict_datasource)
         db.session.commit()
         imported = self.get_datasource(imported_cluster.id)
         self.assert_datasource_equals(datasource, imported)
@@ -340,7 +334,7 @@ class TestDictImportExport(SupersetTestCase):
         datasource, dict_datasource = self.create_druid_datasource(
             "druid_override", id=ID_PREFIX + 3, cols_names=["col1"], metric_names=["m1"]
         )
-        imported_cluster = DruidDatasource.import_from_dict(db.session, dict_datasource)
+        imported_cluster = DruidDatasource.import_from_dict(dict_datasource)
         db.session.commit()
         table_over, table_over_dict = self.create_druid_datasource(
             "druid_override",
@@ -348,9 +342,7 @@ class TestDictImportExport(SupersetTestCase):
             cols_names=["new_col1", "col2", "col3"],
             metric_names=["new_metric1"],
         )
-        imported_over_cluster = DruidDatasource.import_from_dict(
-            db.session, table_over_dict
-        )
+        imported_over_cluster = DruidDatasource.import_from_dict(table_over_dict)
         db.session.commit()
         imported_over = self.get_datasource(imported_over_cluster.id)
         self.assertEqual(imported_cluster.id, imported_over.id)
@@ -366,7 +358,7 @@ class TestDictImportExport(SupersetTestCase):
         datasource, dict_datasource = self.create_druid_datasource(
             "druid_override", id=ID_PREFIX + 3, cols_names=["col1"], metric_names=["m1"]
         )
-        imported_cluster = DruidDatasource.import_from_dict(db.session, dict_datasource)
+        imported_cluster = DruidDatasource.import_from_dict(dict_datasource)
         db.session.commit()
         table_over, table_over_dict = self.create_druid_datasource(
             "druid_override",
@@ -375,7 +367,7 @@ class TestDictImportExport(SupersetTestCase):
             metric_names=["new_metric1"],
         )
         imported_over_cluster = DruidDatasource.import_from_dict(
-            session=db.session, dict_rep=table_over_dict, sync=["metrics", "columns"]
+            dict_rep=table_over_dict, sync=["metrics", "columns"]
         )  # syncing metrics and columns
         db.session.commit()
         imported_over = self.get_datasource(imported_over_cluster.id)
@@ -395,9 +387,7 @@ class TestDictImportExport(SupersetTestCase):
             cols_names=["new_col1", "col2", "col3"],
             metric_names=["new_metric1"],
         )
-        imported = DruidDatasource.import_from_dict(
-            session=db.session, dict_rep=dict_datasource
-        )
+        imported = DruidDatasource.import_from_dict(dict_rep=dict_datasource)
         db.session.commit()
         copy_datasource, dict_cp_datasource = self.create_druid_datasource(
             "copy_cat",
@@ -405,7 +395,7 @@ class TestDictImportExport(SupersetTestCase):
             cols_names=["new_col1", "col2", "col3"],
             metric_names=["new_metric1"],
         )
-        imported_copy = DruidDatasource.import_from_dict(db.session, dict_cp_datasource)
+        imported_copy = DruidDatasource.import_from_dict(dict_cp_datasource)
         db.session.commit()
 
         self.assertEqual(imported.id, imported_copy.id)