Create a psql command \gset to store query results into psql variables.

This eases manipulation of query results in psql scripts.

Pavel Stehule, reviewed by Piyush Newe, Shigeru Hanada, and Tom Lane

9 files changed, 214 insertions, 15 deletions

diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c
index 1e9aa89089e..012cb75e52c 100644
--- a/src/bin/psql/command.c
+++ b/src/bin/psql/command.c
@@ -731,7 +731,7 @@ exec_command(const char *cmd,
 		free(fname);
 	}
 
-	/* \g [filename] means send query, optionally with output to file/pipe */
+	/* \g [filename] -- send query, optionally with output to file/pipe */
 	else if (strcmp(cmd, "g") == 0)
 	{
 		char	   *fname = psql_scan_slash_option(scan_state,
@@ -748,6 +748,22 @@ exec_command(const char *cmd,
 		status = PSQL_CMD_SEND;
 	}
 
+	/* \gset [prefix] -- send query and store result into variables */
+	else if (strcmp(cmd, "gset") == 0)
+	{
+		char	   *prefix = psql_scan_slash_option(scan_state,
+													OT_NORMAL, NULL, false);
+
+		if (prefix)
+			pset.gset_prefix = prefix;
+		else
+		{
+			/* we must set a non-NULL prefix to trigger storing */
+			pset.gset_prefix = pg_strdup("");
+		}
+		status = PSQL_CMD_SEND;
+	}
+
 	/* help */
 	else if (strcmp(cmd, "h") == 0 || strcmp(cmd, "help") == 0)
 	{
diff --git a/src/bin/psql/common.c b/src/bin/psql/common.c
index ab517906fca..a8aa1a2df1b 100644
--- a/src/bin/psql/common.c
+++ b/src/bin/psql/common.c
@@ -43,7 +43,7 @@ pg_strdup(const char *string)
 	if (!string)
 	{
 		psql_error("%s: pg_strdup: cannot duplicate null pointer (internal error)\n",
-				pset.progname);
+				   pset.progname);
 		exit(EXIT_FAILURE);
 	}
 	tmp = strdup(string);
@@ -616,6 +616,65 @@ PrintQueryTuples(const PGresult *results)
 
 
 /*
+ * StoreQueryTuple: assuming query result is OK, save data into variables
+ *
+ * Returns true if successful, false otherwise.
+ */
+static bool
+StoreQueryTuple(const PGresult *result)
+{
+	bool		success = true;
+
+	if (PQntuples(result) < 1)
+	{
+		psql_error("no rows returned for \\gset\n");
+		success = false;
+	}
+	else if (PQntuples(result) > 1)
+	{
+		psql_error("more than one row returned for \\gset\n");
+		success = false;
+	}
+	else
+	{
+		int			i;
+
+		for (i = 0; i < PQnfields(result); i++)
+		{
+			char	   *colname = PQfname(result, i);
+			char	   *varname;
+			char	   *value;
+
+			/* concate prefix and column name */
+			varname = pg_malloc(strlen(pset.gset_prefix) + strlen(colname) + 1);
+			strcpy(varname, pset.gset_prefix);
+			strcat(varname, colname);
+
+			if (!PQgetisnull(result, 0, i))
+				value = PQgetvalue(result, 0, i);
+			else
+			{
+				/* for NULL value, unset rather than set the variable */
+				value = NULL;
+			}
+
+			if (!SetVariable(pset.vars, varname, value))
+			{
+				psql_error("could not set variable \"%s\"\n", varname);
+				free(varname);
+				success = false;
+				break;
+			}
+
+			free(varname);
+		}
+	}
+
+	return success;
+}
+
+
+/*
  * ProcessResult: utility function for use by SendQuery() only
  *
  * When our command string contained a COPY FROM STDIN or COPY TO STDOUT,
@@ -752,7 +811,7 @@ PrintQueryStatus(PGresult *results)
 
 
 /*
- * PrintQueryResults: print out query results as required
+ * PrintQueryResults: print out (or store) query results as required
  *
  * Note: Utility function for use by SendQuery() only.
  *
@@ -770,8 +829,11 @@ PrintQueryResults(PGresult *results)
 	switch (PQresultStatus(results))
 	{
 		case PGRES_TUPLES_OK:
-			/* print the data ... */
-			success = PrintQueryTuples(results);
+			/* store or print the data ... */
+			if (pset.gset_prefix)
+				success = StoreQueryTuple(results);
+			else
+				success = PrintQueryTuples(results);
 			/* if it's INSERT/UPDATE/DELETE RETURNING, also print status */
 			cmdstatus = PQcmdStatus(results);
 			if (strncmp(cmdstatus, "INSERT", 6) == 0 ||
@@ -898,7 +960,7 @@ SendQuery(const char *query)
 		if (on_error_rollback_warning == false && pset.sversion < 80000)
 		{
 			psql_error("The server (version %d.%d) does not support savepoints for ON_ERROR_ROLLBACK.\n",
-					pset.sversion / 10000, (pset.sversion / 100) % 100);
+					   pset.sversion / 10000, (pset.sversion / 100) % 100);
 			on_error_rollback_warning = true;
 		}
 		else
@@ -1046,6 +1108,13 @@ sendquery_cleanup:
 		pset.gfname = NULL;
 	}
 
+	/* reset \gset trigger */
+	if (pset.gset_prefix)
+	{
+		free(pset.gset_prefix);
+		pset.gset_prefix = NULL;
+	}
+
 	return OK;
 }
 
