Rails CSV report streaming

app.rb

# From https://shift.infinite.red/fast-csv-report-generation-with-postgres-in-rails-d444d9b915ab
# Storing here in case the original post ever goes down.

# concern
module DatabaseQueryStreaming

  def stream_query_rows(sql_query, options="WITH CSV HEADER")
    conn = ActiveRecord::Base.connection.raw_connection
    conn.copy_data "COPY (#{sql_query}) TO STDOUT #{options};" do
      while row = conn.get_copy_data
        yield row
      end
    end
  end

end

# model
class ApplicationRecord < ActiveRecord::Base
  extend DatabaseQueryStreaming
  
  self.abstract_class = true
end

# controller
class SalesReportsController < AdminController
  include ActionController::Live
  
  def export
    respond_to do |format|
      format.csv { stream_csv_report }
    end
  end
  
  private
  
  def stream_csv_report
    query = SalesReports.some_complicated_query(query_params).to_sql
    
    query_options = "WITH CSV HEADER"
    # Note that if you have a custom select in your query
    # you may need to generate the header yourself. e.g.
    # => stream.write "Created Date,Ordered Date,Price,# of Items"
    # => query_options = "WITH CSV" # note the lack of 'HEADER'
    stream_file("sales_reports", "csv") do |stream|
      Sale.stream_query_rows(query, query_options) do |row_from_db|
        # row_from_db will be ordered according to the select
        # e.g.
        # => Given: "COPY (SELECT customer.name as name, ordered_at::date, created_at FROM sales INNER JOIN ...) TO STDOUT WITH CSV"
        # => row_from_db will look like 'John Doe,2017-06-16,2016-06-12 13:27:58.580456'
        stream.write row_from_db
      end
    end
  end
  
  def stream_file(filename, extension)
    response.headers["Content-Type"] = "application/octet-stream"
    response.headers["Content-Disposition"] = "attachment; filename=#{filename}.#{extension}"

    yield response.stream
  ensure
    response.stream.close
  end

end

More examples:

• https://medium.com/table-xi/stream-csv-files-in-rails-because-you-can-46c212159ab7
• https://github.com/codeforamerica/drai/pull/141/files