Skip to content

Instantly share code, notes, and snippets.

@stuzero

stuzero/snowflake.rb

Created February 17, 2023 01:52
Show Gist options
  • Save stuzero/90aae045eb85586da171e2fd8b932252 to your computer and use it in GitHub Desktop.
Save stuzero/90aae045eb85586da171e2fd8b932252 to your computer and use it in GitHub Desktop.
Download ZIP
Snowflake SQL API ruby client
Raw
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
@stuzero
Copy link
Author

stuzero commented Feb 17, 2023

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

@stuzero
Copy link
Author

stuzero commented Feb 17, 2023

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

@dxi-stuart
Copy link

This code requires the JWT and Faraday Ruby Gems

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment