Merge pull request #35617 from mattyoho/add-annotation-support-to-relations

Add support for annotating queries generated by ActiveRecord::Relation with SQL comments

15 files changed, 118 insertions, 19 deletions

diff --git a/activerecord/lib/active_record/associations/preloader/through_association.rb b/activerecord/lib/active_record/associations/preloader/through_association.rb
index 32653956b2..db89b77629 100644
--- a/activerecord/lib/active_record/associations/preloader/through_association.rb
+++ b/activerecord/lib/active_record/associations/preloader/through_association.rb
@@ -61,11 +61,15 @@ module ActiveRecord
             scope = through_reflection.klass.unscoped
             options = reflection.options
 
+            values = reflection_scope.values
+            if annotations = values[:annotate]
+              scope.annotate!(*annotations)
+            end
+
             if options[:source_type]
               scope.where! reflection.foreign_type => options[:source_type]
             elsif !reflection_scope.where_clause.empty?
               scope.where_clause = reflection_scope.where_clause
-              values = reflection_scope.values
 
               if includes = values[:includes]
                 scope.includes!(source_reflection.name => includes)
diff --git a/activerecord/lib/active_record/querying.rb b/activerecord/lib/active_record/querying.rb
index 86021f0a80..81ab502824 100644
--- a/activerecord/lib/active_record/querying.rb
+++ b/activerecord/lib/active_record/querying.rb
@@ -15,7 +15,7 @@ module ActiveRecord
       :select, :reselect, :order, :reorder, :group, :limit, :offset, :joins, :left_joins, :left_outer_joins,
       :where, :rewhere, :preload, :eager_load, :includes, :from, :lock, :readonly, :extending, :or,
       :having, :create_with, :distinct, :references, :none, :unscope, :optimizer_hints, :merge, :except, :only,
-      :count, :average, :minimum, :maximum, :sum, :calculate,
+      :count, :average, :minimum, :maximum, :sum, :calculate, :annotate,
       :pluck, :pick, :ids
     ].freeze # :nodoc:
     delegate(*QUERYING_METHODS, to: :all)
diff --git a/activerecord/lib/active_record/relation.rb b/activerecord/lib/active_record/relation.rb
index 6d954a2b2e..36c2422d84 100644
--- a/activerecord/lib/active_record/relation.rb
+++ b/activerecord/lib/active_record/relation.rb
@@ -5,7 +5,7 @@ module ActiveRecord
   class Relation
     MULTI_VALUE_METHODS  = [:includes, :eager_load, :preload, :select, :group,
                             :order, :joins, :left_outer_joins, :references,
-                            :extending, :unscope, :optimizer_hints]
+                            :extending, :unscope, :optimizer_hints, :annotate]
 
     SINGLE_VALUE_METHODS = [:limit, :offset, :lock, :readonly, :reordering,
                             :reverse_order, :distinct, :create_with, :skip_query_cache]
@@ -389,6 +389,8 @@ module ActiveRecord
         stmt.set Arel.sql(klass.sanitize_sql_for_assignment(updates, table.name))
       end
 
+      stmt.comment(*arel.comment_node.values) if arel.comment_node
+
       @klass.connection.update stmt, "#{@klass} Update All"
     end
 
@@ -504,6 +506,7 @@ module ActiveRecord
       stmt.offset(arel.offset)
       stmt.order(*arel.orders)
       stmt.wheres = arel.constraints
+      stmt.comment(*arel.comment_node.values) if arel.comment_node
 
       affected = @klass.connection.delete(stmt, "#{@klass} Destroy")
 
diff --git a/activerecord/lib/active_record/relation/query_methods.rb b/activerecord/lib/active_record/relation/query_methods.rb
index 05ea0850d3..6f0f2125dc 100644
--- a/activerecord/lib/active_record/relation/query_methods.rb
+++ b/activerecord/lib/active_record/relation/query_methods.rb
@@ -349,7 +349,7 @@ module ActiveRecord
     end
 
     VALID_UNSCOPING_VALUES = Set.new([:where, :select, :group, :order, :lock,
-                                     :limit, :offset, :joins, :left_outer_joins,
+                                     :limit, :offset, :joins, :left_outer_joins, :annotate,
                                      :includes, :from, :readonly, :having, :optimizer_hints])
 
     # Removes an unwanted relation that is already defined on a chain of relations.
@@ -948,6 +948,26 @@ module ActiveRecord
       self
     end
 
