apache · john-bodley · Oct 1, 2019 · Sep 5, 2019
diff --git a/superset/security.py b/superset/security.py
@@ -39,7 +39,9 @@
 from superset.exceptions import SupersetSecurityException
 
 if TYPE_CHECKING:
+    from superset.common.query_context import QueryContext
     from superset.models.core import Database, BaseDatasource
+    from superset.viz import BaseViz
 
 from superset.utils.core import DatasourceName  # noqa: E402
 
@@ -783,11 +785,31 @@ def assert_datasource_permission(self, datasource: "BaseDatasource") -> None:
         Assert the the user has permission to access the Superset datasource.
 
         :param datasource: The Superset datasource
-        :rasies SupersetSecurityException: If the user does not have permission
+        :raises SupersetSecurityException: If the user does not have permission
         """
 
         if not self.datasource_access(datasource):
             raise SupersetSecurityException(
                 self.get_datasource_access_error_msg(datasource),
                 self.get_datasource_access_link(datasource),
             )
+
+    def assert_query_context_permission(self, query_context: "QueryContext") -> None:
+        """
+        Assert the the user has permission to access the query context.
+
+        :param query_context: The query context
+        :raises SupersetSecurityException: If the user does not have permission
+        """
+
+        self.assert_datasource_permission(query_context.datasource)
+
+    def assert_viz_permission(self, viz: "BaseViz") -> None:
+        """
+        Assert the the user has permission to access the visualization.
+
+        :param viz: The visualization
+        :raises SupersetSecurityException: If the user does not have permission
+        """
+
+        self.assert_datasource_permission(viz.datasource)
diff --git a/superset/views/api.py b/superset/views/api.py
@@ -41,7 +41,7 @@ def query(self):
         params: query_context: json_blob
         """
         query_context = QueryContext(**json.loads(request.form.get("query_context")))
-        security_manager.assert_datasource_permission(query_context.datasource)
+        security_manager.assert_query_context_permission(query_context)
         payload_json = query_context.get_payload()
         return json.dumps(
             payload_json, default=utils.json_int_dttm_ser, ignore_nan=True

diff --git a/superset/views/core.py b/superset/views/core.py
@@ -168,7 +168,7 @@ def check_datasource_perms(
         force=False,
     )
 
-    security_manager.assert_datasource_permission(viz_obj.datasource)
+    security_manager.assert_viz_permission(viz_obj)
 
 
 def check_slice_perms(self, slice_id):
@@ -177,18 +177,18 @@ def check_slice_perms(self, slice_id):
 
     This function takes `self` since it must have the same signature as the
     the decorated method.
-
     """
+
     form_data, slc = get_form_data(slice_id, use_slice_data=True)
-    datasource_type = slc.datasource.type
-    datasource_id = slc.datasource.id
+
     viz_obj = get_viz(
-        datasource_type=datasource_type,
-        datasource_id=datasource_id,
+        datasource_type=slc.datasource.type,
+        datasource_id=slc.datasource.id,
         form_data=form_data,
         force=False,
     )
-    security_manager.assert_datasource_permission(viz_obj.datasource)
+
+    security_manager.assert_viz_permission(viz_obj)
 
 
 def _deserialize_results_payload(
@@ -2933,7 +2933,7 @@ def slice_query(self, slice_id):
         get the database query string for this slice
         """
         viz_obj = get_viz(slice_id)
-        security_manager.assert_datasource_permission(viz_obj.datasource)
+        security_manager.assert_viz_permission(viz_obj)
         return self.get_query_string_response(viz_obj)
 
     @api

diff --git a/tests/security_tests.py b/tests/security_tests.py
@@ -16,8 +16,10 @@
 # under the License.
 import inspect
 import unittest
+from unittest.mock import Mock, patch
 
-from superset import app, appbuilder, security_manager
+from superset import app, appbuilder, security_manager, viz
+from superset.exceptions import SupersetSecurityException
 from .base_tests import SupersetTestCase
 
 
@@ -324,3 +326,52 @@ def test_views_are_secured(self):
         if unsecured_views:
             view_str = "\n".join([str(v) for v in unsecured_views])
             raise Exception(f"Some views are not secured:\n{view_str}")
+
+
+class SecurityManagerTests(SupersetTestCase):
+    """
+    Testing the Security Manager.
+    """
+
+    @patch("superset.security.SupersetSecurityManager.datasource_access")
+    def test_assert_datasource_permission(self, mock_datasource_access):
+        datasource = self.get_datasource_mock()
+
+        # Datasource with the "datasource_access" permission.
+        mock_datasource_access.return_value = True
+        security_manager.assert_datasource_permission(datasource)
+
+        # Datasource without the "datasource_access" permission.
+        mock_datasource_access.return_value = False
+
+        with self.assertRaises(SupersetSecurityException):
+            security_manager.assert_datasource_permission(datasource)
+
+    @patch("superset.security.SupersetSecurityManager.datasource_access")
+    def test_assert_query_context_permission(self, mock_datasource_access):
+        query_context = Mock()
+        query_context.datasource = self.get_datasource_mock()
+
+        # Query context with the "datasource_access" permission.
+        mock_datasource_access.return_value = True
+        security_manager.assert_query_context_permission(query_context)
+
+        # Query context without the "datasource_access" permission.
+        mock_datasource_access.return_value = False
+
+        with self.assertRaises(SupersetSecurityException):
+            security_manager.assert_query_context_permission(query_context)
+
+    @patch("superset.security.SupersetSecurityManager.datasource_access")
+    def test_assert_viz_permission(self, mock_datasource_access):
+        test_viz = viz.TableViz(self.get_datasource_mock(), form_data={})
+
+        # Visualization with the "datasource_access" permission.
+        mock_datasource_access.return_value = True
+        security_manager.assert_viz_permission(test_viz)
+
+        # Visualization without the "datasource_access" permission.
+        mock_datasource_access.return_value = False
+
+        with self.assertRaises(SupersetSecurityException):
+            security_manager.assert_viz_permission(test_viz)