Merge pull request rails#47762 from Shopify/ar-calculations-ids-for-a-cpk-model

Fix `ActiveRecord::Calculations#ids` for a composite primary key model

Author: eileencodes

activerecord/lib/active_record/relation/calculations.rb

@@ -310,7 +310,7 @@ def async_pick(*column_names)
     #   Person.joins(:companies).ids # SELECT people.id FROM people INNER JOIN companies ON companies.person_id = people.id
     def ids
       if loaded?
-        result = records.pluck(primary_key)
+        result = records.pluck(*Array(primary_key))
         return @async ? Promise::Complete.new(result) : result
       end
 
@@ -319,7 +319,7 @@ def ids
         return relation.ids
       end
 
-      columns = arel_columns([primary_key])
+      columns = arel_columns(Array(primary_key))
       relation = spawn
       relation.select_values = columns
       result = if relation.where_clause.contradiction?

activerecord/test/cases/calculations_test.rb

@@ -24,11 +24,12 @@
 require "models/too_long_table_name"
 require "support/stubs/strong_parameters"
 require "support/async_helper"
+require "models/cpk/book"
 
 class CalculationsTest < ActiveRecord::TestCase
   include AsyncHelper
 
-  fixtures :companies, :accounts, :authors, :author_addresses, :topics, :speedometers, :minivans, :books, :posts, :comments
+  fixtures :companies, :accounts, :authors, :author_addresses, :topics, :speedometers, :minivans, :books, :posts, :comments, :cpk_books
 
   def test_should_sum_field
     assert_equal 318, Account.sum(:credit_limit)
@@ -940,6 +941,25 @@ def test_ids
     assert_equal Company.all.map(&:id).sort, Company.all.ids.sort
   end
 
+  def ids_for_a_composite_primary_key
+    assert_equal Cpk::Book.all.map(&:id).sort, Cpk::Book.all.ids.sort
+  end
+
+  def test_ids_for_a_composite_primary_key_with_scope
+    book = cpk_books(:cpk_great_author_first_book)
+
+    assert_equal [[book.author_id, book.number]], Cpk::Book.all.where(title: book.title).ids
+  end
+
+  def test_ids_for_a_composite_primary_key_on_loaded_relation
+    book = cpk_books(:cpk_great_author_first_book)
+    relation = Cpk::Book.where(title: book.title)
+    relation.to_a
+
+    assert_predicate relation, :loaded?
+    assert_equal [[book.author_id, book.number]], relation.ids
+  end
+
   def test_ids_with_scope
     scoped_ids = [1, 2]
     assert_equal Company.where(id: scoped_ids).map(&:id).sort, Company.where(id: scoped_ids).ids.sort