It seems the new feature \gset doesn't work with tables having none ascii column names:Indeed. The same error is triggered with the \set syntax, which does not involve any query execution.I have added a sentence mentionning the restriction when variables are first discussed in the documentation, see attached patch.
Here is a v10: - does not talk about ASCII variable name constraint, as a patch has been submitted independently to lift this constraint. - rename gcset to cset (compound set, \; + \set), where gset is ; + \set, because "\gcset" looked really strange. - simplify the code a little bit. Also attached is an updated test script. -- Fabien.
diff --git a/doc/src/sgml/ref/pgbench.sgml b/doc/src/sgml/ref/pgbench.sgml
index 1eee8dc..862e1d4 100644
--- a/doc/src/sgml/ref/pgbench.sgml
+++ b/doc/src/sgml/ref/pgbench.sgml
@@ -816,6 +816,51 @@ pgbench <optional> <replaceable>options</> </optional> <replaceable>dbname</>
</para>
<variablelist>
+ <varlistentry id='pgbench-metacommand-gset'>
+ <term>
+ <literal>\cset [<replaceable>prefix</>]</literal> or
+ <literal>\gset [<replaceable>prefix</>]</literal>
+ </term>
+
+ <listitem>
+ <para>
+ These commands may be used to end SQL queries, replacing a semicolon.
+ <literal>\cset</> replaces an embedded semicolon (<literal>\;</>) within
+ a compound SQL command, and <literal>\gset</> replaces a final
+ (<literal>;</>) semicolon which ends the SQL command.
+ </para>
+
+ <para>
+ When these commands are 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</> if provided.
+ </para>
+
+ <para>
+ The following example puts the final account balance from the first query
+ into variable <replaceable>abalance</>, and fills variables
+ <replaceable>one</>, <replaceable>two</> and <replaceable>p_three</> with
+ integers from a compound query.
+<programlisting>
+UPDATE pgbench_accounts
+ SET abalance = abalance + :delta
+ WHERE aid = :aid
+ RETURNING abalance \gset
+-- compound of two queries
+SELECT 1 AS one, 2 AS two \cset
+SELECT 3 AS three \gset p_
+</programlisting>
+ </para>
+
+ <note>
+ <para>
+ <literal>\cset</> and <literal>\gset</> commands do not work when
+ empty SQL queries appear within a compound SQL command.
+ </para>
+ </note>
+ </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 ae36247..1e2402a 100644
--- a/src/bin/pgbench/pgbench.c
+++ b/src/bin/pgbench/pgbench.c
@@ -375,11 +375,14 @@ static const char *QUERYMODE[] = {"simple", "extended", "prepared"};
typedef struct
{
- char *line; /* text of command line */
+ char *line; /* first line for short display */
+ char *lines; /* full multi-line text of command */
int command_num; /* unique index of this Command struct */
int type; /* command type (SQL_COMMAND or META_COMMAND) */
int argc; /* number of command words */
char *argv[MAX_ARGS]; /* command word list */
+ int compound; /* last compound command (number of \;) */
+ char **gset; /* per-compound command prefix */
PgBenchExpr *expr; /* parsed expression, if needed */
SimpleStats stats; /* time spent in this command */
} Command;
@@ -1225,6 +1228,104 @@ 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 **gset)
+{
+ PGresult *res;
+ int compound = 0;
+
+ while ((res = PQgetResult(st->con)) != NULL)
+ {
+ switch (PQresultStatus(res))
+ {
+ case PGRES_COMMAND_OK: /* non-SELECT commands */
+ case PGRES_EMPTY_QUERY: /* may be used for testing no-op overhead */
+ if (gset[compound] != NULL)
+ {
+ fprintf(stderr,
+ "client %d file %d command %d compound %d: "
+ "\\gset expects a row\n",
+ st->id, st->use_file, st->command, compound);
+ st->ecnt++;
+ return false;
+ }
+ break; /* OK */
+
+ case PGRES_TUPLES_OK:
+ if (gset[compound] != NULL)
+ {
+ /* store result into variables */
+ int ntuples = PQntuples(res),
+ nfields = PQnfields(res),
+ f;
+
+ if (ntuples != 1)
+ {
+ fprintf(stderr,
+ "client %d file %d command %d compound %d: "
+ "expecting one row, got %d\n",
+ st->id, st->use_file, st->command, compound, ntuples);
+ st->ecnt++;
+ PQclear(res);
+ discard_response(st);
+ return false;
+ }
+
+ for (f = 0; f < nfields ; f++)
+ {
+ char *varname = PQfname(res, f);
+ if (*gset[compound] != '\0')
+ varname = psprintf("%s%s", gset[compound], varname);
+
+ /* store result as a string */
+ if (!putVariable(st, "gset", varname,
+ PQgetvalue(res, 0, f)))
+ {
+ /* internal error, should it rather abort? */
+ fprintf(stderr,
+ "client %d file %d command %d compound %d: "
+ "error storing into var %s\n",
+ st->id, st->use_file, st->command, compound,
+ varname);
+ st->ecnt++;
+ PQclear(res);
+ discard_response(st);
+ return false;
+ }
+
+ if (*gset[compound] != '\0')
+ free(varname);
+ }
+ }
+ break; /* OK */
+
+ default:
+ /* everything else is unexpected, so probably an error */
+ fprintf(stderr,
+ "client %d file %d aborted in command %d compound %d: %s",
+ st->id, st->use_file, st->command, compound,
+ PQerrorMessage(st->con));
+ st->ecnt++;
+ PQclear(res);
+ discard_response(st);
+ return false;
+ }
+
+ PQclear(res);
+ compound += 1;
+ }
+
+ if (compound == 0)
+ {
+ fprintf(stderr, "client %d command %d: no results\n", st->id, st->command);
+ 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)
@@ -1957,7 +2058,6 @@ evaluateSleep(CState *st, int argc, char **argv, int *usecs)
static void
doCustom(TState *thread, CState *st, StatsData *agg)
{
- PGresult *res;
Command *command;
instr_time now;
bool end_tx_processed = false;
@@ -2295,26 +2395,12 @@ doCustom(TState *thread, CState *st, StatsData *agg)
if (PQisBusy(st->con))
return; /* don't have the whole result yet */
- /*
- * Read and discard the query result;
- */
- res = PQgetResult(st->con);
- switch (PQresultStatus(res))
- {
- case PGRES_COMMAND_OK:
- case PGRES_TUPLES_OK:
- case PGRES_EMPTY_QUERY:
- /* OK */
- PQclear(res);
- discard_response(st);
- st->state = CSTATE_END_COMMAND;
- break;
- default:
- commandFailed(st, PQerrorMessage(st->con));
- PQclear(res);
- st->state = CSTATE_ABORTED;
- break;
- }
+ /* read and discard the query results */
+ if (read_response(st, command->gset))
+ st->state = CSTATE_END_COMMAND;
+ else
+ st->state = CSTATE_ABORTED;
+
break;
/*
@@ -2944,22 +3030,10 @@ syntax_error(const char *source, int lineno,
exit(1);
}
-/*
- * Parse a SQL command; return a Command struct, or NULL if it's a comment
- *
- * On entry, psqlscan.l has collected the command into "buf", so we don't
- * really need to do much here except check for comment and set up a
- * Command struct.
- */
-static Command *
-process_sql_command(PQExpBuffer buf, const char *source)
+static char *
+skip_sql_comments(char *p)
{
- Command *my_command;
- char *p;
- char *nlpos;
-
/* Skip any leading whitespace, as well as "--" style comments */
- p = buf->data;
for (;;)
{
if (isspace((unsigned char) *p))
@@ -2979,17 +3053,79 @@ process_sql_command(PQExpBuffer buf, const char *source)
if (*p == '\0')
return NULL;
+ return p;
+}
+
+/*
+ * Parse a SQL command; return a Command struct, or NULL if it's a comment
+ *
+ * On entry, psqlscan.l has collected the command into "buf", so we don't
+ * really need to do much here except check for comment and set up a
+ * Command struct.
+ */
+static Command *
+create_sql_command(PQExpBuffer buf, const char *source, int compounds)
+{
+ Command *my_command;
+ char *p = skip_sql_comments(buf->data);
+
+ if (p == NULL)
+ return NULL;
+
/* Allocate and initialize Command structure */
my_command = (Command *) pg_malloc0(sizeof(Command));
my_command->command_num = num_commands++;
my_command->type = SQL_COMMAND;
my_command->argc = 0;
+ my_command->compound = compounds;
+ my_command->gset = pg_malloc0(sizeof(char *) * (compounds+1));
initSimpleStats(&my_command->stats);
+ my_command->lines = pg_strdup(p);
+
+ return my_command;
+}
+
+static void
+append_sql_command(Command *my_command, char *more, int compounds)
+{
+ size_t lmore;
+ size_t len = strlen(my_command->lines);
+ int nc;
+
+ Assert(my_command->type == SQL_COMMAND && len > 0);
+
+ more = skip_sql_comments(more);
+
+ if (more == NULL)
+ return;
+
+ lmore = strlen(more);
+ my_command->lines = pg_realloc(my_command->lines, len + lmore + 2);
+ my_command->lines[len] = ';';
+ memcpy(my_command->lines + len + 1, more, lmore + 1);
+
+ nc = my_command->compound + 1 + compounds;
+ my_command->gset =
+ pg_realloc(my_command->gset, sizeof(char *) * (nc+1));
+ memset(my_command->gset + my_command->compound + 1, 0,
+ sizeof(char *) * (compounds + 1));
+ my_command->compound = nc;
+}
+
+static void
+postprocess_sql_command(Command *my_command)
+{
+ char *nlpos;
+ char *p;
+
+ Assert(my_command->type == SQL_COMMAND);
+
/*
* If SQL command is multi-line, we only want to save the first line as
* the "line" label.
*/
+ p = my_command->lines;
nlpos = strchr(p, '\n');
if (nlpos)
{
@@ -3003,19 +3139,17 @@ process_sql_command(PQExpBuffer buf, const char *source)
switch (querymode)
{
case QUERY_SIMPLE:
- my_command->argv[0] = pg_strdup(p);
+ my_command->argv[0] = my_command->lines;
my_command->argc++;
break;
case QUERY_EXTENDED:
case QUERY_PREPARED:
- if (!parseQuery(my_command, p))
+ if (!parseQuery(my_command, my_command->lines))
exit(1);
break;
default:
exit(1);
}
-
- return my_command;
}
/*
@@ -3173,6 +3307,13 @@ 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], "gset") == 0 ||
+ pg_strcasecmp(my_command->argv[0], "cset") == 0)
+ {
+ if (my_command->argc > 2)
+ syntax_error(source, lineno, my_command->line, my_command->argv[0],
+ "at most one argument expected", NULL, -1);
+ }
else
{
syntax_error(source, lineno, my_command->line, my_command->argv[0],
@@ -3196,6 +3337,9 @@ ParseScript(const char *script, const char *desc, int weight)
PQExpBufferData line_buf;
int alloc_num;
int index;
+ bool is_compound = false;
+ int lineno;
+ int start_offset;
#define COMMANDS_ALLOC_NUM 128
alloc_num = COMMANDS_ALLOC_NUM;
@@ -3219,6 +3363,7 @@ ParseScript(const char *script, const char *desc, int weight)
* stdstrings should be true, which is a bit riskier.
*/
psql_scan_setup(sstate, script, strlen(script), 0, true);
+ start_offset = expr_scanner_offset(sstate) - 1;
initPQExpBuffer(&line_buf);
@@ -3228,31 +3373,28 @@ ParseScript(const char *script, const char *desc, int weight)
{
PsqlScanResult sr;
promptStatus_t prompt;
- Command *command;
+ Command *command = NULL;
resetPQExpBuffer(&line_buf);
+ lineno = expr_scanner_get_lineno(sstate, start_offset);
+
+ sstate->semicolons = 0;
sr = psql_scan(sstate, &line_buf, &prompt);
- /* If we collected a SQL command, process that */
- command = process_sql_command(&line_buf, desc);
- if (command)
+ if (is_compound)
{
- ps.commands[index] = command;
- index++;
-
- if (index >= alloc_num)
- {
- alloc_num += COMMANDS_ALLOC_NUM;
- ps.commands = (Command **)
- pg_realloc(ps.commands, sizeof(Command *) * alloc_num);
- }
+ /* a multi-line command ended with \cset */
+ append_sql_command(ps.commands[index-1], line_buf.data,
+ sstate->semicolons);
+ is_compound = false;
}
-
- /* If we reached a backslash, process that */
- if (sr == PSCAN_BACKSLASH)
+ else
{
- command = process_backslash_command(sstate, desc);
+ /* If we collected a new SQL command, process that */
+ command = create_sql_command(&line_buf, desc, sstate->semicolons);
+
+ /* store new command */
if (command)
{
ps.commands[index] = command;
@@ -3267,6 +3409,67 @@ ParseScript(const char *script, const char *desc, int weight)
}
}
+ if (sr == PSCAN_BACKSLASH)
+ {
+ command = process_backslash_command(sstate, desc);
+
+ if (command)
+ {
+ char * bs_cmd = command->argv[0];
+
+ /* merge gset variants into preceeding SQL command */
+ if (pg_strcasecmp(bs_cmd, "gset") == 0 ||
+ pg_strcasecmp(bs_cmd, "cset") == 0)
+ {
+ int cindex;
+ Command *sql_cmd;
+
+ is_compound = bs_cmd[0] == 'c';
+
+ if (index == 0)
+ syntax_error(desc, lineno, NULL, NULL,
+ "\\gset cannot start a script",
+ NULL, -1);
+
+ sql_cmd = ps.commands[index-1];
+
+ if (sql_cmd->type != SQL_COMMAND)
+ syntax_error(desc, lineno, NULL, NULL,
+ "\\gset must follow a SQL command",
+ sql_cmd->line, -1);
+
+ /* this \gset applies to the last sub-command */
+ cindex = sql_cmd->compound;
+
+ if (sql_cmd->gset[cindex] != NULL)
+ syntax_error(desc, lineno, NULL, NULL,
+ "\\gset cannot follow one another",
+ NULL, -1);
+
+ /* get variable prefix */
+ if (command->argc <= 1 || command->argv[1][0] == '\0')
+ sql_cmd->gset[cindex] = "";
+ else
+ sql_cmd->gset[cindex] = command->argv[1];
+
+ /* cleanup unused backslash command */
+ pg_free(command);
+ }
+ else /* any other backslash command is a Command */
+ {
+ ps.commands[index] = command;
+ index++;
+
+ if (index >= alloc_num)
+ {
+ alloc_num += COMMANDS_ALLOC_NUM;
+ ps.commands = (Command **)
+ pg_realloc(ps.commands, sizeof(Command *) * alloc_num);
+ }
+ }
+ }
+ }
+
/* Done if we reached EOF */
if (sr == PSCAN_INCOMPLETE || sr == PSCAN_EOL)
break;
@@ -3274,6 +3477,11 @@ ParseScript(const char *script, const char *desc, int weight)
ps.commands[index] = NULL;
+ /* complete SQL command initializations */
+ while (--index >= 0)
+ if (ps.commands[index]->type == SQL_COMMAND)
+ postprocess_sql_command(ps.commands[index]);
+
addScript(ps);
termPQExpBuffer(&line_buf);
diff --git a/src/bin/pgbench/pgbench.h b/src/bin/pgbench/pgbench.h
index 38b3af5..00bf432 100644
--- a/src/bin/pgbench/pgbench.h
+++ b/src/bin/pgbench/pgbench.h
@@ -11,6 +11,7 @@
#ifndef PGBENCH_H
#define PGBENCH_H
+#include "fe_utils/psqlscan_int.h"
#include "fe_utils/psqlscan.h"
/*
diff --git a/src/fe_utils/psqlscan.l b/src/fe_utils/psqlscan.l
index 27689d7..9974722 100644
--- a/src/fe_utils/psqlscan.l
+++ b/src/fe_utils/psqlscan.l
@@ -678,8 +678,15 @@ other .
* substitution. We want these before {self}, also.
*/
-"\\"[;:] {
- /* Force a semicolon or colon into the query buffer */
+"\\"; {
+ /* Count compound commands */
+ cur_state->semicolons++;
+ /* Force a semicolon into the query buffer */
+ psqlscan_emit(cur_state, yytext + 1, 1);
+ }
+
+"\\": {
+ /* Force a colon into the query buffer */
psqlscan_emit(cur_state, yytext + 1, 1);
}
diff --git a/src/include/fe_utils/psqlscan_int.h b/src/include/fe_utils/psqlscan_int.h
index af62f5e..3fe7789 100644
--- a/src/include/fe_utils/psqlscan_int.h
+++ b/src/include/fe_utils/psqlscan_int.h
@@ -112,6 +112,7 @@ typedef struct PsqlScanStateData
int start_state; /* yylex's starting/finishing state */
int paren_depth; /* depth of nesting in parentheses */
int xcdepth; /* depth of nesting in slash-star comments */
+ int semicolons; /* number of embedded (\;) semi-colons */
char *dolqstart; /* current $foo$ quote start string */
/*
gset-1.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