+    # Adds an SQL comment to queries generated from this relation. For example:
+    #
+    #   User.annotate("selecting user names").select(:name)
+    #   # SELECT "users"."name" FROM "users" /* selecting user names */
+    #
+    #   User.annotate("selecting", "user", "names").select(:name)
+    #   # SELECT "users"."name" FROM "users" /* selecting */ /* user */ /* names */
+    #
+    # The SQL block comment delimiters, "/*" and "*/", will be added automatically.
+    def annotate(*args)
+      check_if_method_has_arguments!(:annotate, args)
+      spawn.annotate!(*args)
+    end
+
+    # Like #annotate, but modifies relation in place.
+    def annotate!(*args) # :nodoc:
+      self.annotate_values += args
+      self
+    end
+
     # Returns the Arel object associated with the relation.
     def arel(aliases = nil) # :nodoc:
       @arel ||= build_arel(aliases)
@@ -1004,6 +1024,7 @@ module ActiveRecord
         arel.distinct(distinct_value)
         arel.from(build_from) unless from_clause.empty?
         arel.lock(lock_value) if lock_value
+        arel.comment(*annotate_values) unless annotate_values.empty?
 
         arel
       end
diff --git a/activerecord/lib/arel/nodes.rb b/activerecord/lib/arel/nodes.rb
index 2f6dd9bc45..f994754620 100644
--- a/activerecord/lib/arel/nodes.rb
+++ b/activerecord/lib/arel/nodes.rb
@@ -61,6 +61,8 @@ require "arel/nodes/outer_join"
 require "arel/nodes/right_outer_join"
 require "arel/nodes/string_join"
 
+require "arel/nodes/comment"
+
 require "arel/nodes/sql_literal"
 
 require "arel/nodes/casted"
diff --git a/activerecord/lib/arel/nodes/comment.rb b/activerecord/lib/arel/nodes/comment.rb
new file mode 100644
index 0000000000..237ff27e7e
--- /dev/null
+++ b/activerecord/lib/arel/nodes/comment.rb
@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+module Arel # :nodoc: all
+  module Nodes
+    class Comment < Arel::Nodes::Node
+      attr_reader :values
+
+      def initialize(values)
+        super()
+        @values = values
+      end
+
+      def initialize_copy(other)
+        super
+        @values = @values.clone
+      end
+
+      def hash
+        [@values].hash
+      end
+
+      def eql?(other)
+        self.class == other.class &&
+          self.values == other.values
+      end
+      alias :== :eql?
+    end
+  end
+end
diff --git a/activerecord/lib/arel/nodes/delete_statement.rb b/activerecord/lib/arel/nodes/delete_statement.rb
index a419975335..56249b2bad 100644
--- a/activerecord/lib/arel/nodes/delete_statement.rb
+++ b/activerecord/lib/arel/nodes/delete_statement.rb
@@ -3,7 +3,7 @@
 module Arel # :nodoc: all
   module Nodes
     class DeleteStatement < Arel::Nodes::Node
-      attr_accessor :left, :right, :orders, :limit, :offset, :key
+      attr_accessor :left, :right, :orders, :limit, :offset, :key, :comment
 
       alias :relation :left
       alias :relation= :left=
@@ -18,16 +18,18 @@ module Arel # :nodoc: all
         @limit = nil
         @offset = nil
         @key = nil
+        @comment = nil
       end
 
       def initialize_copy(other)
         super
         @left  = @left.clone if @left
         @right = @right.clone if @right
+        @comment = @comment.clone if @comment
       end
 
       def hash
-        [self.class, @left, @right, @orders, @limit, @offset, @key].hash
+        [self.class, @left, @right, @orders, @limit, @offset, @key, @comment].hash
       end
 
       def eql?(other)
@@ -37,7 +39,8 @@ module Arel # :nodoc: all
           self.orders == other.orders &&
           self.limit == other.limit &&
           self.offset == other.offset &&
-          self.key == other.key
+          self.key == other.key &&
+          self.comment == other.comment
       end
       alias :== :eql?
     end
diff --git a/activerecord/lib/arel/nodes/insert_statement.rb b/activerecord/lib/arel/nodes/insert_statement.rb
index d28fd1f6c8..8430dd23da 100644
--- a/activerecord/lib/arel/nodes/insert_statement.rb
+++ b/activerecord/lib/arel/nodes/insert_statement.rb
@@ -3,7 +3,7 @@
 module Arel # :nodoc: all
   module Nodes
     class InsertStatement < Arel::Nodes::Node
