Merge pull request rails#47901 from fatkodima/batches-composite-pks

Support batching using composite primary keys

Author: eileencodes

activerecord/lib/active_record/relation/batches.rb

@@ -227,7 +227,7 @@ def in_batches(of: 1000, start: nil, finish: nil, load: false, error_on_ignore:
         batch_limit = remaining if remaining < batch_limit
       end
 
-      relation = relation.reorder(batch_order(order)).limit(batch_limit)
+      relation = relation.reorder(*batch_order(order)).limit(batch_limit)
       relation = apply_limits(relation, start, finish, order)
       relation.skip_query_cache! # Retaining the results in the query cache would undermine the point of batching
       batch_relation = relation
@@ -240,15 +240,15 @@ def in_batches(of: 1000, start: nil, finish: nil, load: false, error_on_ignore:
           yielded_relation = where(primary_key => ids)
           yielded_relation.load_records(records)
         elsif (empty_scope && use_ranges != false) || use_ranges
-          ids = batch_relation.pluck(primary_key)
+          ids = batch_relation.ids
           finish = ids.last
           if finish
             yielded_relation = apply_finish_limit(batch_relation, finish, order)
             yielded_relation = yielded_relation.except(:limit, :order)
             yielded_relation.skip_query_cache!(false)
           end
         else
-          ids = batch_relation.pluck(primary_key)
+          ids = batch_relation.ids
           yielded_relation = where(primary_key => ids)
         end
 
@@ -273,8 +273,8 @@ def in_batches(of: 1000, start: nil, finish: nil, load: false, error_on_ignore:
           end
         end
 
-        batch_relation = relation.where(
-          predicate_builder[primary_key, primary_key_offset, order == :desc ? :lt : :gt]
+        batch_relation = batch_condition(
+          relation, primary_key, primary_key_offset, order == :desc ? :lt : :gt
         )
       end
     end
@@ -287,15 +287,32 @@ def apply_limits(relation, start, finish, order)
       end
 
       def apply_start_limit(relation, start, order)
-        relation.where(predicate_builder[primary_key, start, order == :desc ? :lteq : :gteq])
+        batch_condition(relation, primary_key, start, order == :desc ? :lteq : :gteq)
       end
 
       def apply_finish_limit(relation, finish, order)
-        relation.where(predicate_builder[primary_key, finish, order == :desc ? :gteq : :lteq])
+        batch_condition(relation, primary_key, finish, order == :desc ? :gteq : :lteq)
+      end
+
+      def batch_condition(relation, columns, values, operator)
+        columns = Array(columns)
+        values = Array(values)
+        cursor_positions = columns.zip(values)
+
+        first_clause_column, first_clause_value = cursor_positions.pop
+        where_clause = predicate_builder[first_clause_column, first_clause_value, operator]
+
+        cursor_positions.reverse_each do |column_name, value|
+          where_clause = predicate_builder[column_name, value, operator == :lteq ? :lt : :gt].or(
+            predicate_builder[column_name, value, :eq].and(where_clause)
+          )
+        end
+
+        relation.where(where_clause)
       end
 
       def batch_order(order)
-        table[primary_key].public_send(order)
+        Array(primary_key).map { |column| table[column].public_send(order) }
       end
 
       def act_on_ignored_order(error_on_ignore)

activerecord/test/cases/batches_test.rb

@@ -4,9 +4,10 @@
 require "models/comment"
 require "models/post"
 require "models/subscriber"
+require "models/cpk"
 
 class EachTest < ActiveRecord::TestCase
-  fixtures :posts, :subscribers
+  fixtures :posts, :subscribers, :cpk_orders
 
   def setup
     @posts = Post.order("id asc")
@@ -769,4 +770,30 @@ def test_find_in_batches_should_return_a_sized_enumerator
       end
     end
   end
+
+  test ".find_each iterates over composite primary key" do
+    orders = Cpk::Order.order(*Cpk::Order.primary_key).to_a
+    Cpk::Order.find_each(batch_size: 1).with_index do |order, index|
+      assert_equal orders[index], order
+    end
+  end
+
+  test ".in_batches should start from the start option when using composite primary key" do
+    order = Cpk::Order.second
+    relation = Cpk::Order.in_batches(of: 1, start: order.id).first
+    assert_equal order, relation.first
+  end
+
+  test ".in_batches should end at the finish option when using composite primary key" do
+    order = Cpk::Order.second_to_last
+    relation = Cpk::Order.in_batches(of: 1, finish: order.id).reverse_each.first
+    assert_equal order, relation.last
+  end
+
+  test ".in_batches with scope and using composite primary key" do
+    order1, order2 = Cpk::Order.first(2)
+    shop_id, id = order1.id
+    relation = Cpk::Order.where("shop_id > ? OR shop_id = ? AND id > ?", shop_id, shop_id, id).in_batches(of: 1).first
+    assert_equal order2, relation.first
+  end
 end