gitlab-pg_query 1.3.0

data/lib/pg_query/param_refs.rb

@@ -0,0 +1,45 @@
+class PgQuery
+  def param_refs # rubocop:disable Metrics/CyclomaticComplexity
+    results = []
+
+    treewalker! @tree do |_, _, v|
+      next unless v.is_a?(Hash)
+
+      if v[PARAM_REF]
+        results << { 'location' => v[PARAM_REF]['location'],
+                     'length' => param_ref_length(v[PARAM_REF]) }
+      elsif v[TYPE_CAST]
+        next unless v[TYPE_CAST]['arg'] && v[TYPE_CAST]['typeName']
+
+        p = v[TYPE_CAST]['arg'].delete(PARAM_REF)
+        t = v[TYPE_CAST]['typeName'].delete(TYPE_NAME)
+        next unless p && t
+
+        location = p['location']
+        typeloc  = t['location']
+        typename = t['names']
+        length   = param_ref_length(p)
+
+        if typeloc < location
+          length += location - typeloc
+          location = typeloc
+        end
+
+        results << { 'location' => location, 'length' => length, 'typename' => typename }
+      end
+    end
+
+    results.sort_by! { |r| r['location'] }
+    results
+  end
+
+  private
+
+  def param_ref_length(paramref_node)
+    if paramref_node['number'] == 0 # rubocop:disable Style/NumericPredicate
+      1 # Actually a ? replacement character
+    else
+      ('$' + paramref_node['number'].to_s).size
+    end
+  end
+end

data/lib/pg_query/parse.rb

