psql: add an optional execution-count limit to \watch.

\watch can now be told to stop after N executions of the query.

With the idea that we might want to add more options to \watch
in future, this patch generalizes the command's syntax to a list
of name=value options, with the interval allowed to omit the name
for backwards compatibility.

Andrey Borodin, reviewed by Kyotaro Horiguchi, Nathan Bossart,
Michael Paquier, Yugo Nagata, and myself

Discussion: https://postgr.es/m/CAAhFRxiZ2-n_L1ErMm9AZjgmUK=qS6VHb+0SaMn8sqqbhF7How@mail.gmail.com

6 files changed, 135 insertions, 32 deletions

diff --git a/doc/src/sgml/ref/psql-ref.sgml b/doc/src/sgml/ref/psql-ref.sgml
index 29bbec2188..53875afbf0 100644
--- a/doc/src/sgml/ref/psql-ref.sgml
+++ b/doc/src/sgml/ref/psql-ref.sgml
@@ -3551,12 +3551,16 @@ testdb=&gt; <userinput>\setenv LESS -imx4F</userinput>
 
 
       <varlistentry id="app-psql-meta-command-watch">
-        <term><literal>\watch [ <replaceable class="parameter">seconds</replaceable> ]</literal></term>
+        <term><literal>\watch [ i[nterval]=<replaceable class="parameter">seconds</replaceable> ] [ c[ount]=<replaceable class="parameter">times</replaceable> ] [ <replaceable class="parameter">seconds</replaceable> ]</literal></term>
         <listitem>
         <para>
         Repeatedly execute the current query buffer (as <literal>\g</literal> does)
-        until interrupted or the query fails.  Wait the specified number of
-        seconds (default 2) between executions.  Each query result is
+        until interrupted, or the query fails, or the execution count limit
+        (if given) is reached.  Wait the specified number of
+        seconds (default 2) between executions.  For backwards compatibility,
+        <replaceable class="parameter">seconds</replaceable> can be specified
+        with or without an <literal>interval=</literal> prefix.
+        Each query result is
         displayed with a header that includes the <literal>\pset title</literal>
         string (if any), the time as of query start, and the delay interval.
         </para>
diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c
index d7731234b6..e8f583cac2 100644
--- a/src/bin/psql/command.c
+++ b/src/bin/psql/command.c
@@ -162,7 +162,7 @@ static bool do_connect(enum trivalue reuse_previous_specification,
 static bool do_edit(const char *filename_arg, PQExpBuffer query_buf,
 					int lineno, bool discard_on_quit, bool *edited);
 static bool do_shell(const char *command);
-static bool do_watch(PQExpBuffer query_buf, double sleep);
+static bool do_watch(PQExpBuffer query_buf, double sleep, int iter);
 static bool lookup_object_oid(EditableObjectType obj_type, const char *desc,
 							  Oid *obj_oid);
 static bool get_create_object_cmd(EditableObjectType obj_type, Oid oid,
@@ -2759,7 +2759,8 @@ exec_command_write(PsqlScanState scan_state, bool active_branch,
 }
 
 /*
- * \watch -- execute a query every N seconds
+ * \watch -- execute a query every N seconds.
+ * Optionally, stop after M iterations.
  */
 static backslashResult
 exec_command_watch(PsqlScanState scan_state, bool active_branch,
@@ -2769,32 +2770,109 @@ exec_command_watch(PsqlScanState scan_state, bool active_branch,
 
 	if (active_branch)
 	{
-		char	   *opt = psql_scan_slash_option(scan_state,
-												 OT_NORMAL, NULL, true);
+		bool		have_sleep = false;
+		bool		have_iter = false;
 		double		sleep = 2;
+		int			iter = 0;
 
-		/* Convert optional sleep-length argument */
-		if (opt)
+		/*
+		 * Parse arguments.  We allow either an unlabeled interval or
+		 * "name=value", where name is from the set ('i', 'interval', 'c',
+		 * 'count').
+		 */
+		while (success)
 		{
+			char	   *opt = psql_scan_slash_option(scan_state,
+													 OT_NORMAL, NULL, true);
+			char	   *valptr;
 			char	   *opt_end;
 
-			errno = 0;
-			sleep = strtod(opt, &opt_end);
-			if (sleep < 0 || *opt_end || errno == ERANGE)
+			if (!opt)
+				break;			/* no more arguments */
+
+			valptr = strchr(opt, '=');
+			if (valptr)
 			{
-				pg_log_error("\\watch: incorrect interval value '%s'", opt);
-				free(opt);
-				resetPQExpBuffer(query_buf);
-				psql_scan_reset(scan_state);
-				return PSQL_CMD_ERROR;
+				/* Labeled argument */
+				valptr++;
+				if (strncmp("i=", opt, strlen("i=")) == 0 ||
+					strncmp("interval=", opt, strlen("interval=")) == 0)
+				{
+					if (have_sleep)
+					{
+						pg_log_error("\\watch: interval value is specified more than once");
+						success = false;
+					}
+					else
+					{
+						have_sleep = true;
+						errno = 0;
+						sleep = strtod(valptr, &opt_end);
+						if (sleep < 0 || *opt_end || errno == ERANGE)
+						{
+							pg_log_error("\\watch: incorrect interval value \"%s\"", valptr);
+							success = false;
+						}
+					}
+				}
+				else if (strncmp("c=", opt, strlen("c=")) == 0 ||
+						 strncmp("count=", opt, strlen("count=")) == 0)
+				{
+					if (have_iter)
+					{
+						pg_log_error("\\watch: iteration count is specified more than once");
+						success = false;
+					}
+					else
+					{
+						have_iter = true;
+						errno = 0;
+						iter = strtoint(valptr, &opt_end, 10);
+						if (iter <= 0 || *opt_end || errno == ERANGE)
+						{
+							pg_log_error("\\watch: incorrect iteration count \"%s\"", valptr);
+							success = false;
+						}
+					}
+				}
+				else
+				{
+					pg_log_error("\\watch: unrecognized parameter \"%s\"", opt);
+					success = false;
+				}
+			}
+			else
+			{
+				/* Unlabeled argument: take it as interval */
+				if (have_sleep)
+				{
+					pg_log_error("\\watch: interval value is specified more than once");
+					success = false;
+				}
+				else
+				{
+					have_sleep = true;
+					errno = 0;
+					sleep = strtod(opt, &opt_end);
+					if (sleep < 0 || *opt_end || errno == ERANGE)
+					{
+						pg_log_error("\\watch: incorrect interval value \"%s\"", opt);
+						success = false;
+					}
+				}
 			}
+
 			free(opt);
 		}
 
-		/* If query_buf is empty, recall and execute previous query */
-		(void) copy_previous_query(query_buf, previous_buf);
+		/* If we parsed arguments successfully, do the command */
+		if (success)
+		{
+			/* If query_buf is empty, recall and execute previous query */
+			(void) copy_previous_query(query_buf, previous_buf);
 
-		success = do_watch(query_buf, sleep);
+			success = do_watch(query_buf, sleep, iter);
+		}
 
 		/* Reset the query buffer as though for \r */
 		resetPQExpBuffer(query_buf);
@@ -5071,7 +5149,7 @@ do_shell(const char *command)
  * onto a bunch of exec_command's variables to silence stupider compilers.
  */
 static bool
-do_watch(PQExpBuffer query_buf, double sleep)
+do_watch(PQExpBuffer query_buf, double sleep, int iter)
 {
 	long		sleep_ms = (long) (sleep * 1000);
 	printQueryOpt myopt = pset.popt;
@@ -5204,6 +5282,10 @@ do_watch(PQExpBuffer query_buf, double sleep)
 		if (res <= 0)
 			break;
 
+		/* If we have iteration count, check that it's not exceeded yet */
+		if (iter && (--iter <= 0))
+			break;
+
 		if (pagerpipe && ferror(pagerpipe))
 			break;
 
diff --git a/src/bin/psql/help.c b/src/bin/psql/help.c
index 48fd51592a..ecfb3c099b 100644
--- a/src/bin/psql/help.c
+++ b/src/bin/psql/help.c
@@ -200,7 +200,7 @@ slashUsage(unsigned short int pager)
 	HELP0("  \\gset [PREFIX]         execute query and store result in psql variables\n");
 	HELP0("  \\gx [(OPTIONS)] [FILE] as \\g, but forces expanded output mode\n");
 	HELP0("  \\q                     quit psql\n");
-	HELP0("  \\watch [SEC]           execute query every SEC seconds\n");
+	HELP0("  \\watch [[i=]SEC] [c=N] execute query every SEC seconds, up to N times\n");
 	HELP0("\n");
 
 	HELP0("Help\n");
diff --git a/src/bin/psql/t/001_basic.pl b/src/bin/psql/t/001_basic.pl
index 64ce012062..56b1e3e4a6 100644
--- a/src/bin/psql/t/001_basic.pl
+++ b/src/bin/psql/t/001_basic.pl
@@ -350,21 +350,38 @@ psql_like(
 	'\copy from with DEFAULT'
 );
 
+# Check \watch
+psql_like(
+	$node,
+	'SELECT 1 \watch c=3 i=0.01',
+	qr/1\n1\n1/,
+	'\watch with 3 iterations');
+
 # Check \watch errors
 psql_fails_like(
 	$node,
-	'SELECT 1;\watch -10',
-	qr/incorrect interval value '-10'/,
+	'SELECT 1 \watch -10',
+	qr/incorrect interval value "-10"/,
 	'\watch, negative interval');
 psql_fails_like(
 	$node,
-	'SELECT 1;\watch 10ab',
-	qr/incorrect interval value '10ab'/,
-	'\watch incorrect interval');
+	'SELECT 1 \watch 10ab',
+	qr/incorrect interval value "10ab"/,
+	'\watch, incorrect interval');
+psql_fails_like(
+	$node,
+	'SELECT 1 \watch 10e400',
+	qr/incorrect interval value "10e400"/,
+	'\watch, out-of-range interval');
+psql_fails_like(
+	$node,
+	'SELECT 1 \watch 1 1',
+	qr/interval value is specified more than once/,
+	'\watch, interval value is specified more than once');
 psql_fails_like(
 	$node,
-	'SELECT 1;\watch 10e400',
-	qr/incorrect interval value '10e400'/,
-	'\watch out-of-range interval');
+	'SELECT 1 \watch c=1 c=1',
+	qr/iteration count is specified more than once/,
+	'\watch, iteration count is specified more than once');
 
 done_testing();
diff --git a/src/test/regress/expected/psql.out b/src/test/regress/expected/psql.out
index c00e28361c..956e475447 100644
--- a/src/test/regress/expected/psql.out
+++ b/src/test/regress/expected/psql.out
@@ -4536,7 +4536,7 @@ invalid command \lo
 	\timing arg1
 	\unset arg1
 	\w arg1
-	\watch arg1
+	\watch arg1 arg2
 	\x arg1
 	-- \else here is eaten as part of OT_FILEPIPE argument
 	\w |/no/such/file \else
diff --git a/src/test/regress/sql/psql.sql b/src/test/regress/sql/psql.sql
index 961783d6ea..630f638f02 100644
--- a/src/test/regress/sql/psql.sql
+++ b/src/test/regress/sql/psql.sql
@@ -1022,7 +1022,7 @@ select \if false \\ (bogus \else \\ 42 \endif \\ forty_two;
 	\timing arg1
 	\unset arg1
 	\w arg1
-	\watch arg1
+	\watch arg1 arg2
 	\x arg1
 	-- \else here is eaten as part of OT_FILEPIPE argument
 	\w |/no/such/file \else