#!/bin/bash
#                                                          +----+----+----+----+
#                                                          |    |    |    |    |
# Author: Mark David Scott Cunningham                      | M  | D  | S  | C  |
#                                                          +----+----+----+----+
# Created: 2018-04-18
# Updated: 2018-04-21
#
# Purpose: Report the size of a mysql database or set of databases
#

#------------------------------------------------------------------------------#
# Variable initialization
diskskip=''
linecount=''
hostname=''
username=''
password=''
verbose=''
col=''
list=''
FMT="%11s %11s %11s %-s\n"
FMT2="%11s %-s\n"

#------------------------------------------------------------------------------#
# Utility function for output formatting
dash(){ for ((i=1;i<=$1;i++)); do printf "-"; done; }

#------------------------------------------------------------------------------#
# Help output function
usage(){
  echo "
Usage: $0 [options] [ALL|db1,db2,db3...]

 MySQL Connection Information

  -h | --hostname ... Specify hostname for MySQL connection

  -p | --password ... Specify password for MySQL connection

  -u | --username ... Specify username for MySQL connection

 Table Size Options

  -t | --tablesize .. Print table size information for given databases

  -n | --linecount .. Number of tables to list when printing table size
                      If not specified this defaults to 20 lines of output

  -d | --sortdata ... Sort tables by data size (default if not specified)

  -i | --sortindex .. Sort tables by index size

  -r | --sortrows ... Sort tables by number of rows

 General Flags

  -l | --list ....... Print list of database names

  -v | --verbose .... Verbose mode (Print debug info)

  --help ............ Print this message and quit

Examples:
Database size for all DBs
$0 ALL

Database size for list of DBs
$0 db1,db2,db3

Table size for list of DBs
$0 -t db1,db2,db3

Top 10 largest tables for list of DBs
$0 -t -n10 db1,db2,db3
";
}

#------------------------------------------------------------------------------#
# Database size output function

dbsize(){
  dbname=$1

  dbsize=$($connect -NBe "select SUM(Data_length + Index_length) FROM information_schema.TABLES WHERE TABLE_SCHEMA = \"${dbname}\";");
  if [[ $dbsize == 'NULL' ]]; then dbsize=0; fi

  # Removed as not addding meaningful information
  # free=$($connect -NBe "select SUM(Data_free) FROM information_schema.TABLES WHERE TABLE_SCHEMA = \"${dbname}\";");
  # if [[ $free == 'NULL' ]]; then free=0; fi
  # freem=$(echo "scale=2;${free}/1024/1024" | bc)

  ## Calculate Megabytes from bytes
  dbsizem=$(echo "scale=2;${dbsize}/1024/1024" | bc);

  if [[ $diskskip ]]; then
    printf "$FMT2" "${dbsizem}M" " ${dbname}"
  else
    ## Get size on disk and size of data in the tables
    disksize=$(du -sb ${datadir}${dbname} 2> /dev/null | awk '{print $1}');
    if [[ ! $disksize ]]; then disksize=0; fi

    ## Calculate Megabytes from bytes
    disksizem=$(echo "scale=2;${disksize}/1024/1024" | bc);
    frag=$(echo "scale=2;(${disksize}-${dbsize})/1024/1024" | bc)

    ## Print out the size on disk, size in tables, and the difference
    printf "$FMT" "${disksizem}M" "${dbsizem}M" "${frag}M" " ${dbname}";
  fi
}

#------------------------------------------------------------------------------#
# Table size output function

tablesize(){
dbname=$1
echo

# Output the sorted/formatted table sizes
printf "$FMT" "Rows" "Data_Size" "Index_Size" " Table_Name ($dbname)";
printf "$FMT" "$(dash 11)" "$(dash 11)" "$(dash 11)" "$(dash 44)"

$connect -NBe 'show table status' $dbname\
 | awk '{printf "%11s %10.2fM %10.2fM  %-s\n",$5,($7/1024/1024),($9/1024/1024),$1}'\
 | sort -rnk$col | head -n$top
}

#------------------------------------------------------------------------------#
# Argument Parsing
OPTIONS=$(getopt -o "dh:iln:p:rtu:v" --long "help,sortdata,hostname:,sortindex,list,linecount:,password:,sortrows,tablesize,username:,verbose" -- "$@") # Execute getopt
eval set -- "$OPTIONS" # Magic
while true; do # Evaluate the options for their options
case $1 in
  -d | --sortdata ) col='2' ;;
  -h | --hostname ) hostname="$2"; shift ;;
  -i | --sortindex ) col='3' ;;
  -l | --list ) list=1 ;;
  -n | --linecount ) linecount="$2"; shift ;;
  -p | --password ) password="$2"; shift ;;
  -r | --sortrows ) col='1' ;;
  -t | --tablesize ) table=1 ;;
  -u | --username ) username="$2"; shift ;;
  -v | --verbose ) verbose=1 ;;
  -- ) shift; break ;; # More Magic
  --help ) usage; exit 0 ;; #print-help-info
