pgsqlarbiter 0.2.0

data/lib/pgsqlarbiter/default_query_functions.rb

@@ -0,0 +1,177 @@
+# frozen_string_literal: true
+
+require "set"
+
+module Pgsqlarbiter
+  # Default set of PostgreSQL functions permitted in queries. This includes functions that could
+  # cause resource exhaustion (e.g. generate_series) — resource limits should be enforced elsewhere.
+  #
+  # Excluded: functions not used by regular queries such as pg_sleep, set_config, lo_*, pg_advisory_lock,
+  # pg_notify, sequence functions, and system information functions.
+  DEFAULT_QUERY_FUNCTIONS = Set[
+    # -- Aggregate functions --
+    "array_agg", "avg", "bit_and", "bit_or", "bit_xor",
+    "bool_and", "bool_or", "count", "every",
+    "json_agg", "jsonb_agg", "json_object_agg", "jsonb_object_agg",
+    "max", "min", "range_agg", "range_intersect_agg",
+    "string_agg", "sum", "xmlagg",
+
+    # -- Statistical aggregate functions --
+    "corr", "covar_pop", "covar_samp",
+    "regr_avgx", "regr_avgy", "regr_count", "regr_intercept",
+    "regr_r2", "regr_slope", "regr_sxx", "regr_sxy", "regr_syy",
+    "stddev", "stddev_pop", "stddev_samp",
+    "variance", "var_pop", "var_samp",
+
+    # -- Ordered-set aggregate functions --
+    "mode", "percentile_cont", "percentile_disc",
+
+    # -- Window functions --
+    "row_number", "rank", "dense_rank", "percent_rank", "cume_dist",
+    "ntile", "lag", "lead", "first_value", "last_value", "nth_value",
+
+    # -- Mathematical functions --
+    "abs", "cbrt", "ceil", "ceiling", "degrees", "div",
+    "exp", "factorial", "floor", "gcd", "lcm",
+    "ln", "log", "log10", "min_scale", "mod",
+    "pi", "power", "radians", "random",
+    "round", "scale", "sign", "sqrt",
+    "trim_scale", "trunc", "width_bucket",
+
+    # -- Trigonometric functions --
+    "acos", "acosd", "asin", "asind",
+    "atan", "atan2", "atan2d", "atand",
+    "cos", "cosd", "cot", "cotd",
+    "sin", "sind", "tan", "tand",
+
+    # -- Hyperbolic functions --
+    "sinh", "cosh", "tanh", "asinh", "acosh", "atanh",
+
+    # -- String functions --
+    "ascii", "btrim", "char_length", "character_length",
+    "chr", "concat", "concat_ws",
+    "convert", "convert_from", "convert_to",
+    "decode", "encode", "format",
+    "initcap", "left", "length", "lower",
+    "lpad", "ltrim", "md5",
+    "normalize", "octet_length", "overlay",
+    "parse_ident", "position",
+    "quote_ident", "quote_literal", "quote_nullable",
+    "regexp_count", "regexp_instr", "regexp_like",
+    "regexp_match", "regexp_matches", "regexp_replace",
+    "regexp_split_to_array", "regexp_split_to_table", "regexp_substr",
+    "repeat", "replace", "reverse", "right",
+    "rpad", "rtrim", "split_part",
+    "starts_with", "string_to_array", "string_to_table",
+    "strpos", "substr", "substring",
+    "to_ascii", "to_hex", "translate", "trim",
+    "unicode", "unistr", "upper",
+
+    # -- Binary string functions --
+    "bit_length", "get_bit", "get_byte",
+    "set_bit", "set_byte",
+    "sha224", "sha256", "sha384", "sha512",
+
+    # -- Date/time functions --
+    "age", "clock_timestamp", "date_bin",
+    "date_part", "date_trunc", "extract",
+    "isfinite", "justify_days", "justify_hours", "justify_interval",
+    "make_date", "make_interval", "make_time",
+    "make_timestamp", "make_timestamptz",
+    "now", "statement_timestamp",
+    "timeofday", "transaction_timestamp",
+
+    # -- Formatting functions --
+    "to_char", "to_date", "to_number", "to_timestamp",
+
+    # -- Conditional functions --
+    "coalesce", "nullif", "greatest", "least",
+
+    # -- Comparison functions --
+    "num_nulls", "num_nonnulls",
+
+    # -- JSON/JSONB functions --
+    "to_json", "to_jsonb", "array_to_json", "row_to_json",
+    "json_build_array", "jsonb_build_array",
+    "json_build_object", "jsonb_build_object",
+    "json_object", "jsonb_object",
+    "json_array", "jsonb_array",
+    "json_array_length", "jsonb_array_length",
+    "json_each", "jsonb_each",
+    "json_each_text", "jsonb_each_text",
+    "json_extract_path", "jsonb_extract_path",
+    "json_extract_path_text", "jsonb_extract_path_text",
+    "json_object_keys", "jsonb_object_keys",
+    "json_populate_record", "jsonb_populate_record",
+    "json_populate_recordset", "jsonb_populate_recordset",
+    "json_to_record", "jsonb_to_record",
+    "json_to_recordset", "jsonb_to_recordset",
+    "json_strip_nulls", "jsonb_strip_nulls",
+    "jsonb_set", "jsonb_set_lax", "jsonb_insert",
+    "jsonb_path_exists", "jsonb_path_match",
+    "jsonb_path_query", "jsonb_path_query_array", "jsonb_path_query_first",
+    "jsonb_path_exists_tz", "jsonb_path_match_tz",
+    "jsonb_path_query_tz", "jsonb_path_query_array_tz", "jsonb_path_query_first_tz",
+    "jsonb_pretty",
+    "json_typeof", "jsonb_typeof",
+    "json_array_elements", "jsonb_array_elements",
+    "json_array_elements_text", "jsonb_array_elements_text",
+    "json_scalar", "jsonb_scalar",
+    "json_table",
+
+    # -- Array functions --
+    "array_append", "array_cat", "array_dims", "array_fill",
+    "array_length", "array_lower", "array_ndims",
+    "array_position", "array_positions",
+    "array_prepend", "array_remove", "array_replace",
+    "array_sample", "array_shuffle",
+    "array_to_string", "array_upper",
+    "cardinality", "trim_array", "unnest",
+
+    # -- Range/multirange functions --
+    "isempty", "lower_inc", "upper_inc", "lower_inf", "upper_inf",
+    "range_merge", "multirange",
+    "int4range", "int8range", "numrange",
+    "tsrange", "tstzrange", "daterange",
+    "int4multirange", "int8multirange", "nummultirange",
+    "tsmultirange", "tstzmultirange", "datemultirange",
+
+    # -- Set-returning functions --
+    "generate_series", "generate_subscripts",
+
+    # -- Geometric functions --
+    "area", "center", "diagonal", "diameter", "height",
+    "isclosed", "isopen", "npoints",
+    "pclose", "popen", "radius", "slope", "width",
+    "box", "circle", "line", "lseg", "path", "point", "polygon",
+
+    # -- Network address functions --
+    "abbrev", "broadcast", "family",
+    "host", "hostmask", "inet_merge", "inet_same_family",
+    "masklen", "netmask", "network", "set_masklen",
+
+    # -- Text search functions --
+    "array_to_tsvector", "numnode",
+    "plainto_tsquery", "phraseto_tsquery",
+    "querytree", "setweight", "strip",
+    "to_tsquery", "to_tsvector",
+    "ts_delete", "ts_filter", "ts_headline", "ts_lexize",
+    "ts_rank", "ts_rank_cd", "ts_rewrite",
+    "tsvector_to_array", "websearch_to_tsquery",
+
+    # -- XML functions --
+    "xmlcomment", "xmlconcat", "xmlexists",
+    "xmlelement", "xmlforest", "xmlparse", "xmlroot", "xmlserialize",
+    "xmltable",
+    "xpath", "xpath_exists",
+
+    # -- Grouping function --
+    "grouping",
+
+    # -- Enum functions --
+    "enum_first", "enum_last", "enum_range",
+
+    # -- UUID functions --
+    "gen_random_uuid", "uuidv4", "uuidv7"
+  ].freeze
+end