@@ -0,0 +1,247 @@
+require 'json'
+
+class PgQuery
+  def self.parse(query)
+    tree, stderr = _raw_parse(query)
+
+    begin
+      tree = JSON.parse(tree, max_nesting: 1000)
+    rescue JSON::ParserError
+      raise ParseError.new('Failed to parse JSON', __FILE__, __LINE__, -1)
+    end
+
+    warnings = []
+    stderr.each_line do |line|
+      next unless line[/^WARNING/]
+      warnings << line.strip
+    end
+
+    PgQuery.new(query, tree, warnings)
+  end
+
+  attr_reader :query
+  attr_reader :tree
+  attr_reader :warnings
+
+  def initialize(query, tree, warnings = [])
+    @query = query
+    @tree = tree
+    @warnings = warnings
+    @tables = nil
+    @aliases = nil
+    @cte_names = nil
+  end
+
+  def tables
+    tables_with_types.map { |t| t[:table] }
+  end
+
+  def select_tables
+    tables_with_types.select { |t| t[:type] == :select }.map { |t| t[:table] }
+  end
+
+  def dml_tables
+    tables_with_types.select { |t| t[:type] == :dml }.map { |t| t[:table] }
+  end
+
+  def ddl_tables
+    tables_with_types.select { |t| t[:type] == :ddl }.map { |t| t[:table] }
+  end
+
+  def cte_names
+    load_tables_and_aliases! if @cte_names.nil?
+    @cte_names
+  end
+
+  def aliases
+    load_tables_and_aliases! if @aliases.nil?
+    @aliases
+  end
+
+  def tables_with_types
+    load_tables_and_aliases! if @tables.nil?
+    @tables
+  end
+
+  protected
+
+  def load_tables_and_aliases! # rubocop:disable Metrics/CyclomaticComplexity
+    @tables = [] # types: select, dml, ddl
+    @cte_names = []
+    @aliases = {}
+
+    statements = @tree.dup
+    from_clause_items = [] # types: select, dml, ddl
+    subselect_items = []
+
+    loop do
+      statement = statements.shift
+      if statement
+        case statement.keys[0]
+        when RAW_STMT
+          statements << statement[RAW_STMT][STMT_FIELD]
+        # The following statement types do not modify tables and are added to from_clause_items
+        # (and subsequently @tables)
+        when SELECT_STMT
+          case statement[SELECT_STMT]['op']
+          when 0
+            (statement[SELECT_STMT][FROM_CLAUSE_FIELD] || []).each do |item|
+              if item[RANGE_SUBSELECT]
+                statements << item[RANGE_SUBSELECT]['subquery']
+              else
+                from_clause_items << { item: item, type: :select }
+              end
+            end
+          when 1
+            statements << statement[SELECT_STMT]['larg'] if statement[SELECT_STMT]['larg']
+            statements << statement[SELECT_STMT]['rarg'] if statement[SELECT_STMT]['rarg']
+          end
+
+          if (with_clause = statement[SELECT_STMT]['withClause'])
+            cte_statements, cte_names = statements_and_cte_names_for_with_clause(with_clause)
+            @cte_names.concat(cte_names)
+            statements.concat(cte_statements)
+          end
+        # The following statements modify the contents of a table
+        when INSERT_STMT, UPDATE_STMT, DELETE_STMT
+          value = statement.values[0]
+          from_clause_items << { item: value['relation'], type: :dml }
+          statements << value['selectStmt'] if value.key?('selectStmt')
+          statements << value['withClause'] if value.key?('withClause')
+
+          if (with_clause = value['withClause'])
+            cte_statements, cte_names = statements_and_cte_names_for_with_clause(with_clause)
+            @cte_names.concat(cte_names)
+            statements.concat(cte_statements)
+          end
+        when COPY_STMT
+          from_clause_items << { item: statement.values[0]['relation'], type: :dml } if statement.values[0]['relation']
+          statements << statement.values[0]['query']
+        # The following statement types are DDL (changing table structure)
+        when ALTER_TABLE_STMT, CREATE_STMT
+          from_clause_items << { item: statement.values[0]['relation'], type: :ddl }
+        when CREATE_TABLE_AS_STMT
+          if statement[CREATE_TABLE_AS_STMT]['into'] && statement[CREATE_TABLE_AS_STMT]['into'][INTO_CLAUSE]['rel']
+            from_clause_items << { item: statement[CREATE_TABLE_AS_STMT]['into'][INTO_CLAUSE]['rel'], type: :ddl }
+          end
+          if statement[CREATE_TABLE_AS_STMT]['query']
+            statements << statement[CREATE_TABLE_AS_STMT]['query']
+          end
+        when TRUNCATE_STMT
+          from_clause_items += statement.values[0]['relations'].map { |r| { item: r, type: :ddl } }
+        when VIEW_STMT
+          from_clause_items << { item: statement[VIEW_STMT]['view'], type: :ddl }
+          statements << statement[VIEW_STMT]['query']
+        when VACUUM_STMT, INDEX_STMT, CREATE_TRIG_STMT, RULE_STMT
+          from_clause_items << { item: statement.values[0]['relation'], type: :ddl }
+        when REFRESH_MAT_VIEW_STMT
+          from_clause_items << { item: statement[REFRESH_MAT_VIEW_STMT]['relation'], type: :ddl }
+        when DROP_STMT
+          objects = statement[DROP_STMT]['objects'].map do |obj|
+            if obj.is_a?(Array)
+              obj.map { |obj2| obj2['String'] && obj2['String']['str'] }
+            else
+              obj['String'] && obj['String']['str']
+            end
+          end
+          case statement[DROP_STMT]['removeType']
+          when OBJECT_TYPE_TABLE
+            @tables += objects.map { |r| { table: r.join('.'), type: :ddl } }
+          when OBJECT_TYPE_RULE, OBJECT_TYPE_TRIGGER
+            @tables += objects.map { |r| { table: r[0..-2].join('.'), type: :ddl } }
+          end
+        when GRANT_STMT
+          objects = statement[GRANT_STMT]['objects']
+          case statement[GRANT_STMT]['objtype']
+          when 0 # Column # rubocop:disable Lint/EmptyWhen
+            # FIXME
+          when 1 # Table
+            from_clause_items += objects.map { |o| { item: o, type: :ddl } }
+          when 2 # Sequence # rubocop:disable Lint/EmptyWhen
+            # FIXME
+          end
+        when LOCK_STMT
+          from_clause_items += statement.values[0]['relations'].map { |r| { item: r, type: :ddl } }
+        # The following are other statements that don't fit into query/DML/DDL
+        when EXPLAIN_STMT
+          statements << statement[EXPLAIN_STMT]['query']
+        end
+
+        statement_value = statement.values[0]
+        unless statement.empty?
+          subselect_items.concat(statement_value['targetList']) if statement_value['targetList']
+          subselect_items << statement_value['whereClause'] if statement_value['whereClause']
+          subselect_items.concat(statement_value['sortClause'].collect { |h| h[SORT_BY]['node'] }) if statement_value['sortClause']
+          subselect_items.concat(statement_value['groupClause']) if statement_value['groupClause']
+          subselect_items << statement_value['havingClause'] if statement_value['havingClause']
+        end
+      end
+
+      next_item = subselect_items.shift
+      if next_item
+        case next_item.keys[0]
+        when A_EXPR
+          %w[lexpr rexpr].each do |side|
+            elem = next_item.values[0][side]
+            next unless elem
+            if elem.is_a?(Array)
+              subselect_items += elem
+            else
+              subselect_items << elem
+            end
+          end
+        when BOOL_EXPR
+          subselect_items.concat(next_item.values[0]['args'])
+        when RES_TARGET
+          subselect_items << next_item[RES_TARGET]['val']
+        when SUB_LINK
+          statements << next_item[SUB_LINK]['subselect']
+        end
+      end
+
+      break if subselect_items.empty? && statements.empty?
+    end
+
+    loop do
+      next_item = from_clause_items.shift
+      break unless next_item && next_item[:item]
+
+      case next_item[:item].keys[0]
+      when JOIN_EXPR
+        %w[larg rarg].each do |side|
+          from_clause_items << { item: next_item[:item][JOIN_EXPR][side], type: next_item[:type] }
+        end
+      when ROW_EXPR
+        from_clause_items += next_item[:item][ROW_EXPR]['args'].map { |a| { item: a, type: next_item[:type] } }
+      when RANGE_VAR
+        rangevar = next_item[:item][RANGE_VAR]
+        next if !rangevar['schemaname'] && @cte_names.include?(rangevar['relname'])
+
+        table = [rangevar['schemaname'], rangevar['relname']].compact.join('.')
+        @tables << { table: table, type: next_item[:type] }
+        @aliases[rangevar['alias'][ALIAS]['aliasname']] = table if rangevar['alias']
+      when RANGE_SUBSELECT
+        from_clause_items << { item: next_item[:item][RANGE_SUBSELECT]['subquery'], type: next_item[:type] }
+      when SELECT_STMT
+        from_clause = next_item[:item][SELECT_STMT][FROM_CLAUSE_FIELD]
+        from_clause_items += from_clause.map { |r| { item: r, type: next_item[:type] } } if from_clause
+      end
+    end
+
+    @tables.uniq!
+    @cte_names.uniq!
+  end
+
+  def statements_and_cte_names_for_with_clause(with_clause)
+    statements = []
+    cte_names = []
+
+    with_clause[WITH_CLAUSE]['ctes'].each do |item|
+      next unless item[COMMON_TABLE_EXPR]
+      cte_names << item[COMMON_TABLE_EXPR]['ctename']
+      statements << item[COMMON_TABLE_EXPR]['ctequery']
+    end
+
+    [statements, cte_names]
+  end
+end

