Support nested iteration

CHANGELOG.md

@@ -1,5 +1,6 @@
 ### Master (unreleased)
 
+- [295](https://github.com/Shopify/job-iteration/pull/295) - Support nested iteration
 - [241](https://github.com/Shopify/job-iteration/pull/241) - Require Ruby 2.7+, dropping 2.6 support
 - [241](https://github.com/Shopify/job-iteration/pull/241) - Require Rails 6.0+, dropping 5.2 support
 - [240](https://github.com/Shopify/job-iteration/pull/240) - Allow setting inheritable per-job `job_iteration_max_job_runtime`

README.md

@@ -130,6 +130,27 @@ class CsvJob < ApplicationJob
 end
 ```
 
+```ruby
+class NestedIterationJob < ApplicationJob
+  include JobIteration::Iteration
+
+  def build_enumerator(cursor:)
+    enumerator_builder.nested(
+      [
+        ->(cursor) { enumerator_builder.active_record_on_records(Shop.all, cursor: cursor) },
+        ->(shop, cursor) { enumerator_builder.active_record_on_records(shop.products, cursor: cursor) },
+        ->(_shop, product, cursor) { enumerator_builder.active_record_on_batch_relations(product.product_variants, cursor: cursor) }
+      ],
+      cursor: cursor
+    )
+  end
+
+  def each_iteration(product_variants_relation)
+    # do something
+  end
+end
+```
+
 Iteration hooks into Sidekiq and Resque out of the box to support graceful interruption. No extra configuration is required.
 
 ## Guides

lib/job-iteration/enumerator_builder.rb

@@ -4,6 +4,7 @@
 require_relative "./active_record_enumerator"
 require_relative "./csv_enumerator"
 require_relative "./throttle_enumerator"
+require_relative "./nested_enumerator"
 require "forwardable"
 
 module JobIteration
@@ -146,6 +147,40 @@ def build_csv_enumerator(enumerable, cursor:)
       CsvEnumerator.new(enumerable).rows(cursor: cursor)
     end
 
+    # Builds Enumerator for nested iteration.
+    #
+    # @param enums [Array<Proc>] an Array of Procs, each should return an Enumerator.
+    #   Each proc from enums should accept the yielded items from the parent enumerators
+    #     and the `cursor` as its arguments.
+    #   Each proc's `cursor` argument is its part from the `build_enumerator`'s `cursor` array.
+    # @param cursor [Array<Object>] array of offsets for each of the enums to start iteration from
+    #
+    # @example
+    #   def build_enumerator(cursor:)
+    #     enumerator_builder.nested(
+    #       [
+    #         ->(cursor) {
+    #           enumerator_builder.active_record_on_records(Shop.all, cursor: cursor)
+    #         },
+    #         ->(shop, cursor) {
+    #           enumerator_builder.active_record_on_records(shop.products, cursor: cursor)
+    #         },
+    #         ->(_shop, product, cursor) {
+    #           enumerator_builder.active_record_on_batch_relations(product.product_variants, cursor: cursor)
+    #         }
+    #       ],
+    #       cursor: cursor
+    #     )
+    #   end
+    #
+    #   def each_iteration(product_variants_relation)
+    #     # do something
+    #   end
+    #
+    def build_nested_enumerator(enums, cursor:)
+      NestedEnumerator.new(enums, cursor: cursor).each
+    end
+
     alias_method :once, :build_once_enumerator
     alias_method :times, :build_times_enumerator
     alias_method :array, :build_array_enumerator
@@ -154,6 +189,7 @@ def build_csv_enumerator(enumerable, cursor:)
     alias_method :active_record_on_batch_relations, :build_active_record_enumerator_on_batch_relations
     alias_method :throttle, :build_throttle_enumerator
     alias_method :csv, :build_csv_enumerator
+    alias_method :nested, :build_nested_enumerator
 
     private
 

lib/job-iteration/nested_enumerator.rb

@@ -0,0 +1,40 @@
+# frozen_string_literal: true
+
+module JobIteration
+  # @private
+  class NestedEnumerator
+    def initialize(enums, cursor: nil)
+      unless enums.all?(Proc)
+        raise ArgumentError, "enums must contain only procs/lambdas"
+      end
+
+      if cursor && enums.size != cursor.size
+        raise ArgumentError, "cursor should have one item per enum"
+      end
+
+      @enums = enums
+      @cursor = cursor || Array.new(enums.size)
+    end
+
+    def each(&block)
+      return to_enum unless block_given?
+
+      iterate([], [], 0, &block)
+    end
+
+    private
+
+    def iterate(current_items, current_cursor, index, &block)
+      cursor = @cursor[index]
+      enum = @enums[index].call(*current_items, cursor)
+
+      enum.each do |item, cursor_value|
+        if index == @cursor.size - 1
+          yield item, current_cursor + [cursor_value]
+        else
+          iterate(current_items + [item], current_cursor + [cursor_value], index + 1, &block)
+        end
+      end
+    end
+  end
+end

test/test_helper.rb

@@ -42,6 +42,11 @@ def enqueue_at(job, _delay)
 ActiveJob::Base.queue_adapter = :iteration_test
 
 class Product < ActiveRecord::Base
+  has_many :comments
+end
+
+class Comment < ActiveRecord::Base
+  belongs_to :product
 end
 
 host = ENV["USING_DEV"] == "1" ? "job-iteration.railgun" : "localhost"
@@ -68,9 +73,16 @@ class Product < ActiveRecord::Base
   config.redis = { host: host }
 end
 
-ActiveRecord::Base.connection.create_table(Product.table_name, force: true) do |t|
-  t.string(:name)
-  t.timestamps
+ActiveRecord::Schema.define do
+  create_table(:products, force: true) do |t|
+    t.string(:name)
+    t.timestamps
+  end
+
+  create_table(:comments, force: true) do |t|
+    t.string(:content)
+    t.belongs_to(:product)
+  end
 end
 
 module LoggingHelpers
@@ -118,9 +130,16 @@ class IterationUnitTest < ActiveSupport::TestCase
   end
 
   def insert_fixtures
-    10.times do |n|
-      Product.create!(name: "lipstick #{n}")
-    end
+    now = Time.now
+    products = 10.times.map { |n| { name: "lipstick #{n}", created_at: now - n, updated_at: now - n } }
+    Product.insert_all!(products)
+
+    comments = Product.order(:id).limit(3).map.with_index do |product, index|
+      comments_count = index + 1
+      comments_count.times.map { |n| { content: "#{product.name} comment ##{n}", product_id: product.id } }
+    end.flatten
+
+    Comment.insert_all!(comments)
   end
 
   def truncate_fixtures

test/unit/enumerator_builder_test.rb

@@ -60,6 +60,20 @@ class EnumeratorBuilderTest < ActiveSupport::TestCase
       enumerator_builder(wraps: 0).build_csv_enumerator(CSV.new("test"), cursor: nil)
     end
 
+    test_builder_method(:build_nested_enumerator) do
+      enumerator_builder(wraps: 0).build_nested_enumerator(
+        [
+          ->(cursor) {
+            enumerator_builder.build_active_record_enumerator_on_records(Product.all, cursor: cursor)
+          },
+          ->(product, cursor) {
+            enumerator_builder.build_active_record_enumerator_on_records(product.comments, cursor: cursor)
+          },
+        ],
+        cursor: nil,
+      )
+    end
+
     # checks that all the non-alias methods were tested
     raise "methods not tested: #{methods.inspect}" unless methods.empty?
 

test/unit/nested_enumerator_test.rb

@@ -0,0 +1,81 @@
+# frozen_string_literal: true
+
+require "test_helper"
+
+module JobIteration
+  class NestedEnumeratorTest < IterationUnitTest
+    test "accepts only callables as enums" do
+      error = assert_raises(ArgumentError) do
+        build_enumerator(outer: [[1, 2, 3].each])
+      end
+      assert_equal("enums must contain only procs/lambdas", error.message)
+    end
+
+    test "raises when cursor is not of the same size as enums" do
+      error = assert_raises(ArgumentError) do
+        build_enumerator(cursor: [Product.first.id])
+      end
+      assert_equal("cursor should have one item per enum", error.message)
+    end
+
+    test "yields enumerator when called without a block" do
+      enum = build_enumerator
+      assert enum.is_a?(Enumerator)
+      assert_nil enum.size
+    end
+
+    test "yields every nested record with their cursor position" do
+      enum = build_enumerator
+
+      products = Product.includes(:comments).order(:id).take(3)
+      comments = products.map do |product|
+        product.comments.sort_by(&:id).map { |comment| [comment, [product.id, comment.id]] }
+      end.flatten(1)
+
+      enum.each_with_index do |(comment, cursor), index|
+        expected_comment, expected_cursor = comments[index]
+        assert_equal(expected_comment, comment)
+        assert_equal(expected_cursor, cursor)
+      end
+    end
+
+    test "cursor can be used to resume" do
+      enum = build_enumerator
+      _first_comment, first_cursor = enum.next
+      second_comment, second_cursor = enum.next
+
+      enum = build_enumerator(cursor: first_cursor)
+      assert_equal([second_comment, second_cursor], enum.first)
+    end
+
+    test "doesn't yield anything if contains empty enum" do
+      enum = ->(cursor, _product) { records_enumerator(Comment.none, cursor: cursor) }
+      enum = build_enumerator(inner: enum)
+      assert_empty(enum.to_a)
+    end
+
+    test "works with single level nesting" do
+      enum = build_enumerator(inner: nil)
+      products = Product.order(:id).to_a
+
+      enum.each_with_index do |(product, cursor), index|
+        assert_equal(products[index], product)
+        assert_equal([products[index].id], cursor)
+      end
+    end
+
+    private
+
+    def build_enumerator(
+      outer: ->(cursor) { records_enumerator(Product.all, cursor: cursor) },
+      inner: ->(product, cursor) { records_enumerator(product.comments, cursor: cursor) },
+      cursor: nil
+    )
+      NestedEnumerator.new([outer, inner].compact, cursor: cursor).each
+    end
+
+    def records_enumerator(scope, cursor: nil)
+      ActiveRecordEnumerator.new(scope, cursor: cursor).records
+    end
+  end
+end