Snowflake OAuth for Streamlit

One of the most common questions I get when talking with customers building Streamlit apps with Snowflake is how Streamlit can connect to Snowflake as the visitor of the Streamlit app. In a previous post, I talked about how to do this with username/password authentication via a component that will prompt for the credentials from the visitor and then make (and cache) the connection to Snowflake.

However, a really common request is to use OAuth to do the authentication. I did some work to investigate this, and the result is a new Streamlit component, st_snowauth. (Huge shout-out to Mike Mitrowski for helping me through the details here — I couldn’t have done it without you!)

Now, you can use OAuth to “gate” your Streamlit app until the user logs into Snowflake (via username/password or SSO integration) with code as simple as:

import streamlit as st
from st_snowauth import snowauth_session

st.markdown("## This (and above) is always seen")
session = snowauth_session()
st.markdown("## This (and below) is only seen after authentication")

st_snowauth will return a Snowflake Snowpark session for the authenticated user.

NOTE: Snowflake OAuth does not allow for the user to log in with the default role of ACCOUNTADMIN, SECURITYADMIN, or ORGADMIN. You can further configure Snowflake OAuth to prevent other roles.

st_snowauth is designed to work with Snowflake OAuth with Custom Clients.

Some Details

The basic flow that I have implemented is as follows:

• st_snowauth checks if a Snowpark session is cached in session state
• If it is, then
  - st_snowauth validates that the session is still open
  - st_snowauth returns the cached session and the Streamlit app can continue
• If it is not, then
  - st_snowauth uses the authorization_endpoint to create a link that will direct the user to the Snowflake login page. It will send the redirect_uri, client_id, and state, and request a code in response
  - the user logs into Snowflake, either via username/password or SSO integration,and Snowflake will redirect back to the supplied redirect_uri, returning the code and state as query parameters
  - st_snowauth will collect the code from the query parameters (using st.experimental_get_query_params()), and then use the token_endpoint to exchange the code for a security token.
  - st_snowauth will then use the token to create a Snowpark session, cache it in session state, and return it and the Streamlit app can continue

Snowflake OAuth Configuration

Follow the instructions in the Snowflake documentation to set up a custom client. You will need to specify a redirect_uri, and you should use the URL to the main page of your Streamlit app. For example, if you are running on your laptop (perhaps, during development), you would use http://localhost:8501.

An example command is:

CREATE OR REPLACE SECURITY INTEGRATION oauth_st
  TYPE=OAUTH
  ENABLED=TRUE
  OAUTH_ALLOW_NON_TLS_REDIRECT_URI = TRUE
  OAUTH_CLIENT = CUSTOM
  OAUTH_CLIENT_TYPE='PUBLIC'
  OAUTH_REDIRECT_URI='http://localhost:8501'
  OAUTH_ISSUE_REFRESH_TOKENS = TRUE
  OAUTH_REFRESH_TOKEN_VALIDITY = 86400
;

Once you have created the security integration, you will need to get the client ID and secret. You can retrieve the client ID with the following SQL:

DESCRIBE INTEGRATION oauth_st;

You can retrieve the client secret with the following SQL:

SELECT SYSTEM$SHOW_OAUTH_CLIENT_SECRETS( ‘OAUTH_ST’ );

Setup

The set of configuration parameters we need for st_snowauth are as follows. You can get them from your Snowflake security integration:

• account — the Snowflake account ID
• authorization_endpoint — the URL to use to get an authorization code
• token_endpoint — the URL to use to trade an authorization code for a token
• redirect_uri — the URL that is configured in the security integration as the redirect URL (it should be the URL of the Streamlit app itself)
• client_id — the client ID, as configured in the security integration (see above)
• client_secret — the client secret for the client ID, as returned from SHOW_OAUTH_CLIENT_SECRETS() (see above)

To configure st_snowauth we need 2 things: a set of configuration parameters (as discussed above), and a prompt to display for the login link.

Specifying the Configuration

There are 3 ways to pass the configuration parameters in. First, you can provide a dictionary that contains the configuration parameters:

config_dict = {‘account’: ….}
session = snowauth_session(config_dict)

Second, you can store the parameters in the secrets file (.streamlit/secrets.toml):

[myoauth]
account = "<SNOWFLAKE ACCOUNT ID>"
authorization_endpoint = "<SNOWFLAKE OAUTH AUTH ENDPOINT usually ending in /oauth/authorize>"
token_endpoint = "<SNOWFLAKE OAUTH TOKEN ENDPOINT usually ending in /oauth/token-request>"
redirect_uri = "<REDIRECT URI - this Streamlit's location>"
client_id = "<OAUTH CLIENT ID>"
client_secret = "<OAUTH CLIENT SECRET>"

Then you pass in the name of the block of parameters in st.secrets:

session = snowauth_session(‘myoauth’)

Lastly, by default, st_snowauth will look in st.secrets for the block for parameters named snowauth. So, if you store the configuration in the secrets file under the label “snowauth”, like this:

[snowauth]
account = "<SNOWFLAKE ACCOUNT ID>"
authorization_endpoint = "<SNOWFLAKE OAUTH AUTH ENDPOINT usually ending in /oauth/authorize>"
token_endpoint = "<SNOWFLAKE OAUTH TOKEN ENDPOINT usually ending in /oauth/token-request>"
redirect_uri = "<REDIRECT URI - this Streamlit's location>"
client_id = "<OAUTH CLIENT ID>"
client_secret = "<OAUTH CLIENT SECRET>"

Then you can call st_snowauth with no parameters:

session = snowauth_session()

Specifying the Link Text

The other parameter you can pass into st_snowauth is the text to use in the link. By default, the text is Login to Snowflake, but you can customize this:

session = snowauth_session(label=’Click to login to Snowflake’)

Multipage Apps

st_snowauth works with multipage apps. You need to include the “gate” at the top of each page. If the user has not logged in, they will be presented with the login link. Once they complete the Snowflake login, they will be redirected to the Streamlit app’s main page.

Query Parameters

st_snowauth does work with query parameters. If you navigate to the Streamlit and have query parameters set, st_snowauth will save the parameters that were passed in (leveraging st.experimental_get_query_params()), and when Snowflake redirects back to Streamlit (upon successful login), st_snowauth will re-add the query parameters back to the URL (leveraging st.experimental_set_queyr_params()).