rb_snowflake_client 0.1.2 → 0.2.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 4aa627645364f92fc0e0d9ef949519a7649088e43de1e9ca7480ba284a9efba5
-  data.tar.gz: e18d9a6a1f297cbbb5a10aafcac9dce440c43e0ca430f0c6310c31cea28b4507
+  metadata.gz: 2a93f0b1e503e2d219db8b1d1133395ca484b45aaa1f91ceef013a50a116b973
+  data.tar.gz: c292c5b0d25a921f8deba6882fe288ed084aa07919fa0d056899d20b921612fc
 SHA512:
-  metadata.gz: 80ac86515ff08d723c5eff134ac090b60c079695da9c68cf39b06f8db5d1661f700b78930f5da2c19b3f4686563d4dbcda47df70f9ad35200af9e7582b09ea8b
-  data.tar.gz: d4c10ff8bba434ccd0ee380bbaaf9d14bfc7aed22543c34397fd6e1f83611daab65efeacf8caf241a98fd5e2a0f3591592b47f7c76de4997cfe2ea2fb951a5b8
+  metadata.gz: 28428d4cf4bd70a11fefad327a2ab14deff17e95bf28b999a6b84d67756b35d3207c8224d6a028b9998ca85b62a7226f4df2beddea50f30e9a91350d7df73b4b
+  data.tar.gz: ad55068885a29f1cae0181e52d66764651cb4f7d541e46680bd15644ed8c17f10dda8cb9356ae4948cdf9b249e0c726abc540e26b03e1b76f9f2c538a8251cce

data/Gemfile.lock

@@ -1,7 +1,7 @@
 PATH
   remote: .
   specs:
-    rb_snowflake_client (0.1.2)
+    rb_snowflake_client (0.2.0)
       concurrent-ruby (>= 1.2)
       connection_pool (>= 2.4)
       dotenv (>= 2.8)

data/README.md

@@ -13,6 +13,8 @@ This library is implemented in ruby and while it leverages some libraries that h
 
 # Usage
 
+## Create a client
+
 Add to your Gemfile or use `gem install rb-snowflake-client`
 ```ruby
   gem "rb-snowflake-client"
@@ -27,7 +29,7 @@ require "rb_snowflake_client"
 # see: https://github.com/rinsed-org/pure-ruby-snowflake-client/blob/master/lib/ruby_snowflake/client.rb#L43
 client = RubySnowflake::Client.new(
   "https://yourinstance.region.snowflakecomputing.com", # insert your URL here
-  File.read("secrets/my_key.pem"),                      # path to your private key
+  File.read("secrets/my_key.pem"),                      # your private key in PEM format (scroll down for instructions)
   "snowflake-organization",                             # your account name (doesn't match your URL)
   "snowflake-account",                                  # typically your subdomain
   "snowflake-user",                                     # Your snowflake user
@@ -41,21 +43,22 @@ client = RubySnowflake::Client.new(
 RubySnowflake::Client.from_env
 ```
 Available ENV variables (see below in the config section for details)
-`SNOWFLAKE_URI`
-`SNOWFLAKE_PRIVATE_KEY_PATH`
-   or (use either the key, or the path, key takes precedence if both are provided)
-`SNOWFLAKE_PRIVATE_KEY`
-`SNOWFLAKE_ORGANIZATION`
-`SNOWFLAKE_ACCOUNT`
-`SNOWFLAKE_USER`
-`SNOWFLAKE_DEFAULT_WAREHOUSE`
-`SNOWFLAKE_DEFAULT_DATABASE`
-`SNOWFLAKE_JWT_TOKEN_TTL`
-`SNOWFLAKE_CONNECTION_TIMEOUT`
-`SNOWFLAKE_MAX_CONNECTIONS`
-`SNOWFLAKE_MAX_THREADS_PER_QUERY`
-`SNOWFLAKE_THREAD_SCALE_FACTOR`
-`SNOWFLAKE_HTTP_RETRIES`
+- `SNOWFLAKE_URI`
+- `SNOWFLAKE_PRIVATE_KEY_PATH` or `SNOWFLAKE_PRIVATE_KEY`
+  - Use either the key or the path. Key takes precedence if both are provided.
+- `SNOWFLAKE_ORGANIZATION`
+- `SNOWFLAKE_ACCOUNT`
+- `SNOWFLAKE_USER`
+- `SNOWFLAKE_DEFAULT_WAREHOUSE`
+- `SNOWFLAKE_DEFAULT_DATABASE`
+- `SNOWFLAKE_JWT_TOKEN_TTL`
+- `SNOWFLAKE_CONNECTION_TIMEOUT`
+- `SNOWFLAKE_MAX_CONNECTIONS`
+- `SNOWFLAKE_MAX_THREADS_PER_QUERY`
+- `SNOWFLAKE_THREAD_SCALE_FACTOR`
+- `SNOWFLAKE_HTTP_RETRIES`
+
+## Make queries
 
 Once you have a client, make queries
 ```ruby
@@ -68,44 +71,58 @@ result.each do |row|
   puts row["name"] # or case insensitive strings
   puts row.to_h    # and can produce a hash with keys/values
 end
+```
+
+## Stream results
 
