Merge pull request #46962 from Shopify/foundation-to-support-composite-foreign-keys-in-associations

Support composite foreign keys in associations

activerecord/lib/active_record/associations/association_scope.rb

@@ -35,7 +35,7 @@ module ActiveRecord
         binds = []
         last_reflection = chain.last
 
-        binds << last_reflection.join_id_for(owner)
+        binds.push(*last_reflection.join_id_for(owner))
         if last_reflection.type
           binds << owner.class.polymorphic_name
         end
@@ -56,12 +56,15 @@ module ActiveRecord
         end
 
         def last_chain_scope(scope, reflection, owner)
-          primary_key = reflection.join_primary_key
+          primary_key = Array(reflection.join_primary_key)
-          foreign_key = reflection.join_foreign_key
+          foreign_key = Array(reflection.join_foreign_key)
 
           table = reflection.aliased_table
-          value = transform_value(owner[foreign_key])
+          primary_key_foreign_key_pairs = primary_key.zip(foreign_key)
-          scope = apply_scope(scope, table, primary_key, value)
+          primary_key_foreign_key_pairs.each do |join_key, foreign_key|
+            value = transform_value(owner[foreign_key])
+            scope = apply_scope(scope, table, join_key, value)
+          end
 
           if reflection.type
             polymorphic_type = transform_value(owner.class.polymorphic_name)

activerecord/lib/active_record/associations/belongs_to_association.rb

@@ -131,7 +131,7 @@ module ActiveRecord
         end
 
         def foreign_key_present?
-          owner._read_attribute(reflection.foreign_key)
+          Array(reflection.foreign_key).all? { |fk| owner._read_attribute(fk) }
         end
 
         def invertible_for?(record)

activerecord/lib/active_record/counter_cache.rb

@@ -65,7 +65,7 @@ module ActiveRecord
           updates.merge!(touch_updates)
         end
 
-        unscoped.where(primary_key => object.id).update_all(updates) if updates.any?
+        unscoped.where(primary_key => [object.id]).update_all(updates) if updates.any?
 
         true
       end

activerecord/lib/active_record/reflection.rb

@@ -490,7 +490,12 @@ module ActiveRecord
       end
 
       def foreign_key
-        @foreign_key ||= -(options[:foreign_key]&.to_s || derive_foreign_key)
+        @foreign_key ||= if options[:foreign_key] && options[:foreign_key].is_a?(Array)
+          # composite foreign keys support
+          options[:foreign_key].map { |fk| fk.to_s.freeze }.freeze
+        else
+          -(options[:foreign_key]&.to_s || derive_foreign_key)
+        end
       end
 
       def association_foreign_key
@@ -502,7 +507,13 @@ module ActiveRecord
       end
 
       def active_record_primary_key
-        @active_record_primary_key ||= -(options[:primary_key]&.to_s || primary_key(active_record))
+        @active_record_primary_key ||= if options[:foreign_key] && options[:foreign_key].is_a?(Array)
+          active_record.query_constraints_list
+        else
+          -(options[:primary_key]&.to_s || primary_key(active_record))
+        end
+
+        @active_record_primary_key
       end
 
       def join_primary_key(klass = nil)
@@ -530,7 +541,7 @@ module ActiveRecord
       end
 
       def join_id_for(owner) # :nodoc:
-        owner[join_foreign_key]
+        Array(join_foreign_key).map { |key| owner[key] }
       end
 
       def through_reflection
@@ -764,7 +775,9 @@ module ActiveRecord
 
       # klass option is necessary to support loading polymorphic associations
       def association_primary_key(klass = nil)
-        if primary_key = options[:primary_key]
+        if options[:foreign_key] && options[:foreign_key].is_a?(Array)
+          (klass || self.klass).query_constraints_list
+        elsif primary_key = options[:primary_key]
           @association_primary_key ||= -primary_key.to_s
         else
           primary_key(klass || self.klass)

activerecord/test/cases/associations_test.rb

@@ -34,11 +34,15 @@ require "models/shipping_line"
 require "models/essay"
 require "models/member"
 require "models/membership"
+require "models/sharded/blog"
+require "models/sharded/blog_post"
+require "models/sharded/comment"
+
 
 class AssociationsTest < ActiveRecord::TestCase
   fixtures :accounts, :companies, :developers, :projects, :developers_projects,
            :computers, :people, :readers, :authors, :author_addresses, :author_favorites,
-           :comments, :posts
+           :comments, :posts, :sharded_blogs, :sharded_blog_posts, :sharded_comments
 
   def test_eager_loading_should_not_change_count_of_children
     liquid = Liquid.create(name: "salty")
@@ -124,6 +128,43 @@ class AssociationsTest < ActiveRecord::TestCase
     firm = companies(:first_firm)
     assert_equal [:foo], firm.association_with_references.references_values
   end
