Hello Stephen,

Here's that review.

Thanks for the review.

   <variablelist>
+   <varlistentry id='pgbench-metacommand-gset'>
+    <term>
+     <literal>\cset [<replaceable>prefix</replaceable>]</literal> or
+     <literal>\gset [<replaceable>prefix</replaceable>]</literal>
+    </term>

Seems like this should really be moved down below the section for
'\set'.

Dunno.

I put them there because it is in alphabetical order (for cset at least) and because "set" documentation is heavy about expressions (operators, functions, constants, ...) which have nothing to do with cset/gset, so I felt that having them clearly separated and in abc order was best.

-       char       *line;                       /* text of command line */
+       char       *line;                       /* first line for short display 
*/
+       char       *lines;                      /* full multi-line text of 
command */

Not really a fan of such closely-named variables...  Maybe first_line
instead?

Indeed, looks better.

+                       case PGRES_TUPLES_OK:
+                               if (gset[compound] != NULL)

Probably be good to add some comments here, eh:

/*
* The results are intentionally thrown away if we aren't under a gset
* call.
*/

I added a (shorter) comment.

+                                       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;
+                                       }

Might be interesting to support mutli-row (or no row?) returns in the
future, but not something this patch needs to do, so this looks fine to
me.

It could match PL/pgSQL's INTO, that is first row or NULL if none.

+
+                                               /* store result as a string */
+                                               if (!putVariable(st, "gset", 
varname,
+                                                                               
 PQgetvalue(res, 0, f)))
+                                               {
+                                                       /* internal error, 
should it rather abort? */

I'm a bit on the fence about if we should abort in this case or not.  A
failure here seems likely to happen consistently (whereas the ntuples
case might be a fluke of some kind), which tends to make me think we
should abort, but still thinking about it.

I'm also still thinking about it:-) For me it is an abort, but there is this whole "return false" internal error management in pgbench the purpose of which fails me a little bit, so I stick to that anyway.

+                                               if (*gset[compound] != '\0')
+                                                       free(varname);

A comment here, and above where we're assigning the result of the
psprintf(), to varname probably wouldn't hurt, explaining that the
variable is sometimes pointing into the query result structure and
sometimes not...

I added two comments to avoid a malloc/free when there are no prefixes. ISTM that although it might be a border-line over-optimization case, it is a short one, the free is a few lines away, so it might be ok.

+                               /* read and discard the query results */

That comment doesn't feel quite right now. ;)

Indeed. Changed with "store or discard".


-       /* We don't want to scribble on cmd->argv[0] until done */
-       sql = pg_strdup(cmd->argv[0]);
+       sql = pg_strdup(cmd->lines);

The function-header comment for parseQuery() could really stand to be
improved.

Indeed.

+                               /* merge gset variants into preceeding SQL 
command */
+                               if (pg_strcasecmp(bs_cmd, "gset") == 0 ||
+                                       pg_strcasecmp(bs_cmd, "cset") == 0)
+                               {
+                                                                        "\\gset 
cannot start a script",
+                                                                        "\\gset 
must follow a SQL command",
+                                                                        "\\gset 
cannot follow one another",

These errors should probably be '\\gset and \\cset' or similar, no?
Since we fall into this for both..

Indeed.

Attached an improved version, mostly comments and error message fixes.

I have not changed the 1 row constraint for now. Could be an later extension.

If this patch get through, might be handy for simplifying tests in
current pgbench submissions, especially the "error handling" one.

--
Fabien.
diff --git a/doc/src/sgml/ref/pgbench.sgml b/doc/src/sgml/ref/pgbench.sgml
index 41d9030..dffbae2 100644
--- a/doc/src/sgml/ref/pgbench.sgml
+++ b/doc/src/sgml/ref/pgbench.sgml
@@ -942,6 +942,51 @@ pgbench <optional> <replaceable>options</replaceable> 
</optional> <replaceable>d
   </para>
 
   <variablelist>
+   <varlistentry id='pgbench-metacommand-gset'>
+    <term>
+     <literal>\cset [<replaceable>prefix</replaceable>]</literal> or
+     <literal>\gset [<replaceable>prefix</replaceable>]</literal>
+    </term>
+
+    <listitem>
+     <para>
+      These commands may be used to end SQL queries, replacing a semicolon.
+      <literal>\cset</literal> replaces an embedded semicolon 
(<literal>\;</literal>) within
+      a compound SQL command, and <literal>\gset</literal> replaces a final
+      (<literal>;</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</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>one</replaceable>, <replaceable>two</replaceable> and
+      <replaceable>p_three</replaceable> 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</literal> and <literal>\gset</literal> commands do 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>
diff --git a/src/bin/pgbench/pgbench.c b/src/bin/pgbench/pgbench.c
index fd18568..47bc40c 100644
--- a/src/bin/pgbench/pgbench.c
+++ b/src/bin/pgbench/pgbench.c
@@ -437,12 +437,15 @@ static const char *QUERYMODE[] = {"simple", "extended", 
"prepared"};
 
 typedef struct
 {
-       char       *line;                       /* text of command line */
+       char       *first_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) */
        MetaCommand meta;                       /* meta command identifier, or 
META_NONE */
        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;
@@ -1595,6 +1598,107 @@ valueTruth(PgBenchValue *pval)
        }
 }
 
+/* read all responses from backend, storing into variable or discarding */
+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/cset 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 if 
required */
+                                       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);
+                                               /* prefix varname if required, 
will be freed below */
+                                               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;
+                                               }
+
+                                               /* free varname only if 
allocated because of prefix */
+                                               if (*gset[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)
@@ -2670,7 +2774,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;
@@ -3146,26 +3249,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, "SQL", 
PQerrorMessage(st->con));
-                                               PQclear(res);
-                                               st->state = CSTATE_ABORTED;
-                                               break;
-                               }
+                               /* read, store or discard the query results */
+                               if (read_response(st, command->gset))
+                                       st->state = CSTATE_END_COMMAND;
+                               else
+                                       st->state = CSTATE_ABORTED;
+
                                break;
 
                                /*
@@ -3820,7 +3909,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)
@@ -3828,8 +3917,7 @@ parseQuery(Command *cmd)
        char       *sql,
                           *p;
 
-       /* We don't want to scribble on cmd->argv[0] until done */
-       sql = pg_strdup(cmd->argv[0]);
+       sql = pg_strdup(cmd->lines);
 
        cmd->argc = 1;
 
