How to connect SQLAlchemy to Snowflake database using OAuth2?

Question

I need to connect to Snowflake using SQLAlchemy but the trick is, I need to authenticate using OAuth2. Snowflake documentation only describes connecting using username and password and this cannot be used in the solution I'm building. I can authenticate using Snowflake's python connector but I see no simple path how to glue it with SQLAlchemy. I'd like to know if there is a ready solution before I write a custom interface for this.

Answer 1

Use snowflake.connector.connect to create a PEP-249 Connection to the database - see documentation. Then use param creator of create_engine (docs) - it takes a callable that returns PEP-249 Connection. If you use it then URL param is ignored.

Example code:

def get_connection():
    return snowflake.connector.connect(
        user="<username>",
        host="<hostname>",
        account="<account_identifier>",
        authenticator="oauth",
        token="<oauth_access_token>",
        warehouse="test_warehouse",
        database="test_db",
        schema="test_schema"
    )
engine = create_engine("snowflake://not@used/db", creator=get_connection)

Answer 2

I got this working but just adding more params in the connection URL:

from sqlalchemy.engine import create_engine
import urllib.parse

connection_url = f"snowflake://{user}:@{account}/{database}/{schema}?warehouse={warehouse}&authenticator=oauth&token={urllib.parse.quote(access_token)}"
engine = create_engine(connection_url)

with engine.begin() as connection:
    print(connection.execute('select count(*) from lineitem').fetchone())

If you don't want to be constructing the URL on your own, you can use snowflake.sqlalchemy.URL like this:

from snowflake.sqlalchemy import URL
connection_url = URL(
    user=user,
    authenticator="oauth",
    token=access_token,
    host=host,
    account=account,
    warehouse=warehouse,
    database=database,
    schema=schema
)