btrc-hub/backend/app/services/counts_exporter.rb

require 'csv'

# One read view and a fixed point for the entire download; only one page is held in memory.
class CountsExporter
  EXTRA_KEYS = ['basis', 'format', 'resource', 'headers', 'children', 'order',
                'summary_columns', 'nested_json', 'child_sort', 'labels'].freeze
  attr_reader :format, :resource

  def initialize(params, language_code:)
    allowed = (CountsQuery::KEYS + CountsStateQuery::KEYS + EXTRA_KEYS).uniq
    raise ArgumentError if (params.keys - allowed).any?
    raise ArgumentError unless params.values.all? { _1.is_a?(String) }

    @params = params
    @language = language_code
    @format = params.fetch('format', 'json')
    @resource = params.fetch('resource', 'posts')
    @mode = params.fetch('mode', 'summary')
    @history = params.fetch('basis', 'state') == 'history'
    raise ArgumentError unless ['json', 'csv'].include?(format)
    raise ArgumentError unless ['posts', 'tags'].include?(resource)
    raise ArgumentError unless ['summary', 'details'].include?(@mode)
    raise ArgumentError unless ['state', 'history'].include?(params.fetch('basis', 'state'))

    @query_params = params.except(*EXTRA_KEYS, 'columns', 'page', 'limit')
        .merge('mode' => @mode, 'page' => '1', 'limit' => '100')
    point = CountsStateQuery.new(@query_params.slice('at').merge('resource' => resource),
                                 language_code: @language).at
    @query_params['at'] = point.iso8601(6)
    @query_params['axis'] ||= !@history && resource == 'posts' ? 'tags' : 'month'
    @projection = CountsExportProjection.new(params, resource:,
                                             summary: @mode == 'summary',
                                             tags: !@history && @query_params['axis'] == 'tags')
    if @history
      @query_params.delete('kind')
      history_query
    else
      @query_params.delete('latest')
      @state_query = CountsStateQuery.new(@query_params.merge('resource' => resource),
                                          language_code: @language).validate_query!
    end
  end

  def body
    Enumerator.new do |output|
      ApplicationRecord.transaction(isolation: :repeatable_read) do
        first_row = true
        started = false
        pages do |page|
          unless started
            meta = { resource:, mode: @mode, basis: @history ? 'history' : 'state',
                     at: @query_params['at'], boundary: 'inclusive',
                     filters: @query_params.except('sort', 'page', 'limit', 'mode'),
                     sort: export_sort,
                     child_sort: sort_metadata(@params.fetch('child_sort', '')),
                     columns: @projection.metadata,
                     children: @params['children']&.split(','),
                     field_order: @params.fetch('order', '').split(','), timezone: 'UTC',
                     count: @mode == 'summary' ? page[:group_count] : page[:count],
                     totals: page[:totals], coverage: page[:coverage] }
            output << (format == 'json' ? "{\"metadata\":#{ meta.to_json },\"data\":[" : "\uFEFF")
            if format == 'csv'
              output << CSV.generate_line(
                  @projection.headers.map { csv_value(_1) }, row_sep: "\r\n")
            end
            started = true
          end
          (page[:versions] || page[:state_rows] || page[:groups] || []).each do |row|
            record = @projection.project(row)
            if format == 'json'
              output << ',' unless first_row
              output << record.to_json
            else
              values = @projection.csv_record(record, row).map { csv_value(_1) }
              output << CSV.generate_line(values, row_sep: "\r\n")
            end
            first_row = false
          end
        end
        output << ']}' if format == 'json'
      end
    end
  end

  private

  def export_sort
    sort_metadata(@query_params.fetch('sort', default_sort))
  end

  def sort_metadata(value)
    value.split(',').map do |entry|
      key, direction = entry.split(':', 2)
      { key:, direction: }
    end
  end

  def default_sort
    return 'created_at:desc' if @mode == 'details'

    !@history && @query_params['axis'] == 'tags' ? 'count:desc' : 'group:asc'
  end

  def history_query
    scope = resource == 'posts' ? PostVersion.with_previous_version : TagVersion.all
    CountsQuery.new(scope, resource, @query_params)
  end

  def pages
    unless @history
      @state_query.each_page { yield _1 }
      return
    end
    query = history_query
    total = @mode == 'summary' ? nil : query.count
    loop do
      result = if @mode == 'summary'
                 query.summary(include_count: @query_params['page'] == '1')
               else
                 { count: total,
                   versions: CountsVersionRepr.many(query.details, resource:,
                                                   language_code: @language,
                                                   at: Time.iso8601(@query_params['at'])) }
               end
      total ||= result[:group_count]
      yield result
      break if @query_params['page'].to_i * 100 >= total

      @query_params['page'] = (@query_params['page'].to_i + 1).to_s
      query = history_query
    end
  end

  def csv_value(value)
    text = value.is_a?(Hash) || value.is_a?(Array) ? JSON.generate(value) : value.to_s
    text.match?(/\A\s*[=+@-]/) ? "'#{ text }" : text
  end
end