Merge pull request rails#49113 from adrianna-chang-shopify/ac-infer-pk-id-composite-associations

Infer `primary_key: :id` on associations with composite primary key models

Author: eileencodes

activerecord/lib/active_record/autosave_association.rb

@@ -449,12 +449,16 @@ def save_has_one_association(reflection)
           if autosave && record.marked_for_destruction?
             record.destroy
           elsif autosave != false
-            key = reflection.options[:primary_key] ? _read_attribute(reflection.options[:primary_key].to_s) : id
+            primary_key = Array(compute_primary_key(reflection, self)).map(&:to_s)
+            primary_key_value = primary_key.map { |key| _read_attribute(key) }
 
-            if (autosave && record.changed_for_autosave?) || _record_changed?(reflection, record, key)
+            if (autosave && record.changed_for_autosave?) || _record_changed?(reflection, record, primary_key_value)
               unless reflection.through_reflection
-                Array(key).zip(Array(reflection.foreign_key)).each do |primary_key, foreign_key_column|
-                  record[foreign_key_column] = primary_key
+                foreign_key = Array(reflection.foreign_key)
+                primary_key_foreign_key_pairs = primary_key.zip(foreign_key)
+
+                primary_key_foreign_key_pairs.each do |primary_key, foreign_key|
+                  record[foreign_key] = _read_attribute(primary_key)
                 end
                 association.set_inverse_instance(record)
               end
@@ -534,6 +538,10 @@ def compute_primary_key(reflection, record)
           query_constraints
         elsif record.class.has_query_constraints? && !reflection.options[:foreign_key]
           record.class.query_constraints_list
+        elsif record.class.composite_primary_key?
+          # If record has composite primary key of shape [:<tenant_key>, :id], infer primary_key as :id
+          primary_key = record.class.primary_key
+          primary_key.include?("id") ? "id" : primary_key
         else
           record.class.primary_key
         end

activerecord/lib/active_record/reflection.rb

@@ -500,7 +500,6 @@ def join_table
 
       def foreign_key(infer_from_inverse_of: true)
         @foreign_key ||= if options[:query_constraints]
-          # composite foreign keys support
           options[:query_constraints].map { |fk| fk.to_s.freeze }.freeze
         elsif options[:foreign_key]
           options[:foreign_key].to_s
@@ -533,6 +532,10 @@ def active_record_primary_key
           end
         elsif active_record.has_query_constraints? || options[:query_constraints]
           active_record.query_constraints_list
+        elsif active_record.composite_primary_key?
+          # If active_record has composite primary key of shape [:<tenant_key>, :id], infer primary_key as :id
+          primary_key = primary_key(active_record)
+          primary_key.include?("id") ? "id" : primary_key.freeze
         else
           primary_key(active_record).freeze
         end
@@ -863,6 +866,10 @@ def association_primary_key(klass = nil)
           (klass || self.klass).composite_query_constraints_list
         elsif primary_key = options[:primary_key]
           @association_primary_key ||= -primary_key.to_s
+        elsif (klass || self.klass).composite_primary_key?
+          # If klass has composite primary key of shape [:<tenant_key>, :id], infer primary_key as :id
+          primary_key = (klass || self.klass).primary_key
+          primary_key.include?("id") ? "id" : primary_key
         else
           primary_key(klass || self.klass)
         end

activerecord/test/cases/associations/belongs_to_associations_test.rb

@@ -1784,7 +1784,7 @@ def self.name; "Temp"; end
     end
 
     assert_equal(<<~MESSAGE.squish, error.message)
-      Association Cpk::BrokenBook#order primary key ["shop_id", "id"]
+      Association Cpk::BrokenBook#order primary key ["shop_id", "status"]
       doesn't match with foreign key order_id. Please specify query_constraints, or primary_key and foreign_key values.
     MESSAGE
   end

activerecord/test/cases/associations/has_many_associations_test.rb

@@ -3199,7 +3199,7 @@ def test_key_ensuring_owner_was_is_valid_when_dependent_option_is_destroy_async
     end
 
     assert_equal(<<~MESSAGE.squish, error.message)
