#!/bin/bash
# This script will assist with configuring ProxySQL in combination with Percona XtraDB cluster.
# Version 1.0
###############################################################################################


# Make sure only root can run this script
if [ $(id -u) -ne 0 ]; then
  echo "ERROR: This script must be run as root!" 1>&2
  exit
fi

# Dispay script usage details
usage () {
  echo "Usage: [ options ]"
  echo "Options:"
    echo " --config-file                   Override login credentials from command line and read login credentials from config file."
    echo " --proxysql-user=user_name       User to use when connecting to the ProxySQL service"
    echo " --proxysql-password[=password]  Password to use when connecting to the ProxySQL service"
    echo " --proxysql-port=port_num        Port to use when connecting to the ProxySQL service"
    echo " --proxysql-host=host_name       Hostname to use when connecting to the ProxySQL service"
    echo " --cluster-user=user_name        User to use when connecting to the Percona XtraDB Cluster node"
    echo " --cluster-password[=password]   Password to use when connecting to the Percona XtraDB Cluster node"
    echo " --cluster-port=port_num         Port to use when connecting to the Percona XtraDB Cluster node"
    echo " --cluster-host=host_name        Hostname to use when connecting to the Percona XtraDB Cluster node"
    echo " --monitor-user=user_name        User to use for monitoring Percona XtraDB Cluster nodes through ProxySQL"
    echo " --monitor-password[=password]   Password to for monitoring Percona XtraDB Cluster nodes through ProxySQL"
    echo " --pxc-app-user=user_name        Application user to use when connecting to the Percona XtraDB Cluster node"
    echo " --pxc-app-password[=password]   Application password to use when connecting to the Percona XtraDB Cluster node"
    echo " --enable                        Auto-configure Percona XtraDB Cluster nodes into ProxySQL"
    echo " --disable                       Remove Percona XtraDB Cluster configurations from ProxySQL"
    echo " --galera-check-interval         Interval for monitoring proxysql_galera_checker script(in milliseconds)"
    echo " --mode                          ProxySQL read/write configuration mode, currently it only support 'loadbal' mode" 
    echo " --adduser                       Add Percona XtraDB Cluster application user to ProxySQL database"
}

# Check if we have a functional getopt(1)
if ! getopt --test
  then
  go_out="$(getopt --options= --longoptions=proxysql-user:,proxysql-password::,proxysql-host:,proxysql-port:,cluster-user:,cluster-password::,cluster-host:,cluster-port:,monitor-user:,monitor-password:,pxc-app-user:,pxc-app-password:,galera-check-interval:,mode:,config-file:,enable,disable,adduser,help \
  --name="$(basename "$0")" -- "$@")"
  test $? -eq 0 || exit 1
  eval set -- $go_out
fi

