Ensure that ids_writer and ids_reader is working for CPK

activerecord/lib/active_record/associations/collection_association.rb

@@ -63,8 +63,16 @@ def ids_writer(ids)
         ids = Array(ids).compact_blank
         ids.map! { |i| pk_type.cast(i) }
 
-        records = klass.where(primary_key => ids).index_by do |r|
-          r.public_send(primary_key)
+        records = if klass.composite_primary_key?
+          query_records = ids.map { |values_set| klass.where(primary_key.zip(values_set).to_h) }.inject(&:or)
+
+          query_records.index_by do |r|
+            primary_key.map { |pk| r.public_send(pk) }
+          end
+        else
+          klass.where(primary_key => ids).index_by do |r|
+            r.public_send(primary_key)
+          end
         end.values_at(*ids).compact
 
         if records.size != ids.size

activerecord/lib/active_record/relation/calculations.rb

@@ -269,7 +269,7 @@ def pluck(*column_names)
         relation = apply_join_dependency
         relation.pluck(*column_names)
       else
-        klass.disallow_raw_sql!(column_names)
+        klass.disallow_raw_sql!(column_names.flatten)
         columns = arel_columns(column_names)
         relation = spawn
         relation.select_values = columns

activerecord/test/cases/autosave_association_test.rb

@@ -39,6 +39,7 @@
 require "models/chef"
 require "models/cake_designer"
 require "models/drink_designer"
+require "models/cpk"
 
 class TestAutosaveAssociationsInGeneral < ActiveRecord::TestCase
   def test_autosave_works_even_when_other_callbacks_update_the_parent_model
@@ -596,7 +597,7 @@ def test_valid_adding_with_nested_attributes
 end
 
 class TestDefaultAutosaveAssociationOnAHasManyAssociation < ActiveRecord::TestCase
-  fixtures :companies, :developers
+  fixtures :companies, :developers, :cpk_order_agreements, :cpk_orders, :cpk_books
 
   def test_invalid_adding
     firm = Firm.find(1)
@@ -730,6 +731,37 @@ def test_assign_ids
     assert_includes firm.clients, companies(:second_client)
   end
 
+  def test_assign_ids_with_belongs_to_cpk_model
+    order_agreements = [cpk_order_agreements(:order_agreement_one).id, cpk_order_agreements(:order_agreement_two).id]
+    order = cpk_orders(:cpk_groceries_order_1)
+
+    assert_empty order.order_agreements
+
+    order.order_agreement_ids = order_agreements
+    order.save
+    order.reload
+
+    assert_equal order_agreements, order.order_agreement_ids
+    assert_equal 2, order.order_agreements.length
+    assert_includes order.order_agreements, cpk_order_agreements(:order_agreement_two)
+  end
+
+  def test_assign_ids_with_cpk_for_two_models
+    books = [cpk_books(:cpk_great_author_first_book).id, cpk_books(:cpk_great_author_second_book).id]
+    order = cpk_orders(:cpk_groceries_order_1)
+
+    assert_empty order.books
+
+    order.book_ids = books
+    order.save
+    order.reload
+
+    assert_equal books, order.book_ids
+    assert_equal 2, order.books.length
+    assert_includes order.books, cpk_books(:cpk_great_author_first_book)
+    assert_includes order.books, cpk_books(:cpk_great_author_second_book)
+  end
+
   def test_assign_ids_for_through_a_belongs_to
     firm = Firm.new("name" => "Apple")
     firm.developer_ids = [developers(:david).id, developers(:jamis).id]

activerecord/test/cases/calculations_test.rb

@@ -941,10 +941,14 @@ def test_ids
     assert_equal Company.all.map(&:id).sort, Company.all.ids.sort
   end
 
-  def ids_for_a_composite_primary_key
+  def test_ids_for_a_composite_primary_key
     assert_equal Cpk::Book.all.map(&:id).sort, Cpk::Book.all.ids.sort
   end
 
+  def test_pluck_for_a_composite_primary_key
+    assert_equal Cpk::Book.all.pluck([:author_id, :number]).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)
 

activerecord/test/fixtures/cpk_order_agreements.yml

@@ -0,0 +1,8 @@
+_fixture:
+  model_class: Cpk::OrderAgreement
+
+order_agreement_one:
+  signature: "abc123"
+
+order_agreement_two:
+  signature: "xyz789"

activerecord/test/models/cpk/book.rb

@@ -4,6 +4,8 @@ module Cpk
   class Book < ActiveRecord::Base
     self.table_name = :cpk_books
     self.primary_key = [:author_id, :number]
+
+    belongs_to :order
   end
 
   class BestSeller < Book

activerecord/test/models/cpk/order.rb

@@ -6,5 +6,6 @@ class Order < ActiveRecord::Base
     self.primary_key = [:shop_id, :id]
 
     has_many :order_agreements, primary_key: :id
+    has_many :books
   end
 end

activerecord/test/schema/schema.rb

@@ -244,6 +244,7 @@
     t.integer :number
     t.string :title
     t.integer :revision
+    t.integer :order_id
   end
 
   create_table :cpk_orders, primary_key: [:shop_id, :id], force: true do |t|