mirrors
/
postgres
mirror of https://github.com/postgres/postgres
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

pgbench: add \cset and \gset commands 

These commands allow assignment of values produced by queries to pgbench
variables, where they can be used by further commands.  \gset terminates
a command sequence (just like a bare semicolon); \cset separates
multiple queries in a compound command, like an escaped semicolon (\;).
A prefix can be provided to the \-command and is prepended to the name
of each output column to produce the final variable name.

This feature allows pgbench scripts to react meaningfully to the actual
database contents, allowing more powerful benchmarks to be written.

Authors: Fabien Coelho, Álvaro Herrera
Reviewed-by: Amit Langote <Langote_Amit_f8@lab.ntt.co.jp>
Reviewed-by: Stephen Frost <sfrost@snowman.net>
Reviewed-by: Pavel Stehule <pavel.stehule@gmail.com>
Reviewed-by: Tom Lane <tgl@sss.pgh.pa.us>
Reviewed-by: Tatsuo Ishii <ishii@sraoss.co.jp>
Reviewed-by: Rafia Sabih <rafia.sabih@enterprisedb.com>
Discussion: https://postgr.es/m/alpine.DEB.2.20.1607091005330.3412@sto
This commit is contained in:
Alvaro Herrera 2019-01-10 13:42:20 -03:00
parent e1c1d5444e
commit 6260cc550b
6 changed files with 608 additions and 143 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

85
doc/src/sgml/ref/pgbench.sgml
View File

@ -954,6 +954,91 @@ pgbench <optional> <replaceable>options</replaceable> </optional> <replaceable>d
</para>
<variablelist>
<varlistentry id='pgbench-metacommand-cset'>
<term>
<literal>\cset [<replaceable>prefix</replaceable>]</literal>
</term>
<listitem>
<para>
This command may be used to end SQL queries, replacing an embedded
semicolon (<literal>\;</literal>) within a compound SQL command.
</para>
<para>
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 <replaceable>prefix</replaceable> if provided.
</para>
<para>
The following example sends four queries as one compound SQL command,
inducing one message sent at the protocol level.
The result of the first query is stored into variable <replaceable>one</replaceable>,
the results of the third query are stored into variables <replaceable>z_three</replaceable>
and <replaceable>z_four</replaceable>,
whereas the results of the other queries are discarded.
<programlisting>
-- compound of four queries
SELECT 1 AS one \cset
SELECT 2 AS two \;
SELECT 3 AS three, 4 AS four \cset z_
SELECT 5;
</programlisting>
</para>
<note>
<para>
<literal>\cset</literal> does not work when empty SQL queries appear
within a compound SQL command.
</para>
</note>
</listitem>
</varlistentry>
<varlistentry id='pgbench-metacommand-gset'>
<term>
<literal>\gset [<replaceable>prefix</replaceable>]</literal>
</term>
<listitem>
<para>
This command may be used to end SQL queries, replacing a final semicolon
(<literal>;</literal>).
</para>
<para>
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 <replaceable>prefix</replaceable> if provided.
</para>
<para>
The following example puts the final account balance from the first query
into variable <replaceable>abalance</replaceable>, and fills variables
<replaceable>p_two</replaceable> and <replaceable>p_three</replaceable>
with integers from the last query.
The result of the second query is discarded.
<programlisting>
UPDATE pgbench_accounts
SET abalance = abalance + :delta
WHERE aid = :aid
RETURNING abalance \gset
-- compound of two queries
SELECT 1 \;
SELECT 2 AS two, 3 AS three \gset p_
</programlisting>
</para>
<note>
<para>
<literal>\gset</literal> does not work when empty SQL queries appear
within a compound SQL command.
</para>
</note>
</listitem>
</varlistentry>
<varlistentry>
<term><literal>\if</literal> <replaceable class="parameter">expression</replaceable></term>
<term><literal>\elif</literal> <replaceable class="parameter">expression</replaceable></term>

559
src/bin/pgbench/pgbench.c
View File

