fix: create virtual dataset validation (#26625)

(cherry picked from commit 8e19f59)

superset/commands/dataset/create.py

@@ -25,13 +25,15 @@
 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 @@ def validate(self) -> None:
         ):
             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

superset/commands/dataset/exceptions.py

@@ -144,6 +144,13 @@ def __init__(self) -> None:
         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")

superset/security/manager.py

@@ -1828,6 +1828,8 @@ def raise_for_access(
         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 @@ def raise_for_access(
         :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 @@ def raise_for_access(
         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:

tests/integration_tests/datasets/commands_tests.py

@@ -38,7 +38,7 @@
 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 @@ def test_import_v1_dataset_existing_database(self, mock_g):
             "metadata.yaml": yaml.safe_dump(database_metadata_config),
             "databases/imported_database.yaml": yaml.safe_dump(database_config),
         }
+
         command = ImportDatabasesCommand(contents)
         command.run()
 
@@ -558,34 +559,46 @@ def test_get_table_from_database_error(self, get_table_mock):
                 {"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")
             engine.execute(
                 "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):