pg_pipeline 0.3.3 → 0.3.4

data/lib/pg_pipeline/request.rb

@@ -1,6 +1,6 @@
 # frozen_string_literal: true
 
-require_relative "errors"
+require_relative 'errors'
 
 module PgPipeline
   class Request
@@ -119,7 +119,7 @@ module PgPipeline
 
     def snapshot_params(params)
       return EMPTY_PARAMS if params.nil?
-      raise ArgumentError, "params must be an Array" unless params.is_a?(Array)
+      raise ArgumentError, 'params must be an Array' unless params.is_a?(Array)
       return EMPTY_PARAMS if params.empty?
 
       return params if params.frozen? && immutable_values?(params)
@@ -160,24 +160,23 @@ module PgPipeline
 
     def snapshot_name(name)
       value = name.to_s
-      raise ArgumentError, "statement_name must not be empty" if value.empty?
+      raise ArgumentError, 'statement_name must not be empty' if value.empty?
 
       value.frozen? ? value : value.dup.freeze
     end
 
     def snapshot_param_types(param_types)
       return nil if param_types.nil?
-      raise ArgumentError, "param_types must be an Array or nil" unless param_types.is_a?(Array)
+      raise ArgumentError, 'param_types must be an Array or nil' unless param_types.is_a?(Array)
 
       param_types.map { |oid| oid.nil? ? nil : Integer(oid) }.freeze
     rescue ArgumentError, TypeError
-      raise ArgumentError, "param_types must contain only integer OIDs or nil"
+      raise ArgumentError, 'param_types must contain only integer OIDs or nil'
     end
 
     def transition!(req, from, to)
-      unless req.state == from
-        raise ProtocolError, "invalid request transition #{req.state.inspect} -> #{to.inspect}"
-      end
+      state = req.state
+      raise ProtocolError, "invalid request transition #{state.inspect} -> #{to.inspect}" unless state == from
 
       req.state = to
     end
@@ -205,14 +204,14 @@ module PgPipeline
     end
 
     def query_boundary!(req)
-      raise ProtocolError, "query boundary before query result" unless req.result_seen
-      raise ProtocolError, "duplicate query boundary" if req.query_boundary_seen
+      raise ProtocolError, 'query boundary before query result' unless req.result_seen
+      raise ProtocolError, 'duplicate query boundary' if req.query_boundary_seen
 
       req.query_boundary_seen = true
     end
 
     def finish!(req)
-      raise ProtocolError, "sync arrived before query boundary" unless req.query_boundary_seen
+      raise ProtocolError, 'sync arrived before query boundary' unless req.query_boundary_seen
       return if req.settled
 
       req.settled = true
@@ -237,21 +236,23 @@ module PgPipeline
       req.cancelled = true
       clear_result(req.result)
       req.result = nil
-      req.error.clear_result! if req.error.respond_to?(:clear_result!)
+      error = req.error
+      error.clear_result! if error.respond_to?(:clear_result!)
     end
 
     def wait(req)
       park_waiter(req) unless req.settled
-      raise req.error if req.error
+      error = req.error
+      raise error if error
 
       req.result
     end
 
     def park_waiter(req)
-      raise ProtocolError, "request already has a waiter" if req.waiter
+      raise ProtocolError, 'request already has a waiter' if req.waiter
 
       scheduler = Fiber.scheduler
-      raise Error, "request wait requires an active Fiber scheduler" unless scheduler
+      raise Error, 'request wait requires an active Fiber scheduler' unless scheduler
 
       waiter = Fiber.current
       req.waiter = waiter
@@ -276,8 +277,8 @@ module PgPipeline
     end
 
     def assert_result_slot!(req)
-      raise ProtocolError, "result arrived after query boundary" if req.query_boundary_seen
-      raise ProtocolError, "multiple results for one pipeline unit" if req.result_seen
+      raise ProtocolError, 'result arrived after query boundary' if req.query_boundary_seen
+      raise ProtocolError, 'multiple results for one pipeline unit' if req.result_seen
     end
 
     def clear_result(result)

data/lib/pg_pipeline/runtime/notification.rb

@@ -28,7 +28,8 @@ module PgPipeline
       end
 
       def signal_all