@ -482,6 +482,8 @@ typedef enum MetaCommand
META_SETSHELL, /* \setshell */
META_SHELL, /* \shell */
META_SLEEP, /* \sleep */
META_CSET, /* \cset */
META_GSET, /* \gset */
META_IF, /* \if */
META_ELIF, /* \elif */
META_ELSE, /* \else */
@ -499,16 +501,39 @@ typedef enum QueryMode
static QueryMode querymode = QUERY_SIMPLE;
static const char *QUERYMODE[] = {"simple", "extended", "prepared"};
typedef struct
/*
* struct Command represents one command in a script.
*
* lines The raw, possibly multi-line command text. Variable substitution
* 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
* 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.
* nqueries In a multi-command SQL line, the number of queries.
* varprefix SQL commands terminated with \gset or \cset have this set
* to a non NULL value. If nonempty, it's used to prefix the
* variable name that receives the value.
* varprefix_max Allocated size of the varprefix array.
* expr Parsed expression, if needed.
* stats Time spent in this command.
*/
typedef struct Command
{
char *line; /* text of command line */
int command_num; /* unique index of this Command struct */
int type; /* command type (SQL_COMMAND or META_COMMAND) */
MetaCommand meta; /* meta command identifier, or META_NONE */
int argc; /* number of command words */
char *argv[MAX_ARGS]; /* command word list */
PgBenchExpr *expr; /* parsed expression, if needed */
SimpleStats stats; /* time spent in this command */
PQExpBufferData lines;
char *first_line;
int type;
MetaCommand meta;
int argc;
char *argv[MAX_ARGS];
int nqueries;
char **varprefix;
int varprefix_max;
PgBenchExpr *expr;
SimpleStats stats;
} Command;
typedef struct ParsedScript
@ -521,7 +546,6 @@ typedef struct ParsedScript
static ParsedScript sql_script[MAX_SCRIPTS]; /* SQL script files */
static int num_scripts; /* number of scripts in sql_script[] */
static int num_commands = 0; /* total number of Command structs */
static int64 total_weight = 0;
static int debug = 0; /* debug flag */
@ -587,6 +611,7 @@ static void doLog(TState *thread, CState *st,
static void processXactStats(TState *thread, CState *st, instr_time *now,
bool skipped, StatsData *agg);
static void pgbench_error(const char *fmt,...) pg_attribute_printf(1, 2);
static void allocate_command_varprefix(Command *cmd, int totalqueries);
static void addScript(ParsedScript script);
static void *threadRun(void *arg);
static void finishCon(CState *st);
@ -2569,6 +2594,10 @@ getMetaCommand(const char *cmd)
mc = META_ELSE;
else if (pg_strcasecmp(cmd, "endif") == 0)
mc = META_ENDIF;
else if (pg_strcasecmp(cmd, "cset") == 0)
mc = META_CSET;
else if (pg_strcasecmp(cmd, "gset") == 0)
mc = META_GSET;
else
mc = META_NONE;
return mc;
@ -2796,6 +2825,109 @@ sendCommand(CState *st, Command *command)
return true;
}
/*
* Process query response from the backend.
*
* If varprefix is not NULL, it's the array of variable prefix names where to
* store the results.
*
* Returns true if everything is A-OK, false if any error occurs.
*/
static bool
readCommandResponse(CState *st, char **varprefix)
{
PGresult *res;
int qrynum = 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 (varprefix && varprefix[qrynum] != NULL)
{
fprintf(stderr,
"client %d script %d command %d query %d: expected one row, got %d\n",
st->id, st->use_file, st->command, qrynum, 0);
st->ecnt++;
return false;
}
break;
case PGRES_TUPLES_OK:
if (varprefix && varprefix[qrynum] != NULL)
{
if (PQntuples(res) != 1)
{
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));
st->ecnt++;
PQclear(res);
discard_response(st);
return false;
}
/* store results into variables */
for (int fld = 0; fld < PQnfields(res); fld++)
{
char *varname = PQfname(res, fld);
/* allocate varname only if necessary, freed below */
if (*varprefix[qrynum] != '\0')
varname =
psprintf("%s%s", varprefix[qrynum], varname);
/* store result as a string */
if (!putVariable(st, "gset", varname,
PQgetvalue(res, 0, fld)))
{
/* internal error */
fprintf(stderr,
"client %d script %d command %d query %d: error storing into variable %s\n",
st->id, st->use_file, st->command, qrynum,
varname);
st->ecnt++;
PQclear(res);
discard_response(st);
return false;
}
if (*varprefix[qrynum] != '\0')
pg_free(varname);
}
}
/* otherwise the result is simply thrown away by PQclear below */
break;
default:
/* anything else is unexpected */
fprintf(stderr,
"client %d script %d aborted in command %d query %d: %s",
st->id, st->use_file, st->command, qrynum,
PQerrorMessage(st->con));
st->ecnt++;
PQclear(res);
discard_response(st);
return false;
}
PQclear(res);
qrynum++;
}
if (qrynum == 0)
{
fprintf(stderr, "client %d command %d: no results\n", st->id, st->command);
st->ecnt++;
return false;
}
return true;
}
/*
* Parse the argument to a \sleep command, and return the requested amount
* of delay, in microseconds. Returns true on success, false on error.
@ -2862,8 +2994,6 @@ advanceConnectionState(TState *thread, CState *st, StatsData *agg)
*/
for (;;)
{
PGresult *res;
switch (st->state)
{
/* Select transaction (script) to run. */
@ -3141,24 +3271,11 @@ advanceConnectionState(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, "SQL", PQerrorMessage(st->con));
PQclear(res);
st->state = CSTATE_ABORTED;
break;
}
/* store or discard the query results */
if (readCommandResponse(st, sql_script[st->use_file].commands[st->command]->varprefix))
st->state = CSTATE_END_COMMAND;
else
st->state = CSTATE_ABORTED;
break;
/*
@ -3191,7 +3308,7 @@ advanceConnectionState(TState *thread, CState *st, StatsData *agg)
INSTR_TIME_SET_CURRENT_LAZY(now);
command = sql_script[st->use_file].commands[st->command];
command = sql_script[st->use_file].commands[st->command];
/* XXX could use a mutex here, but we choose not to */
addToSimpleStats(&command->stats,
INSTR_TIME_GET_DOUBLE(now) -
@ -3235,9 +3352,8 @@ advanceConnectionState(TState *thread, CState *st, StatsData *agg)
st->state = CSTATE_CHOOSE_SCRIPT;
/*
* Ensure that we always return on this point, so as to
* avoid an infinite loop if the script only contains meta
* commands.
* Ensure that we always return on this point, so as to avoid
* an infinite loop if the script only contains meta commands.
*/
return;
@ -3976,7 +4092,7 @@ runInitSteps(const char *initialize_steps)
/*
* Replace :param with $n throughout the command's SQL text, which
* is a modifiable string in cmd->argv[0].
* is a modifiable string in cmd->lines.
*/
static bool
parseQuery(Command *cmd)
@ -3984,12 +4100,9 @@ parseQuery(Command *cmd)
char *sql,
*p;
/* We don't want to scribble on cmd->argv[0] until done */
sql = pg_strdup(cmd->argv[0]);
cmd->argc = 1;
p = sql;
p = sql = pg_strdup(cmd->lines.data);
while ((p = strchr(p, ':')) != NULL)
{
char var[13];
@ -4009,7 +4122,7 @@ parseQuery(Command *cmd)
if (cmd->argc >= MAX_ARGS)
{
fprintf(stderr, "statement has too many arguments (maximum is %d): %s\n",
MAX_ARGS - 1, cmd->argv[0]);
MAX_ARGS - 1, cmd->lines.data);
pg_free(name);
return false;
}
@ -4021,7 +4134,7 @@ parseQuery(Command *cmd)
cmd->argc++;
}
pg_free(cmd->argv[0]);
Assert(cmd->argv[0] == NULL);
cmd->argv[0] = sql;
return true;
}
@ -4081,21 +4194,16 @@ syntax_error(const char *source, int lineno,
}
/*
* 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.
* Return a pointer to the start of the SQL command, after skipping over
* whitespace and "--" comments.
* If the end of the string is reached, return NULL.
*/
static Command *
process_sql_command(PQExpBuffer buf, const char *source)
static char *
skip_sql_comments(char *sql_command)
{
Command *my_command;
char *p;
char *nlpos;
char *p = sql_command;
/* Skip any leading whitespace, as well as "--" style comments */
p = buf->data;
for (;;)
{
if (isspace((unsigned char) *p))
@ -4111,41 +4219,153 @@ process_sql_command(PQExpBuffer buf, const char *source)
break;
}
/* If there's nothing but whitespace and comments, we're done */
/* NULL if there's nothing but whitespace and comments */
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 comments and set up a Command
* struct.
*/
static Command *
create_sql_command(PQExpBuffer buf, const char *source, int numqueries)
{
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 = (Command *) pg_malloc(sizeof(Command));
initPQExpBuffer(&my_command->lines);
appendPQExpBufferStr(&my_command->lines, p);
my_command->first_line = NULL; /* this is set later */
my_command->type = SQL_COMMAND;
my_command->meta = META_NONE;
my_command->argc = 0;
memset(my_command->argv, 0, sizeof(my_command->argv));
my_command->nqueries = numqueries;
my_command->varprefix = NULL; /* allocated later, if needed */
my_command->varprefix_max = 0;
my_command->expr = NULL;
initSimpleStats(&my_command->stats);
/*
* Install query text as the sole argv string. If we are using a
* non-simple query mode, we'll extract parameters from it later.
*/
my_command->argv[0] = pg_strdup(p);
my_command->argc = 1;
/*
* If SQL command is multi-line, we only want to save the first line as
* the "line" label.
*/
nlpos = strchr(p, '\n');
if (nlpos)
{
my_command->line = pg_malloc(nlpos - p + 1);
memcpy(my_command->line, p, nlpos - p);
my_command->line[nlpos - p] = '\0';
}
else
my_command->line = pg_strdup(p);
return my_command;
}
/* Free a Command structure and associated data */
static void
free_command(Command *command)
{
termPQExpBuffer(&command->lines);
if (command->first_line)
pg_free(command->first_line);
if (command->argv)
for (int i = 0; i < command->argc; i++)
pg_free(command->argv[i]);
if (command->varprefix)
{
for (int i = 0; i < command->varprefix_max; i++)
if (command->varprefix[i] &&
command->varprefix[i][0] != '\0') /* see ParseScript */
pg_free(command->varprefix[i]);
pg_free(command->varprefix);
}
/*
* It should also free expr recursively, but this is currently not needed
* as only \{g,c}set commands (which do not have an expression) are freed.
*/
pg_free(command);
}
/*
* append "more" text to current compound command which had been interrupted
* by \cset.
*/
static void
append_sql_command(Command *my_command, char *more, int numqueries)
{
Assert(my_command->type == SQL_COMMAND && my_command->lines.len > 0);
more = skip_sql_comments(more);
if (more == NULL)
return;
/* append command text, embedding a ';' in place of the \cset */
appendPQExpBuffer(&my_command->lines, ";\n%s", more);
my_command->nqueries += numqueries;
}
/*
* Once an SQL command is fully parsed, possibly by accumulating several
* parts, complete other fields of the Command structure.
*/
static void
postprocess_sql_command(Command *my_command)
{
char buffer[128];
Assert(my_command->type == SQL_COMMAND);
/* Save the first line for error display. */
strlcpy(buffer, my_command->lines.data, sizeof(buffer));
buffer[strcspn(buffer, "\n\r")] = '\0';
my_command->first_line = pg_strdup(buffer);
/* parse query if necessary */
switch (querymode)
{
case QUERY_SIMPLE:
my_command->argv[0] = my_command->lines.data;
my_command->argc++;
break;
case QUERY_EXTENDED:
case QUERY_PREPARED:
if (!parseQuery(my_command))
exit(1);
break;
default:
exit(1);
}
}
/*
* Determine the command's varprefix size needed and allocate memory for it
*/
static void
allocate_command_varprefix(Command *cmd, int totalqueries)
{
int new_max;
/* sufficient space already allocated? */
if (totalqueries <= cmd->varprefix_max)
return;
/* determine the new array size */
new_max = Max(cmd->varprefix_max, 2);
while (new_max < totalqueries)
new_max *= 2;
/* enlarge the array, zero-initializing the allocated space */
if (cmd->varprefix == NULL)
cmd->varprefix = pg_malloc0(sizeof(char *) * new_max);
else
{
cmd->varprefix = pg_realloc(cmd->varprefix, sizeof(char *) * new_max);
memset(cmd->varprefix + cmd->varprefix_max, 0,
sizeof(char *) * (new_max - cmd->varprefix_max));
}
cmd->varprefix_max = new_max;
}
/*
* Parse a backslash command; return a Command struct, or NULL if comment
*
@ -4177,7 +4397,6 @@ process_backslash_command(PsqlScanState sstate, const char *source)
/* Allocate and initialize Command structure */
my_command = (Command *) pg_malloc0(sizeof(Command));
my_command->command_num = num_commands++;
my_command->type = META_COMMAND;
my_command->argc = 0;
initSimpleStats(&my_command->stats);
@ -4201,7 +4420,7 @@ process_backslash_command(PsqlScanState sstate, const char *source)
if (my_command->meta == META_SET)
{
if (!expr_lex_one_word(sstate, &word_buf, &word_offset))
syntax_error(source, lineno, my_command->line, my_command->argv[0],
syntax_error(source, lineno, my_command->first_line, my_command->argv[0],
"missing argument", NULL, -1);
offsets[j] = word_offset;
@ -4222,10 +4441,11 @@ process_backslash_command(PsqlScanState sstate, const char *source)
my_command->expr = expr_parse_result;
/* Save line, trimming any trailing newline */
my_command->line = expr_scanner_get_substring(sstate,
start_offset,
expr_scanner_offset(sstate),
true);
my_command->first_line =
expr_scanner_get_substring(sstate,
start_offset,
expr_scanner_offset(sstate),
true);
expr_scanner_finish(yyscanner);
@ -4238,7 +4458,7 @@ process_backslash_command(PsqlScanState sstate, const char *source)
while (expr_lex_one_word(sstate, &word_buf, &word_offset))
{
if (j >= MAX_ARGS)
syntax_error(source, lineno, my_command->line, my_command->argv[0],
syntax_error(source, lineno, my_command->first_line, my_command->argv[0],
"too many arguments", NULL, -1);
offsets[j] = word_offset;
@ -4247,19 +4467,20 @@ process_backslash_command(PsqlScanState sstate, const char *source)
}
/* Save line, trimming any trailing newline */
my_command->line = expr_scanner_get_substring(sstate,
start_offset,
expr_scanner_offset(sstate),
true);
my_command->first_line =
expr_scanner_get_substring(sstate,
start_offset,
expr_scanner_offset(sstate),
true);
if (my_command->meta == META_SLEEP)
{
if (my_command->argc < 2)
syntax_error(source, lineno, my_command->line, my_command->argv[0],
syntax_error(source, lineno, my_command->first_line, my_command->argv[0],
"missing argument", NULL, -1);
if (my_command->argc > 3)
syntax_error(source, lineno, my_command->line, my_command->argv[0],
syntax_error(source, lineno, my_command->first_line, my_command->argv[0],
"too many arguments", NULL,
offsets[3] - start_offset);
@ -4288,7 +4509,7 @@ process_backslash_command(PsqlScanState sstate, const char *source)
if (pg_strcasecmp(my_command->argv[2], "us") != 0 &&
pg_strcasecmp(my_command->argv[2], "ms") != 0 &&
pg_strcasecmp(my_command->argv[2], "s") != 0)
syntax_error(source, lineno, my_command->line, my_command->argv[0],
syntax_error(source, lineno, my_command->first_line, my_command->argv[0],
"unrecognized time unit, must be us, ms or s",
my_command->argv[2], offsets[2] - start_offset);
}
@ -4296,25 +4517,31 @@ process_backslash_command(PsqlScanState sstate, const char *source)
else if (my_command->meta == META_SETSHELL)
{
if (my_command->argc < 3)
syntax_error(source, lineno, my_command->line, my_command->argv[0],
syntax_error(source, lineno, my_command->first_line, my_command->argv[0],
"missing argument", NULL, -1);
}
else if (my_command->meta == META_SHELL)
{
if (my_command->argc < 2)
syntax_error(source, lineno, my_command->line, my_command->argv[0],
syntax_error(source, lineno, my_command->first_line, my_command->argv[0],
"missing command", NULL, -1);
}
else if (my_command->meta == META_ELSE || my_command->meta == META_ENDIF)
{
if (my_command->argc != 1)
syntax_error(source, lineno, my_command->line, my_command->argv[0],
syntax_error(source, lineno, my_command->first_line, my_command->argv[0],
"unexpected argument", NULL, -1);
}
else if (my_command->meta == META_CSET || my_command->meta == META_GSET)
{
if (my_command->argc > 2)
syntax_error(source, lineno, my_command->first_line, my_command->argv[0],
"too many arguments", NULL, -1);
}
else
{
/* my_command->meta == META_NONE */
syntax_error(source, lineno, my_command->line, my_command->argv[0],
syntax_error(source, lineno, my_command->first_line, my_command->argv[0],
"invalid command", NULL, -1);
}
@ -4393,6 +4620,9 @@ ParseScript(const char *script, const char *desc, int weight)
PQExpBufferData line_buf;
int alloc_num;
int index;
bool saw_cset = false;
int lineno;
int start_offset;
#define COMMANDS_ALLOC_NUM 128
alloc_num = COMMANDS_ALLOC_NUM;
@ -4416,6 +4646,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);
@ -4425,45 +4656,115 @@ ParseScript(const char *script, const char *desc, int weight)
{
PsqlScanResult sr;
promptStatus_t prompt;
Command *command;
Command *command = NULL;
int semicolons;
resetPQExpBuffer(&line_buf);
lineno = expr_scanner_get_lineno(sstate, start_offset);
sr = psql_scan(sstate, &line_buf, &prompt);
/* If we collected a SQL command, process that */
command = process_sql_command(&line_buf, desc);
if (command)
{
ps.commands[index] = command;
index++;
semicolons = psql_scan_get_escaped_semicolons(sstate);
if (index >= alloc_num)
{
alloc_num += COMMANDS_ALLOC_NUM;
ps.commands = (Command **)
pg_realloc(ps.commands, sizeof(Command *) * alloc_num);
}
if (saw_cset)
{
/* the previous multi-line command ended with \cset */
append_sql_command(ps.commands[index - 1], line_buf.data,
semicolons + 1);
saw_cset = false;
}
else
{
/* If we collected a new SQL command, process that */
command = create_sql_command(&line_buf, desc, semicolons + 1);
/* store new command */
if (command)
ps.commands[index++] = command;
}
/* If we reached a backslash, process that */
if (sr == PSCAN_BACKSLASH)
{
command = process_backslash_command(sstate, desc);
if (command)
{
ps.commands[index] = command;
index++;
if (index >= alloc_num)
/*
* If this is gset/cset, merge into the preceding command. (We
* don't use a command slot in this case).
*/
if (command->meta == META_CSET ||
command->meta == META_GSET)
{
alloc_num += COMMANDS_ALLOC_NUM;
ps.commands = (Command **)
pg_realloc(ps.commands, sizeof(Command *) * alloc_num);
int cindex;
Command *cmd;
/*
* If \cset is seen, set flag to leave the command pending
* for the next iteration to process.
*/
saw_cset = command->meta == META_CSET;
if (index == 0)
syntax_error(desc, lineno, NULL, NULL,
"\\gset/cset cannot start a script",
NULL, -1);
cmd = ps.commands[index - 1];
if (cmd->type != SQL_COMMAND)
syntax_error(desc, lineno, NULL, NULL,
"\\gset/cset must follow a SQL command",
cmd->first_line, -1);
/* this {g,c}set applies to the previous query */
cindex = cmd->nqueries - 1;
/*
* now that we know there's a {g,c}set in this command,
* allocate space for the variable name prefix array.
*/
allocate_command_varprefix(cmd, cmd->nqueries);
/*
* Don't allow the previous command be a gset/cset; that
* would make no sense.
*/
if (cmd->varprefix[cindex] != NULL)
syntax_error(desc, lineno, NULL, NULL,
"\\gset/cset cannot follow one another",
NULL, -1);
/* get variable prefix */
if (command->argc <= 1 || command->argv[1][0] == '\0')
cmd->varprefix[cindex] = ""; /* avoid strdup */
else
cmd->varprefix[cindex] = pg_strdup(command->argv[1]);
/* cleanup unused command */
free_command(command);
continue;
}
/* Attach any other backslash command as a new command */
ps.commands[index++] = command;
}
}
/*
* Since we used a command slot, allocate more if needed. Note we
* always allocate one more in order to accomodate the NULL terminator
* below.
*/
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;
@ -4819,7 +5120,7 @@ printResults(TState *threads, StatsData *total, instr_time total_time,
printf(" %11.3f %s\n",
(cstats->count > 0) ?
1000.0 * cstats->sum / cstats->count : 0.0,
(*commands)->line);
(*commands)->first_line);
}
}
}
@ -5286,28 +5587,18 @@ main(int argc, char **argv)
internal_script_used = true;
}
/* if not simple query mode, parse the script(s) to find parameters */
if (querymode != QUERY_SIMPLE)
{
for (i = 0; i < num_scripts; i++)
{
Command **commands = sql_script[i].commands;
int j;
for (j = 0; commands[j] != NULL; j++)
{
if (commands[j]->type != SQL_COMMAND)
continue;
if (!parseQuery(commands[j]))
exit(1);
}
}
}
/* compute total_weight */
/* complete SQL command initialization and compute total weight */
for (i = 0; i < num_scripts; i++)
{
Command **commands = sql_script[i].commands;
for (int j = 0; commands[j] != NULL; j++)
if (commands[j]->type == SQL_COMMAND)
postprocess_sql_command(commands[j]);
/* cannot overflow: weight is 32b, total_weight 64b */
total_weight += sql_script[i].weight;
}
if (total_weight == 0 && !is_init_mode)
{

80
src/bin/pgbench/t/001_pgbench_with_server.pl
View File

@ -528,6 +528,48 @@ pgbench(
}
});
# working \gset and \cset
pgbench(
'-t 1', 0,
[ qr{type: .*/001_pgbench_gset_and_cset}, qr{processed: 1/1} ],
[ qr{command=3.: int 0\b},
qr{command=5.: int 1\b},
qr{command=6.: int 2\b},
qr{command=8.: int 3\b},
qr{command=9.: int 4\b},
qr{command=10.: int 5\b},
qr{command=12.: int 6\b},
qr{command=13.: int 7\b},
qr{command=14.: int 8\b},
qr{command=16.: int 9\b} ],
'pgbench gset and cset commands',
{ '001_pgbench_gset_and_cset' => q{-- test gset and cset
-- no columns
SELECT \gset
-- one value
SELECT 0 AS i0 \gset
\set i debug(:i0)
-- two values
SELECT 1 AS i1, 2 AS i2 \gset
\set i debug(:i1)
\set i debug(:i2)
-- cset & gset to follow
SELECT :i2 + 1 AS i3, :i2 * :i2 AS i4 \cset
SELECT 5 AS i5 \gset
\set i debug(:i3)
\set i debug(:i4)
\set i debug(:i5)
-- with prefix
SELECT 6 AS i6, 7 AS i7 \cset x_
SELECT 8 AS i8 \gset y_
\set i debug(:x_i6)
\set i debug(:x_i7)
\set i debug(:y_i8)
-- overwrite existing variable
SELECT 0 AS i9, 9 AS i9 \gset
\set i debug(:i9)
} });
# trigger many expression errors
my @errors = (
@ -735,21 +777,45 @@ SELECT LEAST(:i, :i, :i, :i, :i, :i, :i, :i, :i, :i, :i);
[qr{invalid command .* "nosuchcommand"}], q{\nosuchcommand}
],
[ 'misc empty script', 1, [qr{empty command list for script}], q{} ],
[
'bad boolean', 2,
[qr{malformed variable.*trueXXX}], q{\set b :badtrue or true}
],);
[ 'bad boolean', 2,
[qr{malformed variable.*trueXXX}], q{\set b :badtrue or true} ],
# GSET & CSET
[ 'gset no row', 2,
[qr{expected one row, got 0\b}], q{SELECT WHERE FALSE \gset} ],
[ 'cset no row', 2,
[qr{expected one row, got 0\b}], q{SELECT WHERE FALSE \cset
SELECT 1 AS i\gset}, 1 ],
[ 'gset alone', 1, [qr{gset/cset cannot start a script}], q{\gset} ],
[ 'gset no SQL', 1,
[qr{gset/cset must follow a SQL command}], q{\set i +1
\gset} ],
[ 'gset too many arguments', 1,
[qr{too many arguments}], q{SELECT 1 \gset a b} ],
[ 'gset after gset', 1,
[qr{gset/cset cannot follow one another}], q{SELECT 1 AS i \gset
\gset} ],
[ 'gset non SELECT', 2,
[qr{expected one row, got 0}],
q{DROP TABLE IF EXISTS no_such_table \gset} ],
[ 'gset bad default name', 2,
[qr{error storing into variable \?column\?}],
q{SELECT 1 \gset} ],
[ 'gset bad name', 2,
[qr{error storing into variable bad name!}],
q{SELECT 1 AS "bad name!" \gset} ],
);
for my $e (@errors)
{
my ($name, $status, $re, $script) = @$e;
my ($name, $status, $re, $script, $no_prepare) = @$e;
$status != 0 or die "invalid expected status for test \"$name\"";
my $n = '001_pgbench_error_' . $name;
$n =~ s/ /_/g;
pgbench(
'-n -t 1 -M prepared -Dfoo=bla -Dnull=null -Dtrue=true -Done=1 -Dzero=0.0 ' .
'-Dbadtrue=trueXXX -Dmaxint=9223372036854775807 -Dminint=-9223372036854775808',
'-n -t 1 -Dfoo=bla -Dnull=null -Dtrue=true -Done=1 -Dzero=0.0 -Dbadtrue=trueXXX' .
' -Dmaxint=9223372036854775807 -Dminint=-9223372036854775808' .
($no_prepare ? '' : ' -M prepared'),
$status,
[ $status == 1 ? qr{^$} : qr{processed: 0/1} ],
$re,

24
src/fe_utils/psqlscan.l
View File

@ -693,8 +693,15 @@ other .
* substitution. We want these before {self}, also.
*/
"\\"[;:] {
/* Force a semicolon or colon into the query buffer */
"\\"; {
/* Count semicolons in compound commands */
cur_state->escaped_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);
}
@ -1065,6 +1072,9 @@ psql_scan(PsqlScanState state,
/* Set current output target */
state->output_buf = query_buf;
/* Reset number of escaped semicolons seen */
state->escaped_semicolons = 0;
/* Set input source */
if (state->buffer_stack != NULL)
yy_switch_to_buffer(state->buffer_stack->buf, state->scanner);
@ -1208,6 +1218,16 @@ psql_scan_reset(PsqlScanState state)
state->dolqstart = NULL;
}
/*
* Return the number of escaped semicolons in the lexed string seen by the
* previous psql_scan call.
*/
int
psql_scan_get_escaped_semicolons(PsqlScanState state)
{
return state->escaped_semicolons;
}
/*
* Reselect this lexer (psqlscan.l) after using another one.
*

2
src/include/fe_utils/psqlscan.h
View File

@ -90,6 +90,8 @@ extern PsqlScanResult psql_scan(PsqlScanState state,
extern void psql_scan_reset(PsqlScanState state);
extern int psql_scan_get_escaped_semicolons(PsqlScanState state);
extern void psql_scan_reselect_sql_lexer(PsqlScanState state);
extern bool psql_scan_in_quote(PsqlScanState state);

1
src/include/fe_utils/psqlscan_int.h
View File

@ -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 escaped_semicolons; /* number of embedded (\;) semicolons */
char *dolqstart; /* current $foo$ quote start string */
/*