rb_snowflake_client 0.0.6 → 0.1.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 70963fe960bc2ab443a82480ede2636c6c76f2704f1b72f0ddf67c166b3539e5
-  data.tar.gz: 4475b0653841836532b6934ea30630e4f3f89f7923ec4031f99d78d08c43e8ca
+  metadata.gz: f2dfd76d405095487b91fea79a1ab4757cdf80a3058e87d196ae112dcc93b573
+  data.tar.gz: f73bc76ce00aede0edc9a85d1afd62ce3d73a5abbff9c1037db6a271d7d38137
 SHA512:
-  metadata.gz: ba3ed329ca17e395d89e8d10a5e20d00e5b0d00ecf8f6689ad064e8385042070d9821c2944c7a0b6321baa2492c5602e49a66743294c37372ac306e9ffcee238
-  data.tar.gz: 211534933b4a6767c7866b1b926cc54cf309877ac2199d535e6b31b42ccbfcb938533e3a9692e7f4e06cf8aa723a238aa0a763231324d7ec361d0679740b421b
+  metadata.gz: 25f52a2906278092292dade37dc1c46384d2e02bb1770fb7f99b09458e3cfcdb23aaa380e62f98e35463af76b67f9eb122adf046058522c486af10c6cfde0ef2
+  data.tar.gz: '079f280217f4d1d2554f6f895d3c061b93780e812bdf5bbda8e846a4bca1f10666e8c862d6ca00d741cf2cf28565f4994ebe55b8e4e3cc5555ffbbd1ebc217a0'

data/.github/workflows/release-gh-packages.yml

@@ -1,4 +1,4 @@
-name: Release
+name: Release to Github Packages
 on:
   push:
     branches:

data/.github/workflows/release-rubygems.yml

@@ -1,4 +1,4 @@
-name: Release
+name: Release to Rubygems
 on:
   push:
     branches:

data/Gemfile

@@ -12,6 +12,7 @@ gem "dotenv"
 gem "jwt"
 gem "oj"
 gem "rake"
+gem "retryable"
 
 group :development do
   gem "parallel"

data/Gemfile.lock

@@ -1,7 +1,7 @@
 PATH
   remote: .
   specs:
-    rb_snowflake_client (0.0.6)
+    rb_snowflake_client (0.1.0)
 
 GEM
   remote: https://rubygems.org/
@@ -19,6 +19,7 @@ GEM
       coderay (~> 1.1)
       method_source (~> 1.0)
     rake (13.1.0)
+    retryable (3.0.5)
     rspec (3.12.0)
       rspec-core (~> 3.12.0)
       rspec-expectations (~> 3.12.0)
@@ -49,6 +50,7 @@ DEPENDENCIES
   pry
   rake
   rb_snowflake_client!
+  retryable
   rspec
 
 BUNDLED WITH

data/README.md

@@ -32,19 +32,28 @@ client = RubySnowflake::Client.new(
   "snowflake-account",                                  # typically your subdomain
   "snowflake-user",                                     # Your snowflake user
   "some_warehouse",                                     # The name of your warehouse to use by default
+  max_connections: 12,                                  # Config options can be passed in
+  connection_timeout: 45,                               # See below for the full set of options
 )
 
-# alternatively you can use the connect method, which will pull these values from the following environment variables. You can either provide the path to the PEM file, or it's contents in an ENV variable.
-# SNOWFLAKE_URI
-# SNOWFLAKE_PRIVATE_KEY_PATH
-#   or
-# SNOWFLAKE_PRIVATE_KEY
-# SNOWFLAKE_ORGANIZATION
-# SNOWFLAKE_ACCOUNT
-# SNOWFLAKE_USER
-# SNOWFLAKE_DEFAULT_WAREHOUSE
-RubySnowflake::Client.connect
+# alternatively you can use the `from_env` method, which will pull these values from the following environment variables. You can either provide the path to the PEM file, or it's contents in an ENV variable.
+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_JWT_TOKEN_TTL`
+`SNOWFLAKE_CONNECTION_TIMEOUT`
+`SNOWFLAKE_MAX_CONNECTIONS`
+`SNOWFLAKE_MAX_THREADS_PER_QUERY`
+`SNOWFLAKE_THREAD_SCALE_FACTOR`
+`SNOWFLAKE_HTTP_RETRIES`
 
 Once you have a client, make queries
 ```ruby