for arg
do
  case "$arg" in
    -- ) shift; break;;
    --proxysql-user )
    PROXYSQL_USERNAME="$2"
    shift 2
    ;;
    --proxysql-password )
    case "$2" in
      "")
      read -r -s -p  "Enter ProxySQL password:" INPUT_PASS
      if [ -z "$INPUT_PASS" ]; then
        PROXYSQL_PASSWORD=""
	printf "\nContinuing without ProxySQL password...\n";
      else
        PROXYSQL_PASSWORD="$INPUT_PASS"
      fi
      printf "\n"
      ;;
      *)
      PROXYSQL_PASSWORD="$2"
      ;;
    esac
    shift 2
    ;;
    --proxysql-host )
    PROXYSQL_HOSTNAME="$2"
    shift 2
    ;;
    --proxysql-port )
    PROXYSQL_PORT="$2"
    shift 2
    ;;
    --cluster-user )
    CLUSTER_USERNAME="$2"
    shift 2
    ;;
    --cluster-password )
    case "$2" in
      "")
      read -r -s -p  "Enter Cluster password:" INPUT_PASS
      if [ -z "$INPUT_PASS" ]; then
        CLUSTER_PASSWORD=""
	printf "\nContinuing without Cluster password...\n";
      else
        CLUSTER_PASSWORD="$INPUT_PASS"
      fi
      printf "\n"
      ;;
      *)
      CLUSTER_PASSWORD="$2"
      ;;
    esac
    shift 2
    ;;
    --cluster-host )
    CLUSTER_HOSTNAME="$2"
    shift 2
    ;;
    --cluster-port )
    CLUSTER_PORT="$2"
    shift 2
    ;;
    --monitor-user )
    MONITOR_USERNAME="$2"
    shift 2
    ;;
    --monitor-password )
    MONITOR_PASSWORD="$2"
    shift 2
    ;;
    --pxc-app-user )
    PXC_APP_USERNAME="$2"
    shift 2
    ;;
    --pxc-app-password )
    PXC_APP_PASSWORD="$2"
    shift 2
    ;;
    --config-file )
      config_file="$2"
      shift 2
      if [ -z "${config_file}" ]; then
        echo "ERROR: You have not given config file location. Terminating."
        exit 1
      fi
      if [ -e "${config_file}" ]; then
        source "${config_file}"
      else
        echo "ERROR: ${config_file} does not exist. Terminating."
        exit 1
      fi
      ;;
    -e | --enable )
    shift
    enable=1
    ;;
    --adduser )
    shift
    adduser=1
    ;;
    -d | --disable )
    shift
    disable=1
    ;;
    --galera-check-interval )
    GALERA_CHECK_INTERVAL="$2"
    shift 2
    ;;
    --mode )
    RW_MODE="$2"
    shift 2
    ;;
    --help )
    usage
    exit 0
    ;;
  esac
done

if [[ ! -e `which mysql 2> /dev/null` ]] ;then
  echo "mysql client is not found, please install the mysql client package" 
  exit 1
fi

# Check the options gathered from the command line
if [ -z "$PROXYSQL_USERNAME" ];then
  echo "The ProxySQL username is required!"
  usage
  exit 1
fi

if [[ -z "$PROXYSQL_HOSTNAME" ]]; then
  PROXYSQL_HOSTNAME="-h127.0.0.1"
fi

if [[ -z "$PROXYSQL_PORT" ]]; then
  PROXYSQL_PORT="-P6032"
fi

if [ -z "$CLUSTER_USERNAME" ];then
  echo "The Percona XtraDB Cluster username is required!"
  usage
  exit 1
fi

if [[ -z "$CLUSTER_HOSTNAME" ]]; then
  CLUSTER_HOSTNAME="localhost"
fi

if [[ -z "$CLUSTER_PORT" ]]; then
  CLUSTER_PORT="3306"
fi
if [[ -z "$CLUSTER_PASSWORD" ]]; then
  CLUSTER_PASSWORD=""
else
  CLUSTER_PASSWORD="-p$CLUSTER_PASSWORD"
fi

if [[ -z "$GALERA_CHECK_INTERVAL" ]]; then
  GALERA_CHECK_INTERVAL=3000
fi

if [[ -z "$RW_MODE" ]]; then
  RW_MODE="loadbal"
elif [[ "$RW_MODE" != "loadbal" ]];then
  echo "Alert !For now, the only supported mode is 'loadbal'  which will be the default for a load balanced set of evenly weighted read/write nodes."
  echo "Please run with default mode."
  usage
  exit 1
fi

PIDFILE=/tmp/pxc-proxysql-monitor.pid

proxysql_connection_check(){
  PROXYSQL_DSN="-u$PROXYSQL_USERNAME -p$PROXYSQL_PASSWORD -h$PROXYSQL_HOSTNAME -P$PROXYSQL_PORT --protocol=tcp"
  if ! mysql $PROXYSQL_DSN -e"show tables" >/dev/null 2>/dev/null; then 
    echo "Please enter valid ProxySQL connection parameters! Terminating.."
    exit 1
  fi
}

