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 CountsAxes::TooManyAxes
    render json: { error: 'too_many_axes', max_axes: CountsAxes::MAX_AXES }, status: :bad_request
  rescue CountsAxes::PrivateCondition
    render json: { error: 'private_condition' }, status: :forbidden
  rescue CountsTableQuery::TooManyColumns
    render json: { error: 'too_many_columns', max_columns: CountsTableQuery::MAX_COLUMNS },
        status: :unprocessable_entity
  rescue ActiveRecord::StatementTimeout
    render json: { error: 'counts_query_timeout' }, status: :service_unavailable
  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 CountsAxes::TooManyAxes
    render json: { error: 'too_many_axes', max_axes: CountsAxes::MAX_AXES }, status: :bad_request
  rescue CountsAxes::PrivateCondition
    render json: { error: 'private_condition' }, status: :forbidden
  rescue CountsTableQuery::TooManyColumns
    render json: { error: 'too_many_columns', max_columns: CountsTableQuery::MAX_COLUMNS },
        status: :unprocessable_entity
  rescue ActiveRecord::StatementTimeout
    render json: { error: 'counts_query_timeout' }, status: :service_unavailable
  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
    job = params[:job]
    authority = HistoryAccess.user
    first_check = true
    stream = CountsExportDownload::ResponseBody.new(file.fetch('path')) do |status|
      if status == 'active'
        if first_check
          first_check = false
          fresh = authority && User.find_by(id: authority.id)
          next false if authority && (!fresh || fresh.banned?)

          authority = fresh
          HistoryAccess.set(user: authority) do
            CountsExportDownload.find(arguments[:file_token], owner) &&
                CountsExportJobs.stream_active?(job, owner, **arguments)
          end
        else
          CountsExportJobs.stream_active?(job, owner, **arguments)
        end
      else
        HistoryAccess.set(user: authority) do
          CountsExportJobs.download_attempt(job, owner, **arguments, status:)
        end
      end
    end
    unless CountsExportJobs.download_attempt(job, owner, **arguments, status: 'streaming')
      stream.close
      return head :conflict
    end
    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'
    response.headers['Last-Modified'] = File.mtime(file.fetch('path')).httpdate
    response.headers['X-Accel-Buffering'] = 'no'
    self.response_body = stream
  rescue IOError, SystemCallError
    stream&.close
    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 CountsAxes::TooManyAxes
    render json: { error: 'too_many_axes', max_axes: CountsAxes::MAX_AXES }, status: :bad_request
  rescue CountsAxes::PrivateCondition
    render json: { error: 'private_condition' }, status: :forbidden
  rescue CountsTableQuery::TooManyColumns
    render json: { error: 'too_many_columns', max_columns: CountsTableQuery::MAX_COLUMNS },
        status: :unprocessable_entity
  rescue ActiveRecord::StatementTimeout
    render json: { error: 'counts_query_timeout' }, status: :service_unavailable
  rescue ArgumentError, TypeError, KeyError, JSON::ParserError
    head :bad_request
  end
end