From eb9cd2a2a567e5217dd0a22af8b1aad6b2fc19bc Mon Sep 17 00:00:00 2001 From: Sandeep Patel <33354423+suicide11@users.noreply.github.com> Date: Wed, 11 Oct 2023 22:59:52 +0530 Subject: [PATCH] refactor: Issue #25040; Refactored sync_role_definition function in order to reduce number of query. (#25340) --- superset/security/manager.py | 35 ++++++++++++++++++++++++++--------- 1 file changed, 26 insertions(+), 9 deletions(-) diff --git a/superset/security/manager.py b/superset/security/manager.py index dc5e76e18b..2225a0fa3a 100644 --- a/superset/security/manager.py +++ b/superset/security/manager.py @@ -48,7 +48,7 @@ from flask_login import AnonymousUserMixin, LoginManager from jwt.api_jwt import _jwt_global_obj from sqlalchemy import and_, inspect, or_ from sqlalchemy.engine.base import Connection -from sqlalchemy.orm import Session +from sqlalchemy.orm import eagerload, Session from sqlalchemy.orm.mapper import Mapper from sqlalchemy.orm.query import Query as SqlaQuery @@ -732,7 +732,7 @@ class SupersetSecurityManager( # pylint: disable=too-many-public-methods logger.info("Fetching a set of all perms to lookup which ones are missing") all_pvs = set() - for pv in self.get_session.query(self.permissionview_model).all(): + for pv in self._get_all_pvms(): if pv.permission and pv.view_menu: all_pvs.add((pv.permission.name, pv.view_menu.name)) @@ -780,11 +780,13 @@ class SupersetSecurityManager( # pylint: disable=too-many-public-methods self.create_custom_permissions() + pvms = self._get_all_pvms() + # Creating default roles - self.set_role("Admin", self._is_admin_pvm) - self.set_role("Alpha", self._is_alpha_pvm) - self.set_role("Gamma", self._is_gamma_pvm) - self.set_role("sql_lab", self._is_sql_lab_pvm) + self.set_role("Admin", self._is_admin_pvm, pvms) + self.set_role("Alpha", self._is_alpha_pvm, pvms) + self.set_role("Gamma", self._is_gamma_pvm, pvms) + self.set_role("sql_lab", self._is_sql_lab_pvm, pvms) # Configure public role if current_app.config["PUBLIC_ROLE_LIKE"]: @@ -800,6 +802,20 @@ class SupersetSecurityManager( # pylint: disable=too-many-public-methods self.get_session.commit() self.clean_perms() + def _get_all_pvms(self) -> list[PermissionView]: + """ + Gets list of all PVM + """ + pvms = ( + self.get_session.query(self.permissionview_model) + .options( + eagerload(self.permissionview_model.permission), + eagerload(self.permissionview_model.view_menu), + ) + .all() + ) + return [p for p in pvms if p.permission and p.view_menu] + def _get_pvms_from_builtin_role(self, role_name: str) -> list[PermissionView]: """ Gets a list of model PermissionView permissions inferred from a builtin role @@ -860,7 +876,10 @@ class SupersetSecurityManager( # pylint: disable=too-many-public-methods self.get_session.commit() def set_role( - self, role_name: str, pvm_check: Callable[[PermissionView], bool] + self, + role_name: str, + pvm_check: Callable[[PermissionView], bool], + pvms: list[PermissionView], ) -> None: """ Set the FAB permission/views for the role. @@ -870,8 +889,6 @@ class SupersetSecurityManager( # pylint: disable=too-many-public-methods """ logger.info("Syncing %s perms", role_name) - pvms = self.get_session.query(PermissionView).all() - pvms = [p for p in pvms if p.permission and p.view_menu] role = self.add_role(role_name) role_pvms = [ permission_view for permission_view in pvms if pvm_check(permission_view)