pinot-client 1.27.0 → 1.28.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: dfd22c5e76430105c17c3a1279b002084aebb894d4c561913908a4d00038b519
-  data.tar.gz: e668583cdd1afde9b9b35dab28c46b2e707643911a0800cdb98950ac1de2731b
+  metadata.gz: da3d726029acb7167db3fef1a14951e4f4c46beac51f32e88e1e7782eab64fa1
+  data.tar.gz: 4271e56a42a4d269bdf8de68a4c631e8320ec2eec4a36deeb3b635e10e8dc933
 SHA512:
-  metadata.gz: 1aba59c0e39a601dacccfa2f413b9405ad0835e94555272a8f7bf8177b421fd0eac861144fb44d24c69ec41f51834b67a786ed4bd29afb48ae6667671f5ce112
-  data.tar.gz: c81b393e62ffba2390a7c21a52fb0b0a3f0fdfd0421b3a78d9bcbcc772a35f2613e451f58a28cdcb9ecd30c6120aa114f94ae3a80c825594a0cc9baa98a521ca
+  metadata.gz: ef69651b14484fceeec26edca916ede91f04ea3224f24d449e2cf022a3bb847e9ca0d429908daad14e10daa72c9bd4719f4e2400dff06431dca1c6126c90bd7e
+  data.tar.gz: 8ab20c872830e2387245fadebae91e2739b5e30eb959a9769346df4652e0a3e18d673da2e6a3f88ae5753e55e41e7ed21ae709827d920e338e569444ba4a878a

data/lib/pinot/active_support_notifications.rb

@@ -48,7 +48,7 @@ module Pinot
     def self.install!
       return if installed?
 
-      Pinot::Instrumentation.on_query = method(:notify)
+      @listener  = Pinot::Instrumentation.subscribe(method(:notify))
       @installed = true
     end
 
@@ -57,7 +57,8 @@ module Pinot
     end
 
     def self.uninstall!
-      Pinot::Instrumentation.on_query = nil
+      Pinot::Instrumentation.unsubscribe(@listener) if @listener
+      @listener  = nil
       @installed = false
     end
 

data/lib/pinot/instrumentation.rb

@@ -1,32 +1,92 @@
 module Pinot
   # Low-level instrumentation hook that fires after every query executed via
   # Connection#execute_sql. This is the extension point used by
-  # Pinot::ActiveSupportNotifications and any custom observability layer.
+  # Pinot::ActiveSupportNotifications, Pinot::OpenTelemetry, and any custom
+  # observability layer.
   #
-  # Register a single callback:
+  # == Subscribing (multiple listeners supported)
   #
-  #   Pinot::Instrumentation.on_query = ->(event) do
+  #   listener = Pinot::Instrumentation.subscribe(->(event) do
   #     MyMetrics.record(event[:table], event[:duration_ms], event[:success])
+  #   end)
+  #
+  #   # Remove later:
+  #   Pinot::Instrumentation.unsubscribe(listener)
+  #
+  # == Legacy single-callback API (still supported)
+  #
+  #   Pinot::Instrumentation.on_query = ->(event) { ... }
+  #   Pinot::Instrumentation.on_query = nil  # remove
+  #
+  # == Around-execution hook (for OTel and similar span-based tools)
+  #
+  # The `around` hook wraps the entire query execution — the block yields the
+  # query, and the hook is responsible for calling Instrumentation.notify when
+  # done. Only one around hook can be registered at a time.
+  #
+  #   Pinot::Instrumentation.around = ->(table:, query:) do
+  #     MyTracer.in_span("pinot") { yield }
   #   end
   #
-  # The event Hash contains:
+  # == Event Hash keys
+  #
   #   :table        => String  — table name passed to execute_sql
   #   :query        => String  — SQL string
   #   :duration_ms  => Float   — wall-clock time in milliseconds
   #   :success      => Boolean — false when an exception was raised
   #   :error        => Exception or nil — the exception on failure, nil on success
-  #
-  # Only one callback can be registered at a time. Set on_query= to nil to remove it.
   module Instrumentation