pxc_connection_check(){
  DSN=" -u$CLUSTER_USERNAME $CLUSTER_PASSWORD -h$CLUSTER_HOSTNAME -P$CLUSTER_PORT --protocol=tcp"
  if ! mysql $DSN -e"select @@port" >/dev/null 2>/dev/null; then 
    echo "Please enter valid PXC connection parameters! Terminating.."
    exit 1
  fi
}
check_cmd(){
  MPID=$1
  ERROR_MSG=$2
  if [ ${MPID} -ne 0 ]; then echo -e "\nERROR: $ERROR_MSG. Terminating!"; exit 1; fi
}

check_proxysql(){
  if ! pidof proxysql >/dev/null ; then
    echo "ProxySQL is not running, please check the error log at /var/lib/proxysql/proxysql.log"
    exit 1
  fi
}

proxysql_exec() {
  query=$1
  result=$(
    printf "[client]\nuser=${PROXYSQL_USERNAME}\npassword=${PROXYSQL_PASSWORD}\nhost=${PROXYSQL_HOSTNAME}\nport=${PROXYSQL_PORT}\n" | \
      mysql --defaults-file=/dev/stdin --protocol=tcp -e "${query}" 2>/dev/null
  )
}

# Auto configure Percona XtraDB Cluster nodes into ProxySQL
enable_proxysql(){
  # Checking proxysql binary location
  if [[ ! -e $(which proxysql 2> /dev/null) ]]; then
    echo "The proxysql binary was not found, please install the ProxySQL package"  
    exit 1
  elif [[ ! -e $(which proxysql_galera_checker 2> /dev/null) ]] ;then
    echo "The proxysql_galera_checker binary was not found, please check ProxySQL package installation"  
    exit 1
  elif [[ ! -e $(which proxysql_node_monitor 2> /dev/null) ]]; then
    echo "The proxysql_node_monitor binary was not found, please check ProxySQL package installation"  
    exit 1
  else
    PROXYSQL=$(which proxysql)
    PROXYSQL_GALERA_CHECK=$(which proxysql_galera_checker)
    PROXYSQL_NODE_MONITOR=$(which proxysql_node_monitor)
  fi
  # Check for existing proxysql process
  if ! pidof proxysql >/dev/null ; then
    echo "ProxySQL is not running; please start proxysql service"
    exit 1
  fi
  proxysql_connection_check

  #modifying proxysql-admin.cnf file with command line proxysql user credentials if you dont use --config-file option.
  if [ -z "${config_file}" ]; then 
    sed -i "s|[ \t]*PROXYSQL_USERNAME[ \t]*=.*$| PROXYSQL_USERNAME=\"${PROXYSQL_USERNAME}\"|" /etc/proxysql-admin.cnf
    sed -i "s|[ \t]*PROXYSQL_PASSWORD[ \t]*=.*$| PROXYSQL_PASSWORD=\"${PROXYSQL_PASSWORD}\"|" /etc/proxysql-admin.cnf
    sed -i "s|[ \t]*PROXYSQL_HOSTNAME[ \t]*=.*$| PROXYSQL_HOSTNAME=\"${PROXYSQL_HOSTNAME}\"|" /etc/proxysql-admin.cnf
    sed -i "s|[ \t]*PROXYSQL_PORT[ \t]*=.*$| PROXYSQL_PORT=\"${PROXYSQL_PORT}\"|" /etc/proxysql-admin.cnf
  fi

  pxc_connection_check
  echo -e "\nConfiguring ProxySQL monitoring user.."
  if [[ -z $MONITOR_USERNAME ]]; then
    echo -n "Enter ProxySQL monitoring username: "
    read -r MONITOR_USERNAME
      while [[ -z "$MONITOR_USERNAME" ]]
      do
        echo -n "No input entered, Enter ProxySQL monitoring username: "
      read -r MONITOR_USERNAME
    done
  else
    echo -e "ProxySQL monitoring username as per command line is '$MONITOR_USERNAME'"
  fi
  if [[ -z $MONITOR_PASSWORD ]]; then
    read -r -s -p  "Enter ProxySQL monitoring password: " MONITOR_PASSWORD
    while [[ -z "$MONITOR_PASSWORD" ]]
    do
      read -r -s -p  "No input entered, Enter ProxySQL monitoring password: " MONITOR_PASSWORD
    done
  fi

  check_user=`mysql  $DSN -Bse"SELECT user,host FROM mysql.user where user='$MONITOR_USERNAME' and host='%';" 2>/dev/null`
  if [[ -z "$check_user" ]]; then
    mysql  $DSN -e "CREATE USER $MONITOR_USERNAME@'%' IDENTIFIED BY '$MONITOR_PASSWORD';" 2>/dev/null
    check_cmd $?  "Cannot create the ProxySQL monitoring user"
    proxysql_exec "update global_variables set variable_value='$MONITOR_USERNAME' where variable_name='mysql-monitor_username'; update global_variables set variable_value='$MONITOR_PASSWORD' where variable_name='mysql-monitor_password'; "
    check_cmd $?  "Cannot set the mysql-monitor variables in ProxySQL"
    proxysql_exec "LOAD MYSQL VARIABLES TO RUNTIME;SAVE MYSQL VARIABLES TO DISK;"
    echo -e "\n\nUser '$MONITOR_USERNAME'@'%' has been added with USAGE privilege\n"
  else
    echo -e "\n"
    read -p "Monitoring user is already present in Percona XtraDB Cluster. Would you like to proceed with existing username and password [y/n] ? " check_param
    case $check_param in
      y|Y)
        read -r -s -p  "Please enter the password you have assigned to monitoring user '$MONITOR_USERNAME': " MONITOR_PASSWORD
        proxysql_exec "update global_variables set variable_value='$MONITOR_USERNAME' where variable_name='mysql-monitor_username'; update global_variables set variable_value='$MONITOR_PASSWORD' where variable_name='mysql-monitor_password'; "
        check_cmd $?  "Cannot set the mysql-monitor variables in ProxySQL"
        proxysql_exec "LOAD MYSQL VARIABLES TO RUNTIME;SAVE MYSQL VARIABLES TO DISK;"
      ;;
      n|N)
        proxysql_exec "update global_variables set variable_value='$MONITOR_USERNAME' where variable_name='mysql-monitor_username'; update global_variables set variable_value='$MONITOR_PASSWORD' where variable_name='mysql-monitor_password'; "
        check_cmd $?  "Cannot set the mysql-monitor variables in ProxySQL"
        proxysql_exec "LOAD MYSQL VARIABLES TO RUNTIME;SAVE MYSQL VARIABLES TO DISK;"
        echo -e "\n\nUser '$MONITOR_USERNAME'@'%' has been updated in ProxySQL database. Please make sure the credentials are same in Percona XtraDB Cluster\n"
      ;;
      *)
        echo "Please type [y/n]! Terminating.."
        exit 1
      ;;
    esac
  fi

  # Adding Percona XtraDB Cluster nodes to ProxySQL
  echo -e "\nAdding the Percona XtraDB Cluster server nodes to ProxySQL"
  proxysql_exec "DELETE FROM mysql_servers WHERE hostgroup_id=10"
  wsrep_address=(`mysql  $DSN -Bse "show status like 'wsrep_incoming_addresses'" 2>/dev/null | awk '{print $2}' | sed 's|,| |g'`)
  for i in "${wsrep_address[@]}"; do	
    ws_ip=$(echo $i | cut -d':' -f1)
    ws_port=$(echo $i | cut -d':' -f2)
    proxysql_exec "INSERT INTO mysql_servers (hostname,hostgroup_id,port,weight) VALUES ('$ws_ip',10,$ws_port,1000);"
    check_cmd $? "Failed to add the Percona XtraDB Cluster server node $ws_ip:$ws_port"
  done
  proxysql_exec "LOAD MYSQL SERVERS TO RUNTIME; SAVE MYSQL SERVERS TO DISK;"

  # Adding Percona XtraDB Cluster monitoring scripts

  # Adding proxysql galera check scheduler
  proxysql_exec "DELETE FROM SCHEDULER WHERE ID=10;"
  check_cmd $?
  proxysql_exec "INSERT  INTO SCHEDULER (id,active,interval_ms,filename,arg1,arg2,arg3,arg4,arg5) VALUES (10,1,$GALERA_CHECK_INTERVAL,'$PROXYSQL_GALERA_CHECK',10,10,${#wsrep_address[@]},1,'/var/lib/proxysql/proxysql_galera_check.log');"
  check_cmd $? "Failed to add the Percona XtraDB Cluster monitoring scheduler in ProxySQL"

  # Adding Percona XtraDB Cluster membership checking scheduler
  proxysql_exec "DELETE FROM SCHEDULER WHERE ID=11;"
  check_cmd $?
  proxysql_exec "INSERT  INTO SCHEDULER (id,active,interval_ms,filename,arg1,arg2) VALUES (11,1,5000,'$PROXYSQL_NODE_MONITOR',10,'/var/lib/proxysql/proxysql_node_monitor.log');"
  check_cmd $? "Failed to add the Percona XtraDB Cluster membership checking scheduler in ProxySQL"

  proxysql_exec "LOAD SCHEDULER TO RUNTIME;SAVE SCHEDULER TO DISK;"

  echo -e "\nConfiguring Percona XtraDB Cluster application user to connect through ProxySQL"

  if [[ -z $PXC_APP_USERNAME ]]; then
    echo -n "Enter Percona XtraDB Cluster application user name: "
    read -r PXC_APP_USERNAME
    while [[ -z "$PXC_APP_USERNAME" ]]; do
      echo -n "No input entered, Enter Percona XtraDB Cluster application user name: "
      read -r PXC_APP_USERNAME
    done
  else
    echo -e "Percona XtraDB Cluster application user name as per command line is '$PXC_APP_USERNAME'"
  fi

  if [[ -z $PXC_APP_PASSWORD ]]; then
    read -r -s -p  "Enter Percona XtraDB Cluster application user password: " PXC_APP_PASSWORD
    while [[ -z "$PXC_APP_PASSWORD" ]]; do
      read -r -s -p  "No input entered, Enter Percona XtraDB Cluster application user password: " PXC_APP_PASSWORD
    done
  fi

  check_user=`mysql  $DSN -Bse"SELECT user,host FROM mysql.user where user='$PXC_APP_USERNAME' and host='%';" 2>/dev/null`

  if [[ -z "$check_user" ]]; then
    mysql  $DSN -e "CREATE USER $PXC_APP_USERNAME@'%' IDENTIFIED BY '$PXC_APP_PASSWORD';" 2>/dev/null
    check_cmd $? "Cannot add Percona XtraDB Cluster application user : '$PXC_APP_USERNAME' (GRANT)"
    proxysql_exec "INSERT INTO mysql_users (username,password,active,default_hostgroup) values ('$PXC_APP_USERNAME','$PXC_APP_PASSWORD',1,10);LOAD MYSQL USERS TO RUNTIME;SAVE MYSQL USERS FROM RUNTIME;SAVE MYSQL USERS TO DISK;"
    check_cmd $? "Cannot add Percona XtraDB Cluster application user : '$PXC_APP_USERNAME' (mysql_users update)"
    echo -e "\n\nPercona XtraDB Cluster application user '$PXC_APP_USERNAME'@'%' has been added with USAGE privilege, please make sure to grant appropriate privileges\n"
  else
    echo -e "\n\nApplication user '${PXC_APP_USERNAME}'@'%' already present in Percona XtraDB Cluster"
  fi
}

