Merge pull request #7273 from beerlington/foreign_key_model_queries

Convert model name to foreign key in queries

6 files changed, 151 insertions, 5 deletions

diff --git a/activerecord/CHANGELOG.md b/activerecord/CHANGELOG.md
index e94db776a5..818b69b6b8 100644
--- a/activerecord/CHANGELOG.md
+++ b/activerecord/CHANGELOG.md
@@ -1,5 +1,17 @@
 ## Rails 4.0.0 (unreleased) ##
 
+*   Accept belongs_to (including polymorphic) association keys in queries
+
+    The following queries are now equivalent:
+
+        Post.where(:author => author)
+        Post.where(:author_id => author)
+
+        PriceEstimate.where(:estimate_of => treasure)
+        PriceEstimate.where(:estimate_of_type => 'Treasure', :estimate_of_id => treasure)
+
+    *Peter Brown*
+
 *   Use native `mysqldump` command instead of `structure_dump` method
     when dumping the database structure to a sql file. Fixes #5547.
 
diff --git a/activerecord/lib/active_record/relation/predicate_builder.rb b/activerecord/lib/active_record/relation/predicate_builder.rb
index cb8f903474..593de71fb3 100644
--- a/activerecord/lib/active_record/relation/predicate_builder.rb
+++ b/activerecord/lib/active_record/relation/predicate_builder.rb
@@ -1,12 +1,25 @@
 module ActiveRecord
   class PredicateBuilder # :nodoc:
     def self.build_from_hash(engine, attributes, default_table)
-      attributes.map do |column, value|
+      queries = []
+
+      attributes.each do |column, value|
         table = default_table
 
         if value.is_a?(Hash)
           table = Arel::Table.new(column, engine)
-          value.map { |k,v| build(table[k.to_sym], v) }
+
+          value.each do |k,v|
+            if rk = find_reflection_key(column, v, v)
+              if rk[:foreign_type]
+                queries << build(table[rk[:foreign_type]], v.class.base_class)
+              end
+
+              k = rk[:foreign_key]
+            end
+
+            queries << build(table[k.to_sym], v)
+          end
         else
           column = column.to_s
 
@@ -15,9 +28,19 @@ module ActiveRecord
             table = Arel::Table.new(table_name, engine)
           end
 
-          build(table[column.to_sym], value)
+          if rk = find_reflection_key(column, engine, value)
+            if rk[:foreign_type]
+              queries << build(table[rk[:foreign_type]], value.class.base_class)
+            end
+
+            column = rk[:foreign_key]
+          end
+
+          queries << build(table[column.to_sym], value)
         end
-      end.flatten
+      end
+
+      queries
     end
 
     def self.references(attributes)
@@ -31,6 +54,27 @@ module ActiveRecord
       end.compact
     end
 
+    # Find the foreign key when using queries such as:
+    # Post.where(:author => author)
+    #
+    # For polymorphic relationships, find the foreign key and type:
+    # PriceEstimate.where(:estimate_of => treasure)
+    def self.find_reflection_key(parent_column, model, value)
+      # value must be an ActiveRecord object
+      return nil unless value.class < Model::Tag
+
+      if reflection = model.reflections[parent_column.to_sym]
+        if reflection.options[:polymorphic]
+          {
+            :foreign_key  => reflection.foreign_key,
+            :foreign_type => reflection.foreign_type
+          }
+        else
+          { :foreign_key => reflection.foreign_key }
+        end
+      end
+    end
+
     private
       def self.build(attribute, value)
         case value
diff --git a/activerecord/lib/active_record/relation/query_methods.rb b/activerecord/lib/active_record/relation/query_methods.rb
index f6bacf4822..00cdaf686f 100644
--- a/activerecord/lib/active_record/relation/query_methods.rb
+++ b/activerecord/lib/active_record/relation/query_methods.rb
@@ -340,6 +340,24 @@ module ActiveRecord
     #    User.where({ created_at: (Time.now.midnight - 1.day)..Time.now.midnight })
     #    # SELECT * FROM users WHERE (created_at BETWEEN '2012-06-09 07:00:00.000000' AND '2012-06-10 07:00:00.000000')
     #
+    # In the case of a belongs_to relationship, an association key can be used
+    # to specify the model if an ActiveRecord object is used as the value.
+    #
+    #    author = Author.find(1)
+    #
+    #    # The following queries will be equivalent:
+    #    Post.where(:author => author)
+    #    Post.where(:author_id => author)
+    #
+    # This also works with polymorphic belongs_to relationships:
+    #
+    #    treasure = Treasure.create(:name => 'gold coins')
+    #    treasure.price_estimates << PriceEstimate.create(:price => 125)
+    #
+    #    # The following queries will be equivalent:
+    #    PriceEstimate.where(:estimate_of => treasure)
+    #    PriceEstimate.where(:estimate_of_type => 'Treasure', :estimate_of_id => treasure)
+    #
     # === Joins
     #
     # If the relation is the result of a join, you may create a condition which uses any of the
