rb_snowflake_client 1.4.0 → 1.5.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 9c1581fededb6e083bc1f55249d806fca4a63462fafee98f9277986729dd1107
-  data.tar.gz: e810744812f56a65d519d8a30ce96e97d2f1341479ee7aeb2a50766c0f541e90
+  metadata.gz: ddf76c34f14ea4b7192e5bd3dbdb18b7d176e66c6327fab97ef8d4352138d964
+  data.tar.gz: 20223663cc72100e28b22a4ea572611d2b2bc8291519d1cde9f2b536c02865f1
 SHA512:
-  metadata.gz: 7cc78b78f3a5a3d056c826107aab889d23705ea76249800a7c098e43e8515db4fc9f85e11a2fae86cced570b1c5aaa869dded60495467bbb2b52241edd466c2e
-  data.tar.gz: 6b57c2a14a03fe96711d86554eb13b86bd5865a1aafa629a30ea02b2c27fc72f743f8942fc0d4d35b590d824cbd110c2ad3efbe6ffc235037c7039213644e97c
+  metadata.gz: 266de51be70f28c748b703fc55787b0a495404a62226fba151c44dfe00ce31465c3aae5e516ba362cdfa8a22a0d597514d86fd6644f263a046cb275b0a3e3f46
+  data.tar.gz: b0e90b538d9b6557d7de1b58f085570097036df3a9d1ea6be500babb6b56fafb50f21a4d905f1c50e8a274df7ffdd58b6268d65f9dc2220c7e9f70680662479c

data/CHANGELOG.md

@@ -5,6 +5,16 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## Unreleased
+
+## [1.5.0] - 2025-10-14
+### Added
+- Instrumentation feature added for Active Support users
+- Added `query_timeout` as a per-query parameter, allowing timeout override on individual queries
+### Fixed
+- `query_timeout` now properly sends timeout parameter to Snowflake API for server-side enforcement
+- Streaming mode now releases consumed records, fixing memory leak. Note: if you were iterating over streaming results more than once, this is a breaking change (though that was not its intended usage).
+
 ## [1.4.0] - 2025-05-01
 ### Added
 - Enhanced Row API to implement Enumerable interface

data/Gemfile

@@ -8,6 +8,10 @@ gemspec
 gem "bundler"
 gem "rake"
 
+group :development, :test do
+  gem "activesupport"
+end
+
 group :development do
   gem "parallel"
   gem "pry"

data/Gemfile.lock

@@ -1,7 +1,7 @@
 PATH
   remote: .
   specs:
-    rb_snowflake_client (1.4.0)
+    rb_snowflake_client (1.5.0)
       bigdecimal (>= 3.0)
       concurrent-ruby (>= 1.2)
       connection_pool (>= 2.4)
@@ -13,42 +13,66 @@ PATH
 GEM
   remote: https://rubygems.org/
   specs:
-    base64 (0.2.0)
-    bigdecimal (3.1.9)
+    activesupport (8.0.3)
+      base64
+      benchmark (>= 0.3)
+      bigdecimal
+      concurrent-ruby (~> 1.0, >= 1.3.1)
+      connection_pool (>= 2.2.5)
+      drb
+      i18n (>= 1.6, < 2)
+      logger (>= 1.4.2)
+      minitest (>= 5.1)
+      securerandom (>= 0.3)
+      tzinfo (~> 2.0, >= 2.0.5)
+      uri (>= 0.13.1)
+    base64 (0.3.0)
+    benchmark (0.4.1)
+    bigdecimal (3.3.1)
     coderay (1.1.3)
     concurrent-ruby (1.3.5)
-    connection_pool (2.5.3)
-    diff-lcs (1.6.1)
+    connection_pool (2.5.4)
+    diff-lcs (1.6.2)
     dotenv (3.1.8)
-    json (2.11.3)
-    jwt (2.10.1)
+    drb (2.2.3)
+    i18n (1.14.7)
+      concurrent-ruby (~> 1.0)
+    json (2.15.1)
+    jwt (3.1.2)
       base64
+    logger (1.7.0)
     method_source (1.1.0)
+    minitest (5.26.0)
     parallel (1.27.0)
     pry (0.15.2)
       coderay (~> 1.1)
       method_source (~> 1.0)
-    rake (13.2.1)
+    rake (13.3.0)
     retryable (3.0.5)
-    rspec (3.13.0)
+    rspec (3.13.1)
       rspec-core (~> 3.13.0)
       rspec-expectations (~> 3.13.0)
       rspec-mocks (~> 3.13.0)
