#!/bin/sh
#
# postgrespro-setup	Initialization and control operations for Postgres Pro
#

PG_SETUP_VERSION=13.2.1
BINDIR=/opt/pgpro/std-15/bin
SERVICE=postgrespro-std-15
PGENGINE=$BINDIR
# List of unicode locales which should be tried in order of appearance
# if current locale of script and default locale of user postgres do not
# suit our needs.
FALLBACK_LOCALES="ru_RU.UTF-8 en_US.UTF-8 kk_KZ.UTF-8 tr_TR.UTF-8"


USAGE_STRING="Usage: $0 initdb [<initdb-options>]
   or: $0 find-free-port
   or: $0 set-server-port <port>
   or: $0 service (enable|disable)
   or: $0 service (start|stop|condrestart|status)
   or: $0 set GUC-variable value
   or: $0 psql [<psql options>]

Script is aimed to help sysadmin with basic database cluster administration.

Available operation mode:
  initdb [--tune=conf] [<initdb-options>] Create a new Postgres Pro database cluster.
  find-free-port            Search of free server TCP-port number, start with 5432.
  set-server-port <port>    Set server port number.
  service enable            Enable service.
  service disable           Disable service.
  service status            Service status.
  service start             Start service.
  service stop              Stop service.
  service condrestart       Conditional restart of service.
  
  pg-setup psql allows to start SQL shell correctly switching from root
  to postgres user and use port set by pg-setup set-server-port.
"

# note that these options are useful at least for help2man processing
case "$1" in
    --version)
	echo "pg-setup $PG_SETUP_VERSION"
        exit 0
        ;;
    --help|--usage)
        echo "$USAGE_STRING"
        exit 0
        ;;
    "")
    	echo "$USAGE_STRING"
        exit 1
        ;;	
esac

if [ "$(id -u)" -ne 0 ]; then
  echo "$0 must be started as root" >&2
  exit 1
fi

PGDATA=/var/lib/pgpro/std-15/data
DEFCONFIG=/etc/default/postgrespro-std-15
# shellcheck disable=SC1090
[ -f $DEFCONFIG ] && . $DEFCONFIG
export PGDATA
# make safe PATH
PATH=/sbin:/bin:/usr/sbin:/usr/bin:/opt/pgpro/std-15/bin
export PATH

# For SELinux we need to use 'runuser' not 'su'
if [ -x /sbin/runuser ]
then
    SU=/sbin/runuser
else
    SU=su
fi

#
# This function sets confguration value in the postgresql.conf to given
# string. If value should contain quotes, they should be passed  as part
# of string
# i.e. set_conf_value log_destination "'stderr'"
# Keeps end of line comments if it can.
# Removes corresponding paramters from postgresql.auto.conf if any
# FIXME - this function is unable to handle arbitrary includes,
# but possibly would do a right thing, because if directive is in the
# include, it would not find it in the main configuration file 
# and append new value to the end of it after all includes.

set_conf_value() {
	param="$1"
	value="$2"
    config="$PGDATA/postgresql.conf"
    if [ ! -f "$config" ]; then
        echo "Cannot find configuration file for $PGDATA!" >&2
		return 1
    fi

	# Check if there is (possibly commented out) parameter in the config
	if grep -q "^#\\?${param}[ 	]*=" "$config"; then
		# Found one edit
		sed -i "s/^#\\?${param}[ 	]*=[ 	]*[^#]*\\(#.*\\)\\?/$param = $value \\1/" \
			"$config"
	else 
		# Not found, just append
		echo "$param = $value" >> "$config"
	fi
	# Check postgresql.auto.conf
	if [ -f "$PGDATA/postgresql.auto.conf" ] && grep -q "^${param}[ 	]*=" "$PGDATA/postgresql.auto.conf"; then
		#Really we could call sed without previous check, but
		#I want to keep file timestamp if there is nothing to change
		sed -i "/^${param}[ 	]*=/d" "$PGDATA/postgresql.auto.conf";
	fi
}

#
#  This function checks  whether environment variable needs shell
#  quoting. Pass there unquited variable and function will return
#  0 if it got more than one argument and 1 if just one
need_quotes() {
	[ "$#" -gt 1 ] 
}
#
# This function checks whether existing directory does contain any files 
# and subdirectories
#
dir_is_not_empty() {
	[ "$(ls -A "$1")" ]
}

