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

doc/src/sgml/ref/psql-ref.sgml

@@ -1620,6 +1620,46 @@ Tue Oct 26 21:40:57 CEST 1999
         </listitem>
       </varlistentry>
 
+      <varlistentry>
+        <term><literal>\gset</literal> [ <replaceable class="parameter">prefix</replaceable> ]</term>
+
+        <listitem>
+        <para>
+         Sends the current query input buffer to the server and stores the
+         query's output into <application>psql</> variables (see <xref
+         linkend="APP-PSQL-variables" endterm="APP-PSQL-variables-title">).
+         The query to be executed must return exactly one row.  Each column of
+         the row is stored into a separate variable, named the same as the
+         column.  For example:
+<programlisting>
+=&gt; <userinput>SELECT 'hello' AS var1, 10 AS var2</userinput>
+-&gt; <userinput>\gset</userinput>
+=&gt; <userinput>\echo :var1 :var2</userinput>
+hello 10
+</programlisting>
+        </para>
+        <para>
+         If you specify a <replaceable class="parameter">prefix</replaceable>,
+         that string is prepended to the query's column names to create the
+         variable names to use:
+<programlisting>
+=&gt; <userinput>SELECT 'hello' AS var1, 10 AS var2</userinput>
+-&gt; <userinput>\gset result_</userinput>
+=&gt; <userinput>\echo :result_var1 :result_var2</userinput>
+hello 10
+</programlisting>
+        </para>
+        <para>
+         If a column result is NULL, the corresponding variable is unset
+         rather than being set.
+        </para>
+        <para>
+         If the query fails or does not return one row,
+         no variables are changed.
+        </para>
+        </listitem>
+      </varlistentry>
+
       <varlistentry>
         <term><literal>\h</literal> or <literal>\help</literal> <literal>[ <replaceable class="parameter">command</replaceable> ]</literal></term>
         <listitem>

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)
 	{

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);
@@ -615,6 +615,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
  *
@@ -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;
 	}

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"));

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) */

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",

src/test/regress/expected/psql.out

+--
+-- 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

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

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

src/test/regress/sql/psql.sql

+--
+-- 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