pinot-client 1.3.0 → 1.5.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: f6d0692d6f1de1178eb0c27cb2f69e3f6575362b1163e58032538b2b003e13cb
-  data.tar.gz: d9ee2eb20b541cfe54ad640b223425ed16697e869b031349d021ef89c8bab5af
+  metadata.gz: bb872f6cbcc1f34f16a68006eac8acea39229bdc9bd90f9f44a3369ed9c9fb59
+  data.tar.gz: b86cde24abaa6343f9984b3106830e9f00e25656050a4f577629942c3e0067c2
 SHA512:
-  metadata.gz: d66f070cd740df59a8351724b954b902633052e0a94c0cffcc377cba887329709955b0c93075f5f90ea820abe453aa716b4f6eae80600ce90015eda669c33dd0
-  data.tar.gz: '03345279b3e4ae8621fa75a4222dce2988cf61255663b659c7307296a9506680f29f804938ec1b5d28326bf15e4e5c4cf5984e71233b2e0fcf255f6123f0fa36'
+  metadata.gz: afc088a4df6eaa77e6141fa051d7f2e1f3afa1aa8d75555ed2291dfc84fe159a9538ad86621d452bc510066576a331d6e23135aa93529f672bf912ba46ed5b0c
+  data.tar.gz: b34bcb378ab7a6bca383545b79e090273ed0b67d4e5356699461b02bbbaaafd7e7580073e492bbb7f43cf988c77ba500e361e1954c53a990aab5be6d785cae4b

data/README.md

@@ -277,6 +277,10 @@ Environment variables:
 | `BROKER_HOST` | `127.0.0.1` | Pinot broker hostname   |
 | `BROKER_PORT` | `8000`      | Pinot broker HTTP port  |
 
+## Changelog
+
+See [CHANGELOG.md](CHANGELOG.md) for a full history of releases.
+
 ## License
 
 Apache License 2.0 — see [LICENSE](LICENSE).

data/lib/pinot/config.rb

@@ -1,4 +1,14 @@
 module Pinot
+  class ZookeeperConfig
+    attr_accessor :zk_path,           # "host:port" or "host1:port,host2:port/chroot"
+                  :session_timeout_ms # default 30000
+
+    def initialize(zk_path:, session_timeout_ms: 30_000)
+      @zk_path            = zk_path
+      @session_timeout_ms = session_timeout_ms
+    end
+  end
+
   class ControllerConfig
     attr_accessor :controller_address, :update_freq_ms, :extra_controller_api_headers
 
@@ -10,28 +20,55 @@ module Pinot
   end
 
   class ClientConfig
-    attr_accessor :broker_list, :http_timeout, :extra_http_header,
+    attr_accessor :broker_list, :http_timeout, :query_timeout_ms, :extra_http_header,
                   :use_multistage_engine, :controller_config, :logger, :tls_config,
-                  :grpc_config
+                  :grpc_config, :zookeeper_config
 
     def initialize(
       broker_list: [],
       http_timeout: nil,
+      query_timeout_ms: nil,
       extra_http_header: {},
       use_multistage_engine: false,
       controller_config: nil,
       logger: nil,
       tls_config: nil,
-      grpc_config: nil
+      grpc_config: nil,
+      zookeeper_config: nil
     )
       @broker_list = broker_list
       @http_timeout = http_timeout
+      @query_timeout_ms = query_timeout_ms
       @extra_http_header = extra_http_header
       @use_multistage_engine = use_multistage_engine
       @controller_config = controller_config
       @logger = logger
       @tls_config = tls_config
       @grpc_config = grpc_config