initdb(){
	INITDB_OPTIONS=""
	datadir=""
	TUNE=std
	for opt in "$@"; do
		if [ -n "$consume_next" ]; then
			datadir="$opt"
			consume_next=""
			continue
		fi	
		case "$opt" in
		-D|--pgdata)
			if [ -n "$datadir" ]; then
				echo "Duplicate option $opt" 1>&2
				return 1
			fi	
			consume_next=1
		;;
		--pgdata=*)
			if [ -n "$datadir" ]; then
				echo "Duplicate option --pgdata" 1>&2
				return 1
			fi	
			datadir=$(echo "$opt"|sed 's/^[^=]*=//')
		;;
		--auth-local=*)
			if [ -n "$local_auth" ]; then
				echo "Duplicate option --auth-local" 1>&2
				return 1
			fi
			local_auth=$opt
		;;
		--auth-host=*)
			if [ -n "$host_auth" ]; then
				echo "Duplicate option --auth-host" 1>&2
				return 1
			fi
			host_auth=$opt
		;;
		--tune|--auth-local|-auth-host)
			echo "Please use syntax $opt=value" >&2;
			return 1
		;;
		--tune=*)
			TUNE=${opt#*=}
			;;
		*)
			INITDB_OPTIONS="$INITDB_OPTIONS \"$opt\""
		;;	
		esac
	done
	if [ -n "$datadir" ]; then
		# shellcheck disable=SC2166
		if [ "$datadir" != "$PGDATA" -a -f ${DEFCONFIG} ]; then
			echo "${DEFCONFIG} already exists " >&2
			echo "And sets up database location $PGDATA." >&2
			echo "If you want to setup second postgres instance in" >&2
			echo "$datadir" >&2
			echo "use /opt/pgpro/std-15/bin/initdb directly and configure service" >&2
			echo "startup manually." >&2
			return 1
		fi
		PGDATA="$datadir"
		export PGDATA
	fi 
	# Check for relative path 
	if [ ! -d "$PGDATA" ]; then
		mkdir -p "$PGDATA" || return 1
	elif  dir_is_not_empty "$PGDATA" ; then
		echo "Data directory $PGDATA is not empty!" 1>&2
 		return 1
	fi
	# convert path to PGDATA to absolute
	PGDATA="$(sh -c "cd '$PGDATA' && pwd -P")"

 	chown postgres:postgres "$PGDATA"
	chmod go-rwx "$PGDATA"
	# We intend word splitting here. This line checks if
	# PGDATA would be splitted
	# shellcheck disable=SC2086
	if need_quotes $PGDATA ; then
		echo "PGDATA=\"$PGDATA\"" > "${DEFCONFIG}"
	else	
		echo "PGDATA=$PGDATA" > "${DEFCONFIG}" 
	fi
	PGLOG="$(dirname "${PGDATA}")/initdb.$(basename "${PGDATA}").log"
	# Clean up SELinux tagging for PGDATA
	RESTORECON=/sbin/restorecon
	[ -x $RESTORECON ] && $RESTORECON "$PGDATA"
	# Check if locale of this script is suitable for creating database
	# and try to find better one if not

	if [ "$(locale charmap 2>/dev/null)" = "ANSI_X3.4-1968" ]; then
		# If current locale is 7-bit, try to get locale of
		# invoking user or default locale of user postgres, 
		# shellcheck disable=SC2046
		eval $($SU - "${SUDO_USER:-postgres}" -c locale 2>/dev/null)
		[ -z "$LANG" ] || export LANG
		[ -z "$LC_CTYPE" ] || export LC_CTYPE
	elif [ -z "$LANG" ]&&[ -n "$LC_CTYPE" ]; then
		# If LC_CTYPE set to something usable, but LANG is unset,
		# set it equal to LC_CTYPE, because we need LC_COLLATE as well
		# as LC_CTYPE
		LANG="$LC_CTYPE"
		export LANG
	fi
	# If locale is C with some encoding, reset it to just C, because ICU
	# doesn't handle C.UTF-8 and postgres do not map it anymore
	if [ "$(locale 2>/dev/null|sed -n '/LC_CTYPE/{s/^.*="//;s/"$//;s/\..*$//;p}')" = "C" ]; then
		LANG=C; export LANG
		unset LC_CTYPE
		unset LC_COLLATE
		unset LC_ALL
	fi
	# if locale still 7-bit, try fallback locales

	if [ "$(locale charmap 2>/dev/null)" = "ANSI_X3.4-1968" ]; then
		unset LC_CTYPE
		unset LC_ALL
		unset LC_COLLATE
		for L in $FALLBACK_LOCALES; do
			LANG="$L"; export LANG
			[ "$(locale charmap 2>/dev/null)" = "UTF-8" ] && break
		done
	fi
	# if locale still 7-bit, search locale -a output for anything utf8
	if [ "$(locale charmap)" = "ANSI_X3.4-1968" ]; then
		for L in $(locale -a | sed -n 's/\.utf8$/.UTF-8/p'); do
			[ "$L" = "C.UTF-8" ] && continue
			LANG="$L"; export LANG
			[ "$(locale charmap 2>/dev/null)" = "UTF-8" ] && break
		done
	fi	
	# if no utf-8 locale found, fail.
	if [ "$(locale charmap)" = "ANSI_X3.4-1968" ]; then
		echo "No suitable locale found. Please install/create some UTF-8 locale" >&2
		rm -rf "$PGDATA" "$DEFCONFIG"
		return 1
	fi
	LC_CTYPE=$LANG; export LC_CTYPE
	LC_COLLATE=$LANG; export LC_COLLATE
	# Create the initdb log file if needed
	# shellcheck disable=SC2166
	if [  -f "$PGLOG" -o ! -e "$PGLOG" ]; then
		if ! : > "$PGLOG" ; then
			echo "Cannot write installation log file $PGLOG" >&2
			rm -rf "${PGDATA}" "${DEFCONFIG}"
			return 1
		fi
		chown postgres:postgres "$PGLOG"
		chmod go-wx "$PGLOG"
		[ -x $RESTORECON ] && $RESTORECON "$PGLOG"
	else
		echo "Log file $PGLOG exist and is not regular file. Cannot continue." >&2
		rm -rf "${PGDATA}" "${DEFCONFIG}"
		return 1
	fi

	# Initialize the database
	initdbcmd="$PGENGINE/initdb --pgdata='$PGDATA' ${local_auth:---auth-local=peer} ${host_auth:---auth-host=md5} $INITDB_OPTIONS"
	if $SU -l postgres -c "LANG=$LANG $initdbcmd" >> "$PGLOG" 2>&1 < /dev/null; then
		# Create directory for postmaster log files
		mkdir "$PGDATA/log"
		chown postgres:postgres "$PGDATA/log"
		chmod go-rwx "$PGDATA/log"
		[ -x $RESTORECON ] && $RESTORECON "$PGDATA/log"
		# Tune configuration 
		if [ -x "/opt/pgpro/std-15/share/$TUNE.tune" ]; then
			if ! "/opt/pgpro/std-15/share/$TUNE.tune" >> "$PGDATA/postgresql.conf"; then
				rm -rf "${PGDATA}" "${DEFCONFIG}"
				return 1
			fi
			if grep -q "^listen_addresses[      ]*=[   ]*'\\*'" "$PGDATA/postgresql.conf"
			then
				echo "host		all		all		0.0.0.0/0	md5" >> "$PGDATA/pg_hba.conf"
			fi
		fi
		# Set server port
		if grep -q "^listen_addresses[      ]*=[   ]*'\\*'" "$PGDATA/postgresql.conf"; then
		    set_server_port "$(find_free_port 5432 1)"
		else
		    set_server_port "$(find_free_port 5432 0)"
		fi
		# Fix configuration file to enable loging_collector
		if set_conf_value "logging_collector" "on"; then
			echo OK
			return 0
		fi
	fi	
	rm -rf "${PGDATA}" "${DEFCONFIG}"
	echo "failed, see $PGLOG" 1>&2
	return 1
}

