#!/bin/bash
set -eu
### --help Info: скрипт управления репликацией - включение, отключение, статус
### --help Info:
### --help Usage: pg_replication.sh enable|disable|status|resync_replica|get_behind|get_wal_status
### --help Usage:                      [--keep-replaced-db] [--keep-down]
### --help Usage:
### --help Usage:   --keep-replaced-db - сохранить исходную БД, заменённую на реплику
### --help Usage:                БД сохранится в папке резервных копий. Удалять её нужно вручную.
### --help Usage:
### --help Usage:   --keep-down - оставить СУБД выключенной после настройки репликации
### --help Usage:
### --help Example: pg_replication.sh enable
### --help Example: pg_replication.sh enable --keep-replaced-db
### --help Example: pg_replication.sh disable
### --help Example: pg_replication.sh status
### --help Example:
. /opt/fox_utils/crab_sys.sh
[[ $@ =~ --help ]] && sys::usage
sys::arg_parse "$@"

. /opt/CONFIG

declare CMD
declare HA_ROLE
declare HA_COMPANION_IP
declare HOSTNAME
declare SLOT_NAME
declare PAGER
HOSTNAME="$(hostname)"
SLOT_NAME="eva_replica"

# Отключаем пагинатор.
# Проблема в запросе статуса пагинации для логов.
# Если запустить enable или status в маленьком окне терминала,
# psql запустит pager и будет жать его вечно, пока не сделать kill
export PAGER=""


_echo(){
	echo "$HOSTNAME - $@" >&2
	return 0
}


########################################
# Подождёт, пока СУБД на master станет доступна.
# Globals:
#   HA_COMPANION_IP - IP мастера
########################################
wait_for_master(){
	local timeout ret t
	timeout=20
	ret=1
	t=0
	while true; do
		nc -vz "$HA_COMPANION_IP" 5432 && ret=0 || ret=$?
		((ret == 0)) && break
		((t == timeout)) && break
		t=$((t + 1))
		sleep 1
	done
	return $ret
}


########################################
# Убедится, что существует слот репликации
# Globals:
#   SLOT_NAME - имя слота
########################################
ensure_replication_slot(){
	local sql exists
	sql="select active from pg_replication_slots where slot_name='$SLOT_NAME'"
	exists="$(psql -A -t -U postgres postgres -c "$sql")"
	if [[ "${exists:-}" != "f" ]]; then
		# Создаём слот, если его нет
		sql="SELECT pg_create_physical_replication_slot('$SLOT_NAME')"
		psql -U postgres postgres -c "$sql"
	fi
	# Информация для логов
	psql -x -U postgres postgres -c "select * from pg_replication_slots" >&2
	return 0
}


########################################
# Сконфигурирует Мастера: настроит PostgreSQL и создаст слот репликации
# Globals:
#   HA_COMPANION_IP - IP компаньона
########################################
enable_master(){
	local ha_conf hba_conf new_path
	ha_conf="/etc/postgresql/13/main/conf.d/pg_ha.conf"
	hba_conf='/etc/postgresql/13/main/pg_hba.conf'
	_echo "master - конфигурация Postgres conf.d/pg_ha.conf"
	cp /opt/bin/templates/pg_master.conf "$ha_conf"
	ex -s -c "g/# eva_ha/d" -cx "$hba_conf"
	_echo "master - конфигурация Postgres pg_hba.conf"
	echo "host replication postgres $HA_COMPANION_IP/32 trust # eva_ha" >> "$hba_conf"
	_echo "master - старт Postgres"
	/etc/init.d/postgresql restart
	_echo "master - создание слота репликации, если его нет"
	ensure_replication_slot
	[[ -n ${ARG_KEEP_DOWN:-} ]] && /etc/init.d/postgresql stop
	return 0
}


########################################
# Сконфигурирует реплику: настроит PostgreSQL, скопирует БД с пастера, запустит репликацию
# Globals:
#   HA_COMPANION_IP - IP компаньона
#   SLOT_NAME - имя слота репликации на мастере
# Args:
#   ARG_KEEP_REPLACED_DB (--keep-replaced-db) - сохранить текущую БД в архив
########################################
enable_replica(){
	local ha_conf replica_dir original_db new_path
	original_db="/mnt/shared/postgresql/13/main.SWITCH_TO_REPLICATION"
	[[ -d $original_db ]] && return 3

	wait_for_master

	_echo "replica - Конфигурация Postgres"
	ha_conf="/etc/postgresql/13/main/conf.d/pg_ha.conf"
	cp /opt/bin/templates/pg_replica.conf "$ha_conf"
	ex -c "%s/@@@ha_companion_ip%%%/$HA_COMPANION_IP/ge" -cx "$ha_conf"

	_echo "replica - Копия данных с мастера"
	replica_dir="/mnt/shared/eva_pg_backup/pg_basebackup_replica"
	[[ -d $replica_dir ]] && rm -rf --one-file-system "${replica_dir:-_NODIR}"
	pg_basebackup \
		-X stream -F plain -S $SLOT_NAME -h $HA_COMPANION_IP -p 5432  -U postgres -P \
		-D "$replica_dir"
	touch "$replica_dir/standby.signal"
	chown -R postgres "$replica_dir"

	_echo "replica - Переход на репликацию"
	/etc/init.d/postgresql stop
	mv /mnt/shared/postgresql/13/main "$original_db"
	mv "$replica_dir" /mnt/shared/postgresql/13/main
	[[ -z ${ARG_KEEP_DOWN:-} ]] && /etc/init.d/postgresql start

	if [[ -n ${ARG_KEEP_REPLACED_DB:-} ]]; then
		# Нельзя класть в /mnt/shared/eva_pg_backup, т.к. rsync_companion.sh --full удалит
		mkdir -p "/mnt/shared/eva_ha_archive"
		new_path="/mnt/shared/eva_ha_archive/pg.replaced_by_replication.$(date +%Y%m%d_%H%M%S)"
		_echo "replica - старая БД сохранена $new_path"
		mv "$original_db" "$new_path"
	else
		_echo "replica - Удаление старой БД"
		rm -rf --one-file-system "${original_db:-_NODIR}"
	fi
	return 0
}


