Add latest changes from gitlab-org/gitlab@12-10-stable-ee

1 files changed, 112 insertions, 0 deletions

diff --git a/lib/csv_builder.rb b/lib/csv_builder.rb
new file mode 100644
index 00000000000..7df4e3bf85d
--- /dev/null
+++ b/lib/csv_builder.rb
@@ -0,0 +1,112 @@
+# frozen_string_literal: true
+
+# Generates CSV when given a collection and a mapping.
+#
+# Example:
+#
+#     columns = {
+#       'Title' => 'title',
+#       'Comment' => 'comment',
+#       'Author' => -> (post) { post.author.full_name }
+#       'Created At (UTC)' => -> (post) { post.created_at&.strftime('%Y-%m-%d %H:%M:%S') }
+#     }
+#
+#     CsvBuilder.new(@posts, columns).render
+#
+class CsvBuilder
+  attr_reader :rows_written
+
+  #
+  # * +collection+ - The data collection to be used
+  # * +header_to_hash_value+ - A hash of 'Column Heading' => 'value_method'.
+  #
+  # The value method will be called once for each object in the collection, to
+  # determine the value for that row. It can either be the name of a method on
+  # the object, or a lamda to call passing in the object.
+  def initialize(collection, header_to_value_hash)
+    @header_to_value_hash = header_to_value_hash
+    @collection = collection
+    @truncated = false
+    @rows_written = 0
+  end
+
+  # Renders the csv to a string
+  def render(truncate_after_bytes = nil)
+    Tempfile.open(['csv']) do |tempfile|
+      csv = CSV.new(tempfile)
+
+      write_csv csv, until_condition: -> do
+        truncate_after_bytes && tempfile.size > truncate_after_bytes
+      end
+
+      if block_given?
+        yield tempfile
+      else
+        tempfile.rewind
+        tempfile.read
+      end
+    end
+  end
+
+  def truncated?
+    @truncated
+  end
+
+  def rows_expected
+    if truncated? || rows_written == 0
+      @collection.count
+    else
+      rows_written
+    end
+  end
+
+  def status
+    {
+      truncated: truncated?,
+      rows_written: rows_written,
+      rows_expected: rows_expected
+    }
+  end
+
+  private
+
+  def headers
+    @headers ||= @header_to_value_hash.keys
+  end
+
+  def attributes
+    @attributes ||= @header_to_value_hash.values
+  end
+
+  def row(object)
+    attributes.map do |attribute|
+      if attribute.respond_to?(:call)
+        excel_sanitize(attribute.call(object))
+      else
+        excel_sanitize(object.public_send(attribute)) # rubocop:disable GitlabSecurity/PublicSend
+      end
+    end
+  end
+
+  def write_csv(csv, until_condition:)
+    csv << headers
+
+    @collection.find_each do |object|
+      csv << row(object)
+
+      @rows_written += 1
+
+      if until_condition.call
+        @truncated = true
+        break
+      end
+    end
+  end
+
+  def excel_sanitize(line)
+    return if line.nil?
+
+    line = ["'", line].join if line =~ /^[=\+\-@;]/
+    line
+  end
+end