Merge bfff35b6ed into f5843fe588

superset/daos/datasource.py

@@ -26,6 +26,7 @@ from superset.datasets.models import Dataset
 from superset.models.sql_lab import Query, SavedQuery
 from superset.tables.models import Table
 from superset.utils.core import DatasourceType
+from superset.exceptions import SupersetSecurityException
 
 logger = logging.getLogger(__name__)
 
@@ -63,5 +64,10 @@ class DatasourceDAO(BaseDAO[Datasource]):
                 datasource_id,
             )
             raise DatasourceNotFound()
+        
+        try:
+            datasource.raise_for_access()
+        except SupersetSecurityException as ex:
+            raise DatasourceNotFound() from ex
 
         return datasource

superset/datasets/models.py

@@ -116,3 +116,12 @@ class Dataset(AuditMixinNullable, ExtraJSONMixin, ImportExportMixin, Model):
 
     def __repr__(self) -> str:
         return f"<Dataset id={self.id} database_id={self.database_id} {self.name}>"
+    
+    
+    def raise_for_access(self) -> None:
+        """
+        Raise an exception if the user cannot access the Dataset.
+        :raises SupersetSecurityException: If the user cannot access the Dataset
+        """
+
+        security_manager.raise_for_access(query=self)

superset/models/sql_lab.py