data/lib/pg_query/parse_error.rb

@@ -0,0 +1,9 @@
+class PgQuery
+  class ParseError < ArgumentError
+    attr_reader :location
+    def initialize(message, source_file, source_line, location)
+      super("#{message} (#{source_file}:#{source_line})")
+      @location = location
+    end
+  end
+end

data/lib/pg_query/treewalker.rb

@@ -0,0 +1,53 @@
+class PgQuery
+  private
+
+  def treewalker!(normalized_parsetree)
+    exprs = normalized_parsetree.dup.map { |e| [e, []] }
+
+    loop do
+      expr, parent_location = exprs.shift
+
+      if expr.is_a?(Hash)
+        expr.each do |k, v|
+          location = parent_location + [k]
+
+          yield(expr, k, v, location)
+
+          exprs << [v, location] unless v.nil?
+        end
+      elsif expr.is_a?(Array)
+        exprs += expr.map.with_index { |e, idx| [e, parent_location + [idx]] }
+      end
+
+      break if exprs.empty?
+    end
+  end
+
+  def find_tree_location(normalized_parsetree, searched_location)
+    treewalker! normalized_parsetree do |expr, k, v, location|
+      next unless location == searched_location
+      yield(expr, k, v)
+    end
+  end
+
+  def transform_nodes!(parsetree)
+    result = deep_dup(parsetree)
+    exprs = result.dup
+
+    loop do
+      expr = exprs.shift
+
+      if expr.is_a?(Hash)
+        yield(expr) if expr.size == 1 && expr.keys[0][/^[A-Z]+/]
+
+        exprs += expr.values.compact
+      elsif expr.is_a?(Array)
+        exprs += expr
+      end
+
+      break if exprs.empty?
+    end
+
+    result
+  end
+end

data/lib/pg_query/truncate.rb

