Add Relation extensions

1 files changed, 21 insertions, 4 deletions

diff --git a/activerecord/lib/active_record/relation/query_methods.rb b/activerecord/lib/active_record/relation/query_methods.rb
index e224781016..b5e8b7570a 100644
--- a/activerecord/lib/active_record/relation/query_methods.rb
+++ b/activerecord/lib/active_record/relation/query_methods.rb
@@ -10,8 +10,9 @@ module ActiveRecord
 
         next if [:where, :having].include?(query_method)
         class_eval <<-CEVAL
-          def #{query_method}(*args)
+          def #{query_method}(*args, &block)
             new_relation = clone
+            new_relation.send(:apply_modules, Module.new(&block)) if block_given?
             value = Array.wrap(args.flatten).reject {|x| x.blank? }
             new_relation.#{query_method}_values += value if value.present?
             new_relation
@@ -21,8 +22,9 @@ module ActiveRecord
 
       [:where, :having].each do |query_method|
         class_eval <<-CEVAL
-          def #{query_method}(*args)
+          def #{query_method}(*args, &block)
             new_relation = clone
+            new_relation.send(:apply_modules, Module.new(&block)) if block_given?
             value = build_where(*args)
             new_relation.#{query_method}_values += [*value] if value.present?
             new_relation
@@ -34,8 +36,9 @@ module ActiveRecord
         attr_accessor :"#{query_method}_value"
 
         class_eval <<-CEVAL
-          def #{query_method}(value = true)
+          def #{query_method}(value = true, &block)
             new_relation = clone
+            new_relation.send(:apply_modules, Module.new(&block)) if block_given?
             new_relation.#{query_method}_value = value
             new_relation
           end
@@ -43,8 +46,16 @@ module ActiveRecord
       end
     end
 
-    def lock(locks = true)
+    def extending(*modules)
+      new_relation = clone
+      new_relation.send :apply_modules, *modules
+      new_relation
+    end
+
+    def lock(locks = true, &block)
       relation = clone
+      relation.send(:apply_modules, Module.new(&block)) if block_given?
+
       case locks
       when String, TrueClass, NilClass
         clone.tap {|new_relation| new_relation.lock_value = locks || true }
@@ -191,6 +202,12 @@ module ActiveRecord
 
     private
 
+    def apply_modules(modules)
+      values = Array.wrap(modules)
+      @extensions += values if values.present?
+      values.each {|extension| extend(extension) }
+    end
+
     def reverse_sql_order(order_query)
       order_query.to_s.split(/,/).each { |s|
         if s.match(/\s(asc|ASC)$/)