pinot-client 1.24.0 → 1.26.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: b893ea9b2cefc1aca8cf86940399d36be6e071e88adc8023affb660d94ad8ce7
-  data.tar.gz: 539a393551125f7169975b0e641f4d1b60fd0cc6390fc1c9c10b4f86c4036db0
+  metadata.gz: 0c6d47938639a210b77881492e1155d26728be12f987b9721712f16503183ec3
+  data.tar.gz: 9e14c467edbba3ad7fff1e20f6b288ac810895c5069549cd33f08e2972b8d739
 SHA512:
-  metadata.gz: 13d9d34e4a86ac6d62b15db000f4220e28cd5498e78cfdb42b2c9edcfac5645c68332a43048d96bc16089381fc4ea9c3d0ce429a95480c481766d1ebd9dfd6ca
-  data.tar.gz: 679bde8940819c309cbc351712cf2c4d4f3cfa1601e81ae023e5bf03a944404bd907b1f0a6fa44c2d5925909f3b8ce6533adf2a91cfea5aa284be7ab86afc9cf
+  metadata.gz: 73ec7a9ae84e96f71618f430ed542527fb9f9e16f64b1d0d78a1a8bcb0ba1950f6d7117072f0318ffcc08fdb0df2143be268667d6b96e4914d52fc4f0448c0fa
+  data.tar.gz: 9fe83c991b5ef549ce915dbe4af55b50143d17d6d9660915d5fd6308798541f81068e557317b8b88af167b55a63ea7781ef14dfff428cd51b671df7d2a0267b2

data/lib/pinot/connection.rb

@@ -50,6 +50,46 @@ module Pinot
       raise e
     end
 
+    def execute_many(queries, max_concurrency: nil)
+      return [] if queries.empty?
+
+      results  = Array.new(queries.size)
+      # Queue acts as a counting semaphore: pre-filled with N tokens.
+      sem = max_concurrency ? build_semaphore(max_concurrency) : nil
+
+      threads = queries.each_with_index.map do |item, idx|
+        table      = item[:table]           || item["table"]           || ""
+        query      = item[:query]           || item["query"]           || ""
+        timeout_ms = item[:query_timeout_ms] || item["query_timeout_ms"]
+
+        Thread.new do
+          sem&.pop   # acquire
+          begin
+            resp = execute_sql(table, query, query_timeout_ms: timeout_ms)
+            results[idx] = QueryResult.new(table: table, query: query, response: resp, error: nil)
+          rescue => e
+            results[idx] = QueryResult.new(table: table, query: query, response: nil, error: e)
+          ensure
+            sem&.push(:token)  # release
+          end
+        end
+      end
+
+      threads.each(&:join)
+      results
+    end
+
+    def paginate(query, page_size: Paginator::DEFAULT_PAGE_SIZE, table: nil, extra_headers: {})
+      broker = @broker_selector.select_broker(table || "")
+      Paginator.new(
+        @transport.http_client,
+        broker,
+        query,
+        page_size:     page_size,
+        extra_headers: extra_headers
+      )
+    end
+
     def prepare(table, query_template)
       raise ArgumentError, "table name cannot be empty" if table.nil? || table.strip.empty?
       raise ArgumentError, "query template cannot be empty" if query_template.nil? || query_template.strip.empty?
@@ -101,6 +141,12 @@ module Pinot
 
     private
 
+    def build_semaphore(n)
+      q = SizedQueue.new(n)
+      n.times { q.push(:token) }
+      q
+    end
+
     def run_with_circuit_breaker(broker, &block)
       return yield unless @circuit_breaker_registry
       @circuit_breaker_registry.for(broker).call(broker, &block)

data/lib/pinot/paginator.rb

