Adding permission for can_only_access_owned_queries

superset/models/sql_lab.py

@@ -138,6 +138,14 @@ def name(self):
         tab = re.sub(r'\W+', '', tab)
         return f'sqllab_{tab}_{ts}'
 
+    @property
+    def database_name(self):
+        return self.database.name
+
+    @property
+    def username(self):
+        return self.user.username
+
 
 class SavedQuery(Model, AuditMixinNullable, ExtraJSONMixin):
     """ORM model for SQL query"""

superset/security.py

@@ -92,6 +92,7 @@ class SupersetSecurityManager(SecurityManager):
         'schema_access',
         'datasource_access',
         'metric_access',
+        'can_only_access_owned_queries',
     ])
 
     def get_schema_perm(self, database, schema):
@@ -105,6 +106,17 @@ def can_access(self, permission_name, view_name):
             return self.is_item_public(permission_name, view_name)
         return self._has_view_access(user, permission_name, view_name)
 
+    def can_only_access_owned_queries(self) -> bool:
+        """
+        can_access check for custom can_only_access_owned_queries permissions.
+
+        :returns: True if current user can access custom permissions
+        """
+        return self.can_access(
+            'can_only_access_owned_queries',
+            'can_only_access_owned_queries',
+        )
+
     def all_datasource_access(self):
         return self.can_access('all_datasource_access', 'all_datasource_access')
 
@@ -268,6 +280,7 @@ def create_custom_permissions(self):
         # Global perms
         self.merge_perm('all_datasource_access', 'all_datasource_access')
         self.merge_perm('all_database_access', 'all_database_access')
+        self.merge_perm('can_only_access_owned_queries', 'can_only_access_owned_queries')
 
     def create_missing_perms(self):
         """Creates missing perms for datasources, schemas and metrics"""

superset/views/core.py

@@ -22,6 +22,7 @@
 import re
 import time
 import traceback
+from typing import List  # noqa: F401
 from urllib import parse
 
 from flask import (
@@ -82,7 +83,7 @@
     'The access requests seem to have been deleted')
 USER_MISSING_ERR = __('The user seems to have been deleted')
 
