pg_query 1.3.0 → 2.1.0

data/lib/pg_query/param_refs.rb

@@ -1,45 +1,50 @@
-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
+module PgQuery
+  class ParserResult
+    def param_refs # rubocop:disable Metrics/CyclomaticComplexity
+      results = []
+
+      treewalker! @tree do |_, _, node, location|
+        case node
+        when PgQuery::ParamRef
+          # Ignore param refs inside type casts, as these are already handled
+          next if location[-3..-1] == %i[type_cast arg param_ref]
+
+          results << { 'location' => node.location,
+                       'length' => param_ref_length(node) }
+        when PgQuery::TypeCast
+          next unless node.arg && node.type_name
+
+          p = node.arg.param_ref
+          t = node.type_name
+          next unless p && t
+
+          location = p.location
+          typeloc  = t.location
+          length   = param_ref_length(p)
+
+          if location == -1
+            location = typeloc
+          elsif typeloc < location
+            length += location - typeloc
+            location = typeloc
+          end
+
+          results << { 'location' => location, 'length' => length, 'typename' => t.names.map { |n| n.string.str } }
         end
-
-        results << { 'location' => location, 'length' => length, 'typename' => typename }
       end
-    end
 
-    results.sort_by! { |r| r['location'] }
-    results
-  end
+      results.sort_by! { |r| r['location'] }
+      results
+    end
 
-  private
+    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
+    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
 end

data/lib/pg_query/parse.rb

@@ -1,13 +1,11 @@
-require 'json'
-
-class PgQuery
+module PgQuery
   def self.parse(query)
-    tree, stderr = _raw_parse(query)
+    result, stderr = parse_protobuf(query)
 
     begin
-      tree = JSON.parse(tree, max_nesting: 1000)
-    rescue JSON::ParserError
-      raise ParseError.new('Failed to parse JSON', __FILE__, __LINE__, -1)
+      result = PgQuery::ParseResult.decode(result)
+    rescue Google::Protobuf::ParseError => e
+      raise PgQuery::ParseError.new(format('Failed to parse tree: %s', e.message), __FILE__, __LINE__, -1)
     end
 
     warnings = []
@@ -16,232 +14,306 @@ class PgQuery
       warnings << line.strip
     end
 
