From b6b41e70b297f3d4563b9db4e0e394461890a733 Mon Sep 17 00:00:00 2001 From: Daniel Vaz Gaspar Date: Wed, 6 Jul 2022 09:54:18 +0100 Subject: [PATCH] feat: add hooks on set_perm for new data permissions (#20600) * feat: add hooks on set_perm for new data permissions * fix lint --- superset/connectors/sqla/utils.py | 2 +- superset/security/manager.py | 58 ++++++++++++++++++++++- tests/integration_tests/security_tests.py | 39 +++++++++++---- 3 files changed, 88 insertions(+), 11 deletions(-) diff --git a/superset/connectors/sqla/utils.py b/superset/connectors/sqla/utils.py index 69a983156eafb..e58cf797a9931 100644 --- a/superset/connectors/sqla/utils.py +++ b/superset/connectors/sqla/utils.py @@ -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(()) diff --git a/superset/security/manager.py b/superset/security/manager.py index 083bd5b17e6dc..78c83f72ebe27 100644 --- a/superset/security/manager.py +++ b/superset/security/manager.py @@ -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, @@ -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: """ @@ -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 = ( @@ -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 diff --git a/tests/integration_tests/security_tests.py b/tests/integration_tests/security_tests.py index 757e28511660b..81781d16c5d35 100644 --- a/tests/integration_tests/security_tests.py +++ b/tests/integration_tests/security_tests.py @@ -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 @@ -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", @@ -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