Hello Ibrar,

 SELECT 1 AS one \;
 SELECT 2 AS two UNION SELECT 2 \;
 SELECT 3 AS three \aset

will set both "one" and "three", while "two" is not set because there were
two rows. It is a kind of more permissive \gset.

Are you sure two is not set :)?

SELECT 2 AS two UNION SELECT 2;   -- only returns one row.
but
SELECT 2 AS two UNION SELECT 10;  -- returns the two rows.

Indeed, my intension was to show an example like the second.

Is this the expected behavior with \aset?

In my opinion throwing a valid error like "client 0 script 0 command 0 query 0: expected one row, got 2" make more sense.

Hmmm. My intention with \aset is really NOT to throw an error. With pgbench, the existence of the variable can be tested later to know whether it was assigned or not, eg:

  SELECT 1 AS x \;
  -- 2 rows, no assignment
  SELECT 'calvin' AS firstname UNION SELECT 'hobbes' \;
  SELECT 2 AS z \aset
  -- next test is false
  \if :{?firstname}
    ...
  \endif

The rational is that one may want to benefit from combined queries (\;)
which result in less communication thus has lower latency, but still be interested in extracting some results.

The question is what to do if the query returns 0 or >1 rows. If an error is raised, the construct cannot be used for testing whether there is one result or not, eg for a query returning 0 or 1 row, you could not write:

  \set id random(1, :number_of_users)
  SELECT firtname AS fn FROM user WHERE id = :id \aset
  \if :{?fn}
    -- the user exists, proceed with further queries
  \else
    -- no user, maybe it was removed, it is not an error
  \endif

Another option would to just assign the value so that
 - on 0 row no assignment is made, and it can be tested afterwards.
 - on >1 rows the last (first?) value is kept. I took last so to
   ensure that all results are received.

I think that having some permissive behavior allows to write some more interesting test scripts that use combined queries and extract values.

What do you think?

- With \gset

SELECT 2 AS two UNION SELECT 10 \gset
INSERT INTO test VALUES(:two,0,0);

client 0 script 0 command 0 query 0: expected one row, got 2
Run was aborted; the above results are incomplete.

Yes, that is the intented behavior.

- With \aset

SELECT 2 AS two UNION SELECT 10 \aset
INSERT INTO test VALUES(:two,0,0);
[...]
client 0 script 0 aborted in command 1 query 0: ERROR:  syntax error at or near 
":"

Indeed, the user should test whether the variable was assigned before using it if the result is not warranted to return one row.

The new status of this patch is: Waiting on Author

The attached patch implements the altered behavior described above.

--
Fabien.
diff --git a/doc/src/sgml/ref/pgbench.sgml b/doc/src/sgml/ref/pgbench.sgml
index e3b73a4cf5..52949d9a35 100644
--- a/doc/src/sgml/ref/pgbench.sgml
+++ b/doc/src/sgml/ref/pgbench.sgml
@@ -966,18 +966,29 @@ pgbench <optional> <replaceable>options</replaceable> </optional> <replaceable>d
    <varlistentry id='pgbench-metacommand-gset'>
     <term>
      <literal>\gset [<replaceable>prefix</replaceable>]</literal>
+     <literal>\aset [<replaceable>prefix</replaceable>]</literal>
     </term>
 
     <listitem>
      <para>
-      This command may be used to end SQL queries, taking the place of the
+      These commands may be used to end SQL queries, taking the place of the
       terminating 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.
+      When the <literal>\gset</literal> 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>
+      When the <literal>\aset</literal> command is used, all combined SQL queries
+      (separated by <literal>\;</literal>) have their columns stored into variables
+      named after column names, and prefixed with <replaceable>prefix</replaceable>
+      if provided. If a query returns no row, no assignment is made and the variable
+      can be tested for existence to detect this. If a query returns more than one
+      row, the last value is kept.
      </para>
 
      <para>
@@ -986,6 +997,8 @@ pgbench <optional> <replaceable>options</replaceable> </optional> <replaceable>d
       <replaceable>p_two</replaceable> and <replaceable>p_three</replaceable>
       with integers from the third query.
       The result of the second query is discarded.
+      The result of the two last combined queries are stored in variables
+      <replaceable>four</replaceable> and <replaceable>five</replaceable>.
 <programlisting>
 UPDATE pgbench_accounts
   SET abalance = abalance + :delta