@@ -1072,6 +1141,7 @@ ExecQueryUsingCursor(const char *query, double *elapsed_msec)
 	bool		started_txn = false;
 	bool		did_pager = false;
 	int			ntuples;
+	int			fetch_count;
 	char		fetch_cmd[64];
 	instr_time	before,
 				after;
@@ -1119,9 +1189,18 @@ ExecQueryUsingCursor(const char *query, double *elapsed_msec)
 		*elapsed_msec += INSTR_TIME_GET_MILLISEC(after);
 	}
 
+	/*
+	 * In \gset mode, we force the fetch count to be 2, so that we will throw
+	 * the appropriate error if the query returns more than one row.
+	 */
+	if (pset.gset_prefix)
+		fetch_count = 2;
+	else
+		fetch_count = pset.fetch_count;
+
 	snprintf(fetch_cmd, sizeof(fetch_cmd),
 			 "FETCH FORWARD %d FROM _psql_cursor",
-			 pset.fetch_count);
+			 fetch_count);
 
 	/* prepare to write output to \g argument, if any */
 	if (pset.gfname)
@@ -1147,7 +1226,7 @@ ExecQueryUsingCursor(const char *query, double *elapsed_msec)
 		if (pset.timing)
 			INSTR_TIME_SET_CURRENT(before);
 
-		/* get FETCH_COUNT tuples at a time */
+		/* get fetch_count tuples at a time */
 		results = PQexec(pset.db, fetch_cmd);
 
 		if (pset.timing)
@@ -1174,9 +1253,17 @@ ExecQueryUsingCursor(const char *query, double *elapsed_msec)
 			break;
 		}
 
+		if (pset.gset_prefix)
+		{
+			/* StoreQueryTuple will complain if not exactly one row */
+			OK = StoreQueryTuple(results);
+			PQclear(results);
+			break;
+		}
+
 		ntuples = PQntuples(results);
 