-      Association Cpk::BrokenOrder#books primary key ["shop_id", "id"]
+      Association Cpk::BrokenOrder#books primary key ["shop_id", "status"]
       doesn't match with foreign key broken_order_id. Please specify query_constraints, or primary_key and foreign_key values.
     MESSAGE
   end
@@ -3211,7 +3211,7 @@ def test_key_ensuring_owner_was_is_valid_when_dependent_option_is_destroy_async
     end
 
     assert_equal(<<~MESSAGE.squish, error.message)
-    Association Cpk::BrokenOrderWithNonCpkBooks#books primary key [\"shop_id\", \"id\"]
+    Association Cpk::BrokenOrderWithNonCpkBooks#books primary key [\"shop_id\", \"status\"]
     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

@@ -922,7 +922,7 @@ def test_has_one_with_touch_option_on_nonpersisted_built_associations_doesnt_upd
     end
 
     assert_equal(<<~MESSAGE.squish, error.message)
-      Association Cpk::BrokenOrder#book primary key ["shop_id", "id"]
+      Association Cpk::BrokenOrder#book primary key ["shop_id", "status"]
       doesn't match with foreign key broken_order_id. Please specify query_constraints, or primary_key and foreign_key values.
     MESSAGE
   end
@@ -934,7 +934,7 @@ def test_has_one_with_touch_option_on_nonpersisted_built_associations_doesnt_upd
     end
 
     assert_equal(<<~MESSAGE.squish, error.message)
-      Association Cpk::BrokenOrderWithNonCpkBooks#book primary key [\"shop_id\", \"id\"]
+      Association Cpk::BrokenOrderWithNonCpkBooks#book primary key [\"shop_id\", \"status\"]
       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/models/cpk/book.rb

@@ -13,7 +13,7 @@ class BestSeller < Book
   end
 
   class BrokenBook < Book
-    belongs_to :order
+    belongs_to :order, class_name: "Cpk::OrderWithSpecialPrimaryKey"
   end
 
   class BrokenBookWithNonCpkOrder < Book

activerecord/test/models/cpk/order.rb

@@ -9,17 +9,28 @@ class Order < ActiveRecord::Base
 
     alias_attribute :id_value, :id
 
-    has_many :order_agreements, primary_key: :id
+    has_many :order_agreements
     has_many :books, query_constraints: [:shop_id, :order_id]
     has_one :book, query_constraints: [:shop_id, :order_id]
   end
 
   class BrokenOrder < Order
+    self.primary_key = [:shop_id, :status]
+
     has_many :books
     has_one :book
   end
 
+  class OrderWithSpecialPrimaryKey < Order
+    self.primary_key = [:shop_id, :status]
+
+    has_many :books, query_constraints: [:shop_id, :status]
+    has_one :book, query_constraints: [:shop_id, :status]
+  end
+
   class BrokenOrderWithNonCpkBooks < Order
+    self.primary_key = [:shop_id, :status]
+
     has_many :books, class_name: "Cpk::NonCpkBook"
     has_one :book, class_name: "Cpk::NonCpkBook"
   end
@@ -29,7 +40,7 @@ class NonCpkOrder < Order
   end
 
   class OrderWithPrimaryKeyAssociatedBook < Order
-    has_one :book, primary_key: :id, foreign_key: :order_id
+    has_one :book, foreign_key: :order_id
   end
 
   class OrderWithNullifiedBook < Order

activerecord/test/models/cpk/order_agreement.rb

@@ -5,6 +5,6 @@ module Cpk
   class OrderAgreement < ActiveRecord::Base
     self.table_name = :cpk_order_agreements
 
-    belongs_to :order, primary_key: :id # foreign key is derived as `order_id`
+    belongs_to :order
   end
 end

activerecord/test/models/cpk/order_tag.rb

@@ -5,6 +5,6 @@ class OrderTag < ActiveRecord::Base
     self.table_name = :cpk_order_tags
 
     belongs_to :tag
-    belongs_to :order, primary_key: :id
+    belongs_to :order
   end
 end