-    PgQuery.new(query, tree, warnings)
+    PgQuery::ParserResult.new(query, result, 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
+  class ParserResult
+    attr_reader :query
+    attr_reader :tree
+    attr_reader :warnings
 
-  def tables
-    tables_with_types.map { |t| t[:table] }
-  end
+    def initialize(query, tree, warnings = [])
+      @query = query
+      @tree = tree
+      @warnings = warnings
+      @tables = nil
+      @aliases = nil
+      @cte_names = nil
+      @functions = nil
+    end
 
-  def select_tables
-    tables_with_types.select { |t| t[:type] == :select }.map { |t| t[:table] }
-  end
+    def dup_tree
+      ParseResult.decode(ParseResult.encode(@tree))
+    end
 
-  def dml_tables
-    tables_with_types.select { |t| t[:type] == :dml }.map { |t| t[:table] }
-  end
+    def tables
+      tables_with_details.map { |t| t[:name] }.uniq
+    end
 
-  def ddl_tables
-    tables_with_types.select { |t| t[:type] == :ddl }.map { |t| t[:table] }
-  end
+    def select_tables
+      tables_with_details.select { |t| t[:type] == :select }.map { |t| t[:name] }.uniq
+    end
 
-  def cte_names
-    load_tables_and_aliases! if @cte_names.nil?
-    @cte_names
-  end
+    def dml_tables
+      tables_with_details.select { |t| t[:type] == :dml }.map { |t| t[:name] }.uniq
+    end
 
-  def aliases
-    load_tables_and_aliases! if @aliases.nil?
-    @aliases
-  end
+    def ddl_tables
+      tables_with_details.select { |t| t[:type] == :ddl }.map { |t| t[:name] }.uniq
+    end
 
-  def tables_with_types
-    load_tables_and_aliases! if @tables.nil?
-    @tables
-  end
+    # Returns function names, ignoring their argument types. This may be insufficient
+    # if you need to disambiguate two functions with the same name but different argument
+    # types.
+    def functions
+      functions_with_details.map { |f| f[:function] }.uniq
+    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
+    def ddl_functions
+      functions_with_details.select { |f| f[:type] == :ddl }.map { |f| f[:function] }.uniq
+    end
+
+    def call_functions
+      functions_with_details.select { |f| f[:type] == :call }.map { |f| f[:function] }.uniq
+    end
+
+    def cte_names
+      load_objects! if @cte_names.nil?
+      @cte_names
+    end
+
+    def aliases
+      load_objects! if @aliases.nil?
+      @aliases
+    end
+
+    def tables_with_details
+      load_objects! if @tables.nil?
+      @tables
+    end
+
+    def functions_with_details
+      load_objects! if @functions.nil?
+      @functions
+    end
+
+    protected
+
+    def load_objects! # rubocop:disable Metrics/CyclomaticComplexity
+      @tables = [] # types: select, dml, ddl
+      @cte_names = []
+      @aliases = {}
+      @functions = [] # types: call, ddl
+
+      statements = @tree.stmts.dup.to_a.map(&:stmt)
+      from_clause_items = [] # types: select, dml, ddl
+      subselect_items = []
+
+      loop do
+        statement = statements.shift
+        if statement
+          case statement.node
+          when :list
+            statements += statement.list.items
+          # The following statement types do not modify tables and are added to from_clause_items
+          # (and subsequently @tables)
+          when :select_stmt
+            subselect_items.concat(statement.select_stmt.target_list)
+            subselect_items << statement.select_stmt.where_clause if statement.select_stmt.where_clause
+            subselect_items.concat(statement.select_stmt.sort_clause.collect { |h| h.sort_by.node })
+            subselect_items.concat(statement.select_stmt.group_clause)
+            subselect_items << statement.select_stmt.having_clause if statement.select_stmt.having_clause
+
+            case statement.select_stmt.op
+            when :SETOP_NONE
+              (statement.select_stmt.from_clause || []).each do |item|
+                if item.node == :range_subselect
+                  statements << item.range_subselect.subquery
+                else
+                  from_clause_items << { item: item, type: :select }
+                end
+              end
+            when :SETOP_UNION
+              statements << PgQuery::Node.new(select_stmt: statement.select_stmt.larg) if statement.select_stmt.larg
+              statements << PgQuery::Node.new(select_stmt: statement.select_stmt.rarg) if statement.select_stmt.rarg
+            end
+
+            if statement.select_stmt.with_clause
+              cte_statements, cte_names = statements_and_cte_names_for_with_clause(statement.select_stmt.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.public_send(statement.node)
+            from_clause_items << { item: PgQuery::Node.new(range_var: value.relation), type: :dml }
+            statements << value.select_stmt if statement.node == :insert_stmt && value.select_stmt
+
+            if statement.node == :update_stmt
+              value.from_clause.each do |item|
                 from_clause_items << { item: item, type: :select }
               end
+              subselect_items.concat(statement.update_stmt.target_list)
             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']
+            subselect_items << statement.update_stmt.where_clause if statement.node == :update_stmt && statement.update_stmt.where_clause
+            subselect_items << statement.delete_stmt.where_clause if statement.node == :delete_stmt && statement.delete_stmt.where_clause
+
+            if value.with_clause
+              cte_statements, cte_names = statements_and_cte_names_for_with_clause(value.with_clause)
+              @cte_names.concat(cte_names)
+              statements.concat(cte_statements)
+            end
+          when :copy_stmt
+            from_clause_items << { item: PgQuery::Node.new(range_var: statement.copy_stmt.relation), type: :dml } if statement.copy_stmt.relation
+            statements << statement.copy_stmt.query
+          # The following statement types are DDL (changing table structure)
+          when :alter_table_stmt
+            from_clause_items << { item: PgQuery::Node.new(range_var: statement.alter_table_stmt.relation), type: :ddl }
+          when :create_stmt
+            from_clause_items << { item: PgQuery::Node.new(range_var: statement.create_stmt.relation), type: :ddl }
+          when :create_table_as_stmt
+            if statement.create_table_as_stmt.into && statement.create_table_as_stmt.into.rel
+              from_clause_items << { item: PgQuery::Node.new(range_var: statement.create_table_as_stmt.into.rel), type: :ddl }
+            end
+            statements << statement.create_table_as_stmt.query if statement.create_table_as_stmt.query
+          when :truncate_stmt
+            from_clause_items += statement.truncate_stmt.relations.map { |r| { item: r, type: :ddl } }
+          when :view_stmt
+            from_clause_items << { item: PgQuery::Node.new(range_var: statement.view_stmt.view), type: :ddl }
+            statements << statement.view_stmt.query
+          when :index_stmt
+            from_clause_items << { item: PgQuery::Node.new(range_var: statement.index_stmt.relation), type: :ddl }
+          when :create_trig_stmt
+            from_clause_items << { item: PgQuery::Node.new(range_var: statement.create_trig_stmt.relation), type: :ddl }
+          when :rule_stmt
+            from_clause_items << { item: PgQuery::Node.new(range_var: statement.rule_stmt.relation), type: :ddl }
+          when :vacuum_stmt
+            from_clause_items += statement.vacuum_stmt.rels.map { |r| { item: PgQuery::Node.new(range_var: r.vacuum_relation.relation), type: :ddl } if r.node == :vacuum_relation }
+          when :refresh_mat_view_stmt
+            from_clause_items << { item: PgQuery::Node.new(range_var: statement.refresh_mat_view_stmt.relation), type: :ddl }
+          when :drop_stmt
+            objects = statement.drop_stmt.objects.map do |obj|
+              case obj.node
+              when :list
+                obj.list.items.map { |obj2| obj2.string.str if obj2.node == :string }
+              when :string
+                obj.string.str
+              end
+            end
+            case statement.drop_stmt.remove_type
+            when :OBJECT_TABLE
+              @tables += objects.map { |r| { name: r.join('.'), type: :ddl } }
+            when :OBJECT_RULE, :OBJECT_TRIGGER
+              @tables += objects.map { |r| { name: r[0..-2].join('.'), type: :ddl } }
+            when :OBJECT_FUNCTION
+              # Only one function can be dropped in a statement
+              obj = statement.drop_stmt.objects[0].object_with_args
+              @functions << { function: obj.objname[0].string.str, type: :ddl }
+            end
+          when :grant_stmt
+            objects = statement.grant_stmt.objects
+            case statement.grant_stmt.objtype
+            when :OBJECT_COLUMN # Column # rubocop:disable Lint/EmptyWhen
+              # FIXME
+            when :OBJECT_TABLE # Table
+              from_clause_items += objects.map { |o| { item: o, type: :ddl } }
+            when :OBJECT_SEQUENCE # Sequence # rubocop:disable Lint/EmptyWhen
+              # FIXME
+            end
+          when :lock_stmt
+            from_clause_items += statement.lock_stmt.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
+          when :create_function_stmt
+            @functions << {
+              function: statement.create_function_stmt.funcname[0].string.str,
+              type: :ddl
+            }
+          when :rename_stmt
+            if statement.rename_stmt.rename_type == :OBJECT_FUNCTION
+              original_name = statement.rename_stmt.object.object_with_args.objname[0].string.str
+              new_name = statement.rename_stmt.newname
+              @functions += [
+                { function: original_name, type: :ddl },
+                { function: new_name, type: :ddl }
+              ]
             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
+        next_item = subselect_items.shift
+        if next_item
+          case next_item.node
+          when :a_expr
+            %w[lexpr rexpr].each do |side|
+              elem = next_item.a_expr.public_send(side)
+              next unless elem
+              if elem.is_a?(Array) # FIXME: this needs to traverse a list
+                subselect_items += elem
+              else
+                subselect_items << elem
+              end
             end
+          when :bool_expr
+            subselect_items.concat(next_item.bool_expr.args)
+          when :coalesce_expr
+            subselect_items.concat(next_item.coalesce_expr.args)
+          when :min_max_expr
+            subselect_items.concat(next_item.min_max_expr.args)
+          when :res_target
+            subselect_items << next_item.res_target.val
+          when :sub_link
+            statements << next_item.sub_link.subselect
+          when :func_call
+            @functions << {
+              function: next_item.func_call.funcname[0].string.str,
+              type: :call
+            }
           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
+
+        break if subselect_items.empty? && statements.empty?
       end
 
-      break if subselect_items.empty? && statements.empty?
-    end
+      loop do
+        next_item = from_clause_items.shift
+        break unless next_item && next_item[:item]
 
-    loop do
-      next_item = from_clause_items.shift
-      break unless next_item && next_item[:item]
+        case next_item[:item].node
+        when :join_expr
+          from_clause_items << { item: next_item[:item].join_expr.larg, type: next_item[:type] }
+          from_clause_items << { item: next_item[:item].join_expr.rarg, type: next_item[:type] }
+        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.empty? && @cte_names.include?(rangevar.relname)
 
-      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] }
+          table = [rangevar.schemaname, rangevar.relname].reject { |s| s.nil? || s.empty? }.join('.')
+          @tables << {
+            name: table,
+            type: next_item[:type],
+            location: rangevar.location,
+            schemaname: (rangevar.schemaname unless rangevar.schemaname.empty?),
+            relname: rangevar.relname,
+            inh: rangevar.inh
+          }
+          @aliases[rangevar.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
+          from_clause_items += from_clause.map { |r| { item: r, type: next_item[:type] } } if from_clause
         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
+
+      @tables.uniq!
+      @cte_names.uniq!
     end
 
-    @tables.uniq!
-    @cte_names.uniq!
-  end
+    def statements_and_cte_names_for_with_clause(with_clause) # FIXME
+      statements = []
+      cte_names = []
 
-  def statements_and_cte_names_for_with_clause(with_clause)
-    statements = []
-    cte_names = []
+      with_clause.ctes.each do |item|
+        next unless item.node == :common_table_expr
+        cte_names << item.common_table_expr.ctename
+        statements << item.common_table_expr.ctequery
+      end
 
-    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']
+      [statements, cte_names]
     end
-
-    [statements, cte_names]
   end
 end