Hello devs,
I mentionned my intention to add some features to pgbench back in March: https://www.postgresql.org/message-id/alpine.DEB.2.10.1603301618570.5677@stoThe attached patch adds an \into meta command to store results of preceding SELECTs into pgbench variables, so that they can be reused afterwards.
The feature is useful to make more realistic scripts, currently pgbench script cannot really interact with the database as results are discarded.
The chosen syntax is easy to understand and the implementation is quite light, with minimal impact on the code base. I think that this is a reasonnable compromise.
The SELECTs must yield exactly one row, the number of variables must be less than the number of columns.
Also attached a set of test scripts, especially to trigger various error cases.
-- Fabien.
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 <optional> <replaceable>options</> </optional> <replaceable>dbname</>
</para>
<variablelist>
+ <varlistentry id='pgbench-metacommand-into'>
+ <term>
+ <literal>\into <replaceable>var1</> [<replaceable>var2</> ...]</literal>
+ </term>
+
+ <listitem>
+ <para>
+ Stores the first fields of the resulting row from the preceding
+ <command>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.
+ </para>
+
+ <para>
+ Example:
+<programlisting>
+SELECT abalance FROM pgbench_accounts WHERE aid=5432;
+\into balance
+</programlisting>
+ </para>
+ </listitem>
+ </varlistentry>
+
<varlistentry id='pgbench-metacommand-set'>
<term>
<literal>\set <replaceable>varname</> <replaceable>expression</></literal>
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++;
into.sql
Description: application/sql
into-err-1.sql
Description: application/sql
into-err-2.sql
Description: application/sql
into-err-3.sql
Description: application/sql
into-err-4.sql
Description: application/sql
into-err-5.sql
Description: application/sql
into-err-6.sql
Description: application/sql
-- Sent via pgsql-hackers mailing list (pgsql-hackers@postgresql.org) To make changes to your subscription: http://www.postgresql.org/mailpref/pgsql-hackers
