diff --git a/doc/src/sgml/ref/pgbench.sgml b/doc/src/sgml/ref/pgbench.sgml
index 41b3880c91f..58a2aa3bf20 100644
--- a/doc/src/sgml/ref/pgbench.sgml
+++ b/doc/src/sgml/ref/pgbench.sgml
@@ -1057,18 +1057,29 @@ pgbench options d
\gset [prefix]
+ \aset [prefix]
- 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 (;).
- 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 prefix if provided.
+ When the \gset 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 prefix
+ if provided.
+
+
+
+ When the \aset command is used, all combined SQL queries
+ (separated by \;) have their columns stored into variables
+ named after column names, and prefixed with prefix
+ 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.
@@ -1077,6 +1088,8 @@ pgbench options d
p_two and p_three
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
+ four and five.
UPDATE pgbench_accounts
SET abalance = abalance + :delta
@@ -1085,6 +1098,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
diff --git a/src/bin/pgbench/pgbench.c b/src/bin/pgbench/pgbench.c
index b8864c6ae53..e99af801675 100644
--- a/src/bin/pgbench/pgbench.c
+++ b/src/bin/pgbench/pgbench.c
@@ -480,6 +480,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 */
@@ -504,14 +505,16 @@ static const char *QUERYMODE[] = {"simple", "extended", "prepared"};
* not applied.
* first_line A short, single-line extract of 'lines', for error reporting.
* type SQL_COMMAND or META_COMMAND
- * meta The type of meta-command, or META_NONE if command is SQL
+ * meta The type of meta-command, with META_NONE/GSET/ASET if command
+ * is SQL.
* argc Number of arguments of the command, 0 if not yet processed.
* argv Command arguments, the first of which is the command or SQL
* string itself. For SQL commands, after post-processing
* argv[0] is the same as 'lines' with variables substituted.
- * varprefix SQL commands terminated with \gset have this set
+ * varprefix SQL commands terminated with \gset or \aset 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.
*/
@@ -2489,6 +2492,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;
@@ -2711,17 +2716,25 @@ 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 (META_GSET) or *all* commands
+ * (META_ASET).
*
* Returns true if everything is A-OK, false if any error occurs.
*/
static bool
-readCommandResponse(CState *st, char *varprefix)
+readCommandResponse(CState *st, MetaCommand meta, char *varprefix)
{
PGresult *res;
PGresult *next_res;
int qrynum = 0;
+ /*
+ * varprefix should be set only with \gset or \aset, and SQL commands do
+ * not need it.
+ */
+ Assert((meta == META_NONE && varprefix == NULL) ||
+ ((meta == META_GSET || meta == META_ASET) && varprefix != NULL));
+
res = PQgetResult(st->con);
while (res != NULL)
@@ -2736,7 +2749,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 && meta == META_GSET)
{
pg_log_error("client %d script %d command %d query %d: expected one row, got %d",
st->id, st->use_file, st->command, qrynum, 0);
@@ -2745,14 +2758,22 @@ readCommandResponse(CState *st, char *varprefix)
break;
case PGRES_TUPLES_OK:
- if (is_last && varprefix != NULL)
+ if ((is_last && meta == META_GSET) || meta == META_ASET)
{
- if (PQntuples(res) != 1)
+ int ntuples = PQntuples(res);
+
+ if (meta == META_GSET && ntuples != 1)
{
+ /* under \gset, report the error */
pg_log_error("client %d script %d command %d query %d: expected one row, got %d",
st->id, st->use_file, st->command, qrynum, PQntuples(res));
goto error;
}
+ else if (meta == META_ASET && ntuples <= 0)
+ {
+ /* coldly skip empty result under \aset */
+ break;
+ }
/* store results into variables */
for (int fld = 0; fld < PQnfields(res); fld++)
@@ -2763,9 +2784,9 @@ 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, meta == META_ASET ? "aset" : "gset", varname,
+ PQgetvalue(res, ntuples - 1, fld)))
{
/* internal error */
pg_log_error("client %d script %d command %d query %d: error storing into variable %s",
@@ -3181,7 +3202,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]->meta,
+ sql_script[st->use_file].commands[st->command]->varprefix))
st->state = CSTATE_END_COMMAND;
else
st->state = CSTATE_ABORTED;
@@ -4660,7 +4683,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],
@@ -4804,10 +4827,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;
@@ -4830,6 +4853,9 @@ ParseScript(const char *script, const char *desc, int weight)
else
cmd->varprefix = pg_strdup(command->argv[1]);
+ /* update the sql command meta */
+ cmd->meta = command->meta;
+
/* 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 b85a3ac32dd..e85728c3790 100644
--- a/src/bin/pgbench/t/001_pgbench_with_server.pl
+++ b/src/bin/pgbench/t/001_pgbench_with_server.pl
@@ -699,6 +699,51 @@ 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)
+}
+ });
+# \gset cannot accept more than one row, causing command to fail.
+pgbench(
+ '-t 1', 2,
+ [ qr{type: .*/001_pgbench_gset_two_rows}, qr{processed: 0/1} ],
+ [qr{expected one row, got 2\b}],
+ 'pgbench gset command with two rows',
+ {
+ '001_pgbench_gset_two_rows' => q{
+SELECT 5432 AS fail UNION SELECT 5433 ORDER BY 1 \gset
+}
+ });
+
+# working \aset
+# Valid cases.
+pgbench(
+ '-t 1', 0,
+ [ qr{type: .*/001_pgbench_aset}, qr{processed: 1/1} ],
+ [ qr{command=3.: int 8\b}, qr{command=4.: int 7\b} ],
+ 'pgbench aset command',
+ {
+ '001_pgbench_aset' => q{
+-- 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)
+}
+ });
+# Empty result set with \aset, causing command to fail.
+pgbench(
+ '-t 1', 2,
+ [ qr{type: .*/001_pgbench_aset_empty}, qr{processed: 0/1} ],
+ [
+ qr{undefined variable \"i8\"},
+ qr{evaluation of meta-command failed\b}
+ ],
+ 'pgbench aset command with empty result',
+ {
+ '001_pgbench_aset_empty' => q{
+-- empty result
+\; SELECT 5432 AS i8 WHERE FALSE \; \aset
+\set i debug(:i8)
}
});