Merge pull request rails#49655 from Shopify/fix-cpk-update-all-with-join-subquery

Fix `update_all/delete_all` on CPK model relation with join subquery

Author: rafaelfranca

activerecord/lib/active_record/relation.rb

@@ -593,7 +593,12 @@ def update_all(updates)
 
       group_values_arel_columns = arel_columns(group_values.uniq)
       having_clause_ast = having_clause.ast unless having_clause.empty?
-      stmt = arel.compile_update(values, table[primary_key], having_clause_ast, group_values_arel_columns)
+      key = if klass.composite_primary_key?
+        primary_key.map { |pk| table[pk] }
+      else
+        table[primary_key]
+      end
+      stmt = arel.compile_update(values, key, having_clause_ast, group_values_arel_columns)
       klass.connection.update(stmt, "#{klass} Update All").tap { reset }
     end
 
@@ -726,7 +731,12 @@ def delete_all
 
       group_values_arel_columns = arel_columns(group_values.uniq)
       having_clause_ast = having_clause.ast unless having_clause.empty?
-      stmt = arel.compile_delete(table[primary_key], having_clause_ast, group_values_arel_columns)
+      key = if klass.composite_primary_key?
+        primary_key.map { |pk| table[pk] }
+      else
+        table[primary_key]
+      end
+      stmt = arel.compile_delete(key, having_clause_ast, group_values_arel_columns)
 
       klass.connection.delete(stmt, "#{klass} Delete All").tap { reset }
     end

activerecord/lib/arel/tree_manager.rb

@@ -21,7 +21,11 @@ def order(*expr)
       end
 
       def key=(key)
-        @ast.key = Nodes.build_quoted(key)
+        @ast.key = if key.is_a?(Array)
+          key.map { |k| Nodes.build_quoted(k) }
+        else
+          Nodes.build_quoted(key)
+        end
       end
 
       def key

activerecord/lib/arel/visitors/to_sql.rb

@@ -930,7 +930,8 @@ def prepare_update_statement(o)
             stmt.limit = nil
             stmt.offset = nil
             stmt.orders = []
-            stmt.wheres = [Nodes::In.new(o.key, [build_subselect(o.key, o)])]
+            columns = Arel::Nodes::Grouping.new(o.key)
+            stmt.wheres = [Nodes::In.new(columns, [build_subselect(o.key, o)])]
             stmt.relation = o.relation.left if has_join_sources?(o)
             stmt.groups = o.groups unless o.groups.empty?
             stmt.havings = o.havings unless o.havings.empty?

activerecord/test/cases/arel/select_manager_test.rb

@@ -1014,7 +1014,7 @@ def test_join_sources
 
         _(stmt.to_sql).must_be_like %{
           UPDATE "users" SET foo = bar
-          WHERE "users"."id" IN (SELECT "users"."id" FROM "users" LIMIT 1)
+          WHERE ("users"."id") IN (SELECT "users"."id" FROM "users" LIMIT 1)
         }
       end
 
@@ -1028,7 +1028,7 @@ def test_join_sources
 
         _(stmt.to_sql).must_be_like %{
           UPDATE "users" SET foo = bar
-          WHERE "users"."id" IN (SELECT "users"."id" FROM "users" ORDER BY foo)
+          WHERE ("users"."id") IN (SELECT "users"."id" FROM "users" ORDER BY foo)
         }
       end
 
@@ -1053,7 +1053,7 @@ def test_join_sources
         stmt = manager.compile_update({ table[:id] => 1 }, Arel::Attributes::Attribute.new(table, "id"))
 
         _(stmt.to_sql).must_be_like %{
-          UPDATE "users" SET "id" = 1 WHERE "users"."id" IN (SELECT "users"."id" FROM "users" WHERE "users"."foo" = 10 LIMIT 42)
+          UPDATE "users" SET "id" = 1 WHERE ("users"."id") IN (SELECT "users"."id" FROM "users" WHERE "users"."foo" = 10 LIMIT 42)
         }
       end
     end

activerecord/test/cases/relation/delete_all_test.rb

@@ -6,9 +6,10 @@
 require "models/pet"
 require "models/toy"
 require "models/comment"
+require "models/cpk"
 
 class DeleteAllTest < ActiveRecord::TestCase
-  fixtures :authors, :author_addresses, :comments, :posts, :pets, :toys
+  fixtures :authors, :author_addresses, :comments, :posts, :pets, :toys, :cpk_orders, :cpk_order_agreements
 
   def test_destroy_all
     davids = Author.where(name: "David")
@@ -128,4 +129,10 @@ def test_delete_all_with_order_and_limit_and_offset_deletes_subset_only
     assert_raise(ActiveRecord::RecordNotFound) { posts(:thinking) }
     assert posts(:welcome)
   end
+
+  def test_delete_all_composite_model_with_join_subquery
+    agreement = cpk_order_agreements(:order_agreement_three)
+    join_scope = Cpk::Order.joins(:order_agreements).where(order_agreements: { signature: agreement.signature })
+    assert_equal 1, join_scope.delete_all
+  end
 end

activerecord/test/cases/relation/update_all_test.rb

@@ -14,9 +14,11 @@
 require "models/tag"
 require "models/tagging"
 require "models/warehouse_thing"
+require "models/cpk"
 
 class UpdateAllTest < ActiveRecord::TestCase
-  fixtures :authors, :author_addresses, :comments, :developers, :posts, :people, :pets, :toys, :tags, :taggings, "warehouse-things"
+  fixtures :authors, :author_addresses, :comments, :developers, :posts, :people, :pets, :toys, :tags,
+    :taggings, "warehouse-things", :cpk_orders, :cpk_order_agreements
 
   class TopicWithCallbacks < ActiveRecord::Base
     self.table_name = :topics
@@ -300,6 +302,12 @@ def test_klass_level_touch_all
     end
   end
 
+  def test_update_all_composite_model_with_join_subquery
+    agreement = cpk_order_agreements(:order_agreement_three)
+    join_scope = Cpk::Order.joins(:order_agreements).where(order_agreements: { signature: agreement.signature })
+    assert_equal 1, join_scope.update_all(status: "shipped")
+  end
+
   # Oracle UPDATE does not support ORDER BY
   unless current_adapter?(:OracleAdapter)
     def test_update_all_ignores_order_without_limit_from_association