Class: SidekiqIteration::ActiveRecordEnumerator

Inherits:
Object
  • Object
show all
Defined in:
lib/sidekiq_iteration/active_record_enumerator.rb

Constant Summary collapse

SQL_DATETIME_WITH_NSEC = 
"%Y-%m-%d %H:%M:%S.%6N"

Instance Method Summary collapse

Constructor Details

#initialize(relation, columns: nil, batch_size: 100, order: :asc, cursor: nil) ⇒ ActiveRecordEnumerator

Returns a new instance of ActiveRecordEnumerator.



8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
 
# File 'lib/sidekiq_iteration/active_record_enumerator.rb', line 8

def initialize(relation, columns: nil, batch_size: 100, order: :asc, cursor: nil)
  unless relation.is_a?(ActiveRecord::Relation)
    raise ArgumentError, "relation must be an ActiveRecord::Relation"
  end

  if relation.arel.orders.present? || relation.arel.taken.present?
    raise ArgumentError,
      "The relation cannot use ORDER BY or LIMIT due to the way how iteration with a cursor is designed. " \
      "You can use other ways to limit the number of rows, e.g. a WHERE condition on the primary key column."
  end

  @relation = relation
  @primary_key = relation.primary_key
  columns = Array(columns || @primary_key).map(&:to_s)

  if (Array(order) - [:asc, :desc]).any?
    raise ArgumentError, ":order must be :asc or :desc or an array consisting of :asc or :desc, got #{order.inspect}"
  end

  if order.is_a?(Array) && order.size != columns.size
    raise ArgumentError, ":order must include a direction for each batching column"
  end

  @primary_key_index = primary_key_index(columns, relation)
  if @primary_key_index.nil? || (composite_primary_key? && @primary_key_index.any?(nil))
    raise ArgumentError, ":columns must include a primary key columns"
  end

  @batch_size = batch_size
  @order = batch_order(columns, order)
  @cursor = Array(cursor)

  if @cursor.present? && @cursor.size != columns.size
    raise ArgumentError, ":cursor must include values for all the columns from :columns"
  end

  if columns.any?(/\W/)
    arel_columns = columns.map.with_index do |column, i|
      arel_column(column).as("cursor_column_#{i + 1}")
    end
    @cursor_columns = arel_columns.map { |column| column.right.to_s }

    relation =
      if relation.select_values.empty?
        relation.select(@relation.arel_table[Arel.star], arel_columns)
      else
        relation.select(arel_columns)
      end
  else
    @cursor_columns = columns
  end

  @columns = columns
  ordering = @columns.zip(@order).to_h
  @base_relation = relation.reorder(ordering)
  @iteration_count = 0
end

Instance Method Details

#batches ⇒ Object 



77
78
79
80
81
82
83
84
 
# File 'lib/sidekiq_iteration/active_record_enumerator.rb', line 77

def batches
  Enumerator.new(-> { records_size }) do |yielder|
    while (batch = next_batch(load: true))
      increment_iteration
      yielder.yield(batch, cursor_value(batch.last))
    end
  end
end

#records ⇒ Object 



66
67
68
69
70
71
72
73
74
75
 
# File 'lib/sidekiq_iteration/active_record_enumerator.rb', line 66

def records
  Enumerator.new(-> { records_size }) do |yielder|
    batches.each do |batch, _| # rubocop:disable Style/HashEachMethods
      batch.each do |record|
        increment_iteration
        yielder.yield(record, cursor_value(record))
      end
    end
  end
end

#relations ⇒ Object 



86
87
88
89
90
91
92
93
 
# File 'lib/sidekiq_iteration/active_record_enumerator.rb', line 86

def relations
  Enumerator.new(-> { relations_size }) do |yielder|
    while (batch = next_batch(load: false))
      increment_iteration
      yielder.yield(batch, unwrap_array(@cursor))
    end
  end
end