Skip to content

Commit

Permalink
feat: add hooks on set_perm for new data permissions (#20600)
Browse files Browse the repository at this point in the history
* feat: add hooks on set_perm for new data permissions

* fix lint
  • Loading branch information
dpgaspar committed Jul 6, 2022
1 parent 818962c commit f38dd1d
Show file tree
Hide file tree
Showing 3 changed files with 88 additions and 11 deletions.
2 changes: 1 addition & 1 deletion superset/connectors/sqla/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -213,7 +213,7 @@ def find_cached_objects_in_session(
return iter([])
uuids = uuids or []
try:
items = set(session)
items = list(session)
except ObjectDeletedError:
logger.warning("ObjectDeletedError", exc_info=True)
return iter(())
Expand Down
58 changes: 57 additions & 1 deletion superset/security/manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -40,9 +40,11 @@
from flask_appbuilder.security.sqla.models import (
assoc_permissionview_role,
assoc_user_role,
Permission,
PermissionView,
Role,
User,
ViewMenu,
)
from flask_appbuilder.security.views import (
PermissionModelView,
Expand Down Expand Up @@ -931,7 +933,55 @@ def _is_granter_pvm( # pylint: disable=no-self-use

return pvm.permission.name in {"can_override_role_permissions", "can_approve"}

def set_perm( # pylint: disable=unused-argument
def on_permission_after_insert(
self, mapper: Mapper, connection: Connection, target: Permission
) -> None:
"""
Hook that allows for further custom operations when a new permission
is created by set_perm.
Since set_perm is executed by SQLAlchemy after_insert events, we cannot
create new permissions using a session, so any SQLAlchemy events hooked to
`Permission` will not trigger an after_insert.
:param mapper: The table mapper
:param connection: The DB-API connection
:param target: The mapped instance being persisted
"""

def on_view_menu_after_insert(
self, mapper: Mapper, connection: Connection, target: ViewMenu
) -> None:
"""
Hook that allows for further custom operations when a new ViewMenu
is created by set_perm.
Since set_perm is executed by SQLAlchemy after_insert events, we cannot
create new view_menu's using a session, so any SQLAlchemy events hooked to
`ViewMenu` will not trigger an after_insert.
:param mapper: The table mapper
:param connection: The DB-API connection
:param target: The mapped instance being persisted
"""

def on_permission_view_after_insert(
self, mapper: Mapper, connection: Connection, target: PermissionView
) -> None:
"""
Hook that allows for further custom operations when a new PermissionView
is created by set_perm.
Since set_perm is executed by SQLAlchemy after_insert events, we cannot
create new pvms using a session, so any SQLAlchemy events hooked to
`PermissionView` will not trigger an after_insert.
:param mapper: The table mapper
:param connection: The DB-API connection
:param target: The mapped instance being persisted
"""

def set_perm(
self, mapper: Mapper, connection: Connection, target: "BaseDatasource"
) -> None:
"""
Expand Down Expand Up @@ -988,12 +1038,14 @@ def set_perm( # pylint: disable=unused-argument
permission_table.insert().values(name=permission_name)
)
permission = self.find_permission(permission_name)
self.on_permission_after_insert(mapper, connection, permission)
if not view_menu:
view_menu_table = (
self.viewmenu_model.__table__ # pylint: disable=no-member
)
connection.execute(view_menu_table.insert().values(name=view_menu_name))
view_menu = self.find_view_menu(view_menu_name)
self.on_view_menu_after_insert(mapper, connection, view_menu)

if permission and view_menu:
pv = (
Expand All @@ -1010,6 +1062,10 @@ def set_perm( # pylint: disable=unused-argument
permission_id=permission.id, view_menu_id=view_menu.id
)
)
permission = self.find_permission_view_menu(
permission_name, view_menu_name
)
self.on_permission_view_after_insert(mapper, connection, permission)

def raise_for_access(
# pylint: disable=too-many-arguments,too-many-locals
Expand Down
39 changes: 30 additions & 9 deletions tests/integration_tests/security_tests.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@
import unittest
from collections import namedtuple
from unittest import mock
from unittest.mock import Mock, patch
from unittest.mock import Mock, patch, call, ANY
from typing import Any

import jwt
Expand Down Expand Up @@ -157,6 +157,9 @@ def tearDown(self):
session.commit()

def test_set_perm_sqla_table(self):
security_manager.on_view_menu_after_insert = Mock()
security_manager.on_permission_view_after_insert = Mock()

session = db.session
table = SqlaTable(
schema="tmp_schema",
Expand All @@ -172,16 +175,34 @@ def test_set_perm_sqla_table(self):
self.assertEqual(
stored_table.perm, f"[examples].[tmp_perm_table](id:{stored_table.id})"
)
self.assertIsNotNone(
security_manager.find_permission_view_menu(
"datasource_access", stored_table.perm
)

pvm_dataset = security_manager.find_permission_view_menu(
"datasource_access", stored_table.perm
)
pvm_schema = security_manager.find_permission_view_menu(
"schema_access", stored_table.schema_perm
)

self.assertIsNotNone(pvm_dataset)
self.assertEqual(stored_table.schema_perm, "[examples].[tmp_schema]")
self.assertIsNotNone(
security_manager.find_permission_view_menu(
"schema_access", stored_table.schema_perm
)
self.assertIsNotNone(pvm_schema)

# assert on permission hooks
view_menu_dataset = security_manager.find_view_menu(
f"[examples].[tmp_perm_table](id:{stored_table.id})"
)
view_menu_schema = security_manager.find_view_menu(f"[examples].[tmp_schema]")
security_manager.on_view_menu_after_insert.assert_has_calls(
[
call(ANY, ANY, view_menu_dataset),
call(ANY, ANY, view_menu_schema),
]
)
security_manager.on_permission_view_after_insert.assert_has_calls(
[
call(ANY, ANY, pvm_dataset),
call(ANY, ANY, pvm_schema),
]
)

# table name change
Expand Down

0 comments on commit f38dd1d

Please sign in to comment.