+    @listeners = []
+    @around    = nil
+
+    # Add a post-execution listener. Returns the listener so it can be passed
+    # to unsubscribe later.
+    def self.subscribe(listener)
+      @listeners << listener
+      listener
+    end
+
+    # Remove a previously subscribed listener.
+    def self.unsubscribe(listener)
+      @listeners.delete(listener)
+    end
+
+    # Register an around-execution wrapper. Only one wrapper is supported;
+    # the new value replaces any previous one. Set to nil to remove.
+    def self.around=(wrapper)
+      @around = wrapper
+    end
+
+    def self.around
+      @around
+    end
+
+    # Legacy single-callback setter. Replaces all listeners with the given
+    # callback (or clears them when nil).
     def self.on_query=(callback)
-      @on_query = callback
+      @listeners = callback ? [callback] : []
     end
 
+    # Returns the first registered listener (legacy compat).
     def self.on_query
-      @on_query
+      @listeners.first
     end
 
     def self.instrument(table:, query:)
+      if @around
+        @around.call(table: table, query: query) { yield }
+      else
+        _timed_instrument(table: table, query: query) { yield }
+      end
+    end
+
+    # Fire all registered listeners with an event hash. Called by the default
+    # instrument path and by the OTel around wrapper.
+    def self.notify(event)
+      @listeners.each { |l| l.call(event) }
+    end
+
+    private_class_method def self._timed_instrument(table:, query:)
       start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
       result = yield
       duration_ms = (Process.clock_gettime(Process::CLOCK_MONOTONIC) - start) * 1000
@@ -37,9 +97,5 @@ module Pinot
       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/open_telemetry.rb

