Basic implementation of Snowflake pass-through authentication in Superset

docker-superset_superset_config.py

# --------------------------------------
#       Snowflake OAuth
# --------------------------------------

SNOWFLAKE_ACCOUNT = os.getenv('SNOWFLAKE_ACCOUNT') or None

if SNOWFLAKE_ACCOUNT:
    SNOWFLAKE_SECURITY_INTEGRATION_CLIENT_ID = os.getenv('SNOWFLAKE_SECURITY_INTEGRATION_CLIENT_ID')
    SNOWFLAKE_SECURITY_INTEGRATION_CLIENT_SECRET = os.getenv('SNOWFLAKE_SECURITY_INTEGRATION_CLIENT_SECRET')
    SNOWFLAKE_OAUTH_AUTHORIZATION_ENDPOINT = f'https://{SNOWFLAKE_ACCOUNT}.snowflakecomputing.com/oauth/authorize'
    SNOWFLAKE_OAUTH_TOKEN_ENDPOINT = f'https://{SNOWFLAKE_ACCOUNT}.snowflakecomputing.com/oauth/token-request'

    # ------ Models
    import datetime
    import sys
    import json

    from flask import Blueprint, render_template, g, redirect

    from flask_appbuilder import expose, Model
    from flask_appbuilder.api import BaseApi, safe, protect

    from superset.typing import FlaskResponse


    def now():
        import datetime
        return datetime.datetime.utcnow().strftime('%Y-%m-%d %H:%M:%S UTC')


    class AccessToken(Model):
        from sqlalchemy import Column, Integer, ForeignKey, DateTime, Text

        __tablename__ = "access_token"
        type = "table"

        id = Column(Integer, primary_key=True)
        access_token = Column(Text, unique=True, nullable=False)
        refresh_token = Column(Text, unique=True, nullable=False)
        expires_at = Column(DateTime, nullable=False)
        refresh_expires_at = Column(DateTime, nullable=False)
        created_at = Column(DateTime, nullable=False, default=now)
        updated_at = Column(DateTime, nullable=True, default=now)

        user_id = Column(Integer, ForeignKey('ab_user.id'), nullable=False, unique=True)

        def __repr__(self):
            return repr(self.id)


    # ---- DAO class
    from superset.dao.base import BaseDAO
    from superset.extensions import db


    class AccessTokenDAO(BaseDAO):
        model_cls = AccessToken

        @classmethod
        def find_by_user_id(cls, user_id: int) -> Optional[Model]:
            return (
                db.session.query(cls.model_cls).filter_by(user_id=user_id).one_or_none()
            )


    # --- REST API
    class ExternalAuthorizationAPI(BaseApi):
        resource_name = 'external_auth'

        class_permission_name = "external_auth_api"
        method_permission_name = {
            "get_redirect": "access",
        }

        @staticmethod
        def _call_token_request_api(data):
            import base64
            from flask import current_app
            import requests
            import logging

            logging.getLogger("requests").setLevel(logging.DEBUG)

            base_encoding = base64.b64encode('{}:{}'.format(
                current_app.config["SNOWFLAKE_SECURITY_INTEGRATION_CLIENT_ID"],
                current_app.config["SNOWFLAKE_SECURITY_INTEGRATION_CLIENT_SECRET"]).encode()).decode()

            headers = {
                'Authorization': 'Basic {}'.format(base_encoding),
                'Content-type': 'application/x-www-form-urlencoded'
            }

            response_token = requests.post(current_app.config['SNOWFLAKE_OAUTH_TOKEN_ENDPOINT'], data=data,
                                           headers=headers)

            response_token.raise_for_status()

            return response_token.json()

        def refresh_token(self, access_token: AccessToken):

            data = {
                'grant_type': 'refresh_token',
                'refresh_token': access_token.refresh_token
            }

            response = self._call_token_request_api(data)

            expiry_date = datetime.datetime.utcnow() + datetime.timedelta(seconds=response['expires_in'])

            try:
                model = AccessTokenDAO.update(access_token, {
                    'access_token': response['access_token'],
                    'expires_at': expiry_date
                })
                return model
            except:
                logging.exception('Could not refresh token for user %s', access_token.user_id,
                                  exc_info=sys.exc_info())
                raise

        def _handle_sf_authorization_code(self, code):
            response = self._call_token_request_api({
                'grant_type': 'authorization_code',
                'code': code,
                'redirect_uri': f'{os.getenv("SUPERSET_DOMAIN")}/api/v1/external_auth/redirect',
            })

            return response

        @expose('/redirect', methods=["GET"])
        @protect(allow_browser_login=True)
        @safe
        def redirect(self) -> FlaskResponse:
            from superset import appbuilder
            from flask import request, flash

            if not g.user or not g.user.get_id():
                return redirect(appbuilder.get_url_for_login)

            if not request.args.get('code'):
                return json.dumps({'error': 'empty authorization code'}), 400

            response = self._handle_sf_authorization_code(request.args['code'])

            utc_now = datetime.datetime.utcnow()

            expiry_date = utc_now + datetime.timedelta(seconds=response['expires_in'])
            refresh_expiry_date = utc_now + datetime.timedelta(seconds=response['refresh_token_expires_in'])

            if response['username'] != g.user.username:
                flash('Error! User authenticated to Snowflake is not the same as currently logged in', 'danger')
                return redirect('/snowflake_oauth')

            ack = self.__create_token(access_token=response['access_token'],
                                      refresh_token=response['refresh_token'],
                                      expires_at=expiry_date,
                                      refresh_expires_at=refresh_expiry_date,
                                      user=g.user)
            if ack:
                flash('New OAuth token has been added', 'info')
                return redirect("/superset/welcome/")
            else:
                flash('Something went wrong! Couldn\'t get an OAuth token', 'danger')
                return redirect('/snowflake_oauth')

        @staticmethod
        def __create_token(
                access_token: str,
                expires_at: datetime.datetime,
                refresh_expires_at: datetime.datetime,
                refresh_token: str,
                user):

            import sys

            try:
                access_token_instance = AccessTokenDAO.find_by_user_id(user.get_id())

                if access_token_instance:
                    AccessTokenDAO.update(access_token_instance,
                                          {
                                              'access_token': access_token,
                                              'expires_at': expires_at,
                                              'refresh_token': refresh_token,
                                              'refresh_expires_at': refresh_expires_at,
                                              'updated_at': now()

                                          })

                    logging.info('Update token for user %s', user.username)
                else:
                    access_token_instance = AccessTokenDAO.create({
                        'access_token': access_token,
                        'refresh_token': refresh_token,
                        'expires_at': expires_at,
                        'refresh_expires_at': refresh_expires_at,
                        'user_id': user.get_id(),
                    })
                    logging.info('New token added for user %s', user.username)

                return access_token_instance
            except Exception:
                logging.exception('Could not add/update token for user %s', user.username, exc_info=sys.exc_info())
                return False


    # ----------------------------
    #   Custom app initializer
    # ----------------------------

    from superset.app import SupersetAppInitializer


    class MySupersetAppInitializer(SupersetAppInitializer):

        def init_views(self) -> None:
            import logging
            from superset.extensions import appbuilder
            from flask_babel import gettext as __

            appbuilder.add_api(ExternalAuthorizationAPI)

            appbuilder.add_link(
                "Snowflake OAuth",
                label="Snowflake OAuth",
                href="/snowflake_oauth",
                icon="fa-folder-open-o",
                category="Data",
                category_label=__("Data"),
            )

            super().init_views()


    APP_INITIALIZER = MySupersetAppInitializer

    # --------------------------
    #    New flask blueprint
    # ---------------------------
    snowflake_oauth_bp = Blueprint(
        'snowflake_oauth',
        __name__,
        template_folder='templates')


    @snowflake_oauth_bp.route('/snowflake_oauth')
    def snowflake_oauth():
        from superset import appbuilder
        from urllib.parse import urlencode
        from flask import current_app
        from flask_login import current_user

        if not current_user or not current_user.get_id():
            return redirect(appbuilder.get_url_for_login)

        link_params = {
            'client_id': current_app.config['SNOWFLAKE_SECURITY_INTEGRATION_CLIENT_ID'],
            'redirect_uri': f'{os.getenv("SUPERSET_DOMAIN")}/api/v1/external_auth/redirect',
            'response_type': 'code',
            'scope': 'refresh_token'
        }
        query = urlencode(link_params)
        auth_link = f"{current_app.config['SNOWFLAKE_OAUTH_AUTHORIZATION_ENDPOINT']}?{query}"

        return render_template('snowflake_oauth.html',
                               appbuilder=appbuilder,
                               title='Snowflake OAuth',
                               snowflake_auth={
                                   'account': SNOWFLAKE_ACCOUNT,
                                   'auth_link': auth_link
                               }
                               )


    BLUEPRINTS = [snowflake_oauth_bp]

