Merge pull request #50321 from fatkodima/fix-polymorphic-belongs_to-with-query_constraints

eileencodes · eileencodes · commit 59e72b0b69a8 · 2023-12-11T10:12:56.000-05:00
Fix polymorphic `belongs_to` to correctly use parent's `query_constraints`
diff --git a/activerecord/lib/active_record/reflection.rb b/activerecord/lib/active_record/reflection.rb
@@ -870,7 +870,7 @@ def association_class
       def association_primary_key(klass = nil)
         if primary_key = options[:primary_key]
           @association_primary_key ||= -primary_key.to_s
-        elsif !polymorphic? && ((klass || self.klass).has_query_constraints? || options[:query_constraints])
+        elsif (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
diff --git a/activerecord/test/cases/associations_test.rb b/activerecord/test/cases/associations_test.rb
@@ -284,6 +284,14 @@ 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_polymorphic_belongs_to_uses_parent_query_constraints
+    parent_post = sharded_blog_posts(:great_post_blog_one)
+    child_post = Sharded::BlogPost.create!(title: "Child post", blog_id: parent_post.blog_id, parent: parent_post)
+    child_post.reload # reload to forget the parent association
+
+    assert_equal parent_post, child_post.parent
+  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
diff --git a/activerecord/test/models/sharded/blog_post.rb b/activerecord/test/models/sharded/blog_post.rb
@@ -5,9 +5,11 @@ class BlogPost < ActiveRecord::Base
     self.table_name = :sharded_blog_posts
     query_constraints :blog_id, :id
 
+    belongs_to :parent, class_name: name, polymorphic: true
     belongs_to :blog
     has_many :comments
     has_many :delete_comments, class_name: "Sharded::Comment", dependent: :delete_all
+    has_many :children, class_name: name, as: :parent
 
     has_many :blog_post_tags
     has_many :tags, through: :blog_post_tags
diff --git a/activerecord/test/schema/schema.rb b/activerecord/test/schema/schema.rb
@@ -313,6 +313,7 @@
 
   create_table :sharded_blog_posts, force: true do |t|
     t.string :title
+    t.references :parent, polymorphic: true
     t.integer :blog_id
     t.integer :revision
   end
