Merge pull request #48761 from adrianna-chang-shopify/ac-fix-include-method-cpk

Support `include?` and `member?` on composite primary key relations

activerecord/lib/active_record/relation/finder_methods.rb

@@ -355,7 +355,12 @@ module ActiveRecord
       if loaded? || offset_value || limit_value || having_clause.any?
         records.include?(record)
       else
-        record.is_a?(klass) && exists?(record.id)
+        id = if record.class.composite_primary_key?
+          record.class.primary_key.zip(record.id).to_h
+        else
+          record.id
+        end
+        record.is_a?(klass) && exists?(id)
       end
     end
 

activerecord/test/cases/finder_test.rb

@@ -475,6 +475,22 @@ class FinderTest < ActiveRecord::TestCase
     end
   end
 
+  def test_include_on_unloaded_relation_with_composite_primary_key
+    assert_sql(/1 AS one.*LIMIT/) do
+      book = cpk_books(:cpk_great_author_first_book)
+      assert Cpk::Book.where(title: "The first book").include?(book)
+    end
+  end
+
+  def test_include_on_loaded_relation_with_composite_primary_key
+    books = Cpk::Book.where(title: "The first book").load
+    great_author_book = cpk_books(:cpk_great_author_first_book)
+
+    assert_no_queries do
+      assert books.include?(great_author_book)
+    end
+  end
+
   def test_member_on_unloaded_relation_with_match
     assert_sql(/1 AS one.*LIMIT/) do
       assert_equal true, Customer.where(name: "David").member?(customers(:david))
@@ -530,6 +546,22 @@ class FinderTest < ActiveRecord::TestCase
     end
   end
 
+  def test_member_on_unloaded_relation_with_composite_primary_key
+    assert_sql(/1 AS one.*LIMIT/) do
+      book = cpk_books(:cpk_great_author_first_book)
+      assert Cpk::Book.where(title: "The first book").member?(book)
+    end
+  end
+
+  def test_member_on_loaded_relation_with_composite_primary_key
+    books = Cpk::Book.where(title: "The first book").load
+    great_author_book = cpk_books(:cpk_great_author_first_book)
+
+    assert_no_queries do
+      assert books.member?(great_author_book)
+    end
+  end
+
   def test_find_by_array_of_one_id
     assert_kind_of(Array, Topic.find([ 1 ]))
     assert_equal(1, Topic.find([ 1 ]).length)