# Removing PXC configuration from proxysql
disable_proxysql(){
  proxysql_connection_check
  proxysql_exec "DELETE FROM mysql_users WHERE default_hostgroup=10;"
  check_cmd $? "Cannot delete Percona XtraDB Cluster users from ProxySQL"
  proxysql_exec "DELETE FROM mysql_servers WHERE hostgroup_id=10;"
  check_cmd $? "Cannot delete Percona XtraDB Cluster nodes from ProxySQL"
  proxysql_exec "DELETE FROM SCHEDULER WHERE ID IN (10,11);"
  check_cmd $? "Cannot delete Galera checker and node monitoring scheduler from ProxySQL"
  proxysql_exec "LOAD MYSQL USERS TO RUNTIME;SAVE MYSQL USERS TO DISK;LOAD SCHEDULER TO RUNTIME;SAVE SCHEDULER TO DISK;LOAD MYSQL SERVERS TO RUNTIME; SAVE MYSQL SERVERS TO DISK;"
}

adduser(){
  proxysql_connection_check
  pxc_connection_check
  echo -e "\nAdding Percona XtraDB Cluster application user to ProxySQL database"
  echo -n "Enter Percona XtraDB Cluster application user name: "
  read -r PXC_APP_USERNAME
  while [[ -z "$PXC_APP_USERNAME" ]]
  do
    echo -n "No input entered, Enter Percona XtraDB Cluster application user name: "
    read -r PXC_APP_USERNAME
  done
  read -r -s -p  "Enter Percona XtraDB Cluster application user password: " PXC_APP_PASSWORD
  while [[ -z "$PXC_APP_PASSWORD" ]]
  do
    read -r -s -p  "No input entered, Enter Percona XtraDB Cluster application user password: " PXC_APP_PASSWORD
  done

  check_user=`mysql  $PROXYSQL_DSN -Bse"SELECT username FROM mysql_users where username='$PXC_APP_USERNAME'" 2>/dev/null`
  if [[ -z "$check_user" ]]; then
    check_pxc_user=`mysql  $DSN -Bse"SELECT user,host FROM mysql.user where user='$PXC_APP_USERNAME'" 2>/dev/null`
    if [[ -z "$check_pxc_user" ]]; then
      echo -e "\n\n"
      read -p "Application user '$PXC_APP_USERNAME' does not exist in Percona XtraDB Cluster. Would you like to proceed [y/n] ? " check_param
      case $check_param in
        y|Y)
          proxysql_exec "INSERT INTO mysql_users (username,password,active,default_hostgroup) values ('$PXC_APP_USERNAME','$PXC_APP_PASSWORD',1,10);LOAD MYSQL USERS TO RUNTIME;SAVE MYSQL USERS FROM RUNTIME;SAVE MYSQL USERS TO DISK;"
          check_cmd $? "Cannot add Percona XtraDB Cluster application user : '$PXC_APP_USERNAME' to ProxySQL database"
        ;;
        n|N)
          exit 1
        ;;
        *)
          echo "Please type [y/n]! Terminating.."
          exit 1
        ;;
      esac
    else
      proxysql_exec "INSERT INTO mysql_users (username,password,active,default_hostgroup) values ('$PXC_APP_USERNAME','$PXC_APP_PASSWORD',1,10);LOAD MYSQL USERS TO RUNTIME;SAVE MYSQL USERS FROM RUNTIME;SAVE MYSQL USERS TO DISK;"
      check_cmd $? "Cannot add Percona XtraDB Cluster application user : '$PXC_APP_USERNAME' to ProxySQL database"
    fi
  else
    echo -e "\nERROR: Application user '$PXC_APP_USERNAME' already exist in ProxySQL database. Terminating.."
    exit 1
  fi
}

if [ "$enable" == 1 -o "$disable" == 1 -o "$adduser" == 1 ]; then
  if [ "$enable" == 1 ];then
    enable_proxysql
    echo "ProxySQL configuration completed!"
  fi
  if [ "$disable" == 1 ];then  
    disable_proxysql
    echo "ProxySQL configuration removed!"
  fi
  if [ "$adduser" == 1 ];then  
    adduser
    echo -e "\nAdded Percona XtraDB Cluster application user to ProxySQL database!"
  fi
else
  echo "Usage: proxysql-admin <user credentials> {enable|disable}"
  usage
fi