unix_port_free() {
    unix_port=$1
    [ -e "/tmp/.s.PGSQL.$unix_port" ] && return 1
    return 0
}

tcp_port_free() {
    tcp_port=$1
    /usr/bin/timeout 5 /bin/bash -c "cat < /dev/null > /dev/tcp/127.0.0.1/$tcp_port" 2>/dev/null && return 1
    return 0
}

find_free_port() {
    free_port=$1
    check_tcp=$2
    while true; do
	if [ "$check_tcp" = "1" ]; then
	    unix_port_free "$free_port" && tcp_port_free "$free_port" && break
	else
	    unix_port_free "$free_port" && break
	fi
	# shellcheck disable=SC2003
	free_port=$((free_port + 1))
    done
    echo "$free_port"
}

set_server_port(){
    port=$1
	if [ -z "$port" ]; then
	    echo "Error: port value is empty!" >&2
		return 1
	fi
	# shellcheck disable=SC2003
	if ! /usr/bin/expr "$port" + 0 >/dev/null >&2; then
		echo "Bad port value: $port" 1>&2
		return 1
	fi
	# shellcheck disable=SC2166
	if [ "$port" -lt 1024 -o "$port" -gt 65535 ]; then
		echo "Port value $port is out of range (1024-65535)" 1>&2;
		return 1
	fi
	echo "Server will use port $port"
	set_conf_value "port" "$port"
}

