questdb 0.1.0-aarch64-linux-gnu

data/lib/questdb/query.rb

@@ -0,0 +1,429 @@
+# frozen_string_literal: true
+
+#      ___                  _   ____  ____
+#     / _ \ _   _  ___  ___| |_|  _ \| __ )
+#    | | | | | | |/ _ \/ __| __| | | |  _ \
+#    | |_| | |_| |  __/\__ \ |_| |_| | |_) |
+#     \__\_\\__,_|\___||___/\__|____/|____/
+#
+#   Copyright (c) 2014-2019 Appsicle
+#   Copyright (c) 2019-2026 QuestDB
+#
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+#
+
+module QuestDB
+  Column = Struct.new(:name, :kind, :metadata, keyword_init: true) do
+    def initialize(name:, kind:, metadata: {})
+      super(name: String(name).dup.freeze, kind: kind.to_sym, metadata: metadata.dup.freeze)
+      freeze
+    end
+  end
+
+  class Row
+    attr_reader :names, :values
+    def initialize(names, values, name_index = nil)
+      @names = names
+      @values = values.freeze
+      @name_index = name_index || begin
+        index = {}
+        names.each_with_index { |name, offset| index[name] ||= offset }
+        index.freeze
+      end
+      freeze
+    end
+    def [](key)
+      return values.fetch(key) if key.is_a?(Integer)
+      index = @name_index[row_key(key)]
+      index ? values[index] : nil
+    end
+    def fetch(key, *defaults, &block)
+      return values.fetch(key, *defaults, &block) if key.is_a?(Integer)
+      raise ArgumentError, "wrong number of default values" if defaults.length > 1
+      index = @name_index[row_key(key)]
+      return values[index] if index
+      return yield(key) if block
+      return defaults.first unless defaults.empty?
+      raise KeyError, "column not found: #{key.inspect}"
+    end
+    def to_a = values.dup
+    def to_h
+      names.each_with_index.each_with_object({}) { |(name, index), result| result[name] = values[index] }
+    end
+    def deconstruct = values
+
+    private
+
+    def row_key(key) = key.is_a?(Symbol) ? key.name : key.to_s
+  end
+
+  class ResultBatch
+    attr_reader :columns, :rows
+    def self.columns_from_native(native_columns)
+      native_columns.map do |name, kind, metadata|
+        Column.new(name: name, kind: kind, metadata: metadata || {})
+      end.freeze
+    end
+
+    def initialize(native_columns, native_rows)
+      @columns = self.class.columns_from_native(native_columns)
+      names = @columns.map(&:name).freeze
+      name_index = {}
+      names.each_with_index { |name, index| name_index[name] ||= index }
+      name_index.freeze
+      @rows = native_rows.map { |values| Row.new(names, values, name_index) }.freeze
+      freeze
+    end
+    def row_count = rows.length
+    def column_count = columns.length
+    def empty? = rows.empty?
+  end
+
+  class QueryResult
+    include Enumerable
+    attr_reader :columns
+
+    def initialize(reader, native)
+      @reader = reader
+      @native = native
+      @owner_thread = Thread.current
+      @columns = nil
+      @closed = false
+      @consumed = false
+      @cancelled = false
+      @close_reader = false
+    end
+
+    def each_batch
+      return enum_for(__method__) unless block_given?
+      each_native_batch do |native_columns, native_rows|
+        batch = ResultBatch.new(native_columns, native_rows)
+        @columns ||= batch.columns
+        yield batch
+        true
+      end
+      self
+    end
+
+    def each
+      return enum_for(__method__) unless block_given?
+      each_batch { |batch| batch.rows.each { |row| yield row } }
+      self
+    end
+
+    # Lower-allocation streaming for callers that only need positional values.
+    # Arrays are owned copies and frozen before they are yielded; unlike Row,
+    # they do not support name-based lookup.
+    def each_values
+      return enum_for(__method__) unless block_given?
+      each_native_batch do |native_columns, native_rows|
+        @columns ||= ResultBatch.columns_from_native(native_columns)
+        native_rows.each do |values|
+          values.freeze
+          yield values
+        end
+        !native_rows.empty?
+      end
+      self
+    end
+
+    # Whole-result materialization can safely authorize native replay: if a
+    # failover reset occurs, discard the partial prefix before accepting the
+    # replayed batch. Streaming iteration cannot retract yielded rows and raises
+    # instead (see #each_batch).
+    def to_a
+      start_consumption!
+      rows = []
+      reset_sequence = 0
+      begin
+        while (payload = @native.next_batch)
+          native_columns, native_rows, sequence = payload
+          if sequence != reset_sequence
+            rows.clear
+            @columns = nil
+            reset_sequence = sequence
+          end
+          batch = ResultBatch.new(native_columns, native_rows)
+          @columns ||= batch.columns
+          rows.concat(batch.rows)
+        end
+      ensure
+        close_preserving_primary_error
+      end
+      rows
+    end
+
+    def cancel
+      return self if @closed || @cancelled
+      ensure_thread!
+      begin
+        @native.cancel
+      rescue Exception => error
+        begin
+          close
+        rescue Exception
+          # Preserve the cancellation error; native close/drop is best-effort.
+        end
+        raise error
+      end
+      @cancelled = true
+      self
+    end
+
+    def close
+      return self if @closed
+      ensure_thread!
+      reusable = @native.close
+      @closed = true
+      @reader.send(:result_released, self, reusable: reusable)
+      @reader.close if @close_reader && !@reader.closed?
+      self
+    end
+
+    def closed? = @closed || @native.closed?
+    def consumed? = @consumed
+
+    def open
+      raise ArgumentError, "a block is required" unless block_given?
+      begin
+        yield self
+      ensure
+        close
+      end
+    end
+
+    def initialize_copy(*) = raise(TypeError, "QuestDB query results cannot be copied")
+
+    protected
+
+    def close_reader_on_close!
+      @close_reader = true
+      self
+    end
+
+    private
+
+    def each_native_batch
+      start_consumption!
+      reset_sequence = 0
+      yielded = false
+      begin
+        while (payload = @native.next_batch)
+          native_columns, native_rows, sequence = payload
+          if yielded && sequence != reset_sequence
+            raise Ingress::IngressError.new(
+              "query failed over after rows were delivered; streaming replay would duplicate results",
+              code: :failover_would_duplicate,
+              in_doubt: true
+            )
+          end
+          reset_sequence = sequence
+          delivered = yield native_columns, native_rows
+          yielded ||= delivered
+        end
+      ensure
+        close_preserving_primary_error
+      end
+    end
+
+    def start_consumption!
+      ensure_thread!
+      raise Ingress::IngressError.new("query result is single-use", code: :invalid_api_call) if @consumed
+      raise Ingress::IngressError.new("query result is closed", code: :invalid_api_call) if closed?
+      @consumed = true
+    end
+
+    def close_preserving_primary_error
+      primary = $!
+      begin
+        close
+      rescue Exception
+        raise unless primary
+      end
+    end
+
+    def ensure_thread!
+      return if Thread.current.equal?(@owner_thread)
+      raise Ingress::IngressError.new("query result must be used on the thread that created it", code: :invalid_api_call)
+    end
+  end
+
+  class PooledReader
+    def initialize(client, native)
+      @client = client
+      @native = native
+      @owner_thread = Thread.current
+      @active_result = nil
+      @closed = false
+    end
+
+    def query(sql, binds: nil, reset_symbol_dict: true)
+      ensure_open!
+      raise Ingress::IngressError.new("reader already has an active result", code: :invalid_api_call) if @active_result && !@active_result.closed?
+      descriptors = Bindings.normalize(binds)
+      result = QueryResult.new(self, @native.query(String(sql), descriptors, !!reset_symbol_dict))
+      @active_result = result
+      return result unless block_given?
+      begin
+        yield result
+      ensure
+        result.close
+      end
+    end
+
+    def execute(sql, binds: nil)
+      query(sql, binds: binds) { |result| result.each_batch { |_batch| } }
+      nil
+    end
+
+    def server_info
+      ensure_open!
+      raise Ingress::IngressError.new("reader has an active result", code: :invalid_api_call) if @active_result && !@active_result.closed?
+      ServerInfo.from_native(@native.server_info)
+    end
+
+    def close
+      return self if @closed
+      ensure_thread!
+      @active_result.close if @active_result && !@active_result.closed?
+      return self if @closed
+      @native.close(false)
+      @closed = true
+      self
+    end
+
+    def closed? = @closed || @native.closed?
+
+    def open
+      raise ArgumentError, "a block is required" unless block_given?
+      begin
+        yield self
+      ensure
+        close
+      end
+    end
+
+    def initialize_copy(*) = raise(TypeError, "QuestDB pooled readers cannot be copied")
+
+    protected
+
+    def result_released(result, reusable: true)
+      @active_result = nil if @active_result.equal?(result)
+      unless reusable || closed?
+        @native.close(true)
+        @closed = true
+      end
+    end
+
+    private
+
+    def ensure_thread!
+      return if Thread.current.equal?(@owner_thread)
+      raise Ingress::IngressError.new("pooled reader must be used on the thread that borrowed it", code: :invalid_api_call)
+    end
+
+    def ensure_open!
+      ensure_thread!
+      raise Ingress::IngressError.new("pooled reader is closed", code: :invalid_api_call) if closed?
+    end
+  end
+
+  ServerInfo = Struct.new(:role, :raw_role, :epoch, :capabilities, :server_wall_ns, :cluster_id, :node_id, :zone_id, keyword_init: true) do
+    ROLES = %i[standalone primary replica primary_catchup other].freeze
+    def self.from_native(values)
+      new(role: values[0], raw_role: values[1], epoch: values[2], capabilities: values[3],
+          server_wall_ns: values[4], cluster_id: values[5], node_id: values[6], zone_id: values[7])
+    end
+    def initialize(role:, raw_role:, epoch:, capabilities:, server_wall_ns:, cluster_id:, node_id:, zone_id:)
+      super(role: role.to_sym, raw_role: Integer(raw_role), epoch: Integer(epoch), capabilities: Integer(capabilities),
+            server_wall_ns: Integer(server_wall_ns), cluster_id: String(cluster_id).dup.freeze,
+            node_id: String(node_id).dup.freeze, zone_id: zone_id&.then { |v| String(v).dup.freeze })
+      freeze
+    end
+  end
+
+  module Bindings
+    KIND_CODES = {
+      boolean: 0x01, byte: 0x02, short: 0x03, int: 0x04, long: 0x05,
+      float: 0x06, double: 0x07, symbol: 0x09, timestamp: 0x0a,
+      date: 0x0b, uuid: 0x0c, long256: 0x0d, geohash: 0x0e,
+      varchar: 0x0f, timestamp_nanos: 0x10, decimal64: 0x13,
+      decimal128: 0x14, decimal256: 0x15, char: 0x16, binary: 0x17, ipv4: 0x18
+    }.freeze
+    module_function
+
+    def normalize(binds)
+      return [] if binds.nil?
+      raise TypeError, "binds must be an Array" unless binds.is_a?(Array)
+      binds.each_with_index.map { |value, index| normalize_one(value, index + 1) }
+    end
+
+    def normalize_one(value, position)
+      case value
+      when nil then [:null, KIND_CODES[:varchar], nil, nil]
+      when true, false then [:bool, value]
+      when Integer
+        range_i64(value, position)
+        [:i64, value]
+      when Float then [:f64, value]
+      when String then [:varchar, utf8(value)]
+      when Time then [:ts_micros, TimestampMicros.from_time(value).value]
+      when TimestampNanos then [:ts_nanos, value.value]
+      when TimestampMicros then [:ts_micros, value.value]
+      when DateMillis then [:date, value.value]
+      when Char then [:char, value.value]
+      when UUID then [:uuid, value.bytes]
+      when Binary then [:binary, value.bytes]
+      when IPv4 then [:ipv4, value.value]
+      when IPAddr then [:ipv4, IPv4.new(value).value]
+      when Geohash then [:geohash, value.value, value.precision_bits]
+      when Long256 then [:long256, value.bytes]
+      when DecimalValue then decimal_descriptor(value)
+      when BigDecimal then decimal_descriptor(DecimalValue.from_bigdecimal(value))
+      when TypedNull
+        code = KIND_CODES.fetch(value.kind)
+        [:null, code, value.scale, value.precision_bits]
+      else
+        raise TypeError, "unsupported bind $#{position}: #{value.class}"
+      end
+    end
+
+    def decimal_descriptor(value)
+      mantissa = value.mantissa
+      scale = value.scale
+      if mantissa.between?(-(1 << 63), (1 << 63) - 1)
+        [:decimal64, mantissa, scale]
+      elsif mantissa.between?(-(1 << 127), (1 << 127) - 1)
+        unsigned = mantissa % (1 << 128)
+        high = unsigned >> 64
+        high -= 1 << 64 if high >= 1 << 63
+        [:decimal128, unsigned & 0xffff_ffff_ffff_ffff, high, scale]
+      elsif mantissa.between?(-(1 << 255), (1 << 255) - 1)
+        unsigned = mantissa % (1 << 256)
+        [:decimal256, [unsigned.to_s(16).rjust(64, "0")].pack("H*").reverse, scale]
+      else
+        raise RangeError, "decimal mantissa exceeds signed 256-bit range"
+      end
+    end
+
+    def range_i64(value, position)
+      raise RangeError, "bind $#{position} is outside signed 64-bit range" unless value.between?(-(1 << 63), (1 << 63) - 1)
+    end
+
+    def utf8(value)
+      value.encode(Encoding::UTF_8)
+    rescue EncodingError => error
+      raise Ingress::IngressError.new(error.message, code: :invalid_utf8), cause: error
+    end
+  end
+end