diff --git a/doc/src/sgml/ref/pgbench.sgml b/doc/src/sgml/ref/pgbench.sgml
index f3afedb..8a7ad3e 100644
--- a/doc/src/sgml/ref/pgbench.sgml
+++ b/doc/src/sgml/ref/pgbench.sgml
@@ -809,6 +809,29 @@ pgbench options dbname
+
+
+ \into var1 [var2 ...]
+
+
+
+
+ Stores the first fields of the resulting row from the preceding
+ SELECT commands into these variables.
+ The queries must yield exactly one row and the number of provided
+ variables must be less than the total number of columns of the results.
+
+
+
+ Example:
+
+SELECT abalance FROM pgbench_accounts WHERE aid=5432;
+\into balance
+
+
+
+
+
\set varname expression
diff --git a/src/bin/pgbench/pgbench.c b/src/bin/pgbench/pgbench.c
index 87fb006..d9827dc 100644
--- a/src/bin/pgbench/pgbench.c
+++ b/src/bin/pgbench/pgbench.c
@@ -307,6 +307,7 @@ typedef struct
int type; /* command type (SQL_COMMAND or META_COMMAND) */
int argc; /* number of command words */
char *argv[MAX_ARGS]; /* command word list */
+ char *into[MAX_ARGS]; /* NULL-terminated target variables for \into */
PgBenchExpr *expr; /* parsed expression, if needed */
SimpleStats stats; /* time spent in this command */
} Command;
@@ -1148,6 +1149,96 @@ getQueryParams(CState *st, const Command *command, const char **params)
params[i] = getVariable(st, command->argv[i + 1]);
}
+/* read all responses from backend */
+static bool
+read_response(CState *st, char * into[])
+{
+ PGresult *res;
+ int i = 0;
+ bool at_least_one = false;
+
+ while ((res = PQgetResult(st->con)) != NULL)
+ {
+ at_least_one = true;
+
+ switch (PQresultStatus(res))
+ {
+ case PGRES_COMMAND_OK: /* non-SELECT commands */
+ case PGRES_EMPTY_QUERY: /* may be used for testing no-op overhead */
+ break; /* OK */
+
+ case PGRES_TUPLES_OK:
+ if (into != NULL && into[i] != NULL)
+ {
+ /* store result into variables */
+ int ntuples = PQntuples(res),
+ nfields = PQnfields(res),
+ f = 0;
+
+ if (ntuples != 1)
+ {
+ fprintf(stderr,
+ "client %d state %d expecting one row, got %d\n",
+ st->id, st->state, ntuples);
+ st->ecnt++;
+ PQclear(res);
+ discard_response(st);
+ return false;
+ }
+
+ while (into[i] != NULL && f < nfields)
+ {
+ /* store result as a string */
+ if (!putVariable(st, "into", into[i], PQgetvalue(res, 0, f)))
+ {
+ /* internal error, should it rather abort? */
+ fprintf(stderr,
+ "client %d state %d: error storing into var %s\n",
+ st->id, st->state, into[i]);
+ st->ecnt++;
+ PQclear(res);
+ discard_response(st);
+ return false;
+ }
+
+ i++;
+ f++;
+ }
+ }
+ break; /* OK */
+
+ default:
+ /* everything else is unexpected, so probably an error */
+ fprintf(stderr, "client %d aborted in state %d: %s",
+ st->id, st->state, PQerrorMessage(st->con));
+ st->ecnt++;
+ PQclear(res);
+ discard_response(st);
+ return false;
+ }
+
+ PQclear(res);
+ }
+
+ if (!at_least_one)
+ {
+ fprintf(stderr, "client %d state %d: no results\n", st->id, st->state);
+ st->ecnt++;
+ return false;
+ }
+
+ if (into != NULL && into[i] != NULL)
+ {
+ fprintf(stderr,
+ "client %d state %d: missing results to fill into variable %s\n",
+ st->id, st->state, into[i]);
+ st->ecnt++;
+ return false;
+ }
+
+ return true;
+}
+
/* get a value as an int, tell if there is a problem */
static bool
coerceToInt(PgBenchValue *pval, int64 *ival)
@@ -1764,7 +1855,6 @@ chooseScript(TState *thread)
static bool
doCustom(TState *thread, CState *st, StatsData *agg)
{
- PGresult *res;
Command **commands;
bool trans_needs_throttle = false;
instr_time now;
@@ -1891,22 +1981,11 @@ top:
{
/*
* Read and discard the query result; note this is not included in
- * the statement latency numbers.
+ * the statement latency numbers (above), thus if reading the
+ * response fails the transaction is counted nevertheless.
*/
- res = PQgetResult(st->con);
- switch (PQresultStatus(res))
- {
- case PGRES_COMMAND_OK:
- case PGRES_TUPLES_OK:
- break; /* OK */
- default:
- fprintf(stderr, "client %d aborted in state %d: %s",
- st->id, st->state, PQerrorMessage(st->con));
- PQclear(res);
- return clientDone(st);
- }
- PQclear(res);
- discard_response(st);
+ if (!read_response(st, commands[st->state]->into))
+ return clientDone(st);
}
if (commands[st->state + 1] == NULL)
@@ -2930,6 +3009,24 @@ process_backslash_command(PsqlScanState sstate, const char *source)
syntax_error(source, lineno, my_command->line, my_command->argv[0],
"missing command", NULL, -1);
}
+ else if (pg_strcasecmp(my_command->argv[0], "into") == 0)
+ {
+ int i;
+
+ /* at least one variable name must be provided */
+ if (my_command->argc < 2)
+ syntax_error(source, lineno, my_command->line, my_command->argv[0],
+ "missing variable name", NULL, -1);
+
+ /* check that all variable names are valid */
+ for (i = 1; i < my_command->argc; i++)
+ {
+ if (!isLegalVariableName(my_command->argv[i]))
+ syntax_error(source, lineno, my_command->line,
+ my_command->argv[0], "invalid variable name",
+ my_command->argv[i], -1);
+ }
+ }
else
{
syntax_error(source, lineno, my_command->line, my_command->argv[0],
@@ -3010,7 +3107,40 @@ ParseScript(const char *script, const char *desc, int weight)
if (sr == PSCAN_BACKSLASH)
{
command = process_backslash_command(sstate, desc);
- if (command)
+
+ /* special case for \into: merge into preceding SQL command */
+ if (command && command->argc >= 1 &&
+ pg_strcasecmp(command->argv[0], "into") == 0)
+ {
+ if (index == 0)
+ {
+ fprintf(stderr,
+ "meta command \\into cannot start a script\n");
+ exit(1);
+ }
+ else if (ps.commands[index-1]->type != SQL_COMMAND)
+ {
+ fprintf(stderr,
+ "meta command \\into must follow a SELECT (%s)\n",
+ ps.commands[index-1]->line);
+ exit(1);
+ }
+ else if (ps.commands[index-1]->into[0] != NULL)
+ {
+ fprintf(stderr,
+ "SQL command %d already has an associated \\into\n",
+ index-1);
+ exit(1);
+ }
+ /* else (command->argc >= 2): already checked by parser */
+
+ /* into is NULL-terminated thanks to pg_malloc0() */
+ memcpy(ps.commands[index-1]->into,
+ &command->argv[1], sizeof(char *) * (MAX_ARGS-1));
+
+ pg_free(command);
+ }
+ else if (command)
{
ps.commands[index] = command;
index++;