You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by mi...@apache.org on 2024/01/18 18:43:39 UTC

(superset) 06/08: fix: create virtual dataset validation (#26625)

This is an automated email from the ASF dual-hosted git repository.

michaelsmolina pushed a commit to branch 3.1
in repository https://gitbox.apache.org/repos/asf/superset.git

commit 0c569108da6a48253f1e9f82a223da76cb307009
Author: Daniel Vaz Gaspar <da...@gmail.com>
AuthorDate: Wed Jan 17 09:11:15 2024 +0000

    fix: create virtual dataset validation (#26625)
    
    (cherry picked from commit 8e19f59dd276617822d263c700e49386b92d4a6c)
---
 superset/commands/dataset/create.py                |  13 +-
 superset/commands/dataset/exceptions.py            |   7 +
 superset/security/manager.py                       |  16 +
 tests/integration_tests/datasets/commands_tests.py |  45 ++-
 tests/integration_tests/security_tests.py          | 342 ++++++++++-----------
 5 files changed, 233 insertions(+), 190 deletions(-)

diff --git a/superset/commands/dataset/create.py b/superset/commands/dataset/create.py
index 1c354e835f..16b87a567a 100644
--- a/superset/commands/dataset/create.py
+++ b/superset/commands/dataset/create.py
@@ -25,13 +25,15 @@ from superset.commands.base import BaseCommand, CreateMixin
 from superset.commands.dataset.exceptions import (
     DatabaseNotFoundValidationError,
     DatasetCreateFailedError,
+    DatasetDataAccessIsNotAllowed,
     DatasetExistsValidationError,
     DatasetInvalidError,
     TableNotFoundValidationError,
 )
 from superset.daos.dataset import DatasetDAO
 from superset.daos.exceptions import DAOCreateFailedError
-from superset.extensions import db
+from superset.exceptions import SupersetSecurityException
+from superset.extensions import db, security_manager
 
 logger = logging.getLogger(__name__)
 
@@ -82,6 +84,15 @@ class CreateDatasetCommand(CreateMixin, BaseCommand):
         ):
             exceptions.append(TableNotFoundValidationError(table_name))
 
+        if sql:
+            try:
+                security_manager.raise_for_access(
+                    database=database,
+                    sql=sql,
+                    schema=schema,
+                )
+            except SupersetSecurityException as ex:
+                exceptions.append(DatasetDataAccessIsNotAllowed(ex.error.message))
         try:
             owners = self.populate_owners(owner_ids)
             self._properties["owners"] = owners
diff --git a/superset/commands/dataset/exceptions.py b/superset/commands/dataset/exceptions.py
index f135294980..4b8acaca08 100644
--- a/superset/commands/dataset/exceptions.py
+++ b/superset/commands/dataset/exceptions.py
@@ -144,6 +144,13 @@ class OwnersNotFoundValidationError(ValidationError):
         super().__init__([_("Owners are invalid")], field_name="owners")
 
 
+class DatasetDataAccessIsNotAllowed(ValidationError):
+    status = 422
+
+    def __init__(self, message: str) -> None:
+        super().__init__([_(message)], field_name="sql")
+
+
 class DatasetNotFoundError(CommandException):
     status = 404
     message = _("Dataset does not exist")
diff --git a/superset/security/manager.py b/superset/security/manager.py
index 5eb1afdda9..b8978104d9 100644
--- a/superset/security/manager.py
+++ b/superset/security/manager.py
@@ -1828,6 +1828,8 @@ class SupersetSecurityManager(  # pylint: disable=too-many-public-methods
         query_context: Optional["QueryContext"] = None,
         table: Optional["Table"] = None,
         viz: Optional["BaseViz"] = None,
+        sql: Optional[str] = None,
+        schema: Optional[str] = None,
     ) -> None:
         """
         Raise an exception if the user cannot access the resource.
@@ -1838,6 +1840,8 @@ class SupersetSecurityManager(  # pylint: disable=too-many-public-methods
         :param query_context: The query context
         :param table: The Superset table (requires database)
         :param viz: The visualization
+        :param sql: The SQL string (requires database)
+        :param schema: Optional schema name
         :raises SupersetSecurityException: If the user cannot access the resource
         """
 
@@ -1846,7 +1850,19 @@ class SupersetSecurityManager(  # pylint: disable=too-many-public-methods
         from superset.connectors.sqla.models import SqlaTable
         from superset.models.dashboard import Dashboard
         from superset.models.slice import Slice
+        from superset.models.sql_lab import Query
         from superset.sql_parse import Table
+        from superset.utils.core import shortid
+
+        if sql and database:
+            query = Query(
+                database=database,
+                sql=sql,
+                schema=schema,
+                client_id=shortid()[:10],
+                user_id=get_user_id(),
+            )
+            self.get_session.expunge(query)
 
         if database and table or query:
             if query:
diff --git a/tests/integration_tests/datasets/commands_tests.py b/tests/integration_tests/datasets/commands_tests.py
index 1ea554a818..b45bbdb76d 100644
--- a/tests/integration_tests/datasets/commands_tests.py
+++ b/tests/integration_tests/datasets/commands_tests.py
@@ -38,7 +38,7 @@ from superset.commands.importers.exceptions import IncorrectVersionError
 from superset.connectors.sqla.models import SqlaTable
 from superset.models.core import Database
 from superset.models.slice import Slice
-from superset.utils.core import get_example_default_schema
+from superset.utils.core import get_example_default_schema, override_user
 from superset.utils.database import get_example_database
 from tests.integration_tests.base_tests import SupersetTestCase
 from tests.integration_tests.fixtures.birth_names_dashboard import (
@@ -509,6 +509,7 @@ class TestImportDatasetsCommand(SupersetTestCase):
             "metadata.yaml": yaml.safe_dump(database_metadata_config),
             "databases/imported_database.yaml": yaml.safe_dump(database_config),
         }
+
         command = ImportDatabasesCommand(contents)
         command.run()
 
@@ -558,11 +559,7 @@ class TestCreateDatasetCommand(SupersetTestCase):
                 {"table_name": "table", "database": get_example_database().id}
             ).run()
 
-    @patch("superset.security.manager.g")
-    @patch("superset.commands.utils.g")
-    def test_create_dataset_command(self, mock_g, mock_g2):
-        mock_g.user = security_manager.find_user("admin")
-        mock_g2.user = mock_g.user
+    def test_create_dataset_command(self):
         examples_db = get_example_database()
         with examples_db.get_sqla_engine_with_context() as engine:
             engine.execute("DROP TABLE IF EXISTS test_create_dataset_command")
@@ -570,22 +567,38 @@ class TestCreateDatasetCommand(SupersetTestCase):
                 "CREATE TABLE test_create_dataset_command AS SELECT 2 as col"
             )
 