@@ -68,6 +77,27 @@ result.each do |row|
 end
 ```
 
+# 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".
+
+`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.
+
+Example configuration:
+```ruby
+  client = RubySnowflake::Client.from_env
+  client.logger = Rails.logger
+  client.max_connections = 24
+  client.http_retries = 1
+end
+
 # Gotchas
 
 1. Does not yet support multiple statements (work around is to wrap in `BEGIN ... END`)

data/lib/ruby_snowflake/client/threaded_in_memory_strategy.rb

@@ -18,13 +18,18 @@ module RubySnowflake
         end
         futures.each do |future|
           if future.rejected?
-            raise ConnectionStarvedError.new(
-                    "A partition request timed out. This is usually do to using the client in" \
-                    "multiple threads. The client uses a connection thread pool and if too many" \
-                    "requests are all done in threads at the same time, threads can get starved" \
-                    "of access to connections. The solution for this is to spin up a new client" \
-                    "instance with it's own connection pool to snowflake."
-                  )
+            if future.reason.is_a? RubySnowflake::Error
+              raise future.reason
+            else
+              raise ConnectionStarvedError.new(
+                      "A partition request timed out. This is usually do to using the client in" \
+                      "multiple threads. The client uses a connection thread pool and if too many" \
+                      "requests are all done in threads at the same time, threads can get starved" \
+                      "of access to connections. The solution for this is to either increase the " \
+                      "max_connections parameter on the client or create a new client instance" \
+                      "with it's own connection pool to snowflake per thread. Rejection reason: #{future.reason.message}"
+                    )
+            end
           end
           index, partition_data = future.value
           result[index] = partition_data

data/lib/ruby_snowflake/client.rb

@@ -6,9 +6,11 @@ require "concurrent"
 require "connection_pool"
 require "json"
 require "jwt"
+require "logger"
 require "net/http"
 require "oj"
 require "openssl"
+require "retryable"
 require "securerandom"
 require "uri"
 
@@ -33,17 +35,31 @@ module RubySnowflake
   class BadResponseError < Error ; end
   class ConnectionError < Error ; end
   class ConnectionStarvedError < Error ; end
+  class RetryableBadResponseError < Error ; end
   class RequestError < Error ; end
 
-
   class Client
-    JWT_TOKEN_TTL = 3600 # seconds, this is the max supported by snowflake
-    CONNECTION_TIMEOUT = 60 # seconds, how long for a thread to wait for a connection b4 erroring
-    MAX_CONNECTIONS = 8
-    MAX_THREADS = 8
-    THREAD_SCALE_FACTOR = 4 # parition count factor for number of threads (i.e. 2 == once we have 4 partitions, spin up a second thread)
-
-    def self.connect
+    DEFAULT_LOGGER = Logger.new(STDOUT)
+    DEFAULT_LOG_LEVEL = Logger::INFO
+    # seconds (59 min), this is the max supported by snowflake - 1 minute
+    DEFAULT_JWT_TOKEN_TTL = 3540
+    # seconds, how long for a thread to wait for a connection before erroring
+    DEFAULT_CONNECTION_TIMEOUT = 60
+    # default maximum size of the http connection pool
+    DEFAULT_MAX_CONNECTIONS = 16
+    # default maximum size of the thread pool on a single query
+    DEFAULT_MAX_THREADS_PER_QUERY = 8
+    # partition count factor for number of threads
+    # (i.e. 2 == once we have 4 partitions, spin up a second thread)
+    DEFAULT_THREAD_SCALE_FACTOR = 4
+    # 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
+
+    def self.from_env
       private_key = ENV["SNOWFLAKE_PRIVATE_KEY"] || File.read(ENV["SNOWFLAKE_PRIVATE_KEY_PATH"])
 
       new(
@@ -53,10 +69,24 @@ module RubySnowflake
         ENV["SNOWFLAKE_ACCOUNT"],
         ENV["SNOWFLAKE_USER"],
         ENV["SNOWFLAKE_DEFAULT_WAREHOUSE"],
+        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),
       )
     end
 
-    def initialize(uri, private_key, organization, account, user, default_warehouse)
+    def initialize(uri, private_key, organization, account, user, default_warehouse,
+                   logger: DEFAULT_LOGGER,
+                   log_level: DEFAULT_LOG_LEVEL,
+                   jwt_token_ttl: DEFAULT_JWT_TOKEN_TTL,
+                   connection_timeout: DEFAULT_CONNECTION_TIMEOUT,
+                   max_connections: DEFAULT_MAX_CONNECTIONS,
+                   max_threads_per_query: DEFAULT_MAX_THREADS_PER_QUERY,
+                   thread_scale_factor: DEFAULT_THREAD_SCALE_FACTOR,
+                   http_retries: DEFAULT_HTTP_RETRIES)
       @base_uri = uri
       @private_key_pem = private_key
       @organization = organization
@@ -65,6 +95,16 @@ module RubySnowflake
       @default_warehouse = default_warehouse
       @public_key_fingerprint = public_key_fingerprint(@private_key_pem)
 
+      # 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)
@@ -84,13 +124,18 @@ module RubySnowflake
           Oj.dump(request_body)
         )
       end
-      handle_errors(response)
       retreive_result_set(response, streaming)
     end
 
+    def self.env_option(env_var_name, default_value)
+      value = ENV[env_var_name]
+      value.nil? || value.empty? ? default_value : ENV[env_var_name].to_i
+    end
+    private_class_method :env_option
+
     private
       def connection_pool
-        @connection_pool ||= ConnectionPool.new(size: MAX_CONNECTIONS, timeout: CONNECTION_TIMEOUT) do
+        @connection_pool ||= ConnectionPool.new(size: @max_connections, timeout: @connection_timeout) do
           HttpConnectionWrapper.new(hostname, port).start
         end
       end
@@ -108,7 +153,7 @@ module RubySnowflake
 
         @token_semaphore.acquire do
           now = Time.now.to_i
-          @token_expires_at = now + JWT_TOKEN_TTL
+          @token_expires_at = now + @jwt_token_ttl
 
           private_key = OpenSSL::PKey.read(@private_key_pem)
 
@@ -127,13 +172,6 @@ module RubySnowflake
         Time.now.to_i > @token_expires_at
       end
 
-      def handle_errors(response)
-        if response.code != "200"
-          raise BadResponseError.new({}),
-            "Bad response! Got code: #{response.code}, w/ message #{response.body}"
-        end
-      end
-
       def request_with_auth_and_headers(connection, request_class, path, body=nil)
         uri = URI.parse("#{@base_uri}#{path}")
         request = request_class.new(uri)
@@ -143,11 +181,43 @@ module RubySnowflake
         request["X-Snowflake-Authorization-Token-Type"] = "KEYPAIR_JWT"
         request.body = body unless body.nil?
 
-        response = nil
-        bm = Benchmark.measure { response = connection.request(request) }
-        puts "HTTP Request time: #{bm.real}"
-        handle_errors(response)
-        response
+        Retryable.retryable(tries: @http_retries + 1,
+                            on: RetryableBadResponseError,
+                            log_method: retryable_log_method) do
+          response = nil
+          bm = Benchmark.measure { response = connection.request(request) }
+          logger.debug { "HTTP Request time: #{bm.real}" }
+          raise_on_bad_response(response)
+          response
+        end
+      end
+
+      def raise_on_bad_response(response)
+        return if response.code == "200"
+
+        # there are a class of errors we want to retry rather than just giving up
+        if retryable_http_response_code?(response.code)
+          raise RetryableBadResponseError.new({}),
+                "Retryable bad response! Got code: #{response.code}, w/ message #{response.body}"
+
+        else # not one we should retry
+          raise BadResponseError.new({}),
+            "Bad response! Got code: #{response.code}, w/ message #{response.body}"
+        end
+      end
+
+      # shamelessly stolen from the battle tested python client
+      # https://github.com/snowflakedb/snowflake-connector-python/blob/eceed981f93e29d2f4663241253b48340389f4ef/src/snowflake/connector/network.py#L191
+      def retryable_http_response_code?(code)
+        # retry (in order): bad request, forbidden (token expired in flight), method not allowed,
+        #   request timeout, too many requests, anything in the 500 range (504 is fairly common)
+        [400, 403, 405, 408, 429, 504].include?(code.to_i) || (500..599).include?(code)
+      end
+
+      def retryable_log_method
+        @retryable_log_method ||= proc do |retries, error|
+          logger.info("Retry attempt #{retries} because #{error.message}")
+        end
       end
 
       def retreive_result_set(response, streaming)
@@ -177,14 +247,14 @@ module RubySnowflake
 
         partition_json = nil
         bm = Benchmark.measure { partition_json = Oj.load(partition_response.body, oj_options) }
-        puts "JSON parsing took: #{bm.real}"
+        logger.debug { "JSON parsing took: #{bm.real}" }
         partition_data = partition_json["data"]
 
         partition_data
       end
 
       def number_of_threads_to_use(partition_count)
-        [[1, (partition_count / THREAD_SCALE_FACTOR.to_f).ceil].max, MAX_THREADS].min
+        [[1, (partition_count / @thread_scale_factor.to_f).ceil].max, @max_threads_per_query].min
       end
 
       def oj_options

data/lib/ruby_snowflake/version.rb

@@ -1,3 +1,3 @@
 module RubySnowflake
-  VERSION = '0.0.6'
+  VERSION = '0.1.0'
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: rb_snowflake_client
 version: !ruby/object:Gem::Version
-  version: 0.0.6
+  version: 0.1.0
 platform: ruby
 authors:
 - Rinsed
 autorequire: 
 bindir: bin
 cert_chain: []
-date: 2023-11-28 00:00:00.000000000 Z
+date: 2023-11-29 00:00:00.000000000 Z
 dependencies: []
 description: 'Using the HTTP V2 Api for Snowflake runs queries & creates native Ruby
   objects.