-FORM_DATA_KEY_BLACKLIST = []
+FORM_DATA_KEY_BLACKLIST: List[str] = []
 if not config.get('ENABLE_JAVASCRIPT_CONTROLS'):
     FORM_DATA_KEY_BLACKLIST = [
         'js_tooltip',
@@ -2766,10 +2767,21 @@ def queries(self, last_updated_ms):
     @has_access
     @expose('/search_queries')
     @log_this
-    def search_queries(self):
-        """Search for queries."""
+    def search_queries(self) -> Response:
+        """
+        Search for previously run sqllab queries. Used for Sqllab Query Search
+        page /superset/sqllab#search.
+
+        Custom permission can_only_search_queries_owned restricts queries
+        to only queries run by current user.
+
+        :returns: Response with list of sql query dicts
+        """
         query = db.session.query(Query)
-        search_user_id = request.args.get('user_id')
+        if security_manager.can_only_access_owned_queries():
+            search_user_id = g.user.get_user_id()
+        else:
+            search_user_id = request.args.get('user_id')
         database_id = request.args.get('database_id')
         search_text = request.args.get('search_text')
         status = request.args.get('status')
@@ -2778,7 +2790,7 @@ def search_queries(self):
         to_time = request.args.get('to')
 
         if search_user_id:
-            # Filter on db Id
+            # Filter on user_id
             query = query.filter(Query.user_id == search_user_id)
 
         if database_id:

superset/views/sql_lab.py

@@ -15,16 +15,38 @@
 # specific language governing permissions and limitations
 # under the License.
 # pylint: disable=C,R,W
+from typing import Callable
+
 from flask import g, redirect
 from flask_appbuilder import expose
 from flask_appbuilder.models.sqla.interface import SQLAInterface
 from flask_appbuilder.security.decorators import has_access
 from flask_babel import gettext as __
 from flask_babel import lazy_gettext as _
+from flask_sqlalchemy import BaseQuery
 
-from superset import appbuilder
+from superset import appbuilder, security_manager
 from superset.models.sql_lab import Query, SavedQuery
-from .base import BaseSupersetView, DeleteMixin, SupersetModelView
+from .base import BaseSupersetView, DeleteMixin, SupersetFilter, SupersetModelView
+
+
+class QueryFilter(SupersetFilter):
+    def apply(
+            self,
+            query: BaseQuery,
+            func: Callable) -> BaseQuery:
+        """
+        Filter queries to only those owned by current user if
+        can_only_access_owned_queries permission is set.
+
+        :returns: query
+        """
+        if security_manager.can_only_access_owned_queries():
+            query = (
+                query
+                .filter(Query.user_id == g.user.get_user_id())
+            )
+        return query
 
 
 class QueryView(SupersetModelView):
@@ -35,10 +57,12 @@ class QueryView(SupersetModelView):
     add_title = _('Add Query')
     edit_title = _('Edit Query')
 
-    list_columns = ['user', 'database', 'status', 'start_time', 'end_time']
+    list_columns = ['username', 'database_name', 'status', 'start_time', 'end_time']
+    base_filters = [['id', QueryFilter, lambda: []]]
     label_columns = {
         'user': _('User'),
-        'database': _('Database'),
+        'username': _('User'),
+        'database_name': _('Database'),
         'status': _('Status'),
         'start_time': _('Start Time'),
         'end_time': _('End Time'),

tests/sqllab_tests.py

@@ -53,9 +53,10 @@ def run_some_queries(self):
         self.logout()
 
     def tearDown(self):
+        self.logout()
         db.session.query(Query).delete()
         db.session.commit()
-        self.logout()
+        db.session.close()
 
     def test_sql_json(self):
         self.login('admin')
@@ -223,6 +224,46 @@ def test_search_query_on_time(self):
         data = json.loads(resp)
         self.assertEquals(2, len(data))
 
+    def test_search_query_with_owner_only_perms(self) -> None:
+        """
+        Test a search query with can_only_access_owned_queries perm added to
+        Admin and make sure only Admin queries show up.
+        """
+        session = db.session
+
+        # Add can_only_access_owned_queries perm to Admin user
+        owned_queries_view = security_manager.find_permission_view_menu(
+            'can_only_access_owned_queries',
+            'can_only_access_owned_queries',
+        )
+        security_manager.add_permission_role(
+            security_manager.find_role('Admin'),
+            owned_queries_view,
+        )
+        session.commit()
+
+        # Test search_queries for Admin user
+        self.run_some_queries()
+        self.login('admin')
+
+        user_id = security_manager.find_user('admin').id
+        data = self.get_json_resp('/superset/search_queries')
+        self.assertEquals(2, len(data))
+        user_ids = {k['userId'] for k in data}
+        self.assertEquals(set([user_id]), user_ids)
+
+        # Remove can_only_access_owned_queries from Admin
+        owned_queries_view = security_manager.find_permission_view_menu(
+            'can_only_access_owned_queries',
+            'can_only_access_owned_queries',
+        )
+        security_manager.del_permission_role(
+            security_manager.find_role('Admin'),
+            owned_queries_view,
+        )
+
+        session.commit()
+
     def test_alias_duplicate(self):
         self.run_sql(
             'SELECT username as col, id as col, username FROM ab_user',
@@ -309,6 +350,68 @@ def test_sql_limit(self):
             query_limit=test_limit)
         self.assertEquals(len(data['data']), test_limit)
 
+    def test_queryview_filter(self) -> None:
+        """
+        Test queryview api without can_only_access_owned_queries perm added to
+        Admin and make sure all queries show up.
+        """
+        self.run_some_queries()
+        self.login(username='admin')
+
+        url = '/queryview/api/read'
+        data = self.get_json_resp(url)
+        admin = security_manager.find_user('admin')
+        gamma_sqllab = security_manager.find_user('gamma_sqllab')
+        self.assertEquals(3, len(data['result']))
+        user_queries = [
+            result.get('username') for result in data['result']
+        ]
+        assert admin.username in user_queries
+        assert gamma_sqllab.username in user_queries
+
+    def test_queryview_filter_owner_only(self) -> None:
+        """
+        Test queryview api with can_only_access_owned_queries perm added to
+        Admin and make sure only Admin queries show up.
+        """
+        session = db.session
+
+        # Add can_only_access_owned_queries perm to Admin user
+        owned_queries_view = security_manager.find_permission_view_menu(
+            'can_only_access_owned_queries',
+            'can_only_access_owned_queries',
+        )
+        security_manager.add_permission_role(
+            security_manager.find_role('Admin'),
+            owned_queries_view,
+        )
+        session.commit()
+
+        # Test search_queries for Admin user
+        self.run_some_queries()
+        self.login('admin')
+
+        url = '/queryview/api/read'
+        data = self.get_json_resp(url)
+        admin = security_manager.find_user('admin')
+        self.assertEquals(2, len(data['result']))
+        all_admin_user_queries = all([
+            result.get('username') == admin.username for result in data['result']
+        ])
+        assert all_admin_user_queries is True
+
+        # Remove can_only_access_owned_queries from Admin
+        owned_queries_view = security_manager.find_permission_view_menu(
+            'can_only_access_owned_queries',
+            'can_only_access_owned_queries',
+        )
+        security_manager.del_permission_role(
+            security_manager.find_role('Admin'),
+            owned_queries_view,
+        )
+
+        session.commit()
+
 
 if __name__ == '__main__':
     unittest.main()