Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Ezaf-3082 - customize Alpha role to be able create database connection #7

Merged
merged 1 commit into from
Sep 6, 2023
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 31 additions & 2 deletions superset/header_auth_security_manager.py
Original file line number Diff line number Diff line change
@@ -1,14 +1,14 @@
import jwt
from flask import Request, flash, g, redirect, request, session
from flask_appbuilder._compat import as_unicode
from flask_appbuilder.security.manager import AUTH_REMOTE_USER
from flask_appbuilder.security.views import AuthView
from flask_appbuilder.utils.base import get_safe_redirect
from flask_appbuilder.views import expose
from flask_login import login_user, logout_user
from superset.security.manager import SupersetSecurityManager
from werkzeug.sansio.utils import get_current_url
from werkzeug.wrappers import Response as WerkzeugResponse
from flask_appbuilder.security.manager import AUTH_REMOTE_USER
from superset.security.manager import SupersetSecurityManager


# Custom Authenticator based on principals in request headers
Expand Down Expand Up @@ -59,6 +59,33 @@ def login(self) -> WerkzeugResponse:
next_url = request.args.get("next", "")
return redirect(get_safe_redirect(next_url))

def __get_or_create_custom_role(self, role_name: str):
ab_security_manager = self.appbuilder.sm

custom_alpha_role = ab_security_manager.find_role(role_name)
if custom_alpha_role:
return

alpha_role = ab_security_manager.find_role("Alpha")
if alpha_role:
alpha_permissions = alpha_role.permissions
custom_alpha_role = ab_security_manager.add_role(
role_name,
alpha_permissions
)

if custom_alpha_role is None:
raise Exception(f"Cannot create {role_name} role")

write_db_perm = ab_security_manager.find_permission_view_menu('can_write', 'Database')
if write_db_perm:
custom_alpha_role.permissions.append(write_db_perm)
ab_security_manager.get_session.commit()
else:
raise Exception("'can_write Database' permission does not exist")
else:
raise Exception("Alpha role not found")

def __get_or_create_user(self, username):
ab_security_manager = self.appbuilder.sm
user = ab_security_manager.find_user(username)
Expand All @@ -78,6 +105,8 @@ def __get_or_create_user(self, username):
else:
# The default authentication role should be defined in helm/superset/values.yaml as AUTH_USER_REGISTRATION_ROLE
role_name = ab_security_manager.auth_user_registration_role

self.__get_or_create_custom_role(role_name)

user = ab_security_manager.add_user(
username=username,
Expand Down
Loading