-      attr_accessor :relation, :columns, :values, :select
+      attr_accessor :relation, :columns, :values, :select, :comment
 
       def initialize
         super()
@@ -11,6 +11,7 @@ module Arel # :nodoc: all
         @columns  = []
         @values   = nil
         @select   = nil
+        @comment  = nil
       end
 
       def initialize_copy(other)
@@ -18,10 +19,11 @@ module Arel # :nodoc: all
         @columns = @columns.clone
         @values =  @values.clone if @values
         @select =  @select.clone if @select
+        @comment =  @comment.clone if @comment
       end
 
       def hash
-        [@relation, @columns, @values, @select].hash
+        [@relation, @columns, @values, @select, @comment].hash
       end
 
       def eql?(other)
@@ -29,7 +31,8 @@ module Arel # :nodoc: all
           self.relation == other.relation &&
           self.columns == other.columns &&
           self.select == other.select &&
-          self.values == other.values
+          self.values == other.values &&
+          self.comment == other.comment
       end
       alias :== :eql?
     end
diff --git a/activerecord/lib/arel/nodes/select_core.rb b/activerecord/lib/arel/nodes/select_core.rb
index 5df6ac8412..b6154b7ff4 100644
--- a/activerecord/lib/arel/nodes/select_core.rb
+++ b/activerecord/lib/arel/nodes/select_core.rb
@@ -3,7 +3,7 @@
 module Arel # :nodoc: all
   module Nodes
     class SelectCore < Arel::Nodes::Node
-      attr_accessor :projections, :wheres, :groups, :windows
+      attr_accessor :projections, :wheres, :groups, :windows, :comment
       attr_accessor :havings, :source, :set_quantifier, :optimizer_hints
 
       def initialize
@@ -18,6 +18,7 @@ module Arel # :nodoc: all
         @groups          = []
         @havings         = []
         @windows         = []
+        @comment         = nil
       end
 
       def from
@@ -39,12 +40,13 @@ module Arel # :nodoc: all
         @groups      = @groups.clone
         @havings     = @havings.clone
         @windows     = @windows.clone
+        @comment     = @comment.clone if @comment
       end
 
       def hash
         [
           @source, @set_quantifier, @projections, @optimizer_hints,
-          @wheres, @groups, @havings, @windows
+          @wheres, @groups, @havings, @windows, @comment
         ].hash
       end
 
@@ -57,7 +59,8 @@ module Arel # :nodoc: all
           self.wheres == other.wheres &&
           self.groups == other.groups &&
           self.havings == other.havings &&
-          self.windows == other.windows
+          self.windows == other.windows &&
+          self.comment == other.comment
       end
       alias :== :eql?
     end
diff --git a/activerecord/lib/arel/nodes/update_statement.rb b/activerecord/lib/arel/nodes/update_statement.rb
index cfaa19e392..015bcd7613 100644
--- a/activerecord/lib/arel/nodes/update_statement.rb
+++ b/activerecord/lib/arel/nodes/update_statement.rb
@@ -3,7 +3,7 @@
 module Arel # :nodoc: all
   module Nodes
     class UpdateStatement < Arel::Nodes::Node
-      attr_accessor :relation, :wheres, :values, :orders, :limit, :offset, :key
+      attr_accessor :relation, :wheres, :values, :orders, :limit, :offset, :key, :comment
 
       def initialize
         @relation = nil
@@ -13,16 +13,18 @@ module Arel # :nodoc: all
         @limit    = nil
         @offset   = nil
         @key      = nil
+        @comment  = nil
       end
 
       def initialize_copy(other)
         super
         @wheres = @wheres.clone
         @values = @values.clone
+        @comment = @comment.clone if @comment
       end
 
       def hash
-        [@relation, @wheres, @values, @orders, @limit, @offset, @key].hash
+        [@relation, @wheres, @values, @orders, @limit, @offset, @key, @comment].hash
       end
 
       def eql?(other)
@@ -33,7 +35,8 @@ module Arel # :nodoc: all
           self.orders == other.orders &&
           self.limit == other.limit &&
           self.offset == other.offset &&
-          self.key == other.key
+          self.key == other.key &&
+          self.comment == other.comment
       end
       alias :== :eql?
     end
