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