Expand db:stats rake task to include table + index size

lib/tasks/db.rake

@@ -7,27 +7,40 @@ require_relative "support/db_backup_utils"
 namespace :db do
   desc "Database Stats"
   task stats: [:environment] do
-    heading_fmt = "%-30s %8s"
-    row_fmt = "%-30s %8d"
-    puts format(heading_fmt, "Table", "Count")
-    puts format(heading_fmt, "-" * 30, "-" * 8)
+    heading_fmt = "%-30s %8s %11s"
+    row_fmt     = "%-30s %8d %8.1f MB"
+
+    puts format(heading_fmt, "Table", "Count", "Size")
+    puts format(heading_fmt, "-" * 30, "-" * 8, "-" * 11)
+
     db_name = ActiveRecord::Base.connection.current_database
-    ActiveRecord::Base.connection.tables.each do |t|
-      sql = if ENV["EXACT_COUNT"]
-              "select count(*) from #{t}"
-            else
-              "SELECT TABLE_ROWS FROM information_schema.TABLES WHERE TABLE_SCHEMA = '#{db_name}' AND TABLE_NAME = '#{t}'"
-            end
-      res = ActiveRecord::Base.connection.exec_query(sql)
-      row = res.rows.first
-      if row[0] < 100 && ENV["EXACT_COUNT"].nil?
-        # fast enough to get an exact count
-        sql = "select count(*) from #{t}"
-        res = ActiveRecord::Base.connection.exec_query(sql)
-        row = res.rows.first
-      end
-      puts format(row_fmt, t, row[0])
+    conn    = ActiveRecord::Base.connection
+
+    sql = <<~SQL.squish
+      SELECT TABLE_NAME, TABLE_ROWS, (DATA_LENGTH + INDEX_LENGTH) / 1024.0 / 1024.0
+      FROM information_schema.TABLES
+      WHERE TABLE_SCHEMA = '#{db_name}'
+    SQL
+    table_stats = conn.exec_query(sql).rows.to_h { |name, rows, mb| [name, [rows, mb]] }
+
+    total_mb = 0.0
+    total_count = 0
+    conn.tables.each do |t|
+      approx_count, mb = table_stats[t]
+
+      count = if ENV["EXACT_COUNT"] || approx_count.to_i < 100
+                conn.exec_query("SELECT COUNT(*) FROM #{t}").rows.first[0]
+              else
+                approx_count
+              end
+
+      total_mb += mb
+      total_count += count
+      puts format(row_fmt, t, count, mb.to_f)
     end
+    puts format(heading_fmt, "", "=" * 8, "=" * 11)
+    puts format(row_fmt, "TOTAL", total_count, total_mb.to_f)
+
     puts "",
          "Status as of #{Time.new.utc}",
          ENV["EXACT_COUNT"] ? "(exact count)" : "(approximate count - set EXACT_COUNT=1 to get exact)"