@@ -0,0 +1,123 @@
+module Pinot
+  # Implements Pinot's server-side cursor API.
+  #
+  # The broker stores the full result set and returns slices on demand.
+  # All fetch requests after the first must go to the same broker that
+  # owns the cursor state (brokerHost:brokerPort from the initial response).
+  #
+  # Usage:
+  #   paginator = conn.paginate("SELECT * FROM t LIMIT 10000", page_size: 100)
+  #   paginator.each_page { |resp| process(resp.result_table) }
+  #   paginator.each_row  { |row|  puts row.map(&:to_s).join(", ") }
+  #   paginator.delete    # optional early cleanup; also called automatically on exhaustion
+  class Paginator
+    include Enumerable
+
+    DEFAULT_PAGE_SIZE = 1000
+
+    def initialize(http_client, broker_address, query, page_size:, extra_headers: {})
+      raise ArgumentError, "page_size must be a positive integer" unless page_size.is_a?(Integer) && page_size > 0
+
+      @http_client    = http_client
+      @broker_address = broker_address
+      @query          = query
+      @page_size      = page_size
+      @extra_headers  = extra_headers
+
+      @request_id  = nil
+      @cursor_base = nil  # "http://host:port" — set after first response
+      @exhausted   = false
+    end
+
+    # Yields each page as a BrokerResponse. Returns an Enumerator without a block.
+    def each_page
+      return enum_for(:each_page) unless block_given?
+
+      # Submit the query and get the first page + cursor metadata
+      first = submit_cursor
+      return if first.result_table.nil? || first.result_table.rows.empty?
+
+      yield first
+
+      fetched = first.num_rows || first.result_table.rows.size
+      total   = first.num_rows_result_set || 0
+
+      while fetched < total
+        page = fetch_page(fetched)
+        rows = page.result_table&.rows || []
+        break if rows.empty?
+
+        yield page
+
+        fetched += rows.size
+        break if rows.size < @page_size
+      end
+
+      delete
+    end
+
+    # Yields each row Array across all pages. Returns an Enumerator without a block.
+    # Aliased as #each so Enumerable methods (.map, .select, .to_a, etc.) work.
+    def each(&block)
+      return enum_for(:each) unless block_given?
+
+      each_page do |response|
+        response.result_table.rows.each(&block)
+      end
+    end
+
+    alias each_row each
+
+    # Delete the cursor from the broker early (also called automatically after exhaustion).
+    def delete
+      return unless @request_id && @cursor_base
+
+      url = "#{@cursor_base}/responseStore/#{@request_id}"
+      @http_client.delete(url, headers: json_headers)
+      @request_id = nil
+    rescue StandardError
+      # best-effort; cursor will expire naturally
+    end
+
+    private
+
+    def submit_cursor
+      base   = broker_base(@broker_address)
+      url    = "#{base}/query/sql?getCursor=true&numRows=#{@page_size}"
+      body   = JSON.generate("sql" => @query)
+      resp   = @http_client.post(url, body: body, headers: json_headers)
+
+      raise TransportError, "cursor submit returned HTTP #{resp.code}" unless resp.code.to_i == 200
+
+      parsed = BrokerResponse.from_json(resp.body)
+
+      @request_id  = parsed.request_id
+      @cursor_base = broker_base_from_response(parsed) || base
+
+      parsed
+    end
+
+    def fetch_page(offset)
+      url  = "#{@cursor_base}/responseStore/#{@request_id}/results?offset=#{offset}&numRows=#{@page_size}"
+      resp = @http_client.get(url, headers: json_headers)
+
+      raise TransportError, "cursor fetch returned HTTP #{resp.code}" unless resp.code.to_i == 200
+
+      BrokerResponse.from_json(resp.body)
+    end
+
+    def broker_base(address)
+      return address if address.start_with?("http://", "https://")
+      "http://#{address}"
+    end
+
+    def broker_base_from_response(resp)
+      return nil unless resp.broker_host && resp.broker_port
+      "http://#{resp.broker_host}:#{resp.broker_port}"
+    end
+
+    def json_headers
+      { "Content-Type" => "application/json; charset=utf-8" }.merge(@extra_headers)
+    end
+  end
+end

