Support ordering using multiple directions for ActiveRecord enumerators

CHANGELOG.md

@@ -1,5 +1,11 @@
 ## master (unreleased)
 
+- Support ordering using multiple directions for ActiveRecord enumerators
+
+    ```ruby
+    active_record_records_enumerator(..., columns: [:shop_id, :id], order: [:asc, :desc])
+    ```
+
 - Support iterating over ActiveRecord models with composite primary keys
 
 - Use Arel to generate SQL in ActiveRecord enumerator

lib/sidekiq_iteration/active_record_enumerator.rb

@@ -17,21 +17,24 @@ def initialize(relation, columns: nil, batch_size: 100, order: :asc, cursor: nil
       end
 
       @relation = relation
+      @primary_key = relation.primary_key
+      columns = Array(columns || @primary_key).map(&:to_s)
 
-      unless order == :asc || order == :desc
-        raise ArgumentError, ":order must be :asc or :desc, got #{order.inspect}"
+      if (Array(order) - [:asc, :desc]).any?
+        raise ArgumentError, ":order must be :asc or :desc or an array consisting of :asc or :desc, got #{order.inspect}"
       end
 
-      @primary_key = relation.primary_key
-      columns = Array(columns || @primary_key).map(&:to_s)
+      if order.is_a?(Array) && order.size != columns.size
+        raise ArgumentError, ":order must include a direction for each batching column"
+      end
 
       @primary_key_index = primary_key_index(columns, relation)
       if @primary_key_index.nil? || (composite_primary_key? && @primary_key_index.any?(nil))
         raise ArgumentError, ":columns must include a primary key columns"
       end
 
       @batch_size = batch_size
-      @order = order
+      @order = batch_order(columns, order)
       @cursor = Array(cursor)
 
       if @cursor.present? && @cursor.size != columns.size
@@ -55,7 +58,7 @@ def initialize(relation, columns: nil, batch_size: 100, order: :asc, cursor: nil
       end
 
       @columns = columns
-      ordering = @columns.to_h { |column| [column, @order] }
+      ordering = @columns.zip(@order).to_h
       @base_relation = relation.reorder(ordering)
       @iteration_count = 0
     end
@@ -105,6 +108,14 @@ def primary_key_index(columns, relation)
         end
       end
 
+      def batch_order(columns, order)
+        if order.is_a?(Array)
+          order
+        else
+          [order] * columns.size
+        end
+      end
+
       def arel_column(column)
         if column.include?(".")
           Arel.sql(column)
@@ -210,19 +221,18 @@ def column_value(value)
       end
 
       def cursor_operators
-        leading_operator = @order == :asc ? :gt : :lt
-
         # Start from the record pointed by cursor when just starting.
-        last_operator =
-          if @order == :asc
-            first_iteration? ? :gteq : :gt
+        @columns.zip(@order).map do |column, order|
+          if column == @columns.last
+            if order == :asc
+              first_iteration? ? :gteq : :gt
+            else
+              first_iteration? ? :lteq : :lt
+            end
           else
-            first_iteration? ? :lteq : :lt
+            order == :asc ? :gt : :lt
           end
-
-        operators = [leading_operator] * (@columns.count - 1)
-        operators << last_operator
-        operators
+        end
       end
 
       def increment_iteration

lib/sidekiq_iteration/enumerators.rb

@@ -27,7 +27,7 @@ def array_enumerator(array, cursor:)
     # @option options :columns [Array<String, Symbol>, String, Symbol] used to build the actual query for iteration,
     #   defaults to primary key
     # @option options :batch_size [Integer] (100) size of the batch
-    # @option options :order [:asc, :desc] (:asc) specifies iteration order
+    # @option options :order [:asc, :desc, Array<:asc, :desc>] (:asc) specifies iteration order
     #
     # +columns:+ argument is used to build the actual query for iteration. +columns+: defaults to primary key:
     #

test/active_record_enumerator_test.rb

@@ -161,7 +161,7 @@ class ActiveRecordEnumeratorTest < TestCase
       assert_equal(expected_ids, records.map(&:id))
     end
 
-    test "order is configurable" do
+    test ":order with single direction" do
       enum = build_enumerator(order: :desc).batches
       product_batches = Product.order(id: :desc).take(4).in_groups_of(2).map { |products| [products, products.last.id] }
 
@@ -170,11 +170,32 @@ class ActiveRecordEnumeratorTest < TestCase
       end
     end
 
+    test ":order with multiple directions" do
+      enum = build_enumerator(relation: Order.all, order: [:asc, :desc]).records
+      orders = Order.order(shop_id: :asc, id: :desc).to_a
+      assert_equal(orders.size, enum.size)
+
+      enum.each_with_index do |(element, cursor), index|
+        order = orders[index]
+        assert_equal(order, element)
+        assert_equal([order.shop_id, order[:id]], cursor)
+      end
+    end
+
     test "raises on invalid order" do
-      error = assert_raises(ArgumentError) do
+      assert_raises_with_message(ArgumentError, ":order must be :asc or :desc") do
         build_enumerator(order: :invalid)
       end
-      assert_equal ":order must be :asc or :desc, got :invalid", error.message
+
+      assert_raises_with_message(ArgumentError, ":order must be :asc or :desc") do
+        build_enumerator(order: [:asc, :invalid])
+      end
+    end
+
+    test "raises on mismatched order array size" do
+      assert_raises_with_message(ArgumentError, ":order must include a direction for each batching column") do
+        build_enumerator(order: [:asc, :desc])
+      end
     end
 
     test "can be resumed" do