Define #count automatically for batch collections using the enumerator

Co-authored-by: Étienne Barrié <etienne.barrie@shopify.com>

app/jobs/concerns/maintenance_tasks/task_job_concern.rb

@@ -32,6 +32,7 @@ def retry_on(*, **)
     def build_enumerator(_run, cursor:)
       cursor ||= @run.cursor
       collection = @task.collection
+      @enumerator = nil
 
       collection_enum = case collection
       when ActiveRecord::Relation
@@ -43,13 +44,13 @@ def build_enumerator(_run, cursor:)
             a batch enumerator with the "start" or "finish" options.
           MSG
         end
-        enumerator_builder.active_record_on_batch_relations(
+        # For now, only support automatic count based on the enumerator for
+        # batches
+        @enumerator = enumerator_builder.active_record_on_batch_relations(
           collection.relation,
           cursor: cursor,
           batch_size: collection.batch_size,
         )
-        @run.update!(tick_total: enumerator.size)
-        enumerator
       when Array
         enumerator_builder.build_array_enumerator(collection, cursor: cursor)
       when CSV
@@ -102,8 +103,9 @@ def before_perform
     end
 
     def on_start
-      @run.tick_total = @task.count unless @run.tick_total
-      @run.update!(started_at: Time.now)
+      count = @task.count
+      count = @enumerator&.size if count == :no_count
+      @run.update!(started_at: Time.now, tick_total: count)
     end
 
     def on_complete

app/tasks/maintenance_tasks/task.rb

@@ -131,6 +131,7 @@ def process(_item)
     #
     # @return [Integer, nil]
     def count
+      :no_count
     end
   end
 end

test/dummy/app/tasks/maintenance/update_posts_in_batches_task.rb

@@ -5,10 +5,6 @@ def collection
       Post.in_batches(of: 5)
     end
 
-    def count
-      Post.count
-    end
-
     def process(batch_of_posts)
       batch_of_posts.update_all(content: "New content added on #{Time.now.utc}")
     end

test/tasks/maintenance_tasks/task_test.rb

@@ -54,9 +54,9 @@ class TaskTest < ActiveSupport::TestCase
       assert_equal 2, Maintenance::TestTask.count
     end
 
-    test "#count is nil by default" do
+    test "#count is :no_count by default" do
       task = Task.new
-      assert_nil task.count
+      assert_equal(:no_count, task.count)
     end
 
     test "#collection raises NoMethodError" do