-    rspec-core (3.13.3)
+    rspec-core (3.13.5)
       rspec-support (~> 3.13.0)
-    rspec-expectations (3.13.3)
+    rspec-expectations (3.13.5)
       diff-lcs (>= 1.2.0, < 2.0)
       rspec-support (~> 3.13.0)
-    rspec-mocks (3.13.2)
+    rspec-mocks (3.13.5)
       diff-lcs (>= 1.2.0, < 2.0)
       rspec-support (~> 3.13.0)
-    rspec-support (3.13.2)
+    rspec-support (3.13.6)
+    securerandom (0.4.1)
+    tzinfo (2.0.6)
+      concurrent-ruby (~> 1.0)
+    uri (1.0.4)
 
 PLATFORMS
   arm64-darwin-22
   ruby
 
 DEPENDENCIES
+  activesupport
   bundler
   parallel
   pry

data/README.md

@@ -132,6 +132,14 @@ Queries by default use the primary role assigned to the account. If there are mu
 client.query("SELECT * FROM BIGTABLE", role: "MY_ROLE")
 ```
 
+## Query timeout
+
+You can override the query timeout on a per-query basis. The timeout is specified in seconds and will be enforced by both Snowflake server-side and the client-side polling mechanism.
+
+```ruby
+client.query("SELECT * FROM BIGTABLE", query_timeout: 30)
+```
+
 ## Binding parameters
 
 Say we have `BIGTABLE` with a `data` column of a type `VARIANT`.
@@ -140,16 +148,45 @@ Say we have `BIGTABLE` with a `data` column of a type `VARIANT`.
 json_string = '{"valid": "json"}'
 query = "insert into BIGTABLE(data) select parse_json(?)"
 bindings = {
-  "1": {
-    "type": "TEXT",
-    "value": json_string
-  }
-}
+      "1" => {
+        "type" => "TEXT",
+        "value" => "Other Event"
+      }
+    }
 client.query(query, bindings: bindings)
 ```
 
 For additional information about binding parameters refer to snowflake documentation: https://docs.snowflake.com/en/developer-guide/sql-api/submitting-requests#using-bind-variables-in-a-statement
 
