bzync-nextsql 0.1.0 → 0.2.0

data/lib/nextsql/pool.rb

@@ -0,0 +1,215 @@
+# frozen_string_literal: true
+
+require "monitor"
+
+module NextSQL
+  # Thread-safe connection pool for NextSQL Ruby driver.
+  #
+  #   pool = NextSQL.create_pool(cfg, min_size: 2, max_size: 10)
+  #   pool.with do |conn|
+  #     res = conn.exec("SELECT * FROM users")
+  #   end
+  #
+  #   # Or directly on the pool:
+  #   pool.exec("INSERT INTO logs VALUES ($1)", ["event"])
+  #   pool.close
+  class ConnectionPool
+    class PooledItem
+      attr_reader :conn
+      attr_accessor :last_used
+
+      def initialize(conn)
+        @conn = conn
+        @last_used = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+      end
+    end
+
+    def initialize(config, min_size: 1, max_size: 10, timeout: 30.0, max_idle: 300.0, statement_cache_size: 50)
+      raise Error.new("invalid_argument", "min_size must be >= 0") if min_size < 0
+      raise Error.new("invalid_argument", "max_size must be >= 1") if max_size < 1
+      raise Error.new("invalid_argument", "max_size must be >= min_size") if max_size < min_size
+      raise Error.new("invalid_argument", "statement_cache_size must be >= 0") if statement_cache_size < 0
+
+      @config = config.dup
+      @config.statement_cache_size = statement_cache_size if @config.respond_to?(:statement_cache_size=)
+      @min_size = min_size
+      @max_size = max_size
+      @timeout = timeout
+      @max_idle = max_idle
+
+      @monitor = Monitor.new
+      @cond = @monitor.new_cond
+      @idle = []
+      @in_use = {}.compare_by_identity
+      @closed = false
+
+      if min_size > 0
+        min_size.times do
+          begin
+            conn = Connection.connect(@config)
+            @idle << PooledItem.new(conn)
+          rescue StandardError
+            close
+            raise
+          end
+        end
+      end
+    end
+
+    def size
+      @monitor.synchronize { @idle.size + @in_use.size }
+    end
+
+    def idle_count
+      @monitor.synchronize { @idle.size }
+    end
+
+    def in_use_count
+      @monitor.synchronize { @in_use.size }
+    end
+
+    def closed?
+      @monitor.synchronize { @closed }
+    end
+
+    # Acquires a connection from the pool, waiting up to +timeout+ seconds if
+    # all connections are currently in use.
+    def acquire(timeout: nil)
+      wait_time = timeout || @timeout
+      deadline = wait_time > 0 ? Process.clock_gettime(Process::CLOCK_MONOTONIC) + wait_time : Float::INFINITY
+
+      @monitor.synchronize do
+        loop do
+          raise Error.new("unavailable", "connection pool is closed") if @closed
+
+          now = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+          while (item = @idle.pop)
+            if @max_idle > 0 && (now - item.last_used) > @max_idle
+              safe_close(item.conn)
+              next
+            end
+            if !alive?(item.conn)
+              safe_close(item.conn)
+              next
+            end
+            @in_use[item.conn] = item
+            return item.conn
+          end
+
+          if (@idle.size + @in_use.size) < @max_size
+            conn = Connection.connect(@config)
+            item = PooledItem.new(conn)
+            @in_use[conn] = item
+            return conn
+          end
+
+          remaining = deadline - Process.clock_gettime(Process::CLOCK_MONOTONIC)
+          if remaining <= 0
+            raise Error.new("resource_exhausted", "connection pool exhausted, wait timeout exceeded")
+          end
+
+          @cond.wait([remaining, 0.5].min)
+        end
+      end
+    end
+
+    # Returns a connection to the pool.
+    def release(conn)
+      @monitor.synchronize do
+        item = @in_use.delete(conn)
+        return if item.nil?
+
+        if @closed
+          safe_close(conn)
+          @cond.signal
+          return
+        end
+
+        unless clean_connection(item)
+          safe_close(conn)
+          @cond.signal
+          return
+        end
+
+        item.last_used = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+        @idle << item
+        @cond.signal
+      end
+    end
+
+    # Checks out a connection for the duration of the block and guarantees its
+    # release even if the block raises.
+    def with(timeout: nil)
+      conn = acquire(timeout: timeout)
+      begin
+        yield conn
+      ensure
+        release(conn)
+      end
+    end
+
+    # Executes a statement on a pooled connection and returns the result.
+    def exec(sql, params = nil)
+      with { |conn| conn.exec(sql, params) }
+    end
+
+    # Queries a statement on a pooled connection, buffering all rows into
+    # a Result before releasing the connection.
+    def query(sql, params = nil)
+      with do |conn|
+        rows = conn.query(sql, params)
+        rows.collect
+      end
+    end
+
+    # Closes all idle and active connections in the pool.
+    def close
+      @monitor.synchronize do
+        return if @closed
+
+        @closed = true
+        @idle.each { |item| safe_close(item.conn) }
+        @idle.clear
+
+        @in_use.each_key { |conn| safe_close(conn) }
+        @in_use.clear
+
+        @cond.broadcast
+      end
+    end
+
+    private
+
+    def safe_close(conn)
+      conn.close
+    rescue StandardError
+      nil
+    end
+
+    def alive?(conn)
+      if conn.instance_variable_defined?(:@sock) && conn.instance_variable_get(:@sock).nil?
+        return false
+      end
+      true
+    rescue StandardError
+      true
+    end
+
+    def clean_connection(item)
+      item.conn.exec("ROLLBACK")
+      true
+    rescue StandardError
+      false
+    end
+  end
+
+  def self.create_pool(config, min_size: 1, max_size: 10, timeout: 30.0, max_idle: 300.0)
+    ConnectionPool.new(
+      config,
+      min_size: min_size,
+      max_size: max_size,
+      timeout: timeout,
+      max_idle: max_idle,
+    )
+  end
+end