-		if (ntuples < pset.fetch_count)
+		if (ntuples < fetch_count)
 		{
 			/* this is the last result set, so allow footer decoration */
 			my_popt.topt.stop_table = true;
@@ -1214,7 +1301,7 @@ ExecQueryUsingCursor(const char *query, double *elapsed_msec)
 		 * writing things to the stream, we presume $PAGER has disappeared and
 		 * stop bothering to pull down more data.
 		 */
-		if (ntuples < pset.fetch_count || cancel_pressed || flush_error ||
+		if (ntuples < fetch_count || cancel_pressed || flush_error ||
 			ferror(pset.queryFout))
 			break;
 	}
diff --git a/src/bin/psql/help.c b/src/bin/psql/help.c
index fd7effa87b3..43cb550bd24 100644
--- a/src/bin/psql/help.c
+++ b/src/bin/psql/help.c
@@ -165,13 +165,14 @@ slashUsage(unsigned short int pager)
 
 	currdb = PQdb(pset.db);
 
-	output = PageOutput(94, pager);
+	output = PageOutput(95, pager);
 
 	/* if you add/remove a line here, change the row count above */
 
 	fprintf(output, _("General\n"));
 	fprintf(output, _("  \\copyright             show PostgreSQL usage and distribution terms\n"));
 	fprintf(output, _("  \\g [FILE] or ;         execute query (and send results to file or |pipe)\n"));
+	fprintf(output, _("  \\gset [PREFIX]         execute query and store results in psql variables\n"));
 	fprintf(output, _("  \\h [NAME]              help on syntax of SQL commands, * for all commands\n"));
 	fprintf(output, _("  \\q                     quit psql\n"));
 	fprintf(output, "\n");
@@ -261,7 +262,7 @@ slashUsage(unsigned short int pager)
 				currdb);
 	else
 		fprintf(output, _("  \\c[onnect] [DBNAME|- USER|- HOST|- PORT|-]\n"
-						  "                         connect to new database (currently no connection)\n")),
+						  "                         connect to new database (currently no connection)\n"));
 	fprintf(output, _("  \\encoding [ENCODING]   show or set client encoding\n"));
 	fprintf(output, _("  \\password [USERNAME]   securely change the password for a user\n"));
 	fprintf(output, _("  \\conninfo              display information about current connection\n"));
diff --git a/src/bin/psql/settings.h b/src/bin/psql/settings.h
index a80fb990c16..e78aa9a4999 100644
--- a/src/bin/psql/settings.h
+++ b/src/bin/psql/settings.h
@@ -73,6 +73,7 @@ typedef struct _psqlSettings
 	printQueryOpt popt;
 
 	char	   *gfname;			/* one-shot file output argument for \g */
+	char	   *gset_prefix;	/* one-shot prefix argument for \gset */
 
 	bool		notty;			/* stdin or stdout is not a tty (as determined
 								 * on startup) */
diff --git a/src/bin/psql/tab-complete.c b/src/bin/psql/tab-complete.c
index f7d84669f97..09396ca5900 100644
--- a/src/bin/psql/tab-complete.c
+++ b/src/bin/psql/tab-complete.c
@@ -856,7 +856,7 @@ psql_completion(char *text, int start, int end)
 		"\\dF", "\\dFd", "\\dFp", "\\dFt", "\\dg", "\\di", "\\dl", "\\dL",
 		"\\dn", "\\do", "\\dp", "\\drds", "\\ds", "\\dS", "\\dt", "\\dT", "\\dv", "\\du",
 		"\\e", "\\echo", "\\ef", "\\encoding",
-		"\\f", "\\g", "\\h", "\\help", "\\H", "\\i", "\\ir", "\\l",
+		"\\f", "\\g", "\\gset", "\\h", "\\help", "\\H", "\\i", "\\ir", "\\l",
 		"\\lo_import", "\\lo_export", "\\lo_list", "\\lo_unlink",
 		"\\o", "\\p", "\\password", "\\prompt", "\\pset", "\\q", "\\qecho", "\\r",
 		"\\set", "\\sf", "\\t", "\\T",
