diff --git a/lib/job-iteration/active_record_cursor.rb b/lib/job-iteration/active_record_cursor.rb index 10a8f4ef..240fad57 100644 --- a/lib/job-iteration/active_record_cursor.rb +++ b/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? diff --git a/lib/job-iteration/active_record_enumerator.rb b/lib/job-iteration/active_record_enumerator.rb index 363a4ecf..22bfa41b 100644 --- a/lib/job-iteration/active_record_enumerator.rb +++ b/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