Attached a v26 with I hope everything ok, but for the documentation that I'm unsure how to improve.

Attached v27 is the same but with an improved documentation where full examples, with and without prefix, are provided for both cset & gset.

--
Fabien.
diff --git a/doc/src/sgml/ref/pgbench.sgml b/doc/src/sgml/ref/pgbench.sgml
index b5e3a62a33..cc369c423e 100644
--- a/doc/src/sgml/ref/pgbench.sgml
+++ b/doc/src/sgml/ref/pgbench.sgml
@@ -954,6 +954,91 @@ pgbench <optional> <replaceable>options</replaceable> </optional> <replaceable>d
   </para>
 
   <variablelist>
+   <varlistentry id='pgbench-metacommand-cset'>
+    <term>
+     <literal>\cset [<replaceable>prefix</replaceable>]</literal>
+    </term>
+
+    <listitem>
+     <para>
+      This command may be used to end SQL queries, replacing an embedded
+      semicolon (<literal>\;</literal>) within a compound SQL command.
+     </para>
+
+     <para>
+      When this command is used, the preceding SQL query is expected to
+      return one row, the columns of which are stored into variables named after
+      column names, and prefixed with <replaceable>prefix</replaceable> if provided.
+     </para>
+
+     <para>
+      The following example sends four queries as one compound SQL command,
+      inducing one message sent at the protocol level.
+      The result of the first query is stored into variable <replaceable>one</replaceable>,
+      the results of the third query are stored into variables <replaceable>z_three</replaceable>
+      and <replaceable>z_four</replaceable>,
+      whereas the results of the other queries are discarded.
+<programlisting>
+-- compound of four queries
+SELECT 1 AS one \cset
+SELECT 2 AS two \;
+SELECT 3 AS three, 4 AS four \cset z_
+SELECT 5;
+</programlisting>
+     </para>
+
+     <note>
+      <para>
+        <literal>\cset</literal> does not work when empty SQL queries appear
+        within a compound SQL command.
+      </para>
+     </note>
+    </listitem>
+   </varlistentry>
+
+   <varlistentry id='pgbench-metacommand-gset'>
+    <term>
+     <literal>\gset [<replaceable>prefix</replaceable>]</literal>
+    </term>
+
+    <listitem>
+     <para>
+      This commands may be used to end SQL queries, replacing a final semicolon
+      (<literal>;</literal>). 
+     </para>
+
+     <para>
+      When this command is used, the preceding SQL query is expected to
+      return one row, the columns of which are stored into variables named after
+      column names, and prefixed with <replaceable>prefix</replaceable> if provided.
+     </para>
+
+     <para>
+      The following example puts the final account balance from the first query
+      into variable <replaceable>abalance</replaceable>, and fills variables
+      <replaceable>p_two</replaceable> and <replaceable>p_three</replaceable> 
+      with integers from the last query.
+      The result of the second query is discarded.
+<programlisting>
+UPDATE pgbench_accounts
+  SET abalance = abalance + :delta
+  WHERE aid = :aid
+  RETURNING abalance \gset
+-- compound of two queries
+SELECT 1 \;
+SELECT 2 AS two, 3 AS three \gset p_
+</programlisting>
+     </para>
+
+     <note>
+      <para>
+        <literal>\gset</literal> does not work when empty SQL queries appear
+        within a compound SQL command.
+      </para>
+     </note>
+    </listitem>
+   </varlistentry>
+
    <varlistentry>
     <term><literal>\if</literal> <replaceable class="parameter">expression</replaceable></term>
     <term><literal>\elif</literal> <replaceable class="parameter">expression</replaceable></term>
diff --git a/src/bin/pgbench/pgbench.c b/src/bin/pgbench/pgbench.c
index 649297ae4f..4b8b6ba10f 100644
--- a/src/bin/pgbench/pgbench.c
+++ b/src/bin/pgbench/pgbench.c
@@ -499,14 +499,17 @@ typedef enum QueryMode
 static QueryMode querymode = QUERY_SIMPLE;
 static const char *QUERYMODE[] = {"simple", "extended", "prepared"};
 
-typedef struct
+typedef struct Command
 {
-	char	   *line;			/* text of command line */
-	int			command_num;	/* unique index of this Command struct */
+	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 */
 } Command;
@@ -521,7 +524,6 @@ typedef struct ParsedScript
 
 static ParsedScript sql_script[MAX_SCRIPTS];	/* SQL script files */
 static int	num_scripts;		/* number of scripts in sql_script[] */
