Improve cpk validation check

activerecord/lib/active_record/associations.rb

@@ -190,9 +190,9 @@ class CompositePrimaryKeyMismatchError < ActiveRecordError # :nodoc:
     def initialize(reflection = nil)
       if reflection
         if reflection.has_one? || reflection.collection?
-          super("Association #{reflection.active_record}##{reflection.name} primary key #{reflection.active_record_primary_key} doesn't match with foreign key #{reflection.foreign_key}. Please specify query_constraints.")
+          super("Association #{reflection.active_record}##{reflection.name} primary key #{reflection.active_record_primary_key} doesn't match with foreign key #{reflection.foreign_key}. Please specify query_constraints, or primary_key and foreign_key values.")
         else
-          super("Association #{reflection.active_record}##{reflection.name} primary key #{reflection.association_primary_key} doesn't match with foreign key #{reflection.foreign_key}. Please specify query_constraints.")
+          super("Association #{reflection.active_record}##{reflection.name} primary key #{reflection.association_primary_key} doesn't match with foreign key #{reflection.foreign_key}. Please specify query_constraints, or primary_key and foreign_key values.")
         end
       else
         super("Association primary key doesn't match with foreign key.")

activerecord/lib/active_record/reflection.rb

@@ -542,7 +542,7 @@ def join_foreign_key
       def check_validity!
         check_validity_of_inverse!
 
-        if !polymorphic? && klass.composite_primary_key?
+        if !polymorphic? && (klass.composite_primary_key? || active_record.composite_primary_key?)
           if (has_one? || collection?) && Array(active_record_primary_key).length != Array(foreign_key).length
             raise CompositePrimaryKeyMismatchError.new(self)
           elsif belongs_to? && Array(association_primary_key).length != Array(foreign_key).length
@@ -805,7 +805,7 @@ def association_class
       # klass option is necessary to support loading polymorphic associations
       def association_primary_key(klass = nil)
         if options[:query_constraints]
-          (klass || self.klass).query_constraints_list
+          (klass || self.klass).composite_query_constraints_list
         elsif primary_key = options[:primary_key]
           @association_primary_key ||= -primary_key.to_s
         else

activerecord/test/cases/associations/belongs_to_associations_test.rb

@@ -1777,17 +1777,37 @@ def self.name; "Temp"; end
     ActiveRecord.belongs_to_required_validates_foreign_key = original_value
   end
 
-  test "composite primary key malformed association" do
+  test "composite primary key malformed association class" do
     error = assert_raises(ActiveRecord::CompositePrimaryKeyMismatchError) do
       book = Cpk::BrokenBook.new(title: "Some book", order: Cpk::Order.new(id: [1, 2]))
       book.save!
     end
 
     assert_equal(<<~MESSAGE.squish, error.message)
       Association Cpk::BrokenBook#order primary key ["shop_id", "id"]
-      doesn't match with foreign key order_id. Please specify query_constraints.
+      doesn't match with foreign key order_id. Please specify query_constraints, or primary_key and foreign_key values.
     MESSAGE
   end
+
+  test "composite primary key malformed association owner class" do
+    error = assert_raises(ActiveRecord::CompositePrimaryKeyMismatchError) do
+      book = Cpk::BrokenBookWithNonCpkOrder.new(title: "Some book", order: Cpk::NonCpkOrder.new(id: 1))
+      book.save!
+    end
+
+    assert_equal(<<~MESSAGE.squish, error.message)
+      Association Cpk::BrokenBookWithNonCpkOrder#order primary key ["id"]
+      doesn't match with foreign key ["shop_id", "order_id"]. Please specify query_constraints, or primary_key and foreign_key values.
+    MESSAGE
+  end
+
+  test "association with query constraints assigns id on replacement" do
+    book = Cpk::NonCpkBook.create!(id: 1, author_id: 2, non_cpk_order: Cpk::NonCpkOrder.new)
+    other_order = Cpk::NonCpkOrder.create!
+    book.non_cpk_order = other_order
+
+    assert_equal(other_order.id, book.order_id)
+  end
 end
 
 class BelongsToWithForeignKeyTest < ActiveRecord::TestCase

activerecord/test/cases/associations/has_many_associations_test.rb

@@ -3174,15 +3174,27 @@ def test_key_ensuring_owner_was_is_valid_when_dependent_option_is_destroy_async
     end
   end
 
