Respect custom primary keys for associations in `Relation#where`

While we query the proper columns, we go through normal handling for
converting the value to a primitive which assumes it should use the
table's primary key. If the association specifies a different value (and
we know that we're working with an association), we should use the
custom primary key instead.

Fixes #18813.

4 files changed, 65 insertions, 2 deletions

diff --git a/activerecord/CHANGELOG.md b/activerecord/CHANGELOG.md
index 0105dfa78c..8808574df5 100644
--- a/activerecord/CHANGELOG.md
+++ b/activerecord/CHANGELOG.md
@@ -1,3 +1,9 @@
+*   Respect custom primary keys for associations when calling `Relation#where`
+
+    Fixes #18813.
+
+    *Sean Griffin*
+
 *   Fixed several edge cases which could result in a counter cache updating
     twice or not updating at all for `has_many` and `has_many :through`.
 
diff --git a/activerecord/lib/active_record/relation/predicate_builder/association_query_handler.rb b/activerecord/lib/active_record/relation/predicate_builder/association_query_handler.rb
index aabcf20c1d..159889d3b8 100644
--- a/activerecord/lib/active_record/relation/predicate_builder/association_query_handler.rb
+++ b/activerecord/lib/active_record/relation/predicate_builder/association_query_handler.rb
@@ -31,7 +31,14 @@ module ActiveRecord
       end
 
       def ids
-        value
+        case value
+        when Relation
+          value.select(primary_key)
+        when Array
+          value.map { |v| convert_to_id(v) }
+        else
+          convert_to_id(value)
+        end
       end
 
       def base_class
@@ -42,6 +49,10 @@ module ActiveRecord
 
       private
 
+      def primary_key
+        associated_table.association_primary_key(base_class)
+      end
+
       def polymorphic_base_class_from_value
         case value
         when Relation
@@ -53,6 +64,15 @@ module ActiveRecord
           value.class.base_class
         end
       end
+
+      def convert_to_id(value)
+        case value
+        when Base
+          value._read_attribute(primary_key)
+        else
+          value
+        end
+      end
     end
   end
 end
diff --git a/activerecord/lib/active_record/table_metadata.rb b/activerecord/lib/active_record/table_metadata.rb
index 6c8792ee80..3dd6321a97 100644
--- a/activerecord/lib/active_record/table_metadata.rb
+++ b/activerecord/lib/active_record/table_metadata.rb
@@ -1,6 +1,7 @@
 module ActiveRecord
   class TableMetadata # :nodoc:
     delegate :foreign_type, :foreign_key, to: :association, prefix: true
+    delegate :association_primary_key, to: :association
 
     def initialize(klass, arel_table, association = nil)
       @klass = klass
diff --git a/activerecord/test/cases/relation/where_test.rb b/activerecord/test/cases/relation/where_test.rb
index 537937decd..6af31017d6 100644
--- a/activerecord/test/cases/relation/where_test.rb
+++ b/activerecord/test/cases/relation/where_test.rb
@@ -5,6 +5,7 @@ require "models/cake_designer"
 require "models/chef"
 require "models/comment"
 require "models/edge"
+require "models/essay"
 require "models/post"
 require "models/price_estimate"
 require "models/topic"
@@ -13,7 +14,7 @@ require "models/vertex"
 
 module ActiveRecord
   class WhereTest < ActiveRecord::TestCase
-    fixtures :posts, :edges, :authors, :binaries
+    fixtures :posts, :edges, :authors, :binaries, :essays
 
     def test_where_copies_bind_params
       author = authors(:david)
@@ -240,5 +241,40 @@ module ActiveRecord
       count = Binary.where(:data => 0).count
       assert_equal 0, count
     end
+
+    def test_where_on_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
+      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
+      author = authors(:david)
+
+      assert_queries(1) do
+        Essay.where(writer: Author.where(id: author.id)).to_a
+      end
+    end
+
+    def test_where_on_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
+      essay = Essay.where(writer: ["David"]).first
+
+      assert_equal essays(:david_modest_proposal), essay
+    end
   end
 end