lib/job-iteration/enumerator_builder.rb in job-iteration-1.1.3 vs lib/job-iteration/enumerator_builder.rb in job-iteration-1.1.4

@@ -1,8 +1,9 @@
 # frozen_string_literal: true
 require_relative "./active_record_enumerator"
 require_relative "./csv_enumerator"
+require_relative "./throttle_enumerator"
 require "forwardable"
 
 module JobIteration
   class EnumeratorBuilder
     extend Forwardable
@@ -59,20 +60,10 @@
         end
 
       wrap(self, enumerable.each_with_index.drop(drop).to_enum { enumerable.size })
     end
 
-    # Builds Enumerator from a lock queue instance that belongs to a job.
-    # The helper is only to be used from jobs that use LockQueue module.
-    def build_lock_queue_enumerator(lock_queue, at_most_once:)
-      unless lock_queue.is_a?(BackgroundQueue::LockQueue::RedisQueue) ||
-          lock_queue.is_a?(BackgroundQueue::LockQueue::RolloutRedisQueue)
-        raise ArgumentError, "an argument to #build_lock_queue_enumerator must be a LockQueue"
-      end
-      wrap(self, BackgroundQueue::LockQueueEnumerator.new(lock_queue, at_most_once: at_most_once).to_enum)
-    end
-
     # Builds Enumerator from Active Record Relation. Each Enumerator tick moves the cursor one row forward.
     #
     # +columns:+ argument is used to build the actual query for iteration. +columns+: defaults to primary key:
     #
     #   1) SELECT * FROM users ORDER BY id LIMIT 100
@@ -117,14 +108,24 @@
         **args
       ).batches
       wrap(self, enum)
     end
 
+    def build_throttle_enumerator(enum, throttle_on:, backoff:)
+      JobIteration::ThrottleEnumerator.new(
+        enum,
+        @job,
+        throttle_on: throttle_on,
+        backoff: backoff
+      ).to_enum
+    end
+
     alias_method :once, :build_once_enumerator
     alias_method :times, :build_times_enumerator
     alias_method :array, :build_array_enumerator
     alias_method :active_record_on_records, :build_active_record_enumerator_on_records
     alias_method :active_record_on_batches, :build_active_record_enumerator_on_batches
+    alias_method :throttle, :build_throttle_enumerator
 
     private
 
     def build_active_record_enumerator(scope, cursor:, **args)
       unless scope.is_a?(ActiveRecord::Relation)