+
+  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_post_blog_one)
+
+    assert_equal(blog_post, comment.blog_post)
+  end
+
+  def test_belongs_to_a_model_with_composite_primary_key_uses_composite_pk_in_sql
+    comment = sharded_comments(:great_comment_blog_post_one)
+
+    sql = capture_sql do
+      comment.blog_post
+    end.first
+
+    assert_match(/#{Regexp.escape(Sharded::BlogPost.connection.quote_table_name("sharded_blog_posts.blog_id"))} =/, sql)
+    assert_match(/#{Regexp.escape(Sharded::BlogPost.connection.quote_table_name("sharded_blog_posts.id"))} =/, sql)
+  end
+
+  def test_has_many_association_with_composite_foreign_key_loads_records
+    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_post_blog_one)
+
+    sql = capture_sql do
+      blog_post.comments.to_a
+    end.first
+
+    assert_match(/#{Regexp.escape(Sharded::Comment.connection.quote_table_name("sharded_comments.blog_post_id"))} =/, sql)
+    assert_match(/#{Regexp.escape(Sharded::Comment.connection.quote_table_name("sharded_comments.blog_id"))} =/, sql)
+  end
 end
 
 class AssociationProxyTest < ActiveRecord::TestCase

activerecord/test/fixtures/sharded_blog_posts.yml

@@ -0,0 +1,10 @@
+_fixture:
+  model_class: Sharded::BlogPost
+
+great_post_blog_one:
+  title: "My first post!"
+  blog_id: <%= ActiveRecord::FixtureSet.identify(:sharded_blog_one) %>
+
+great_post_blog_two:
+  title: "My first post!"
+  blog_id: <%= ActiveRecord::FixtureSet.identify(:sharded_blog_two) %>

activerecord/test/fixtures/sharded_blogs.yml

@@ -0,0 +1,8 @@
+_fixture:
+  model_class: Sharded::Blog
+
+sharded_blog_one:
+  name: "Blog One"
+
+sharded_blog_two:
+  name: "Blog two"

activerecord/test/fixtures/sharded_comments.yml

@@ -0,0 +1,17 @@
+_fixture:
+  model_class: Sharded::Comment
+
+great_comment_blog_post_one:
+  body: "I really enjoyed the post!"
+  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_post_blog_one) %>
+  blog_id: <%= ActiveRecord::FixtureSet.identify(:sharded_blog_one) %>
+
+great_comment_blog_post_two:
+  body: "I really enjoyed the post!"
+  blog_post_id: <%= ActiveRecord::FixtureSet.identify(:great_post_blog_two) %>
+  blog_id: <%= ActiveRecord::FixtureSet.identify(:sharded_blog_two) %>

activerecord/test/models/sharded/blog.rb

@@ -0,0 +1,12 @@
+# frozen_string_literal: true
+
+# Models under `Sharded` namespace represent an application that can be sharded by a `blog_id` column.
+# `Blog` model plays the role of a tenant in the application.
+# Being sharded by the `blog_id` means that queries to the database include the `blog_id` column in the clauses
+# which serves as a sharding key allows presumed sharding implementation to route the query to a correct shard.
+
+module Sharded
+  class Blog < ActiveRecord::Base
+    self.table_name = :sharded_blogs
+  end
+end

activerecord/test/models/sharded/blog_post.rb

@@ -0,0 +1,11 @@
+# frozen_string_literal: true
+
+module Sharded
+  class BlogPost < ActiveRecord::Base
+    self.table_name = :sharded_blog_posts
+    query_constraints :blog_id, :id
+
+    belongs_to :blog
+    has_many :comments, foreign_key: [:blog_id, :blog_post_id]
+  end
+end

activerecord/test/models/sharded/comment.rb

@@ -0,0 +1,11 @@
+# frozen_string_literal: true
+
+module Sharded
+  class Comment < ActiveRecord::Base
+    self.table_name = :sharded_comments
+    query_constraints :blog_id, :id
+
+    belongs_to :blog_post, foreign_key: [:blog_id, :blog_post_id]
+    belongs_to :blog
+  end
+end

activerecord/test/schema/schema.rb

@@ -243,6 +243,21 @@ ActiveRecord::Schema.define do
     t.index [:clothing_type, :color], unique: true
   end
 
+  create_table :sharded_blogs, force: true do |t|
+    t.string :name
+  end
+
+  create_table :sharded_blog_posts, force: true do |t|
+    t.string :title
+    t.integer :blog_id
+  end
+
+  create_table :sharded_comments, force: true do |t|
+    t.string :body
+    t.integer :blog_post_id
+    t.integer :blog_id
+  end
+
   create_table :clubs, force: true do |t|
     t.string :name
     t.integer :category_id