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

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;
 	}