Fix derived foreign key to return correctly when association id is part of query constraints

Author: iamradioactive

activerecord/lib/active_record/reflection.rb

@@ -804,6 +804,8 @@ def derive_fk_query_constraints(foreign_key)
             MSG
           end
 
+          return foreign_key if primary_query_constraints.include?(foreign_key)
+
           first_key, last_key = primary_query_constraints
 
           if first_key == owner_pk

activerecord/test/cases/reflection_test.rb

@@ -653,6 +653,14 @@ def test_association_primary_key_uses_explicit_primary_key_option_as_first_prior
     assert_equal "id", actual
   end
 
+  def test_belongs_to_reflection_with_query_constraints_infers_correct_foreign_key
+    blog_foreign_key = Sharded::Comment.reflect_on_association(:blog).foreign_key
+    blog_post_foreign_key = Sharded::Comment.reflect_on_association(:blog_post).foreign_key
+
+    assert_equal "blog_id", blog_foreign_key
+    assert_equal ["blog_id", "blog_post_id"], blog_post_foreign_key
+  end
+
   private
     def assert_reflection(klass, association, options)
       assert reflection = klass.reflect_on_association(association)