diff --git a/activerecord/lib/arel/select_manager.rb b/activerecord/lib/arel/select_manager.rb
index 32286b67f4..4e9f527235 100644
--- a/activerecord/lib/arel/select_manager.rb
+++ b/activerecord/lib/arel/select_manager.rb
@@ -244,6 +244,15 @@ module Arel # :nodoc: all
       @ctx.source
     end
 
+    def comment(*values)
+      @ctx.comment = Nodes::Comment.new(values)
+      self
+    end
+
+    def comment_node
+      @ctx.comment
+    end
+
     private
       def collapse(exprs)
         exprs = exprs.compact
diff --git a/activerecord/lib/arel/tree_manager.rb b/activerecord/lib/arel/tree_manager.rb
index 0476399618..326c4f995c 100644
--- a/activerecord/lib/arel/tree_manager.rb
+++ b/activerecord/lib/arel/tree_manager.rb
@@ -36,6 +36,11 @@ module Arel # :nodoc: all
         @ast.wheres << expr
         self
       end
+
+      def comment(*values)
+        @ast.comment = Nodes::Comment.new(values)
+        self
+      end
     end
 
     attr_reader :ast
diff --git a/activerecord/lib/arel/visitors/depth_first.rb b/activerecord/lib/arel/visitors/depth_first.rb
index 109afb7402..d696edc507 100644
--- a/activerecord/lib/arel/visitors/depth_first.rb
+++ b/activerecord/lib/arel/visitors/depth_first.rb
@@ -181,6 +181,10 @@ module Arel # :nodoc: all
           visit o.limit
         end
 
+        def visit_Arel_Nodes_Comment(o)
+          visit o.values
+        end
+
         def visit_Array(o)
           o.each { |i| visit i }
         end
diff --git a/activerecord/lib/arel/visitors/dot.rb b/activerecord/lib/arel/visitors/dot.rb
index 37803ce0c0..ecc386de07 100644
--- a/activerecord/lib/arel/visitors/dot.rb
+++ b/activerecord/lib/arel/visitors/dot.rb
@@ -234,6 +234,10 @@ module Arel # :nodoc: all
         end
         alias :visit_Set :visit_Array
 
+        def visit_Arel_Nodes_Comment(o)
+          visit_edge(o, "values")
+        end
+
         def visit_edge(o, method)
           edge(method) { visit o.send(method) }
         end
diff --git a/activerecord/lib/arel/visitors/to_sql.rb b/activerecord/lib/arel/visitors/to_sql.rb
index 1630226085..4192d9efdc 100644
--- a/activerecord/lib/arel/visitors/to_sql.rb
+++ b/activerecord/lib/arel/visitors/to_sql.rb
@@ -35,6 +35,7 @@ module Arel # :nodoc: all
           collect_nodes_for o.wheres, collector, " WHERE ", " AND "
           collect_nodes_for o.orders, collector, " ORDER BY "
           maybe_visit o.limit, collector
+          maybe_visit o.comment, collector
         end
 
         def visit_Arel_Nodes_UpdateStatement(o, collector)
@@ -47,6 +48,7 @@ module Arel # :nodoc: all
           collect_nodes_for o.wheres, collector, " WHERE ", " AND "
           collect_nodes_for o.orders, collector, " ORDER BY "
           maybe_visit o.limit, collector
+          maybe_visit o.comment, collector
         end
 
         def visit_Arel_Nodes_InsertStatement(o, collector)
@@ -62,9 +64,9 @@ module Arel # :nodoc: all
             maybe_visit o.values, collector
           elsif o.select
             maybe_visit o.select, collector
-          else
-            collector
           end
+
+          maybe_visit o.comment, collector
         end
 
         def visit_Arel_Nodes_Exists(o, collector)
@@ -162,7 +164,7 @@ module Arel # :nodoc: all
           collect_nodes_for o.havings, collector, " HAVING ", " AND "
           collect_nodes_for o.windows, collector, " WINDOW "
 
-          collector
+          maybe_visit o.comment, collector
         end
 
         def visit_Arel_Nodes_OptimizerHints(o, collector)
@@ -170,6 +172,10 @@ module Arel # :nodoc: all
           collector << "/*+ #{hints} */"
         end
 
+        def visit_Arel_Nodes_Comment(o, collector)
+          collector << o.values.map { |v| "/* #{sanitize_as_sql_comment(v)} */" }.join(" ")
+        end
+
         def collect_nodes_for(nodes, collector, spacer, connector = ", ")
           unless nodes.empty?
             collector << spacer