pgbench-aset-3.patch

diff --git a/doc/src/sgml/ref/pgbench.sgml b/doc/src/sgml/ref/pgbench.sgml
index e3b73a4cf5..52949d9a35 100644
--- a/doc/src/sgml/ref/pgbench.sgml
+++ b/doc/src/sgml/ref/pgbench.sgml
@@ -966,18 +966,29 @@ pgbench <optional> <replaceable>options</replaceable> </optional> <replaceable>d
    <varlistentry id='pgbench-metacommand-gset'>
     <term>
      <literal>\gset [<replaceable>prefix</replaceable>]</literal>
+     <literal>\aset [<replaceable>prefix</replaceable>]</literal>
     </term>
 
     <listitem>
      <para>
-      This command may be used to end SQL queries, taking the place of the
+      These commands may be used to end SQL queries, taking the place of the
       terminating semicolon (<literal>;</literal>).
      </para>
 
      <para>
-      When this command is used, the preceding SQL query is expected to
-      return one row, the columns of which are stored into variables named after
-      column names, and prefixed with <replaceable>prefix</replaceable> if provided.
+      When the <literal>\gset</literal> command is used, the preceding SQL query is
+      expected to return one row, the columns of which are stored into variables
+      named after column names, and prefixed with <replaceable>prefix</replaceable>
+      if provided.
+     </para>
+
+     <para>
+      When the <literal>\aset</literal> command is used, all combined SQL queries
+      (separated by <literal>\;</literal>) have their columns stored into variables
+      named after column names, and prefixed with <replaceable>prefix</replaceable>
+      if provided. If a query returns no row, no assignment is made and the variable
+      can be tested for existence to detect this. If a query returns more than one
+      row, the last value is kept.
      </para>
 
      <para>
@@ -986,6 +997,8 @@ pgbench <optional> <replaceable>options</replaceable> </optional> <replaceable>d
       <replaceable>p_two</replaceable> and <replaceable>p_three</replaceable>
       with integers from the third query.
       The result of the second query is discarded.
+      The result of the two last combined queries are stored in variables
+      <replaceable>four</replaceable> and <replaceable>five</replaceable>.
 <programlisting>
 UPDATE pgbench_accounts
   SET abalance = abalance + :delta
@@ -994,6 +1007,7 @@ UPDATE pgbench_accounts
 -- compound of two queries
 SELECT 1 \;
 SELECT 2 AS two, 3 AS three \gset p_
+SELECT 4 AS four \; SELECT 5 AS five \aset
 </programlisting>
      </para>
     </listitem>
diff --git a/src/bin/pgbench/pgbench.c b/src/bin/pgbench/pgbench.c
index 8b84658ccd..477b6c873b 100644
--- a/src/bin/pgbench/pgbench.c
+++ b/src/bin/pgbench/pgbench.c
@@ -461,6 +461,7 @@ typedef enum MetaCommand
 	META_SHELL,					/* \shell */
 	META_SLEEP,					/* \sleep */
 	META_GSET,					/* \gset */
+	META_ASET,					/* \aset */
 	META_IF,					/* \if */
 	META_ELIF,					/* \elif */
 	META_ELSE,					/* \else */
@@ -493,6 +494,7 @@ static const char *QUERYMODE[] = {"simple", "extended", "prepared"};
  * varprefix 	SQL commands terminated with \gset have this set
  *				to a non NULL value.  If nonempty, it's used to prefix the
  *				variable name that receives the value.
+ * aset			do gset on all possible queries of a combined query (\;).
  * expr			Parsed expression, if needed.
  * stats		Time spent in this command.
  */
@@ -505,6 +507,7 @@ typedef struct Command
 	int			argc;
 	char	   *argv[MAX_ARGS];
 	char	   *varprefix;
+	bool		aset;
 	PgBenchExpr *expr;
 	SimpleStats stats;
 } Command;
@@ -2479,6 +2482,8 @@ getMetaCommand(const char *cmd)
 		mc = META_ENDIF;
 	else if (pg_strcasecmp(cmd, "gset") == 0)
 		mc = META_GSET;
+	else if (pg_strcasecmp(cmd, "aset") == 0)
+		mc = META_ASET;
 	else
 		mc = META_NONE;
 	return mc;
@@ -2710,12 +2715,12 @@ sendCommand(CState *st, Command *command)
  * Process query response from the backend.
  *
  * If varprefix is not NULL, it's the variable name prefix where to store
- * the results of the *last* command.
+ * the results of the *last* command (gset) or *all* commands (aset).
  *
  * Returns true if everything is A-OK, false if any error occurs.
  */
 static bool
