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

johnbodley pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-superset.git


The following commit(s) were added to refs/heads/master by this push:
     new aefef9c  chore(security): Updating assert logic (#10034)
aefef9c is described below

commit aefef9ca55aeef24950547211c730b84736cfd2c
Author: John Bodley <4567245+john-bod...@users.noreply.github.com>
AuthorDate: Tue Jun 23 20:49:39 2020 -0700

    chore(security): Updating assert logic (#10034)
    
    * chore(security): Updating assert logic
    
    * Deprecating rejected_tables
    
    Co-authored-by: John Bodley <john.bod...@airbnb.com>
---
 UPDATING.md                        |   2 +
 superset/charts/api.py             |   4 +-
 superset/common/query_context.py   |   9 ++
 superset/connectors/base/models.py |  10 +++
 superset/models/sql_lab.py         |   9 ++
 superset/security/manager.py       | 180 +++++++++++++++++++++----------------
 superset/views/api.py              |   7 +-
 superset/views/core.py             |  49 +++++-----
 superset/views/utils.py            |  16 ++--
 superset/viz.py                    |   9 ++
 superset/viz_sip38.py              |   9 ++
 tests/security_tests.py            | 117 ++++++++++++++++++------
 12 files changed, 278 insertions(+), 143 deletions(-)

diff --git a/UPDATING.md b/UPDATING.md
index 739c3b3..f5a3644 100644
--- a/UPDATING.md
+++ b/UPDATING.md
@@ -23,6 +23,8 @@ assists people when migrating to a new version.
 
 ## Next
 
+* [10034](https://github.com/apache/incubator-superset/pull/10034): Deprecates 
the public security manager  `assert_datasource_permission`, 
`assert_query_context_permission`, `assert_viz_permission`, and 
`rejected_tables` methods with the `raise_for_access` method which also handles 
assertion logic for SQL tables.
+
 * [10031](https://github.com/apache/incubator-superset/pull/10030): Renames 
the following public security manager methods: `can_access_datasource` to 
`can_access_table`, `all_datasource_access` to `can_access_all_datasources`, 
`all_database_access` to `can_access_all_databases`, `database_access` to 
`can_access_database`, `schema_access` to `can_access_schema`, and
 `datasource_access` to `can_access_datasource`. Regrettably it is not viable 
to provide aliases for the deprecated methods as this would result in a name 
clash. Finally the `can_access_table` (previously `can_access_database`) method 
signature has changed, i.e., the optional `schema` argument no longer exists.
 
diff --git a/superset/charts/api.py b/superset/charts/api.py
index d309ef1..c6ff00b 100644
--- a/superset/charts/api.py
+++ b/superset/charts/api.py
@@ -53,7 +53,7 @@ from superset.charts.schemas import (
 )
 from superset.constants import RouteMethod
 from superset.exceptions import SupersetSecurityException
-from superset.extensions import event_logger, security_manager
+from superset.extensions import event_logger
 from superset.models.slice import Slice
 from superset.tasks.thumbnails import cache_chart_thumbnail
 from superset.utils.core import ChartDataResultFormat, json_int_dttm_ser
@@ -454,7 +454,7 @@ class ChartRestApi(BaseSupersetModelRestApi):
         except KeyError:
             return self.response_400(message="Request is incorrect")
         try:
-            security_manager.assert_query_context_permission(query_context)
+            query_context.raise_for_access()
         except SupersetSecurityException:
             return self.response_401()
         payload = query_context.get_payload()
diff --git a/superset/common/query_context.py b/superset/common/query_context.py
index 7cb82a4..fb70f0a 100644
--- a/superset/common/query_context.py
+++ b/superset/common/query_context.py
@@ -286,3 +286,12 @@ class QueryContext:
             "stacktrace": stacktrace,
             "rowcount": len(df.index),
         }
+
+    def raise_for_access(self) -> None:
+        """
+        Raise an exception if the user cannot access the resource.
+
+        :raises SupersetSecurityException: If the user cannot access the 
resource
+        """
+
+        security_manager.raise_for_access(query_context=self)
diff --git a/superset/connectors/base/models.py 
b/superset/connectors/base/models.py
index f2dad88..1f6f032 100644
--- a/superset/connectors/base/models.py
+++ b/superset/connectors/base/models.py
@@ -23,6 +23,7 @@ from sqlalchemy import and_, Boolean, Column, Integer, 
String, Text
 from sqlalchemy.ext.declarative import declared_attr
 from sqlalchemy.orm import foreign, Query, relationship, RelationshipProperty
 
+from superset import security_manager
 from superset.constants import NULL_STRING
 from superset.models.helpers import AuditMixinNullable, ImportMixin, 
QueryResult
 from superset.models.slice import Slice
@@ -496,6 +497,15 @@ class BaseDatasource(
             return NotImplemented
         return self.uid == other.uid
 
+    def raise_for_access(self) -> None:
+        """
+        Raise an exception if the user cannot access the resource.
+
+        :raises SupersetSecurityException: If the user cannot access the 
resource
+        """
+
+        security_manager.raise_for_access(datasource=self)
+
 
 class BaseColumn(AuditMixinNullable, ImportMixin):
     """Interface for column"""
diff --git a/superset/models/sql_lab.py b/superset/models/sql_lab.py
index 9c3c239..8a90e9a 100644
--- a/superset/models/sql_lab.py
+++ b/superset/models/sql_lab.py
@@ -148,6 +148,15 @@ class Query(Model, ExtraJSONMixin):
     def username(self) -> str:
         return self.user.username
 
+    def raise_for_access(self) -> None:
+        """
+        Raise an exception if the user cannot access the resource.
+
+        :raises SupersetSecurityException: If the user cannot access the 
resource
+        """
+
+        security_manager.raise_for_access(query=self)
+
 
 class SavedQuery(Model, AuditMixinNullable, ExtraJSONMixin):
     """ORM model for SQL query"""
diff --git a/superset/security/manager.py b/superset/security/manager.py
index 88d7493..9a001cf 100644
--- a/superset/security/manager.py
+++ b/superset/security/manager.py
@@ -17,7 +17,7 @@
 # pylint: disable=C,R,W
 """A set of constants and methods to manage permissions and security"""
 import logging
-from typing import Any, Callable, List, Optional, Set, Tuple, TYPE_CHECKING, 
Union
+from typing import Any, Callable, cast, List, Optional, Set, Tuple, 
TYPE_CHECKING, Union
 
 from flask import current_app, g
 from flask_appbuilder import Model
@@ -38,7 +38,7 @@ from flask_appbuilder.widgets import ListWidget
 from sqlalchemy import or_
 from sqlalchemy.engine.base import Connection
 from sqlalchemy.orm.mapper import Mapper
-from sqlalchemy.orm.query import Query
+from sqlalchemy.orm.query import Query as SqlaQuery
 
 from superset import sql_parse
 from superset.connectors.connector_registry import ConnectorRegistry
@@ -52,6 +52,7 @@ if TYPE_CHECKING:
     from superset.connectors.base.models import BaseDatasource
     from superset.connectors.druid.models import DruidCluster
     from superset.models.core import Database
+    from superset.models.sql_lab import Query
     from superset.sql_parse import Table
     from superset.viz import BaseViz
 
@@ -215,30 +216,32 @@ class SupersetSecurityManager(SecurityManager):
 
     def can_access_all_datasources(self) -> bool:
         """
-        Return True if the user can access all Superset datasources, False 
otherwise.
+        Return True if the user can fully access all the Superset datasources, 
False
+        otherwise.
 
-        :returns: Whether the user can access all Superset datasources
+        :returns: Whether the user can fully access all Superset datasources
         """
 
         return self.can_access("all_datasource_access", 
"all_datasource_access")
 
     def can_access_all_databases(self) -> bool:
         """
-        Return True if the user can access all Superset databases, False 
otherwise.
+        Return True if the user can fully access all the Superset databases, 
False
+        otherwise.
 
-        :returns: Whether the user can access all Superset databases
+        :returns: Whether the user can fully access all Superset databases
         """
 
         return self.can_access("all_database_access", "all_database_access")
 
     def can_access_database(self, database: Union["Database", "DruidCluster"]) 
-> bool:
         """
-        Return True if the user can access the Superset database, False 
otherwise.
+        Return True if the user can fully access the Superset database, False 
otherwise.
 
         Note for Druid the database is akin to the Druid cluster.
 
         :param database: The Superset database
-        :returns: Whether the user can access the Superset database
+        :returns: Whether the user can fully access the Superset database
         """
 
         return (
@@ -249,14 +252,14 @@ class SupersetSecurityManager(SecurityManager):
 
     def can_access_schema(self, datasource: "BaseDatasource") -> bool:
         """
-        Return True if the user can access the schema associated with the 
Superset
+        Return True if the user can fully access the schema associated with 
the Superset
         datasource, False otherwise.
 
         Note for Druid datasources the database and schema are akin to the 
Druid cluster
         and datasource name prefix respectively, i.e., [schema.]datasource.
 
         :param datasource: The Superset datasource
-        :returns: Whether the user can access the datasource's schema
+        :returns: Whether the user can fully access the datasource's schema
         """
 
         return (
@@ -267,15 +270,19 @@ class SupersetSecurityManager(SecurityManager):
 
     def can_access_datasource(self, datasource: "BaseDatasource") -> bool:
         """
-        Return True if the user can access the Superset datasource, False 
otherwise.
+        Return True if the user can fully access of the Superset datasource, 
False
+        otherwise.
 
         :param datasource: The Superset datasource
-        :returns: Whether the user can access the Superset datasource
+        :returns: Whether the user can fully access the Superset datasource
         """
 
-        return self.can_access_schema(datasource) or self.can_access(
-            "datasource_access", datasource.perm or ""
-        )
+        try:
+            self.raise_for_access(datasource=datasource)
+        except SupersetSecurityException:
+            return False
+
+        return True
 
     def get_datasource_access_error_msg(self, datasource: "BaseDatasource") -> 
str:
         """
@@ -369,45 +376,12 @@ class SupersetSecurityManager(SecurityManager):
         :returns: Whether the user can access the SQL table
         """
 
-        from superset import db
-        from superset.connectors.sqla.models import SqlaTable
-
-        if self.can_access_database(database):
-            return True
-
-        schema_perm = self.get_schema_perm(database, schema=table.schema)
-        if schema_perm and self.can_access("schema_access", schema_perm):
-            return True
-
-        datasources = SqlaTable.query_datasources_by_name(
-            db.session, database, table.table, schema=table.schema
-        )
-        for datasource in datasources:
-            if self.can_access("datasource_access", datasource.perm):
-                return True
-        return False
-
-    def rejected_tables(
-        self, sql: str, database: "Database", schema: str
-    ) -> Set["Table"]:
-        """
-        Return the list of rejected SQL tables.
-
-        :param sql: The SQL statement
-        :param database: The SQL database
-        :param schema: The SQL database schema
-        :returns: The rejected tables
-        """
-
-        from superset.sql_parse import Table
+        try:
+            self.raise_for_access(database=database, table=table)
+        except SupersetSecurityException:
+            return False
 
-        return {
-            table
-            for table in sql_parse.ParsedQuery(sql).tables
-            if not self.can_access_table(
-                database, Table(table.table, table.schema or schema)
-            )
-        }
+        return True
 
     def get_public_role(self) -> Optional[Any]:  # Optional[self.role_model]
         from superset import conf
@@ -860,45 +834,92 @@ class SupersetSecurityManager(SecurityManager):
                     )
                 )
 
-    def assert_datasource_permission(self, datasource: "BaseDatasource") -> 
None:
+    def raise_for_access(
+        self,
+        database: Optional["Database"] = None,
+        datasource: Optional["BaseDatasource"] = None,
+        query: Optional["Query"] = None,
+        query_context: Optional["QueryContext"] = None,
+        table: Optional["Table"] = None,
+        viz: Optional["BaseViz"] = None,
+    ) -> None:
         """
-        Assert the the user has permission to access the Superset datasource.
+        Raise an exception if the user cannot access the resource.
 
+        :param database: The Superset database
         :param datasource: The Superset datasource
-        :raises SupersetSecurityException: If the user does not have permission
+        :param query: The SQL Lab query
+        :param query_context: The query context
+        :param table: The Superset table (requires database)
+        :param viz: The visualization
+        :raises SupersetSecurityException: If the user cannot access the 
resource
         """
 
-        if not self.can_access_datasource(datasource):
-            raise SupersetSecurityException(
-                self.get_datasource_access_error_object(datasource),
-            )
+        from superset.connectors.sqla.models import SqlaTable
+        from superset.sql_parse import Table
 
-    def assert_query_context_permission(self, query_context: "QueryContext") 
-> None:
-        """
-        Assert the the user has permission to access the query context.
+        if database and table or query:
+            if query:
+                database = query.database
 
-        :param query_context: The query context
-        :raises SupersetSecurityException: If the user does not have permission
-        """
+            database = cast("Database", database)
 
-        self.assert_datasource_permission(query_context.datasource)
+            if self.can_access_database(database):
+                return
 
-    def assert_viz_permission(self, viz: "BaseViz") -> None:
-        """
-        Assert the the user has permission to access the visualization.
+            if query:
+                tables = {
+                    Table(table_.table, table_.schema or query.schema)
+                    for table_ in sql_parse.ParsedQuery(query.sql).tables
+                }
+            elif table:
+                tables = {table}
 
-        :param viz: The visualization
-        :raises SupersetSecurityException: If the user does not have permission
-        """
+            denied = set()
+
+            for table_ in tables:
+                schema_perm = self.get_schema_perm(database, 
schema=table_.schema)
+
+                if not (schema_perm and self.can_access("schema_access", 
schema_perm)):
+                    datasources = SqlaTable.query_datasources_by_name(
+                        self.get_session, database, table_.table, 
schema=table_.schema
+                    )
+
+                    # Access to any datasource is suffice.
+                    for datasource in datasources:
+                        if self.can_access("datasource_access", 
datasource.perm):
+                            break
+                    else:
+                        denied.add(table_)
 
-        self.assert_datasource_permission(viz.datasource)
+            if denied:
+                raise SupersetSecurityException(
+                    self.get_table_access_error_object(denied),
+                )
+
+        if datasource or query_context or viz:
+            if query_context:
+                datasource = query_context.datasource
+            elif viz:
+                datasource = viz.datasource
+
+            assert datasource
+
+            if not (
+                self.can_access_schema(datasource)
+                or self.can_access("datasource_access", datasource.perm or "")
+            ):
+                raise SupersetSecurityException(
+                    self.get_datasource_access_error_object(datasource),
+                )
 
-    def get_rls_filters(self, table: "BaseDatasource") -> List[Query]:
+    def get_rls_filters(self, table: "BaseDatasource") -> List[SqlaQuery]:
         """
-        Retrieves the appropriate row level security filters for the current 
user and the passed table.
+        Retrieves the appropriate row level security filters for the current 
user and
+        the passed table.
 
         :param table: The table to check against
-        :returns: A list of filters.
+        :returns: A list of filters
         """
         if hasattr(g, "user") and hasattr(g.user, "id"):
             from superset import db
@@ -935,10 +956,11 @@ class SupersetSecurityManager(SecurityManager):
 
     def get_rls_ids(self, table: "BaseDatasource") -> List[int]:
         """
-        Retrieves the appropriate row level security filters IDs for the 
current user and the passed table.
+        Retrieves the appropriate row level security filters IDs for the 
current user
+        and the passed table.
 
         :param table: The table to check against
-        :returns: A list of IDs.
+        :returns: A list of IDs
         """
         ids = [f.id for f in self.get_rls_filters(table)]
         ids.sort()  # Combinations rather than permutations
diff --git a/superset/views/api.py b/superset/views/api.py
index d370598..a5090b3 100644
--- a/superset/views/api.py
+++ b/superset/views/api.py
@@ -20,7 +20,7 @@ from flask import request
 from flask_appbuilder import expose
 from flask_appbuilder.security.decorators import has_access_api
 
-from superset import db, event_logger, security_manager
+from superset import db, event_logger
 from superset.common.query_context import QueryContext
 from superset.legacy import update_time_range
 from superset.models.slice import Slice
@@ -39,10 +39,11 @@ class Api(BaseSupersetView):
         """
         Takes a query_obj constructed in the client and returns payload data 
response
         for the given query_obj.
-        params: query_context: json_blob
+
+        raises SupersetSecurityException: If the user cannot access the 
resource
         """
         query_context = 
QueryContext(**json.loads(request.form["query_context"]))
-        security_manager.assert_query_context_permission(query_context)
+        query_context.raise_for_access()
         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
index e08d92a..4146914 100755
--- a/superset/views/core.py
+++ b/superset/views/core.py
@@ -69,6 +69,7 @@ from superset.exceptions import (
     CertificateException,
     DatabaseNotFound,
     SupersetException,
+    SupersetSecurityException,
     SupersetTimeoutException,
 )
 from superset.jinja_context import get_template_processor
@@ -728,7 +729,8 @@ class Superset(BaseSupersetView):
         :param datasource_type: Type of datasource e.g. table
         :param datasource_id: Datasource id
         :param column: Column name to retrieve values for
-        :return:
+        :returns: The Flask response
+        :raises SupersetSecurityException: If the user cannot access the 
resource
         """
         # TODO: Cache endpoint by user, datasource and column
         datasource = ConnectorRegistry.get_datasource(
@@ -736,7 +738,8 @@ class Superset(BaseSupersetView):
         )
         if not datasource:
             return json_error_response(DATASOURCE_MISSING_ERR)
-        security_manager.assert_datasource_permission(datasource)
+
+        datasource.raise_for_access()
         payload = json.dumps(
             datasource.values_for_column(column, 
config["FILTER_SELECT_ROW_LIMIT"]),
             default=utils.json_int_dttm_ser,
@@ -1986,14 +1989,10 @@ class Superset(BaseSupersetView):
                 status=404,
             )
 
-        rejected_tables = security_manager.rejected_tables(
-            query.sql, query.database, query.schema
-        )
-        if rejected_tables:
-            return json_errors_response(
-                
[security_manager.get_table_access_error_object(rejected_tables)],
-                status=403,
-            )
+        try:
+            query.raise_for_access()
+        except SupersetSecurityException as ex:
+            return json_errors_response([ex.error], status=403)
 
         payload = utils.zlib_decompress(blob, decode=not 
results_backend_use_msgpack)
         obj = _deserialize_results_payload(
@@ -2291,14 +2290,12 @@ class Superset(BaseSupersetView):
 
         logger.info(f"Triggering query_id: {query_id}")
 
-        rejected_tables = security_manager.rejected_tables(sql, mydb, schema)
-        if rejected_tables:
+        try:
+            query.raise_for_access()
+        except SupersetSecurityException as ex:
             query.status = QueryStatus.FAILED
             session.commit()
-            return json_errors_response(
-                
[security_manager.get_table_access_error_object(rejected_tables)],
-                status=403,
-            )
+            return json_errors_response([ex.error], status=403)
 
         try:
             template_processor = get_template_processor(
@@ -2345,12 +2342,12 @@ class Superset(BaseSupersetView):
         logger.info("Exporting CSV file [{}]".format(client_id))
         query = db.session.query(Query).filter_by(client_id=client_id).one()
 
-        rejected_tables = security_manager.rejected_tables(
-            query.sql, query.database, query.schema
-        )
-        if rejected_tables:
-            flash(security_manager.get_table_access_error_msg(rejected_tables))
+        try:
+            query.raise_for_access()
+        except SupersetSecurityException as ex:
+            flash(ex.error.message)
             return redirect("/")
+
         blob = None
         if results_backend and query.results_key:
             logger.info(
@@ -2399,6 +2396,13 @@ class Superset(BaseSupersetView):
     @expose("/fetch_datasource_metadata")
     @event_logger.log_this
     def fetch_datasource_metadata(self) -> FlaskResponse:
+        """
+        Fetch the datasource metadata.
+
+        :returns: The Flask response
+        :raises SupersetSecurityException: If the user cannot access the 
resource
+        """
+
         datasource_id, datasource_type = 
request.args["datasourceKey"].split("__")
         datasource = ConnectorRegistry.get_datasource(
             datasource_type, datasource_id, db.session
@@ -2407,8 +2411,7 @@ class Superset(BaseSupersetView):
         if not datasource:
             return json_error_response(DATASOURCE_MISSING_ERR)
 
-        # Check permission for datasource
-        security_manager.assert_datasource_permission(datasource)
+        datasource.raise_for_access()
         return json_success(json.dumps(datasource.data))
 
     @has_access_api
diff --git a/superset/views/utils.py b/superset/views/utils.py
index 6248865..2a8b2cc 100644
--- a/superset/views/utils.py
+++ b/superset/views/utils.py
@@ -28,14 +28,7 @@ from flask_appbuilder.security.sqla import models as 
ab_models
 from flask_appbuilder.security.sqla.models import User
 
 import superset.models.core as models
-from superset import (
-    app,
-    dataframe,
-    db,
-    is_feature_enabled,
-    result_set,
-    security_manager,
-)
+from superset import app, dataframe, db, is_feature_enabled, result_set
 from superset.connectors.connector_registry import ConnectorRegistry
 from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
 from superset.exceptions import SupersetException, SupersetSecurityException
@@ -433,7 +426,7 @@ def check_datasource_perms(
         force=False,
     )
 
-    security_manager.assert_viz_permission(viz_obj)
+    viz_obj.raise_for_access()
 
 
 def check_slice_perms(_self: Any, slice_id: int) -> None:
@@ -442,6 +435,9 @@ def check_slice_perms(_self: Any, slice_id: int) -> None:
 
     This function takes `self` since it must have the same signature as the
     the decorated method.
+
+    :param slice_id: The slice ID
+    :raises SupersetSecurityException: If the user cannot access the resource
     """
 
     form_data, slc = get_form_data(slice_id, use_slice_data=True)
@@ -454,7 +450,7 @@ def check_slice_perms(_self: Any, slice_id: int) -> None:
             force=False,
         )
 
-        security_manager.assert_viz_permission(viz_obj)
+        viz_obj.raise_for_access()
 
 
 def _deserialize_results_payload(
diff --git a/superset/viz.py b/superset/viz.py
index 154841c..79f7a1d 100644
--- a/superset/viz.py
+++ b/superset/viz.py
@@ -552,6 +552,15 @@ class BaseViz:
     def json_data(self) -> str:
         return json.dumps(self.data)
 
+    def raise_for_access(self) -> None:
+        """
+        Raise an exception if the user cannot access the resource.
+
+        :raises SupersetSecurityException: If the user cannot access the 
resource
+        """
+
+        security_manager.raise_for_access(viz=self)
+
 
 class TableViz(BaseViz):
 
diff --git a/superset/viz_sip38.py b/superset/viz_sip38.py
index a3e00de..9ef4841 100644
--- a/superset/viz_sip38.py
+++ b/superset/viz_sip38.py
@@ -592,6 +592,15 @@ class BaseViz:
     def json_data(self):
         return json.dumps(self.data)
 
+    def raise_for_access(self) -> None:
+        """
+        Raise an exception if the user cannot access the resource.
+
+        :raises SupersetSecurityException: If the user cannot access the 
resource
+        """
+
+        security_manager.raise_for_access(viz=self)
+
 
 class TableViz(BaseViz):
 
diff --git a/tests/security_tests.py b/tests/security_tests.py
index 9426036..a36a499 100644
--- a/tests/security_tests.py
+++ b/tests/security_tests.py
@@ -26,9 +26,11 @@ import tests.test_app
 from superset import app, appbuilder, db, security_manager, viz
 from superset.connectors.druid.models import DruidCluster, DruidDatasource
 from superset.connectors.sqla.models import RowLevelSecurityFilter, SqlaTable
+from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
 from superset.exceptions import SupersetSecurityException
 from superset.models.core import Database
 from superset.models.slice import Slice
+from superset.sql_parse import Table
 from superset.utils.core import get_example_database
 
 from .base_tests import SupersetTestCase
@@ -774,48 +776,111 @@ class SecurityManagerTests(SupersetTestCase):
     Testing the Security Manager.
     """
 
-    @patch("superset.security.SupersetSecurityManager.can_access_datasource")
-    def test_assert_datasource_permission(self, mock_can_access_datasource):
+    @patch("superset.security.SupersetSecurityManager.raise_for_access")
+    def test_can_access_datasource(self, mock_raise_for_access):
         datasource = self.get_datasource_mock()
 
-        # Datasource with the "datasource_access" permission.
-        mock_can_access_datasource.return_value = True
-        security_manager.assert_datasource_permission(datasource)
+        mock_raise_for_access.return_value = None
+        
self.assertTrue(security_manager.can_access_datasource(datasource=datasource))
 
-        # Datasource without the "datasource_access" permission.
-        mock_can_access_datasource.return_value = False
+        mock_raise_for_access.side_effect = SupersetSecurityException(
+            SupersetError(
+                "dummy",
+                SupersetErrorType.DATASOURCE_SECURITY_ACCESS_ERROR,
+                ErrorLevel.ERROR,
+            )
+        )
+
+        
self.assertFalse(security_manager.can_access_datasource(datasource=datasource))
+
+    @patch("superset.security.SupersetSecurityManager.raise_for_access")
+    def test_can_access_table(self, mock_raise_for_access):
+        database = get_example_database()
+        table = Table("bar", "foo")
+
+        mock_raise_for_access.return_value = None
+        self.assertTrue(security_manager.can_access_table(database, table))
+
+        mock_raise_for_access.side_effect = SupersetSecurityException(
+            SupersetError(
+                "dummy",
+                SupersetErrorType.TABLE_SECURITY_ACCESS_ERROR,
+                ErrorLevel.ERROR,
+            )
+        )
+
+        self.assertFalse(security_manager.can_access_table(database, table))
+
+    @patch("superset.security.SupersetSecurityManager.can_access")
+    @patch("superset.security.SupersetSecurityManager.can_access_schema")
+    def test_raise_for_access_datasource(self, mock_can_access_schema, 
mock_can_access):
+        datasource = self.get_datasource_mock()
+
+        mock_can_access_schema.return_value = True
+        security_manager.raise_for_access(datasource=datasource)
+
+        mock_can_access.return_value = False
+        mock_can_access_schema.return_value = False
+
+        with self.assertRaises(SupersetSecurityException):
+            security_manager.raise_for_access(datasource=datasource)
+
+    @patch("superset.security.SupersetSecurityManager.can_access")
+    def test_raise_for_access_query(self, mock_can_access):
+        query = Mock(
+            database=get_example_database(), schema="bar", sql="SELECT * FROM 
foo"
+        )
+
+        mock_can_access.return_value = True
+        security_manager.raise_for_access(query=query)
+
+        mock_can_access.return_value = False
+
+        with self.assertRaises(SupersetSecurityException):
+            security_manager.raise_for_access(query=query)
+
+    @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
+    ):
+        query_context = Mock(datasource=self.get_datasource_mock())
+
+        mock_can_access_schema.return_value = True
+        security_manager.raise_for_access(query_context=query_context)
+
+        mock_can_access.return_value = False
+        mock_can_access_schema.return_value = False
 
         with self.assertRaises(SupersetSecurityException):
-            security_manager.assert_datasource_permission(datasource)
+            security_manager.raise_for_access(query_context=query_context)
 
-    @patch("superset.security.SupersetSecurityManager.can_access_datasource")
-    def test_assert_query_context_permission(self, mock_can_access_datasource):
-        query_context = Mock()
-        query_context.datasource = self.get_datasource_mock()
+    @patch("superset.security.SupersetSecurityManager.can_access")
+    def test_raise_for_access_table(self, mock_can_access):
+        database = get_example_database()
+        table = Table("bar", "foo")
 
-        # Query context with the "datasource_access" permission.
-        mock_can_access_datasource.return_value = True
-        security_manager.assert_query_context_permission(query_context)
+        mock_can_access.return_value = True
+        security_manager.raise_for_access(database=database, table=table)
 
-        # Query context without the "datasource_access" permission.
-        mock_can_access_datasource.return_value = False
+        mock_can_access.return_value = False
 
         with self.assertRaises(SupersetSecurityException):
-            security_manager.assert_query_context_permission(query_context)
+            security_manager.raise_for_access(database=database, table=table)
 
-    @patch("superset.security.SupersetSecurityManager.can_access_datasource")
-    def test_assert_viz_permission(self, mock_can_access_datasource):
+    @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):
         test_viz = viz.TableViz(self.get_datasource_mock(), form_data={})
 
-        # Visualization with the "datasource_access" permission.
-        mock_can_access_datasource.return_value = True
-        security_manager.assert_viz_permission(test_viz)
+        mock_can_access_schema.return_value = True
+        security_manager.raise_for_access(viz=test_viz)
 
-        # Visualization without the "datasource_access" permission.
-        mock_can_access_datasource.return_value = False
+        mock_can_access.return_value = False
+        mock_can_access_schema.return_value = False
 
         with self.assertRaises(SupersetSecurityException):
-            security_manager.assert_viz_permission(test_viz)
+            security_manager.raise_for_access(viz=test_viz)
 
 
 class RowLevelSecurityTests(SupersetTestCase):

Reply via email to