+      @zookeeper_config = zookeeper_config
+    end
+
+    def validate!
+      sources = [
+        !broker_list.empty?,
+        !controller_config.nil?,
+        !zookeeper_config.nil?,
+        !grpc_config.nil?
+      ].count(true)
+
+      if sources == 0
+        raise ConfigurationError, "ClientConfig requires at least one of: broker_list, controller_config, zookeeper_config, or grpc_config"
+      end
+
+      if !http_timeout.nil? && http_timeout <= 0
+        raise ConfigurationError, "http_timeout must be positive, got: #{http_timeout}"
+      end
+
+      if !query_timeout_ms.nil? && query_timeout_ms <= 0
+        raise ConfigurationError, "query_timeout_ms must be positive, got: #{query_timeout_ms}"
+      end
+
+      self
     end
   end
 end

data/lib/pinot/connection.rb

@@ -23,9 +23,11 @@ module Pinot
     end
 
     def execute_sql(table, query)
-      logger.debug "Executing SQL on table=#{table}: #{query}"
-      broker = @broker_selector.select_broker(table)
-      @transport.execute(broker, build_request(query))
+      Pinot::Instrumentation.instrument(table: table, query: query) do
+        logger.debug "Executing SQL on table=#{table}: #{query}"
+        broker = @broker_selector.select_broker(table)
+        @transport.execute(broker, build_request(query))
+      end
     rescue => e
       raise "unable to execute SQL on table #{table}: #{e.message}"
     end

data/lib/pinot/connection_factory.rb

@@ -12,6 +12,8 @@ module Pinot
   end
 
   def self.from_config(config, http_client: nil)
+    config.validate!
+
     if config.grpc_config
       transport = GrpcTransport.new(config.grpc_config)
       selector  = SimpleBrokerSelector.new(config.grpc_config.broker_list)
@@ -27,6 +29,25 @@ module Pinot
       return conn
     end
 