-        table = CreateDatasetCommand(
-            {"table_name": "test_create_dataset_command", "database": examples_db.id}
-        ).run()
-        fetched_table = (
-            db.session.query(SqlaTable)
-            .filter_by(table_name="test_create_dataset_command")
-            .one()
-        )
-        self.assertEqual(table, fetched_table)
-        self.assertEqual([owner.username for owner in table.owners], ["admin"])
+        with override_user(security_manager.find_user("admin")):
+            table = CreateDatasetCommand(
+                {
+                    "table_name": "test_create_dataset_command",
+                    "database": examples_db.id,
+                }
+            ).run()
+            fetched_table = (
+                db.session.query(SqlaTable)
+                .filter_by(table_name="test_create_dataset_command")
+                .one()
+            )
+            self.assertEqual(table, fetched_table)
+            self.assertEqual([owner.username for owner in table.owners], ["admin"])
 
         db.session.delete(table)
         with examples_db.get_sqla_engine_with_context() as engine:
             engine.execute("DROP TABLE test_create_dataset_command")
         db.session.commit()
 
+    def test_create_dataset_command_not_allowed(self):
+        examples_db = get_example_database()
+        with override_user(security_manager.find_user("gamma")):
+            with self.assertRaises(DatasetInvalidError):
+                _ = CreateDatasetCommand(
+                    {
+                        "sql": "select * from ab_user",
+                        "database": examples_db.id,
+                        "table_name": "exp1",
+                    }
+                ).run()
+
 
 class TestDatasetWarmUpCacheCommand(SupersetTestCase):
     def test_warm_up_cache_command_table_not_found(self):
