pgbench: add \cset and \gset commands

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

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

Authors: Fabien Coelho, Álvaro Herrera
Reviewed-by: Amit Langote <Langote_Amit_f8@lab.ntt.co.jp>
Reviewed-by: Stephen Frost <sfrost@snowman.net>
Reviewed-by: Pavel Stehule <pavel.stehule@gmail.com>
Reviewed-by: Tom Lane <tgl@sss.pgh.pa.us>
Reviewed-by: Tatsuo Ishii <ishii@sraoss.co.jp>
Reviewed-by: Rafia Sabih <rafia.sabih@enterprisedb.com>
Discussion: https://postgr.es/m/alpine.DEB.2.20.1607091005330.3412@sto

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 command may be used to end SQL queries, replacing a final semicolon
+      (<literal>;</literal>).
+     </para>
+
+     <para>
+      When this command is used, the preceding SQL query is expected to
+      return one row, the columns of which are stored into variables named after
+      column names, and prefixed with <replaceable>prefix</replaceable> if provided.
+     </para>
+
+     <para>
+      The following example puts the final account balance from the first query
+      into variable <replaceable>abalance</replaceable>, and fills variables
+      <replaceable>p_two</replaceable> and <replaceable>p_three</replaceable>
+      with integers from the last query.
+      The result of the second query is discarded.
+<programlisting>
+UPDATE pgbench_accounts
+  SET abalance = abalance + :delta
+  WHERE aid = :aid
+  RETURNING abalance \gset
+-- compound of two queries
+SELECT 1 \;
+SELECT 2 AS two, 3 AS three \gset p_
+</programlisting>
+     </para>
+
+     <note>
+      <para>
+        <literal>\gset</literal> does not work when empty SQL queries appear
+        within a compound SQL command.
+      </para>
+     </note>
+    </listitem>
+   </varlistentry>
+
    <varlistentry>
     <term><literal>\if</literal> <replaceable class="parameter">expression</replaceable></term>
     <term><literal>\elif</literal> <replaceable class="parameter">expression</replaceable></term>

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{expected one row, got 0\b}], q{SELECT WHERE FALSE \gset} ],
+	[   'cset no row',                    2,
+		[qr{expected one row, got 0\b}], q{SELECT WHERE FALSE \cset
+SELECT 1 AS i\gset}, 1 ],
+	[ 'gset alone', 1, [qr{gset/cset cannot start a script}], q{\gset} ],
+	[   'gset no SQL',                        1,
+		[qr{gset/cset must follow a SQL command}], q{\set i +1
+\gset} ],
+	[   'gset too many arguments',                   1,
+		[qr{too many arguments}], q{SELECT 1 \gset a b} ],
+	[   'gset after gset',                        1,
+	    [qr{gset/cset cannot follow one another}], q{SELECT 1 AS i \gset
+\gset} ],
+	[   'gset non SELECT', 2,
+		[qr{expected one row, got 0}],
+		q{DROP TABLE IF EXISTS no_such_table \gset} ],
+	[   'gset bad default name', 2,
+		[qr{error storing into variable \?column\?}],
+		q{SELECT 1 \gset} ],
+	[   'gset bad name', 2,
+		[qr{error storing into variable bad name!}],
+		q{SELECT 1 AS "bad name!" \gset} ],
+	);
 
 for my $e (@errors)
 {
-	my ($name, $status, $re, $script) = @$e;
+	my ($name, $status, $re, $script, $no_prepare) = @$e;
 	$status != 0 or die "invalid expected status for test \"$name\"";
 	my $n = '001_pgbench_error_' . $name;
 	$n =~ s/ /_/g;
 	pgbench(
-		'-n -t 1 -M prepared -Dfoo=bla -Dnull=null -Dtrue=true -Done=1 -Dzero=0.0 ' .
-		'-Dbadtrue=trueXXX -Dmaxint=9223372036854775807 -Dminint=-9223372036854775808',
+		'-n -t 1 -Dfoo=bla -Dnull=null -Dtrue=true -Done=1 -Dzero=0.0 -Dbadtrue=trueXXX' .
+		' -Dmaxint=9223372036854775807 -Dminint=-9223372036854775808' .
+			($no_prepare ? '' : ' -M prepared'),
 		$status,
 		[ $status == 1 ? qr{^$} : qr{processed: 0/1} ],
 		$re,

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.
  *

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);

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 */
 
 	/*