pg_pipeline 0.2.1

data/lib/pg_pipeline/session_guard.rb

@@ -0,0 +1,213 @@
+# frozen_string_literal: true
+
+require_relative "errors"
+
+module PgPipeline
+  module SessionGuard
+    module_function
+
+    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
+    }.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
+    }.freeze
+
+    def assert_multiplexable!(sql, mode: :default)
+      reason = unsafe_reason(sql, mode: mode)
+      return true unless reason
+
+      raise UnsafeMultiplexError,
+            "refusing to multiplex SQL (#{reason}); the shared path only accepts " \
+            "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
+
+    GUARD_CACHE_LIMIT = 2048
+
+    def unsafe_reason(sql, mode: :default)
+      mode = normalize_mode!(mode)
+      key = sql.to_s
+      cache = guard_cache[mode]
+      return cache[key] if cache.key?(key)
+
+      reason = compute_unsafe_reason(key, mode)
+      cache.clear if cache.size >= GUARD_CACHE_LIMIT
+      cache[key] = reason
+      reason
+    end
+
+    def guard_cache
+      @guard_cache ||= { default: {}, strict: {} }
+    end
+
+    def compute_unsafe_reason(sql, mode)
+      code = code_only(sql)
+      lead = code[/\A\s*([a-zA-Z_]+)/, 1]&.downcase
+
+      return "empty" unless lead
+      return "leading:#{lead}" unless ALLOWED_LEADING.include?(lead)
+      return "multiple-statements" if multiple_statements?(code)
+
+      FORBIDDEN_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
+
+    def normalize_mode!(mode)
+      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(", ")}"
+    end
+
+    def code_only(sql)
+      source = sql.to_s.b
+      output = String.new(capacity: source.bytesize, encoding: Encoding::BINARY)
+      index = 0
+      block_depth = 0
+
+      while index < source.bytesize
+        if block_depth.positive?
+          if source.byteslice(index, 2) == "/*"
+            block_depth += 1
+            output << "  "
+            index += 2
+          elsif source.byteslice(index, 2) == "*/"
+            block_depth -= 1
+            output << "  "
+            index += 2
+          else
+            output << (source.getbyte(index) == 10 ? "\n" : " ")
+            index += 1
+          end
+          next
+        end
+
+        if source.byteslice(index, 2) == "--"
+          newline = source.index("\n", index + 2)
+          if newline
+            output << " " * (newline - index) << "\n"
+            index = newline + 1
+          else
+            output << " " * (source.bytesize - index)
+            break
+          end
+          next
+        end
+
+        if source.byteslice(index, 2) == "/*"
+          block_depth = 1
+          output << "  "
+          index += 2
+          next
+        end
+
+        byte = source.getbyte(index)
+        if byte == 39
+          index = mask_quoted(source, output, index, 39, escape_backslash: escape_string_prefix?(source, index))
+          next
+        end
+
+        if byte == 34
+          index = mask_quoted(source, output, index, 34, escape_backslash: false)
+          next
+        end
+
+        if byte == 36
+          remainder = source.byteslice(index, source.bytesize - index)
+          tag = remainder.match(/\A\$(?:[A-Za-z_][A-Za-z0-9_]*)?\$/)&.[](0)
+          if tag
+            closing = source.index(tag, index + tag.bytesize)
+            finish = closing ? closing + tag.bytesize : source.bytesize
+            output << " " * (finish - index)
+            index = finish
+            next
+          end
+        end
+
+        output << source.getbyte(index)
+        index += 1
+      end
+
+      output
+    end
+
+    def mask_quoted(source, output, index, quote_byte, escape_backslash:)
+      output << " "
+      index += 1
+
+      while index < source.bytesize
+        byte = source.getbyte(index)
+
+        if escape_backslash && byte == 92
+          output << " "
+          index += 1
+          if index < source.bytesize
+            output << (source.getbyte(index) == 10 ? "\n" : " ")
+            index += 1
+          end
+        elsif byte == quote_byte
+          if source.getbyte(index + 1) == quote_byte
+            output << "  "
+            index += 2
+          else
+            output << " "
+            return index + 1
+          end
+        else
+          output << (byte == 10 ? "\n" : " ")
+          index += 1
+        end
+      end
+
+      index
+    end
+    private_class_method :mask_quoted
+
+    def escape_string_prefix?(source, quote_index)
+      return false if quote_index.zero?
+
+      marker = source.getbyte(quote_index - 1)
+      return false unless marker == 69 || marker == 101
+
+      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)
+    end
+    private_class_method :escape_string_prefix?
+
+    def multiple_statements?(code)
+      semicolons = []
+      code.each_char.with_index { |char, i| semicolons << i if char == ";" }
+      return false if semicolons.empty?
+
+      last_non_space = code.rstrip.length - 1
+      semicolons.length > 1 || semicolons.first != last_non_space
+    end
+  end
+end

data/lib/pg_pipeline/transaction.rb

