table_size.rake

if Module.const_defined?('ActionView')
  include ::ActionView::Helpers::NumberHelper

  namespace :db do
    namespace :tables do
      desc "how big is each table?"
      task :size => :environment do
        sql = "SELECT tablename FROM pg_tables WHERE schemaname = 'public' ORDER BY tablename;"
        tables = ActiveRecord::Base.connection.execute(sql)
        sizes = []
        tables.each do |table|
          name = table['tablename']
          sql = "SELECT pg_total_relation_size('#{name}');"
          res = ActiveRecord::Base.connection.execute(sql)
          sizes << Map.new({name: name, relation_size: res[0]['pg_total_relation_size'].to_i})
        end
      
        if ENV['SORT']
          key, direction = *ENV["SORT"].split(" ")
          sort = key if %w(name relation_size).include?(key)
          if sort
            sizes.sort_by!(&sort.to_sym)
            sizes.reverse! if direction == 'desc'
          else
            puts "Valid sort keys are 'name' and 'relation_size'.\nSorting by relation_size by default."
            sizes.sort_by!(&:relation_size)
          end
        else
          sizes.sort_by!(&:relation_size)
        end
      
        sizes.each{|table| puts "#{table.name} #{number_to_human_size(table.relation_size).rjust(40-table.name.length)}"}
        puts "Total in all tables: #{(number_to_human_size(sizes.map(&:relation_size).sum)).rjust(20)}"
      end
    end
  end
end