1# frozen_string_literal: true
2
3# Generates CSV when given a collection and a mapping.
4#
5# Example:
6#
7#     columns = {
8#       'Title' => 'title',
9#       'Comment' => 'comment',
10#       'Author' => -> (post) { post.author.full_name }
11#       'Created At (UTC)' => -> (post) { post.created_at&.strftime('%Y-%m-%d %H:%M:%S') }
12#     }
13#
14#     CsvBuilder.new(@posts, columns).render
15#
16class CsvBuilder
17  DEFAULT_ORDER_BY = 'id'
18  DEFAULT_BATCH_SIZE = 1000
19  PREFIX_REGEX = /\A[=\+\-@;]/.freeze
20
21  attr_reader :rows_written
22
23  #
24  # * +collection+ - The data collection to be used
25  # * +header_to_hash_value+ - A hash of 'Column Heading' => 'value_method'.
26  #
27  # The value method will be called once for each object in the collection, to
28  # determine the value for that row. It can either be the name of a method on
29  # the object, or a lamda to call passing in the object.
30  def initialize(collection, header_to_value_hash)
31    @header_to_value_hash = header_to_value_hash
32    @collection = collection
33    @truncated = false
34    @rows_written = 0
35  end
36
37  # Renders the csv to a string
38  def render(truncate_after_bytes = nil)
39    Tempfile.open(['csv']) do |tempfile|
40      csv = CSV.new(tempfile)
41
42      write_csv csv, until_condition: -> do
43        truncate_after_bytes && tempfile.size > truncate_after_bytes
44      end
45
46      if block_given?
47        yield tempfile
48      else
49        tempfile.rewind
50        tempfile.read
51      end
52    end
53  end
54
55  def truncated?
56    @truncated
57  end
58
59  def rows_expected
60    if truncated? || rows_written == 0
61      @collection.count
62    else
63      rows_written
64    end
65  end
66
67  def status
68    {
69      truncated: truncated?,
70      rows_written: rows_written,
71      rows_expected: rows_expected
72    }
73  end
74
75  protected
76
77  def each(&block)
78    @collection.find_each(&block) # rubocop: disable CodeReuse/ActiveRecord
79  end
80
81  private
82
83  def headers
84    @headers ||= @header_to_value_hash.keys
85  end
86
87  def attributes
88    @attributes ||= @header_to_value_hash.values
89  end
90
91  def row(object)
92    attributes.map do |attribute|
93      if attribute.respond_to?(:call)
94        excel_sanitize(attribute.call(object))
95      else
96        excel_sanitize(object.public_send(attribute)) # rubocop:disable GitlabSecurity/PublicSend
97      end
98    end
99  end
100
101  def write_csv(csv, until_condition:)
102    csv << headers
103
104    each do |object|
105      csv << row(object)
106
107      @rows_written += 1
108
109      if until_condition.call
110        @truncated = true
111        break
112      end
113    end
114  end
115
116  def excel_sanitize(line)
117    return if line.nil?
118    return line unless line.is_a?(String) && line.match?(PREFIX_REGEX)
119
120    ["'", line].join
121  end
122end
123