diff --git a/doc/src/sgml/ref/pgbench.sgml b/doc/src/sgml/ref/pgbench.sgml
index ef22a484e7..200f6b0010 100644
--- a/doc/src/sgml/ref/pgbench.sgml
+++ b/doc/src/sgml/ref/pgbench.sgml
@@ -966,18 +966,27 @@ 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 \;) which return just one row see their
+ columns stored into variables named after column names, and prefixed with
+ prefix if provided. Other queries are ignored.
@@ -986,6 +995,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
@@ -994,6 +1005,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 8b84658ccd..6895187455 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;
@@ -2715,7 +2720,7 @@ sendCommand(CState *st, Command *command)
* 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,10 +2750,15 @@ 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)
{
+ /* under \aset, ignore results which do not return one row */
+ if (is_aset)
+ break;
+
+ /* else 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));
@@ -2765,7 +2775,7 @@ readCommandResponse(CState *st, char *varprefix)
varname = psprintf("%s%s", varprefix, varname);
/* store result as a string */
- if (!putVariable(st, "gset", varname,
+ if (!putVariable(st, is_aset ? "aset" : "gset", varname,
PQgetvalue(res, 0, fld)))
{
/* internal error */
@@ -3189,7 +3199,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 +4137,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 +4366,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 +4511,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 +4536,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..87ef7d9136 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 6\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
+SELECT 8 AS i6 UNION SELECT 9 \aset
+\set i debug(:i6)
+\set i debug(:i7)
}
});