btrc-hub/backend/app/controllers/counts_exports_controller.rb

class CountsExportsController < ApplicationController
  wrap_parameters false

  def create_job
    locale = resolve_locale!
    payload = request.request_parameters.merge(request.query_parameters)
    session[:counts_export_owner] ||= SecureRandom.hex(32)
    job = CountsExportJobs.start(payload, session[:counts_export_owner], locale.language_code)
    response.headers['Cache-Control'] = 'private, no-store'
    render json: { job: }, status: :accepted
  rescue CountsExportJobs::WorkerUnavailable
    render json: { error: 'export_worker_unavailable' }, status: :service_unavailable
  rescue CountsExportDownload::Busy
    head :too_many_requests
  rescue ArgumentError, TypeError, KeyError, JSON::ParserError
    head :bad_request
  end

  def job_status
    job_state(cancel: false)
  end

  def cancel_job
    job_state(cancel: true)
  end

  def start_download
    result = CountsExportJobs.download_attempt(params[:token], session[:counts_export_owner])
    return head :not_found unless result

    response.headers['Cache-Control'] = 'private, no-store'
    render json: result
  end

  def prepare
    locale = resolve_locale!
    payload = request.request_parameters.merge(request.query_parameters)
    exporter = CountsExporter.new(payload, language_code: locale.language_code)
    session[:counts_export_owner] ||= SecureRandom.hex(32)
    token = CountsExportDownload.prepare(exporter, session[:counts_export_owner])
    response.headers['Cache-Control'] = 'private, no-store'
    render json: { token: }
  rescue CountsExportDownload::Busy
    head :too_many_requests
  rescue CountsExportDownload::TooLarge
    head 413
  rescue CountsExportDownload::TimedOut
    head :gateway_timeout
  rescue CountsExportDownload::StorageFull, Errno::ENOSPC
    head :insufficient_storage
  rescue ArgumentError, TypeError, KeyError, JSON::ParserError
    head :bad_request
  end

  def download
    if params[:job].present? || params[:attempt].present?
      return head :method_not_allowed unless request.get?

      return confirmed_download
    end
    file = CountsExportDownload.find(params[:token], session[:counts_export_owner])
    return head :not_found unless file

    response.headers['Cache-Control'] = 'private, no-store'
    send_file file.fetch('path'),
              filename: "counts-#{ file.fetch('resource') }.#{ file.fetch('format') }",
              type: file['format'] == 'json' ? 'application/json' : 'text/csv',
              disposition: 'attachment'
  end

  def confirmed_download
    owner = session[:counts_export_owner]
    arguments = { file_token: params[:token], attempt: params[:attempt] }
    return head :not_found unless arguments[:attempt].present? &&
        CountsExportJobs.download_attempt(params[:job], owner, **arguments)

    file = CountsExportDownload.find(params[:token], owner)
    unless file
      CountsExportJobs.download_attempt(params[:job], owner, **arguments, status: 'unavailable')
      return head :not_found
    end
    # Open before committing the response, but never buffer the entire export.
    stream = File.open(file.fetch('path'), 'rb')
    response.headers['Cache-Control'] = 'private, no-store'
    send_file_headers! filename: "counts-#{ file.fetch('resource') }.#{ file.fetch('format') }",
        type: file['format'] == 'json' ? 'application/json' : 'text/csv',
        disposition: 'attachment'
    job = params[:job]
    authority = HistoryAccess.user
    self.response_body = Enumerator.new do |output|
      begin
        first = true
        while (chunk = stream.read(64 * 1024))
          output << chunk
          if first
            HistoryAccess.set(user: authority) do
              CountsExportJobs.download_attempt(job, owner, **arguments, status: 'started')
            end
            first = false
          end
        end
      ensure
        stream.close
      end
    end
  rescue IOError, SystemCallError
    CountsExportJobs.download_attempt(params[:job], owner, **arguments, status: 'failed')
    head :service_unavailable
  end
  private :confirmed_download

  def job_state(cancel:)
    result = CountsExportJobs.status(params[:token], session[:counts_export_owner], cancel:)
    return head :not_found unless result

    response.headers['Cache-Control'] = 'private, no-store'
    render json: result
  end
  private :job_state

  def index
    locale = resolve_locale!
    payload = request.request_parameters.except('authenticity_token')
        .merge(request.query_parameters)
    exporter = CountsExporter.new(payload,
                                  language_code: locale.language_code)
    session[:counts_export_owner] ||= SecureRandom.hex(32)
    token = CountsExportDownload.prepare(exporter, session[:counts_export_owner])
    file = CountsExportDownload.find(token, session[:counts_export_owner])
    return head :not_found unless file

    send_file file.fetch('path'),
        filename: "counts-#{ exporter.resource }.#{ exporter.format }",
        type: exporter.format == 'json' ? 'application/json' : 'text/csv',
        disposition: 'attachment'
  rescue ArgumentError, TypeError, KeyError, JSON::ParserError
    head :bad_request
  end
end