enable_service() {
    if [ -n "$SCTL" ]; then
	$SCTL enable $SERVICE
	return $?
    else
	URCD=/usr/sbin/update-rc.d
	if [ -x $URCD ]; then
	    $URCD $SERVICE defaults
	else
	    /sbin/chkconfig --add $SERVICE
	    /sbin/chkconfig --level 35 $SERVICE on
	fi
	return $?
    fi
}

disable_service() {
    if [ -n "$SCTL" ]; then
	$SCTL disable $SERVICE
	return $?
    else
	URCD=/usr/sbin/update-rc.d
	if [ -x $URCD ]; then
	    $URCD -f $SERVICE remove
	else
	    /sbin/chkconfig --del $SERVICE
	fi
	return $?
    fi
}

start_service() {
    if [ -n "$SCTL" ]; then
	$SCTL is-active --quiet $SERVICE || $SCTL start $SERVICE
	return $?
    else
	SS=/usr/sbin/service
	if [ -x /sbin/service ]; then
	    SS=/sbin/service
	fi
	$SS $SERVICE start
	return $?
    fi
}

status_service() {
    if [ -n "$SCTL" ]; then
	$SCTL status $SERVICE
	return $?
    else
	SS=/usr/sbin/service
	if [ -x /sbin/service ]; then
	    SS=/sbin/service
	fi
	$SS $SERVICE status
	return $?
    fi
}

stop_service() {
    if [ -n "$SCTL" ]; then
	$SCTL stop $SERVICE
	return $?
    else
	SS=/usr/sbin/service
	if [ -x /sbin/service ]; then
	    SS=/sbin/service
	fi
	$SS $SERVICE stop
	return $?
    fi
}

condrestart_service() {
    if [ -n "$SCTL" ]; then
	if $SCTL 1>/dev/null 2>/dev/null; then
	    $SCTL condrestart $SERVICE
	    return $?
	fi
    else
	SS=/usr/sbin/service
	if [ -x /sbin/service ]; then
	    SS=/sbin/service
	fi
	$SS $SERVICE condrestart
	return $?
    fi
}

service() {
	if [ ! -f ${DEFCONFIG} ]; then
		echo "Default database not found. Use $0 initdb to create it" 2>&1
		if [ "$1" = "enable" ]; then
			return 1
		else
			return 0
		fi
	fi	
	SCTL=""
	# shellcheck disable=2166
	if [ -d "/run/systemd/system" -a -x "/bin/systemctl" ]; then
		SCTL=/bin/systemctl
	fi	
	export SCTL
	cmd=$1
	shift
    case "${cmd}" in
	start|stop|enable|disable|status|condrestart)
	    "${cmd}_service" "$@"

	    ;;
	*)
            echo >&2 "$USAGE_STRING"
            exit 2
    esac
}

command="$1" 
shift

case "$command" in
    initdb)
	    # We want wordsplitting here
	    # shellcheck disable=SC2086
        initdb "$@" $PGSETUP_INITDB_OPTIONS || exit 1
        ;;
    find-free-port)
		find_free_port 5432 1 || exit 1
	;;
    set-server-port)
		set_server_port "$@" || exit 1
	;;
    service)
		service "$@" || exit 1
	;;
	set)
		set_conf_value "$@" || exit 1 
	;;
	psql)
		if [ ! -f "$PGDATA/postgresql.conf" ]; then
			echo "Database in $PGDATA not found. Do you need pg-setup initdb first?" 1>&2
			exit 1;
		fi	
		PORT=$(sed -n 's/^port[ 	]*=[ 	]*\([0-9]\+\)/\1/p' "$PGDATA/postgresql.conf")
		$SU	 -l postgres -c "exec ${BINDIR}/psql ${PORT:+-p ${PORT}} $*"
	;;
    *)
        echo >&2 "$USAGE_STRING"
        exit 2
	;;
esac
