Enable Okta Login for Superset

README.md

This Gist demonstrates how to enable Okta/OpenID Connect login for Superset (and with slight modifications, other Flask-Appbuilder apps).

All you need to do is to update superset_config.py as following and make sure your Cliend ID (OKTA_KEY) and Secret (OKTA_SECRET) are put into the env.

superset_config.py

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
import os
import logging

from flask_appbuilder.security.manager import AUTH_OAUTH

logging.getLogger('requests').setLevel(logging.CRITICAL)
logging.getLogger('werkzeug').setLevel(logging.CRITICAL)

def get_env_variable(var_name, default=None):
    """Get the environment variable or raise exception."""
    try:
        return os.environ[var_name]
    except KeyError:
        if default is not None:
            return default
        else:
            error_msg = 'The environment variable {} was missing, abort...'\
                        .format(var_name)
            raise EnvironmentError(error_msg)


PUBLIC_ROLE_LIKE_GAMMA = True

# ====== Start Okta Login ===========
AUTH_TYPE = AUTH_OAUTH
AUTH_USER_REGISTRATION = True  # allow self-registration (login creates a user)
AUTH_USER_REGISTRATION_ROLE = "Gamma"  # default is a Gamma user

# OKTA_BASE_URL must be
#    https://{yourOktaDomain}/oauth2/v1/ (okta authorization server)
# Cannot be
#    https://{yourOktaDomain}/oauth2/default/v1/ (custom authorization server)
# Otherwise you won't be able to obtain Groups info.
OKTA_BASE_URL = get_env_variable('OKTA_BASE_URL')
OKTA_KEY = get_env_variable('OKTA_KEY')
OKTA_SECRET = get_env_variable('OKTA_SECRET')
OAUTH_PROVIDERS = [{
    'name':'okta',
    'token_key': 'access_token', # Name of the token in the response of access_token_url
    'icon':'fa-circle-o',   # Icon for the provider
    'remote_app': {
        'consumer_key': OKTA_KEY,  # Client Id (Identify Superset application)
        'consumer_secret': OKTA_SECRET, # Secret for this Client Id (Identify Superset application)
        'request_token_params': {
            'scope': 'openid email profile groups'
        },
        'access_token_method': 'POST',    # HTTP Method to call access_token_url
        'base_url': OKTA_BASE_URL,
        'access_token_url': OKTA_BASE_URL + 'token',
        'authorize_url': OKTA_BASE_URL + 'authorize'
    }
}]

from superset.security import SupersetSecurityManager

logger = logging.getLogger('okta_login')

class CustomSsoSecurityManager(SupersetSecurityManager):

    def oauth_user_info(self, provider, response=None):
        if provider == 'okta':
            res = self.appbuilder.sm.oauth_remotes[provider].get('userinfo')
            if res.status != 200:
                logger.error('Failed to obtain user info: %s', res.data)
                return
            me = res.data
            logger.debug(" user_data: %s", me)
            prefix = 'Superset '
            groups = [
                x.replace(prefix, '').strip() for x in me['groups']
                if x.startswith(prefix)
            ]
            return {
                'username' : me['preferred_username'],
                'name' : me['name'],
                'email' : me['email'],
                'first_name': me['given_name'],
                'last_name': me['family_name'],
                'roles': groups,
            }

    def auth_user_oauth(self, userinfo):
        user = super(CustomSsoSecurityManager, self).auth_user_oauth(userinfo)
        roles = [self.find_role(x) for x in userinfo['roles']]
        roles = [x for x in roles if x is not None]
        user.roles = roles
        logger.debug(' Update <User: %s> role to %s', user.username, roles)
        self.update_user(user)  # update user roles
        return user

CUSTOM_SECURITY_MANAGER = CustomSsoSecurityManager


# ====== End Okta Login ============

CSRF_ENABLED = True

POSTGRES_USER = get_env_variable('POSTGRES_USER')
POSTGRES_PASSWORD = get_env_variable('POSTGRES_PASSWORD')
POSTGRES_HOST = get_env_variable('POSTGRES_HOST')
POSTGRES_PORT = get_env_variable('POSTGRES_PORT')
POSTGRES_DB = get_env_variable('POSTGRES_DB')

# The SQLAlchemy connection string.
SQLALCHEMY_DATABASE_URI = 'postgresql://%s:%s@%s:%s/%s' % (POSTGRES_USER,
                                                           POSTGRES_PASSWORD,
                                                           POSTGRES_HOST,
                                                           POSTGRES_PORT,
                                                           POSTGRES_DB)

REDIS_HOST = get_env_variable('REDIS_HOST')
REDIS_PORT = get_env_variable('REDIS_PORT')

class CeleryConfig(object):
    BROKER_URL = 'redis://%s:%s/0' % (REDIS_HOST, REDIS_PORT)
    CELERY_IMPORTS = ('superset.sql_lab', )
    CELERY_RESULT_BACKEND = 'redis://%s:%s/1' % (REDIS_HOST, REDIS_PORT)
    CELERY_ANNOTATIONS = {'tasks.add': {'rate_limit': '10/s'}}
    CELERY_TASK_PROTOCOL = 1


CELERY_CONFIG = CeleryConfig