-static int	num_commands = 0;	/* total number of Command structs */
 static int64 total_weight = 0;
 
 static int	debug = 0;			/* debug flag */
@@ -1732,6 +1734,107 @@ 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)
@@ -2862,8 +2965,6 @@ advanceConnectionState(TState *thread, CState *st, StatsData *agg)
 	 */
 	for (;;)
 	{
-		PGresult   *res;
-
 		switch (st->state)
 		{
 				/* Select transaction (script) to run.  */
@@ -3141,24 +3242,11 @@ advanceConnectionState(TState *thread, CState *st, StatsData *agg)
 				if (PQisBusy(st->con))
 					return;		/* don't have the whole result yet */
 
-				/* Read and discard the query result */
-				res = PQgetResult(st->con);
-				switch (PQresultStatus(res))
-				{
-					case PGRES_COMMAND_OK:
-					case PGRES_TUPLES_OK:
-					case PGRES_EMPTY_QUERY:
-						/* OK */
-						PQclear(res);
-						discard_response(st);
-						st->state = CSTATE_END_COMMAND;
-						break;
-					default:
-						commandFailed(st, "SQL", PQerrorMessage(st->con));
-						PQclear(res);
-						st->state = CSTATE_ABORTED;
-						break;
-				}
+				/* read, store or discard the query results */
+				if (read_response(st, sql_script[st->use_file].commands[st->command]->prefix))
+					st->state = CSTATE_END_COMMAND;
+				else
+					st->state = CSTATE_ABORTED;
 				break;
 
 				/*
@@ -3191,7 +3279,7 @@ advanceConnectionState(TState *thread, CState *st, StatsData *agg)
 
 					INSTR_TIME_SET_CURRENT_LAZY(now);
 
-					command	= sql_script[st->use_file].commands[st->command];
+					command = sql_script[st->use_file].commands[st->command];
 					/* XXX could use a mutex here, but we choose not to */
 					addToSimpleStats(&command->stats,
 									 INSTR_TIME_GET_DOUBLE(now) -
@@ -3235,9 +3323,8 @@ advanceConnectionState(TState *thread, CState *st, StatsData *agg)
 				st->state = CSTATE_CHOOSE_SCRIPT;
 
 				/*
-				 * Ensure that we always return on this point, so as to
-				 * avoid an infinite loop if the script only contains meta
-				 * commands.
+				 * Ensure that we always return on this point, so as to avoid
+				 * an infinite loop if the script only contains meta commands.
 				 */
 				return;
 
@@ -3976,7 +4063,7 @@ runInitSteps(const char *initialize_steps)
 
 /*
  * Replace :param with $n throughout the command's SQL text, which
- * is a modifiable string in cmd->argv[0].
+ * is a modifiable string in cmd->lines.
  */
 static bool
 parseQuery(Command *cmd)
@@ -3984,12 +4071,9 @@ parseQuery(Command *cmd)
 	char	   *sql,
 			   *p;
 
-	/* We don't want to scribble on cmd->argv[0] until done */
-	sql = pg_strdup(cmd->argv[0]);
-
 	cmd->argc = 1;
 
-	p = sql;
+	p = sql = pg_strdup(cmd->lines.data);
 	while ((p = strchr(p, ':')) != NULL)
 	{
 		char		var[13];
@@ -4009,7 +4093,7 @@ parseQuery(Command *cmd)
 		if (cmd->argc >= MAX_ARGS)
 		{
 			fprintf(stderr, "statement has too many arguments (maximum is %d): %s\n",
-					MAX_ARGS - 1, cmd->argv[0]);
+					MAX_ARGS - 1, cmd->lines.data);
 			pg_free(name);
 			return false;
 		}
@@ -4021,7 +4105,7 @@ parseQuery(Command *cmd)
 		cmd->argc++;
 	}
 
-	pg_free(cmd->argv[0]);
+	Assert(cmd->argv[0] == NULL);
 	cmd->argv[0] = sql;
 	return true;
 }
@@ -4081,21 +4165,16 @@ syntax_error(const char *source, int lineno,
 }
 
 /*
- * Parse a SQL command; return a Command struct, or NULL if it's a comment
- *
- * On entry, psqlscan.l has collected the command into "buf", so we don't
- * really need to do much here except check for comment and set up a
- * Command struct.
+ * Return a pointer to the start of the SQL command, after skipping over
+ * whitespace and "--" comments.
+ * If the end of the string is reached, return NULL.
  */
-static Command *
-process_sql_command(PQExpBuffer buf, const char *source)
+static char *
+skip_sql_comments(char *sql_command)
 {
-	Command    *my_command;
-	char	   *p;
-	char	   *nlpos;
+	char	   *p = sql_command;
 
 	/* Skip any leading whitespace, as well as "--" style comments */
-	p = buf->data;
 	for (;;)
 	{
 		if (isspace((unsigned char) *p))
@@ -4111,41 +4190,114 @@ process_sql_command(PQExpBuffer buf, const char *source)
 			break;
 	}
 
-	/* If there's nothing but whitespace and comments, we're done */
+	/* NULL if there's nothing but whitespace and comments */
 	if (*p == '\0')
 		return NULL;
 
+	return p;
+}
+
+/*
+ * Parse a SQL command; return a Command struct, or NULL if it's a comment
+ *
+ * On entry, psqlscan.l has collected the command into "buf", so we don't
+ * really need to do much here except check for comment and set up a
+ * Command struct.
+ */
+static Command *
+create_sql_command(PQExpBuffer buf, const char *source, int numqueries)
+{
+	Command    *my_command;
+	char	   *p = skip_sql_comments(buf->data);
+
+	if (p == NULL)
+		return NULL;
+
 	/* Allocate and initialize Command structure */
-	my_command = (Command *) pg_malloc0(sizeof(Command));
-	my_command->command_num = num_commands++;
+	my_command = (Command *) pg_malloc(sizeof(Command));
+	initPQExpBuffer(&my_command->lines);
+	appendPQExpBufferStr(&my_command->lines, p);
+	my_command->first_line = NULL;		/* this is set later */
 	my_command->type = SQL_COMMAND;
 	my_command->meta = META_NONE;
+	my_command->argc = 0;
+	memset(my_command->argv, 0, sizeof(my_command->argv));
+	my_command->nqueries = numqueries;
+	my_command->prefix_size = numqueries + 8;
+	my_command->prefix = pg_malloc0(sizeof(char *) * my_command->prefix_size);
+	my_command->expr = NULL;
 	initSimpleStats(&my_command->stats);
 
-	/*
-	 * Install query text as the sole argv string.  If we are using a
-	 * non-simple query mode, we'll extract parameters from it later.
-	 */
-	my_command->argv[0] = pg_strdup(p);
-	my_command->argc = 1;
-
-	/*
-	 * If SQL command is multi-line, we only want to save the first line as
-	 * the "line" label.
-	 */
-	nlpos = strchr(p, '\n');
-	if (nlpos)
-	{
-		my_command->line = pg_malloc(nlpos - p + 1);
-		memcpy(my_command->line, p, nlpos - p);
-		my_command->line[nlpos - p] = '\0';
-	}
-	else
-		my_command->line = pg_strdup(p);
-
 	return my_command;
 }
 
+/*
+ * 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 numqueries)
+{
+	int			nq;
+
+	Assert(my_command->type == SQL_COMMAND && my_command->lines.len > 0);
+
+	more = skip_sql_comments(more);
+	if (more == NULL)
+		return;
+
+	/* append command text, embedding a ';' in place of the \cset */
+	appendPQExpBuffer(&my_command->lines, ";\n%s", more);
+
+	/* 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;
+}
+
+/*
+ * Once an SQL command is fully parsed, possibly by accumulating several
+ * parts, complete other fields of the Command structure.
+ */
+static void
+postprocess_sql_command(Command *my_command)
+{
+	char	buffer[128];
+
+	Assert(my_command->type == SQL_COMMAND);
+
+	/* Save the first line for error display. */
+	strlcpy(buffer, my_command->lines.data, sizeof(buffer));
+	buffer[strcspn(buffer, "\n\r")] = '\0';
+	my_command->first_line = pg_strdup(buffer);
+
+	/* parse query if necessary */
+	switch (querymode)
+	{
+		case QUERY_SIMPLE:
+			my_command->argv[0] = my_command->lines.data;
+			my_command->argc++;
+			break;
+		case QUERY_EXTENDED:
+		case QUERY_PREPARED:
+			if (!parseQuery(my_command))
+				exit(1);
+			break;
+		default:
+			exit(1);
+	}
+}
+
 /*
  * Parse a backslash command; return a Command struct, or NULL if comment
  *
@@ -4177,7 +4329,6 @@ process_backslash_command(PsqlScanState sstate, const char *source)
 
 	/* Allocate and initialize Command structure */
 	my_command = (Command *) pg_malloc0(sizeof(Command));
-	my_command->command_num = num_commands++;
 	my_command->type = META_COMMAND;
 	my_command->argc = 0;
 	initSimpleStats(&my_command->stats);
@@ -4201,7 +4352,7 @@ process_backslash_command(PsqlScanState sstate, const char *source)
 		if (my_command->meta == META_SET)
 		{
 			if (!expr_lex_one_word(sstate, &word_buf, &word_offset))
-				syntax_error(source, lineno, my_command->line, my_command->argv[0],
+				syntax_error(source, lineno, my_command->first_line, my_command->argv[0],
 							 "missing argument", NULL, -1);
 
 			offsets[j] = word_offset;
@@ -4222,10 +4373,11 @@ process_backslash_command(PsqlScanState sstate, const char *source)
 		my_command->expr = expr_parse_result;
 
 		/* Save line, trimming any trailing newline */
-		my_command->line = expr_scanner_get_substring(sstate,
-													  start_offset,
-													  expr_scanner_offset(sstate),
-													  true);
+		my_command->first_line =
+			expr_scanner_get_substring(sstate,
+									   start_offset,
+									   expr_scanner_offset(sstate),
+									   true);
 
 		expr_scanner_finish(yyscanner);
 
@@ -4238,7 +4390,7 @@ process_backslash_command(PsqlScanState sstate, const char *source)
 	while (expr_lex_one_word(sstate, &word_buf, &word_offset))
 	{
 		if (j >= MAX_ARGS)
-			syntax_error(source, lineno, my_command->line, my_command->argv[0],
+			syntax_error(source, lineno, my_command->first_line, my_command->argv[0],
 						 "too many arguments", NULL, -1);
 
 		offsets[j] = word_offset;
@@ -4247,19 +4399,20 @@ process_backslash_command(PsqlScanState sstate, const char *source)
 	}
 
 	/* Save line, trimming any trailing newline */
-	my_command->line = expr_scanner_get_substring(sstate,
-												  start_offset,
-												  expr_scanner_offset(sstate),
-												  true);
+	my_command->first_line =
+		expr_scanner_get_substring(sstate,
+								   start_offset,
+								   expr_scanner_offset(sstate),
+								   true);
 
 	if (my_command->meta == META_SLEEP)
 	{
 		if (my_command->argc < 2)
-			syntax_error(source, lineno, my_command->line, my_command->argv[0],
+			syntax_error(source, lineno, my_command->first_line, my_command->argv[0],
 						 "missing argument", NULL, -1);
 
 		if (my_command->argc > 3)
-			syntax_error(source, lineno, my_command->line, my_command->argv[0],
+			syntax_error(source, lineno, my_command->first_line, my_command->argv[0],
 						 "too many arguments", NULL,
 						 offsets[3] - start_offset);
 
@@ -4288,7 +4441,7 @@ process_backslash_command(PsqlScanState sstate, const char *source)
 			if (pg_strcasecmp(my_command->argv[2], "us") != 0 &&
 				pg_strcasecmp(my_command->argv[2], "ms") != 0 &&
 				pg_strcasecmp(my_command->argv[2], "s") != 0)
-				syntax_error(source, lineno, my_command->line, my_command->argv[0],
+				syntax_error(source, lineno, my_command->first_line, my_command->argv[0],
 							 "unrecognized time unit, must be us, ms or s",
 							 my_command->argv[2], offsets[2] - start_offset);
 		}
@@ -4296,25 +4449,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);
 	}
 
@@ -4393,6 +4553,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;
@@ -4416,6 +4579,7 @@ ParseScript(const char *script, const char *desc, int weight)
 	 * stdstrings should be true, which is a bit riskier.
 	 */
 	psql_scan_setup(sstate, script, strlen(script), 0, true);
+	start_offset = expr_scanner_offset(sstate) - 1;
 
 	initPQExpBuffer(&line_buf);
 
@@ -4425,31 +4589,28 @@ ParseScript(const char *script, const char *desc, int weight)
 	{
 		PsqlScanResult sr;
 		promptStatus_t prompt;
-		Command    *command;
+		Command    *command = NULL;
+		int			semicolons;
 
 		resetPQExpBuffer(&line_buf);
+		lineno = expr_scanner_get_lineno(sstate, start_offset);
 
 		sr = psql_scan(sstate, &line_buf, &prompt);
+		semicolons = psql_scan_get_escaped_semicolons(sstate);
 
-		/* 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);
-			}
+			/* the previous multi-line command ended with \cset */
+			append_sql_command(ps.commands[index - 1], line_buf.data,
+							   semicolons + 1);
+			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, semicolons + 1);
+
+			/* store new command */
 			if (command)
 			{
 				ps.commands[index] = command;
@@ -4464,6 +4625,68 @@ 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)
+				{
+					bool		is_gset = bs_cmd[0] == 'g';
+					int			cindex;
+					Command    *sql_cmd;
+
+					is_compound = !is_gset;
+
+					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->nqueries - 1;
+
+					if (sql_cmd->prefix[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] = "";
+					else
+						sql_cmd->prefix[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;
@@ -4819,7 +5042,7 @@ printResults(TState *threads, StatsData *total, instr_time total_time,
 					printf("   %11.3f  %s\n",
 						   (cstats->count > 0) ?
 						   1000.0 * cstats->sum / cstats->count : 0.0,
-						   (*commands)->line);
+						   (*commands)->first_line);
 				}
 			}
 		}
@@ -5286,28 +5509,18 @@ main(int argc, char **argv)
 		internal_script_used = true;
 	}
 
-	/* if not simple query mode, parse the script(s) to find parameters */
-	if (querymode != QUERY_SIMPLE)
-	{
-		for (i = 0; i < num_scripts; i++)
-		{
-			Command   **commands = sql_script[i].commands;
-			int			j;
-
-			for (j = 0; commands[j] != NULL; j++)
-			{
-				if (commands[j]->type != SQL_COMMAND)
-					continue;
-				if (!parseQuery(commands[j]))
-					exit(1);
-			}
-		}
-	}
-
-	/* compute total_weight */
+	/* complete SQL command initialization and compute total weight */
 	for (i = 0; i < num_scripts; i++)
+	{
+		Command   **commands = sql_script[i].commands;
+
+		for (int j = 0; commands[j] != NULL; j++)
+			if (commands[j]->type == SQL_COMMAND)
+				postprocess_sql_command(commands[j]);
+
 		/* cannot overflow: weight is 32b, total_weight 64b */
 		total_weight += sql_script[i].weight;
+	}
 
 	if (total_weight == 0 && !is_init_mode)
 	{
diff --git a/src/bin/pgbench/t/001_pgbench_with_server.pl b/src/bin/pgbench/t/001_pgbench_with_server.pl
index 89678e7f3f..0eaf18c2f8 100644
--- a/src/bin/pgbench/t/001_pgbench_with_server.pl
+++ b/src/bin/pgbench/t/001_pgbench_with_server.pl
@@ -528,6 +528,48 @@ pgbench(
 }
 	});
 
+# working \gset and \cset
+pgbench(
+	'-t 1', 0,
+	[ qr{type: .*/001_pgbench_gset_and_cset}, qr{processed: 1/1} ],
+	[   qr{command=3.: int 0\b},
+		qr{command=5.: int 1\b},
+		qr{command=6.: int 2\b},
+		qr{command=8.: int 3\b},
+		qr{command=9.: int 4\b},
+		qr{command=10.: int 5\b},
+		qr{command=12.: int 6\b},
+		qr{command=13.: int 7\b},
+		qr{command=14.: int 8\b},
+		qr{command=16.: int 9\b} ],
+	'pgbench gset and cset commands',
+	{   '001_pgbench_gset_and_cset' => q{-- test gset and cset
+-- no columns
+SELECT \gset
+-- one value
+SELECT 0 AS i0 \gset
+\set i debug(:i0)
+-- two values
+SELECT 1 AS i1, 2 AS i2 \gset
+\set i debug(:i1)
+\set i debug(:i2)
+-- cset & gset to follow
+SELECT :i2 + 1 AS i3, :i2 * :i2 AS i4 \cset
+  SELECT 5 AS i5 \gset
+\set i debug(:i3)
+\set i debug(:i4)
+\set i debug(:i5)
+-- with prefix
+SELECT 6 AS i6, 7 AS i7 \cset x_
+  SELECT 8 AS i8 \gset y_
+\set i debug(:x_i6)
+\set i debug(:x_i7)
+\set i debug(:y_i8)
+-- overwrite existing variable
+SELECT 0 AS i9, 9 AS i9 \gset
+\set i debug(:i9)
+} });
+
 # trigger many expression errors
 my @errors = (
 
@@ -735,21 +777,45 @@ SELECT LEAST(:i, :i, :i, :i, :i, :i, :i, :i, :i, :i, :i);
 		[qr{invalid command .* "nosuchcommand"}], q{\nosuchcommand}
 	],
 	[ 'misc empty script', 1, [qr{empty command list for script}], q{} ],
-	[
-		'bad boolean',                     2,
-		[qr{malformed variable.*trueXXX}], q{\set b :badtrue or true}
-	],);
+	[   'bad boolean',                     2,
+		[qr{malformed variable.*trueXXX}], q{\set b :badtrue or true} ],
 
+	# GSET & CSET
+	[   'gset no row',                    2,
+		[qr{expecting 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
+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', 2,
+		[qr{gset/cset expects a row}],
+		q{DROP TABLE IF EXISTS no_such_table \gset} ],
+	[   'gset bad default name', 2,
+		[qr{error storing into var \?column\?}],
+		q{SELECT 1 \gset} ],
+	[   'gset bad name', 2,
+		[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;
 	$status != 0 or die "invalid expected status for test \"$name\"";
 	my $n = '001_pgbench_error_' . $name;
 	$n =~ s/ /_/g;
 	pgbench(
-		'-n -t 1 -M prepared -Dfoo=bla -Dnull=null -Dtrue=true -Done=1 -Dzero=0.0 ' .
-		'-Dbadtrue=trueXXX -Dmaxint=9223372036854775807 -Dminint=-9223372036854775808',
+		'-n -t 1 -Dfoo=bla -Dnull=null -Dtrue=true -Done=1 -Dzero=0.0 -Dbadtrue=trueXXX' .
+		' -Dmaxint=9223372036854775807 -Dminint=-9223372036854775808' .
+			($no_prepare ? '' : ' -M prepared'),
 		$status,
 		[ $status == 1 ? qr{^$} : qr{processed: 0/1} ],
 		$re,
diff --git a/src/fe_utils/psqlscan.l b/src/fe_utils/psqlscan.l
index 9ea32c319f..321744cddb 100644
--- a/src/fe_utils/psqlscan.l
+++ b/src/fe_utils/psqlscan.l
@@ -693,8 +693,15 @@ other			.
 	 * substitution.  We want these before {self}, also.
 	 */
 
-"\\"[;:]		{
-					/* Force a semicolon or colon into the query buffer */
+"\\";			{
+					/* Count semicolons in compound commands */
+					cur_state->escaped_semicolons++;
+					/* Force a semicolon into the query buffer */
+					psqlscan_emit(cur_state, yytext + 1, 1);
+				}
+
+"\\":			{
+					/* Force a colon into the query buffer */
 					psqlscan_emit(cur_state, yytext + 1, 1);
 				}
 
@@ -1065,6 +1072,9 @@ psql_scan(PsqlScanState state,
 	/* Set current output target */
 	state->output_buf = query_buf;
 
+	/* Reset number of escaped semicolons seen */
+	state->escaped_semicolons = 0;
+
 	/* Set input source */
 	if (state->buffer_stack != NULL)
 		yy_switch_to_buffer(state->buffer_stack->buf, state->scanner);
@@ -1208,6 +1218,16 @@ psql_scan_reset(PsqlScanState state)
 	state->dolqstart = NULL;
 }
 
+/*
+ * Return the number of escaped semicolons in the lexed string seen by the
+ * previous psql_scan call.
+ */
+int
+psql_scan_get_escaped_semicolons(PsqlScanState state)
+{
+	return state->escaped_semicolons;
+}
+
 /*
  * Reselect this lexer (psqlscan.l) after using another one.
  *
diff --git a/src/include/fe_utils/psqlscan.h b/src/include/fe_utils/psqlscan.h
index 1cf5b2e7fa..d6fef9ff77 100644
--- a/src/include/fe_utils/psqlscan.h
+++ b/src/include/fe_utils/psqlscan.h
@@ -90,6 +90,8 @@ extern PsqlScanResult psql_scan(PsqlScanState state,
 
 extern void psql_scan_reset(PsqlScanState state);
 
+extern int psql_scan_get_escaped_semicolons(PsqlScanState state);
+
 extern void psql_scan_reselect_sql_lexer(PsqlScanState state);
 
 extern bool psql_scan_in_quote(PsqlScanState state);
diff --git a/src/include/fe_utils/psqlscan_int.h b/src/include/fe_utils/psqlscan_int.h
index 42a738f422..752cc9406a 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			escaped_semicolons;	/* number of embedded (\;) semicolons */
 	char	   *dolqstart;		/* current $foo$ quote start string */
 
 	/*

Reply via email to