@@ -462,6 +462,15 @@ class SavedQuery(
     def last_run_delta_humanized(self) -> str:
         return self._last_run_delta_humanized
 
+    
+    def raise_for_access(self) -> None:
+        """
+        Raise an exception if the user cannot access the SavedQuery.
+        :raises SupersetSecurityException: If the user cannot access the SavedQuery
+        """
+
+        security_manager.raise_for_access(query=self)
+
 
 class TabState(AuditMixinNullable, ExtraJSONMixin, Model):
     __tablename__ = "tab_state"

superset/tables/models.py

@@ -44,6 +44,7 @@ from superset.models.helpers import (
 )
 from superset.sql_parse import Table as TableName
 from superset.superset_typing import ResultSetColumnType
+from superset import security_manager
 
 if TYPE_CHECKING:
     from superset.datasets.models import Dataset
@@ -204,3 +205,11 @@ class Table(AuditMixinNullable, ExtraJSONMixin, ImportExportMixin, Model):
                 session.add(new_table)
 
         return all_tables
+    
+    def raise_for_access(self) -> None:
+        """
+        Raise an exception if the user cannot access the Table.
+        :raises SupersetSecurityException: If the user cannot access the Table
+        """
+
+        security_manager.raise_for_access(query=self)

tests/unit_tests/datasource/dao_tests.py

@@ -16,12 +16,14 @@
 # under the License.
 
 from collections.abc import Iterator
+from flask_appbuilder.security.sqla.models import User
 
 import pytest
+from pytest_mock import MockFixture
 from sqlalchemy.orm.session import Session
 
 from superset.utils.core import DatasourceType
-
+from superset.utils.core import override_user
 
 @pytest.fixture
 def session_with_data(session: Session) -> Iterator[Session]:
@@ -99,9 +101,12 @@ FROM my_catalog.my_schema.my_table
     yield session
 
 
-def test_get_datasource_sqlatable(session_with_data: Session) -> None:
+def test_get_datasource_sqlatable(mocker: MockFixture, session_with_data: Session) -> None:
     from superset.connectors.sqla.models import SqlaTable
     from superset.daos.datasource import DatasourceDAO
+    from superset import security_manager
+
+    mocker.patch.object(security_manager, "can_access", return_value=True)
 
     result = DatasourceDAO.get_datasource(
         datasource_type=DatasourceType.TABLE,
@@ -112,11 +117,29 @@ def test_get_datasource_sqlatable(session_with_data: Session) -> None:
     assert "my_sqla_table" == result.table_name
     assert isinstance(result, SqlaTable)
 
+def test_get_datasource_sqlatable_false(mocker: MockFixture, session_with_data: Session) -> None:
+    from superset.daos.datasource import DatasourceDAO
+    from superset import security_manager
+    from superset.daos.exceptions import DatasourceNotFound
 
-def test_get_datasource_query(session_with_data: Session) -> None:
+
+    mocker.patch.object(security_manager, "can_access", return_value=False)
+
+    with pytest.raises(DatasourceNotFound):
+        with override_user(User()):
+            DatasourceDAO.get_datasource(
+                datasource_type=DatasourceType.TABLE,
+                datasource_id=1,
+                session=session_with_data,
+            )
+
+def test_get_datasource_query(mocker: MockFixture, session_with_data: Session) -> None:
     from superset.daos.datasource import DatasourceDAO
     from superset.models.sql_lab import Query
+    from superset import security_manager
 
+    mocker.patch.object(security_manager, "can_access", return_value=True)
+    
     result = DatasourceDAO.get_datasource(
         datasource_type=DatasourceType.QUERY, datasource_id=1
     )
@@ -124,10 +147,30 @@ def test_get_datasource_query(session_with_data: Session) -> None:
     assert result.id == 1
     assert isinstance(result, Query)
 
+def test_get_datasource_query_false(mocker: MockFixture, session_with_data: Session) -> None:
+    from superset.daos.datasource import DatasourceDAO
+    from superset.models.sql_lab import Query
+    from superset import security_manager
+    from superset.daos.exceptions import DatasourceNotFound
 
-def test_get_datasource_saved_query(session_with_data: Session) -> None:
+
+    mocker.patch.object(security_manager, "can_access", return_value=False)
+
+    with pytest.raises(DatasourceNotFound):
+        with override_user(User()):
+            DatasourceDAO.get_datasource(
+                datasource_type=DatasourceType.QUERY,
+                datasource_id=1,
+                session=session_with_data,
+            )
+
+
+def test_get_datasource_saved_query(mocker: MockFixture, session_with_data: Session) -> None:
     from superset.daos.datasource import DatasourceDAO
     from superset.models.sql_lab import SavedQuery
+    from superset import security_manager
+
+    mocker.patch.object(security_manager, "can_access", return_value=True)
 
     result = DatasourceDAO.get_datasource(
         datasource_type=DatasourceType.SAVEDQUERY,
@@ -137,10 +180,27 @@ def test_get_datasource_saved_query(session_with_data: Session) -> None:
     assert result.id == 1
     assert isinstance(result, SavedQuery)
 
+def test_get_datasource_saved_query_false(mocker: MockFixture, session_with_data: Session) -> None:
+    from superset.daos.datasource import DatasourceDAO
+    from superset import security_manager
+    from superset.daos.exceptions import DatasourceNotFound
 
-def test_get_datasource_sl_table(session_with_data: Session) -> None:
+    mocker.patch.object(security_manager, "can_access", return_value=False)
+
+    with pytest.raises(DatasourceNotFound):
+        with override_user(User()):
+            DatasourceDAO.get_datasource(
+                datasource_type=DatasourceType.SAVEDQUERY,
+                datasource_id=1,
+                session=session_with_data,
+            )
+
+def test_get_datasource_sl_table(mocker: MockFixture, session_with_data: Session) -> None:
     from superset.daos.datasource import DatasourceDAO
     from superset.tables.models import Table
+    from superset import security_manager
+
+    mocker.patch.object(security_manager, "can_access", return_value=True)
 
     result = DatasourceDAO.get_datasource(
         datasource_type=DatasourceType.SLTABLE,
@@ -150,10 +210,12 @@ def test_get_datasource_sl_table(session_with_data: Session) -> None:
     assert result.id == 1
     assert isinstance(result, Table)
 
-
-def test_get_datasource_sl_dataset(session_with_data: Session) -> None:
+def test_get_datasource_sl_dataset(mocker: MockFixture, session_with_data: Session) -> None:
     from superset.daos.datasource import DatasourceDAO
     from superset.datasets.models import Dataset
+    from superset import security_manager
+
+    mocker.patch.object(security_manager, "can_access", return_value=True)
 
     result = DatasourceDAO.get_datasource(
         datasource_type=DatasourceType.DATASET,
@@ -163,11 +225,13 @@ def test_get_datasource_sl_dataset(session_with_data: Session) -> None:
     assert result.id == 1
     assert isinstance(result, Dataset)
 
-
-def test_get_datasource_w_str_param(session_with_data: Session) -> None:
+def test_get_datasource_w_str_param(mocker: MockFixture, session_with_data: Session) -> None:
     from superset.connectors.sqla.models import SqlaTable
     from superset.daos.datasource import DatasourceDAO
     from superset.tables.models import Table
+    from superset import security_manager
+
+    mocker.patch.object(security_manager, "can_access", return_value=True)
 
     assert isinstance(
         DatasourceDAO.get_datasource(