@@ -0,0 +1,60 @@
+class PgQuery
+  PossibleTruncation = Struct.new(:location, :node_type, :length, :is_array)
+
+  A_TRUNCATED = 'A_Truncated'.freeze
+
+  # Truncates the query string to be below the specified length, first trying to
+  # omit less important parts of the query, and only then cutting off the end.
+  def truncate(max_length)
+    output = deparse(@tree)
+
+    # Early exit if we're already below the max length
+    return output if output.size <= max_length
+
+    truncations = find_possible_truncations
+
+    # Truncate the deepest possible truncation that is the longest first
+    truncations.sort_by! { |t| [-t.location.size, -t.length] }
+
+    tree = deep_dup(@tree)
+    truncations.each do |truncation|
+      next if truncation.length < 3
+
+      find_tree_location(tree, truncation.location) do |expr, k|
+        expr[k] = { A_TRUNCATED => nil }
+        expr[k] = [expr[k]] if truncation.is_array
+      end
+
+      output = deparse(tree)
+      return output if output.size <= max_length
+    end
+
+    # We couldn't do a proper smart truncation, so we need a hard cut-off
+    output[0..max_length - 4] + '...'
+  end
+
+  private
+
+  def find_possible_truncations
+    truncations = []
+
+    treewalker! @tree do |_expr, k, v, location|
+      case k
+      when TARGET_LIST_FIELD
+        length = deparse([{ SELECT_STMT => { k => v } }]).size - 7 # 'SELECT '.size
+
+        truncations << PossibleTruncation.new(location, TARGET_LIST_FIELD, length, true)
+      when 'whereClause'
+        length = deparse([{ SELECT_STMT => { k => v } }]).size
+
+        truncations << PossibleTruncation.new(location, 'whereClause', length, false)
+      when 'ctequery'
+        truncations << PossibleTruncation.new(location, 'ctequery', deparse([v]).size, false)
+      when 'cols'
+        truncations << PossibleTruncation.new(location, 'cols', deparse(v).size, true)
+      end
+    end
+
+    truncations
+  end
+end

data/lib/pg_query/version.rb

@@ -0,0 +1,3 @@
+class PgQuery
+  VERSION = '1.3.0'.freeze
+end

metadata

@@ -0,0 +1,130 @@
+--- !ruby/object:Gem::Specification
+name: gitlab-pg_query
+version: !ruby/object:Gem::Version
+  version: 1.3.0
+platform: ruby
+authors:
+- Lukas Fittl
+autorequire: 
+bindir: bin
+cert_chain: []
+date: 2020-10-23 00:00:00.000000000 Z
+dependencies:
+- !ruby/object:Gem::Dependency
+  name: rake-compiler
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '0'
+- !ruby/object:Gem::Dependency
+  name: rspec
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '3.0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '3.0'
+- !ruby/object:Gem::Dependency
+  name: rubocop
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - '='
+      - !ruby/object:Gem::Version
+        version: 0.49.1
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - '='
+      - !ruby/object:Gem::Version
+        version: 0.49.1
+- !ruby/object:Gem::Dependency
+  name: rubocop-rspec
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - '='
+      - !ruby/object:Gem::Version
+        version: 1.15.1
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - '='
+      - !ruby/object:Gem::Version
+        version: 1.15.1
+description: Parses SQL queries using a copy of the PostgreSQL server query parser
+email: lukas@fittl.com
+executables: []
+extensions:
+- ext/pg_query/extconf.rb
+extra_rdoc_files:
+- CHANGELOG.md
+- README.md
+files:
+- CHANGELOG.md
+- LICENSE
+- README.md
+- Rakefile
+- ext/pg_query/extconf.rb
+- ext/pg_query/pg_query_ruby.c
+- ext/pg_query/pg_query_ruby.h
+- ext/pg_query/pg_query_ruby.sym
+- lib/pg_query.rb
+- lib/pg_query/deep_dup.rb
+- lib/pg_query/deparse.rb
+- lib/pg_query/deparse/alter_table.rb
+- lib/pg_query/deparse/interval.rb
+- lib/pg_query/deparse/keywords.rb
+- lib/pg_query/deparse/rename.rb
+- lib/pg_query/filter_columns.rb
+- lib/pg_query/fingerprint.rb
+- lib/pg_query/legacy_parsetree.rb
+- lib/pg_query/node_types.rb
+- lib/pg_query/param_refs.rb
+- lib/pg_query/parse.rb
+- lib/pg_query/parse_error.rb
+- lib/pg_query/treewalker.rb
+- lib/pg_query/truncate.rb
+- lib/pg_query/version.rb
+homepage: http://github.com/lfittl/pg_query
+licenses:
+- BSD-3-Clause
+metadata: {}
+post_install_message: 
+rdoc_options:
+- "--main"
+- README.md
+- "--exclude"
+- ext/
+require_paths:
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '0'
+required_rubygems_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '0'
+requirements: []
+rubygems_version: 3.0.3
+signing_key: 
+specification_version: 4
+summary: PostgreSQL query parsing and normalization library
+test_files: []