Wrap enumerators to avoid JobIteration deprecation

Shopify · Dec 16, 2024 · 22ef7cb · 22ef7cb
1 parent 925a013
commit 22ef7cb
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/app/jobs/concerns/maintenance_tasks/task_job_concern.rb b/app/jobs/concerns/maintenance_tasks/task_job_concern.rb
@@ -33,6 +33,7 @@ def retry_on(*, **)
     def build_enumerator(_run, cursor:)
       cursor ||= @run.cursor
       self.cursor_position = cursor
+      enumerator_builder = self.enumerator_builder
       @collection_enum = @task.enumerator_builder(cursor: cursor)
 
       @collection_enum ||= case (collection = @task.collection)
@@ -75,6 +76,9 @@ def build_enumerator(_run, cursor:)
         MSG
       end
 
+      unless @collection_enum.is_a?(JobIteration.enumerator_builder::Wrapper)
+        @collection_enum = enumerator_builder.wrap(enumerator_builder, @collection_enum)
+      end
       throttle_enumerator(@collection_enum)
     end