@@ -0,0 +1,160 @@
+require "pinot"
+
+module Pinot
+  # Opt-in OpenTelemetry bridge.
+  #
+  # == Setup
+  #
+  # Add to an initializer after the opentelemetry SDK is configured:
+  #
+  #   require "pinot/open_telemetry"
+  #   Pinot::OpenTelemetry.install!
+  #
+  # == What it does
+  #
+  # Each call to Connection#execute_sql (and anything built on top of it —
+  # execute_sql_with_params, execute_many, PreparedStatement) creates an OTel
+  # span named "pinot.query" with attributes following the OpenTelemetry
+  # semantic conventions for database spans:
+  #
+  #   db.system        = "pinot"
+  #   db.statement     = "<sql>"        (the full SQL string)
+  #   db.name          = "<table>"      (the Pinot table name)
+  #   db.operation     = "SELECT"       (first token of the SQL)
+  #
+  # On failure the span is marked with error status and the exception is
+  # recorded on the span.
+  #
+  # == Trace-context propagation
+  #
+  # When installed, every outbound HTTP request to a broker is injected with
+  # W3C Trace Context headers (traceparent / tracestate) so distributed traces
+  # flow through Pinot. This relies on OpenTelemetry.propagation being
+  # configured (the default SDK sets this up automatically).
+  #
+  # == Feature flag
+  #
+  # The bridge can be toggled at runtime without reinstalling:
+  #
+  #   Pinot::OpenTelemetry.enabled = false   # pause tracing (e.g. in tests)
+  #   Pinot::OpenTelemetry.enabled = true    # resume
+  #   Pinot::OpenTelemetry.enabled?          # => true / false
+  #
+  # == Lifecycle
+  #
+  #   Pinot::OpenTelemetry.install!    # idempotent
+  #   Pinot::OpenTelemetry.installed?  # => true
+  #   Pinot::OpenTelemetry.uninstall!  # removes hooks; leaves transport unpatched
+  #
+  # Note: this gem does NOT depend on opentelemetry-api or opentelemetry-sdk.
+  # Both must be present and initialized before install! is called.
+  module OpenTelemetry
+    SPAN_NAME    = "pinot.query"
+    DB_SYSTEM    = "pinot"
+    TRACER_NAME  = "pinot-client"
+
+    @installed = false
+    @enabled   = true
+
+    # Enable or disable tracing at runtime without uninstalling.
+    def self.enabled=(val)
+      @enabled = val ? true : false
+    end
+
+    def self.enabled?
+      @enabled
+    end
+
+    def self.install!
+      return if installed?
+
+      _install_around_hook
+      _patch_transport
+      @installed = true
+    end
+
+    def self.installed?
+      @installed
+    end
+
+    def self.uninstall!
+      ::Pinot::Instrumentation.around = nil
+      @installed = false
+      # Note: JsonHttpTransport prepend is permanent once applied (Ruby limitation).
+      # Disable the propagator by unsetting the flag — it no-ops when disabled.
+    end
+
+    # -------------------------------------------------------------------------
+    # Internal helpers
+    # -------------------------------------------------------------------------
+
+    def self._install_around_hook
+      ::Pinot::Instrumentation.around = method(:_around)
+    end
+    private_class_method :_install_around_hook
+
+    def self._around(table:, query:)
+      unless @enabled
+        # Bypass tracing; still time and notify listeners.
+        return ::Pinot::Instrumentation.send(:_timed_instrument, table: table, query: query) { yield }
+      end
+
+      tracer = ::OpenTelemetry.tracer_provider.tracer(TRACER_NAME, ::Pinot::VERSION)
+      attrs  = _span_attributes(table, query)
+
+      tracer.in_span(SPAN_NAME, attributes: attrs, kind: :client) do |span|
+        start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+        begin
+          result = yield
+          duration_ms = (Process.clock_gettime(Process::CLOCK_MONOTONIC) - start) * 1000
+          ::Pinot::Instrumentation.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
+          span.record_exception(e)
+          span.status = ::OpenTelemetry::Trace::Status.error(e.message)
+          ::Pinot::Instrumentation.notify(
+            table: table, query: query, duration_ms: duration_ms, success: false, error: e
+          )
+          raise
+        end
+      end
+    end
+    private_class_method :_around
+
+    def self._span_attributes(table, query)
+      attrs = {
+        "db.system"    => DB_SYSTEM,
+        "db.statement" => query,
+        "db.name"      => table.to_s
+      }
+      op = query.to_s.lstrip.split(/\s+/, 2).first&.upcase
+      attrs["db.operation"] = op if op && !op.empty?
+      attrs
+    end
+    private_class_method :_span_attributes
+
+    # Patch JsonHttpTransport to inject W3C trace-context headers into every
+    # outbound HTTP request. Applied once via Module#prepend.
+    def self._patch_transport
+      return if ::Pinot::JsonHttpTransport.ancestors.include?(TraceContextInjector)
+      ::Pinot::JsonHttpTransport.prepend(TraceContextInjector)
+    end
+    private_class_method :_patch_transport
+
+    # Prepended into JsonHttpTransport. Injects traceparent/tracestate into
+    # request headers when the bridge is enabled and a current span exists.
+    module TraceContextInjector
+      def execute(broker_address, request, extra_request_headers: {})
+        otel = ::Pinot::OpenTelemetry
+        return super unless otel.installed? && otel.enabled?
+
+        carrier = {}
+        ::OpenTelemetry.propagation.inject(carrier)
+        super(broker_address, request, extra_request_headers: extra_request_headers.merge(carrier))
+      end
+    end
+  end
+end

data/lib/pinot/version.rb

@@ -1,3 +1,3 @@
 module Pinot
-  VERSION = "1.27.0"
+  VERSION = "1.28.0"
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: pinot-client
 version: !ruby/object:Gem::Version
-  version: 1.27.0
+  version: 1.28.0
 platform: ruby
 authors:
 - Xiang Fu
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2026-04-19 00:00:00.000000000 Z
+date: 2026-04-20 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: logger
@@ -116,6 +116,7 @@ files:
 - lib/pinot/grpc_transport.rb
 - lib/pinot/instrumentation.rb
 - lib/pinot/logger.rb
+- lib/pinot/open_telemetry.rb
 - lib/pinot/paginator.rb
 - lib/pinot/prepared_statement.rb
 - lib/pinot/proto/broker_service.proto