data/lib/pgsqlarbiter/error.rb

@@ -0,0 +1,18 @@
+# frozen_string_literal: true
+
+module Pgsqlarbiter
+  # Base error class for all pgsqlarbiter errors.
+  class Error < StandardError; end
+
+  # Raised when the lexer encounters invalid syntax (e.g. unterminated strings or unexpected characters).
+  class LexError < Error; end
+
+  # Raised when the analyzer encounters invalid or unparseable SQL structure.
+  class ParseError < Error; end
+
+  # Raised when the SQL contains more than one statement.
+  class MultipleStatementsError < Error; end
+
+  # Raised when the statement type is not a supported DML type (e.g. DDL, DCL, or TCL).
+  class DisallowedStatementError < Error; end
+end

data/lib/pgsqlarbiter/keywords.rb

@@ -0,0 +1,66 @@
+# frozen_string_literal: true
+
+require "set"
+
+module Pgsqlarbiter
+  module Keywords
+    ALLOWED_STATEMENT_TYPES = Set[
+      "SELECT", "INSERT", "UPDATE", "DELETE", "MERGE", "VALUES", "WITH"
+    ].freeze
+
+    # Keywords that take parentheses but are NOT function calls
+    NON_FUNCTION_KEYWORDS = Set[
+      "EXISTS", "CASE", "CAST", "IN", "NOT", "ANY", "ALL", "SOME",
+      "ARRAY", "ROW", "VALUES", "LATERAL", "TABLE"
+    ].freeze
+
+    # Keywords that ARE function calls despite being reserved words
+    FUNCTION_KEYWORDS = Set[
+      "COALESCE", "NULLIF", "GREATEST", "LEAST", "EXTRACT",
+      "TRIM", "SUBSTRING", "OVERLAY", "POSITION", "NORMALIZE", "GROUPING",
+      "XMLELEMENT", "XMLFOREST", "XMLPARSE", "XMLROOT", "XMLSERIALIZE"
+    ].freeze
+
+    # Complete set of keywords the lexer recognizes
+    ALL = Set[
+      # Statement types
+      "SELECT", "INSERT", "UPDATE", "DELETE", "MERGE", "VALUES", "WITH",
+      # Disallowed statement types (rejected by analyzer)
+      "CREATE", "DROP", "ALTER", "TRUNCATE", "GRANT", "REVOKE",
+      "SHOW", "SET", "RESET",
+      "BEGIN", "START", "COMMIT", "ROLLBACK", "SAVEPOINT", "RELEASE",
+      "PREPARE", "EXECUTE", "DEALLOCATE",
+      "LISTEN", "NOTIFY", "UNLISTEN",
+      "LOAD", "COPY", "VACUUM", "ANALYZE", "CLUSTER", "REINDEX",
+      "LOCK", "DISCARD", "COMMENT", "SECURITY", "REASSIGN", "REFRESH",
+      "IMPORT", "CALL", "DO", "EXPLAIN",
+      # Structural keywords
+      "FROM", "JOIN", "INNER", "LEFT", "RIGHT", "FULL", "CROSS",
+      "NATURAL", "OUTER", "ON", "USING", "INTO", "AS",
+      "WHERE", "GROUP", "HAVING", "ORDER", "LIMIT", "OFFSET", "FETCH",
+      "UNION", "INTERSECT", "EXCEPT",
+      "ALL", "DISTINCT", "LATERAL", "ONLY", "TABLE",
+      "RETURNING", "RECURSIVE", "COLUMNS",
+      "NOT", "MATERIALIZED",
+      "MATCHED", "WHEN", "THEN", "BY", "CONFLICT",
+      "AND", "OR", "IS", "IN", "BETWEEN",
+      "LIKE", "ILIKE", "SIMILAR",
+      "CASE", "CAST", "END", "ELSE",
+      "EXISTS", "ANY", "SOME", "ARRAY", "ROW",
+      "WINDOW", "OVER", "PARTITION", "WITHIN", "FILTER",
+      "NOTHING",
+      "FOR", "IF", "ELSE", "TRUE", "FALSE", "NULL",
+      "ASC", "DESC", "NULLS", "FIRST", "LAST",
+      # Function keywords
+      "COALESCE", "NULLIF", "GREATEST", "LEAST", "EXTRACT",
+      "TRIM", "SUBSTRING", "OVERLAY", "POSITION", "NORMALIZE", "GROUPING",
+      "XMLELEMENT", "XMLFOREST", "XMLPARSE", "XMLROOT", "XMLSERIALIZE",
+      # Type keywords (common)
+      "INT", "INTEGER", "BIGINT", "SMALLINT", "REAL", "FLOAT",
+      "DOUBLE", "PRECISION", "NUMERIC", "DECIMAL",
+      "CHAR", "CHARACTER", "VARCHAR", "TEXT",
+      "BOOLEAN", "DATE", "TIME", "TIMESTAMP", "INTERVAL",
+      "BOTH", "LEADING", "TRAILING"
+    ].freeze
+  end
+end

