From 849dec458ef15794689224544a5c1a1f3fec4b9e Mon Sep 17 00:00:00 2001
From: Rick Song <ricksong@squareup.com>
Date: Thu, 28 Jan 2016 20:31:37 -0800
Subject: [PATCH] Makes a minor fix to PredicateBuilder respect custom primary
 keys when calling `Relation#where`

---
 activerecord/CHANGELOG.md                     |  6 +++
 .../relation/predicate_builder.rb             |  3 +-
 .../test/cases/relation/where_test.rb         | 46 +++++++++++++++++--
 3 files changed, 49 insertions(+), 6 deletions(-)

diff --git a/activerecord/CHANGELOG.md b/activerecord/CHANGELOG.md
index 0dc04845dee..f859e4dc07c 100644
--- a/activerecord/CHANGELOG.md
+++ b/activerecord/CHANGELOG.md
@@ -1,3 +1,9 @@
+*   Fixes custom primary keys for associations when calling `Relation#where`
+
+    Fixes #23327.
+
+    *Rick Song*
+
 *   No longer pass deprecated option `-i` to `pg_dump`.
 
     *Paul Sadauskas*
diff --git a/activerecord/lib/active_record/relation/predicate_builder.rb b/activerecord/lib/active_record/relation/predicate_builder.rb
index f98cd137a4c..f9f879018e7 100644
--- a/activerecord/lib/active_record/relation/predicate_builder.rb
+++ b/activerecord/lib/active_record/relation/predicate_builder.rb
@@ -56,7 +56,8 @@ module ActiveRecord
       # For polymorphic relationships, find the foreign key and type:
       # PriceEstimate.where(estimate_of: treasure)
       if klass && reflection = klass._reflect_on_association(column.to_sym)
-        if reflection.polymorphic? && base_class = polymorphic_base_class_from_value(value)
+        base_class = polymorphic_base_class_from_value(value)
+        if reflection.polymorphic? && base_class
           queries << build(table[reflection.foreign_type], base_class)
         end
 
diff --git a/activerecord/test/cases/relation/where_test.rb b/activerecord/test/cases/relation/where_test.rb
index d0445901dcd..9f348080838 100644
--- a/activerecord/test/cases/relation/where_test.rb
+++ b/activerecord/test/cases/relation/where_test.rb
@@ -2,6 +2,7 @@ require "cases/helper"
 require "models/author"
 require "models/binary"
 require "models/cake_designer"
+require "models/category"
 require "models/chef"
 require "models/comment"
 require "models/edge"
@@ -14,7 +15,7 @@ require "models/vertex"
 
 module ActiveRecord
   class WhereTest < ActiveRecord::TestCase
-    fixtures :posts, :edges, :authors, :binaries, :essays
+    fixtures :posts, :edges, :authors, :categories, :binaries, :essays
 
     def test_where_copies_bind_params
       author = authors(:david)
@@ -207,20 +208,55 @@ module ActiveRecord
     end
 
     def test_where_on_association_with_custom_primary_key
+      category = categories(:general)
+      essay = Essay.where(category: category).first
+
+      assert_equal essays(:david_modest_proposal), essay
+    end
+
+    def test_where_on_association_with_custom_primary_key_with_relation
+      category = categories(:general)
+      essay = Essay.where(category: Category.where(id: category.id)).first
+
+      assert_equal essays(:david_modest_proposal), essay
+    end
+
+    def test_where_on_association_with_relation_performs_subselect_not_two_queries
+      category = categories(:general)
+
+      assert_queries(1) do
+        Essay.where(category: Category.where(id: category.id)).to_a
+      end
+    end
+
+    def test_where_on_association_with_custom_primary_key_with_array_of_base
+      category = categories(:general)
+      essay = Essay.where(category: [category]).first
+
+      assert_equal essays(:david_modest_proposal), essay
+    end
+
+    def test_where_on_association_with_custom_primary_key_with_array_of_ids
+      essay = Essay.where(category: ["General"]).first
+
+      assert_equal essays(:david_modest_proposal), essay
+    end
+
+    def test_where_on_polymorphic_association_with_custom_primary_key
       author = authors(:david)
       essay = Essay.where(writer: author).first
 
       assert_equal essays(:david_modest_proposal), essay
     end
 
-    def test_where_on_association_with_custom_primary_key_with_relation
+    def test_where_on_polymorphic_association_with_custom_primary_key_with_relation
       author = authors(:david)
       essay = Essay.where(writer: Author.where(id: author.id)).first
 
       assert_equal essays(:david_modest_proposal), essay
     end
 
-    def test_where_on_association_with_relation_performs_subselect_not_two_queries
+    def test_where_on_polymorphic_association_with_relation_performs_subselect_not_two_queries
       author = authors(:david)
 
       assert_queries(1) do
@@ -228,14 +264,14 @@ module ActiveRecord
       end
     end
 
-    def test_where_on_association_with_custom_primary_key_with_array_of_base
+    def test_where_on_polymorphic_association_with_custom_primary_key_with_array_of_base
       author = authors(:david)
       essay = Essay.where(writer: [author]).first
 
       assert_equal essays(:david_modest_proposal), essay
     end
 
-    def test_where_on_association_with_custom_primary_key_with_array_of_ids
+    def test_where_on_polymorphic_association_with_custom_primary_key_with_array_of_ids
       essay = Essay.where(writer: ["David"]).first
 
       assert_equal essays(:david_modest_proposal), essay