diff --git a/tests/integration_tests/security_tests.py b/tests/integration_tests/security_tests.py
index 9eaabf3680..90a24cbaff 100644
--- a/tests/integration_tests/security_tests.py
+++ b/tests/integration_tests/security_tests.py
@@ -43,6 +43,7 @@ from superset.utils.core import (
     DatasourceType,
     backend,
     get_example_default_schema,
+    override_user,
 )
 from superset.utils.database import get_example_database
 from superset.utils.urls import get_url_host
@@ -1192,37 +1193,31 @@ class TestRolePermission(SupersetTestCase):
             self.assertEqual(schemas, ["1"])
         delete_schema_perm("[examples].[1]")
 
-    @patch("superset.utils.core.g")
-    @patch("superset.security.manager.g")
-    def test_schemas_accessible_by_user_datasource_access(self, mock_sm_g, mock_g):
+    def test_schemas_accessible_by_user_datasource_access(self):
         # User has schema access to the datasource temp_schema.wb_health_population in examples DB.
-        mock_g.user = mock_sm_g.user = security_manager.find_user("gamma")
+        database = get_example_database()
         with self.client.application.test_request_context():
-            database = get_example_database()
-            schemas = security_manager.get_schemas_accessible_by_user(
-                database, ["temp_schema", "2", "3"]
-            )
-            self.assertEqual(schemas, ["temp_schema"])
+            with override_user(security_manager.find_user("gamma")):
+                schemas = security_manager.get_schemas_accessible_by_user(
+                    database, ["temp_schema", "2", "3"]
+                )
+                self.assertEqual(schemas, ["temp_schema"])
 
-    @patch("superset.utils.core.g")
-    @patch("superset.security.manager.g")
-    def test_schemas_accessible_by_user_datasource_and_schema_access(
-        self, mock_sm_g, mock_g
-    ):
+    def test_schemas_accessible_by_user_datasource_and_schema_access(self):
         # User has schema access to the datasource temp_schema.wb_health_population in examples DB.
         create_schema_perm("[examples].[2]")
-        mock_g.user = mock_sm_g.user = security_manager.find_user("gamma")
         with self.client.application.test_request_context():
             database = get_example_database()
-            schemas = security_manager.get_schemas_accessible_by_user(
-                database, ["temp_schema", "2", "3"]
-            )
-            self.assertEqual(schemas, ["temp_schema", "2"])
-        vm = security_manager.find_permission_view_menu(
-            "schema_access", "[examples].[2]"
-        )
-        self.assertIsNotNone(vm)
-        delete_schema_perm("[examples].[2]")
+            with override_user(security_manager.find_user("gamma")):
+                schemas = security_manager.get_schemas_accessible_by_user(
+                    database, ["temp_schema", "2", "3"]
+                )
+                self.assertEqual(schemas, ["temp_schema", "2"])
+                vm = security_manager.find_permission_view_menu(
+                    "schema_access", "[examples].[2]"
+                )
+                self.assertIsNotNone(vm)
+                delete_schema_perm("[examples].[2]")
 
     @pytest.mark.usefixtures("load_world_bank_dashboard_with_slices")
     def test_gamma_user_schema_access_to_dashboards(self):