data/lib/pgsqlarbiter/lexer.rb

@@ -0,0 +1,272 @@
+# frozen_string_literal: true
+
+require "strscan"
+
+module Pgsqlarbiter
+  # SQL lexer that converts a query string into an array of {Token} objects.
+  #
+  # Handles all PostgreSQL token types including keywords, identifiers (plain and
+  # double-quoted), strings (single-quoted, dollar-quoted, and prefixed), numbers,
+  # parameters, operators, and punctuation.
+  class Lexer
+    include TokenType
+
+    # Tokenize a SQL query string.
+    #
+    # @param sql [String] the SQL string to tokenize
+    # @return [Array<Token>] list of tokens ending with an EOF token
+    # @raise [LexError] on invalid syntax such as unexpected characters or unterminated
+    #   strings/comments
+    def tokenize(sql)
+      @scanner = StringScanner.new(sql)
+      @tokens = []
+
+      until @scanner.eos?
+        scan_token
+      end
+
+      @tokens << Token.new(type: EOF, value: nil, position: @scanner.pos)
+      @tokens
+    end
+
+    private
+
+    def scan_token
+      pos = @scanner.pos
+
+      # 1. Whitespace — skip
+      return if @scanner.skip(/\s+/)
+
+      # 2. Line comment — skip
+      return if @scanner.skip(/--[^\n]*/)
+
+      # 3. Block comment — skip (handle nesting)
+      if @scanner.scan(/\/\*/)
+        scan_block_comment(pos)
+        return
+      end
+
+      # 4. Dollar-quoted string
+      if @scanner.check(/\$([\p{L}_][\p{L}\d_]*)?\$/)
+        scan_dollar_string(pos)
+        return
+      end
+
+      # 5. Parameter placeholder ($1, $2, ...)
+      if (m = @scanner.scan(/\$\d+/))
+        @tokens << Token.new(type: PARAM, value: m, position: pos)
+        return
+      end
+
+      # 6. Single-quoted strings (including E'', B'', X'', N'' prefixes)
+      # Prefix detection is handled in the identifier branch (step 12)
+      if @scanner.check(/'/)
+        scan_single_quoted_string(pos, prefix: nil)
+        return
+      end
+
+      # 7. Double-quoted identifier
+      if @scanner.check(/"/)
+        scan_quoted_identifier(pos, unicode: false)
+        return
+      end
+
+      # 8. Numbers
+      if (m = @scanner.scan(/0[xX][0-9a-fA-F_]+/))
+        @tokens << Token.new(type: NUMBER, value: m, position: pos)
+        return
+      end
+      if (m = @scanner.scan(/0[oO][0-7_]+/))
+        @tokens << Token.new(type: NUMBER, value: m, position: pos)
+        return
+      end
+      if (m = @scanner.scan(/0[bB][01_]+/))
+        @tokens << Token.new(type: NUMBER, value: m, position: pos)
+        return
+      end
+      if (m = @scanner.scan(/\d[\d_]*\.[\d_]+(?:[eE][+-]?\d[\d_]*)?/))
+        @tokens << Token.new(type: NUMBER, value: m, position: pos)
+        return
+      end
+      if (m = @scanner.scan(/\d[\d_]*[eE][+-]?\d[\d_]*/))
+        @tokens << Token.new(type: NUMBER, value: m, position: pos)
+        return
+      end
+      if (m = @scanner.scan(/\d[\d_]*/))
+        # Check this isn't followed by dot+digits (which would be a decimal)
+        if @scanner.check(/\.[\d_]/)
+          m += @scanner.scan(/\.[\d_]+(?:[eE][+-]?\d[\d_]*)?/)
+        end
+        @tokens << Token.new(type: NUMBER, value: m, position: pos)
+        return
+      end
+      if (m = @scanner.scan(/\.[\d_]+(?:[eE][+-]?\d[\d_]*)?/))
+        @tokens << Token.new(type: NUMBER, value: m, position: pos)
+        return
+      end
+
+      # 9. Typecast ::
+      if @scanner.scan(/::/)
+        @tokens << Token.new(type: TYPECAST, value: "::", position: pos)
+        return
+      end
+
+      # 10. Single-char punctuation
+      ch = @scanner.peek(1)
+      case ch
+      when "("
+        @scanner.getch
+        @tokens << Token.new(type: LPAREN, value: "(", position: pos)
+        return
+      when ")"
+        @scanner.getch
+        @tokens << Token.new(type: RPAREN, value: ")", position: pos)
+        return
+      when "["
+        @scanner.getch
+        @tokens << Token.new(type: LBRACKET, value: "[", position: pos)
+        return
+      when "]"
+        @scanner.getch
+        @tokens << Token.new(type: RBRACKET, value: "]", position: pos)
+        return
+      when ","
+        @scanner.getch
+        @tokens << Token.new(type: COMMA, value: ",", position: pos)
+        return
+      when ";"
+        @scanner.getch
+        @tokens << Token.new(type: SEMICOLON, value: ";", position: pos)
+        return
+      when "*"
+        @scanner.getch
+        @tokens << Token.new(type: STAR, value: "*", position: pos)
+        return
+      when "."
+        @scanner.getch
+        @tokens << Token.new(type: DOT, value: ".", position: pos)
+        return
+      end
+
+      # 11. Multi-char operators
+      if (m = @scanner.scan(%r{[+\-/<>=~!@#%^&|`?]+}))
+        @tokens << Token.new(type: OP, value: m, position: pos)
+        return
+      end
+
+      # 12. Unquoted identifier / keyword (with string prefix detection)
+      if (m = @scanner.scan(/[\p{L}_][\p{L}\d_]*/))
+        lower = m.downcase
+
+        # Check for U& prefix for unicode strings/identifiers
+        if lower == "u" && @scanner.check(/&['"]/i)
+          @scanner.scan(/&/)
+          if @scanner.check(/'/)
+            scan_single_quoted_string(pos, prefix: "U&")
+          else
+            scan_quoted_identifier(pos, unicode: true)
+          end
+          return
+        end
+
+        # Check for string prefixes: E, B, X, N immediately followed by '
+        if %w[e b x n].include?(lower) && @scanner.check(/'/)
+          scan_single_quoted_string(pos, prefix: lower)
+          return
+        end
+
+        upper = m.upcase
+        if Keywords::ALL.include?(upper)
+          @tokens << Token.new(type: KEYWORD, value: upper, position: pos)
+        else
+          @tokens << Token.new(type: IDENT, value: lower, position: pos)
+        end
+        return
+      end
+
+      # 13. Single colon (not part of ::)
+      if @scanner.scan(/:/)
+        @tokens << Token.new(type: OP, value: ":", position: pos)
+        return
+      end
+
+      raise LexError, "unexpected character #{@scanner.peek(1).inspect} at position #{pos}"
+    end
+
+    def scan_block_comment(start_pos)
+      depth = 1
+      until @scanner.eos?
+        if @scanner.scan(/\/\*/)
+          depth += 1
+        elsif @scanner.scan(/\*\//)
+          depth -= 1
+          return if depth == 0
+        else
+          @scanner.getch
+        end
+      end
+      raise LexError, "unterminated block comment starting at position #{start_pos}"
+    end
+
+    def scan_dollar_string(start_pos)
+      @scanner.scan(/\$([\p{L}_][\p{L}\d_]*)?\$/)
+      tag = @scanner.matched
+      content = +""
+      until @scanner.eos?
+        idx = @scanner.rest.index(tag)
+        if idx
+          content << @scanner.rest[0, idx]
+          @scanner.pos += idx + tag.length
+          @tokens << Token.new(type: STRING, value: content, position: start_pos)
+          return
+        else
+          content << @scanner.rest
+          @scanner.terminate
+        end
+      end
+      raise LexError, "unterminated dollar-quoted string starting at position #{start_pos}"
+    end
+
+    def scan_single_quoted_string(start_pos, prefix:)
+      @scanner.scan(/'/)
+      escape_mode = (prefix == "e") # E-strings support backslash escapes
+      content = +""
+
+      until @scanner.eos?
+        if escape_mode && @scanner.scan(/\\/)
+          if @scanner.eos?
+            raise LexError, "unterminated string starting at position #{start_pos}"
+          end
+          content << "\\" << @scanner.getch
+        elsif @scanner.scan(/''/)
+          content << "''"
+        elsif @scanner.scan(/'/)
+          @tokens << Token.new(type: STRING, value: content, position: start_pos)
+          return
+        else
+          content << @scanner.getch
+        end
+      end
+
+      raise LexError, "unterminated string starting at position #{start_pos}"
+    end
+
+    def scan_quoted_identifier(start_pos, unicode:)
+      @scanner.scan(/"/)
+      content = +""
+
+      until @scanner.eos?
+        if @scanner.scan(/""/)
+          content << '"'
+        elsif @scanner.scan(/"/)
+          @tokens << Token.new(type: QUOTED_IDENT, value: content, position: start_pos)
+          return
+        else
+          content << @scanner.getch
+        end
+      end
+
+      raise LexError, "unterminated quoted identifier starting at position #{start_pos}"
+    end
+  end
+end

data/lib/pgsqlarbiter/token.rb

@@ -0,0 +1,34 @@
+# frozen_string_literal: true
+
+module Pgsqlarbiter
+  # Immutable token produced by the {Lexer}.
+  #
+  # @!attribute [r] type
+  #   @return [Symbol] token type (one of the {TokenType} constants)
+  # @!attribute [r] value
+  #   @return [String, nil] the token text (+nil+ for EOF)
+  # @!attribute [r] position
+  #   @return [Integer] character offset in the original SQL string
+  Token = Data.define(:type, :value, :position)
+
+  # Constants for all token types produced by the {Lexer}.
+  module TokenType
+    KEYWORD      = :keyword
+    IDENT        = :ident
+    QUOTED_IDENT = :quoted_ident
+    STRING       = :string
+    NUMBER       = :number
+    PARAM        = :param
+    LPAREN       = :lparen
+    RPAREN       = :rparen
+    LBRACKET     = :lbracket
+    RBRACKET     = :rbracket
+    COMMA        = :comma
+    DOT          = :dot
+    SEMICOLON    = :semicolon
+    STAR         = :star
+    TYPECAST     = :typecast
+    OP           = :op
+    EOF          = :eof
+  end
+end

data/lib/pgsqlarbiter/verdict.rb

@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+module Pgsqlarbiter
+  # Immutable result of judging a SQL query against an {Arbiter}'s rules.
+  #
+  # A +Verdict+ tells you whether a query is allowed and, if not, exactly which
+  # checks failed. Use {#allowed?} for a quick boolean, {#reasons} for
+  # human-readable denial strings, or the individual fields for programmatic
+  # branching.
+  #
+  # @!attribute [r] allowed
+  #   @return [Boolean] +true+ when the query passes all checks
+  # @!attribute [r] statement_type_allowed
+  #   @return [Boolean] +true+ when the statement type is in the whitelist
+  # @!attribute [r] statement_type
+  #   @return [Symbol] the query's actual statement type
+  # @!attribute [r] disallowed_tables
+  #   @return [Array<String>] tables referenced by the query that are not whitelisted
+  # @!attribute [r] disallowed_functions
+  #   @return [Array<String>] functions called by the query that are not whitelisted
+  Verdict = Data.define(:allowed, :statement_type_allowed, :statement_type,
+                        :disallowed_tables, :disallowed_functions) do
+    alias_method :allowed?, :allowed
+    alias_method :statement_type_allowed?, :statement_type_allowed
+
+    # Human-readable denial reasons. Empty when the query is allowed.
+    #
+    # @return [Array<String>] frozen list of reason strings
+    def reasons
+      r = []
+      r << "statement type :#{statement_type} is not allowed" unless statement_type_allowed
+      disallowed_tables.each { |t| r << "table #{t.inspect} is not allowed" }
+      disallowed_functions.each { |f| r << "function #{f.inspect} is not allowed" }
+      r.freeze
+    end
+  end
+end

data/lib/pgsqlarbiter/version.rb

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