Associations::Preloader supports preloading associations with composite keys

activerecord/lib/active_record/associations/preloader/association.rb

@@ -36,7 +36,12 @@ def load_records_in_batch(loaders)
           end
 
           def load_records_for_keys(keys, &block)
-            scope.where(association_key_name => keys).load(&block)
+            if association_key_name.is_a?(Array)
+              or_scope = keys.map { |values_set| scope.klass.where(association_key_name.zip(values_set).to_h) }.inject(&:or)
+              scope.merge(or_scope)
+            else
+              scope.where(association_key_name => keys)
+            end.load(&block)
           end
         end
 
@@ -151,7 +156,7 @@ def loader_query
 
         def owners_by_key
           @owners_by_key ||= owners.each_with_object({}) do |owner, result|
-            key = convert_key(owner[owner_key_name])
+            key = derive_key(owner, owner_key_name)
             (result[key] ||= []) << owner if key
           end
         end
@@ -169,7 +174,7 @@ def scope
         end
 
         def set_inverse(record)
-          if owners = owners_by_key[convert_key(record[association_key_name])]
+          if owners = owners_by_key[derive_key(record, association_key_name)]
             # Processing only the first owner
             # because the record is modified but not an owner
             association = owners.first.association(reflection.name)
@@ -182,11 +187,10 @@ def load_records(raw_records = nil)
           # #compare_by_identity makes such owners different hash keys
           @records_by_owner = {}.compare_by_identity
           raw_records ||= loader_query.records_for([self])
-
           @preloaded_records = raw_records.select do |record|
             assignments = false
 
-            owners_by_key[convert_key(record[association_key_name])]&.each do |owner|
+            owners_by_key[derive_key(record, association_key_name)]&.each do |owner|
               entries = (@records_by_owner[owner] ||= [])
 
               if reflection.collection? || entries.empty?
@@ -206,7 +210,7 @@ def associate_records_from_unscoped(unscoped_records)
           return if reflection.collection?
 
           unscoped_records.select { |r| r[association_key_name].present? }.each do |record|
-            owners = owners_by_key[convert_key(record[association_key_name])]
+            owners = owners_by_key[derive_key(record, association_key_name)]
             owners&.each_with_index do |owner, i|
               association = owner.association(reflection.name)
               association.target = record
@@ -246,6 +250,14 @@ def key_conversion_required?
             @key_conversion_required
           end
 
+          def derive_key(owner, key)
+            if key.is_a?(Array)
+              key.map { |k| convert_key(owner[k]) }
+            else
+              convert_key(owner[key])
+            end
+          end
+
           def convert_key(key)
             if key_conversion_required?
               key.to_s

activerecord/test/cases/associations/eager_test.rb

@@ -33,6 +33,9 @@
 require "models/pirate"
 require "models/matey"
 require "models/parrot"
+require "models/sharded/blog"
+require "models/sharded/blog_post"
+require "models/sharded/comment"
 
 class EagerLoadingTooManyIdsTest < ActiveRecord::TestCase
   fixtures :citations
@@ -51,7 +54,7 @@ class EagerAssociationTest < ActiveRecord::TestCase
             :companies, :accounts, :tags, :taggings, :ratings, :people, :readers, :categorizations,
             :owners, :pets, :author_favorites, :jobs, :references, :subscribers, :subscriptions, :books,
             :developers, :projects, :developers_projects, :members, :memberships, :clubs, :sponsors,
-            :pirates, :mateys
+            :pirates, :mateys, :sharded_blogs, :sharded_blog_posts, :sharded_comments
 
   def test_eager_with_has_one_through_join_model_with_conditions_on_the_through
     member = Member.all.merge!(includes: :favorite_club).find(members(:some_other_guy).id)
@@ -1675,6 +1678,25 @@ def test_preloading_has_many_through_with_custom_scope
     end
   end
 