-        pending, @waiters = @waiters, []
+        pending = @waiters
+        @waiters = []
         pending.each { |waiter| Runtime.wake_dequeued(waiter, self) }
         nil
       end

data/lib/pg_pipeline/runtime/semaphore.rb

@@ -5,7 +5,7 @@ module PgPipeline
     class Semaphore
       def initialize(limit)
         @limit = Integer(limit)
-        raise ArgumentError, "limit must be >= 0" if @limit.negative?
+        raise ArgumentError, 'limit must be >= 0' if @limit.negative?
 
         @available = @limit
         @waiting = WaitList.new
@@ -71,7 +71,7 @@ module PgPipeline
         def empty? = @head.nil?
 
         def push(node)
-          raise ArgumentError, "node already queued" if node.list
+          raise ArgumentError, 'node already queued' if node.list
 
           node.list = self
           node.prev = @tail
@@ -108,7 +108,7 @@ module PgPipeline
           @list = @prev = @next = nil
           @granted = false
           @scheduler = Fiber.scheduler
-          @waiter = {fiber: fiber, scheduler: @scheduler, ready: false, blocker: blocker, queued: true}
+          @waiter = { fiber: fiber, scheduler: @scheduler, ready: false, blocker: blocker, queued: true }
         end
 
         def granted? = @granted

data/lib/pg_pipeline/runtime/task.rb

@@ -6,7 +6,7 @@ module PgPipeline
       attr_reader :fiber, :name, :error
 
       def self.spawn(name: nil, &block)
-        raise ArgumentError, "block required" unless block
+        raise ArgumentError, 'block required' unless block
 
         task = new(name: name)
         task.__send__(:start, &block)
@@ -14,8 +14,13 @@ module PgPipeline
       end
 
       def initialize(name: nil)
-        @fiber, @scheduler, @blocker, @result, @error = nil, nil, nil, nil, nil
-        @done, @cancelled = false, false
+        @fiber = nil
+        @scheduler = nil
+        @blocker = nil
+        @result = nil
+        @error = nil
+        @done = false
+        @cancelled = false
 
         @name = name
         @waiters = []
@@ -24,7 +29,7 @@ module PgPipeline
       def wait(timeout = nil)
         unless @done
           deadline = Runtime.deadline_for(timeout)
-          raise TimeoutError, "task did not finish in time" unless join(deadline)
+          raise TimeoutError, 'task did not finish in time' unless join(deadline)
         end
         raise @error if @error
 
@@ -45,7 +50,7 @@ module PgPipeline
       def cancelled? = @cancelled
 
       def raise_if_cancelled!
-        raise Cancel, "task stopped" if @cancelled
+        raise Cancel, 'task stopped' if @cancelled
       end
 
       def enter_block(waiter)
@@ -68,7 +73,7 @@ module PgPipeline
           this.__send__(:adopt_fiber, Fiber.current)
           Thread.current[Runtime::CURRENT_TASK_KEY] = this
           begin
-            this.__send__(:complete, block.call, nil)
+            this.__send__(:complete, yield, nil)
           rescue Cancel
             this.__send__(:complete, nil, nil)
           rescue Exception => e
@@ -89,7 +94,7 @@ module PgPipeline
         return false unless fiber&.alive?
         return false unless scheduler&.respond_to?(:fiber_interrupt)
 
-        result = scheduler.fiber_interrupt(fiber, Cancel.new("task stopped"))
+        result = scheduler.fiber_interrupt(fiber, Cancel.new('task stopped'))
         result != false
       rescue FiberError
         false
@@ -119,7 +124,8 @@ module PgPipeline
       end
 
       def wake_waiters
-        pending, @waiters = @waiters, []
+        pending = @waiters
+        @waiters = []
         pending.each { |waiter| Runtime.wake_dequeued(waiter, self) }
       end
     end

data/lib/pg_pipeline/runtime.rb

@@ -1,8 +1,8 @@
 # frozen_string_literal: true
 
-require "timeout"
+require 'timeout'
 
-require_relative "errors"
+require_relative 'errors'
 
 module PgPipeline
   module Runtime
@@ -14,7 +14,7 @@ module PgPipeline
 
     CURRENT_TASK_KEY = :pg_pipeline_current_task
     WAITER_KEY = :pg_pipeline_waiter
