Fix SQL Lab schema permission checks (#9756)

Co-authored-by: bogdan kyryliuk <bogdankyryliuk@dropbox.com>

superset/security/manager.py

@@ -333,7 +333,7 @@ def can_access_datasource(
         if self.database_access(database) or self.all_datasource_access():
             return True
 
-        schema_perm = self.get_schema_perm(database, schema)
+        schema_perm = self.get_schema_perm(database, schema=table.schema or schema)
         if schema_perm and self.can_access("schema_access", schema_perm):
             return True
 
@@ -356,7 +356,6 @@ def rejected_tables(
         :param schema: The SQL database schema
         :returns: The rejected tables
         """
-
         query = sql_parse.ParsedQuery(sql)
 
         return {

tests/base_tests.py

@@ -18,7 +18,7 @@
 """Unit tests for Superset"""
 import imp
 import json
-from typing import Dict, Union
+from typing import Dict, Union, List
 from unittest.mock import Mock, patch
 
 import pandas as pd
@@ -58,6 +58,24 @@ def __init__(self, *args, **kwargs):
     def create_app(self):
         return app
 
+    @staticmethod
+    def create_user_with_roles(username: str, roles: List[str]):
+        user_to_create = security_manager.find_user(username)
+        if not user_to_create:
+            security_manager.add_user(
+                username,
+                username,
+                username,
+                f"{username}@superset.com",
+                security_manager.find_role("Gamma"),  # it needs a role
+                password="general",
+            )
+            db.session.commit()
+            user_to_create = security_manager.find_user(username)
+            assert user_to_create
+        user_to_create.roles = [security_manager.find_role(r) for r in roles]
+        db.session.commit()
+
     @staticmethod
     def create_user(
         username: str,
@@ -240,6 +258,7 @@ def run_sql(
         sql_editor_id=None,
         select_as_cta=False,
         tmp_table_name=None,
+        schema=None,
     ):
         if user_name:
             self.logout()
@@ -256,6 +275,8 @@ def run_sql(
             json_payload["tmp_table_name"] = tmp_table_name
         if select_as_cta:
             json_payload["select_as_cta"] = select_as_cta
+        if schema:
+            json_payload["schema"] = schema
 
         resp = self.get_json_resp(
             "/superset/sql_json/", raise_on_error=False, json_=json_payload

tests/sqllab_tests.py

@@ -24,9 +24,8 @@
 import prison
 
 import tests.test_app
-from superset import config, db, security_manager
+from superset import db, security_manager
 from superset.connectors.sqla.models import SqlaTable
-from superset.dataframe import df_to_records
 from superset.db_engine_specs import BaseEngineSpec
 from superset.models.sql_lab import Query
 from superset.result_set import SupersetResultSet
@@ -122,30 +121,65 @@ def test_sql_json_has_access(self):
         examples_db_permission_view = security_manager.add_permission_view_menu(
             "database_access", examples_db.perm
         )
-
-        astronaut = security_manager.add_role("Astronaut")
+        astronaut = security_manager.add_role("ExampleDBAccess")
         security_manager.add_permission_role(astronaut, examples_db_permission_view)
-        # Astronaut role is Gamma + sqllab + db permissions
-        for perm in security_manager.find_role("Gamma").permissions:
-            security_manager.add_permission_role(astronaut, perm)
-        for perm in security_manager.find_role("sql_lab").permissions:
-            security_manager.add_permission_role(astronaut, perm)
-
-        gagarin = security_manager.find_user("gagarin")
-        if not gagarin:
-            security_manager.add_user(
-                "gagarin",
-                "Iurii",
-                "Gagarin",
-                "gagarin@cosmos.ussr",
-                astronaut,
-                password="general",
-            )
-        data = self.run_sql(QUERY_1, "3", user_name="gagarin")
+        # Gamma user, with sqllab and db permission
+        self.create_user_with_roles("Gagarin", ["ExampleDBAccess", "Gamma", "sql_lab"])
+
+        data = self.run_sql(QUERY_1, "1", user_name="Gagarin")
         db.session.query(Query).delete()
         db.session.commit()
         self.assertLess(0, len(data["data"]))
 
+    def test_sql_json_schema_access(self):
+        examples_db = get_example_database()
+        db_backend = examples_db.backend
+        if db_backend == "sqlite":
+            # sqlite doesn't support database creation
+            return
+
+        sqllab_test_db_schema_permission_view = security_manager.add_permission_view_menu(
+            "schema_access", f"[{examples_db.name}].[sqllab_test_db]"
+        )
+        schema_perm_role = security_manager.add_role("SchemaPermission")
+        security_manager.add_permission_role(
+            schema_perm_role, sqllab_test_db_schema_permission_view
+        )
+        self.create_user_with_roles(
+            "SchemaUser", ["SchemaPermission", "Gamma", "sql_lab"]
+        )
+
+        db.session.execute(
+            "CREATE TABLE IF NOT EXISTS sqllab_test_db.test_table AS SELECT 1 as c1, 2 as c2"
+        )
+
+        data = self.run_sql(
+            "SELECT * FROM sqllab_test_db.test_table", "3", user_name="SchemaUser"
+        )
+        self.assertEqual(1, len(data["data"]))
+
+        data = self.run_sql(
+            "SELECT * FROM sqllab_test_db.test_table",
+            "4",
+            user_name="SchemaUser",
+            schema="sqllab_test_db",
+        )
+        self.assertEqual(1, len(data["data"]))
+
+        # postgres needs a schema as a part of the table name.
+        if db_backend == "mysql":
+            data = self.run_sql(
+                "SELECT * FROM test_table",
+                "5",
+                user_name="SchemaUser",
+                schema="sqllab_test_db",
+            )
+            self.assertEqual(1, len(data["data"]))
+
+        db.session.query(Query).delete()
+        db.session.execute("DROP TABLE IF EXISTS sqllab_test_db.test_table")
+        db.session.commit()
+
     def test_queries_endpoint(self):
         self.run_some_queries()