Add database role support for enumerator in batches

Author: etravassos

lib/job-iteration/active_record_cursor.rb

@@ -64,7 +64,7 @@ def update_from_record(record)
       end
     end
 
-    def next_batch(batch_size)
+    def next_batch(batch_size, database_role: nil)
       return if @reached_end
 
       relation = @base_relation.limit(batch_size)
@@ -74,7 +74,13 @@ def next_batch(batch_size)
       end
 
       records = relation.uncached do
-        relation.to_a
+        if database_role.present?
+          ActiveRecord::Base.connected_to(role: database_role) do
+            relation.to_a
+          end
+        else
+          relation.to_a
+        end
       end
 
       update_from_record(records.last) unless records.empty?

lib/job-iteration/active_record_enumerator.rb

@@ -7,9 +7,10 @@ module JobIteration
   class ActiveRecordEnumerator
     SQL_DATETIME_WITH_NSEC = "%Y-%m-%d %H:%M:%S.%N"
 
-    def initialize(relation, columns: nil, batch_size: 100, cursor: nil)
+    def initialize(relation, columns: nil, batch_size: 100, cursor: nil, database_role: nil)
       @relation = relation
       @batch_size = batch_size
+      @database_role = database_role
       @columns = if columns
         Array(columns)
       else
@@ -31,7 +32,7 @@ def records
     def batches
       cursor = finder_cursor
       Enumerator.new(method(:size)) do |yielder|
-        while (records = cursor.next_batch(@batch_size))
+        while (records = cursor.next_batch(@batch_size, database_role: @database_role))
           yielder.yield(records, cursor_value(records.last)) if records.any?
         end
       end