@@ -1642,25 +1637,42 @@ class TestSecurityManager(SupersetTestCase):
         with self.assertRaises(SupersetSecurityException):
             security_manager.raise_for_access(query=query)
 
-    @patch("superset.security.manager.g")
+    def test_raise_for_access_sql_fails(self):
+        with override_user(security_manager.find_user("gamma")):
+            with self.assertRaises(SupersetSecurityException):
+                security_manager.raise_for_access(
+                    database=get_example_database(),
+                    schema="bar",
+                    sql="SELECT * FROM foo",
+                )
+
+    @patch("superset.security.SupersetSecurityManager.is_owner")
+    @patch("superset.security.SupersetSecurityManager.can_access")
+    def test_raise_for_access_sql(self, mock_can_access, mock_is_owner):
+        mock_can_access.return_value = True
+        mock_is_owner.return_value = True
+        with override_user(security_manager.find_user("gamma")):
+            security_manager.raise_for_access(
+                database=get_example_database(), schema="bar", sql="SELECT * FROM foo"
+            )
+
     @patch("superset.security.SupersetSecurityManager.is_owner")
     @patch("superset.security.SupersetSecurityManager.can_access")
     @patch("superset.security.SupersetSecurityManager.can_access_schema")
     def test_raise_for_access_query_context(
-        self, mock_can_access_schema, mock_can_access, mock_is_owner, mock_g
+        self, mock_can_access_schema, mock_can_access, mock_is_owner
     ):
         query_context = Mock(datasource=self.get_datasource_mock(), form_data={})
 
         mock_can_access_schema.return_value = True
         security_manager.raise_for_access(query_context=query_context)
 
-        mock_g.user = security_manager.find_user("gamma")
         mock_can_access.return_value = False
         mock_can_access_schema.return_value = False
         mock_is_owner.return_value = False
-
-        with self.assertRaises(SupersetSecurityException):
-            security_manager.raise_for_access(query_context=query_context)
+        with override_user(security_manager.find_user("gamma")):
+            with self.assertRaises(SupersetSecurityException):
+                security_manager.raise_for_access(query_context=query_context)
 
     @patch("superset.security.SupersetSecurityManager.can_access")
     def test_raise_for_access_table(self, mock_can_access):
@@ -1675,30 +1687,27 @@ class TestSecurityManager(SupersetTestCase):
         with self.assertRaises(SupersetSecurityException):
             security_manager.raise_for_access(database=database, table=table)
 
-    @patch("superset.security.manager.g")
     @patch("superset.security.SupersetSecurityManager.is_owner")
     @patch("superset.security.SupersetSecurityManager.can_access")
     @patch("superset.security.SupersetSecurityManager.can_access_schema")
     def test_raise_for_access_viz(
-        self, mock_can_access_schema, mock_can_access, mock_is_owner, mock_g
+        self, mock_can_access_schema, mock_can_access, mock_is_owner
     ):
         test_viz = viz.TimeTableViz(self.get_datasource_mock(), form_data={})
 
         mock_can_access_schema.return_value = True
         security_manager.raise_for_access(viz=test_viz)
 
-        mock_g.user = security_manager.find_user("gamma")
         mock_can_access.return_value = False
         mock_can_access_schema.return_value = False
         mock_is_owner.return_value = False
-
-        with self.assertRaises(SupersetSecurityException):
-            security_manager.raise_for_access(viz=test_viz)
+        with override_user(security_manager.find_user("gamma")):
+            with self.assertRaises(SupersetSecurityException):
+                security_manager.raise_for_access(viz=test_viz)
 
     @pytest.mark.usefixtures("load_birth_names_dashboard_with_slices")
     @pytest.mark.usefixtures("load_world_bank_dashboard_with_slices")
     @with_feature_flags(DASHBOARD_RBAC=True)
-    @patch("superset.security.manager.g")
     @patch("superset.security.SupersetSecurityManager.is_owner")
     @patch("superset.security.SupersetSecurityManager.can_access")
     @patch("superset.security.SupersetSecurityManager.can_access_schema")
