Merge pull request #47996 from Shopify/pm/cpk-query-by-single-record

Extend query-association interface for composite models
rails · Apr 20, 2023 · d8a8df3 · d8a8df3
2 parents c978035 + 45da2cf
commit d8a8df3
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 2 deletions.
diff --git a/activerecord/lib/active_record/relation/predicate_builder/association_query_value.rb b/activerecord/lib/active_record/relation/predicate_builder/association_query_value.rb
@@ -10,7 +10,10 @@ def initialize(associated_table, value)
 
       def queries
         if associated_table.join_foreign_key.is_a?(Array)
-          ids.map { |ids_set| associated_table.join_foreign_key.zip(ids_set).to_h }
+          id_list = ids
+          id_list = id_list.pluck(primary_key) if id_list.is_a?(Relation)
+
+          id_list.map { |ids_set| associated_table.join_foreign_key.zip(ids_set).to_h }
         else
           [ associated_table.join_foreign_key => ids ]
         end
@@ -26,7 +29,7 @@ def ids
           when Array
             value.map { |v| convert_to_id(v) }
           else
-            convert_to_id(value)
+            [convert_to_id(value)]
           end
         end
 

diff --git a/activerecord/test/cases/associations_test.rb b/activerecord/test/cases/associations_test.rb
@@ -166,6 +166,23 @@ def test_querying_by_whole_associated_records_using_query_constraints
     assert_equal(expected_posts.map(&:id).sort, blog_posts.map(&:id).sort)
   end
 
+  def test_querying_by_single_associated_record_works_using_query_constraints
+    comments = [sharded_comments(:great_comment_blog_post_one), sharded_comments(:great_comment_blog_post_two)]
+
+    blog_posts = Sharded::BlogPost.where(comments: comments.last).to_a
+
+    expected_posts = [sharded_blog_posts(:great_post_blog_two)]
+    assert_equal(expected_posts.map(&:id).sort, blog_posts.map(&:id).sort)
+  end
+
+  def test_querying_by_relation_with_composite_key
+    expected_posts = [sharded_blog_posts(:great_post_blog_one), sharded_blog_posts(:great_post_blog_two)]
+
+    blog_posts = Sharded::BlogPost.where(comments: Sharded::Comment.where(body: "I really enjoyed the post!")).to_a
+
+    assert_equal(expected_posts.map(&:id).sort, blog_posts.map(&:id).sort)
+  end
+
   def test_has_many_association_with_composite_foreign_key_loads_records
     blog_post = sharded_blog_posts(:great_post_blog_one)