Merge pull request rails#49088 from eileencodes/derive-foreign-key-query-constraints

Allow Active Record to derive the association query constraints

Author: eileencodes

activerecord/lib/active_record/autosave_association.rb

@@ -532,6 +532,8 @@ def compute_primary_key(reflection, record)
           primary_key_options
         elsif reflection.options[:query_constraints] && (query_constraints = record.class.query_constraints_list)
           query_constraints
+        elsif record.class.has_query_constraints? && !reflection.options[:foreign_key]
+          record.class.query_constraints_list
         else
           record.class.primary_key
         end

activerecord/lib/active_record/reflection.rb

@@ -502,8 +502,16 @@ 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
         else
-          -(options[:foreign_key]&.to_s || derive_foreign_key(infer_from_inverse_of: infer_from_inverse_of))
+          derived_fk = derive_foreign_key(infer_from_inverse_of: infer_from_inverse_of)
+
+          if active_record.has_query_constraints?
+            derived_fk = derive_fk_query_constraints(active_record, derived_fk)
+          end
+
+          derived_fk
         end
       end
 
@@ -759,6 +767,50 @@ def derive_foreign_key(infer_from_inverse_of: true)
           end
         end
 
+        def derive_fk_query_constraints(klass, foreign_key)
+          primary_query_constraints = klass.query_constraints_list
+          owner_pk = klass.primary_key
+
+          if primary_query_constraints.size != 2
+            raise ArgumentError, <<~MSG.squish
+              The query constraints list on the `#{klass}` model has more than 2
+              attributes. Active Record is unable to derive the query constraints
+              for the association. You need to explicitly define the query constraints
+              for this association.
+            MSG
+          end
+
+          if !primary_query_constraints.include?(owner_pk)
+            raise ArgumentError, <<~MSG.squish
+              The query constraints on the `#{klass}` model does not include the primary
+              key so Active Record is unable to derive the foreign key constraints for
+              the association. You need to explicitly define the query constraints for this
+              association.
+            MSG
+          end
+
+          # The primary key and foreign key are both already in the query constraints
+          # so we don't want to derive the key. In this case we want a single key.
+          if primary_query_constraints.include?(owner_pk) && primary_query_constraints.include?(foreign_key)
+            return foreign_key
+          end
+
+          first_key, last_key = primary_query_constraints
+
+          if first_key == owner_pk
+            [foreign_key, last_key.to_s]
+          elsif last_key == owner_pk
+            [first_key.to_s, foreign_key]
+          else
+            raise ArgumentError, <<~MSG.squish
+              Active Record couldn't correctly interpret the query constraints
+              for the `#{klass}` model. The query constraints on `#{klass}` are
+              `#{primary_query_constraints}` and the foreign key is `#{foreign_key}`.
+              You need to explicitly set the query constraints for this association.
+            MSG
+          end
+        end
+
         def derive_join_table
           ModelSchema.derive_join_table_name active_record.table_name, klass.table_name
         end

activerecord/test/cases/associations_test.rb

@@ -45,7 +45,7 @@ class AssociationsTest < ActiveRecord::TestCase
   fixtures :accounts, :companies, :developers, :projects, :developers_projects,
            :computers, :people, :readers, :authors, :author_addresses, :author_favorites,
            :comments, :posts, :sharded_blogs, :sharded_blog_posts, :sharded_comments, :sharded_tags, :sharded_blog_posts_tags,
-           :cpk_orders
+           :cpk_orders, :cpk_books, :cpk_reviews
 
   def test_eager_loading_should_not_change_count_of_children
     liquid = Liquid.create(name: "salty")
@@ -225,6 +225,18 @@ def test_has_many_association_from_a_model_with_query_constraints_different_from
     assert_equal(expected_comments.sort, comments.sort)
   end
 