+## Instrumentation
+
+If ActiveSupport is available, this library additionally emits [notification events](https://api.rubyonrails.org/classes/ActiveSupport/Notifications.html) around queries. You can subscribe to those to track timing, query counts, etc.
+
+* `rb_snowflake_client.snowflake_query.finish`: published at query end
+
+Events receive a payload with the following properties:
+* `database`: snowflake database
+* `schema`: snowflake schema
+* `warehouse`: snowflake warehouse
+* `query_id`: random UUID for the query
+* `query_name`: argument passed to query/fetch
+* `exception`: present if the query raised an error, see [Notifications documentation](https://api.rubyonrails.org/classes/ActiveSupport/Notifications.html#module-ActiveSupport::Notifications-label-Subscribers) for details
+* `exception_object`: present if the query raised an error, see [Notifications documentation](https://api.rubyonrails.org/classes/ActiveSupport/Notifications.html#module-ActiveSupport::Notifications-label-Subscribers) for details
+
+An example integration with [Datadog](https://www.rubydoc.info/gems/datadog) might look like this:
+
+```ruby
+ActiveSupport::Notifications.subscribe("rb_snowflake_client.snowflake_query.finish") do |name, start, finish, id, payload|
+  span = Datadog::Tracing.trace(payload[:query_name] || "snowflake_query",
+    resource: "snowflake",
+    start_time: start,
+    tags: payload,
+    type: Datadog::Tracing::Metadata::Ext::AppTypes::TYPE_DB)
+  
+  span.finish(finish)
+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". Configuration options can only be set on initialization through `new` or `from_env`.

data/lib/ruby_snowflake/client.rb

@@ -12,6 +12,13 @@ require "retryable"
 require "securerandom"
 require "uri"
 
+begin
+  require "active_support"
+  require "active_support/notifications"
+rescue LoadError
+  # This isn't required
+end
+
 require_relative "client/http_connection_wrapper"
 require_relative "client/key_pair_jwt_auth_manager"
 require_relative "client/single_thread_in_memory_strategy"
@@ -143,31 +150,35 @@ module RubySnowflake
       @_enable_polling_queries = false
     end
 
-    def query(query, warehouse: nil, streaming: false, database: nil, schema: nil, bindings: nil, role: nil)
+    def query(query, warehouse: nil, streaming: false, database: nil, schema: nil, bindings: nil, role: nil, query_name: nil, query_timeout: nil)
       warehouse ||= @default_warehouse
       database ||= @default_database
       role ||= @default_role
+      query_timeout ||= @query_timeout
 
-      query_start_time = Time.now.to_i
-      response = nil
-      connection_pool.with do |connection|
-        request_body = {
-          "warehouse" => warehouse&.upcase,
-          "schema" => schema&.upcase,
-          "database" =>  database&.upcase,
-          "statement" => query,
-          "bindings" => bindings,
-          "role" => role
-        }
-
-        response = request_with_auth_and_headers(
-          connection,
-          Net::HTTP::Post,
-          "/api/v2/statements?requestId=#{SecureRandom.uuid}&async=#{@_enable_polling_queries}",
-          request_body.to_json
-        )
+      with_instrumentation({ database:, schema:, warehouse:, query_name: }) do
+        query_start_time = Time.now.to_i
+        response = nil
+        connection_pool.with do |connection|
+          request_body = {
+            "warehouse" => warehouse&.upcase,
+            "schema" => schema&.upcase,
+            "database" =>  database&.upcase,
+            "statement" => query,
+            "bindings" => bindings,
+            "role" => role,
+            "timeout" => query_timeout
+          }
+
+          response = request_with_auth_and_headers(
+            connection,
+            Net::HTTP::Post,
+            "/api/v2/statements?requestId=#{SecureRandom.uuid}&async=#{@_enable_polling_queries}",
+            request_body.to_json
+          )
+        end
+        retrieve_result_set(query_start_time, query, response, streaming, query_timeout)
       end
-      retrieve_result_set(query_start_time, query, response, streaming)
     end
 
     alias fetch query
@@ -251,7 +262,7 @@ module RubySnowflake
         end
       end
 
-      def poll_for_completion_or_timeout(query_start_time, query, statement_handle)
+      def poll_for_completion_or_timeout(query_start_time, query, statement_handle, query_timeout)
         first_data_json_body = nil
 
         connection_pool.with do |connection|
@@ -259,7 +270,7 @@ module RubySnowflake
             sleep POLLING_INTERVAL
 
             elapsed_time = Time.now.to_i - query_start_time
-            if elapsed_time > @query_timeout
+            if elapsed_time > query_timeout
               cancelled = attempt_to_cancel_and_silence_errors(connection, statement_handle)
               raise QueryTimeoutError.new("Query timed out. Query cancelled? #{cancelled}; Duration: #{elapsed_time}; Query: '#{query}'")
             end
@@ -287,12 +298,12 @@ module RubySnowflake
         false
       end
 
-      def retrieve_result_set(query_start_time, query, response, streaming)
+      def retrieve_result_set(query_start_time, query, response, streaming, query_timeout)
         json_body = JSON.parse(response.body, JSON_PARSE_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)
+          result_response = poll_for_completion_or_timeout(query_start_time, query, statement_handle, query_timeout)
           json_body = JSON.parse(result_response.body, JSON_PARSE_OPTIONS)
         end
 
@@ -329,5 +340,15 @@ 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 with_instrumentation(tags, &block)
+        return block.call unless defined?(::ActiveSupport) && ::ActiveSupport
+
+        ::ActiveSupport::Notifications.instrument(
+            "rb_snowflake_client.snowflake_query.finish",
+            tags.merge(query_id: SecureRandom.uuid)) do
+          block.call
+        end
+      end
   end
 end

data/lib/ruby_snowflake/streaming_result.rb

@@ -27,9 +27,19 @@ module RubySnowflake
         if data[index].is_a? Concurrent::Future
           data[index] = data[index].value # wait for it to finish
         end
+
         data[index].each do |row|
           yield wrap_row(row)
         end
+
+        # After iterating over the current partition, clear the data to release memory
+        data[index].clear
+
+        # Reassign to a symbol so:
+        # - When looking at the list of partitions in `data` it is easier to detect
+        # - Will raise an exception if `data.each` is attempted to be called again
+        # - It won't trigger prefetch detection as `next_index`
+        data[index] = :finished
       end
     end
 

data/lib/ruby_snowflake/version.rb

@@ -1,3 +1,3 @@
 module RubySnowflake
-  VERSION = "1.4.0"
+  VERSION = "1.5.0"
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: rb_snowflake_client
 version: !ruby/object:Gem::Version
-  version: 1.4.0
+  version: 1.5.0
 platform: ruby
 authors:
 - Rinsed
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2025-05-01 00:00:00.000000000 Z
+date: 2025-10-14 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: bigdecimal