Here are my proposed final changes. I noticed that you were allocating
the prefixes for all cases even when there were no cset/gset in the
command; I changed it to delay the allocation until needed. I also
noticed you were skipping the Meta enum dance for no good reason; added
that makes conditionals simpler. The read_response routine seemed
misplaced and I gave it a name in a style closer to the rest of the
pgbench code. Also, you were missing to free the ->lines pqexpbuffer
when the command is discarded. I grant that the free_command() stuff
might be bogus since it's only tested with a command that's barely
initialized, but it seems better to make it bogus in this way (fixable
if we ever extend its use) than to forever leak memory silently.
I didn't test this beyond running "make check".
--
Álvaro Herrera https://www.2ndQuadrant.com/
PostgreSQL Development, 24x7 Support, Remote DBA, Training & Services
diff --git a/src/bin/pgbench/pgbench.c b/src/bin/pgbench/pgbench.c
index 4b8b6ba10f..3e111876fb 100644
--- a/src/bin/pgbench/pgbench.c
+++ b/src/bin/pgbench/pgbench.c
@@ -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,19 +501,39 @@ typedef enum QueryMode
static QueryMode querymode = QUERY_SIMPLE;
static const char *QUERYMODE[] = {"simple", "extended", "prepared"};
+/*
+ * 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
{
- PQExpBufferData lines; /* raw command text (possibly multi-line) */
- char *first_line; /* first line for short display */
- 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 */
- int nqueries; /* number of compounds within an SQL command */
- int prefix_size; /* allocated size of prefix, >= nqueries */
- char **prefix; /* per-compound command prefix for [cg]set */
- 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
@@ -589,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);
@@ -1734,107 +1757,6 @@ valueTruth(PgBenchValue *pval)
}
}
-/* read all responses from backend, storing into variable or discarding */
-static bool
-read_response(CState *st, char **prefix)
-{
- 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 (prefix[compound] != NULL)
- {
- fprintf(stderr,
- "client %d file %d command %d compound %d: "
- "\\gset/cset expects a row\n",
- st->id, st->use_file, st->command, compound);
- st->ecnt++;
- return false;
- }
- break; /* OK */
-
- case PGRES_TUPLES_OK:
- if (prefix[compound] != NULL)
- {
- /* store result into variables if required */
- int ntuples = PQntuples(res),
- nfields = PQnfields(res);
-
- 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 (int f = 0; f < nfields; f++)
- {
- char *varname = PQfname(res, f);
-
- /* prefix varname if required, will be freed below */
- if (*prefix[compound] != '\0')
- varname = psprintf("%s%s", prefix[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;
- }
-
- /* free varname only if allocated because of prefix */
- if (*prefix[compound] != '\0')
- free(varname);
- }
- }
- /* otherwise the result is simply thrown away by PQclear below */
- 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)
@@ -2672,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;
@@ -2900,6 +2826,109 @@ sendCommand(CState *st, Command *command)
}
/*
+ * 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.
*/
@@ -3242,8 +3271,8 @@ advanceConnectionState(TState *thread, CState *st, StatsData *agg)
if (PQisBusy(st->con))
return; /* don't have the whole result yet */
- /* read, store or discard the query results */
- if (read_response(st, sql_script[st->use_file].commands[st->command]->prefix))
+ /* 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;
@@ -4201,8 +4230,8 @@ skip_sql_comments(char *sql_command)
* 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.
+ * 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)
@@ -4217,29 +4246,44 @@ create_sql_command(PQExpBuffer buf, const char *source, int numqueries)
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->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->prefix_size = numqueries + 8;
- my_command->prefix = pg_malloc0(sizeof(char *) * my_command->prefix_size);
+ my_command->varprefix = NULL; /* allocated later, if needed */
+ my_command->varprefix_max = 0;
my_command->expr = NULL;
initSimpleStats(&my_command->stats);
return my_command;
}
+/* Free a Command structure and associated data */
+static void
+free_command(Command *command)
+{
+ pg_free(command->lines.data);
+ 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)
+ pg_free(command->varprefix);
+ if (command->expr)
+ pg_free(command->expr);
+ pg_free(command);
+}
+
/*
- * append "more" text to current compound command which may have been
- * interrupted by \cset.
+ * 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)
{
- int nq;
-
Assert(my_command->type == SQL_COMMAND && my_command->lines.len > 0);
more = skip_sql_comments(more);
@@ -4249,20 +4293,7 @@ append_sql_command(Command *my_command, char *more, int numqueries)
/* append command text, embedding a ';' in place of the \cset */
appendPQExpBuffer(&my_command->lines, ";\n%s", more);
- /* update number of commands and extend array of prefixes */
- nq = my_command->nqueries + numqueries;
- if (nq > my_command->prefix_size)
- {
- my_command->prefix_size = 2 * nq;
- my_command->prefix =
- pg_realloc(my_command->prefix,
- sizeof(char *) * my_command->prefix_size);
- }
-
- /* fill with NULL */
- memset(my_command->prefix + my_command->nqueries, 0,
- sizeof(char *) * numqueries);
- my_command->nqueries = nq;
+ my_command->nqueries += numqueries;
}
/*
@@ -4272,7 +4303,7 @@ append_sql_command(Command *my_command, char *more, int numqueries)
static void
postprocess_sql_command(Command *my_command)
{
- char buffer[128];
+ char buffer[128];
Assert(my_command->type == SQL_COMMAND);
@@ -4299,6 +4330,35 @@ postprocess_sql_command(Command *my_command)
}
/*
+ * 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;
+
+ /* and 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
*
* At call, we have scanned only the initial backslash.
@@ -4464,12 +4524,11 @@ 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 (pg_strcasecmp(my_command->argv[0], "gset") == 0 ||
- pg_strcasecmp(my_command->argv[0], "cset") == 0)
+ 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],
- "at most one argument expected", NULL, -1);
+ "too many arguments", NULL, -1);
}
else
{
@@ -4553,7 +4612,7 @@ ParseScript(const char *script, const char *desc, int weight)
PQExpBufferData line_buf;
int alloc_num;
int index;
- bool is_compound = false;
+ bool saw_cset = false;
int lineno;
int start_offset;
@@ -4596,14 +4655,15 @@ ParseScript(const char *script, const char *desc, int weight)
lineno = expr_scanner_get_lineno(sstate, start_offset);
sr = psql_scan(sstate, &line_buf, &prompt);
+
semicolons = psql_scan_get_escaped_semicolons(sstate);
- if (is_compound)
+ if (saw_cset)
{
/* the previous multi-line command ended with \cset */
append_sql_command(ps.commands[index - 1], line_buf.data,
semicolons + 1);
- is_compound = false;
+ saw_cset = false;
}
else
{
@@ -4612,81 +4672,91 @@ ParseScript(const char *script, const char *desc, int weight)
/* store new command */
if (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);
- }
- }
+ ps.commands[index++] = command;
}
+ /* If we reached a backslash, process that */
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)
+ /*
+ * 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)
{
- bool is_gset = bs_cmd[0] == 'g';
int cindex;
- Command *sql_cmd;
+ Command *cmd;
- is_compound = !is_gset;
+ /*
+ * 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);
- sql_cmd = ps.commands[index - 1];
+ cmd = ps.commands[index - 1];
- if (sql_cmd->type != SQL_COMMAND)
+ if (cmd->type != SQL_COMMAND)
syntax_error(desc, lineno, NULL, NULL,
"\\gset/cset must follow a SQL command",
- sql_cmd->first_line, -1);
+ cmd->first_line, -1);
- /* this \gset applies to the last sub-command */
- cindex = sql_cmd->nqueries - 1;
+ /* this {g,c}set applies to the previous query */
+ cindex = cmd->nqueries - 1;
- if (sql_cmd->prefix[cindex] != NULL)
+ /*
+ * 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')
- sql_cmd->prefix[cindex] = "";
+ cmd->varprefix[cindex] = "";
else
- sql_cmd->prefix[cindex] = command->argv[1];
+ cmd->varprefix[cindex] = pg_strdup(command->argv[1]);
- /* cleanup unused backslash command */
- pg_free(command);
- }
- else /* any other backslash command is a Command */
- {
- ps.commands[index] = command;
- index++;
+ /* cleanup unused command */
+ free_command(command);
- if (index >= alloc_num)
- {
- alloc_num += COMMANDS_ALLOC_NUM;
- ps.commands = (Command **)
- pg_realloc(ps.commands, sizeof(Command *) * alloc_num);
- }
+ 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;
diff --git a/src/bin/pgbench/t/001_pgbench_with_server.pl b/src/bin/pgbench/t/001_pgbench_with_server.pl
index 0eaf18c2f8..c87748086a 100644
--- a/src/bin/pgbench/t/001_pgbench_with_server.pl
+++ b/src/bin/pgbench/t/001_pgbench_with_server.pl
@@ -782,27 +782,27 @@ SELECT LEAST(:i, :i, :i, :i, :i, :i, :i, :i, :i, :i, :i);
# GSET & CSET
[ 'gset no row', 2,
- [qr{expecting one row, got 0\b}], q{SELECT WHERE FALSE \gset} ],
+ [qr{expected one row, got 0\b}], q{SELECT WHERE FALSE \gset} ],
[ 'cset no row', 2,
- [qr{expecting one row, got 0\b}], q{SELECT WHERE FALSE \cset
+ [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 args', 1,
- [qr{at most one argument expected}], q{SELECT 1 \gset a b} ],
+ [ '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{gset/cset expects a row}],
+ [qr{expected one row, got 0}],
q{DROP TABLE IF EXISTS no_such_table \gset} ],
[ 'gset bad default name', 2,
- [qr{error storing into var \?column\?}],
+ [qr{error storing into variable \?column\?}],
q{SELECT 1 \gset} ],
[ 'gset bad name', 2,
- [qr{error storing into var bad name!}],
+ [qr{error storing into variable bad name!}],
q{SELECT 1 AS "bad name!" \gset} ],
);