@@ -1707,7 +1716,6 @@ class TestSecurityManager(SupersetTestCase):
         mock_can_access_schema,
         mock_can_access,
         mock_is_owner,
-        mock_g,
     ):
         births = self.get_dash_by_slug("births")
         girls = self.get_slice("Girls", db.session, expunge_from_session=False)
@@ -1731,161 +1739,156 @@ class TestSecurityManager(SupersetTestCase):
             }
         )
 
-        mock_g.user = security_manager.find_user("gamma")
         mock_is_owner.return_value = False
         mock_can_access.return_value = False
         mock_can_access_schema.return_value = False
+        with override_user(security_manager.find_user("gamma")):
+            for kwarg in ["query_context", "viz"]:
+                births.roles = []
+
+                # No dashboard roles.
+                with self.assertRaises(SupersetSecurityException):
+                    security_manager.raise_for_access(
+                        **{
+                            kwarg: Mock(
+                                datasource=birth_names,
+                                form_data={
+                                    "dashboardId": births.id,
+                                    "slice_id": girls.id,
+                                },
+                            )
+                        }
+                    )
 
-        for kwarg in ["query_context", "viz"]:
-            births.roles = []
-
-            # No dashboard roles.
-            with self.assertRaises(SupersetSecurityException):
-                security_manager.raise_for_access(
-                    **{
-                        kwarg: Mock(
-                            datasource=birth_names,
-                            form_data={
-                                "dashboardId": births.id,
-                                "slice_id": girls.id,
-                            },
-                        )
-                    }
-                )
+                births.roles = [self.get_role("Gamma")]
+
+                # Undefined dashboard.
+                with self.assertRaises(SupersetSecurityException):
+                    security_manager.raise_for_access(
+                        **{
+                            kwarg: Mock(
+                                datasource=birth_names,
+                                form_data={},
+                            )
+                        }
+                    )
 
-            births.roles = [self.get_role("Gamma")]
+                # Undefined dashboard chart.
+                with self.assertRaises(SupersetSecurityException):
+                    security_manager.raise_for_access(
+                        **{
+                            kwarg: Mock(
+                                datasource=birth_names,
+                                form_data={"dashboardId": births.id},
+                            )
+                        }
+                    )
 
-            # Undefined dashboard.
-            with self.assertRaises(SupersetSecurityException):
-                security_manager.raise_for_access(
-                    **{
-                        kwarg: Mock(
-                            datasource=birth_names,
-                            form_data={},
-                        )
-                    }
-                )
+                # Ill-defined dashboard chart.
+                with self.assertRaises(SupersetSecurityException):
+                    security_manager.raise_for_access(
+                        **{
+                            kwarg: Mock(
+                                datasource=birth_names,
+                                form_data={
+                                    "dashboardId": births.id,
+                                    "slice_id": treemap.id,
+                                },
+                            )
+                        }
+                    )
 
-            # Undefined dashboard chart.
-            with self.assertRaises(SupersetSecurityException):
-                security_manager.raise_for_access(
-                    **{
-                        kwarg: Mock(
-                            datasource=birth_names,
-                            form_data={"dashboardId": births.id},
-                        )
-                    }
-                )
+                # Dashboard chart not associated with said datasource.
+                with self.assertRaises(SupersetSecurityException):
+                    security_manager.raise_for_access(
+                        **{
+                            kwarg: Mock(
+                                datasource=birth_names,
+                                form_data={
+                                    "dashboardId": world_health.id,
+                                    "slice_id": treemap.id,
+                                },
+                            )
+                        }
+                    )
 