@@ -3853,7 +3941,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);
                        pg_free(name);
                        return false;
                }
@@ -3865,7 +3953,7 @@ parseQuery(Command *cmd)
                cmd->argc++;
        }
 
-       pg_free(cmd->argv[0]);
+       Assert(cmd->argv[0] == NULL);
        cmd->argv[0] = sql;
        return true;
 }
@@ -3924,22 +4012,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))
@@ -3959,35 +4035,111 @@ 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->meta = META_NONE;
+       my_command->argc = 0;
+       my_command->compound = compounds;
+       my_command->gset = pg_malloc0(sizeof(char *) * (compounds+1));
        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;
+       my_command->lines = pg_strdup(p);
+
+       return my_command;
+}
+
+/*
+ * append "more" text to current compound command which may have been
+ * interrupted by \cset.
+ */
+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;
+
+       /* append command text, embedding a ';' in place of the \cset */
+       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);
+
+       /* update number of compounds and extend array of prefixes */
+       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.
+        * the "line" label for display.
         */
+       p = my_command->lines;
        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';
+               my_command->first_line = pg_malloc(nlpos - p + 1);
+               memcpy(my_command->first_line, p, nlpos - p);
+               my_command->first_line[nlpos - p] = '\0';
        }
        else
-               my_command->line = pg_strdup(p);
+               my_command->first_line = pg_strdup(p);
 
-       return my_command;
+       /* parse query if necessary */
+       switch (querymode)
+       {
+               case QUERY_SIMPLE:
+                       my_command->argv[0] = my_command->lines;
+                       my_command->argc++;
+                       break;
+               case QUERY_EXTENDED:
+               case QUERY_PREPARED:
+                       if (!parseQuery(my_command))
+                               exit(1);
+                       break;
+               default:
+                       exit(1);
+       }
 }
 
 /*
@@ -4045,7 +4197,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;
@@ -4066,10 +4218,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);
 
@@ -4082,7 +4235,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;
@@ -4091,19 +4244,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);
 
@@ -4132,7 +4286,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);
                }
@@ -4140,25 +4294,32 @@ 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 (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->first_line, 
my_command->argv[0],
+                                                "at most one argument 
expected", 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);
        }
 
@@ -4235,6 +4396,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;
@@ -4258,6 +4422,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);
 
@@ -4267,31 +4432,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;
@@ -4306,6 +4468,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/cset 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/cset must follow a SQL command",
+                                                                        
sql_cmd->first_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/cset 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;
@@ -4661,7 +4884,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);
                                }
                        }
                }
@@ -5129,28 +5352,19 @@ 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 initializations and collect total weight */
        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)
+                               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)
        {
diff --git a/src/bin/pgbench/pgbench.h b/src/bin/pgbench/pgbench.h
index 6983865..c349477 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/bin/pgbench/t/001_pgbench_with_server.pl 
b/src/bin/pgbench/t/001_pgbench_with_server.pl
index be08b20..7f86df5 100644
--- a/src/bin/pgbench/t/001_pgbench_with_server.pl
+++ b/src/bin/pgbench/t/001_pgbench_with_server.pl
@@ -474,6 +474,48 @@ pgbench(
 \shell echo shell-echo-output
 } });
 
+# 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 = (
 
@@ -626,16 +668,44 @@ 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', 0, [qr{malformed variable.*trueXXX}], q{\set b 
:badtrue or true} ],
-    );
 
+       # GSET & CSET
+       [   'gset no row',                    0,
+               [qr{expecting one row, got 0\b}], q{SELECT WHERE FALSE \gset} ],
+       [   'cset no row',                    0,
+               [qr{expecting 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 after gset',                        1,
+           [qr{gset/cset cannot follow one another}], q{SELECT 1 AS i \gset
+\gset} ],
+       [   'gset non SELECT',
+               0,
+               [qr{gset/cset expects a row}],
+               q{DROP TABLE IF EXISTS no_such_table \gset} ],
+       [   'gset bad default name',
+               0,
+               [qr{error storing into var \?column\?}],
+               q{SELECT 1 \gset} ],
+       [   'gset bad name',
+               0,
+               [qr{error storing into var 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;
        my $n = '001_pgbench_error_' . $name;
        $n =~ s/ /_/g;
        pgbench(
-               '-n -t 1 -Dfoo=bla -Dnull=null -Dtrue=true -Done=1 -Dzero=0.0 
-Dbadtrue=trueXXX -M prepared',
+               '-n -t 1 -Dfoo=bla -Dnull=null -Dtrue=true -Done=1 -Dzero=0.0 
-Dbadtrue=trueXXX' .
+                       ($no_prepare ? '' : ' -M prepared'),
                $status,
                [ $status ? qr{^$} : qr{processed: 0/1} ],
                $re,
diff --git a/src/fe_utils/psqlscan.l b/src/fe_utils/psqlscan.l
index 1cc587b..efca525 100644
--- a/src/fe_utils/psqlscan.l
+++ b/src/fe_utils/psqlscan.l
@@ -680,8 +680,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 0be0db6..8ef4abd 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 */
 
        /*

Reply via email to