rb_snowflake_client 0.2.0 → 0.3.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 2a93f0b1e503e2d219db8b1d1133395ca484b45aaa1f91ceef013a50a116b973
-  data.tar.gz: c292c5b0d25a921f8deba6882fe288ed084aa07919fa0d056899d20b921612fc
+  metadata.gz: b345626b574463c788bc7ff125c6297eac071ce62cd77fadb77c5da896f6089d
+  data.tar.gz: 3bed923cc293d33af04eb846241637d6c72a8bef8917958e50b013fa8d120f70
 SHA512:
-  metadata.gz: 28428d4cf4bd70a11fefad327a2ab14deff17e95bf28b999a6b84d67756b35d3207c8224d6a028b9998ca85b62a7226f4df2beddea50f30e9a91350d7df73b4b
-  data.tar.gz: ad55068885a29f1cae0181e52d66764651cb4f7d541e46680bd15644ed8c17f10dda8cb9356ae4948cdf9b249e0c726abc540e26b03e1b76f9f2c538a8251cce
+  metadata.gz: 952d6ffe14f350158abbe2342815687e36564233f7de8a452c27174f2f483f76c32e331502a1a35e7e20a654f11a74c109f7018410a9a3add7f7fa42ea3e32c7
+  data.tar.gz: 3dcc696ab574bebb3b0416dbc9595df656c6547f7b4a36b70d32d63030dfedc54d9f84e8b9838c0fd22c3cb5a9b9ce0efdc9c74b78e3bb41a57e58c29757c120

data/Gemfile.lock

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

data/README.md

@@ -37,6 +37,7 @@ client = RubySnowflake::Client.new(
   "some_database",                                      # The name of the database in the context of which the queries will run
   max_connections: 12,                                  # Config options can be passed in
   connection_timeout: 45,                               # See below for the full set of options
+  query_timeout: 1200,                                  # how long to wait for queries, in seconds
 )
 
 # 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.
@@ -57,6 +58,7 @@ Available ENV variables (see below in the config section for details)
 - `SNOWFLAKE_MAX_THREADS_PER_QUERY`
 - `SNOWFLAKE_THREAD_SCALE_FACTOR`
 - `SNOWFLAKE_HTTP_RETRIES`
+- `SNOWFLAKE_QUERY_TIMEOUT`
 
 ## Make queries
 
@@ -115,6 +117,7 @@ The client supports the following configuration options, each with their own get
 - `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.
+- `query_timeout` - By default the client will wait 10 minutes (600s) for a query to finish, you can change this default, will also set this limit in the query for snowflake to obey. Set in seconds.
 
 Example configuration:
 ```ruby

data/lib/ruby_snowflake/client.rb

@@ -12,7 +12,6 @@ require "retryable"
 require "securerandom"
 require "uri"
 
-
 require_relative "client/http_connection_wrapper"
 require_relative "client/key_pair_jwt_auth_manager"
 require_relative "client/single_thread_in_memory_strategy"
@@ -36,6 +35,7 @@ module RubySnowflake
   class ConnectionStarvedError < Error ; end
   class RetryableBadResponseError < Error ; end
   class RequestError < Error ; end
+  class QueryTimeoutError < Error ;  end
 
   class Client
     DEFAULT_LOGGER = Logger.new(STDOUT)
@@ -53,11 +53,16 @@ module RubySnowflake
     DEFAULT_THREAD_SCALE_FACTOR = 4
     # how many times to retry common retryable HTTP responses (i.e. 429, 504)
     DEFAULT_HTTP_RETRIES = 2
+    # how long to wait to allow a query to complete, in seconds
+    DEFAULT_QUERY_TIMEOUT = 600 # 10 minutes
 
     OJ_OPTIONS = { :bigdecimal_load => :bigdecimal }.freeze
+    VALID_RESPONSE_CODES = %w(200 202).freeze
+    POLLING_RESPONSE_CODE = "202"
+    POLLING_INTERVAL = 2 # seconds
 
     # can't be set after initialization
-    attr_reader :connection_timeout, :max_connections, :logger, :max_threads_per_query, :thread_scale_factor, :http_retries
+    attr_reader :connection_timeout, :max_connections, :logger, :max_threads_per_query, :thread_scale_factor, :http_retries, :query_timeout
 
     def self.from_env(logger: DEFAULT_LOGGER,
                       log_level: DEFAULT_LOG_LEVEL,
@@ -66,7 +71,8 @@ module RubySnowflake
                       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))
+                      http_retries: env_option("SNOWFLAKE_HTTP_RETRIES", DEFAULT_HTTP_RETRIES),
+                      query_timeout: env_option("SNOWFLAKE_QUERY_TIMEOUT", DEFAULT_QUERY_TIMEOUT))
       private_key = ENV["SNOWFLAKE_PRIVATE_KEY"] || File.read(ENV["SNOWFLAKE_PRIVATE_KEY_PATH"])
 
       new(
@@ -85,6 +91,7 @@ module RubySnowflake
         max_threads_per_query: max_threads_per_query,
         thread_scale_factor: thread_scale_factor,
         http_retries: http_retries,
+        query_timeout: query_timeout,
       )
     end
 
