Allow \watch queries to stop on minimum rows returned

When running a repeat query with \watch in psql, it can be
helpful to be able to stop the watch process when the query
no longer returns the expected amount of rows.  An example
would be to watch for the presence of a certain event in
pg_stat_activity and stopping when the event is no longer
present, or to watch an index creation and stop when the
index is created.

This adds a min_rows=MIN parameter to \watch which can be
set to a non-negative integer, and the watch query will
stop executing when it returns less than MIN rows.

Author: Greg Sabino Mullane <htamfids@gmail.com>
Reviewed-by: Michael Paquier <michael@paquier.xyz>
Reviewed-by: Daniel Gustafsson <daniel@yesql.se>
Discussion: https://postgr.es/m/CAKAnmmKStATuddYxP71L+p0DHtp9Rvjze3XRoy0Dyw67VQ45UA@mail.gmail.com

1 files changed, 14 insertions, 5 deletions

diff --git a/src/bin/psql/common.c b/src/bin/psql/common.c
index 10ad1f2538d..ede197bebeb 100644
--- a/src/bin/psql/common.c
+++ b/src/bin/psql/common.c
@@ -36,6 +36,7 @@ static int	ExecQueryAndProcessResults(const char *query,
 									   double *elapsed_msec,
 									   bool *svpt_gone_p,
 									   bool is_watch,
+									   int min_rows,
 									   const printQueryOpt *opt,
 									   FILE *printQueryFout);
 static bool command_no_begin(const char *query);
@@ -632,7 +633,7 @@ PSQLexec(const char *query)
  * e.g., because of the interrupt, -1 on error.
  */
 int
-PSQLexecWatch(const char *query, const printQueryOpt *opt, FILE *printQueryFout)
+PSQLexecWatch(const char *query, const printQueryOpt *opt, FILE *printQueryFout, int min_rows)
 {
 	bool		timing = pset.timing;
 	double		elapsed_msec = 0;
@@ -646,7 +647,7 @@ PSQLexecWatch(const char *query, const printQueryOpt *opt, FILE *printQueryFout)
 
 	SetCancelConn(pset.db);
 
-	res = ExecQueryAndProcessResults(query, &elapsed_msec, NULL, true, opt, printQueryFout);
+	res = ExecQueryAndProcessResults(query, &elapsed_msec, NULL, true, min_rows, opt, printQueryFout);
 
 	ResetCancelConn();
 
@@ -1134,7 +1135,7 @@ SendQuery(const char *query)
 			 pset.crosstab_flag || !is_select_command(query))
 	{
 		/* Default fetch-it-all-and-print mode */
-		OK = (ExecQueryAndProcessResults(query, &elapsed_msec, &svpt_gone, false, NULL, NULL) > 0);
+		OK = (ExecQueryAndProcessResults(query, &elapsed_msec, &svpt_gone, false, 0, NULL, NULL) > 0);
 	}
 	else
 	{
@@ -1415,11 +1416,12 @@ DescribeQuery(const char *query, double *elapsed_msec)
 static int
 ExecQueryAndProcessResults(const char *query,
 						   double *elapsed_msec, bool *svpt_gone_p,
-						   bool is_watch,
+						   bool is_watch, int min_rows,
 						   const printQueryOpt *opt, FILE *printQueryFout)
 {
 	bool		timing = pset.timing;
 	bool		success;
+	bool		return_early = false;
 	instr_time	before,
 				after;
 	PGresult   *result;
@@ -1461,6 +1463,10 @@ ExecQueryAndProcessResults(const char *query,
 
 	/* first result */
 	result = PQgetResult(pset.db);
+	if (min_rows > 0 && PQntuples(result) < min_rows)
+	{
+		return_early = true;
+	}
 
 	while (result != NULL)
 	{
@@ -1683,7 +1689,10 @@ ExecQueryAndProcessResults(const char *query,
 	if (!CheckConnection())
 		return -1;
 
-	return cancel_pressed ? 0 : success ? 1 : -1;
+	if (cancel_pressed || return_early)
+		return 0;
+
+	return success ? 1 : -1;
 }