Export AWS Athena usage history as CSV

export_athena_usage.rb

require 'csv'
require 'aws-sdk-athena'

def export_athena_usage(limit:, region:, filter_scanned: 10**9)
  athena = Aws::Athena::Client.new region: region
  next_token = nil
  fields = %w[query_execution_id query work_group]
  stats_fields = %w[data_scanned_in_bytes total_execution_time_in_millis]
  scanned = 0
  found = 0
  CSV.open("athena_results_#{Time.now.iso8601[0, 19].tr(':','-')}.csv", 'w') do |csv|
    csv << ['time', 'state', *fields, *stats_fields]
    loop do
      res = athena.list_query_executions(next_token: next_token)
      ids = res.query_execution_ids
      next_token = res.next_token
      data = athena.batch_get_query_execution(query_execution_ids: ids).query_executions
      last_time = nil
      data.each do |d|
        stats = d.statistics
        if stats&.data_scanned_in_bytes && stats.data_scanned_in_bytes > filter_scanned
          csv << [d.status.submission_date_time, d.status.state, *fields.map { |f| d[f] }, *stats_fields.map { |f| stats[f] }]
          found += 1
        end
        scanned += 1
        last_time = d.status.submission_date_time
      end

      print "\r#{scanned}, #{found}, #{last_time}   "
      break if scanned >= limit || !next_token
    end
  end
end