Validate composite key length when owner record class has composite key

Catches more bugs in CPK assocaitions by validating the shape of keys on
both ends of the association.

Co-authored-by: Nikita Vasilevsky <[email protected]>

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

activerecord/test/cases/associations/belongs_to_associations_test.rb

@@ -1777,7 +1777,7 @@ 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!
@@ -1788,6 +1788,18 @@ def self.name; "Temp"; end
       doesn't match with foreign key order_id. Please specify query_constraints.
     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
+
+    # TODO: Improve this message
+    assert_equal(<<~MESSAGE.chomp, error.message)
+      Association Cpk::BrokenBookWithNonCpkOrder#order primary key  doesn't match with foreign key ["shop_id", "order_id"]. Please specify query_constraints.
+    MESSAGE
+  end
 end
 
 class BelongsToWithForeignKeyTest < ActiveRecord::TestCase

activerecord/test/cases/associations/has_many_associations_test.rb

@@ -3174,7 +3174,7 @@ 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!
@@ -3186,6 +3186,18 @@ def test_key_ensuring_owner_was_is_valid_when_dependent_option_is_destroy_async
     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.
+    MESSAGE
+  end
+
   private
     def force_signal37_to_load_all_clients_of_firm
       companies(:first_firm).clients_of_firm.load_target

activerecord/test/cases/associations/has_one_associations_test.rb

@@ -958,7 +958,7 @@ 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!
@@ -969,4 +969,16 @@ class SpecialContent < ActiveRecord::Base
       doesn't match with foreign key broken_order_id. Please specify query_constraints.
     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.
+    MESSAGE
+  end
 end

activerecord/test/models/cpk/book.rb

@@ -17,6 +17,14 @@ 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
+  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