+    if config.zookeeper_config
+      selector = ZookeeperBrokerSelector.new(zk_path: config.zookeeper_config.zk_path)
+      selector.init
+
+      inner = http_client || HttpClient.new(timeout: config.http_timeout, tls_config: config.tls_config)
+      transport = JsonHttpTransport.new(
+        http_client: inner,
+        extra_headers: config.extra_http_header || {},
+        logger: config.logger
+      )
+
+      return Connection.new(
+        transport: transport,
+        broker_selector: selector,
+        use_multistage_engine: config.use_multistage_engine || false,
+        logger: config.logger
+      )
+    end
+
     inner = http_client || HttpClient.new(timeout: config.http_timeout, tls_config: config.tls_config)
 
     transport = JsonHttpTransport.new(

data/lib/pinot/instrumentation.rb

@@ -0,0 +1,36 @@
+module Pinot
+  module Instrumentation
+    # Called around every query execution.
+    # Implement by setting Pinot::Instrumentation.on_query = proc { |event| ... }
+    # event is a Hash:
+    #   :table        => String
+    #   :query        => String
+    #   :duration_ms  => Float
+    #   :success      => Boolean
+    #   :error        => Exception or nil
+
+    def self.on_query=(callback)
+      @on_query = callback
+    end
+
+    def self.on_query
+      @on_query
+    end
+
+    def self.instrument(table:, query:)
+      start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+      result = yield
+      duration_ms = (Process.clock_gettime(Process::CLOCK_MONOTONIC) - start) * 1000
+      notify(table: table, query: query, duration_ms: duration_ms, success: true, error: nil)
+      result
+    rescue => e
+      duration_ms = (Process.clock_gettime(Process::CLOCK_MONOTONIC) - start) * 1000
+      notify(table: table, query: query, duration_ms: duration_ms, success: false, error: e)
+      raise
+    end
+
+    def self.notify(event)
+      @on_query&.call(event)
+    end
+  end
+end

data/lib/pinot/version.rb

@@ -1,3 +1,3 @@
 module Pinot
-  VERSION = "1.3.0"
+  VERSION = "1.5.0"
 end

data/lib/pinot/zookeeper_broker_selector.rb

@@ -0,0 +1,80 @@
+require "json"
+require "set"
+require_relative "table_aware_broker_selector"
+require_relative "errors"
+
+module Pinot
+  class ZookeeperBrokerSelector < TableAwareBrokerSelector
+    # ZK path where Pinot stores broker external view
+    BROKER_EXTERNAL_VIEW_PATH = "/EXTERNALVIEW/brokerResource"
+
+    def initialize(zk_path:, zk_client: nil)
+      super()
+      @zk_path   = zk_path      # e.g. "localhost:2181"
+      @zk_client = zk_client    # injectable for testing
+    end
+
+    def init
+      @zk = @zk_client || build_zk_client
+      fetch_and_update
+      setup_watcher
+    end
+
+    private
+
+    def build_zk_client
+      begin
+        require "zk"
+      rescue LoadError
+        raise ConfigurationError, "The 'zk' gem is required to use ZookeeperBrokerSelector. Add it to your Gemfile: gem \"zk\""
+      end
+      ZK.new(@zk_path)
+    end
+
+    def fetch_and_update
+      data, _stat = @zk.get(BROKER_EXTERNAL_VIEW_PATH)
+      parsed = JSON.parse(data)
+      all_brokers, table_map = parse_external_view(parsed)
+      update_broker_data(all_brokers, table_map)
+    end
+
+    def setup_watcher
+      @zk.register(BROKER_EXTERNAL_VIEW_PATH) do |event|
+        fetch_and_update rescue nil
+        setup_watcher  # re-register watch after each trigger
+      end
+      # Set initial watch
+      @zk.exists?(BROKER_EXTERNAL_VIEW_PATH, watch: true)
+    end
+
+    def parse_external_view(data)
+      # Pinot external view format:
+      # { "mapFields": { "tableName_OFFLINE": { "Broker_host_port": "ONLINE" } } }
+      map_fields = data["mapFields"] || {}
+
+      all_brokers = Set.new
+      table_map   = {}
+
+      map_fields.each do |raw_table, broker_map|
+        # Strip _OFFLINE / _REALTIME suffix from the table key, matching Go's extractTableName
+        table_name = raw_table.sub(/_OFFLINE\z/, "").sub(/_REALTIME\z/, "")
+        brokers = []
+        broker_map.each do |broker_key, state|
+          next unless state == "ONLINE"
+          # Broker key format: Broker_<hostname>_<port>
+          # Use the last segment as port and the second-to-last as host
+          parts = broker_key.split("_")
+          next if parts.length < 2
+          port = parts.last
+          next unless port =~ /\A\d+\z/
+          host = parts[-2]
+          brokers << "#{host}:#{port}"
+        end
+        table_map[table_name] = brokers
+        all_brokers.merge(brokers)
+      end
+
+      [all_brokers.to_a, table_map]
+    end
+  end
+end

data/lib/pinot.rb

@@ -5,6 +5,7 @@ require "bigdecimal"
 require "securerandom"
 
 require_relative "pinot/errors"
+require_relative "pinot/instrumentation"
 require_relative "pinot/version"
 require_relative "pinot/logger"
 require_relative "pinot/config"
@@ -27,3 +28,8 @@ begin
 rescue LoadError
   # grpc gem not available; GrpcTransport disabled
 end
+
+begin
+  require_relative "pinot/zookeeper_broker_selector"
+rescue LoadError
+end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: pinot-client
 version: !ruby/object:Gem::Version
-  version: 1.3.0
+  version: 1.5.0
 platform: ruby
 authors:
 - Xiang Fu
@@ -84,6 +84,7 @@ files:
 - lib/pinot/errors.rb
 - lib/pinot/grpc_config.rb
 - lib/pinot/grpc_transport.rb
+- lib/pinot/instrumentation.rb
 - lib/pinot/logger.rb
 - lib/pinot/prepared_statement.rb
 - lib/pinot/proto/broker_service.proto
@@ -96,6 +97,7 @@ files:
 - lib/pinot/tls_config.rb
 - lib/pinot/transport.rb
 - lib/pinot/version.rb
+- lib/pinot/zookeeper_broker_selector.rb
 homepage: https://github.com/startreedata/ruby-pinot-client
 licenses:
 - Apache-2.0