-    DEADLINE_MESSAGE = "execution expired"
+    DEADLINE_MESSAGE = 'execution expired'
 
     MISSING_TIMEOUT_HOOK_WARNING = <<~MESSAGE
       pg_pipeline: %s does not implement #timeout_after.
@@ -32,7 +32,7 @@ module PgPipeline
     end
 
     def scheduler!
-      Fiber.scheduler or raise Error, "operation requires an active Fiber scheduler"
+      Fiber.scheduler or raise Error, 'operation requires an active Fiber scheduler'
     end
 
     def native_timeouts?(target = Fiber.scheduler)
@@ -48,9 +48,10 @@ module PgPipeline
         fiber: nil, scheduler: nil, ready: false, blocker: nil
       }
 
-      if waiter[:blocker]
-        raise Error, "waiter already parked on #{waiter[:blocker].class}; " \
-                     "a fiber may only park in one place at a time"
+      current_blocker = waiter[:blocker]
+      if current_blocker
+        raise Error, "waiter already parked on #{current_blocker.class}; " \
+                     'a fiber may only park in one place at a time'
       end
 
       waiter[:fiber] = Fiber.current
@@ -116,19 +117,19 @@ module PgPipeline
       return nil if timeout.nil?
 
       seconds = Float(timeout)
-      raise ArgumentError, "timeout must be non-negative and finite" unless seconds.finite? && seconds >= 0
+      raise ArgumentError, 'timeout must be non-negative and finite' unless seconds.finite? && seconds >= 0
 
       monotonic_now + seconds
     end
 
-    def with_timeout(duration)
+    def with_timeout(duration, &block)
       return yield if duration.nil?
 
       seconds = Float(duration)
-      raise ArgumentError, "timeout must be non-negative and finite" unless seconds.finite? && seconds >= 0
+      raise ArgumentError, 'timeout must be non-negative and finite' unless seconds.finite? && seconds >= 0
 
       begin
-        arm_deadline(seconds) { yield }
+        arm_deadline(seconds, &block)
       rescue Deadline => e
         raise TimeoutError, e.message
       end
@@ -136,9 +137,7 @@ module PgPipeline
 
     def arm_deadline(seconds, &block)
       target = Fiber.scheduler
-      if native_timeouts?(target)
-        return target.timeout_after(seconds, Deadline, DEADLINE_MESSAGE, &block)
-      end
+      return target.timeout_after(seconds, Deadline, DEADLINE_MESSAGE, &block) if native_timeouts?(target)
 
       warn_missing_timeout_hook(target) if target
       ::Timeout.timeout(seconds, Deadline, DEADLINE_MESSAGE, &block)
@@ -154,7 +153,7 @@ module PgPipeline
   end
 end
 
-require_relative "runtime/notification"
-require_relative "runtime/queue"
-require_relative "runtime/semaphore"
-require_relative "runtime/task"
+require_relative 'runtime/notification'
+require_relative 'runtime/queue'
+require_relative 'runtime/semaphore'
+require_relative 'runtime/task'

data/lib/pg_pipeline/server_caps.rb

@@ -1,8 +1,8 @@
 # frozen_string_literal: true
 
-require "pg"
+require 'pg'
 
-require_relative "errors"
+require_relative 'errors'
 
 module PgPipeline
   class ServerCaps
@@ -59,7 +59,7 @@ module PgPipeline
       return true if supported?(caps)
 
       raise UnsupportedServerError,
-            "pg_pipeline requires libpq >= 14, PostgreSQL protocol v3, and " \
+            'pg_pipeline requires libpq >= 14, PostgreSQL protocol v3, and ' \
             "ruby-pg pipeline bindings; libpq=#{caps.libpq_version}, " \
             "protocol=#{caps.protocol_version}, pipeline_api=#{caps.pipeline_api}"
     end

data/lib/pg_pipeline/session.rb

@@ -1,6 +1,6 @@
 # frozen_string_literal: true
 
-require_relative "errors"
+require_relative 'errors'
 
 module PgPipeline
   class Session
@@ -30,10 +30,10 @@ module PgPipeline
 
     def ensure_active!(session)
       conn = connection(session)