data/lib/nextsql/protocol.rb

@@ -53,6 +53,12 @@ module NextSQL
     AUTH_PASSWORD = 1
     AUTH_PASSWORD_KEY = 2
     FLAG_CANCEL = 1
+    # FLAG_PUBLIC_ERROR_CODES asks the server for the stable ERR_* error
+    # taxonomy (docs/error-codes.md). A server that does not implement it
+    # ignores the bit and keeps the NSQL v1 error shape, which decode_error
+    # still reads, so setting it is safe against any server version. The server
+    # echoes it in the hello-ok flags when it was accepted.
+    FLAG_PUBLIC_ERROR_CODES = 2
     FLAG_NULL = 0x01
 
     KIND_UUID = 1
@@ -313,9 +319,16 @@ module NextSQL
     end
 
     def decode_hello_ok(b)
-      raise ProtocolError, "bad hello-ok length" unless b.bytesize == 11
-
-      [u16(b, 0), b.getbyte(2), b.byteslice(3, 8)]
+      raise ProtocolError, "bad hello-ok length" unless [11, 13].include?(b.bytesize)
+
+      flags = 0
+      if b.bytesize == 13
+        flags = u16(b, 11)
+        # The server omits the field when it accepted no capability, so a
+        # present-but-zero field is a second encoding of the v1 hello-ok.
+        raise ProtocolError, "empty hello-ok flags" if flags.zero?
+      end
+      [u16(b, 0), b.getbyte(2), b.byteslice(3, 8), flags]
     end
 
     def encode_query(sql, params)
@@ -946,8 +959,16 @@ module NextSQL
 
     def decode_error(b)
       code, off = read_u16_string(b, 0, MAX_NAME)
-      msg, = read_u16_string(b, off, MAX_NAME)
-      Error.new(code, msg)
+      msg, off = read_u16_string(b, off, MAX_NAME)
+      # Optional trailing field, present only from a server that accepted
+      # FLAG_PUBLIC_ERROR_CODES. Its absence is normal -- an older server
+      # ignores the request bit -- so this must never be required.
+      public_code = ""
+      if off < b.bytesize
+        public_code, = read_u16_string(b, off, MAX_NAME)
+        raise ProtocolError, "empty public error code" if public_code.empty?
+      end
+      Error.new(code, msg, public_code)
     end
 
     def encode_set_read_consistency(mode, max_staleness_ms)

data/lib/nextsql/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module NextSQL
-  VERSION = "0.1.0"
+  VERSION = "0.2.0"
 end

data/lib/nextsql.rb

@@ -25,8 +25,10 @@
 require_relative "nextsql/version"
 require_relative "nextsql/errors"
 require_relative "nextsql/protocol"
+require_relative "nextsql/field_encryption"
 require_relative "nextsql/client"
 require_relative "nextsql/cluster"
+require_relative "nextsql/pool"
 
 module NextSQL
   READ_STRONG = Protocol::READ_STRONG
@@ -65,4 +67,9 @@ module NextSQL
   def self.connect_cluster(cfg)
     Cluster.connect(cfg)
   end
+
+  # Creates a connection pool managing a set of connections to a NextSQL server.
+  def self.create_pool(cfg, **kwargs)
+    ConnectionPool.new(cfg, **kwargs)
+  end
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: bzync-nextsql
 version: !ruby/object:Gem::Version
-  version: 0.1.0
+  version: 0.2.0
 platform: ruby
 authors:
 - Rayan Reynaldo
@@ -10,8 +10,8 @@ cert_chain: []
 date: 1980-01-02 00:00:00.000000000 Z
 dependencies: []
 description: 'Native-protocol client for NextSQL: TLS 1.3, prepared statements, follower-read
-  routing, and NSCE1 client-side field encryption. Keys and passwords are never accepted
-  in a URL.'
+  routing, and NSCE1/NSCE2 client-side field encryption. Keys and passwords are never
+  accepted in a URL.'
 executables: []
 extensions: []
 extra_rdoc_files: []
@@ -22,6 +22,8 @@ files:
 - lib/nextsql/client.rb
 - lib/nextsql/cluster.rb
 - lib/nextsql/errors.rb
+- lib/nextsql/field_encryption.rb
+- lib/nextsql/pool.rb
 - lib/nextsql/protocol.rb
 - lib/nextsql/version.rb
 homepage: https://nextsql.bzync.com/docs/drivers
@@ -47,7 +49,7 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 4.0.20
+rubygems_version: 4.0.23
 specification_version: 4
 summary: Official NextSQL Ruby driver. Speaks the native NSQL v1 protocol.
 test_files: []