-            # Ill-defined dashboard chart.
-            with self.assertRaises(SupersetSecurityException):
+                # Dashboard chart associated with said datasource.
                 security_manager.raise_for_access(
                     **{
                         kwarg: Mock(
                             datasource=birth_names,
                             form_data={
                                 "dashboardId": births.id,
-                                "slice_id": treemap.id,
-                            },
-                        )
-                    }
-                )
-
-            # Dashboard chart not associated with said datasource.
-            with self.assertRaises(SupersetSecurityException):
-                security_manager.raise_for_access(
-                    **{
-                        kwarg: Mock(
-                            datasource=birth_names,
-                            form_data={
-                                "dashboardId": world_health.id,
-                                "slice_id": treemap.id,
+                                "slice_id": girls.id,
                             },
                         )
                     }
                 )
 
-            # Dashboard chart associated with said datasource.
-            security_manager.raise_for_access(
-                **{
-                    kwarg: Mock(
-                        datasource=birth_names,
-                        form_data={
-                            "dashboardId": births.id,
-                            "slice_id": girls.id,
-                        },
+                # Ill-defined native filter.
+                with self.assertRaises(SupersetSecurityException):
+                    security_manager.raise_for_access(
+                        **{
+                            kwarg: Mock(
+                                datasource=birth_names,
+                                form_data={
+                                    "dashboardId": births.id,
+                                    "type": "NATIVE_FILTER",
+                                },
+                            )
+                        }
                     )
-                }
-            )
 
-            # Ill-defined native filter.
-            with self.assertRaises(SupersetSecurityException):
-                security_manager.raise_for_access(
-                    **{
-                        kwarg: Mock(
-                            datasource=birth_names,
-                            form_data={
-                                "dashboardId": births.id,
-                                "type": "NATIVE_FILTER",
-                            },
-                        )
-                    }
-                )
+                # Native filter not associated with said datasource.
+                with self.assertRaises(SupersetSecurityException):
+                    security_manager.raise_for_access(
+                        **{
+                            kwarg: Mock(
+                                datasource=birth_names,
+                                form_data={
+                                    "dashboardId": births.id,
+                                    "native_filter_id": "NATIVE_FILTER-IJKLMNOP",
+                                    "type": "NATIVE_FILTER",
+                                },
+                            )
+                        }
+                    )
 
-            # Native filter not associated with said datasource.
-            with self.assertRaises(SupersetSecurityException):
+                # Native filter associated with said datasource.
                 security_manager.raise_for_access(
                     **{
                         kwarg: Mock(
                             datasource=birth_names,
                             form_data={
                                 "dashboardId": births.id,
-                                "native_filter_id": "NATIVE_FILTER-IJKLMNOP",
+                                "native_filter_id": "NATIVE_FILTER-ABCDEFGH",
                                 "type": "NATIVE_FILTER",
                             },
                         )
                     }
                 )
 
-            # Native filter associated with said datasource.
-            security_manager.raise_for_access(
-                **{
-                    kwarg: Mock(
-                        datasource=birth_names,
-                        form_data={
-                            "dashboardId": births.id,
-                            "native_filter_id": "NATIVE_FILTER-ABCDEFGH",
-                            "type": "NATIVE_FILTER",
-                        },
-                    )
-                }
-            )
-
-        db.session.expunge_all()
+            db.session.expunge_all()
 
-    @patch("superset.security.manager.g")
-    def test_get_user_roles(self, mock_g):
+    def test_get_user_roles(self):
         admin = security_manager.find_user("admin")
-        mock_g.user = admin
-        roles = security_manager.get_user_roles()
-        self.assertEqual(admin.roles, roles)
+        with override_user(admin):
+            roles = security_manager.get_user_roles()
+            self.assertEqual(admin.roles, roles)
 
-    @patch("superset.security.manager.g")
-    def test_get_anonymous_roles(self, mock_g):
-        mock_g.user = security_manager.get_anonymous_user()
-        roles = security_manager.get_user_roles()
-        self.assertEqual([security_manager.get_public_role()], roles)
+    def test_get_anonymous_roles(self):
+        with override_user(security_manager.get_anonymous_user()):
+            roles = security_manager.get_user_roles()
+            self.assertEqual([security_manager.get_public_role()], roles)
 
 
 class TestDatasources(SupersetTestCase):
-    @patch("superset.security.manager.g")
     @patch("superset.security.SupersetSecurityManager.can_access_database")
     @patch("superset.security.SupersetSecurityManager.get_session")
     def test_get_user_datasources_admin(
-        self, mock_get_session, mock_can_access_database, mock_g
+        self, mock_get_session, mock_can_access_database
     ):
         Datasource = namedtuple("Datasource", ["database", "schema", "name"])
-        mock_g.user = security_manager.find_user("admin")
         mock_can_access_database.return_value = True
         mock_get_session.query.return_value.filter.return_value.all.return_value = []
 
@@ -1897,23 +1900,20 @@ class TestDatasources(SupersetTestCase):
                 Datasource("database1", "schema1", "table2"),
                 Datasource("database2", None, "table1"),
             ]