-      raise Error, "session handle is no longer active" unless session.active? && conn
+      raise Error, 'session handle is no longer active' unless session.active? && conn
       return if Fiber.current.equal?(session.owner_fiber)
 
-      raise Error, "session handle is fiber-local and cannot be used from another fiber"
+      raise Error, 'session handle is fiber-local and cannot be used from another fiber'
     end
 
     def query(session, sql, params)
@@ -43,19 +43,21 @@ module PgPipeline
 
     def exec(session, sql, params = nil)
       ensure_active!(session)
+      conn = connection(session)
       if params.nil?
-        connection(session).exec(sql)
+        conn.exec(sql)
       else
-        connection(session).exec_params(sql, params)
+        conn.exec_params(sql, params)
       end
     end
 
     def prepare(session, name, sql, param_types)
       ensure_active!(session)
+      conn = connection(session)
       if param_types
-        connection(session).prepare(name, sql, param_types)
+        conn.prepare(name, sql, param_types)
       else
-        connection(session).prepare(name, sql)
+        conn.prepare(name, sql)
       end
     end
 

data/lib/pg_pipeline/session_guard.rb

@@ -1,6 +1,6 @@
 # frozen_string_literal: true
 
-require_relative "errors"
+require_relative 'errors'
 
 module PgPipeline
   module SessionGuard
