Disable use of start and finish options on batch enumerator

Shopify · May 25, 2021 · 736550a · 736550a
1 parent f920669
commit 736550a
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 5 deletions.
diff --git a/app/jobs/concerns/maintenance_tasks/task_job_concern.rb b/app/jobs/concerns/maintenance_tasks/task_job_concern.rb
@@ -37,21 +37,26 @@ def build_enumerator(_run, cursor:)
       when ActiveRecord::Relation
         enumerator_builder.active_record_on_records(collection, cursor: cursor)
       when ActiveRecord::Batches::BatchEnumerator
+        if collection.instance_variable_get(:@start) ||
+          collection.instance_variable_get(:@finish)
+          raise ArgumentError, "#{@task.class.name}#collection cannot support "\
+            "a batch enumerator with the +start+ or +finish+ options."
+        end
         relation = collection.instance_variable_get(:@relation)
         batch_size = collection.instance_variable_get(:@of)
-        enumerator_builder.active_record_on_batches(
+        enumerator_builder.active_record_on_batch_relations(
           relation,
           cursor: cursor,
           batch_size: batch_size,
-          as_relation: true
         )
       when Array
         enumerator_builder.build_array_enumerator(collection, cursor: cursor)
       when CSV
         JobIteration::CsvEnumerator.new(collection).rows(cursor: cursor)
       else
         raise ArgumentError, "#{@task.class.name}#collection must be either "\
-          "an Active Record Relation, Array, or CSV."
+          "an Active Record Relation, ActiveRecord::Batches::BatchEnumerator, "\
+          "Array, or CSV."
       end
 
       @task.throttle_conditions.reduce(collection_enum) do |enum, condition|

diff --git a/test/jobs/maintenance_tasks/task_job_test.rb b/test/jobs/maintenance_tasks/task_job_test.rb
@@ -247,8 +247,9 @@ class << self
       assert_predicate @run, :errored?
       assert_equal "ArgumentError", @run.error_class
       assert_empty @run.backtrace
-      expected_message = "Maintenance::TestTask#collection "\
-        "must be either an Active Record Relation, Array, or CSV."
+      expected_message = "Maintenance::TestTask#collection must be either an "\
+        "Active Record Relation, ActiveRecord::Batches::BatchEnumerator, "\
+        "Array, or CSV."
       assert_equal expected_message, @run.error_message
     end
 
@@ -368,5 +369,23 @@ class << self
 
       assert_equal Post.count, run.reload.tick_count
     end
+
+    test ".perform_now raises if +start+ or +finish+ options are used on batch enumerator" do
+      batch_enumerator = Post.in_batches(of: 5, start: 1, finish: 10)
+
+      Maintenance::UpdatePostsInBatchesTask.any_instance
+        .expects(:collection).returns(batch_enumerator)
+
+      run = Run.create!(task_name: "Maintenance::UpdatePostsInBatchesTask")
+      TaskJob.perform_now(run)
+
+      assert_predicate run.reload, :errored?
+      assert_equal "ArgumentError", run.error_class
+      assert_empty run.backtrace
+      expected_message = "Maintenance::UpdatePostsInBatchesTask#collection "\
+        "cannot support a batch enumerator with the +start+ or +finish+ "\
+        "options."
+      assert_equal expected_message, run.error_message
+    end
   end
 end