+            with override_user(security_manager.find_user("admin")):
+                datasources = security_manager.get_user_datasources()
+                assert sorted(datasources) == [
+                    Datasource("database1", "schema1", "table1"),
+                    Datasource("database1", "schema1", "table2"),
+                    Datasource("database2", None, "table1"),
+                ]
 
-            datasources = security_manager.get_user_datasources()
-
-        assert sorted(datasources) == [
-            Datasource("database1", "schema1", "table1"),
-            Datasource("database1", "schema1", "table2"),
-            Datasource("database2", None, "table1"),
-        ]
-
-    @patch("superset.security.manager.g")
     @patch("superset.security.SupersetSecurityManager.can_access_database")
     @patch("superset.security.SupersetSecurityManager.get_session")
     def test_get_user_datasources_gamma(
-        self, mock_get_session, mock_can_access_database, mock_g
+        self, mock_get_session, mock_can_access_database
     ):
         Datasource = namedtuple("Datasource", ["database", "schema", "name"])
-        mock_g.user = security_manager.find_user("gamma")
         mock_can_access_database.return_value = False
         mock_get_session.query.return_value.filter.return_value.all.return_value = []
 
@@ -1925,19 +1925,16 @@ class TestDatasources(SupersetTestCase):
                 Datasource("database1", "schema1", "table2"),
                 Datasource("database2", None, "table1"),
             ]
+            with override_user(security_manager.find_user("gamma")):
+                datasources = security_manager.get_user_datasources()
+                assert datasources == []
 
-            datasources = security_manager.get_user_datasources()
-
-        assert datasources == []
-
-    @patch("superset.security.manager.g")
     @patch("superset.security.SupersetSecurityManager.can_access_database")
     @patch("superset.security.SupersetSecurityManager.get_session")
     def test_get_user_datasources_gamma_with_schema(
-        self, mock_get_session, mock_can_access_database, mock_g
+        self, mock_get_session, mock_can_access_database
     ):
         Datasource = namedtuple("Datasource", ["database", "schema", "name"])
-        mock_g.user = security_manager.find_user("gamma")
         mock_can_access_database.return_value = False
 
         mock_get_session.query.return_value.filter.return_value.all.return_value = [
@@ -1953,13 +1950,12 @@ class TestDatasources(SupersetTestCase):
                 Datasource("database1", "schema1", "table2"),
                 Datasource("database2", None, "table1"),
             ]
-
-            datasources = security_manager.get_user_datasources()
-
-        assert sorted(datasources) == [
-            Datasource("database1", "schema1", "table1"),
-            Datasource("database1", "schema1", "table2"),
-        ]
+            with override_user(security_manager.find_user("gamma")):
+                datasources = security_manager.get_user_datasources()
+                assert sorted(datasources) == [
+                    Datasource("database1", "schema1", "table1"),
+                    Datasource("database1", "schema1", "table2"),
+                ]
 
 
 class FakeRequest: