Destroy associations async respect query constraints

The codepaths related to destroying associations asynchronously now
consider when query constraints are present. In most cases, this means
interpreting the primary key as an array of columns, and identifying
associated records by a tuple of these columns, where previously this
would've been a single ID. In each of the callsites, we use branching.
This is done to maintain backwards compatibility and ensure the
signature of the destroy job remains stable: it has consumers outside of
Rails.

activerecord/lib/active_record/associations/belongs_to_association.rb

@@ -11,8 +11,13 @@ module ActiveRecord
         when :destroy
           raise ActiveRecord::Rollback unless target.destroy
         when :destroy_async
-          id = owner.public_send(reflection.foreign_key.to_sym)
-          primary_key_column = reflection.active_record_primary_key.to_sym
+          if reflection.foreign_key.is_a?(Array)
+            primary_key_column = reflection.active_record_primary_key.map(&:to_sym)
+            id = reflection.foreign_key.map { |col| owner.public_send(col.to_sym) }
+          else
+            primary_key_column = reflection.active_record_primary_key.to_sym
+            id = owner.public_send(reflection.foreign_key.to_sym)
+          end
 
           enqueue_destroy_association(
             owner_model_name: owner.class.to_s,

activerecord/lib/active_record/associations/has_many_association.rb

@@ -33,10 +33,12 @@ module ActiveRecord
 
           unless target.empty?
             association_class = target.first.class
-            primary_key_column = association_class.primary_key.to_sym
-
-            ids = target.collect do |assoc|
-              assoc.public_send(primary_key_column)
+            if association_class.query_constraints_list
+              primary_key_column = association_class.query_constraints_list.map(&:to_sym)
+              ids = target.collect { |assoc| primary_key_column.map { |col| assoc.public_send(col) } }
+            else
+              primary_key_column = association_class.primary_key.to_sym
+              ids = target.collect { |assoc| assoc.public_send(primary_key_column) }
             end
 
             ids.each_slice(owner.class.destroy_association_async_batch_size || ids.size) do |ids_batch|

activerecord/lib/active_record/associations/has_one_association.rb

@@ -33,8 +33,13 @@ module ActiveRecord
             target.destroy
             throw(:abort) unless target.destroyed?
           when :destroy_async
-            primary_key_column = target.class.primary_key.to_sym
-            id = target.public_send(primary_key_column)
+            if target.class.query_constraints_list
+              primary_key_column = target.class.query_constraints_list.map(&:to_sym)
+              id = primary_key_column.map { |col| target.public_send(col) }
+            else
+              primary_key_column = target.class.primary_key.to_sym
+              id = target.public_send(primary_key_column)
+            end
 
             enqueue_destroy_association(
               owner_model_name: owner.class.to_s,

activerecord/lib/active_record/destroy_association_async_job.rb

@@ -23,8 +23,15 @@ module ActiveRecord
         raise DestroyAssociationAsyncError, "owner record not destroyed"
       end
 
-      association_model.where(association_primary_key_column => association_ids).find_each do |r|
-        r.destroy
+      if association_model.query_constraints_list
+        association_ids
+          .map { |assoc_ids| association_model.where(association_primary_key_column.zip(assoc_ids).to_h) }
+          .inject(&:or)
+          .find_each { |r| r.destroy }
+      else
+        association_model.where(association_primary_key_column => association_ids).find_each do |r|
+          r.destroy
+        end
       end
     end
 

activerecord/test/activejob/destroy_association_async_test.rb

@@ -18,6 +18,12 @@ require "models/dl_keyed_has_one"
 require "models/dl_keyed_join"
 require "models/dl_keyed_has_many"
 require "models/dl_keyed_has_many_through"
+require "models/sharded/blog_post_destroy_async"
+require "models/sharded/comment_destroy_async"
+require "models/sharded/tag"
+require "models/sharded/blog_post"
+require "models/sharded/blog_post_tag"
+require "models/sharded/blog"
 
 class DestroyAssociationAsyncTest < ActiveRecord::TestCase
   self.use_transactional_tests = false
@@ -43,6 +49,39 @@ class DestroyAssociationAsyncTest < ActiveRecord::TestCase
     BookDestroyAsync.delete_all
   end
 
+  test "destroying a record destroys has_many :through associated by composite primary key using a job" do
+    blog = Sharded::Blog.create!
+    blog_post = Sharded::BlogPostDestroyAsync.create!(blog_id: blog.id)
+
+    tag1 = Sharded::Tag.create!(name: "Short Read", blog_id: blog.id)
+    tag2 = Sharded::Tag.create!(name: "Science", blog_id: blog.id)
+
+    blog_post.tags << [tag1, tag2]
+
+    blog_post.save!
+
+    assert_enqueued_jobs 1, only: ActiveRecord::DestroyAssociationAsyncJob do
+      blog_post.destroy
+    end
+
+    sql = capture_sql do
+      assert_difference -> { Sharded::Tag.count }, -2 do
+        perform_enqueued_jobs only: ActiveRecord::DestroyAssociationAsyncJob
+      end
+    end
+
+    delete_sqls = sql.select { |sql| sql.start_with?("DELETE") }
+    assert_equal 2, delete_sqls.count
+
+    delete_sqls.each do |sql|
+      assert_match(/#{Regexp.escape(Sharded::Tag.connection.quote_table_name("sharded_tags.blog_id"))} =/, sql)
+    end
+  ensure
+    Sharded::Tag.delete_all
+    Sharded::BlogPostDestroyAsync.delete_all
+    Sharded::Blog.delete_all
+  end
+
   test "destroying a scoped has_many through only deletes within the scope deleted" do
     tag = Tag.create!(name: "Der be treasure")
     tag2 = Tag.create!(name: "Der be rum")
@@ -128,6 +167,33 @@ class DestroyAssociationAsyncTest < ActiveRecord::TestCase
     BookDestroyAsync.delete_all
   end
 
+  test "belongs to associated by composite primary key" do
+    blog = Sharded::Blog.create!
+    blog_post = Sharded::BlogPostDestroyAsync.create!(blog_id: blog.id)
+    comment = Sharded::CommentDestroyAsync.create!(body: "Great post! :clap:")
+
+    comment.blog_post = blog_post
+    comment.save!
+
+    assert_enqueued_jobs 1, only: ActiveRecord::DestroyAssociationAsyncJob do
+      comment.destroy
+    end
+
+    sql = capture_sql do
+      assert_difference -> { Sharded::BlogPostDestroyAsync.count }, -1 do
+        perform_enqueued_jobs only: ActiveRecord::DestroyAssociationAsyncJob
+      end
+    end
+
+    delete_sqls = sql.select { |sql| sql.start_with?("DELETE") }
+    assert_equal 1, delete_sqls.count
+    assert_match(/#{Regexp.escape(Sharded::BlogPost.connection.quote_table_name("sharded_blog_posts.blog_id"))} =/, delete_sqls.first)
+  ensure
+    Sharded::BlogPostDestroyAsync.delete_all
+    Sharded::CommentDestroyAsync.delete_all
+    Sharded::Blog.delete_all
+  end
+
   test "enqueues belongs_to to be deleted with custom primary key" do
     belongs = DlKeyedBelongsTo.create!
     parent = DestroyAsyncParent.create!
@@ -218,6 +284,39 @@ class DestroyAssociationAsyncTest < ActiveRecord::TestCase
     DestroyAsyncParent.delete_all
   end
 
+  test "has_many associated with composite primary key" do
+    blog = Sharded::Blog.create!
+    blog_post = Sharded::BlogPostDestroyAsync.create!(blog_id: blog.id)
+
+    comment1 = Sharded::CommentDestroyAsync.create!(body: "Great post! :clap:")
+    comment2 = Sharded::CommentDestroyAsync.create!(body: "Terrible post! :thumbs-down:")
+
+    blog_post.comments << [comment1, comment2]
+
+    blog_post.save!
+
+    assert_enqueued_jobs 1, only: ActiveRecord::DestroyAssociationAsyncJob do
+      blog_post.destroy
+    end
+
+    sql = capture_sql do
+      assert_difference -> { Sharded::CommentDestroyAsync.count }, -2 do
+        perform_enqueued_jobs only: ActiveRecord::DestroyAssociationAsyncJob
+      end
+    end
+
+    delete_sqls = sql.select { |sql| sql.start_with?("DELETE") }
+    assert_equal 2, delete_sqls.count
+
+    delete_sqls.each do |sql|
+      assert_match(/#{Regexp.escape(Sharded::Tag.connection.quote_table_name("sharded_comments.blog_id"))} =/, sql)
+    end
+  ensure
+    Sharded::CommentDestroyAsync.delete_all
+    Sharded::BlogPostDestroyAsync.delete_all
+    Sharded::Blog.delete_all
+  end
+
   test "not enqueue the job if transaction is not committed" do
     dl_keyed_has_many = DlKeyedHasMany.new
     parent = DestroyAsyncParent.create!
@@ -327,8 +426,8 @@ class DestroyAssociationAsyncTest < ActiveRecord::TestCase
       raise ActiveRecord::Rollback
     end
     assert_no_enqueued_jobs only: ActiveRecord::DestroyAssociationAsyncJob
+  ensure
+    Tag.delete_all
+    BookDestroyAsync.delete_all
   end
-ensure
-  Tag.delete_all
-  BookDestroyAsync.delete_all
 end

activerecord/test/models/sharded/blog_post_destroy_async.rb

@@ -0,0 +1,14 @@
+# frozen_string_literal: true
+
+module Sharded
+  class BlogPostDestroyAsync < ActiveRecord::Base
+    self.table_name = :sharded_blog_posts
+    query_constraints :blog_id, :id
+
+    belongs_to :blog
+    has_many :comments, dependent: :destroy_async, query_constraints: [:blog_id, :blog_post_id], class_name: "Sharded::CommentDestroyAsync"
+
+    has_many :blog_post_tags, query_constraints: [:blog_id, :blog_post_id], class_name: "Sharded::BlogPostTag"
+    has_many :tags, through: :blog_post_tags, dependent: :destroy_async, class_name: "Sharded::Tag"
+  end
+end

activerecord/test/models/sharded/comment_destroy_async.rb

@@ -0,0 +1,12 @@
+# frozen_string_literal: true
+
+module Sharded
+  class CommentDestroyAsync < ActiveRecord::Base
+    self.table_name = :sharded_comments
+    query_constraints :blog_id, :id
+
+    belongs_to :blog_post, dependent: :destroy_async, query_constraints: [:blog_id, :blog_post_id], class_name: "Sharded::BlogPostDestroyAsync"
+    belongs_to :blog_post_by_id, class_name: "Sharded::BlogPostDestroyAsync", foreign_key: :blog_post_id
+    belongs_to :blog
+  end
+end