data/lib/pinot/query_result.rb

@@ -0,0 +1,11 @@
+module Pinot
+  QueryResult = Struct.new(:table, :query, :response, :error, keyword_init: true) do
+    def success?
+      error.nil?
+    end
+
+    def error?
+      !error.nil?
+    end
+  end
+end

data/lib/pinot/response.rb

@@ -188,7 +188,11 @@ module Pinot
                 :num_docs_scanned, :num_entries_scanned_in_filter,
                 :num_entries_scanned_post_filter, :total_docs,
                 :time_used_ms, :min_consuming_freshness_time_ms,
-                :num_groups_limit_reached
+                :num_groups_limit_reached,
+                # cursor fields — only present when getCursor=true
+                :request_id, :num_rows_result_set, :offset, :num_rows,
+                :broker_host, :broker_port,
+                :submission_time_ms, :expiration_time_ms
 
     def self.from_json(json_str)
       hash = JSON.parse(json_str)
@@ -215,6 +219,19 @@ module Pinot
       @total_docs = hash["totalDocs"] || 0
       @time_used_ms = hash["timeUsedMs"] || 0
       @min_consuming_freshness_time_ms = hash["minConsumingFreshnessTimeMs"] || 0
+
+      @request_id          = hash["requestId"]
+      @num_rows_result_set = hash["numRowsResultSet"]
+      @offset              = hash["offset"]
+      @num_rows            = hash["numRows"]
+      @broker_host         = hash["brokerHost"]
+      @broker_port         = hash["brokerPort"]
+      @submission_time_ms  = hash["submissionTimeMs"]
+      @expiration_time_ms  = hash["expirationTimeMs"]
+    end
+
+    def cursor?
+      !@request_id.nil?
     end
   end
 end

data/lib/pinot/transport.rb

@@ -40,6 +40,15 @@ module Pinot
       end
     end
 
+    def delete(url, headers: {})
+      uri = URI.parse(url)
+      with_connection(url) do |http|
+        req = Net::HTTP::Delete.new(uri.request_uri)
+        headers.each { |k, v| req[k] = v }
+        http.request(req)
+      end
+    end
+
     def close
       @reaper.kill rescue nil
       @pool_mutex.synchronize do
@@ -183,6 +192,8 @@ module Pinot
     # 250 = ExecutionTimeoutError (server-side), 400 = BrokerTimeoutError.
     TIMEOUT_ERROR_CODES = [250, 400].freeze
 
+    attr_reader :http_client
+
     def initialize(http_client:, extra_headers: {}, timeout_ms: nil, logger: nil,
                    max_retries: 0, retry_interval_ms: 200)
       @http_client = http_client

data/lib/pinot/version.rb

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

data/lib/pinot.rb

@@ -19,6 +19,8 @@ require_relative "pinot/controller_response"
 require_relative "pinot/controller_based_broker_selector"
 require_relative "pinot/transport"
 require_relative "pinot/circuit_breaker"
+require_relative "pinot/query_result"
+require_relative "pinot/paginator"
 require_relative "pinot/connection"
 require_relative "pinot/prepared_statement"
 require_relative "pinot/connection_factory"

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: pinot-client
 version: !ruby/object:Gem::Version
-  version: 1.24.0
+  version: 1.26.0
 platform: ruby
 authors:
 - Xiang Fu
@@ -115,10 +115,12 @@ files:
 - lib/pinot/grpc_transport.rb
 - lib/pinot/instrumentation.rb
 - lib/pinot/logger.rb
+- lib/pinot/paginator.rb
 - lib/pinot/prepared_statement.rb
 - lib/pinot/proto/broker_service.proto
 - lib/pinot/proto/broker_service_pb.rb
 - lib/pinot/proto/broker_service_services_pb.rb
+- lib/pinot/query_result.rb
 - lib/pinot/request.rb
 - lib/pinot/response.rb
 - lib/pinot/simple_broker_selector.rb