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

5 files changed, 118 insertions, 13 deletions

diff --git a/src/bin/psql/command.c b/src/bin/psql/command.c
index 1e9aa89089..012cb75e52 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 ab517906fc..a8aa1a2df1 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 fd7effa87b..43cb550bd2 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 a80fb990c1..e78aa9a499 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 f7d84669f9..09396ca590 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",