Snowflake SQL API ruby client

snowflake.rb

class Snowflake
    require 'digest'
    require 'jwt'
    require 'securerandom'

    class Client
        attr_reader :authorization, :expiration
        attr_accessor :account, :region, :user, :private_key_path

        def initialize(attributes={})
            @account = attributes[:account]
            @region = attributes[:region]
            @user = attributes[:user]
            @private_key_path = attributes[:private_key_path]
            @authorization = nil
            @expiration = nil
        end
        
        def make_query(attributes={})
            attributes = attributes.merge(client: self)
            Query.new(attributes)
        end

        private 

        def get_token
            private_key_file = File.read(self.private_key_path)
            private_key = OpenSSL::PKey::RSA.new(private_key_file)
            public_key = private_key.public_key
            public_key_fp = Base64.encode64(Digest::SHA256.digest(public_key.to_der)) # Public Key Fingerprint
    
            qualified_username = self.account.upcase + '.' + self.user.upcase
            issuer = qualified_username + '.' + 'SHA256:' + public_key_fp.strip
            issue_time = Time.now.to_i 
            expiration_time = issue_time + 1*60*60 # 1 hour from now
    
            payload = {
                iss: issuer,
                sub: qualified_username,
                iat: issue_time,
                exp: expiration_time
            }
            return JWT.encode payload, private_key, 'RS256', { typ: 'JWT' }
        end
    
        def authorize
            # Only get a token when there is none, or the current one expires in less than 5 minutes
            if (not @authorization) or (@expiration - Time.now < 5*60) 
                @authorization = 'Bearer ' + self.get_token
                @expiration =Time.at(JWT.decode(@authorization[7..-1],nil,false).first['exp'])
            end
        end

    end

    class Query
        attr_reader :response
        attr_accessor :statement

        def initialize(attributes={})
            @client = attributes[:client]
            @warehouse = attributes[:warehouse]
            @database = attributes[:database]
            @schema = attributes[:schema]
            @statement = attributes[:statement]
            @role = attributes[:role] 
            @response = nil   
        end

        def execute
            @client.send :authorize # Get or Refresh JWT if rquired
            endpoint = '/api/v2/statements'
            requestId = SecureRandom.uuid
            baseUrl = 'https://' + @client.account + '.' + @client.region + '.snowflakecomputing.com'
            
            url = baseUrl + endpoint
            params = {requestId: requestId, retry: 'true'}
    
            body = {
                statement: @statement,
                timeout: 60,
                database: @database,
                schema: @schema,
                warehouse: @warehouse,
                role: @role
            }.to_json
    
            connection = Faraday.new(
                url: baseUrl,
                params: params,
                headers: {
                    'Content-type'=> 'application/json',
                    'Authorization' => @client.authorization,
                    'Accept' => 'application/json',
                    'X-Snowflake-Authorization-Token-Type' => 'KEYPAIR_JWT'
                }
            )
            response = connection.post(endpoint, body)
            @response = {
                status: response.status,
                headers: response.headers,
                body: response.body
            }
        end
    end

    private_constant :Query

end

Create a client

snowflake_client = Snowflake.new(
    account: 'myorg',
    region: 'us-east-1',
    user: 'my_user',
    private_key_path: '/home/ubuntu/.ssh/rsa_key.p8'
)

Create a query

my_query = snowflake_client.Query.new(
    statement: 'SELECT * FROM User;',
    database: 'TEST_DB',
    schema: 'PUBLIC',
    warehouse: 'ANALYTICS_WH_XL',
    role: 'ANALYST_ROLE'
)

Run the query

my_query.execute

Get the results

my_query.response

https://docs.snowflake.com/en/user-guide/key-pair-auth

This code requires the JWT and Faraday Ruby Gems