Hello Stephen,

Attached is v18, another basic rebase after some perl automatic reindentation.

--
Fabien.
diff --git a/doc/src/sgml/ref/pgbench.sgml b/doc/src/sgml/ref/pgbench.sgml
index e4b37dd..28a1387 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 78b8f17..7bf4331 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)
@@ -2679,7 +2783,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;
@@ -3172,26 +3275,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;
 
                                /*
@@ -3846,7 +3935,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)
@@ -3854,8 +3943,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;
 
@@ -3879,7 +3967,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;
                }
@@ -3891,7 +3979,7 @@ parseQuery(Command *cmd)
                cmd->argc++;
        }
 
-       pg_free(cmd->argv[0]);
+       Assert(cmd->argv[0] == NULL);
        cmd->argv[0] = sql;
        return true;
 }
@@ -3950,22 +4038,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))
@@ -3985,35 +4061,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);
+       }
 }
 
 /*
@@ -4071,7 +4223,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;
@@ -4092,10 +4244,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);
 
@@ -4108,7 +4261,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;
@@ -4117,19 +4270,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);
 
@@ -4158,7 +4312,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);
                }
@@ -4166,25 +4320,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);
        }
 
@@ -4263,6 +4424,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;
@@ -4286,6 +4450,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);
 
@@ -4295,31 +4460,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;
@@ -4334,6 +4496,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;
@@ -4689,7 +4912,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);
                                }
                        }
                }
@@ -5159,28 +5382,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 947f13d..c1e86da 100644
--- a/src/bin/pgbench/t/001_pgbench_with_server.pl
+++ b/src/bin/pgbench/t/001_pgbench_with_server.pl
@@ -484,6 +484,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 = (
 
@@ -618,16 +660,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',                     0,
-               [qr{malformed variable.*trueXXX}], q{\set b :badtrue or true} 
],);
+               [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