Updating permission when refreshing druid datasource (#2655)

* Updating permission when refreshing druid datasource

* Adding test

* Fix style

* Deletion view_menu after db, table, cluster, ds deletion

* Update table model

* Linting

* Override _delete instead of post_delete

* fix

* lint

* fix multi delete

* fix

* Refactoring

* Amending

superset/connectors/druid/models.py

@@ -108,6 +108,9 @@ class DruidCluster(Model, AuditMixinNullable):
     def perm(self):
         return "[{obj.cluster_name}].(id:{obj.id})".format(obj=self)
 
+    def get_perm(self):
+        return self.perm
+
     @property
     def name(self):
         return self.verbose_name if self.verbose_name else self.cluster_name

superset/connectors/druid/views.py

@@ -3,7 +3,7 @@ import logging
 
 import sqlalchemy as sqla
 
-from flask import Markup, flash, redirect
+from flask import Markup, flash, redirect, abort
 from flask_appbuilder import CompactCRUDMixin, expose
 from flask_appbuilder.models.sqla.interface import SQLAInterface
 
@@ -136,6 +136,8 @@ class DruidClusterModelView(SupersetModelView, DeleteMixin):  # noqa
     def pre_update(self, cluster):
         self.pre_add(cluster)
 
+    def _delete(self, pk):
+        DeleteMixin._delete(self, pk)
 
 appbuilder.add_view(
     DruidClusterModelView,
@@ -231,6 +233,9 @@ class DruidDatasourceModelView(SupersetModelView, DeleteMixin):  # noqa
     def post_update(self, datasource):
         self.post_add(datasource)
 
+    def _delete(self, pk):
+        DeleteMixin._delete(self, pk)
+
 appbuilder.add_view(
     DruidDatasourceModelView,
     "Druid Datasources",

superset/connectors/sqla/views.py

@@ -3,7 +3,7 @@ import logging
 
 from past.builtins import basestring
 
-from flask import Markup, flash, redirect
+from flask import Markup, flash, redirect, abort
 from flask_appbuilder import CompactCRUDMixin, expose
 from flask_appbuilder.models.sqla.interface import SQLAInterface
 import sqlalchemy as sa
@@ -242,6 +242,9 @@ class TableModelView(SupersetModelView, DeleteMixin):  # noqa
     def post_update(self, table):
         self.post_add(table, flash_message=False)
 
+    def _delete(self, pk):
+        DeleteMixin._delete(self, pk)
+
     @expose('/edit/<pk>', methods=['GET', 'POST'])
     @has_access
     def edit(self, pk):

superset/models/helpers.py

@@ -10,6 +10,7 @@ from flask import escape, Markup
 from flask_appbuilder.models.mixins import AuditMixin
 from flask_appbuilder.models.decorators import renders
 from superset.utils import QueryStatus
+from superset import sm
 
 
 class ImportMixin(object):
@@ -117,7 +118,44 @@ class QueryResult(object):
         self.error_message = error_message
 
 
+def merge_perm(sm, permission_name, view_menu_name, connection):
+
+    permission = sm.find_permission(permission_name)
+    view_menu = sm.find_view_menu(view_menu_name)
+    pv = None
+
+    if not permission:
+        permission_table = sm.permission_model.__table__
+        connection.execute(
+            permission_table.insert()
+            .values(name=permission_name)
+        )
+    if not view_menu:
+        view_menu_table = sm.viewmenu_model.__table__
+        connection.execute(
+            view_menu_table.insert()
+            .values(name=view_menu_name)
+        )
+
+    permission = sm.find_permission(permission_name)
+    view_menu = sm.find_view_menu(view_menu_name)
+
+    if permission and view_menu:
+        pv = sm.get_session.query(sm.permissionview_model).filter_by(
+            permission=permission, view_menu=view_menu).first()
+    if not pv and permission and view_menu:
+        permission_view_table = sm.permissionview_model.__table__
+        connection.execute(
+            permission_view_table.insert()
+            .values(
+                permission_id=permission.id,
+                view_menu_id=view_menu.id
+                )
+        )
+
+
 def set_perm(mapper, connection, target):  # noqa
+
     if target.perm != target.get_perm():
         link_table = target.__table__
         connection.execute(
@@ -125,3 +163,6 @@ def set_perm(mapper, connection, target):  # noqa
             .where(link_table.c.id == target.id)
             .values(perm=target.get_perm())
         )
+
+    # add to view menu if not already exists
+    merge_perm(sm, 'datasource_access', target.get_perm(), connection)

superset/views/base.py

@@ -3,7 +3,7 @@ import json
 import logging
 import traceback
 
-from flask import g, redirect, Response
+from flask import g, redirect, Response, flash, abort
 from flask_babel import gettext as __
 
 from flask_appbuilder import BaseView
@@ -207,6 +207,51 @@ def validate_json(form, field):  # noqa
 
 
 class DeleteMixin(object):
+    def _delete(self, pk):
+        """
+            Delete function logic, override to implement diferent logic
+            deletes the record with primary_key = pk
+
+            :param pk:
+                record primary key to delete
+        """
+        item = self.datamodel.get(pk, self._base_filters)
+        if not item:
+            abort(404)
+        try:
+            self.pre_delete(item)
+        except Exception as e:
+            flash(str(e), "danger")
+        else:
+            view_menu = sm.find_view_menu(item.get_perm())
+            pvs = sm.get_session.query(sm.permissionview_model).filter_by(
+                view_menu=view_menu).all()
+
+            schema_view_menu = None
+            if hasattr(item, 'schema_perm'):
+                schema_view_menu = sm.find_view_menu(item.schema_perm)
+
+                pvs.extend(sm.get_session.query(
+                    sm.permissionview_model).filter_by(
+                    view_menu=schema_view_menu).all())
+
+            if self.datamodel.delete(item):
+                self.post_delete(item)
+
+                for pv in pvs:
+                    sm.get_session.delete(pv)
+
+                if view_menu:
+                    sm.get_session.delete(view_menu)
+
+                if schema_view_menu:
+                    sm.get_session.delete(schema_view_menu)
+
+                sm.get_session.commit()
+
+            flash(*self.datamodel.message)
+            self.update_redirect()
+
     @action(
         "muldelete",
         __("Delete"),
@@ -215,7 +260,15 @@ class DeleteMixin(object):
         single=False
     )
     def muldelete(self, items):
-        self.datamodel.delete_all(items)
+        if not items:
+            abort(404)
+        for item in items:
+            try:
+                self.pre_delete(item)
+            except Exception as e:
+                flash(str(e), "danger")
+            else:
+                self._delete(item.id)
         self.update_redirect()
         return redirect(self.get_redirect())
 

superset/views/core.py

@@ -16,7 +16,7 @@ import traceback
 import sqlalchemy as sqla
 
 from flask import (
-    g, request, redirect, flash, Response, render_template, Markup)
+    g, request, redirect, flash, Response, render_template, Markup, abort)
 from flask_appbuilder import expose
 from flask_appbuilder.actions import action
 from flask_appbuilder.models.sqla.interface import SQLAInterface
@@ -252,6 +252,8 @@ class DatabaseView(SupersetModelView, DeleteMixin):  # noqa
     def pre_update(self, db):
         self.pre_add(db)
 
+    def _delete(self, pk):
+        DeleteMixin._delete(self, pk)
 
 appbuilder.add_link(
     'Import Dashboards',

tests/druid_tests.py

@@ -276,6 +276,49 @@ class DruidTests(SupersetTestCase):
         self.assertIn('datasource_for_gamma', resp)
         self.assertNotIn('datasource_not_for_gamma', resp)
 
+    @patch('superset.connectors.druid.models.PyDruid')
+    def test_sync_druid_perm(self, PyDruid):
+        self.login(username='admin')
+        instance = PyDruid.return_value
+        instance.time_boundary.return_value = [
+            {'result': {'maxTime': '2016-01-01'}}]
+        instance.segment_metadata.return_value = SEGMENT_METADATA
+
+        cluster = (
+            db.session
+            .query(DruidCluster)
+            .filter_by(cluster_name='test_cluster')
+            .first()
+        )
+        if cluster:
+            db.session.delete(cluster)
+        db.session.commit()
+
+        cluster = DruidCluster(
+            cluster_name='test_cluster',
+            coordinator_host='localhost',
+            coordinator_port=7979,
+            broker_host='localhost',
+            broker_port=7980,
+            metadata_last_refreshed=datetime.now())
+
+        db.session.add(cluster)
+        cluster.get_datasources = Mock(return_value=['test_datasource'])
+        cluster.get_druid_version = Mock(return_value='0.9.1')
+
+        cluster.refresh_datasources()
+        datasource_id = cluster.datasources[0].id
+        db.session.commit()
+
+        view_menu_name = cluster.datasources[0].get_perm()
+        view_menu = sm.find_view_menu(view_menu_name)
+        permission = sm.find_permission("datasource_access")
+
+        pv = sm.get_session.query(sm.permissionview_model).filter_by(
+            permission=permission, view_menu=view_menu).first()
+        assert pv is not None
+
+
 
 if __name__ == '__main__':
     unittest.main()

tests/sqllab_tests.py

@@ -59,7 +59,9 @@ class SqlLabTests(SupersetTestCase):
         main_db_permission_view = (
             db.session.query(ab_models.PermissionView)
             .join(ab_models.ViewMenu)
+            .join(ab_models.Permission)
             .filter(ab_models.ViewMenu.name == '[main].(id:1)')
+            .filter(ab_models.Permission.name == 'database_access')
             .first()
         )
         astronaut = sm.add_role("Astronaut")