+  def test_query_constraints_over_three_without_defining_explicit_foreign_key_query_constraints_raises
+    Sharded::BlogPostWithRevision.has_many :comments_without_query_constraints, primary_key: [:blog_id, :id], class_name: "Comment"
+    blog_post = sharded_blog_posts(:great_post_blog_one)
+    blog_post = Sharded::BlogPostWithRevision.find(blog_post.id)
+
+    error = assert_raises ArgumentError do
+      blog_post.comments_without_query_constraints.to_a
+    end
+
+    assert_equal "The query constraints list on the `Sharded::BlogPostWithRevision` model has more than 2 attributes. Active Record is unable to derive the query constraints for the association. You need to explicitly define the query constraints for this association.", error.message
+  end
+
   def test_model_with_composite_query_constraints_has_many_association_sql
     blog_post = sharded_blog_posts(:great_post_blog_one)
 
@@ -315,6 +327,20 @@ def test_assign_composite_foreign_key_belongs_to_association
     assert_equal(another_blog.id, comment.blog_id)
   end
 
+  def test_query_constraints_that_dont_include_the_primary_key_raise
+    original = Sharded::BlogPost.instance_variable_get(:@query_constraints_list)
+    Sharded::BlogPost.query_constraints :title, :revision
+    Sharded::BlogPost.has_many :comments_without_query_constraints, primary_key: [:blog_id, :id], class_name: "Comment"
+    blog_post = sharded_blog_posts(:great_post_blog_one)
+
+    error = assert_raises ArgumentError do
+      blog_post.comments_without_query_constraints.to_a
+    end
+
+    assert_equal "The query constraints on the `Sharded::BlogPost` model does not include the primary key so Active Record is unable to derive the foreign key constraints for the association. You need to explicitly define the query constraints for this association.", error.message
+  ensure
+    Sharded::BlogPost.instance_variable_set(:@query_constraints_list, original)
+  end
 
   def test_assign_belongs_to_cpk_model_by_id_attribute
     order = cpk_orders(:cpk_groceries_order_1)

activerecord/test/models/sharded/blog_post.rb

@@ -6,10 +6,10 @@ class BlogPost < ActiveRecord::Base
     query_constraints :blog_id, :id
 
     belongs_to :blog
-    has_many :comments, query_constraints: [:blog_id, :blog_post_id]
-    has_many :delete_comments, class_name: "Sharded::Comment", query_constraints: [:blog_id, :blog_post_id], dependent: :delete_all
+    has_many :comments
+    has_many :delete_comments, class_name: "Sharded::Comment", dependent: :delete_all
 
-    has_many :blog_post_tags, query_constraints: [:blog_id, :blog_post_id]
+    has_many :blog_post_tags
     has_many :tags, through: :blog_post_tags
   end
 end

activerecord/test/models/sharded/blog_post_tag.rb

@@ -3,8 +3,9 @@
 module Sharded
   class BlogPostTag < ActiveRecord::Base
     self.table_name = :sharded_blog_posts_tags
+    query_constraints :blog_id, :id
 
-    belongs_to :blog_post, query_constraints: [:blog_id, :blog_post_id]
-    belongs_to :tag, query_constraints: [:blog_id, :tag_id]
+    belongs_to :blog_post
+    belongs_to :tag
   end
 end

activerecord/test/models/sharded/comment.rb

@@ -5,7 +5,7 @@ class Comment < ActiveRecord::Base
     self.table_name = :sharded_comments
     query_constraints :blog_id, :id
 
-    belongs_to :blog_post, query_constraints: [:blog_id, :blog_post_id]
+    belongs_to :blog_post
     belongs_to :blog_post_by_id, class_name: "Sharded::BlogPost", foreign_key: :blog_post_id
     belongs_to :blog
   end

activerecord/test/models/sharded/tag.rb

@@ -5,7 +5,7 @@ class Tag < ActiveRecord::Base
     self.table_name = :sharded_tags
     query_constraints :blog_id, :id
 
-    has_many :blog_post_tags, query_constraints: [:blog_id, :tag_id]
+    has_many :blog_post_tags
     has_many :blog_posts, through: :blog_post_tags
   end
 end