Updating permission when refreshing druid datasource

superset/models/helpers.py

@@ -10,7 +10,7 @@
 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):
     def override(self, obj):
@@ -116,12 +116,48 @@ def __init__(  # noqa
         self.status = status
         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(
             link_table.update()
             .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)

tests/druid_tests.py

@@ -276,6 +276,49 @@ def test_filter_druid_datasource(self):
         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 @@ def test_sql_json_has_access(self):
         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")