-# You can also stream results and not hold them all in memory.
-# The client will prefetch the next data partition only. If you
-# have some IO in your processing there should usually be data
-# available for you.
+You can also stream results and not hold them all in memory. The client will prefetch the next data partition only. If you have some IO in your processing there should usually be data available for you.
+
+```ruby
 result = client.query("SELECT * FROM HUGETABLE", streaming: true)
 result.each do |row|
   puts row
 end
+```
 
+## Switching databases
 
-# You can also overwrite the database specified in the initializer, and run
-# your query with a different context.
+You can also overwrite the database specified in the initializer, and run your query with a different context.
+
+```ruby
 result = client.query("SELECT * FROM SECRET_TABLE", database: "OTHER_DB")
 result.each do |row|
   puts row
 end
 ```
 
+## Switching warehouses
+
+Clients are not warehouse specific, you can override the default warehouse per query
+
+```ruby
+client.query("SELECT * FROM BIGTABLE", warehouse: "FAST_WH")
+```
+
 # Configuration Options
 
-The client supports the following configuration options, each with their own getter/setter except connection pool options which must be set at construction. Additionally, all except logger can be configured with environment variables (see above, but the pattern is like: "SNOWFLAKE_HTTP_RETRIES".
+The client supports the following configuration options, each with their own getter/setter except connection pool options which must be set at construction. Additionally, all except logger can be configured with environment variables (see above, but the pattern is like: "SNOWFLAKE_HTTP_RETRIES". Configuration options can only be set on initiialization through `new` or `from_env`.
 
-`logger` - takes any ruby logger (by default it's a std lib Logger.new(STDOUT), set at DEBUG level. Not available as an ENV variable config option
-`log_level` - takes a log level, type is dependent on logger, for the default ruby Logger, use a level like `Logger::WARN`. Not available as an ENV variable config option.
-`jwt_token_ttl` - The time to live set on JWT token in seconds, defaults to 3540 (59 minutes, the longest Snowflake supports is 60)
-`connection_timeout` - The amount of time in seconds that the client's connection pool will wait before erroring in handing out a valid connection, defaults to 60 seconds
-`max_connections` - The maximum number of http connections to hold open in the connection pool. If you use the client in a threaded context, you may need to increase this to be threads * client.max_threads_per_query, defaults to 16. Can only be set on initialization.
-`max_threads_per_query` - The maximum number of threads the client should use to retreive data, per query, defaults to 8. If you want the client to act in a single threaded way, set this to 1
-`thread_scale_factor` - When downloading a result set into memory, thread count is calculated by dividing a query's partition count by this number. For details on implementation see the code in `client.rb`.
-`http_retries` - By default the client will retry common typically transient errors (http responses) twice, you can change the number of retries with this.
+- `logger` - takes any ruby logger (by default it's a std lib Logger.new(STDOUT), set at DEBUG level. Not available as an ENV variable config option
+- `log_level` - takes a log level, type is dependent on logger, for the default ruby Logger, use a level like `Logger::WARN`. Not available as an ENV variable config option.
+- `jwt_token_ttl` - The time to live set on JWT token in seconds, defaults to 3540 (59 minutes, the longest Snowflake supports is 60).
+- `connection_timeout` - The amount of time in seconds that the client's connection pool will wait before erroring in handing out a valid connection, defaults to 60 seconds
+- `max_connections` - The maximum number of http connections to hold open in the connection pool. If you use the client in a threaded context, you may need to increase this to be threads * client.max_threads_per_query, defaults to 16.
+- `max_threads_per_query` - The maximum number of threads the client should use to retreive data, per query, defaults to 8. If you want the client to act in a single threaded way, set this to 1
+- `thread_scale_factor` - When downloading a result set into memory, thread count is calculated by dividing a query's partition count by this number. For details on implementation see the code in `client.rb`.
+- `http_retries` - By default the client will retry common typically transient errors (http responses) twice, you can change the number of retries with this.
 
 Example configuration:
 ```ruby
-  client = RubySnowflake::Client.from_env
-  client.logger = Rails.logger
-  client.max_connections = 24
-  client.http_retries = 1
+  client = RubySnowflake::Client.from_env(
+    logger: Rails.logger
+    max_connections: 24
+    http_retries 1
+  )
 end
 ```
 

data/lib/ruby_snowflake/client/http_connection_wrapper.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 module RubySnowflake
   class Client
     class HttpConnectionWrapper

data/lib/ruby_snowflake/client/key_pair_jwt_auth_manager.rb

@@ -0,0 +1,59 @@
+# frozen_string_literal: true
+
+require "jwt"
+require "openssl"
+require "concurrent"
+
+module RubySnowflake
+  class Client
+    class KeyPairJwtAuthManager
+      # requires text of a PEM formatted RSA private key
+      def initialize(organization, account, user, private_key, jwt_token_ttl)
+        @organization = organization
+        @account = account
+        @user = user
+        @private_key_pem = private_key
+        @jwt_token_ttl = jwt_token_ttl
+
+        # start with an expired value to force creation
+        @token_expires_at = Time.now.to_i - 1
+        @token_semaphore = Concurrent::Semaphore.new(1)
+      end
+
+      def jwt_token
+        return @token unless jwt_token_expired?
+
+        @token_semaphore.acquire do
+          now = Time.now.to_i
+          @token_expires_at = now + @jwt_token_ttl
+
+          private_key = OpenSSL::PKey.read(@private_key_pem)
+
+          payload = {
+            :iss => "#{@organization.upcase}-#{@account.upcase}.#{@user}.#{public_key_fingerprint}",
+            :sub => "#{@organization.upcase}-#{@account.upcase}.#{@user}",
+            :iat => now,
+            :exp => @token_expires_at
+          }
+
+          @token = JWT.encode payload, private_key, "RS256"
+        end
+      end
+
+      private
+        def jwt_token_expired?
+          Time.now.to_i > @token_expires_at
+        end
+
+        def public_key_fingerprint
+          reutrn @public_key_fingerprint unless @public_key_fingerprint.nil?
+
+          public_key_der = OpenSSL::PKey::RSA.new(@private_key_pem).public_key.to_der
+          digest = OpenSSL::Digest::SHA256.new.digest(public_key_der)
+          fingerprint = Base64.strict_encode64(digest)
+
+          @public_key_fingerprint = "SHA256:#{fingerprint}"
+        end
+    end
+  end
+end

data/lib/ruby_snowflake/client.rb

@@ -5,22 +5,21 @@ require "benchmark"
 require "concurrent"
 require "connection_pool"
 require "json"
-require "jwt"
 require "logger"
 require "net/http"
 require "oj"
-require "openssl"
 require "retryable"
 require "securerandom"
 require "uri"
 
 
-require_relative "result"
-require_relative "streaming_result"
 require_relative "client/http_connection_wrapper"
+require_relative "client/key_pair_jwt_auth_manager"
 require_relative "client/single_thread_in_memory_strategy"
 require_relative "client/streaming_result_strategy"
 require_relative "client/threaded_in_memory_strategy"
+require_relative "result"
+require_relative "streaming_result"
 
 module RubySnowflake
   class Error < StandardError
@@ -55,11 +54,19 @@ module RubySnowflake
     # how many times to retry common retryable HTTP responses (i.e. 429, 504)
     DEFAULT_HTTP_RETRIES = 2
 
-    # can't be set after initialization
-    attr_reader :connection_timeout, :max_connections
-    attr_accessor :logger, :jwt_token_ttl, :max_threads_per_query, :thread_scale_factor, :http_retries
+    OJ_OPTIONS = { :bigdecimal_load => :bigdecimal }.freeze
 
-    def self.from_env
+    # can't be set after initialization
+    attr_reader :connection_timeout, :max_connections, :logger, :max_threads_per_query, :thread_scale_factor, :http_retries
+
+    def self.from_env(logger: DEFAULT_LOGGER,
+                      log_level: DEFAULT_LOG_LEVEL,
+                      jwt_token_ttl: env_option("SNOWFLAKE_JWT_TOKEN_TTL", DEFAULT_JWT_TOKEN_TTL),
+                      connection_timeout: env_option("SNOWFLAKE_CONNECTION_TIMEOUT", DEFAULT_CONNECTION_TIMEOUT ),
+                      max_connections: env_option("SNOWFLAKE_MAX_CONNECTIONS", DEFAULT_MAX_CONNECTIONS ),
+                      max_threads_per_query: env_option("SNOWFLAKE_MAX_THREADS_PER_QUERY", DEFAULT_MAX_THREADS_PER_QUERY),
+                      thread_scale_factor: env_option("SNOWFLAKE_THREAD_SCALE_FACTOR", DEFAULT_THREAD_SCALE_FACTOR),
+                      http_retries: env_option("SNOWFLAKE_HTTP_RETRIES", DEFAULT_HTTP_RETRIES))
       private_key = ENV["SNOWFLAKE_PRIVATE_KEY"] || File.read(ENV["SNOWFLAKE_PRIVATE_KEY_PATH"])
 
       new(
@@ -70,12 +77,14 @@ module RubySnowflake
         ENV["SNOWFLAKE_USER"],
         ENV["SNOWFLAKE_DEFAULT_WAREHOUSE"],
         ENV["SNOWFLAKE_DEFAULT_DATABASE"],
-        jwt_token_ttl: env_option("SNOWFLAKE_JWT_TOKEN_TTL", DEFAULT_JWT_TOKEN_TTL),
-        connection_timeout: env_option("SNOWFLAKE_CONNECTION_TIMEOUT", DEFAULT_CONNECTION_TIMEOUT ),
-        max_connections: env_option("SNOWFLAKE_MAX_CONNECTIONS", DEFAULT_MAX_CONNECTIONS ),
-        max_threads_per_query: env_option("SNOWFLAKE_MAX_THREADS_PER_QUERY", DEFAULT_MAX_THREADS_PER_QUERY),
-        thread_scale_factor: env_option("SNOWFLAKE_THREAD_SCALE_FACTOR", DEFAULT_THREAD_SCALE_FACTOR),
-        http_retries: env_option("SNOWFLAKE_HTTP_RETRIES", DEFAULT_HTTP_RETRIES),
+        logger: logger,
+        log_level: log_level,
+        jwt_token_ttl: jwt_token_ttl,
+        connection_timeout: connection_timeout,
+        max_connections: max_connections,
+        max_threads_per_query: max_threads_per_query,
+        thread_scale_factor: thread_scale_factor,
+        http_retries: http_retries,
       )
     end
 
@@ -91,27 +100,19 @@ module RubySnowflake
       http_retries: DEFAULT_HTTP_RETRIES
     )
       @base_uri = uri
-      @private_key_pem = private_key
-      @organization = organization
-      @account = account
-      @user = user
+      @key_pair_jwt_auth_manager =
+        KeyPairJwtAuthManager.new(organization, account, user, private_key, jwt_token_ttl)
       @default_warehouse = default_warehouse
-      @public_key_fingerprint = public_key_fingerprint(@private_key_pem)
       @default_database = default_database
 
       # set defaults for config settings
       @logger = logger
       @logger.level = log_level
-      @jwt_token_ttl = jwt_token_ttl
       @connection_timeout = connection_timeout
       @max_connections = max_connections
       @max_threads_per_query = max_threads_per_query
       @thread_scale_factor = thread_scale_factor
       @http_retries = http_retries
-
-      # start with an expired value to force creation
-      @token_expires_at = Time.now.to_i - 1
-      @token_semaphore = Concurrent::Semaphore.new(1)
     end
 
     def query(query, warehouse: nil, streaming: false, database: nil)
@@ -157,36 +158,19 @@ module RubySnowflake
         @port ||= URI.parse(@base_uri).port
       end
 
-      def jwt_token
-        return @token unless jwt_token_expired?
-
-        @token_semaphore.acquire do
-          now = Time.now.to_i
-          @token_expires_at = now + @jwt_token_ttl
-
-          private_key = OpenSSL::PKey.read(@private_key_pem)
-
-          payload = {
-            :iss => "#{@organization.upcase}-#{@account.upcase}.#{@user}.#{@public_key_fingerprint}",
-            :sub => "#{@organization.upcase}-#{@account.upcase}.#{@user}",
-            :iat => now,
-            :exp => @token_expires_at
-          }
-
-          @token = JWT.encode payload, private_key, "RS256"
+      def handle_errors(response)
+        if response.code != "200"
+          raise BadResponseError.new({}),
+            "Bad response! Got code: #{response.code}, w/ message #{response.body}"
         end
       end
 
-      def jwt_token_expired?
-        Time.now.to_i > @token_expires_at
-      end
-
       def request_with_auth_and_headers(connection, request_class, path, body=nil)
         uri = URI.parse("#{@base_uri}#{path}")
         request = request_class.new(uri)
         request["Content-Type"] = "application/json"
         request["Accept"] = "application/json"
-        request["Authorization"] = "Bearer #{jwt_token}"
+        request["Authorization"] = "Bearer #{@key_pair_jwt_auth_manager.jwt_token}"
         request["X-Snowflake-Authorization-Token-Type"] = "KEYPAIR_JWT"
         request.body = body unless body.nil?
 
@@ -230,7 +214,7 @@ module RubySnowflake
       end
 
       def retreive_result_set(response, streaming)
-        json_body = Oj.load(response.body, oj_options)
+        json_body = Oj.load(response.body, OJ_OPTIONS)
         statement_handle = json_body["statementHandle"]
         num_threads = number_of_threads_to_use(json_body["resultSetMetaData"]["partitionInfo"].size)
         retreive_proc = ->(index) { retreive_partition_data(statement_handle, index) }
@@ -254,8 +238,8 @@ module RubySnowflake
           )
         end
 
-        partition_json = nil
-        bm = Benchmark.measure { partition_json = Oj.load(partition_response.body, oj_options) }
+        partition_json = {}
+        bm = Benchmark.measure { partition_json = Oj.load(partition_response.body, OJ_OPTIONS) }
         logger.debug { "JSON parsing took: #{bm.real}" }
         partition_data = partition_json["data"]
 
@@ -265,17 +249,5 @@ module RubySnowflake
       def number_of_threads_to_use(partition_count)
         [[1, (partition_count / @thread_scale_factor.to_f).ceil].max, @max_threads_per_query].min
       end
-
-      def oj_options
-        { :bigdecimal_load => :bigdecimal }
-      end
-
-      def public_key_fingerprint(private_key_pem_string)
-        public_key_der = OpenSSL::PKey::RSA.new(private_key_pem_string).public_key.to_der
-        digest = OpenSSL::Digest::SHA256.new.digest(public_key_der)
-        fingerprint = Base64.strict_encode64(digest)
-
-        "SHA256:#{fingerprint}"
-      end
   end
 end

data/lib/ruby_snowflake/row.rb

@@ -31,14 +31,19 @@ module RubySnowflake
         else
           BigDecimal(@data[index]).round(@row_types[index][:scale])
         end
+
+      # snowflake treats these all as 64 bit IEEE 754 floating point numbers, and will we too
       when :float, :double, :"double precision", :real
-        # snowflake treats these all as 64 bit IEEE 754 floating point numbers, and will we too
         Float(@data[index])
-      when :time, :datetime, :timestamp, :timestamp_ltz, :timestamp_ntz
+
+      # Despite snowflake indicating that it sends the offset in minutes, the actual time in UTC
+      # is always sent in the first half of the data. If an offset is sent it looks like:
+      #   "1641008096.123000000 1980"
+      # If there isn't one, it's just like this:
+      #   "1641065696.123000000"
+      # in all cases, the actual time, in UTC is the float value, and the offset is ignorable
+      when :time, :datetime, :timestamp, :timestamp_ntz, :timestamp_ltz, :timestamp_tz
         Time.strptime(@data[index], TIME_FORMAT).utc
-      when :timestamp_tz
-        timestamp, offset_minutes = @data[index].split(" ")
-        Time.strptime(@data[index], TIME_FORMAT).utc - (Integer(offset_minutes) * 60)
       else
         @data[index]
       end

data/lib/ruby_snowflake/version.rb

@@ -1,3 +1,3 @@
 module RubySnowflake
-  VERSION = '0.1.2'
+  VERSION = '0.2.0'
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: rb_snowflake_client
 version: !ruby/object:Gem::Version
-  version: 0.1.2
+  version: 0.2.0
 platform: ruby
 authors:
 - Rinsed
 autorequire: 
 bindir: bin
 cert_chain: []
-date: 2023-12-04 00:00:00.000000000 Z
+date: 2023-12-07 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: concurrent-ruby
@@ -117,6 +117,7 @@ files:
 - lib/rb_snowflake_client.rb
 - lib/ruby_snowflake/client.rb
 - lib/ruby_snowflake/client/http_connection_wrapper.rb
+- lib/ruby_snowflake/client/key_pair_jwt_auth_manager.rb
 - lib/ruby_snowflake/client/single_thread_in_memory_strategy.rb
 - lib/ruby_snowflake/client/streaming_result_strategy.rb
 - lib/ruby_snowflake/client/threaded_in_memory_strategy.rb