docker-superset_templates_snowflake_oauth.html

{% extends "superset/basic.html" %}

{% block title %}
  {{ title }}
{% endblock %}

{% block body %}
<div id="app" class="container" xmlns="http://www.w3.org/1999/html">
    <div class="container">
      <div class="row">
        {% include "superset/flash_wrapper.html" %}
        {% if snowflake_auth %}
        <div class="container">
          <div class="mainbox col-md-6 col-md-offset-3 col-sm-8 col-sm-offset-2">
            <div class="panel panel-primary">
              <div class="panel-heading">
                <div class="panel-title">Snowflake OAuth</div>
              </div>
              <div class="panel-body" style="padding-top:30px">
                <div class="help-block">Click to login to Snowflake account <em>{{ snowflake_auth.account }}</em>
                </div>
                <div class="help-block">This connection will be valid for 8 hours.</div>
                <br>
                <a href="{{ snowflake_auth.auth_link }}" class="btn btn-primary">Connect</a>
              </div>
            </div>
          </div>
        </div>
        {% endif %}
      </div>
    </div>
</div>
{% endblock %}

superset.sql

------------------------------------------------------
-- Integration for OAuth in Superset prod
-----------------------------------------------------

-- With ACCOUNTADMIN role

CREATE SECURITY INTEGRATION superset_prod
    TYPE = OAUTH
    ENABLED = TRUE
    OAUTH_CLIENT = CUSTOM
    OAUTH_CLIENT_TYPE='CONFIDENTIAL'
    oauth_issue_refresh_tokens=True
    OAUTH_REFRESH_TOKEN_VALIDITY=28800 -- 8hs
    OAUTH_REDIRECT_URI = 'https://superset.yourdomain.com/api/v1/external_auth/redirect'
    OAUTH_ALLOW_NON_TLS_REDIRECT_URI=False
    COMMENT = 'OAUTH with Snowflake from production superset';

-- Get client ID and secret by running
SELECT SYSTEM$SHOW_OAUTH_CLIENT_SECRETS('SUPERSET_PROD');

https://docs.snowflake.com/en/user-guide/oauth-custom.html

could you tell me how to use this code for snowflake connection in sperset?

@myu65 do you want to connect via OAuth2?

yes.
I just find PR for oauth for db.

apache/superset#27631

@myu65 that PR is just the foundation, so far it has only been implemented in GSheets and Trino for Apache Superset.

For Snowflake there are 2 options:

1. The company I work for (https://preset.io/) offers a hosted Superset that includes OAuth2 support for BigQuery, Databricks, and Snowflake. We offer a free tier (no credit card, free forever) that you can use to test, and might be enough for your needs.
2. You can implement the logic in Apache Superset, based on SIP-85 above. If you do that, please tag me as a reviewer in the PR.

Thank you for your answer.

I’ll try the second option. This is my first trial to develop apache superset code.