From dcfddc1f90090ad84a0fc7ff80844c704c71b4cc Mon Sep 17 00:00:00 2001
From: Nikita Vasilevsky <nikita.vasilevsky@shopify.com>
Date: Wed, 13 Sep 2023 21:00:39 +0000
Subject: [PATCH] Explicit `primary_key:` option should always take priority in
 associations

This commit fixes the issue where the `primary_key:` option was ignored
if the associated model had a `query_constraints` configured.
Now `primary_key:` option always takes priority and only if there is no
`primary_key:` option, the `query_constraints` are used to determine
the `association_primary_key` value.
---
 activerecord/lib/active_record/reflection.rb | 6 +++---
 activerecord/test/cases/associations_test.rb | 7 +++++++
 activerecord/test/cases/reflection_test.rb   | 6 ++++++
 activerecord/test/models/sharded/comment.rb  | 2 +-
 4 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/activerecord/lib/active_record/reflection.rb b/activerecord/lib/active_record/reflection.rb
index 87795a7dcd124..372aa2b6070ab 100644
--- a/activerecord/lib/active_record/reflection.rb
+++ b/activerecord/lib/active_record/reflection.rb
@@ -862,10 +862,10 @@ def association_class
 
       # klass option is necessary to support loading polymorphic associations
       def association_primary_key(klass = nil)
-        if !polymorphic? && ((klass || self.klass).has_query_constraints? || options[:query_constraints])
-          (klass || self.klass).composite_query_constraints_list
-        elsif primary_key = options[:primary_key]
+        if primary_key = options[:primary_key]
           @association_primary_key ||= -primary_key.to_s
+        elsif !polymorphic? && ((klass || self.klass).has_query_constraints? || options[:query_constraints])
+          (klass || self.klass).composite_query_constraints_list
         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
diff --git a/activerecord/test/cases/associations_test.rb b/activerecord/test/cases/associations_test.rb
index a79c555b85cb6..3808bb01bd10f 100644
--- a/activerecord/test/cases/associations_test.rb
+++ b/activerecord/test/cases/associations_test.rb
@@ -260,6 +260,13 @@ def test_belongs_to_association_does_not_use_parent_query_constraints_if_not_con
     assert_equal(blog_post, comment.blog_post_by_id)
   end
 
+  def test_preloads_model_with_query_constraints_by_explicitly_configured_fk_and_pk
+    comment = sharded_comments(:great_comment_blog_post_one)
+    comments = Sharded::Comment.where(id: comment.id).preload(:blog_post_by_id).to_a
+    comment = comments.first
+    assert_equal(comment.blog_post_by_id, comment.blog_post)
+  end
+
   def test_append_composite_foreign_key_has_many_association
     blog_post = sharded_blog_posts(:great_post_blog_one)
     comment = Sharded::Comment.new(body: "Great post! :clap:")
diff --git a/activerecord/test/cases/reflection_test.rb b/activerecord/test/cases/reflection_test.rb
index feee7ad5d8602..6566c2a3f671b 100644
--- a/activerecord/test/cases/reflection_test.rb
+++ b/activerecord/test/cases/reflection_test.rb
@@ -29,6 +29,7 @@
 require "models/recipe"
 require "models/user_with_invalid_relation"
 require "models/hardback"
+require "models/sharded/comment"
 
 class ReflectionTest < ActiveRecord::TestCase
   include ActiveRecord::Reflection
@@ -614,6 +615,11 @@ def test_automatic_inverse_does_not_suppress_name_error_from_incidental_code
     end
   end
 
+  def test_association_primary_key_uses_explicit_primary_key_option_as_first_priority
+    actual = Sharded::Comment.reflect_on_association(:blog_post_by_id).association_primary_key
+    assert_equal "id", actual
+  end
+
   private
     def assert_reflection(klass, association, options)
       assert reflection = klass.reflect_on_association(association)
diff --git a/activerecord/test/models/sharded/comment.rb b/activerecord/test/models/sharded/comment.rb
index cd8a339efc20c..1f2d7e6978c7f 100644
--- a/activerecord/test/models/sharded/comment.rb
+++ b/activerecord/test/models/sharded/comment.rb
@@ -6,7 +6,7 @@ class Comment < ActiveRecord::Base
     query_constraints :blog_id, :id
 
     belongs_to :blog_post
-    belongs_to :blog_post_by_id, class_name: "Sharded::BlogPost", foreign_key: :blog_post_id
+    belongs_to :blog_post_by_id, class_name: "Sharded::BlogPost", foreign_key: :blog_post_id, primary_key: :id
     belongs_to :blog
   end
 end