########################################
# Отключит репликацию. Для мастера и реплики - одна функция.
# Globals:
#   SLOT_NAME - имя слота репликации
########################################
do_disable(){
	local sql exists pg_ha_conf hba_conf standby_file f in_recovery

	hba_conf='/etc/postgresql/13/main/pg_hba.conf'
	ex -s -c "g/# eva_ha/d" -cx "$hba_conf"
	/etc/init.d/postgresql restart

	if pgrep postgres; then
		in_recovery="$(psql -A -t  -U postgres postgres -c "select pg_is_in_recovery()")"
		if [[ $in_recovery == f ]]; then
			sql="select active from pg_replication_slots where slot_name='$SLOT_NAME'"
			exists="$(psql -A -t -U postgres postgres -c "$sql")"
			if [[ -n "${exists:-}" ]]; then
				_echo "Удаляем слот репликации $SLOT_NAME"
				psql -U postgres -c "select pg_drop_replication_slot('$SLOT_NAME');"
			fi
			# Информация для логов
			psql -x -U postgres postgres -c "select * from pg_replication_slots"
		elif [[ $in_recovery == t ]]; then
			_echo "promote replica to master"
			psql -U postgres postgres -c "select pg_promote()"
		fi
	fi

	pg_ha_conf="/etc/postgresql/13/main/conf.d/pg_ha.conf"
	# Его удалит pg_promote. На всякий случайЮ проверяем.
	standby_file="/mnt/shared/postgresql/13/main/standby.signal"
	for f in $pg_ha_conf $standby_file; do
		if [[ -f $f ]]; then
			_echo "Удаляем $f"
			rm -f "$f"
		fi
	done
	[[ -n ${ARG_KEEP_DOWN:-} ]] && /etc/init.d/postgresql stop
	return 0
}


########################################
# Включит репликацию
# Globals:
#   HA_ROLE - роль текущего инстанса, master или replica
########################################
do_enable(){
	if [[ $HA_ROLE == master ]]; then
		enable_master
	elif [[ $HA_ROLE == replica ]]; then
		enable_replica
	else
		_echo "ERROR в /opt/CONFIG не указан HA_ROLE"
	fi
	return 0
}


########################################
# Заново скопирует БД с мастера.
# Функция нужна, чтобы перегрузить БД на реплике, после восстановления мастера из бекапа.
# Globals:
#   HA_ROLE - роль текущего инстанса, master или replica
# Returns:
#   4 - если запустить на ВМ с ролью, отличной от replica
########################################
do_resync_replica(){
	if [[ ${HA_ROLE:-} != replica ]]; then
		echo "ERR! HA_ROLE=\"$HA_ROLE\". Функцию можно выполнить только на replica."
		return 4
	fi
	wait_for_master
	enable_replica
	[[ -n ${ARG_KEEP_DOWN:-} ]] && /etc/init.d/postgresql stop
	return 0
}


########################################
# Выведет информацию по статусу репликации.
# TODO: Выводит детально инфу из БД служебными функциями. Возможно, стоит докрутить аналитику.
#       В аналитике писать какие-то понятные статусы: ок, не ок и т.д, обращать на что-то
#       внимание инженера.
#       Сейчас единственная аналитика - replica_bytes_behind (насколько байт отстала реплика).
# Globals:
#   HA_ROLE - роль текущего инстанса, master или replica
# Outputs:
#   Статус репликации, вывод встроенных функций PostgreSQL
########################################
do_status(){
	_echo status
	if [[ $HA_ROLE == master ]]; then
		psql -x -U postgres postgres -c "SELECT *,
		pg_wal_lsn_diff(pg_current_wal_lsn(), replay_lsn) as replica_bytes_behind
		FROM pg_stat_replication;"
	elif [[ $HA_ROLE == replica ]]; then
		psql -x -U postgres postgres -c "SELECT * FROM pg_stat_wal_receiver;"
	else
		_echo "ERROR в /opt/CONFIG не указан HA_ROLE"
	fi
	return 0
}


########################################
# Выведет информацию, насколько реплика отстала от мастера, в байтах.
# Вспомогательная функция для скриптов, должна выводить только число, и больше ничего.
# Globals:
#   HA_COMPANION_IP - IP компаньона
# Outputs:
#   Отставание реплики от мастера, в байтах
########################################
do_get_behind(){
	local sql
	sql="SELECT pg_wal_lsn_diff(pg_current_wal_lsn(), replay_lsn)"
	sql+=" FROM pg_stat_replication where client_addr='$HA_COMPANION_IP';"
	/opt/bin/sqlexec -t -A "$sql"
	return 0
}


########################################
# Выведет информацию о состоянии WAL для слота репликации.
# https://www.postgresql.org/docs/13/view-pg-replication-slots.html
# Вспомогательная функция для скриптов, должна выводить только число, и больше ничего.
# Globals:
#   SLOT_NAME - имя слота репликации
# Outputs:
#   Статус WAL слота репликации
########################################
do_get_wal_status(){
	local sql
	sql="select wal_status from pg_replication_slots where slot_name='$SLOT_NAME'"
	/opt/bin/sqlexec -t -A "$sql"
	return 0
}


main(){
	CMD="$ARG_1"
	"do_$CMD"
	return 0
}


main "$@"

exit 0