+  test "preloading belongs_to association associated by a composite query_constraints" do
+    blog_ids = [sharded_blogs(:sharded_blog_one).id, sharded_blogs(:sharded_blog_two).id]
+    posts = Sharded::BlogPost.where(blog_id: blog_ids).includes(:comments).to_a
+    assert posts.all? { |post| post.comments.loaded? }
+
+    post = posts.find { |post| post.id == sharded_blog_posts(:great_post_blog_one).id }
+    expected_comments = Sharded::Comment.where(blog_id: post.blog_id, blog_post_id: post.id).to_a
+    assert_equal(post.comments.sort, expected_comments.sort)
+  end
+
+  test "reloading has_many association associated by a composite query_constraints" do
+    blog_ids = [sharded_blogs(:sharded_blog_one).id, sharded_blogs(:sharded_blog_two).id]
+    comments = Sharded::Comment.where(blog_id: blog_ids).includes(:blog_post).to_a
+    assert comments.all? { |comment| comment.association(:blog_post).loaded? }
+
+    comment = comments.find { |comment| comment.id == sharded_comments(:great_comment_blog_post_one).id }
+    assert_equal(comment.blog_post, sharded_blog_posts(:great_post_blog_one))
+  end
+
   private
     def find_all_ordered(klass, include = nil)
       klass.order("#{klass.table_name}.#{klass.primary_key}").includes(include).to_a

activerecord/test/cases/associations/inner_join_association_test.rb

@@ -216,8 +216,8 @@ def test_find_with_conditions_on_through_reflection
   end
 
   test "joins a belongs_to association with a composite foreign key" do
-    first_post_comments = Sharded::Comment.joins(:blog_post).where(blog_post: { title: "My first post!" }).to_a
-    expected_blog_post = sharded_blog_posts(:great_blog_post_one)
+    first_post_comments = Sharded::Comment.joins(:blog_post).where(blog_post: { title: "My first post in my Blog1!" }).to_a
+    expected_blog_post = sharded_blog_posts(:great_post_blog_one)
 
     assert_not_empty comments
     assert_equal(expected_blog_post.comments.to_a.sort, first_post_comments.sort)

activerecord/test/cases/associations_test.rb

@@ -131,7 +131,7 @@ def test_association_with_references
 
   def test_belongs_to_a_model_with_composite_foreign_key_finds_associated_record
     comment = sharded_comments(:great_comment_blog_post_one)
-    blog_post = sharded_blog_posts(:great_blog_post_one)
+    blog_post = sharded_blog_posts(:great_post_blog_one)
 
     assert_equal(blog_post, comment.blog_post)
   end
@@ -148,15 +148,15 @@ def test_belongs_to_a_model_with_composite_primary_key_uses_composite_pk_in_sql
   end
 
   def test_has_many_association_with_composite_foreign_key_loads_records
-    blog_post = sharded_blog_posts(:great_blog_post_one)
+    blog_post = sharded_blog_posts(:great_post_blog_one)
 
     comments = blog_post.comments.to_a
     assert_includes(comments, sharded_comments(:wow_comment_blog_post_one))
     assert_includes(comments, sharded_comments(:great_comment_blog_post_one))
   end
 
   def test_model_with_composite_query_constraints_has_many_association_sql
-    blog_post = sharded_blog_posts(:great_blog_post_one)
+    blog_post = sharded_blog_posts(:great_post_blog_one)
 
     sql = capture_sql do
       blog_post.comments.to_a
@@ -167,7 +167,7 @@ def test_model_with_composite_query_constraints_has_many_association_sql
   end
 
   def test_append_composite_foreign_key_has_many_association
-    blog_post = sharded_blog_posts(:great_blog_post_one)
+    blog_post = sharded_blog_posts(:great_post_blog_one)
     comment = Sharded::Comment.new(body: "Great post! :clap:")
     comment.save
     blog_post.comments << comment
@@ -206,7 +206,7 @@ def test_assign_composite_foreign_key_belongs_to_association
   end
 
   def test_append_composite_foreign_key_has_many_association_with_autosave
