#!/bin/bash
set -o nounset

##########################################
# Defaults
##########################################
: ${recipient:=}
: ${host:=localhost}
: ${port:=5432}
: ${databases:=}
: ${user:=$(id -un)}
: ${location:=/tmp}
: ${schemas:=}
: ${num_backups:=1}
: ${password:=}
: ${prompt_for_password:=}

PSQL="$(command -v -p psql)"
PGDUMP="$(command -v -p pg_dump)"
GPG="$(command -v -p gpg)"

##########################################
# Utility functions and their setup
##########################################
tmpdir=$(mktemp -d)
cleanup_pids=""

function log () {
  local level=${1}
  shift

  local end="" start=""

  if [[ -x $(which tput) && $(tput setaf 1 &> /dev/null; echo $?) == 0 ]]; then
    end="$(tput sgr0)"
    if [[ ${level} == "info" ]]; then
      start="$(tput setaf 2)"
    elif [[ ${level} == "warn" ]]; then
      start="$(tput setaf 3)"
    elif [[ ${level} == "error" ]]; then
      start="$(tput setaf 1)"
    fi
  fi

  echo -e "${start}${@}${end}" >&2
}

function cleanup () {
  for pid in ${cleanup_pids}; do
    kill -9 ${pid} &> /dev/null
  done
  cleanup_pids=""

  rm -rf ${tmpdir}/*
}
trap cleanup SIGINT SIGQUIT SIGTERM

function cleanup-on-exit () {
  cleanup
  rm -rf ${tmpdir}
}
trap cleanup-on-exit SIGKILL EXIT

function rotate-backups () {
  # start by removing backups numbered greater than $num_backups
  for file in $(ls -1 ${dump_file_base}.[0-9]* 2>/dev/null); do
    if (( ${file##*.} >= ${num_backups} )); then
      rm -f $file
    fi
  done

  # rotate backups up one
  for ((num = ${num_backups} - 1; num > 0; num--)); do
    if [[ -f ${dump_file_base}.${num} ]]; then
      ((numplus = ${num} + 1))
      mv ${dump_file_base}.${num} ${dump_file_base}.${numplus}
    fi
  done

  # rotate the base file up
  if ((${num_backups} > 0)); then
    if [[ -f ${dump_file_base} ]]; then
      mv ${dump_file_base} ${dump_file_base}.1
    fi
  fi

  # move the new backup into place
  mv ${dump_file} ${dump_file_base}
}

function usage () {
  cat <<- EOT
  Usage :  $(basename ${0}) [options] [--]

  Options:
    -l   Directory to place backup file
    -e   GPG recipent key
    -u   User
    -h   Database host
    -p   Port number
    -d   Database to backup
         Multiple -d options can be given
    -n   Restrict backup to a given schema (ex: public)
         Multiple -n options can be given
         Keep in mind that any schemas given will apply for all databases
    -w   Prompt for password
    -k   Number of old backups to keep (default is one)

EOT
}

##########################################
# Arguement parsing
##########################################
# Require options...
if [[ ${#} == 0 ]]; then
  usage
  exit 1
fi

while getopts "k:e:h:d:u:l:p:n:w" opt; do
  case "${opt}" in
    e)
      [[ ! ${OPTARG} =~ ^- ]] && recipient="${OPTARG}"
      ;;

    h)
      [[ ! ${OPTARG} =~ ^- ]] && host="${OPTARG}"
      ;;

    d)
      [[ ! ${OPTARG} =~ ^- ]] && databases="${databases:-} ${OPTARG}"
      ;;

    u)
      [[ ! ${OPTARG} =~ ^- ]] && user="${OPTARG}"
      ;;

    l)
      [[ ! ${OPTARG} =~ ^- ]] && location="${OPTARG}"
      ;;

    p)
      [[ ! ${OPTARG} =~ ^- ]] && port="${OPTARG}"
      ;;

    n)
      [[ ! ${OPTARG} =~ ^- ]] && schemas="${schemas:-} ${OPTARG}"
      ;;

    w)
      prompt_for_password="-W"
      ;;

    k)
      [[ ! ${OPTARG} =~ ^- ]] && num_backups="${OPTARG}"
      ;;

    \?)
      usage
      exit 1
      ;;
  esac
done
shift $(($OPTIND-1))

##########################################
# Option checking
##########################################
if [[ -z ${PGDUMP} || ! -x ${PGDUMP} ]]; then
  log error "pg_dump not found"
  _exit_error=2
fi

if [[ -z ${databases} ]]; then
  log error "Database was not specified"
  _exit_error=2
fi

if [[ -n ${recipient} ]]; then
  if [[ ! -x ${GPG} ]]; then
    log error "GPG exectuable not available, caonnot encrypt backup."
    _exit_error=2
  fi
  if [[ ! $(${GPG} --list-keys | grep uid | grep "${recipient}") ]]; then
    log error "GPG encryption key not found for: ${recipient}"
    _exit_error=2
  fi
fi

if [[ ! -d ${location} ]]; then
  log error "Backup location must be a directory"
  _exit_error=2
elif [[ ! -w ${location} ]]; then
  log error "Backup location is not writable (insufficient permissions)"
  _exit_error=2
fi

if [[ -n ${_exit_error:-} ]]; then
  echo
  usage
  exit ${_exit_error}
fi

##########################################
# Begin the actual backup
##########################################
if [[ -n ${prompt_for_password} ]]; then
  read -p "Enter psql password for ${user}: " -s password
  echo
  export PGPASSWORD="${password}"
fi

schema_opts=""
for schema in ${schemas}; do
  schema_opts="${schema_opts} -n ${schema}"
done

if [[ -n ${recipient} ]]; then
  log info "Encrytping backups for: ${recipient}\n"
fi

for database in ${databases}; do
  retcode=0
  dump_file_base="${location}/${database}.pg_dump${recipient:+.gpg}"
  dump_file="${dump_file_base}.new"

  log info $(date)
  log info "Dumping ${database}@${host}:${port} as ${user} to ${dump_file_base}"

  PGDUMP_COMMAND="${PGDUMP} --ignore-version --format=custom --username ${user} \
    --host ${host} --port ${port} ${schema_opts} ${database}"

  # Check that we can actually dump
  ${PGDUMP_COMMAND} -w --schema-only &> /dev/null
  if [[ $? != 0 ]]; then
    log error "Cannot connect to ${database}@${host}:${port} as ${user}"
    log error "  You probably need to provide a password"
    log error $(date): Backup failed
    continue
  fi

  if [[ -n ${recipient:-} ]]; then
    ${GPG} --default-recipient ${recipient} -o ${dump_file} \
      -e <(${PGDUMP_COMMAND}) 2> ${tmpdir}/gpg_error.log &
    gpg_pid=${!}
    cleanup_pids+=" ${gpg_pid}"
    while $(ps -p ${gpg_pid} &> /dev/null); do
      sleep 1 || { retcode=1; break; }
      if $(grep -q -i 'no space left on device' ${tmpdir}/gpg_error.log); then
        log error "Could not write to ${location}: insufficient space"
        cleanup
        retcode=1
      fi
    done
  else
    ${PGDUMP_COMMAND} -f ${dump_file}
    retcode=${?}
  fi

  if [[ ${retcode} != 0 ]]; then
    rm -f ${dump_file}
    log error "Encountered error during dump, removing partial dump file"
    log error $(date): Backup failed
  else
    # Rotate the old backups now that we know we can connect
    log info "Rotating backups for ${database}, keeping ${num_backups}"
    rotate-backups

    log info "$(date): Backup completed"
  fi

  # make sure that our tempdir is clean for the next db
  cleanup
done