diff --git a/activerecord/test/cases/relation/where_test.rb b/activerecord/test/cases/relation/where_test.rb
index 90c690e266..76339fe41b 100644
--- a/activerecord/test/cases/relation/where_test.rb
+++ b/activerecord/test/cases/relation/where_test.rb
@@ -1,9 +1,77 @@
 require "cases/helper"
+require 'models/author'
+require 'models/price_estimate'
+require 'models/treasure'
 require 'models/post'
 
 module ActiveRecord
   class WhereTest < ActiveRecord::TestCase
-    fixtures :posts
+    fixtures :posts, :authors
+
+    def test_belongs_to_shallow_where
+      author = Post.first.author
+      query_with_id = Post.where(:author_id => author)
+      query_with_assoc = Post.where(:author => author)
+
+      assert_equal query_with_id.to_sql, query_with_assoc.to_sql
+    end
+
+    def test_belongs_to_nested_where
+      author = Post.first.author
+      query_with_id = Author.where(:posts => {:author_id => author}).joins(:posts)
+      query_with_assoc = Author.where(:posts => {:author => author}).joins(:posts)
+
+      assert_equal query_with_id.to_sql, query_with_assoc.to_sql
+    end
+
+    def test_polymorphic_shallow_where
+      treasure = Treasure.create(:name => 'gold coins')
+      treasure.price_estimates << PriceEstimate.create(:price => 125)
+
+      query_by_column = PriceEstimate.where(:estimate_of_type => 'Treasure', :estimate_of_id => treasure)
+      query_by_model = PriceEstimate.where(:estimate_of => treasure)
+
+      assert_equal query_by_column.to_sql, query_by_model.to_sql
+    end
+
+    def test_polymorphic_sti_shallow_where
+      treasure = HiddenTreasure.create!(:name => 'gold coins')
+      treasure.price_estimates << PriceEstimate.create!(:price => 125)
+
+      query_by_column = PriceEstimate.where(:estimate_of_type => 'Treasure', :estimate_of_id => treasure)
+      query_by_model = PriceEstimate.where(:estimate_of => treasure)
+
+      assert_equal query_by_column.to_sql, query_by_model.to_sql
+    end
+
+    def test_polymorphic_nested_where
+      estimate = PriceEstimate.create :price => 125
+      treasure = Treasure.create :name => 'Booty'
+
+      treasure.price_estimates << estimate
+
+      query_by_column = Treasure.where(:price_estimates => {:estimate_of_type => 'Treasure', :estimate_of_id => treasure}).joins(:price_estimates)
+      query_by_model = Treasure.where(:price_estimates => {:estimate_of => treasure}).joins(:price_estimates)
+
+      assert_equal treasure, query_by_column.first
+      assert_equal treasure, query_by_model.first
+      assert_equal query_by_column.to_a, query_by_model.to_a
+    end
+
+    def test_polymorphic_sti_nested_where
+      estimate = PriceEstimate.create :price => 125
+      treasure = HiddenTreasure.create!(:name => 'gold coins')
+      treasure.price_estimates << PriceEstimate.create!(:price => 125)
+
+      treasure.price_estimates << estimate
+
+      query_by_column = Treasure.where(:price_estimates => {:estimate_of_type => 'Treasure', :estimate_of_id => treasure}).joins(:price_estimates)
+      query_by_model = Treasure.where(:price_estimates => {:estimate_of => treasure}).joins(:price_estimates)
+
+      assert_equal treasure, query_by_column.first
+      assert_equal treasure, query_by_model.first
+      assert_equal query_by_column.to_a, query_by_model.to_a
+    end
 
     def test_where_error
       assert_raises(ActiveRecord::StatementInvalid) do
diff --git a/activerecord/test/models/treasure.rb b/activerecord/test/models/treasure.rb
index 2a98e74f2c..e864295acf 100644
--- a/activerecord/test/models/treasure.rb
+++ b/activerecord/test/models/treasure.rb
@@ -6,3 +6,6 @@ class Treasure < ActiveRecord::Base
 
   accepts_nested_attributes_for :looter
 end
+
+class HiddenTreasure < Treasure
+end
diff --git a/activerecord/test/schema/schema.rb b/activerecord/test/schema/schema.rb
index b4e611cb09..798ea20efc 100644
--- a/activerecord/test/schema/schema.rb
+++ b/activerecord/test/schema/schema.rb
@@ -693,6 +693,7 @@ ActiveRecord::Schema.define do
 
   create_table :treasures, :force => true do |t|
     t.column :name, :string
+    t.column :type, :string
     t.column :looter_id, :integer
     t.column :looter_type, :string
   end