@@ -0,0 +1,92 @@
+# frozen_string_literal: true
+
+require "pg"
+
+require_relative "session"
+
+module PgPipeline
+  class Transaction < Session
+    attr_accessor :open, :savepoint_seq
+
+    def initialize(conn)
+      super
+      @open = false
+      @savepoint_seq = 0
+    end
+
+    def run(&block)
+      TransactionOps.run(self, &block)
+    end
+
+    def savepoint(name = nil, &block)
+      TransactionOps.savepoint(self, name, &block)
+    end
+
+    def open? = @open
+  end
+
+  module TransactionOps
+    module_function
+
+    def run(tx)
+      SessionOps.ensure_active!(tx)
+      raise Error, "transaction is already open" if tx.open
+
+      conn = SessionOps.connection(tx)
+      conn.exec("BEGIN")
+      tx.open = true
+      begin
+        result = yield tx
+        conn.exec("COMMIT")
+        tx.open = false
+        result
+      rescue Exception
+        # Only roll back the transaction we opened. Early guard errors
+        # ("already open", inactive handle) must not hit this path.
+        rollback_quietly(tx)
+        raise
+      end
+    end
+
+    def savepoint(tx, name)
+      SessionOps.ensure_active!(tx)
+      raise Error, "savepoint requires an open transaction" unless tx.open
+
+      conn = SessionOps.connection(tx)
+      tx.savepoint_seq += 1
+      point = name || "pgp_sp_#{tx.savepoint_seq}"
+      ident = conn.quote_ident(point)
+
+      conn.exec("SAVEPOINT #{ident}")
+      begin
+        result = yield tx
+        conn.exec("RELEASE SAVEPOINT #{ident}")
+        result
+      rescue Exception
+        rollback_to_savepoint(tx, ident)
+        raise
+      end
+    end
+
+    def rollback_to_savepoint(tx, ident)
+      conn = SessionOps.connection(tx)
+      return unless conn
+
+      conn.exec("ROLLBACK TO SAVEPOINT #{ident}")
+      conn.exec("RELEASE SAVEPOINT #{ident}")
+    rescue PG::Error
+      nil
+    end
+
+    def rollback_quietly(tx)
+      return unless tx.open
+
+      conn = SessionOps.connection(tx)
+      conn&.exec("ROLLBACK")
+    rescue PG::Error
+      nil
+    ensure
+      tx.open = false
+    end
+  end
+end

data/lib/pg_pipeline/version.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module PgPipeline
+  VERSION = "0.2.1"
+end

data/lib/pg_pipeline.rb

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+require_relative "pg_pipeline/version"
+require_relative "pg_pipeline/errors"
+require_relative "pg_pipeline/server_caps"
+require_relative "pg_pipeline/session_guard"
+require_relative "pg_pipeline/session"
+require_relative "pg_pipeline/transaction"
+require_relative "pg_pipeline/request"
+require_relative "pg_pipeline/connection_driver"
+require_relative "pg_pipeline/pool"
+require_relative "pg_pipeline/client"
+
+module PgPipeline
+end

metadata

@@ -0,0 +1,155 @@
+--- !ruby/object:Gem::Specification
+name: pg_pipeline
+version: !ruby/object:Gem::Version
+  version: 0.2.1
+platform: ruby
+authors:
+- Roman Hajdarov
+bindir: bin
+cert_chain: []
+date: 1980-01-02 00:00:00.000000000 Z
+dependencies:
+- !ruby/object:Gem::Dependency
+  name: async
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '2.42'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '2.42'
+- !ruby/object:Gem::Dependency
+  name: pg
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '1.5'
+    - - "<"
+      - !ruby/object:Gem::Version
+        version: '2'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '1.5'
+    - - "<"
+      - !ruby/object:Gem::Version
+        version: '2'
+- !ruby/object:Gem::Dependency
+  name: rake
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '13.0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '13.0'
+- !ruby/object:Gem::Dependency
+  name: rspec
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '3.13'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '3.13'
+- !ruby/object:Gem::Dependency
+  name: ruby-prof
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '1.7'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '1.7'
+- !ruby/object:Gem::Dependency
+  name: stackprof
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '0.2'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '0.2'
+description: |
+  A driver-adjacent Ruby control-plane over ruby-pg/libpq. It multiplexes
+  independent, session-neutral extended-protocol operations from many Async
+  fibers onto a small number of PostgreSQL connections while keeping explicit
+  transactions and session-changing work on exclusive pinned connections.
+  Control-plane only: all wire work stays in libpq. Zero lines of C.
+email:
+- romanhajdarov@gmail.com
+executables: []
+extensions: []
+extra_rdoc_files: []
+files:
+- CHANGELOG.md
+- DESIGN.md
+- LICENSE.txt
+- README.md
+- lib/pg_pipeline.rb
+- lib/pg_pipeline/bounded_queue.rb
+- lib/pg_pipeline/client.rb
+- lib/pg_pipeline/connection_driver.rb
+- lib/pg_pipeline/errors.rb
+- lib/pg_pipeline/pool.rb
+- lib/pg_pipeline/request.rb
+- lib/pg_pipeline/server_caps.rb
+- lib/pg_pipeline/session.rb
+- lib/pg_pipeline/session_guard.rb
+- lib/pg_pipeline/transaction.rb
+- lib/pg_pipeline/version.rb
+homepage: https://github.com/roman-haidarov/pg_pipeline
+licenses:
+- MIT
+metadata:
+  homepage_uri: https://github.com/roman-haidarov/pg_pipeline
+  source_code_uri: https://github.com/roman-haidarov/pg_pipeline/tree/main
+  changelog_uri: https://github.com/roman-haidarov/pg_pipeline/blob/main/CHANGELOG.md
+  bug_tracker_uri: https://github.com/roman-haidarov/pg_pipeline/issues
+rdoc_options: []
+require_paths:
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '3.3'
+required_rubygems_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '0'
+requirements: []
+rubygems_version: 3.6.7
+specification_version: 4
+summary: Async-native PostgreSQL pipeline multiplexing on top of ruby-pg
+test_files: []