-readCommandResponse(CState *st, char *varprefix)
+readCommandResponse(CState *st, char *varprefix, bool is_aset)
 {
 	PGresult   *res;
 	PGresult   *next_res;
@@ -2735,7 +2740,7 @@ readCommandResponse(CState *st, char *varprefix)
 		{
 			case PGRES_COMMAND_OK:	/* non-SELECT commands */
 			case PGRES_EMPTY_QUERY: /* may be used for testing no-op overhead */
-				if (is_last && varprefix != NULL)
+				if (is_last && varprefix != NULL && !is_aset)
 				{
 					fprintf(stderr,
 							"client %d script %d command %d query %d: expected one row, got %d\n",
@@ -2745,16 +2750,23 @@ readCommandResponse(CState *st, char *varprefix)
 				break;
 
 			case PGRES_TUPLES_OK:
-				if (is_last && varprefix != NULL)
+				if (varprefix != NULL && (is_last || is_aset))
 				{
-					if (PQntuples(res) != 1)
+					int ntuples		= PQntuples(res);
+
+					if (!is_aset && ntuples != 1)
 					{
+						/* under \gset, report the error */
 						fprintf(stderr,
 								"client %d script %d command %d query %d: expected one row, got %d\n",
 								st->id, st->use_file, st->command, qrynum, PQntuples(res));
 						goto error;
 					}
 
+					/* skip empty result under \aset */
+					if (ntuples <= 0)
+						break;
+
 					/* store results into variables */
 					for (int fld = 0; fld < PQnfields(res); fld++)
 					{
@@ -2764,15 +2776,14 @@ readCommandResponse(CState *st, char *varprefix)
 						if (*varprefix != '\0')
 							varname = psprintf("%s%s", varprefix, varname);
 
-						/* store result as a string */
-						if (!putVariable(st, "gset", varname,
-										 PQgetvalue(res, 0, fld)))
+						/* store last row result as a string */
+						if (!putVariable(st, is_aset ? "aset" : "gset", varname,
+										 PQgetvalue(res, ntuples - 1, fld)))
 						{
 							/* internal error */
 							fprintf(stderr,
 									"client %d script %d command %d query %d: error storing into variable %s\n",
-									st->id, st->use_file, st->command, qrynum,
-									varname);
+									st->id, st->use_file, st->command, qrynum, varname);
 							goto error;
 						}
 
@@ -3189,7 +3200,9 @@ advanceConnectionState(TState *thread, CState *st, StatsData *agg)
 					return;		/* don't have the whole result yet */
 
 				/* store or discard the query results */
-				if (readCommandResponse(st, sql_script[st->use_file].commands[st->command]->varprefix))
+				if (readCommandResponse(st,
+										sql_script[st->use_file].commands[st->command]->varprefix,
+										sql_script[st->use_file].commands[st->command]->aset))
 					st->state = CSTATE_END_COMMAND;
 				else
 					st->state = CSTATE_ABORTED;
@@ -4125,6 +4138,7 @@ create_sql_command(PQExpBuffer buf, const char *source)
 	my_command->argc = 0;
 	memset(my_command->argv, 0, sizeof(my_command->argv));
 	my_command->varprefix = NULL;	/* allocated later, if needed */
+	my_command->aset = false;
 	my_command->expr = NULL;
 	initSimpleStats(&my_command->stats);
 
@@ -4353,7 +4367,7 @@ process_backslash_command(PsqlScanState sstate, const char *source)
 			syntax_error(source, lineno, my_command->first_line, my_command->argv[0],
 						 "unexpected argument", NULL, -1);
 	}
-	else if (my_command->meta == META_GSET)
+	else if (my_command->meta == META_GSET || my_command->meta == META_ASET)
 	{
 		if (my_command->argc > 2)
 			syntax_error(source, lineno, my_command->first_line, my_command->argv[0],
@@ -4498,10 +4512,10 @@ ParseScript(const char *script, const char *desc, int weight)
 			if (command)
 			{
 				/*
-				 * If this is gset, merge into the preceding command. (We
-				 * don't use a command slot in this case).
+				 * If this is gset or aset, merge into the preceding command.
+				 * (We don't use a command slot in this case).
 				 */
-				if (command->meta == META_GSET)
+				if (command->meta == META_GSET || command->meta == META_ASET)
 				{
 					Command    *cmd;
 
@@ -4523,6 +4537,7 @@ ParseScript(const char *script, const char *desc, int weight)
 						cmd->varprefix = pg_strdup("");
 					else
 						cmd->varprefix = pg_strdup(command->argv[1]);
+					cmd->aset = (command->meta == META_ASET);
 
 					/* cleanup unused command */
 					free_command(command);
diff --git a/src/bin/pgbench/t/001_pgbench_with_server.pl b/src/bin/pgbench/t/001_pgbench_with_server.pl
index dc2c72fa92..62cc0bb1ac 100644
--- a/src/bin/pgbench/t/001_pgbench_with_server.pl
+++ b/src/bin/pgbench/t/001_pgbench_with_server.pl
@@ -538,7 +538,7 @@ pgbench(
 }
 	});
 
-# working \gset
+# working \gset and \aset
 pgbench(
 	'-t 1', 0,
 	[ qr{type: .*/001_pgbench_gset}, qr{processed: 1/1} ],
@@ -548,9 +548,11 @@ pgbench(
 		qr{command=6.: int 2\b},
 		qr{command=8.: int 3\b},
 		qr{command=10.: int 4\b},
-		qr{command=12.: int 5\b}
+		qr{command=12.: int 5\b},
+		qr{command=15.: int 8\b},
+		qr{command=16.: int 7\b}
 	],
-	'pgbench gset command',
+	'pgbench gset & aset commands',
 	{
 		'001_pgbench_gset' => q{-- test gset
 -- no columns
@@ -571,6 +573,12 @@ SELECT 0 AS i4, 4 AS i4 \gset
 -- work on the last SQL command under \;
 \; \; SELECT 0 AS i5 \; SELECT 5 AS i5 \; \; \gset
 \set i debug(:i5)
+-- test aset, which applies to a combined query
+\; SELECT 6 AS i6 \; SELECT 7 AS i7 \; \aset
+-- unless it returns more than one row, last is kept
+SELECT 8 AS i6 UNION SELECT 9 ORDER BY 1 DESC \aset
+\set i debug(:i6)
+\set i debug(:i7)
 }
 	});