postgres_container_restore

#!/bin/bash

while getopts 'f:h:u:p:n:P:s:i:' flag; do
  case "${flag}" in
    f) DUMPFILE="${OPTARG}" ;;
    h) DBHOST="${OPTARG}" ;;
    u) DBUSER="${OPTARG}" ;;
    p) DBPASSWORD="${OPTARG}" ;;
    n) DBNAME="${OPTARG}" ;;
    P) DBPORT="${OPTARG}" ;;
    s) SSHHOST="${OPTARG}" ;;
    i) SSHKEY="${OPTARG}" ;;
  esac
done

DBPORT="${DBPORT:-5432}"
SSH_PID=""

if [ -z "$DUMPFILE" ] || [ -z "$DBHOST" ] || [ -z "$DBUSER" ] || [ -z "$DBPASSWORD" ] || [ -z "$DBNAME" ]; then
  echo "postgres_container_restore [OPTIONS...]"
  echo ""
  echo "Restore a PostgreSQL database from a backup file, replacing all existing data."
  echo ""
  echo "Flags:"
  echo "  -f    Path to the .dump.gz backup file"
  echo "  -h    Hostname or IP of the PostgreSQL server"
  echo "  -u    PostgreSQL username"
  echo "  -p    PostgreSQL password"
  echo "  -n    Target database name"
  echo "  -P    PostgreSQL port (default: 5432)"
  echo "  -s    SSH host to tunnel through (e.g. user@hostname)"
  echo "  -i    SSH identity file (private key)"
  echo ""
  echo "Example:"
  echo "  postgres_container_restore -f /var/backup/postgres/mydb/mydb_20240101-120000.dump.gz -h 127.0.0.1 -u myuser -p mypassword -n mydb -s deploy@10.0.0.1"
  exit 1
fi

if [ ! -f "$DUMPFILE" ] || [ ! -r "$DUMPFILE" ]; then
  echo "Dump file ${DUMPFILE} does not exist or is not readable"
  exit 2
fi

if ! hash pg_restore 2>/dev/null; then
  echo "pg_restore is required, please install postgresql-client"
  exit 3
fi

if ! hash psql 2>/dev/null; then
  echo "psql is required, please install postgresql-client"
  exit 3
fi

if ! hash gunzip 2>/dev/null; then
  echo "gunzip is required, please install gzip"
  exit 3
fi

if [ -n "$SSHHOST" ] && ! hash ssh 2>/dev/null; then
  echo "ssh is required for tunnelling"
  exit 3
fi

export PGPASSWORD="${DBPASSWORD}"

CONN_HOST="${DBHOST}"
CONN_PORT="${DBPORT}"

if [ -n "$SSHHOST" ]; then
  LOCAL_PORT=$(( RANDOM + 49152 ))
  SSH_KEY_OPT=""
  [ -n "$SSHKEY" ] && SSH_KEY_OPT="-i ${SSHKEY}"

  echo "Opening SSH tunnel to ${SSHHOST} (localhost:${LOCAL_PORT} -> ${DBHOST}:${DBPORT})..."
  ssh -N -o ExitOnForwardFailure=yes -o ConnectTimeout=10 ${SSH_KEY_OPT} \
    -L "${LOCAL_PORT}:${DBHOST}:${DBPORT}" "${SSHHOST}" &
  SSH_PID=$!
  sleep 2

  if ! kill -0 "${SSH_PID}" 2>/dev/null; then
    echo "Failed to establish SSH tunnel to ${SSHHOST}"
    exit 4
  fi

  CONN_HOST="127.0.0.1"
  CONN_PORT="${LOCAL_PORT}"
fi

trap '[ -n "${SSH_PID}" ] && kill "${SSH_PID}" 2>/dev/null' EXIT

echo "Terminating existing connections to ${DBNAME}..."
psql -h "${CONN_HOST}" -p "${CONN_PORT}" -U "${DBUSER}" -d postgres \
  -c "SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = '${DBNAME}' AND pid <> pg_backend_pid();" \
  > /dev/null

echo "Dropping database ${DBNAME} if it exists..."
psql -h "${CONN_HOST}" -p "${CONN_PORT}" -U "${DBUSER}" -d postgres \
  -c "DROP DATABASE IF EXISTS \"${DBNAME}\";"

if [ $? -ne 0 ]; then
  >&2 echo "Failed to drop database ${DBNAME}"
  exit 5
fi

echo "Creating database ${DBNAME}..."
psql -h "${CONN_HOST}" -p "${CONN_PORT}" -U "${DBUSER}" -d postgres \
  -c "CREATE DATABASE \"${DBNAME}\";"

if [ $? -ne 0 ]; then
  >&2 echo "Failed to create database ${DBNAME}"
  exit 6
fi

echo "Restoring ${DUMPFILE} into database ${DBNAME}..."
gunzip -c "${DUMPFILE}" | pg_restore -h "${CONN_HOST}" -p "${CONN_PORT}" -U "${DBUSER}" \
  -d "${DBNAME}" --no-owner --role="${DBUSER}"
PIPE_STATUS=("${PIPESTATUS[@]}")

if [ ${PIPE_STATUS[0]} -eq 0 ] && [ ${PIPE_STATUS[1]} -eq 0 ]; then
  echo "Restore of database ${DBNAME} completed successfully"
else
  >&2 echo "Restore of database ${DBNAME} failed"
  exit 7
fi

exit 0