@@ -97,7 +104,8 @@ module RubySnowflake
       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
+      http_retries: DEFAULT_HTTP_RETRIES,
+      query_timeout: DEFAULT_QUERY_TIMEOUT
     )
       @base_uri = uri
       @key_pair_jwt_auth_manager =
@@ -113,26 +121,33 @@ module RubySnowflake
       @max_threads_per_query = max_threads_per_query
       @thread_scale_factor = thread_scale_factor
       @http_retries = http_retries
+      @query_timeout = query_timeout
+
+      # Do NOT use normally, this exists for tests so we can reliably trigger the polling
+      # response workflow from snowflake in tests
+      @_enable_polling_queries = false
     end
 
     def query(query, warehouse: nil, streaming: false, database: nil)
       warehouse ||= @default_warehouse
       database ||= @default_database
 
+      query_start_time = Time.now.to_i
       response = nil
       connection_pool.with do |connection|
         request_body = {
-          "statement" => query, "warehouse" => warehouse, "database" =>  database
+          "statement" => query, "warehouse" => warehouse,
+          "database" =>  database, "timeout" => @query_timeout
         }
 
         response = request_with_auth_and_headers(
           connection,
           Net::HTTP::Post,
-          "/api/v2/statements?requestId=#{SecureRandom.uuid}",
+          "/api/v2/statements?requestId=#{SecureRandom.uuid}&async=#{@_enable_polling_queries}",
           Oj.dump(request_body)
         )
       end
-      retreive_result_set(response, streaming)
+      retreive_result_set(query_start_time, query, response, streaming)
     end
 
     alias fetch query
@@ -158,13 +173,6 @@ module RubySnowflake
         @port ||= URI.parse(@base_uri).port
       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)
@@ -186,7 +194,7 @@ module RubySnowflake
       end
 
       def raise_on_bad_response(response)
-        return if response.code == "200"
+        return if VALID_RESPONSE_CODES.include? response.code
 
         # there are a class of errors we want to retry rather than just giving up
         if retryable_http_response_code?(response.code)
@@ -213,9 +221,50 @@ module RubySnowflake
         end
       end
 
-      def retreive_result_set(response, streaming)
+      def poll_for_completion_or_timeout(query_start_time, query, statement_handle)
+        first_data_json_body = nil
+
+        connection_pool.with do |connection|
+          loop do
+            sleep POLLING_INTERVAL
+
+            if Time.now.to_i - query_start_time > @query_timeout
+              cancelled = attempt_to_cancel_and_silence_errors(connection, statement_handle)
+              raise QueryTimeoutError.new("Query timed out. Query cancelled? #{cancelled} Query: #{query}")
+            end
+
+            poll_response = request_with_auth_and_headers(connection, Net::HTTP::Get,
+                                                          "/api/v2/statements/#{statement_handle}")
+            if poll_response.code == POLLING_RESPONSE_CODE
+              next
+            else
+              return poll_response
+            end
+          end
+        end
+      end
+
+      def attempt_to_cancel_and_silence_errors(connection, statement_handle)
+        cancel_response = request_with_auth_and_headers(connection, Net::HTTP::Post,
+                                                        "/api/v2/#{statement_handle}/cancel")
+        true
+      rescue Error => error
+        if error.is_a?(BadResponseError) && error.message.include?("404")
+          return true # snowflake cancelled it before we did
+        end
+        @logger.error("Error on attempting to cancel query #{statement_handle}, will raise a QueryTimeoutError")
+        false
+      end
+
+      def retreive_result_set(query_start_time, query, response, streaming)
         json_body = Oj.load(response.body, OJ_OPTIONS)
         statement_handle = json_body["statementHandle"]
+
+        if response.code == POLLING_RESPONSE_CODE
+          result_response = poll_for_completion_or_timeout(query_start_time, query, statement_handle)
+          json_body = Oj.load(result_response.body, OJ_OPTIONS)
+        end
+
         num_threads = number_of_threads_to_use(json_body["resultSetMetaData"]["partitionInfo"].size)
         retreive_proc = ->(index) { retreive_partition_data(statement_handle, index) }
 

data/lib/ruby_snowflake/version.rb

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

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: rb_snowflake_client
 version: !ruby/object:Gem::Version
-  version: 0.2.0
+  version: 0.3.0
 platform: ruby
 authors:
 - Rinsed
 autorequire: 
 bindir: bin
 cert_chain: []
-date: 2023-12-07 00:00:00.000000000 Z
+date: 2023-12-08 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: concurrent-ruby