-    blog_post = sharded_blog_posts(:great_blog_post_one)
+    blog_post = sharded_blog_posts(:great_post_blog_one)
     comment = Sharded::Comment.new(body: "Great post! :clap:")
     blog_post.comments << comment
 
@@ -539,7 +539,8 @@ def test_associations_raise_with_name_error_if_associated_to_classes_that_do_not
 end
 
 class PreloaderTest < ActiveRecord::TestCase
-  fixtures :posts, :comments, :books, :authors, :tags, :taggings, :essays, :categories, :author_addresses
+  fixtures :posts, :comments, :books, :authors, :tags, :taggings, :essays, :categories,
+    :author_addresses, :sharded_blog_posts, :sharded_comments
 
   def test_preload_with_scope
     post = posts(:welcome)
@@ -1126,6 +1127,26 @@ def test_preload_wont_set_the_wrong_target
       assert_not_equal some_other_record, post.author
     end
   end
+
+  def test_preload_has_many_association_with_composite_foreign_key
+    blog_post = sharded_blog_posts(:great_post_blog_one)
+    blog_posts = [blog_post, sharded_blog_posts(:great_post_blog_two)]
+
+    ::ActiveRecord::Associations::Preloader.new(records: blog_posts, associations: [:comments]).call
+
+    assert blog_post.association(:comments).loaded?
+    assert_includes(blog_post.comments.to_a, sharded_comments(:great_comment_blog_post_one))
+  end
+
+  def test_preload_belongs_to_association_with_composite_foreign_key
+    comment = sharded_comments(:great_comment_blog_post_one)
+    comments = [comment, sharded_comments(:great_comment_blog_post_two)]
+
+    ActiveRecord::Associations::Preloader.new(records: comments, associations: :blog_post).call
+
+    assert comment.association(:blog_post).loaded?
+    assert_equal sharded_blog_posts(:great_post_blog_one), comment.blog_post
+  end
 end
 
 class GeneratedMethodsTest < ActiveRecord::TestCase

activerecord/test/fixtures/sharded_blog_posts.yml

@@ -1,10 +1,14 @@
 _fixture:
   model_class: Sharded::BlogPost
 
-great_blog_post_one:
-  title: "My first post!"
+great_post_blog_one:
+  title: "My first post in my Blog1!"
   blog_id: <%= ActiveRecord::FixtureSet.identify(:sharded_blog_one) %>
 
-great_blog_post_two:
-  title: "My second post!"
+second_post_blog_one:
+  title: "This is my second post in my Blog1!"
+  blog_id: <%= ActiveRecord::FixtureSet.identify(:sharded_blog_one) %>
+
+great_post_blog_two:
+  title: "My first post in my Blog2!"
   blog_id: <%= ActiveRecord::FixtureSet.identify(:sharded_blog_two) %>

activerecord/test/fixtures/sharded_comments.yml

@@ -3,17 +3,17 @@ _fixture:
 
 great_comment_blog_post_one:
   body: "I really enjoyed the post!"
-  blog_post_id: <%= ActiveRecord::FixtureSet.identify(:great_blog_post_one) %>
+  blog_post_id: <%= ActiveRecord::FixtureSet.identify(:great_post_blog_one) %>
   blog_id: <%= ActiveRecord::FixtureSet.identify(:sharded_blog_one) %>
 
 wow_comment_blog_post_one:
   body: "Wow!"
-  blog_post_id: <%= ActiveRecord::FixtureSet.identify(:great_blog_post_one) %>
+  blog_post_id: <%= ActiveRecord::FixtureSet.identify(:great_post_blog_one) %>
   blog_id: <%= ActiveRecord::FixtureSet.identify(:sharded_blog_one) %>
 
 unique_comment_blog_post_one:
   body: "Your first blog post is great!"
-  blog_post_id: <%= ActiveRecord::FixtureSet.identify(:great_blog_post_one) %>
+  blog_post_id: <%= ActiveRecord::FixtureSet.identify(:great_post_blog_one) %>
   blog_id: <%= ActiveRecord::FixtureSet.identify(:sharded_blog_one) %>
 
 great_comment_blog_post_two: