pinot-client 1.22.0 → 1.24.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 8fa5d08fb08bee353252e2da56932948941faa5fdbdd4f3410e5ff8887755e04
-  data.tar.gz: c8bb9cbe06c5d209d8f9f63031776fc604a153730b4be207b7fd6d8a3b69e5f0
+  metadata.gz: b893ea9b2cefc1aca8cf86940399d36be6e071e88adc8023affb660d94ad8ce7
+  data.tar.gz: 539a393551125f7169975b0e641f4d1b60fd0cc6390fc1c9c10b4f86c4036db0
 SHA512:
-  metadata.gz: 6d5c85b83effd2ce352b09c03d350eedd817415a18994ed23ef7be972f80dfb8d72cbbbace9c83bfddedb095b6695af0b8b4538f2f0023ee40844e6b554849a9
-  data.tar.gz: c450ffb393a8798ae3ea3e608cf2cb976bdf5aeb6f7145121764c41bed33069807fdb3087d02c3963552f79568fb6f2020088794cccf47f97218d10fe7041c99
+  metadata.gz: 13d9d34e4a86ac6d62b15db000f4220e28cd5498e78cfdb42b2c9edcfac5645c68332a43048d96bc16089381fc4ea9c3d0ce429a95480c481766d1ebd9dfd6ca
+  data.tar.gz: 679bde8940819c309cbc351712cf2c4d4f3cfa1601e81ae023e5bf03a944404bd907b1f0a6fa44c2d5925909f3b8ce6533adf2a91cfea5aa284be7ab86afc9cf

data/lib/pinot/connection.rb

@@ -27,13 +27,14 @@ module Pinot
       @trace = false
     end
 
-    def execute_sql(table, query, query_timeout_ms: nil)
+    def execute_sql(table, query, query_timeout_ms: nil, headers: {})
       Pinot::Instrumentation.instrument(table: table, query: query) do
         logger.debug "Executing SQL on table=#{table}: #{query}"
         broker = @broker_selector.select_broker(table)
         effective_timeout = query_timeout_ms || @query_timeout_ms
         run_with_circuit_breaker(broker) do
-          @transport.execute(broker, build_request(query, timeout_ms: effective_timeout))
+          @transport.execute(broker, build_request(query, timeout_ms: effective_timeout),
+                             extra_request_headers: headers)
         end
       end
     end
@@ -42,11 +43,10 @@ module Pinot
       execute_sql(table, query, query_timeout_ms: timeout_ms)
     end
 
-    def execute_sql_with_params(table, query_pattern, params, query_timeout_ms: nil)
+    def execute_sql_with_params(table, query_pattern, params, query_timeout_ms: nil, headers: {})
       query = format_query(query_pattern, params)
-      execute_sql(table, query, query_timeout_ms: query_timeout_ms)
+      execute_sql(table, query, query_timeout_ms: query_timeout_ms, headers: headers)
     rescue => e
-      # Re-raise format errors directly (they already have the right message)
       raise e
     end
 

data/lib/pinot/prepared_statement.rb

@@ -60,7 +60,7 @@ module Pinot
       nil
     end
 
-    def execute
+    def execute(headers: {})
       @mutex.synchronize do
         raise PreparedStatementClosedError, "prepared statement is closed" if @closed
         @parameters.each_with_index do |p, i|
@@ -72,10 +72,10 @@ module Pinot
       rescue => e
         raise "failed to build query: #{e.message}"
       end
-      @connection.execute_sql(@table, query)
+      @connection.execute_sql(@table, query, headers: headers)
     end
 
-    def execute_with_params(*params)
+    def execute_with_params(*params, headers: {})
       @mutex.synchronize { raise PreparedStatementClosedError, "prepared statement is closed" if @closed }
       if params.length != @param_count
         raise "expected #{@param_count} parameters, got #{params.length}"
@@ -85,7 +85,7 @@ module Pinot
       rescue => e
         raise "failed to build query: #{e.message}"
       end
-      @connection.execute_sql(@table, query)
+      @connection.execute_sql(@table, query, headers: headers)
     end
 
     def clear_parameters

data/lib/pinot/transport.rb

@@ -171,6 +171,7 @@ module Pinot
 
     # HTTP status codes that map to specific error classes and are safe to retry
     HTTP_ERROR_MAP = {
+      "408" => QueryTimeoutError,
       "429" => RateLimitError,
       "503" => BrokerUnavailableError,
       "504" => BrokerUnavailableError
@@ -178,6 +179,10 @@ module Pinot
 
     RETRYABLE_HTTP_ERRORS = [RateLimitError, BrokerUnavailableError].freeze
 
+    # Pinot exception errorCode values that indicate query timeout.
+    # 250 = ExecutionTimeoutError (server-side), 400 = BrokerTimeoutError.
+    TIMEOUT_ERROR_CODES = [250, 400].freeze
+
     def initialize(http_client:, extra_headers: {}, timeout_ms: nil, logger: nil,
                    max_retries: 0, retry_interval_ms: 200)
       @http_client = http_client
@@ -188,7 +193,7 @@ module Pinot
       @retry_interval_ms = retry_interval_ms
     end
 
-    def execute(broker_address, request)
+    def execute(broker_address, request, extra_request_headers: {})
       logger.debug "Pinot query to #{broker_address}: #{request.query}"
 
       attempts = 0
@@ -202,6 +207,7 @@ module Pinot
         headers = DEFAULT_HEADERS
           .merge(@extra_headers)
           .merge("X-Correlation-Id" => SecureRandom.uuid)
+          .merge(extra_request_headers)
 
         resp = @http_client.post(url, body: body, headers: headers)
 
@@ -215,18 +221,24 @@ module Pinot
           raise TransportError, "http exception with HTTP status code #{resp.code}"
         end
 
-        begin
+        broker_response = begin
           BrokerResponse.from_json(resp.body)
         rescue JSON::ParserError => e
           raise e.message
         end
+
+        if (timeout_ex = broker_response.exceptions.find { |ex| TIMEOUT_ERROR_CODES.include?(ex.error_code) })
+          raise QueryTimeoutError, timeout_ex.message
+        end
+
+        broker_response
       rescue *RETRYABLE_HTTP_ERRORS, *RETRYABLE_ERRORS => e
         if attempts < max_attempts
           sleep_ms = (@retry_interval_ms || 200) * (2 ** (attempts - 1))
           sleep(sleep_ms / 1000.0)
           retry
         end
-        raise
+        raise Net::ReadTimeout === e || Net::WriteTimeout === e ? QueryTimeoutError.new(e.message) : e
       end
     end
 

data/lib/pinot/version.rb

@@ -1,3 +1,3 @@
 module Pinot
-  VERSION = "1.22.0"
+  VERSION = "1.24.0"
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: pinot-client
 version: !ruby/object:Gem::Version
-  version: 1.22.0
+  version: 1.24.0
 platform: ruby
 authors:
 - Xiang Fu