esac;
shift;
done

#------------------------------------------------------------------------------#
# Check for existing mysql credential config files
mysql=$(which mysql)

# Make sure user is not trying to override an existing credentials file
if [[ -x $mysql && -f ~/.my.cnf && ! $hostname && ! $username && ! $password ]]; then
  connect="$mysql";
  if [[ $UID != 0 ]]; then
    echo -e "\nConfiguration .my.cnf found, but not running as root. \nYou will likely have access to a limited number of databases.\n";
  fi

elif [[ -x $mysql ]]; then

  # If no .my.cnf found and no credentials provided, print warning and then prompt for credentials.
  if [[ ! -f ~/.my.cnf ]]; then
    echo -e "\nNo .my.cnf credentials file found in home directory\nPlease enter your credentials below.\n"
  fi

  # Write out temporary .my.cnf (avoids warnings about using password at the command line being insecure)
  tmpfile="~/mysqlsize.tmp"
  if [[ ! $hostname ]]; then read -p "Hostname: " hostname; fi
  if [[ ! $username ]]; then read -p "Username: " username; fi
  if [[ ! $password ]]; then read -sp "Password: " password; fi
    echo "[mysql]" > $tmpfile
    echo "host=$hostname" >> $tmpfile
    echo "user=$username" >> $tmpfile
    echo "password=$password" >> $tmpfile
  connect="$mysql --defaults-file=$tmpfile"
  echo
else
  echo -e "\nIt does not appear that the mysql client is installed.\n"; exit 1;
fi

#------------------------------------------------------------------------------#
## Get list of databases (if ALL, then generate the total list, otherwise use given list or prompt for input)
dbs="$@"
if [[ $dbs == ALL ]]; then
  databases=$($connect -NBe 'show databases' | grep -Ev '^performance_schema$|^information_schema$')

# Cleanup comma separated list if provided
elif [[ $dbs ]]; then
  databases=$(echo $dbs | sed 's/,/ /g')

# Prompt for missing input if running the tablesize function
elif [[ $table && ! $dbs ]]; then
  read -p "Database: " dbs
  if [[ ! $dbs ]]; then
    usage; exit 2; #no-db-specified
  else
    databases=$(echo $dbs | sed 's/,/ /g')
  fi

# If not using the list utilty then assume error and print usage
elif [[ ! $list ]]; then
  usage; exit 1; #no-input
fi

#------------------------------------------------------------------------------#
## Call apropriate function based on input

# If list then just list out database names
if [[ $list ]]; then
  echo
  $connect -NBe 'show databases' | grep -Ev '^performance_schema$|^information_schema$'
  echo

# If table and not list then run the tablesize function
elif [[ $table ]]; then
  # List out database table sizes
  if [[ ! $col ]]; then col='2'; fi
  if [[ $linecount =~ [0-9]{1,} ]]; then top="$linecount"; else top='20'; fi
  for x in $databases; do
    tablesize $x
  done; echo

# If not list and not table then run dbsize function
elif [[ ! $table ]]; then
  ## List out database sizes

  # Check what hostname MySQL is reporting
  dbhost=$($connect -NBe 'select @@hostname')

  if [[ $HOSTNAME != $dbhost ]]; then
    echo -e "\nIt looks like you're connecting to a remote DB server. \nSize on disk and size comparison will be skipped.\n"
    diskskip=1
    printf "$FMT2" " MySQLData" " Database"
    printf "$FMT2" "$(dash 11)" "$(dash 68)"

  else
    # Check where the datadir is for using du to check size on disk
    datadir=$($connect -NBe 'select @@datadir')
    printf "\n$FMT" " On-Disk" " MySQLData" " Size-Diff" " Database"
    printf "$FMT" "$(dash 11)" "$(dash 11)" "$(dash 11)" "$(dash 44)"
  fi

  for x in $databases; do
    dbsize $x
  done; echo

fi

#------------------------------------------------------------------------------#
# Verbose output (Debugging)
if [[ $verbose ]]; then
echo "Debugging output"
echo "$(dash 40)"
echo "col = $col"
echo "dbs = $dbs"
echo "list = $list"
echo "table = $table"
echo "localhost = $HOSTNAME"
echo "dbhost = $dbhost"
echo "hostname = $hostname"
echo "username = $username"
echo "password = $password"
echo "tmpfile = $tmpfile"
fi

#------------------------------------------------------------------------------#
# Cleanup
if [[ -f $tmpfile ]]; then rm -f $tmpfile; fi