@@ -994,6 +1007,7 @@ UPDATE pgbench_accounts
 -- compound of two queries
 SELECT 1 \;
 SELECT 2 AS two, 3 AS three \gset p_
+SELECT 4 AS four \; SELECT 5 AS five \aset
 </programlisting>
      </para>
     </listitem>
diff --git a/src/bin/pgbench/pgbench.c b/src/bin/pgbench/pgbench.c
index 8b84658ccd..477b6c873b 100644
--- a/src/bin/pgbench/pgbench.c
+++ b/src/bin/pgbench/pgbench.c
@@ -461,6 +461,7 @@ typedef enum MetaCommand
 	META_SHELL,					/* \shell */
 	META_SLEEP,					/* \sleep */
 	META_GSET,					/* \gset */
+	META_ASET,					/* \aset */
 	META_IF,					/* \if */
 	META_ELIF,					/* \elif */
 	META_ELSE,					/* \else */
@@ -493,6 +494,7 @@ static const char *QUERYMODE[] = {"simple", "extended", "prepared"};
  * varprefix 	SQL commands terminated with \gset have this set
  *				to a non NULL value.  If nonempty, it's used to prefix the
  *				variable name that receives the value.
+ * aset			do gset on all possible queries of a combined query (\;).
  * expr			Parsed expression, if needed.
  * stats		Time spent in this command.
  */
@@ -505,6 +507,7 @@ typedef struct Command
 	int			argc;
 	char	   *argv[MAX_ARGS];
 	char	   *varprefix;
+	bool		aset;
 	PgBenchExpr *expr;
 	SimpleStats stats;
 } Command;
@@ -2479,6 +2482,8 @@ getMetaCommand(const char *cmd)
 		mc = META_ENDIF;
 	else if (pg_strcasecmp(cmd, "gset") == 0)
 		mc = META_GSET;
+	else if (pg_strcasecmp(cmd, "aset") == 0)
+		mc = META_ASET;
 	else
 		mc = META_NONE;
 	return mc;
@@ -2710,12 +2715,12 @@ sendCommand(CState *st, Command *command)
  * Process query response from the backend.
  *
  * If varprefix is not NULL, it's the variable name prefix where to store
- * the results of the *last* command.
+ * the results of the *last* command (gset) or *all* commands (aset).
  *
  * Returns true if everything is A-OK, false if any error occurs.
  */
 static bool
