Snowflake Resource with OAuth

snowflake_oauth.py

from dagster_snowflake import SnowflakeConnection
from typing import Mapping
from dagster._annotations import public
from dagster import resource, StringSource


class SnowflakeOAuthConnection(SnowflakeConnection):

    def __init__(self, config: Mapping[str, str], log):

        self.conn_args = {
            k: config.get(k)
            for k in (
                "account",
                "user",
                "host",
                "database",
                "schema",
                "role",
                "warehouse",
                "autocommit"
            )
            if config.get(k) is not None
        }

        self.conn_args['authenticator'] = 'oauth'
        self.autocommit = self.conn_args.get("autocommit", False)
        self.log = log

    
    @public
    def get_connection(self, raw_conn: bool = True):

        self.conn_args['token'] = token_generator(...)

        return super().get_connection(raw_conn)

@resource(
    config_schema = {
         "account": StringSource,
         "user": StringSource,
         "host": StringSource,
         "database": StringSource,
         "schema": StringSource,
         "role": StringSource,
         "warehouse": StringSource,
         "autocommit": bool
    }
)
def snowflake_oauth_resource(context): 
    return SnowflakeOAuthConnection(context.resource_config, context.log)