pinot-client 1.25.0 → 1.26.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 1f76e9f295ba992d3a6c2ec7a4dc738230c8676bda0f64bafb3aa85af1530c34
-  data.tar.gz: 4ea8dfd7e0c530f38c860928c18c1facc76d5fabb7667cf7b802864629365dd1
+  metadata.gz: 0c6d47938639a210b77881492e1155d26728be12f987b9721712f16503183ec3
+  data.tar.gz: 9e14c467edbba3ad7fff1e20f6b288ac810895c5069549cd33f08e2972b8d739
 SHA512:
-  metadata.gz: 2a07d7e620bf3adbbcd2f809c37854a1c96da7e1bbb416cbe994d6ce7ea5d17620839aad66b2059e5f1f0de683cd5a405868d1729b5da7dbba19ff820c1ac704
-  data.tar.gz: 2da45bfbb16a63ed13c72fdec7c2684511d27caac9a211155135a3697c88cbf36686a7ded2d569859cb751e8dac59a04b7d91aaf06b3d06136461a217166924e
+  metadata.gz: 73ec7a9ae84e96f71618f430ed542527fb9f9e16f64b1d0d78a1a8bcb0ba1950f6d7117072f0318ffcc08fdb0df2143be268667d6b96e4914d52fc4f0448c0fa
+  data.tar.gz: 9fe83c991b5ef549ce915dbe4af55b50143d17d6d9660915d5fd6308798541f81068e557317b8b88af167b55a63ea7781ef14dfff428cd51b671df7d2a0267b2

data/lib/pinot/connection.rb

@@ -79,6 +79,17 @@ module Pinot
       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?

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/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.25.0"
+  VERSION = "1.26.0"
 end

data/lib/pinot.rb

@@ -20,6 +20,7 @@ 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.25.0
+  version: 1.26.0
 platform: ruby
 authors:
 - Xiang Fu
@@ -115,6 +115,7 @@ 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