From 49648b04fc381c40e2f9911838f57cbd3735903f Mon Sep 17 00:00:00 2001
From: Benjamin Renard <brenard@easter-eggs.com>
Date: Mon, 3 Jun 2024 15:24:38 +0200
Subject: [PATCH] Add -e parameter to allow to specify expected replication
 state

---
 check_pg_streaming_replication | 19 ++++++++++++++-----
 1 file changed, 14 insertions(+), 5 deletions(-)

diff --git a/check_pg_streaming_replication b/check_pg_streaming_replication
index 680aee7..a6423c4 100755
--- a/check_pg_streaming_replication
+++ b/check_pg_streaming_replication
@@ -41,10 +41,13 @@ PG_DB=""
 CHECK_CUR_MASTER_LSN=1
 REPLAY_WARNING_DELAY=3
 REPLAY_CRITICAL_DELAY=5
+EXPECTED_SYNC_STATE=sync
 
 DEBUG=0
 
 function usage () {
+	ERROR="$1"
+	[ -n "$ERROR" ] && echo -e "$ERROR\n"
 	cat << EOF
 Usage: $0 [-d] [-h] [options]
 	-u pg_user		Specify local Postgres user (Default: try to auto-detect or use $DEFAULT_PG_USER)
@@ -64,13 +67,14 @@ Usage: $0 [-d] [-h] [options]
 				of the last received LSN (Default: $CHECK_CUR_MASTER_LSN)
 	-w replay_warn_delay	Specify the replay warning delay in second (Default: $REPLAY_WARNING_DELAY)
 	-c replay_crit_delay	Specify the replay critical delay in second (Default: $REPLAY_CRITICAL_DELAY)
+	-e expected_sync_state	The expected replication state ('sync' or 'async', default: $EXPECTED_SYNC_STATE)
 	-d			Debug mode
 	-h 			Show this message
 EOF
-	exit 0
+	[ -n "$ERROR" ] && exit 1 || exit 0
 }
 
-while getopts "hu:b:B:V:m:r:U:p:D:C:w:c:d" OPTION
+while getopts "hu:b:B:V:m:r:U:p:D:C:w:c:e:d" OPTION
 do
 	case $OPTION in
 		u)
@@ -109,6 +113,11 @@ do
 		c)
 			REPLAY_CRITICAL_DELAY=$OPTARG
 		;;
+		e)
+			[ "$OPTARG" != "sync" -a "$OPTARG" != "async" ] && \
+				usage "Invalid expected replication state '$OPTARG'. Possible values: sync or async."
+			EXPECTED_SYNC_STATE=$OPTARG
+		;;
 		d)
 			DEBUG=1
 		;;
@@ -255,7 +264,7 @@ RECOVERY_MODE=0
 if [ -f $RECOVERY_CONF ]
 then
 	debug "File recovery.conf found. Hot-standby mode."
-	
+
 	# Check recovery mode
 	if [ $RECOVERY_MODE -ne 1 ]
 	then
@@ -354,9 +363,9 @@ then
 
 	M_CUR_SYNC_STATE=$( echo "$M_CUR_REPL_STATE_INFO"|cut -d'|' -f2 )
 	debug "Master current sync state: $M_CUR_SYNC_STATE"
-	if [ "$M_CUR_SYNC_STATE" != "sync" ]
+	if [ "$M_CUR_SYNC_STATE" != "$EXPECTED_SYNC_STATE" ]
 	then
-		echo "CRITICAL: this host is not synchronized according to master host (current sync state = '$M_CUR_SYNC_STATE')"
+		echo "CRITICAL: unexpected replication state '$M_CUR_SYNC_STATE' (expected state = '$EXPECTED_SYNC_STATE')"
 		exit 2
 	fi