#!/bin/bash
PGDIR=/home/ants/code/compiles/pg-9.2
BINDIR="${PGDIR}/bin"
BASEPATH=/media/VMs/dbs/replicatest

N1DATA="${BASEPATH}/n1"
N1PORT="5432"

N2DATA="${BASEPATH}/n2"
N2PORT="5433"

N3DATA="${BASEPATH}/n3"
N3PORT="5434"

ARCHIVEDIR="${BASEPATH}/archive"

PAGER="cat"

cleanup() {
	echo "Cleanup"
	$BINDIR/pg_ctl -D "${N1DATA}" -m fast stop
	$BINDIR/pg_ctl -D "${N2DATA}" -m fast stop
	$BINDIR/pg_ctl -D "${N3DATA}" -m fast stop
	rm -rf "${N1DATA}" "${N2DATA}" "${N3DATA}" "${ARCHIVEDIR}"
}

make_recoveryconf() {
	(cat <<-RECOVERYCONF
	restore_command = 'cp "${ARCHIVEDIR}/%f" %p'
	recovery_target_timeline = 'latest'
	standby_mode = on
	primary_conninfo = 'host=localhost port=${2}'
	trigger_file = 'failover.trigger'
	RECOVERYCONF
	) > "${1}/recovery.conf"
}

verify_up() {
	$BINDIR/psql -p $1 -c "SELECT max(ts) FROM test" || exit
}
verify_up_as_master() {
	until $BINDIR/psql -p $1 -c "INSERT INTO test (ts) VALUES (CURRENT_TIMESTAMP) RETURNING ts"; do
		echo "  ..waiting"
		sleep 1
    done
}
replication_status() {
	while [[ -z $( psql -p $1 -t -c "SELECT state FROM pg_stat_replication" ) ]]; do
	   echo "  ..waiting"
	   sleep 1
	done
	$BINDIR/psql -p $1 -c "SELECT * FROM pg_stat_replication" || exit
}

create_cluster() {
	echo "Create archive"
	mkdir $ARCHIVEDIR

	echo "Init db new cluster"
	$BINDIR/initdb "${N1DATA}"

	echo "Set up N1 to be master"
	(
	cat <<-PGCONF
		port = ${N1PORT}
		wal_level = hot_standby
		archive_mode = on
		archive_command = 'bash -c "if echo %f | grep -q history; then cp %p ${ARCHIVEDIR}; fi"'
		max_wal_senders = 2
		wal_keep_segments = 10
		hot_standby = on
	PGCONF
	) >> "${N1DATA}/postgresql.conf"
	(
	cat <<-HBACONF
		local   replication	 all								 trust
		host	replication	 all		 127.0.0.1/32			trust
	HBACONF
	) >> "${N1DATA}/pg_hba.conf"

	echo "Start master and create a db for our user"
	$BINDIR/pg_ctl -D "${N1DATA}" -l "${N1DATA}/pg.log" -w start
	$BINDIR/createdb || exit
	$BINDIR/psql -p $N1PORT -c "CREATE TABLE test (ts timestamptz)"
	verify_up_as_master $N1PORT

	echo "Create a new base backup for n2"
	sh -c "sleep 5; psql -p ${N1PORT} -c 'SELECT pg_switch_xlog()'" & # Switch required for streaming to complete
	$BINDIR/pg_basebackup -p ${N1PORT} -D ${N2DATA} --xlog=stream --checkpoint=fast -P -v || exit

	echo "Adjust N2 conf to be slave"
	sed -i "s/port = .*/port = ${N2PORT}/" ${N2DATA}/postgresql.conf
	make_recoveryconf "${N2DATA}" ${N1PORT}

	echo "Start n2"
	$BINDIR/pg_ctl -D "${N2DATA}" -l "${N2DATA}/pg.log" -w start
	echo "Verify n2 is up"
	verify_up $N2PORT

	echo "Create a new base backup for n3 from n2"
	sh -c "sleep 5; psql -p ${N1PORT} -c 'CHECKPOINT; SELECT pg_switch_xlog()'" & # Checkpoint and switch sometimes required for streaming to complete
	$BINDIR/pg_basebackup -p ${N2PORT} -D ${N3DATA} --xlog=stream --checkpoint=fast -P -v || exit

	echo "Adjust n3 conf to be slave of n2"
	sed -i "s/port = .*/port = ${N3PORT}/" ${N3DATA}/postgresql.conf
	make_recoveryconf "${N3DATA}" ${N2PORT}

	echo "Start n3"
	$BINDIR/pg_ctl -D "${N3DATA}" -l "${N3DATA}/pg.log" -w start
	echo "Verify n3 is up"
	verify_up $N3PORT
}

test_failover() {
	echo "Stop n1"
	$BINDIR/pg_ctl -D "${N1DATA}" -m fast stop

	echo "Promote n2"
	$BINDIR/pg_ctl -D "${N2DATA}" promote
	verify_up_as_master $N2PORT

	echo "Waiting and looking at replication status at n2"
	replication_status $N2PORT
	verify_up $N3PORT

	echo "Adjust n1 conf to be slave of n3"
	make_recoveryconf "${N1DATA}" ${N3PORT}

	echo "Start n1"
	$BINDIR/pg_ctl -D "${N1DATA}" -l "${N1DATA}/pg.log" -w start
	echo "Create data on master"
	verify_up_as_master $N2PORT
	verify_up $N1PORT

	echo "Replication status at n3"
	$BINDIR/psql -p $N2PORT -c "SELECT pg_switch_xlog()" # Required for replication to reconnect, otherwise gets stuck at record with zero length
	replication_status $N3PORT

	echo "State on n2 (master)"
	verify_up_as_master $N2PORT
	sleep 1
	echo "State on n3 (replica)"
	verify_up $N3PORT
	echo "State on n1 (cascading replica)"
	verify_up $N1PORT
}

test_self_replication() {
	$BINDIR/pg_ctl -D "${N3DATA}" -m fast stop
	make_recoveryconf "${N3DATA}" ${N3PORT}
	$BINDIR/pg_ctl -D "${N3DATA}" -l "${N3DATA}/pg.log" -w start
	replication_status $N3PORT
	verify_up_as_master $N3PORT		
}

cleanup
create_cluster
test_failover
test_self_replication