diff --git a/src/test/regress/expected/psql.out b/src/test/regress/expected/psql.out
new file mode 100644
index 00000000000..7dd5c7b51b9
--- /dev/null
+++ b/src/test/regress/expected/psql.out
@@ -0,0 +1,54 @@
+--
+-- Tests for psql features that aren't closely connected to any
+-- specific server features
+--
+-- \gset
+select 10 as test01, 20 as test02, 'Hello' as test03 \gset pref01_
+\echo :pref01_test01 :pref01_test02 :pref01_test03
+10 20 Hello
+-- should fail: bad variable name
+select 10 as "bad name"
+\gset
+could not set variable "bad name"
+-- multiple backslash commands in one line
+select 1 as x, 2 as y \gset pref01_ \\ \echo :pref01_x
+1
+select 3 as x, 4 as y \gset pref01_ \echo :pref01_x \echo :pref01_y
+3
+4
+select 5 as x, 6 as y \gset pref01_ \\ \g \echo :pref01_x :pref01_y
+ x | y 
+---+---
+ 5 | 6
+(1 row)
+
+5 6
+select 7 as x, 8 as y \g \gset pref01_ \echo :pref01_x :pref01_y
+ x | y 
+---+---
+ 7 | 8
+(1 row)
+
+7 8
+-- NULL should unset the variable
+\set var2 xyz
+select 1 as var1, NULL as var2, 3 as var3 \gset
+\echo :var1 :var2 :var3
+1 :var2 3
+-- \gset requires just one tuple
+select 10 as test01, 20 as test02 from generate_series(1,3) \gset
+more than one row returned for \gset
+select 10 as test01, 20 as test02 from generate_series(1,0) \gset
+no rows returned for \gset
+-- \gset should work in FETCH_COUNT mode too
+\set FETCH_COUNT 1
+select 1 as x, 2 as y \gset pref01_ \\ \echo :pref01_x
+1
+select 3 as x, 4 as y \gset pref01_ \echo :pref01_x \echo :pref01_y
+3
+4
+select 10 as test01, 20 as test02 from generate_series(1,3) \gset
+more than one row returned for \gset
+select 10 as test01, 20 as test02 from generate_series(1,0) \gset
+no rows returned for \gset
+\unset FETCH_COUNT
diff --git a/src/test/regress/parallel_schedule b/src/test/regress/parallel_schedule
index bdcf3a6a559..d3def07f92c 100644
--- a/src/test/regress/parallel_schedule
+++ b/src/test/regress/parallel_schedule
@@ -88,7 +88,7 @@ test: privileges security_label collate
 # ----------
 # Another group of parallel tests
 # ----------
-test: misc alter_generic
+test: alter_generic misc psql
 
 # rules cannot run concurrently with any test that creates a view
 test: rules
diff --git a/src/test/regress/serial_schedule b/src/test/regress/serial_schedule
index c7c2ed0f6a0..7059fca092b 100644
--- a/src/test/regress/serial_schedule
+++ b/src/test/regress/serial_schedule
@@ -95,8 +95,9 @@ test: prepared_xacts
 test: privileges
 test: security_label
 test: collate
-test: misc
 test: alter_generic
+test: misc
+test: psql
 test: rules
 test: event_trigger
 test: select_views
diff --git a/src/test/regress/sql/psql.sql b/src/test/regress/sql/psql.sql
new file mode 100644
index 00000000000..50ee02a7f79
--- /dev/null
+++ b/src/test/regress/sql/psql.sql
@@ -0,0 +1,39 @@
+--
+-- Tests for psql features that aren't closely connected to any
+-- specific server features
+--
+
+-- \gset
+
+select 10 as test01, 20 as test02, 'Hello' as test03 \gset pref01_
+
+\echo :pref01_test01 :pref01_test02 :pref01_test03
+
+-- should fail: bad variable name
+select 10 as "bad name"
+\gset
+
+-- multiple backslash commands in one line
+select 1 as x, 2 as y \gset pref01_ \\ \echo :pref01_x
+select 3 as x, 4 as y \gset pref01_ \echo :pref01_x \echo :pref01_y
+select 5 as x, 6 as y \gset pref01_ \\ \g \echo :pref01_x :pref01_y
+select 7 as x, 8 as y \g \gset pref01_ \echo :pref01_x :pref01_y
+
+-- NULL should unset the variable
+\set var2 xyz
+select 1 as var1, NULL as var2, 3 as var3 \gset
+\echo :var1 :var2 :var3
+
+-- \gset requires just one tuple
+select 10 as test01, 20 as test02 from generate_series(1,3) \gset
+select 10 as test01, 20 as test02 from generate_series(1,0) \gset
+
+-- \gset should work in FETCH_COUNT mode too
+\set FETCH_COUNT 1
+
+select 1 as x, 2 as y \gset pref01_ \\ \echo :pref01_x
+select 3 as x, 4 as y \gset pref01_ \echo :pref01_x \echo :pref01_y
+select 10 as test01, 20 as test02 from generate_series(1,3) \gset
+select 10 as test01, 20 as test02 from generate_series(1,0) \gset
+
+\unset FETCH_COUNT