-readCommandResponse(CState *st, char *varprefix)
+readCommandResponse(CState *st, char *varprefix, bool is_aset)
 {
 	PGresult   *res;
 	PGresult   *next_res;
@@ -2735,7 +2740,7 @@ readCommandResponse(CState *st, char *varprefix)
 		{
 			case PGRES_COMMAND_OK:	/* non-SELECT commands */
 			case PGRES_EMPTY_QUERY: /* may be used for testing no-op overhead */
-				if (is_last && varprefix != NULL)
+				if (is_last && varprefix != NULL && !is_aset)
 				{
 					fprintf(stderr,
 							"client %d script %d command %d query %d: expected one row, got %d\n",
@@ -2745,16 +2750,23 @@ readCommandResponse(CState *st, char *varprefix)
 				break;
 
 			case PGRES_TUPLES_OK:
-				if (is_last && varprefix != NULL)
+				if (varprefix != NULL && (is_last || is_aset))
 				{
-					if (PQntuples(res) != 1)
+					int ntuples		= PQntuples(res);
+
+					if (!is_aset && ntuples != 1)
 					{
+						/* under \gset, report the error */
 						fprintf(stderr,
 								"client %d script %d command %d query %d: expected one row, got %d\n",
 								st->id, st->use_file, st->command, qrynum, PQntuples(res));
 						goto error;
 					}
 
+					/* skip empty result under \aset */
+					if (ntuples <= 0)
+						break;
+
 					/* store results into variables */
 					for (int fld = 0; fld < PQnfields(res); fld++)
 					{
@@ -2764,15 +2776,14 @@ readCommandResponse(CState *st, char *varprefix)
 						if (*varprefix != '\0')
 							varname = psprintf("%s%s", varprefix, varname);
 
-						/* store result as a string */
-						if (!putVariable(st, "gset", varname,
-										 PQgetvalue(res, 0, fld)))
+						/* store last row result as a string */
+						if (!putVariable(st, is_aset ? "aset" : "gset", varname,
+										 PQgetvalue(res, ntuples - 1, fld)))
 						{
 							/* internal error */
 							fprintf(stderr,
 									"client %d script %d command %d query %d: error storing into variable %s\n",
-									st->id, st->use_file, st->command, qrynum,
-									varname);
+									st->id, st->use_file, st->command, qrynum, varname);
 							goto error;
 						}
 
@@ -3189,7 +3200,9 @@ advanceConnectionState(TState *thread, CState *st, StatsData *agg)
 					return;		/* don't have the whole result yet */
 
 				/* store or discard the query results */
-				if (readCommandResponse(st, sql_script[st->use_file].commands[st->command]->varprefix))
+				if (readCommandResponse(st,
+										sql_script[st->use_file].commands[st->command]->varprefix,
+										sql_script[st->use_file].commands[st->command]->aset))
 					st->state = CSTATE_END_COMMAND;
 				else
 					st->state = CSTATE_ABORTED;
@@ -4125,6 +4138,7 @@ create_sql_command(PQExpBuffer buf, const char *source)
 	my_command->argc = 0;
 	memset(my_command->argv, 0, sizeof(my_command->argv));
 	my_command->varprefix = NULL;	/* allocated later, if needed */
+	my_command->aset = false;
 	my_command->expr = NULL;
 	initSimpleStats(&my_command->stats);
 
@@ -4353,7 +4367,7 @@ process_backslash_command(PsqlScanState sstate, const char *source)
 			syntax_error(source, lineno, my_command->first_line, my_command->argv[0],
 						 "unexpected argument", NULL, -1);
 	}
-	else if (my_command->meta == META_GSET)
+	else if (my_command->meta == META_GSET || my_command->meta == META_ASET)
 	{
 		if (my_command->argc > 2)
 			syntax_error(source, lineno, my_command->first_line, my_command->argv[0],
@@ -4498,10 +4512,10 @@ ParseScript(const char *script, const char *desc, int weight)
 			if (command)
 			{
 				/*
-				 * If this is gset, merge into the preceding command. (We
-				 * don't use a command slot in this case).
+				 * If this is gset or aset, merge into the preceding command.
+				 * (We don't use a command slot in this case).
 				 */
-				if (command->meta == META_GSET)
+				if (command->meta == META_GSET || command->meta == META_ASET)
 				{
 					Command    *cmd;
 
@@ -4523,6 +4537,7 @@ ParseScript(const char *script, const char *desc, int weight)
 						cmd->varprefix = pg_strdup("");
 					else
 						cmd->varprefix = pg_strdup(command->argv[1]);
+					cmd->aset = (command->meta == META_ASET);
 
 					/* cleanup unused command */
 					free_command(command);
diff --git a/src/bin/pgbench/t/001_pgbench_with_server.pl b/src/bin/pgbench/t/001_pgbench_with_server.pl
index dc2c72fa92..62cc0bb1ac 100644
--- a/src/bin/pgbench/t/001_pgbench_with_server.pl
+++ b/src/bin/pgbench/t/001_pgbench_with_server.pl
@@ -538,7 +538,7 @@ pgbench(
 }
 	});
 
-# working \gset
+# working \gset and \aset
 pgbench(
 	'-t 1', 0,
 	[ qr{type: .*/001_pgbench_gset}, qr{processed: 1/1} ],
@@ -548,9 +548,11 @@ pgbench(
 		qr{command=6.: int 2\b},
 		qr{command=8.: int 3\b},
 		qr{command=10.: int 4\b},
-		qr{command=12.: int 5\b}
+		qr{command=12.: int 5\b},
+		qr{command=15.: int 8\b},
+		qr{command=16.: int 7\b}
 	],
-	'pgbench gset command',
+	'pgbench gset & aset commands',
 	{
 		'001_pgbench_gset' => q{-- test gset
 -- no columns
@@ -571,6 +573,12 @@ SELECT 0 AS i4, 4 AS i4 \gset
 -- work on the last SQL command under \;
 \; \; SELECT 0 AS i5 \; SELECT 5 AS i5 \; \; \gset
 \set i debug(:i5)
+-- test aset, which applies to a combined query
+\; SELECT 6 AS i6 \; SELECT 7 AS i7 \; \aset
+-- unless it returns more than one row, last is kept
+SELECT 8 AS i6 UNION SELECT 9 ORDER BY 1 DESC \aset
+\set i debug(:i6)
+\set i debug(:i7)
 }
 	});
 

Reply via email to