@@ -9,27 +9,27 @@ module PgPipeline
     VALID_MODES = %i[default strict].freeze
     ALLOWED_LEADING = %w[select insert update delete merge values with].freeze
     FORBIDDEN_PATTERNS = {
-      "set_config" => /\bset_config\s*\(/i,
-      "setseed" => /\bsetseed\s*\(/i,
-      "currval" => /\bcurrval\s*\(/i,
-      "lastval" => /\blastval\s*\(/i,
-      "session-advisory-lock" => /\bpg_(?:try_)?advisory_lock(?:_shared)?\s*\(/i,
-      "session-advisory-unlock" => /\bpg_advisory_unlock(?:_shared|_all)?\s*\(/i,
-      "select-into-temp" => /\binto\s+(?:(?:global|local)\s+)?temp(?:orary)?\b/i,
-      "select-into-pg-temp" => /\binto\s+(?:table\s+)?pg_temp(?:_\d+)?\s*\./i
+      'set_config' => /\bset_config\s*\(/i,
+      'setseed' => /\bsetseed\s*\(/i,
+      'currval' => /\bcurrval\s*\(/i,
+      'lastval' => /\blastval\s*\(/i,
+      'session-advisory-lock' => /\bpg_(?:try_)?advisory_lock(?:_shared)?\s*\(/i,
+      'session-advisory-unlock' => /\bpg_advisory_unlock(?:_shared|_all)?\s*\(/i,
+      'select-into-temp' => /\binto\s+(?:(?:global|local)\s+)?temp(?:orary)?\b/i,
+      'select-into-pg-temp' => /\binto\s+(?:table\s+)?pg_temp(?:_\d+)?\s*\./i
     }.freeze
     STRICT_FORBIDDEN = {
-      "nextval" => /\bnextval\s*\(/i,
-      "setval" => /\bsetval\s*\(/i,
-      "setseed" => /\bsetseed\s*\(/i,
-      "set_config" => /\bset_config\s*\(/i,
-      "session-advisory-lock" => /\bpg_(?:try_)?advisory_lock(?:_shared)?\s*\(/i,
-      "session-advisory-unlock" => /\bpg_advisory_unlock(?:_shared|_all)?\s*\(/i,
-      "pg_export_snapshot" => /\bpg_export_snapshot\s*\(/i
+      'nextval' => /\bnextval\s*\(/i,
+      'setval' => /\bsetval\s*\(/i,
+      'setseed' => /\bsetseed\s*\(/i,
+      'set_config' => /\bset_config\s*\(/i,
+      'session-advisory-lock' => /\bpg_(?:try_)?advisory_lock(?:_shared)?\s*\(/i,
+      'session-advisory-unlock' => /\bpg_advisory_unlock(?:_shared|_all)?\s*\(/i,
+      'pg_export_snapshot' => /\bpg_export_snapshot\s*\(/i
     }.freeze
     PATTERN_PREFILTER = /\binto\b/i
     DOLLAR_QUOTE_TAG = /\A\$(?:(?:[A-Za-z_]|[^\x00-\x7F])(?:[A-Za-z0-9_]|[^\x00-\x7F])*)?\$/
-    NEEDS_MASK = /['"]|--|\/\*|\$(?:(?:[A-Za-z_]|[^\x00-\x7F])(?:[A-Za-z0-9_]|[^\x00-\x7F])*)?\$/
+    NEEDS_MASK = %r{['"]|--|/\*|\$(?:(?:[A-Za-z_]|[^\x00-\x7F])(?:[A-Za-z0-9_]|[^\x00-\x7F])*)?\$}
     LEADING_KEYWORD = /\A\s*([a-zA-Z_]+)/
     WHITESPACE_BYTES = [9, 10, 11, 12, 13, 32].freeze
 
@@ -43,7 +43,7 @@ module PgPipeline
 
       raise UnsafeMultiplexError,
             "refusing to multiplex SQL (#{reason}); the shared path only accepts " \
-            "session-neutral operations. Use Client#session for session work or " \
+            'session-neutral operations. Use Client#session for session work or ' \
             "Client#transaction for an explicit transaction.\n" \
             "  offending SQL: #{sql.to_s.strip[0, 160]}"
     end
@@ -75,20 +75,29 @@ module PgPipeline
       code = NEEDS_MASK.match?(sql) ? code_only(sql) : sql
       lead = code[LEADING_KEYWORD, 1]&.downcase
 
-      return "empty" unless lead
+      return 'empty' unless lead
       return "leading:#{lead}" unless ALLOWED_LEADING.include?(lead)
-      return "multiple-statements" if multiple_statements?(code)
-      return nil unless code.include?("(") || code.match?(PATTERN_PREFILTER)
+      return 'multiple-statements' if multiple_statements?(code)
+      return nil unless code.include?('(') || code.match?(PATTERN_PREFILTER)
 
-      FORBIDDEN_PATTERNS.each do |name, pattern|
+      forbidden_reason(code) || strict_forbidden_reason(code, mode)
+    end
+
+    def forbidden_reason(code)
+      pattern_reason(code, FORBIDDEN_PATTERNS)
+    end
+
+    def strict_forbidden_reason(code, mode)
+      return unless mode == :strict
+
+      reason = pattern_reason(code, STRICT_FORBIDDEN)
+      "strict:#{reason}" if reason
+    end
+
+    def pattern_reason(code, patterns)
+      patterns.each do |name, pattern|
         return name if code.match?(pattern)
       end
-      if mode == :strict
-        STRICT_FORBIDDEN.each do |name, pattern|
-          return "strict:#{name}" if code.match?(pattern)
-        end
-      end
-
       nil
     end
 
@@ -96,7 +105,7 @@ module PgPipeline
       normalized = mode.respond_to?(:to_sym) ? mode.to_sym : mode
       return normalized if VALID_MODES.include?(normalized)
 
-      raise ArgumentError, "guard must be one of: #{VALID_MODES.map(&:inspect).join(", ")}"
+      raise ArgumentError, "guard must be one of: #{VALID_MODES.map(&:inspect).join(', ')}"
     end
 
     def code_only(sql)
@@ -104,43 +113,19 @@ module PgPipeline
       size = source.bytesize
       output = String.new(capacity: size, encoding: Encoding::BINARY)
       index = 0
-      block_depth = 0
 
       while index < size
         byte = source.getbyte(index)
-        nxt = index + 1 < size ? source.getbyte(index + 1) : nil
-        if block_depth.positive?
-          if byte == 47 && nxt == 42
-            block_depth += 1
-            output << "  "
-            index += 2
-          elsif byte == 42 && nxt == 47
-            block_depth -= 1
-            output << "  "
-            index += 2
-          else
-            output << (byte == 10 ? 10 : 32)
-            index += 1
-          end
-          next
-        end
+        next_index = index + 1
+        nxt = next_index < size ? source.getbyte(next_index) : nil
 
         if byte == 45 && nxt == 45
-          newline = source.index("\n", index + 2)
-          if newline
-            output << (" " * (newline - index)) << "\n"
-            index = newline + 1
-          else
-            output << (" " * (size - index))
-            break
-          end
+          index = mask_line_comment(source, output, index)
           next
         end
 
         if byte == 47 && nxt == 42
-          block_depth = 1
-          output << "  "
-          index += 2
+          index = mask_block_comment(source, output, index)
           next
         end
 
@@ -155,12 +140,8 @@ module PgPipeline
         end
 
         if byte == 36
-          remainder = source.byteslice(index, size - index)
-          tag = DOLLAR_QUOTE_TAG.match(remainder)&.[](0)
-          if tag
-            closing = source.index(tag, index + tag.bytesize)
-            finish = closing ? closing + tag.bytesize : size
-            output << (" " * (finish - index))
+          finish = mask_dollar_quote(source, output, index, size)
+          if finish
             index = finish
             next
           end
@@ -173,29 +154,79 @@ module PgPipeline
       output
     end
 
+    def mask_line_comment(source, output, index)
+      newline = source.index("\n", index + 2)
+      if newline
+        output << (' ' * (newline - index)) << "\n"
+        newline + 1
+      else
+        output << (' ' * (source.bytesize - index))
+        source.bytesize
+      end
+    end
+
+    def mask_block_comment(source, output, index)
+      depth = 1
+      index = append_comment_delimiter(output, index)
+
+      while index < source.bytesize
+        byte = source.getbyte(index)
+        nxt = source.getbyte(index + 1)
+        if byte == 47 && nxt == 42
+          depth += 1
+          index = append_comment_delimiter(output, index)
+        elsif byte == 42 && nxt == 47
+          depth -= 1
+          index = append_comment_delimiter(output, index)
+          break if depth.zero?
+        else
+          output << (byte == 10 ? 10 : 32)
+          index += 1
+        end
+      end
+
+      index
+    end
+
+    def append_comment_delimiter(output, index)
+      output << '  '
+      index + 2
+    end
+
+    def mask_dollar_quote(source, output, index, size)
+      remainder = source.byteslice(index, size - index)
+      tag = DOLLAR_QUOTE_TAG.match(remainder)&.[](0)
+      return unless tag
+
+      closing = source.index(tag, index + tag.bytesize)
+      finish = closing ? closing + tag.bytesize : size
+      output << (' ' * (finish - index))
+      finish
+    end
+
     def mask_quoted(source, output, index, quote_byte, escape_backslash:)
-      output << " "
+      output << ' '
       index += 1
 
       while index < source.bytesize
         byte = source.getbyte(index)
         if escape_backslash && byte == 92
-          output << " "
+          output << ' '
           index += 1
           if index < source.bytesize
-            output << (source.getbyte(index) == 10 ? "\n" : " ")
+            output << (source.getbyte(index) == 10 ? "\n" : ' ')
             index += 1
           end
         elsif byte == quote_byte
           if source.getbyte(index + 1) == quote_byte
-            output << "  "
+            output << '  '
             index += 2
           else
-            output << " "
+            output << ' '
             return index + 1
           end
         else
-          output << (byte == 10 ? "\n" : " ")
+          output << (byte == 10 ? "\n" : ' ')
           index += 1
         end
       end
@@ -207,15 +238,19 @@ module PgPipeline
       return false if quote_index.zero?
 
       marker = source.getbyte(quote_index - 1)
-      return false unless marker == 69 || marker == 101
+      return false unless [69, 101].include?(marker)
 
       before = quote_index >= 2 ? source.getbyte(quote_index - 2) : nil
-      before.nil? || !(before.between?(48, 57) || before.between?(65, 90) || before.between?(97, 122) || before == 95)
+      !sql_identifier_byte?(before)
+    end
+
+    def sql_identifier_byte?(byte)
+      byte && (byte.between?(48, 57) || byte.between?(65, 90) || byte.between?(97, 122) || byte == 95)
     end
-    private_class_method :escape_string_prefix?
+    private_class_method :escape_string_prefix?, :sql_identifier_byte?
 
     def multiple_statements?(code)
-      first = code.index(";")
+      first = code.index(';')
       return false unless first
 
       index = first + 1