-  test "composite primary key malformed association" do
+  test "composite primary key malformed association class" do
     error = assert_raises(ActiveRecord::CompositePrimaryKeyMismatchError) do
       order = Cpk::BrokenOrder.new(id: [1, 2], books: [Cpk::Book.new(title: "Some book")])
       order.save!
     end
 
     assert_equal(<<~MESSAGE.squish, error.message)
       Association Cpk::BrokenOrder#books primary key ["shop_id", "id"]
-      doesn't match with foreign key broken_order_id. Please specify query_constraints.
+      doesn't match with foreign key broken_order_id. Please specify query_constraints, or primary_key and foreign_key values.
+    MESSAGE
+  end
+
+  test "composite primary key malformed association owner class" do
+    error = assert_raises(ActiveRecord::CompositePrimaryKeyMismatchError) do
+      order = Cpk::BrokenOrderWithNonCpkBooks.new(id: [1, 2], books: [Cpk::NonCpkBook.new(title: "Some book")])
+      order.save!
+    end
+
+    assert_equal(<<~MESSAGE.squish, error.message)
+    Association Cpk::BrokenOrderWithNonCpkBooks#books primary key [\"shop_id\", \"id\"]
+    doesn't match with foreign key broken_order_with_non_cpk_books_id. Please specify query_constraints, or primary_key and foreign_key values.
     MESSAGE
   end
 

activerecord/test/cases/associations/has_one_associations_test.rb

@@ -958,15 +958,27 @@ class SpecialContent < ActiveRecord::Base
     end
   end
 
-  test "composite primary key malformed association" do
+  test "composite primary key malformed association class" do
     error = assert_raises(ActiveRecord::CompositePrimaryKeyMismatchError) do
       order = Cpk::BrokenOrder.new(id: [1, 2], book: Cpk::Book.new(title: "Some book"))
       order.save!
     end
 
     assert_equal(<<~MESSAGE.squish, error.message)
       Association Cpk::BrokenOrder#book primary key ["shop_id", "id"]
-      doesn't match with foreign key broken_order_id. Please specify query_constraints.
+      doesn't match with foreign key broken_order_id. Please specify query_constraints, or primary_key and foreign_key values.
+    MESSAGE
+  end
+
+  test "composite primary key malformed association owner class" do
+    error = assert_raises(ActiveRecord::CompositePrimaryKeyMismatchError) do
+      order = Cpk::BrokenOrderWithNonCpkBooks.new(id: [1, 2], book: Cpk::NonCpkBook.new(title: "Some book"))
+      order.save!
+    end
+
+    assert_equal(<<~MESSAGE.squish, error.message)
+      Association Cpk::BrokenOrderWithNonCpkBooks#book primary key [\"shop_id\", \"id\"]
+      doesn't match with foreign key broken_order_with_non_cpk_books_id. Please specify query_constraints, or primary_key and foreign_key values.
     MESSAGE
   end
 end

activerecord/test/models/cpk/book.rb

@@ -3,7 +3,6 @@
 module Cpk
   class Book < ActiveRecord::Base
     self.table_name = :cpk_books
-
     belongs_to :order, autosave: true, query_constraints: [:shop_id, :order_id]
     belongs_to :author, class_name: "Cpk::Author"
 
@@ -17,6 +16,16 @@ class BrokenBook < Book
     belongs_to :order
   end
 
+  class BrokenBookWithNonCpkOrder < Book
+    belongs_to :order, class_name: "Cpk::NonCpkOrder", query_constraints: [:shop_id, :order_id]
+  end
+
+  class NonCpkBook < Book
+    self.primary_key = :id
+
+    belongs_to :non_cpk_order, query_constraints: [:order_id]
+  end
+
   class NullifiedBook < Book
     has_one :chapter, query_constraints: [:author_id, :book_id], dependent: :nullify
   end

activerecord/test/models/cpk/order.rb

@@ -17,6 +17,15 @@ class BrokenOrder < Order
     has_one :book
   end
 
+  class BrokenOrderWithNonCpkBooks < Order
+    has_many :books, class_name: "Cpk::NonCpkBook"
+    has_one :book, class_name: "Cpk::NonCpkBook"
+  end
+
+  class NonCpkOrder < Order
+    self.primary_key = :id
+  end
+
   class OrderWithPrimaryKeyAssociatedBook < Order
     has_one :book, primary_key: :id, foreign_key: :order_id
   end