pg_query 2.0.1 → 2.1.1

data/lib/pg_query/parse.rb

@@ -4,8 +4,8 @@ module PgQuery
 
     begin
       result = PgQuery::ParseResult.decode(result)
-    rescue Google::Protobuf::ParseError
-      raise PgQuery::ParseError.new('Failed to parse tree', __FILE__, __LINE__, -1)
+    rescue Google::Protobuf::ParseError => e
+      raise PgQuery::ParseError.new(format('Failed to parse tree: %s', e.message), __FILE__, __LINE__, -1)
     end
 
     warnings = []
@@ -29,6 +29,7 @@ module PgQuery
       @tables = nil
       @aliases = nil
       @cte_names = nil
+      @functions = nil
     end
 
     def dup_tree
@@ -51,27 +52,52 @@ module PgQuery
       tables_with_details.select { |t| t[:type] == :ddl }.map { |t| t[:name] }.uniq
     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
+
+    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_tables_and_aliases! if @cte_names.nil?
+      load_objects! if @cte_names.nil?
       @cte_names
     end
 
     def aliases
-      load_tables_and_aliases! if @aliases.nil?
+      load_objects! if @aliases.nil?
       @aliases
     end
 
     def tables_with_details
-      load_tables_and_aliases! if @tables.nil?
+      load_objects! if @tables.nil?
       @tables
     end
 
+    def functions_with_details
+      load_objects! if @functions.nil?
+      @functions
+    end
+
     protected
 
-    def load_tables_and_aliases! # rubocop:disable Metrics/CyclomaticComplexity
+    # Parses the query and finds table and function references
+    #
+    # Note we use ".to_ary" on arrays from the Protobuf library before
+    # passing them to concat, because of https://bugs.ruby-lang.org/issues/18140
+    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
@@ -86,20 +112,16 @@ module PgQuery
           # 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.concat(statement.select_stmt.target_list.to_ary)
             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.concat(statement.select_stmt.group_clause.to_ary)
             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
+                from_clause_items << { item: item, type: :select }
               end
             when :SETOP_UNION
               statements << PgQuery::Node.new(select_stmt: statement.select_stmt.larg) if statement.select_stmt.larg
@@ -117,7 +139,13 @@ module PgQuery
             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
 
-            subselect_items.concat(statement.update_stmt.target_list) if statement.node == :update_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.to_ary)
+            end
+
             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
 
@@ -168,6 +196,10 @@ module PgQuery
               @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
@@ -184,12 +216,28 @@ module PgQuery
           # 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
         end
 
         next_item = subselect_items.shift
         if next_item
           case next_item.node
+          when :list
+            subselect_items += next_item.list.items
           when :a_expr
             %w[lexpr rexpr].each do |side|
               elem = next_item.a_expr.public_send(side)
@@ -201,47 +249,54 @@ module PgQuery
               end
             end
           when :bool_expr
-            subselect_items.concat(next_item.bool_expr.args)
+            subselect_items.concat(next_item.bool_expr.args.to_ary)
+          when :coalesce_expr
+            subselect_items.concat(next_item.coalesce_expr.args.to_ary)
+          when :min_max_expr
+            subselect_items.concat(next_item.min_max_expr.args.to_ary)
           when :res_target
             subselect_items << next_item.res_target.val
           when :sub_link
             statements << next_item.sub_link.subselect
+          when :func_call
+            subselect_items.concat(next_item.func_call.args.to_ary)
+            @functions << {
+              function: next_item.func_call.funcname.map { |f| f.string.str }.join('.'),
+              type: :call
+            }
           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].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)
-
-          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
+        if 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)
+
+            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
+            statements << next_item[:item].range_subselect.subquery
+          when :range_function
+            subselect_items += next_item[:item].range_function.functions
+          end
         end
+
+        break if subselect_items.empty? && statements.empty? && from_clause_items.empty?
       end
 
       @tables.uniq!